Copy Laneformer repo with BF16 weights

.eval_results/humaneval.yaml

@@ -0,0 +1,10 @@
+- dataset:
+    id: openai/openai_humaneval
+    task_id: default
+  value: 53.05
+  date: "2026-06-11"
+- dataset:
+    id: evalplus/humanevalplus
+    task_id: default
+  value: 50.00
+  date: "2026-06-11"

LICENSE

@@ -0,0 +1,222 @@
+Kog LaneformerTP Repository License Notice
+=========================================
+
+This repository contains materials under different licenses.
+
+1. Kog-owned materials
+----------------------
+Unless a file states otherwise, the model weights, Hugging Face custom modeling and configuration code, model configuration files, metadata files, documentation, and README/model card are licensed under the Apache License, Version 2.0 below.
+
+Apache License
+Version 2.0, January 2004
+http://www.apache.org/licenses/
+
+TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+1. Definitions.
+
+"License" shall mean the terms and conditions for use, reproduction,
+and distribution as defined by Sections 1 through 9 of this document.
+
+"Licensor" shall mean the copyright owner or entity authorized by
+the copyright owner that is granting the License.
+
+"Legal Entity" shall mean the union of the acting entity and all
+other entities that control, are controlled by, or are under common
+control with that entity. For the purposes of this definition,
+"control" means (i) the power, direct or indirect, to cause the
+direction or management of such entity, whether by contract or
+otherwise, or (ii) ownership of fifty percent (50%) or more of the
+outstanding shares, or (iii) beneficial ownership of such entity.
+
+"You" (or "Your") shall mean an individual or Legal Entity
+exercising permissions granted by this License.
+
+"Source" form shall mean the preferred form for making modifications,
+including but not limited to software source code, documentation
+source, and configuration files.
+
+"Object" form shall mean any form resulting from mechanical
+transformation or translation of a Source form, including but
+not limited to compiled object code, generated documentation,
+and conversions to other media types.
+
+"Work" shall mean the work of authorship, whether in Source or
+Object form, made available under the License, as indicated by a
+copyright notice that is included in or attached to the work
+(an example is provided in the Appendix below).
+
+"Derivative Works" shall mean any work, whether in Source or Object
+form, that is based on (or derived from) the Work and for which the
+editorial revisions, annotations, elaborations, or other modifications
+represent, as a whole, an original work of authorship. For the purposes
+of this License, Derivative Works shall not include works that remain
+separable from, or merely link (or bind by name) to the interfaces of,
+the Work and Derivative Works thereof.
+
+"Contribution" shall mean any work of authorship, including
+the original version of the Work and any modifications or additions
+to that Work or Derivative Works thereof, that is intentionally
+submitted to Licensor for inclusion in the Work by the copyright owner
+or by an individual or Legal Entity authorized to submit on behalf of
+the copyright owner. For the purposes of this definition, "submitted"
+means any form of electronic, verbal, or written communication sent
+to the Licensor or its representatives, including but not limited to
+communication on electronic mailing lists, source code control systems,
+and issue tracking systems that are managed by, or on behalf of, the
+Licensor for the purpose of discussing and improving the Work, but
+excluding communication that is conspicuously marked or otherwise
+designated in writing by the copyright owner as "Not a Contribution."
+
+"Contributor" shall mean Licensor and any individual or Legal Entity
+on behalf of whom a Contribution has been received by Licensor and
+subsequently incorporated within the Work.
+
+2. Grant of Copyright License. Subject to the terms and conditions of
+this License, each Contributor hereby grants to You a perpetual,
+worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+copyright license to reproduce, prepare Derivative Works of,
+publicly display, publicly perform, sublicense, and distribute the
+Work and such Derivative Works in Source or Object form.
+
+3. Grant of Patent License. Subject to the terms and conditions of
+this License, each Contributor hereby grants to You a perpetual,
+worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+(except as stated in this section) patent license to make, have made,
+use, offer to sell, sell, import, and otherwise transfer the Work,
+where such license applies only to those patent claims licensable
+by such Contributor that are necessarily infringed by their
+Contribution(s) alone or by combination of their Contribution(s)
+with the Work to which such Contribution(s) was submitted. If You
+institute patent litigation against any entity (including a
+cross-claim or counterclaim in a lawsuit) alleging that the Work
+or a Contribution incorporated within the Work constitutes direct
+or contributory patent infringement, then any patent licenses
+granted to You under this License for that Work shall terminate
+as of the date such litigation is filed.
+
+4. Redistribution. You may reproduce and distribute copies of the
+Work or Derivative Works thereof in any medium, with or without
+modifications, and in Source or Object form, provided that You
+meet the following conditions:
+
+(a) You must give any other recipients of the Work or
+    Derivative Works a copy of this License; and
+
+(b) You must cause any modified files to carry prominent notices
+    stating that You changed the files; and
+
+(c) You must retain, in the Source form of any Derivative Works
+    that You distribute, all copyright, patent, trademark, and
+    attribution notices from the Source form of the Work,
+    excluding those notices that do not pertain to any part of
+    the Derivative Works; and
+
+(d) If the Work includes a "NOTICE" text file as part of its
+    distribution, then any Derivative Works that You distribute must
+    include a readable copy of the attribution notices contained
+    within such NOTICE file, excluding those notices that do not
+    pertain to any part of the Derivative Works, in at least one
+    of the following places: within a NOTICE text file distributed
+    as part of the Derivative Works; within the Source form or
+    documentation, if provided along with the Derivative Works; or,
+    within a display generated by the Derivative Works, if and
+    wherever such third-party notices normally appear. The contents
+    of the NOTICE file are for informational purposes only and
+    do not modify the License. You may add Your own attribution
+    notices within Derivative Works that You distribute, alongside
+    or as an addendum to the NOTICE text from the Work, provided
+    that such additional attribution notices cannot be construed
+    as modifying the License.
+
+You may add Your own copyright statement to Your modifications and
+may provide additional or different license terms and conditions
+for use, reproduction, or distribution of Your modifications, or
+for any such Derivative Works as a whole, provided Your use,
+reproduction, and distribution of the Work otherwise complies with
+the conditions stated in this License.
+
+5. Submission of Contributions. Unless You explicitly state otherwise,
+any Contribution intentionally submitted for inclusion in the Work
+by You to the Licensor shall be under the terms and conditions of
+this License, without any additional terms or conditions.
+Notwithstanding the above, nothing herein shall supersede or modify
+the terms of any separate license agreement you may have executed
+with Licensor regarding such Contributions.
+
+6. Trademarks. This License does not grant permission to use the trade
+names, trademarks, service marks, or product names of the Licensor,
+except as required for reasonable and customary use in describing the
+origin of the Work and reproducing the content of the NOTICE file.
+
+7. Disclaimer of Warranty. Unless required by applicable law or
+agreed to in writing, Licensor provides the Work (and each
+Contributor provides its Contributions) on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+implied, including, without limitation, any warranties or conditions
+of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+PARTICULAR PURPOSE. You are solely responsible for determining the
+appropriateness of using or redistributing the Work and assume any
+risks associated with Your exercise of permissions under this License.
+
+8. Limitation of Liability. In no event and under no legal theory,
+whether in tort (including negligence), contract, or otherwise,
+unless required by applicable law (such as deliberate and grossly
+negligent acts) or agreed to in writing, shall any Contributor be
+liable to You for damages, including any direct, indirect, special,
+incidental, or consequential damages of any character arising as a
+result of this License or out of the use or inability to use the
+Work (including but not limited to damages for loss of goodwill,
+work stoppage, computer failure or malfunction, or any and all
+other commercial damages or losses), even if such Contributor
+has been advised of the possibility of such damages.
+
+9. Accepting Warranty or Additional Liability. While redistributing
+the Work or Derivative Works thereof, You may choose to offer,
+and charge a fee for, acceptance of support, warranty, indemnity,
+or other liability obligations and/or rights consistent with this
+License. However, in accepting such obligations, You may act only
+on Your own behalf and on Your sole responsibility, not on behalf
+of any other Contributor, and only if You agree to indemnify,
+defend, and hold each Contributor harmless for any liability
+incurred by, or claims asserted against, such Contributor by reason
+of your accepting any such warranty or additional liability.
+
+END OF TERMS AND CONDITIONS
+
+APPENDIX: How to apply the Apache License to your work.
+
+To apply the Apache License to your work, attach the following
+boilerplate notice, with the fields enclosed by brackets "[]"
+replaced with your own identifying information. (Do not include
+the brackets!)  The text should be enclosed in the appropriate
+comment syntax for the file format. We also recommend that a
+file or class name and description of purpose be included on the
+same "printed page" as the copyright notice for easier
+identification within third-party archives.
+
+Copyright 2026 Kog AI
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+
+2. Third-party tokenizer materials
+----------------------------------
+The tokenizer files in this repository are based on the Llama 2 tokenizer and are not licensed under the Apache License, Version 2.0 above. These tokenizer materials include, without limitation, tokenizer.model, tokenizer.json, tokenizer_config.json, special_tokens_map.json, and any other file whose purpose is to reproduce or configure the Llama 2 tokenizer.
+
+The Llama 2 tokenizer materials are distributed under the LLAMA 2 Community License Agreement. See THIRD_PARTY_LICENSES/LLAMA2_LICENSE for the full license text.
+
+Required attribution notice for Llama 2 materials:
+
+"Llama 2 is licensed under the LLAMA 2 Community License, Copyright (c) Meta Platforms, Inc. All Rights Reserved."
+
+Users and redistributors are responsible for complying with the LLAMA 2 Community License Agreement, including the Acceptable Use Policy referenced by that agreement.

NOTICE

@@ -0,0 +1,11 @@
+Llama 2 tokenizer notice
+
+This repository includes tokenizer artifacts based on the Llama 2 tokenizer.
+
+Llama 2 is licensed under the LLAMA 2 Community License,
+Copyright (c) Meta Platforms, Inc. All Rights Reserved.
+
+The Llama 2 tokenizer materials are distributed under the LLAMA 2 Community License Agreement.
+See THIRD_PARTY_LICENSES/LLAMA2_LICENSE for the full license text and consult Meta's Llama 2 Acceptable Use Policy referenced by that license.
+
+All Kog-owned model weights, Hugging Face custom modeling/configuration code, documentation, and configuration files are licensed separately under the Apache License, Version 2.0 unless a file states otherwise.

README.md

@@ -0,0 +1,395 @@
+---
+library_name: transformers
+pipeline_tag: text-generation
+license: other
+license_name: kog-apache-2.0-plus-llama2-tokenizer
+license_link: https://huggingface.co/kogai/laneformer-2b-it/blob/main/LICENSE
+# The repository is multi-license: Kog-owned model/code/docs are Apache-2.0;
+# tokenizer artifacts based on the Llama 2 tokenizer are under the Llama 2 Community License.
+tags:
+  - pytorch
+  - causal-lm
+  - text-generation
+  - laneformer
+  - delayed-tensor-parallelism
+  - dtp
+language:
+  - en
+datasets:
+  - nvidia/Nemotron-Pretraining-Dataset-sample
+  - nvidia/Nemotron-Pretraining-Legal-v1
+  - nvidia/Nemotron-Pretraining-Specialized-v1.2
+  - nvidia/Nemotron-Pretraining-Code-v3
+  - nvidia/Nemotron-Pretraining-Specialized-v1.1
+  - nvidia/Nemotron-CC-Code-v1
+  - nvidia/Nemotron-CC-v2.1
+  - nvidia/Nemotron-Pretraining-Code-v2
+  - nvidia/Nemotron-Pretraining-Specialized-v1
+  - nvidia/Nemotron-CC-Math-v1
+  - nvidia/Nemotron-CC-v2
+  - nvidia/Nemotron-Pretraining-SFT-v1
+  - nvidia/Nemotron-Pretraining-Code-v1
+model_type: laneformer_tp
+parameters: 2B
+widget:
+  - text: "Explain the A* algorithm in details, then write a full Python implementation from scratch."
+  - text: "Explain the Union-Find / Disjoint Set data structure in detail, then write a complete implementation with path compression and union by rank."
+  - text: "Explain how an LRU cache works, then write a full implementation in Python without using built-in cache libraries"
+  - text: "Explain how a Trie works, then write a complete implementation supporting insert, search, and prefix queries."
+  - text: "Explain topological sorting in detail, then write a full implementation using both DFS and Kahn's algorithm."
+  - text: "Explain how a binary heap works, then write a complete min-heap implementation from scratch in Python."
+  - text: "Explain Dijkstra's algorithm in detail, then write a full Python implementation from scratch."
+---
+
+# Kog Laneformer 2B
+
+Kog Laneformer 2B is a 2B-class decoder-only causal language model built around **Laneformer** and **Delayed Tensor Parallelism (DTP)**, a latency-oriented Transformer variant designed to overlap tensor-parallel communication with useful computation and weight streaming.
+
+The model is intended to make our Laneformer architecture available in the Hugging Face ecosystem for research, inspection, fine-tuning experiments, and compatibility testing. It is also the architecture family used in our public inference-engine preview, where single-request decoding speeds is **3,000 output tokens/s/request on 8x AMD MI300X** and **2,100 output tokens/s/request on 8x NVIDIA H200** in FP16 without speculative decoding.
+
+Those numbers are **Kog Inference Engine (KIE) benchmark results**, not expected performance from the generic Transformers reference path.
+
+For background, see:
+
+- [Real-time LLM Inference on Standard Datacenter GPUs (3,000 tokens/s per request)](https://blog.kog.ai/real-time-llm-inference-on-standard-gpus-3-000-tokens-s-per-request/)
+- [Delayed Tensor Parallelism for Faster Transformer Inference](https://blog.kog.ai/delayed-tensor-parallelism-for-faster-transformer-inference/)
+
+> **Custom code notice:** this model architecture is not currently part of upstream Transformers. Loading with `AutoConfig`, `AutoModel`, or `AutoModelForCausalLM` requires `trust_remote_code=True`. For production or security-sensitive use, review the modeling code and pin a specific Hub commit hash with `revision="<commit-hash>"`.
+
+## Why this model exists
+
+Many LLM serving systems optimize aggregate throughput across many concurrent requests. Our Laneformer DTP work targets a different regime: **low-batch, single-request decode speed**, which is important for agentic coding loops, real-time copilots, voice assistants, and other sequential workflows where each generated step gates the next one.
+
+At batch size 1, autoregressive decoding is often constrained by memory bandwidth, synchronization, kernel launch overheads, and communication latency rather than raw FLOPs. Laneformer uses a lane-structured architecture and delayed communication pattern so tensor-parallel work can be organized around the latency structure of full-node GPU inference.
+
+## Model details
+
+| Field | Value |
+|---|---:|
+| Model family | Laneformer |
+| HF model type | `laneformer_tp` |
+| Architecture class | `LaneformerTPForCausalLM` |
+| Task | Decoder-only causal language modeling |
+| Parameters | ~2.3B |
+| Hidden size | 3072 |
+| Layers | 15 |
+| Attention heads | 32 |
+| KV heads | 16 |
+| Intermediate size | 12288 |
+| Context length | 4096 tokens |
+| Vocabulary size | 32000 |
+| Tokenizer | Llama 2 tokenizer |
+| Number of lanes | 8 |
+| Broadcast delay | 2 |
+| LM head type | `vocab_parallel` |
+| Sliding window | 2048 tokens |
+| Sliding-window layers | 0-9 |
+| Full-attention layers | 10-14 |
+| RoPE theta | 10000 |
+| Tied word embeddings | false |
+| Expected Transformers version | `>=4.57.1` |
+| Published weight dtype | BFloat16 |
+
+## Architecture overview
+
+Laneformer follows a Llama-style decoder-only architecture, but restructures tensor-parallel communication around **lanes**. In standard tensor parallelism, each attention or MLP module typically requires communication before downstream computation can proceed. In Delayed Tensor Parallelism, local outputs are communicated asynchronously and consumed several modules later, allowing communication latency to be hidden behind subsequent computation and weight streaming.
+
+This release candidate uses:
+
+- 8 lanes.
+- Broadcast delay of 2.
+- Grouped-query attention with 32 query heads and 16 KV heads.
+- RoPE positional embeddings.
+- 10 sliding-window attention layers followed by 5 full-attention layers.
+- A vocab-parallel language-modeling head.
+- Llama 2 tokenizer vocabulary and special-token convention.
+
+The Hugging Face implementation is a reference/compatibility implementation. It is useful for loading, generation, inspection, and downstream experimentation, but it is not the same execution path as our production inference engine, which uses a latency-optimized runtime and low-level GPU kernels optimization.
+
+## Installation
+
+### Nvidia
+
+```bash
+uv venv --python 3.12
+source .venv/bin/activate
+uv pip install torch --index-url https://download.pytorch.org/whl/cu130
+uv pip install "transformers>=4.57.1" accelerate safetensors sentencepiece protobuf
+```
+
+### AMD
+
+```bash
+uv venv --python 3.12
+source .venv/bin/activate
+uv pip install torch --index-url https://download.pytorch.org/whl/rocm7.2
+uv pip install "transformers>=4.57.1" accelerate safetensors sentencepiece protobuf
+```
+
+## Usage
+
+```python
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, set_seed
+
+model_id = "kogai/laneformer-2b-it"
+revision = "main"
+seed = 42
+temperature = 0.8
+
+set_seed(seed)
+print(f"Using random seed {seed}.")
+
+print(f"Loading tokenizer for {model_id} ({revision=})...")
+tokenizer = AutoTokenizer.from_pretrained(
+    model_id,
+    trust_remote_code=True,
+    revision=revision,
+)
+print("Tokenizer loaded.")
+
+print("Loading model. This may take a few minutes on first run while files download...")
+model = AutoModelForCausalLM.from_pretrained(
+    model_id,
+    trust_remote_code=True,
+    revision=revision,
+    dtype="bfloat16",
+    device_map="auto",
+)
+print(f"Model loaded on {model.device}.")
+
+# The Llama 2 tokenizer has no native pad token. If this repo sets PAD to EOS,
+# always pass attention_mask for batched generation.
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+    print("Tokenizer had no pad token; using EOS as the pad token.")
+
+prompts = [
+    "Explain how a binary heap works, then write a complete min-heap implementation from scratch in Python.",
+]
+
+print("\nPrompt:")
+print(prompts[0])
+print("\nTokenizing prompt...")
+inputs = tokenizer(
+    prompts,
+    return_tensors="pt",
+    padding=True,
+    truncation=True,
+    max_length=4096,
+).to(model.device)
+print(f"Prompt token count: {inputs.input_ids.shape[-1]}")
+
+print(f"Generating response with sampling enabled ({temperature=})...")
+with torch.no_grad():
+    output_ids = model.generate(
+        **inputs,
+        max_new_tokens=1024,
+        do_sample=True,
+        temperature=temperature,
+        use_cache=True,
+        pad_token_id=tokenizer.pad_token_id,
+        eos_token_id=tokenizer.eos_token_id,
+    )
+generated = output_ids[:, inputs.input_ids.shape[-1]:]
+
+print("\nGenerated response:")
+print("-" * 80)
+print(tokenizer.batch_decode(generated, skip_special_tokens=True)[0])
+```
+
+### Safer custom-code loading
+
+Because this model uses remote custom code, pin a reviewed commit hash for reproducible and safer usage:
+
+```python
+commit_hash = "<reviewed-commit-hash>"
+model = AutoModelForCausalLM.from_pretrained(
+    "kogai/laneformer-2b-it",
+    trust_remote_code=True,
+    revision=commit_hash,
+    dtype="auto",
+    device_map="auto",
+)
+```
+
+## Tokenizer
+
+This repository uses the **Llama 2 tokenizer** with a 32,000-token vocabulary. The expected special-token convention is:
+
+| Token | ID |
+|---|---:|
+| `<unk>` | 0 |
+| `<s>` | 1 |
+| `</s>` | 2 |
+
+The model config uses `eos_token_id = 2` and `pad_token_id = 2`.
+
+## Training data and procedure
+
+This checkpoint was trained with a staged language-modeling recipe designed to first learn broad language/modeling capabilities and then continue training on a code- and reasoning-focused mixture. The TorchTitan training configs use sequence length 4096 and global batch size 1536, or 6,291,456 tokens per optimizer step.
+
+### Training stages
+
+| Stage | Focus | Framework | Tokens | Steps | Approx. tokens / step | Notes |
+|---|---|---|---:|---:|---:|---|
+| Pre-training | General language-model pre-training | TorchTitan | ~4T | 620,000 | ~6.29M | NVIDIA Nemotron pre-training datasets; final checkpoint selected. |
+| Mid-training | Code- and reasoning-focused continuation training | TorchTitan | ~2T | 310,000 | ~6.29M | NVIDIA Nemotron pre-training datasets; initialized from the Pre-training checkpoint. |
+| Post-training | SFT, instruction tuning | Hugging Face Transformers | ~210M | 200 | ~1.05M | Custom data mixture; Initialized from the mid-training checkpoint. |
+
+### Data sources
+
+Both TorchTitan stages depend primarily on [NVIDIA Nemotron pre-training datasets](https://huggingface.co/collections/nvidia/nemotron-pre-training-datasets). Pre-training uses a Nemotron-CC-v2-centered mixture for broad language-model pretraining, with web, synthetic web, QA, math, SFT-style, and code components. Mid-training continues from the pre-training checkpoint on a more code- and reasoning-heavy Nemotron-derived mixture, increasing the emphasis on code metadata, synthetic code data, Common Crawl code pages, math, STEM, and reasoning data.
+
+### Training infrastructure
+
+| Field | Value |
+|---|---|
+| Pre-training framework | TorchTitan |
+| Post-training framework | Hugging Face Transformers |
+| Training hardware | 24 nodes × 8 NVIDIA H100 GPUs = 192 H100 GPUs |
+| Training clusters | Scaleway cluster; ADASTRA cluster |
+| Cloud/provider/region | France |
+| Training time | ~21 days |
+| Precision | FP32/BF16 mixed precision |
+| Parallelism | Data-parallel |
+| Optimizer | AdamW |
+| Learning-rate schedule | WSD |
+| Sequence length during training | 4096 |
+| Checkpoint selection | final step |
+
+## Inference and performance
+
+Kog Inference Engine running the Laneformer 2B model public-preview numbers:
+
+| Setting | Reported speed | Notes |
+|---|---:|---|
+| 8x AMD MI300X | 3,000 output tokens/s/request | FP16, batch size 1, no speculative decoding |
+| 8x NVIDIA H200 | 2,100 output tokens/s/request | FP16, batch size 1, no speculative decoding |
+
+These benchmark numbers are for Kog's optimized inference stack, not the plain Hugging Face Transformers implementation in this repository. This preview does not rely on quantization, speculative decoding, pruning, early exit, or KV-cache compression to reach this speed.
+
+## Evaluation
+
+The following internal Kog evaluations were run in June 2026 with `lm-evaluation-harness` 0.4.12 against a local Kog serving endpoint for `Kog Laneformer 2B` @ batch size 1 in FP16. These numbers evaluate the served preview model.
+
+All values are percentages.
+
+| Benchmark | Metric | Value | Samples | Setting |
+|---|---:|---:|---:|---|
+| HumanEval | pass@1 | 53.05 | 164 | 0-shot generation; greedy decoding; `max_gen_toks=1024` |
+| HumanEval+ | pass@1 | 50.00 | 164 | Same as HumanEval pass@1, using `evalplus/humanevalplus` |
+| HumanEval | pass@16 | 77.44 | 164 | 16 sampled completions; `temperature=0.8`; `top_p=1.0`; `max_gen_toks=1024` |
+| IFEval | prompt-level strict accuracy | 12.57 | 541 | 0-shot chat generation; greedy decoding; `temperature=0`; `max_gen_toks=1280` |
+| IFEval | instruction-level strict accuracy | 23.74 | 541 | Same IFEval run |
+| ARC-Challenge | normalized accuracy | 31.06 | 1,172 | 0-shot multiple-choice logprobs |
+| ARC-Easy | normalized accuracy | 47.90 | 2,376 | 0-shot multiple-choice logprobs |
+| MMLU-Pro | exact match, custom extract | 11.77 | 12,032 | 5-shot chain-of-thought generation; greedy decoding; |
+| GPQA Diamond | normalized accuracy | 25.25 | 198 | 0-shot multiple-choice logprobs |
+| GPQA Main | normalized accuracy | 25.22 | 448 | 0-shot multiple-choice logprobs |
+| GPQA Extended | normalized accuracy | 21.98 | 546 | 0-shot multiple-choice logprobs |
+
+Long-context checks were also run with RULER-style synthetic tasks at 2,048 and 4,096 tokens using 0-shot greedy chat generation (`temperature=0`, remote tokenizer, `max_length=4096`). Values below are string-match scores.
+
+| Task | 2,048 tokens | 4,096 tokens | Effective samples |
+|---|---:|---:|---:|
+| NIAH single 1 | 100.00 | 100.00 | 500 |
+| NIAH single 2 | 100.00 | 100.00 | 500 |
+| NIAH single 3 | 99.80 | 91.60 | 500 |
+| NIAH multikey 1 | 99.20 | 83.40 | 500 |
+| NIAH multikey 2 | 97.80 | 60.00 | 500 |
+| NIAH multikey 3 | 91.80 | 91.20 | 500 |
+| NIAH multiquery | 95.40 | 69.70 | 500 |
+| NIAH multivalue | 94.60 | 82.20 | 500 |
+| RULER variable tracking | 77.92 | 3.92 | 500 |
+| RULER common words extraction | 24.72 | 78.14 | 500 |
+| RULER frequent words extraction | 78.80 | 64.73 | 500 |
+
+TODO (I'll add all benchmark and remove this sentence): HellaSwag, standard MMLU, and GSM8K are not reported for this release.
+
+## Intended use
+
+This model is intended for:
+
+- Research and experimentation with latency-oriented Transformer architectures.
+- Evaluation of Laneformer / Delayed Tensor Parallelism design choices.
+- Causal language modeling and code-generation experiments.
+- Fine-tuning experiments, subject to all applicable license terms.
+- Inference-system and attention-backend testing.
+
+This model is not automatically suitable for high-stakes settings such as medical, legal, financial, employment, education, public-sector, law-enforcement, or safety-critical decision-making.
+
+## Out-of-scope use
+
+Do not use this model or tokenizer in ways that violate the repository licenses, the Llama 2 Community License, the Llama 2 Acceptable Use Policy, applicable law, privacy rights, or safety policies. Do not present generated outputs as factual without independent verification. Do not use the model as the sole decision-maker in high-stakes workflows.
+
+## Limitations and risks
+
+- This is a 2B-class model and is not a frontier general-purpose assistant.
+- The model is mid-trained and post-trained for code generation only.
+- The model may generate incorrect, insecure, biased, toxic, or misleading content.
+- The model may hallucinate facts, APIs, package names, citations, and code behavior.
+- The base model may not follow instructions reliably unless it has been instruction-tuned.
+- The public HF implementation prioritizes compatibility and inspection, not the full Kog Inference Engine latency path.
+- Performance may depend on the exact PyTorch, Transformers, device, dtype, and attention backend used.
+- Because this repo uses custom code, downstream users should review code before execution and pin a revision in production.
+
+## License
+
+This repository is **multi-license**.
+
+### Kog-owned materials: Apache License 2.0
+
+Unless a file states otherwise, the model weights, Hugging Face custom modeling and configuration code, model configuration files, metadata files, documentation, and README/model card are released under the **Apache License 2.0**.
+
+See [`LICENSE`](./LICENSE) for the full Apache License 2.0 text.
+
+### Tokenizer materials: Llama 2 Community License
+
+The tokenizer files are based on the **Llama 2 tokenizer** and are not licensed under Apache License 2.0. These tokenizer materials include, without limitation:
+
+- `tokenizer.model`
+- `tokenizer.json`
+- `tokenizer_config.json`
+- `special_tokens_map.json`
+- any other file whose purpose is to reproduce or configure the Llama 2 tokenizer
+
+These tokenizer materials are distributed under the **LLAMA 2 Community License Agreement**. See [`THIRD_PARTY_LICENSES/LLAMA2_LICENSE`](./THIRD_PARTY_LICENSES/LLAMA2_LICENSE) and [`NOTICE`](./NOTICE).
+
+Required Llama 2 attribution notice:
+
+> Llama 2 is licensed under the LLAMA 2 Community License, Copyright (c) Meta Platforms, Inc. All Rights Reserved.
+
+Users and redistributors are responsible for complying with the Llama 2 Community License and its Acceptable Use Policy.
+
+## Citation
+
+If you use this model or architecture, please cite the model repository and the Kog technical blog posts:
+
+```bibtex
+@misc{kog_laneformer_tp_2b_2026,
+  title        = {Kog Laneformer 2B},
+  author       = {Kog AI},
+  year         = {2026},
+  publisher    = {Hugging Face},
+  howpublished = {\url{https://huggingface.co/kogai/laneformer-2b-it}}
+}
+
+@misc{kog_real_time_llm_inference_2026,
+  title        = {Real-time LLM Inference on Standard Datacenter GPUs (3,000 tokens/s per request)},
+  author       = {Kog Team},
+  year         = {2026},
+  howpublished = {\url{https://blog.kog.ai/real-time-llm-inference-on-standard-gpus-3-000-tokens-s-per-request/}}
+}
+
+@misc{kog_delayed_tensor_parallelism_2026,
+  title        = {Delayed Tensor Parallelism for Faster Transformer Inference},
+  author       = {Kog Team},
+  year         = {2026},
+  howpublished = {\url{https://blog.kog.ai/delayed-tensor-parallelism-for-faster-transformer-inference/}}
+}
+```
+
+## Contact
+
+For questions about this model, open a discussion or issue on the Hugging Face repository, or contact Kog AI through the channels listed on [kog.ai](https://www.kog.ai/).

THIRD_PARTY_LICENSES/LLAMA2_LICENSE

@@ -0,0 +1,125 @@
+LLAMA 2 COMMUNITY LICENSE AGREEMENT	
+Llama 2 Version Release Date: July 18, 2023
+
+"Agreement" means the terms and conditions for use, reproduction, distribution and 
+modification of the Llama Materials set forth herein.
+
+"Documentation" means the specifications, manuals and documentation 
+accompanying Llama 2 distributed by Meta at ai.meta.com/resources/models-and-
+libraries/llama-downloads/.
+
+"Licensee" or "you" means you, or your employer or any other person or entity (if 
+you are entering into this Agreement on such person or entity's behalf), of the age 
+required under applicable laws, rules or regulations to provide legal consent and that 
+has legal authority to bind your employer or such other person or entity if you are 
+entering in this Agreement on their behalf.
+
+"Llama 2" means the foundational large language models and software and 
+algorithms, including machine-learning model code, trained model weights, 
+inference-enabling code, training-enabling code, fine-tuning enabling code and other 
+elements of the foregoing distributed by Meta at ai.meta.com/resources/models-and-
+libraries/llama-downloads/.
+
+"Llama Materials" means, collectively, Meta's proprietary Llama 2 and 
+Documentation (and any portion thereof) made available under this Agreement.
+
+"Meta" or "we" means Meta Platforms Ireland Limited (if you are located in or, if you 
+are an entity, your principal place of business is in the EEA or Switzerland) and Meta 
+Platforms, Inc. (if you are located outside of the EEA or Switzerland). 
+
+By clicking "I Accept" below or by using or distributing any portion or element of the 
+Llama Materials, you agree to be bound by this Agreement.
+
+1. License Rights and Redistribution. 
+
+      a. Grant of Rights. You are granted a non-exclusive, worldwide, non-
+transferable and royalty-free limited license under Meta's intellectual property or 
+other rights owned by Meta embodied in the Llama Materials to use, reproduce, 
+distribute, copy, create derivative works of, and make modifications to the Llama 
+Materials.  
+      
+      b. Redistribution and Use.  
+
+            i. If you distribute or make the Llama Materials, or any derivative works 
+thereof, available to a third party, you shall provide a copy of this Agreement to such 
+third party. 
+            ii.  If you receive Llama Materials, or any derivative works thereof, from 
+a Licensee as part of an integrated end user product, then Section 2 of this 
+Agreement will not apply to you. 
+
+            iii. You must retain in all copies of the Llama Materials that you 
+distribute the following attribution notice within a "Notice" text file distributed as a 
+part of such copies: "Llama 2 is licensed under the LLAMA 2 Community License, 
+Copyright (c) Meta Platforms, Inc. All Rights Reserved."
+
+            iv. Your use of the Llama Materials must comply with applicable laws 
+and regulations (including trade compliance laws and regulations) and adhere to the 
+Acceptable Use Policy for the Llama Materials (available at 
+https://ai.meta.com/llama/use-policy), which is hereby incorporated by reference into 
+this Agreement.
+
+            v. You will not use the Llama Materials or any output or results of the 
+Llama Materials to improve any other large language model (excluding Llama 2 or 
+derivative works thereof).  
+
+2. Additional Commercial Terms. If, on the Llama 2 version release date, the 
+monthly active users of the products or services made available by or for Licensee, 
+or Licensee's affiliates, is greater than 700 million monthly active users in the 
+preceding calendar month, you must request a license from Meta, which Meta may 
+grant to you in its sole discretion, and you are not authorized to exercise any of the 
+rights under this Agreement unless or until Meta otherwise expressly grants you 
+such rights.
+            
+3. Disclaimer of Warranty. UNLESS REQUIRED BY APPLICABLE LAW, THE 
+LLAMA MATERIALS AND ANY OUTPUT AND RESULTS THEREFROM ARE 
+PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, 
+EITHER EXPRESS OR IMPLIED, INCLUDING, WITHOUT LIMITATION, ANY 
+WARRANTIES OF TITLE, NON-INFRINGEMENT, MERCHANTABILITY, OR 
+FITNESS FOR A PARTICULAR PURPOSE. YOU ARE SOLELY RESPONSIBLE 
+FOR DETERMINING THE APPROPRIATENESS OF USING OR REDISTRIBUTING 
+THE LLAMA MATERIALS AND ASSUME ANY RISKS ASSOCIATED WITH YOUR 
+USE OF THE LLAMA MATERIALS AND ANY OUTPUT AND RESULTS.
+
+4. Limitation of Liability. IN NO EVENT WILL META OR ITS AFFILIATES BE 
+LIABLE UNDER ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, TORT, 
+NEGLIGENCE, PRODUCTS LIABILITY, OR OTHERWISE, ARISING OUT OF THIS 
+AGREEMENT, FOR ANY LOST PROFITS OR ANY INDIRECT, SPECIAL, 
+CONSEQUENTIAL, INCIDENTAL, EXEMPLARY OR PUNITIVE DAMAGES, EVEN 
+IF META OR ITS AFFILIATES HAVE BEEN ADVISED OF THE POSSIBILITY OF 
+ANY OF THE FOREGOING.
+ 
+5. Intellectual Property.
+
+      a. No trademark licenses are granted under this Agreement, and in 
+connection with the Llama Materials, neither Meta nor Licensee may use any name 
+or mark owned by or associated with the other or any of its affiliates, except as 
+required for reasonable and customary use in describing and redistributing the 
+Llama Materials.
+
+      b. Subject to Meta's ownership of Llama Materials and derivatives made by or 
+for Meta, with respect to any derivative works and modifications of the Llama 
+Materials that are made by you, as between you and Meta, you are and will be the 
+owner of such derivative works and modifications.
+
+      c. If you institute litigation or other proceedings against Meta or any entity 
+(including a cross-claim or counterclaim in a lawsuit) alleging that the Llama 
+Materials or Llama 2 outputs or results, or any portion of any of the foregoing, 
+constitutes infringement of intellectual property or other rights owned or licensable 
+by you, then any licenses granted to you under this Agreement shall terminate as of 
+the date such litigation or claim is filed or instituted. You will indemnify and hold 
+harmless Meta from and against any claim by any third party arising out of or related 
+to your use or distribution of the Llama Materials.
+
+6. Term and Termination. The term of this Agreement will commence upon your 
+acceptance of this Agreement or access to the Llama Materials and will continue in 
+full force and effect until terminated in accordance with the terms and conditions 
+herein. Meta may terminate this Agreement if you are in breach of any term or 
+condition of this Agreement. Upon termination of this Agreement, you shall delete 
+and cease use of the Llama Materials. Sections 3, 4 and 7 shall survive the 
+termination of this Agreement. 
+
+7. Governing Law and Jurisdiction. This Agreement will be governed and 
+construed under the laws of the State of California without regard to choice of law 
+principles, and the UN Convention on Contracts for the International Sale of Goods 
+does not apply to this Agreement. The courts of California shall have exclusive 
+jurisdiction of any dispute arising out of this Agreement. 

__init__.py

@@ -0,0 +1,11 @@
+from transformers import AutoConfig
+
+from .configuration_laneformer import LaneformerTPConfig
+from .modeling_laneformer import LaneformerTPForCausalLM, LaneformerTPModel
+
+LaneformerTPConfig.register_for_auto_class()
+LaneformerTPModel.register_for_auto_class("AutoModel")
+LaneformerTPForCausalLM.register_for_auto_class("AutoModelForCausalLM")
+AutoConfig.register("laneformer_tp", LaneformerTPConfig)
+
+__all__ = ["LaneformerTPConfig", "LaneformerTPModel", "LaneformerTPForCausalLM"]

chat_template.jinja

@@ -0,0 +1,21 @@
+{{- bos_token }}
+{%- for message in messages %}
+    {%- if message['role'] == 'system' %}
+        {#- Ignore system messages entirely in this template. #}
+    {%- elif message['role'] == 'user' %}
+        {{- '<|start_header_id|>user<|end_header_id|>\n\n' }}
+        {{- message['content'] | trim }}
+        {{- eos_token }}
+    {%- elif message['role'] == 'assistant' %}
+        {{- '<|start_header_id|>assistant<|end_header_id|>\n\n' }}
+        {%- generation -%}
+        {{- message['content'] | trim }}
+        {{- eos_token }}
+        {%- endgeneration -%}
+    {%- else %}
+        {{- raise_exception("Unsupported role for no-system/no-tools template: " + message['role']) }}
+    {%- endif %}
+{%- endfor %}
+{%- if add_generation_prompt %}
+    {{- '<|start_header_id|>assistant<|end_header_id|>\n\n' }}
+{%- endif %}

config.json

@@ -0,0 +1,66 @@
+{
+  "model_type": "laneformer_tp",
+  "architectures": [
+    "LaneformerTPForCausalLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_laneformer.LaneformerTPConfig",
+    "AutoModel": "modeling_laneformer.LaneformerTPModel",
+    "AutoModelForCausalLM": "modeling_laneformer.LaneformerTPForCausalLM"
+  },
+  "hidden_size": 3072,
+  "num_hidden_layers": 15,
+  "num_attention_heads": 32,
+  "num_key_value_heads": 16,
+  "intermediate_size": 12288,
+  "max_position_embeddings": 8192,
+  "max_seq_len": 8192,
+  "eos_token_id": 2,
+  "pad_token_id": 2,
+  "vocab_size": 32000,
+  "rope_theta": 10000,
+  "norm_eps": 1e-05,
+  "use_cache": false,
+  "num_lanes": 8,
+  "broadcast_delay": 2,
+  "use_attention_comm": true,
+  "use_mlp_comm": true,
+  "use_early_comm": false,
+  "lm_head_type": "vocab_parallel",
+  "pre_norm_lane_agg": false,
+  "replicated_rmsn_scale": true,
+  "tie_word_embeddings": false,
+  "sliding_window": 2048,
+  "swa_layers": [
+    0,
+    1,
+    2,
+    3,
+    4,
+    5,
+    6,
+    7,
+    8,
+    9
+  ],
+  "layer_types": [
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention"
+  ],
+  "_attn_implementation": "flex_attention",
+  "dtype": "bfloat16",
+  "transformers_version": "4.57.6"
+}

configuration_laneformer.py

@@ -0,0 +1,292 @@
+from __future__ import annotations
+
+from typing import Any, Optional
+
+from transformers import PretrainedConfig
+
+
+_VALID_LM_HEAD_TYPES = {"replicate", "lane", "vocab_parallel"}
+_VALID_ROPE_SCALING_KEYS = {
+    "scaling_factor",
+    "low_freq_factor",
+    "high_freq_factor",
+    "original_max_position_embeddings",
+}
+
+
+def _infer_layer_types(
+    *,
+    num_hidden_layers: int,
+    sliding_window: Optional[int],
+    swa_layers: list[int],
+    use_sliding_window: Optional[bool],
+) -> list[str]:
+    if (
+        sliding_window is None
+        or int(sliding_window) <= 0
+        or use_sliding_window is False
+    ):
+        return ["full_attention"] * int(num_hidden_layers)
+
+    sliding_layer_indices = {int(layer) for layer in swa_layers}
+    return [
+        "sliding_attention"
+        if layer_index in sliding_layer_indices
+        else "full_attention"
+        for layer_index in range(int(num_hidden_layers))
+    ]
+
+
+def _validate_positive_int(name: str, value: int) -> int:
+    value = int(value)
+    if value <= 0:
+        raise ValueError(f"{name} must be a positive integer, got {value}")
+    return value
+
+
+def _validate_nonnegative_int(name: str, value: int) -> int:
+    value = int(value)
+    if value < 0:
+        raise ValueError(f"{name} must be >= 0, got {value}")
+    return value
+
+
+def _validate_positive_float(name: str, value: float) -> float:
+    value = float(value)
+    if value <= 0.0:
+        raise ValueError(f"{name} must be > 0, got {value}")
+    return value
+
+
+def _validate_probability(name: str, value: float) -> float:
+    value = float(value)
+    if not 0.0 <= value < 1.0:
+        raise ValueError(f"{name} must satisfy 0.0 <= {name} < 1.0, got {value}")
+    return value
+
+
+def _validate_divisible(name: str, numerator: int, denominator: int) -> None:
+    if int(numerator) % int(denominator) != 0:
+        raise ValueError(f"{name} must be divisible by {denominator}, got {numerator}")
+
+
+def _is_power_of_two(value: int) -> bool:
+    return value > 0 and (value & (value - 1)) == 0
+
+
+def _is_rocm_torch_build() -> bool:
+    try:
+        import torch
+    except ImportError:
+        return False
+
+    return getattr(torch.version, "hip", None) is not None
+
+
+def _validate_swa_layers(swa_layers: list[int], num_hidden_layers: int) -> None:
+    for layer in swa_layers:
+        if not 0 <= layer < num_hidden_layers:
+            raise ValueError(
+                "swa_layers must satisfy 0 <= layer < num_hidden_layers, "
+                f"got layer {layer} with num_hidden_layers={num_hidden_layers}"
+            )
+
+
+def _validate_rope_scaling(rope_scaling: Optional[dict[str, Any]]) -> None:
+    if rope_scaling is None:
+        return
+    if not isinstance(rope_scaling, dict):
+        raise ValueError("rope_scaling must be a dictionary")
+
+    unknown_keys = set(rope_scaling) - _VALID_ROPE_SCALING_KEYS
+    if unknown_keys:
+        raise ValueError(f"Unknown rope_scaling keys: {sorted(unknown_keys)}")
+
+    scaling_factor = float(rope_scaling.get("scaling_factor", 1.0))
+    low_freq_factor = float(rope_scaling.get("low_freq_factor", 1.0))
+    high_freq_factor = float(rope_scaling.get("high_freq_factor", 1.0))
+    original_max_position_embeddings = int(
+        rope_scaling.get("original_max_position_embeddings", 4096)
+    )
+
+    if scaling_factor <= 0:
+        raise ValueError("rope_scaling.scaling_factor must be > 0")
+    if low_freq_factor <= 0 or high_freq_factor <= 0:
+        raise ValueError("rope_scaling frequency factors must be > 0")
+    if high_freq_factor == low_freq_factor:
+        raise ValueError(
+            "rope_scaling.high_freq_factor and low_freq_factor must differ"
+        )
+    if original_max_position_embeddings <= 0:
+        raise ValueError("rope_scaling.original_max_position_embeddings must be > 0")
+
+
+class LaneformerTPConfig(PretrainedConfig):
+    model_type = "laneformer_tp"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    def __init__(
+        self,
+        hidden_size: int = 4096,
+        num_hidden_layers: int = 32,
+        num_attention_heads: int = 32,
+        num_key_value_heads: Optional[int] = None,
+        intermediate_size: int = 4 * 4096,
+        max_position_embeddings: int = 131072,
+        eos_token_id: int = 0,
+        pad_token_id: Optional[int] = None,
+        vocab_size: int = 128256,
+        rope_theta: float = 10000.0,
+        norm_eps: float = 1e-5,
+        multiple_of: int = 32,
+        depth_init: bool = True,
+        use_flex_attn: bool = False,
+        attn_mask_type: str = "causal",
+        initializer_range: float = 0.02,
+        attention_dropout: float = 0.0,
+        use_cache: bool = False,
+        # laneformer-specific extras (HF is fine with extra fields):
+        num_lanes: int = 1,
+        broadcast_delay: int = 1,
+        use_attention_comm: bool = True,
+        use_mlp_comm: bool = True,
+        use_early_comm: bool = True,
+        lm_head_type: str = "replicate",
+        pre_norm_lane_agg: bool = False,
+        replicated_rmsn_scale: bool = True,
+        tie_word_embeddings: bool = False,
+        sliding_window: Optional[int] = None,
+        sliding_window_n_layers: int = 0,
+        swa_layers: Optional[list[int]] = None,
+        rope_scaling: Optional[dict[str, Any]] = None,
+        **kwargs,
+    ):
+        if pad_token_id is None:
+            pad_token_id = eos_token_id
+        if num_key_value_heads is None:
+            num_key_value_heads = num_attention_heads
+
+        hidden_size = _validate_positive_int("hidden_size", hidden_size)
+        num_hidden_layers = _validate_positive_int(
+            "num_hidden_layers", num_hidden_layers
+        )
+        num_attention_heads = _validate_positive_int(
+            "num_attention_heads", num_attention_heads
+        )
+        num_key_value_heads = _validate_positive_int(
+            "num_key_value_heads", num_key_value_heads
+        )
+        intermediate_size = _validate_positive_int(
+            "intermediate_size", intermediate_size
+        )
+        max_position_embeddings = _validate_positive_int(
+            "max_position_embeddings", max_position_embeddings
+        )
+        vocab_size = _validate_positive_int("vocab_size", vocab_size)
+        num_lanes = _validate_positive_int("num_lanes", num_lanes)
+        rope_theta = _validate_positive_float("rope_theta", rope_theta)
+        norm_eps = _validate_positive_float("norm_eps", norm_eps)
+        multiple_of = _validate_positive_int("multiple_of", multiple_of)
+        attention_dropout = _validate_probability(
+            "attention_dropout", attention_dropout
+        )
+        broadcast_delay = _validate_nonnegative_int("broadcast_delay", broadcast_delay)
+        sliding_window_n_layers = _validate_nonnegative_int(
+            "sliding_window_n_layers", sliding_window_n_layers
+        )
+        if sliding_window_n_layers > num_hidden_layers:
+            raise ValueError(
+                "sliding_window_n_layers must satisfy "
+                "0 <= sliding_window_n_layers <= num_hidden_layers, "
+                f"got {sliding_window_n_layers} with "
+                f"num_hidden_layers={num_hidden_layers}"
+            )
+        _validate_divisible("hidden_size", hidden_size, num_attention_heads)
+        _validate_divisible(
+            "num_attention_heads", num_attention_heads, num_key_value_heads
+        )
+        _validate_divisible("hidden_size", hidden_size, num_lanes)
+        _validate_divisible("intermediate_size", intermediate_size, num_lanes)
+        head_dim = hidden_size // num_attention_heads
+        if head_dim % 2 != 0:
+            raise ValueError(
+                "hidden_size // num_attention_heads must be even for RoPE, "
+                f"got {head_dim}"
+            )
+        _validate_divisible(
+            "num_key_value_heads * head_dim",
+            num_key_value_heads * head_dim,
+            num_lanes,
+        )
+        if lm_head_type not in _VALID_LM_HEAD_TYPES:
+            raise ValueError(
+                f"lm_head_type must be one of {sorted(_VALID_LM_HEAD_TYPES)}, "
+                f"got {lm_head_type!r}"
+            )
+        if lm_head_type == "vocab_parallel":
+            _validate_divisible("vocab_size", vocab_size, num_lanes)
+        if sliding_window is not None:
+            sliding_window = _validate_positive_int("sliding_window", sliding_window)
+        if tie_word_embeddings:
+            raise ValueError("tie_word_embeddings is not supported for LaneformerTP")
+        _validate_rope_scaling(rope_scaling)
+        if swa_layers is None:
+            swa_layers = list(range(sliding_window_n_layers))
+        else:
+            swa_layers = sorted({int(layer) for layer in swa_layers})
+        _validate_swa_layers(swa_layers, num_hidden_layers)
+
+        super().__init__(
+            eos_token_id=eos_token_id,
+            pad_token_id=pad_token_id,
+            vocab_size=vocab_size,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+        if (
+            self._attn_implementation == "flex_attention"
+            and _is_rocm_torch_build()
+            and not _is_power_of_two(head_dim)
+        ):
+            self._attn_implementation = "sdpa"
+
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.num_key_value_heads = num_key_value_heads
+        self.intermediate_size = intermediate_size
+        self.max_position_embeddings = max_position_embeddings
+        self.rope_theta = rope_theta
+        self.norm_eps = norm_eps
+        self.multiple_of = multiple_of
+        self.depth_init = depth_init
+        self.use_flex_attn = use_flex_attn
+        self.attn_mask_type = attn_mask_type
+        self.initializer_range = initializer_range
+        self.attention_dropout = attention_dropout
+        self.use_cache = use_cache
+        self.sliding_window = sliding_window
+        self.rope_scaling = rope_scaling
+        self.swa_layers = swa_layers
+        # legacy field kept for backward compatibility
+        self.sliding_window_n_layers = len(self.swa_layers)
+        if getattr(self, "layer_types", None) is None:
+            self.layer_types = _infer_layer_types(
+                num_hidden_layers=num_hidden_layers,
+                sliding_window=sliding_window,
+                swa_layers=self.swa_layers,
+                use_sliding_window=getattr(self, "use_sliding_window", None),
+            )
+
+        # custom fields
+        self.num_lanes = num_lanes
+
+        self.use_attention_comm = use_attention_comm
+        self.use_mlp_comm = use_mlp_comm
+        # legacy aggregate flag kept for compatibility
+        self.use_comm = bool(self.use_attention_comm and self.use_mlp_comm)
+        self.use_early_comm = use_early_comm
+        self.broadcast_delay = broadcast_delay
+        self.lm_head_type = lm_head_type
+        self.pre_norm_lane_agg = pre_norm_lane_agg
+        self.replicated_rmsn_scale = replicated_rmsn_scale

generation_config.json

@@ -0,0 +1,6 @@
+{
+  "_from_model_config": true,
+  "eos_token_id": 2,
+  "transformers_version": "4.57.6",
+  "pad_token_id": 2
+}

model-00001-of-00002.safetensors

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

model-00002-of-00002.safetensors

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

model.safetensors.index.json

@@ -0,0 +1,146 @@
+{
+  "metadata": {
+    "total_parameters": 2320069632,
+    "total_size": 4640139264
+  },
+  "weight_map": {
+    "lm_head.weight": "model-00002-of-00002.safetensors",
+    "model.layers.0.attention.wk.weight": "model-00001-of-00002.safetensors",
+    "model.layers.0.attention.wo.weight": "model-00001-of-00002.safetensors",
+    "model.layers.0.attention.wq.weight": "model-00001-of-00002.safetensors",
+    "model.layers.0.attention.wv.weight": "model-00001-of-00002.safetensors",
+    "model.layers.0.attention_norm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.0.feed_forward.w1.weight": "model-00001-of-00002.safetensors",
+    "model.layers.0.feed_forward.w2.weight": "model-00001-of-00002.safetensors",
+    "model.layers.0.feed_forward.w3.weight": "model-00001-of-00002.safetensors",
+    "model.layers.0.ffn_norm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.1.attention.wk.weight": "model-00001-of-00002.safetensors",
+    "model.layers.1.attention.wo.weight": "model-00001-of-00002.safetensors",
+    "model.layers.1.attention.wq.weight": "model-00001-of-00002.safetensors",
+    "model.layers.1.attention.wv.weight": "model-00001-of-00002.safetensors",
+    "model.layers.1.attention_norm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.1.feed_forward.w1.weight": "model-00001-of-00002.safetensors",
+    "model.layers.1.feed_forward.w2.weight": "model-00001-of-00002.safetensors",
+    "model.layers.1.feed_forward.w3.weight": "model-00001-of-00002.safetensors",
+    "model.layers.1.ffn_norm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.10.attention.wk.weight": "model-00002-of-00002.safetensors",
+    "model.layers.10.attention.wo.weight": "model-00002-of-00002.safetensors",
+    "model.layers.10.attention.wq.weight": "model-00002-of-00002.safetensors",
+    "model.layers.10.attention.wv.weight": "model-00002-of-00002.safetensors",
+    "model.layers.10.attention_norm.weight": "model-00002-of-00002.safetensors",
+    "model.layers.10.feed_forward.w1.weight": "model-00002-of-00002.safetensors",
+    "model.layers.10.feed_forward.w2.weight": "model-00002-of-00002.safetensors",
+    "model.layers.10.feed_forward.w3.weight": "model-00002-of-00002.safetensors",
+    "model.layers.10.ffn_norm.weight": "model-00002-of-00002.safetensors",
+    "model.layers.11.attention.wk.weight": "model-00002-of-00002.safetensors",
+    "model.layers.11.attention.wo.weight": "model-00002-of-00002.safetensors",
+    "model.layers.11.attention.wq.weight": "model-00002-of-00002.safetensors",
+    "model.layers.11.attention.wv.weight": "model-00002-of-00002.safetensors",
+    "model.layers.11.attention_norm.weight": "model-00002-of-00002.safetensors",
+    "model.layers.11.feed_forward.w1.weight": "model-00002-of-00002.safetensors",
+    "model.layers.11.feed_forward.w2.weight": "model-00002-of-00002.safetensors",
+    "model.layers.11.feed_forward.w3.weight": "model-00002-of-00002.safetensors",
+    "model.layers.11.ffn_norm.weight": "model-00002-of-00002.safetensors",
+    "model.layers.12.attention.wk.weight": "model-00002-of-00002.safetensors",
+    "model.layers.12.attention.wo.weight": "model-00002-of-00002.safetensors",
+    "model.layers.12.attention.wq.weight": "model-00002-of-00002.safetensors",
+    "model.layers.12.attention.wv.weight": "model-00002-of-00002.safetensors",
+    "model.layers.12.attention_norm.weight": "model-00002-of-00002.safetensors",
+    "model.layers.12.feed_forward.w1.weight": "model-00002-of-00002.safetensors",
+    "model.layers.12.feed_forward.w2.weight": "model-00002-of-00002.safetensors",
+    "model.layers.12.feed_forward.w3.weight": "model-00002-of-00002.safetensors",
+    "model.layers.12.ffn_norm.weight": "model-00002-of-00002.safetensors",
+    "model.layers.13.attention.wk.weight": "model-00002-of-00002.safetensors",
+    "model.layers.13.attention.wo.weight": "model-00002-of-00002.safetensors",
+    "model.layers.13.attention.wq.weight": "model-00002-of-00002.safetensors",
+    "model.layers.13.attention.wv.weight": "model-00002-of-00002.safetensors",
+    "model.layers.13.attention_norm.weight": "model-00002-of-00002.safetensors",
+    "model.layers.13.feed_forward.w1.weight": "model-00002-of-00002.safetensors",
+    "model.layers.13.feed_forward.w2.weight": "model-00002-of-00002.safetensors",
+    "model.layers.13.feed_forward.w3.weight": "model-00002-of-00002.safetensors",
+    "model.layers.13.ffn_norm.weight": "model-00002-of-00002.safetensors",
+    "model.layers.14.attention.wk.weight": "model-00002-of-00002.safetensors",
+    "model.layers.14.attention.wo.weight": "model-00002-of-00002.safetensors",
+    "model.layers.14.attention.wq.weight": "model-00002-of-00002.safetensors",
+    "model.layers.14.attention.wv.weight": "model-00002-of-00002.safetensors",
+    "model.layers.14.attention_norm.weight": "model-00002-of-00002.safetensors",
+    "model.layers.14.feed_forward.w1.weight": "model-00002-of-00002.safetensors",
+    "model.layers.14.feed_forward.w2.weight": "model-00002-of-00002.safetensors",
+    "model.layers.14.feed_forward.w3.weight": "model-00002-of-00002.safetensors",
+    "model.layers.14.ffn_norm.weight": "model-00002-of-00002.safetensors",
+    "model.layers.2.attention.wk.weight": "model-00001-of-00002.safetensors",
+    "model.layers.2.attention.wo.weight": "model-00001-of-00002.safetensors",
+    "model.layers.2.attention.wq.weight": "model-00001-of-00002.safetensors",
+    "model.layers.2.attention.wv.weight": "model-00001-of-00002.safetensors",
+    "model.layers.2.attention_norm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.2.feed_forward.w1.weight": "model-00001-of-00002.safetensors",
+    "model.layers.2.feed_forward.w2.weight": "model-00001-of-00002.safetensors",
+    "model.layers.2.feed_forward.w3.weight": "model-00001-of-00002.safetensors",
+    "model.layers.2.ffn_norm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.3.attention.wk.weight": "model-00001-of-00002.safetensors",
+    "model.layers.3.attention.wo.weight": "model-00001-of-00002.safetensors",
+    "model.layers.3.attention.wq.weight": "model-00001-of-00002.safetensors",
+    "model.layers.3.attention.wv.weight": "model-00001-of-00002.safetensors",
+    "model.layers.3.attention_norm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.3.feed_forward.w1.weight": "model-00001-of-00002.safetensors",
+    "model.layers.3.feed_forward.w2.weight": "model-00001-of-00002.safetensors",
+    "model.layers.3.feed_forward.w3.weight": "model-00001-of-00002.safetensors",
+    "model.layers.3.ffn_norm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.4.attention.wk.weight": "model-00001-of-00002.safetensors",
+    "model.layers.4.attention.wo.weight": "model-00001-of-00002.safetensors",
+    "model.layers.4.attention.wq.weight": "model-00001-of-00002.safetensors",
+    "model.layers.4.attention.wv.weight": "model-00001-of-00002.safetensors",
+    "model.layers.4.attention_norm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.4.feed_forward.w1.weight": "model-00001-of-00002.safetensors",
+    "model.layers.4.feed_forward.w2.weight": "model-00001-of-00002.safetensors",
+    "model.layers.4.feed_forward.w3.weight": "model-00001-of-00002.safetensors",
+    "model.layers.4.ffn_norm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.5.attention.wk.weight": "model-00001-of-00002.safetensors",
+    "model.layers.5.attention.wo.weight": "model-00001-of-00002.safetensors",
+    "model.layers.5.attention.wq.weight": "model-00001-of-00002.safetensors",
+    "model.layers.5.attention.wv.weight": "model-00001-of-00002.safetensors",
+    "model.layers.5.attention_norm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.5.feed_forward.w1.weight": "model-00001-of-00002.safetensors",
+    "model.layers.5.feed_forward.w2.weight": "model-00001-of-00002.safetensors",
+    "model.layers.5.feed_forward.w3.weight": "model-00001-of-00002.safetensors",
+    "model.layers.5.ffn_norm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.6.attention.wk.weight": "model-00001-of-00002.safetensors",
+    "model.layers.6.attention.wo.weight": "model-00001-of-00002.safetensors",
+    "model.layers.6.attention.wq.weight": "model-00001-of-00002.safetensors",
+    "model.layers.6.attention.wv.weight": "model-00001-of-00002.safetensors",
+    "model.layers.6.attention_norm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.6.feed_forward.w1.weight": "model-00001-of-00002.safetensors",
+    "model.layers.6.feed_forward.w2.weight": "model-00001-of-00002.safetensors",
+    "model.layers.6.feed_forward.w3.weight": "model-00001-of-00002.safetensors",
+    "model.layers.6.ffn_norm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.7.attention.wk.weight": "model-00001-of-00002.safetensors",
+    "model.layers.7.attention.wo.weight": "model-00001-of-00002.safetensors",
+    "model.layers.7.attention.wq.weight": "model-00001-of-00002.safetensors",
+    "model.layers.7.attention.wv.weight": "model-00001-of-00002.safetensors",
+    "model.layers.7.attention_norm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.7.feed_forward.w1.weight": "model-00001-of-00002.safetensors",
+    "model.layers.7.feed_forward.w2.weight": "model-00001-of-00002.safetensors",
+    "model.layers.7.feed_forward.w3.weight": "model-00001-of-00002.safetensors",
+    "model.layers.7.ffn_norm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.8.attention.wk.weight": "model-00001-of-00002.safetensors",
+    "model.layers.8.attention.wo.weight": "model-00002-of-00002.safetensors",
+    "model.layers.8.attention.wq.weight": "model-00001-of-00002.safetensors",
+    "model.layers.8.attention.wv.weight": "model-00001-of-00002.safetensors",
+    "model.layers.8.attention_norm.weight": "model-00002-of-00002.safetensors",
+    "model.layers.8.feed_forward.w1.weight": "model-00002-of-00002.safetensors",
+    "model.layers.8.feed_forward.w2.weight": "model-00002-of-00002.safetensors",
+    "model.layers.8.feed_forward.w3.weight": "model-00002-of-00002.safetensors",
+    "model.layers.8.ffn_norm.weight": "model-00002-of-00002.safetensors",
+    "model.layers.9.attention.wk.weight": "model-00002-of-00002.safetensors",
+    "model.layers.9.attention.wo.weight": "model-00002-of-00002.safetensors",
+    "model.layers.9.attention.wq.weight": "model-00002-of-00002.safetensors",
+    "model.layers.9.attention.wv.weight": "model-00002-of-00002.safetensors",
+    "model.layers.9.attention_norm.weight": "model-00002-of-00002.safetensors",
+    "model.layers.9.feed_forward.w1.weight": "model-00002-of-00002.safetensors",
+    "model.layers.9.feed_forward.w2.weight": "model-00002-of-00002.safetensors",
+    "model.layers.9.feed_forward.w3.weight": "model-00002-of-00002.safetensors",
+    "model.layers.9.ffn_norm.weight": "model-00002-of-00002.safetensors",
+    "model.norm.weight": "model-00002-of-00002.safetensors",
+    "model.tok_embeddings.weight": "model-00001-of-00002.safetensors"
+  }
+}

modeling_laneformer.py

@@ -0,0 +1,1047 @@
+"""
+Hugging Face-compatible Laneformer implementation.
+
+This file defines:
+- LaneformerTPPreTrainedModel
+- LaneformerTPModel (backbone)
+- LaneformerTPForCausalLM (head + loss)
+
+It mirrors the TorchTitan Laneformer structure so that state_dict keys largely match,
+which makes converting checkpoints trivial (often identity mapping).
+"""
+
+from __future__ import annotations
+
+import inspect
+import math
+from dataclasses import dataclass
+from enum import Enum
+from functools import partial
+from typing import Optional, Tuple, Union, Unpack
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from transformers.cache_utils import Cache, DynamicCache
+from transformers.generation import GenerationMixin
+from transformers.masking_utils import (
+    create_causal_mask,
+    create_sliding_window_causal_mask,
+)
+from transformers.modeling_outputs import (
+    BaseModelOutputWithPast,
+    CausalLMOutputWithPast,
+)
+from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from transformers.utils import auto_docstring, can_return_tuple, TransformersKwargs
+
+from .configuration_laneformer import LaneformerTPConfig
+
+
+# ============================================================
+# Rotary helpers (mirrors your implementation)
+# ============================================================
+@dataclass
+class RoPEScalingArgs:
+    scaling_factor: float = 1.0
+    low_freq_factor: float = 1.0
+    high_freq_factor: float = 1.0
+    original_max_position_embeddings: int = 4096
+
+
+def _rope_scaling_from_config(
+    config: LaneformerTPConfig,
+) -> RoPEScalingArgs | None:
+    if config.rope_scaling is None:
+        return None
+    try:
+        args = RoPEScalingArgs(**config.rope_scaling)
+    except TypeError as exc:
+        raise ValueError(f"Invalid rope_scaling configuration: {exc}") from exc
+
+    if args.scaling_factor <= 0:
+        raise ValueError("rope_scaling.scaling_factor must be > 0")
+    if args.low_freq_factor <= 0 or args.high_freq_factor <= 0:
+        raise ValueError("rope_scaling frequency factors must be > 0")
+    if args.high_freq_factor == args.low_freq_factor:
+        raise ValueError(
+            "rope_scaling.high_freq_factor and low_freq_factor must differ"
+        )
+    if args.original_max_position_embeddings <= 0:
+        raise ValueError("rope_scaling.original_max_position_embeddings must be > 0")
+    return args
+
+
+def compute_rope_freqs(
+    dim: int,
+    theta: float = 10000.0,
+    scaling_args: RoPEScalingArgs | None = None,
+    device: torch.device | None = None,
+) -> torch.Tensor:
+    freqs = 1.0 / (
+        theta ** (torch.arange(0, dim, 2, device=device)[: (dim // 2)].float() / dim)
+    )
+    if scaling_args is not None:
+        scaling_factor = scaling_args.scaling_factor
+        low_freq_factor = scaling_args.low_freq_factor
+        high_freq_factor = scaling_args.high_freq_factor
+        original_max_position_embeddings = scaling_args.original_max_position_embeddings
+        wavelen = 2 * math.pi / freqs
+        high_freq_wavelen = original_max_position_embeddings / high_freq_factor
+        low_freq_wavelen = original_max_position_embeddings / low_freq_factor
+        freqs = torch.where(wavelen > low_freq_wavelen, freqs / scaling_factor, freqs)
+        smooth_factor = (
+            original_max_position_embeddings / wavelen - low_freq_factor
+        ) / (high_freq_factor - low_freq_factor)
+        smoothed_freqs = (
+            1 - smooth_factor
+        ) * freqs / scaling_factor + smooth_factor * freqs
+        is_medium_freqs = ~(wavelen < high_freq_wavelen) * ~(wavelen > low_freq_wavelen)
+        freqs = torch.where(is_medium_freqs, smoothed_freqs, freqs)
+    return freqs
+
+
+def precompute_freqs_cis(
+    dim: int,
+    end: int,
+    theta: float = 10000.0,
+    scaling_args: RoPEScalingArgs | None = None,
+) -> torch.Tensor:
+    freqs = compute_rope_freqs(dim, theta, scaling_args)
+    t = torch.arange(end, device=freqs.device)
+    freqs = torch.outer(t, freqs).float()
+    freqs_cis = torch.polar(torch.ones_like(freqs), freqs)  # complex64
+    return freqs_cis
+
+
+def freqs_cis_for_position_ids(
+    position_ids: torch.Tensor,
+    dim: int,
+    theta: float = 10000.0,
+    scaling_args: RoPEScalingArgs | None = None,
+) -> torch.Tensor:
+    freqs = compute_rope_freqs(
+        dim,
+        theta,
+        scaling_args,
+        device=position_ids.device,
+    )
+    angles = position_ids.float().unsqueeze(-1) * freqs
+    return torch.polar(torch.ones_like(angles), angles)
+
+
+def reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor) -> torch.Tensor:
+    ndim = x.ndim
+    assert ndim > 1
+    seqlen = x.shape[1]
+    freqs_cis = freqs_cis[0:seqlen]
+    assert freqs_cis.shape == (seqlen, x.shape[-1])
+    shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]
+    return freqs_cis.view(*shape)
+
+
+def apply_rotary_emb(
+    xq: torch.Tensor,
+    xk: torch.Tensor,
+    freqs_cis: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
+    xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))
+    # freqs_cis = reshape_for_broadcast(freqs_cis, xq_)
+    xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3)
+    xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3)
+    return xq_out.type_as(xq), xk_out.type_as(xk)
+
+
+def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
+    bs, slen, n_kv_heads, head_dim = x.shape
+    if n_rep == 1:
+        return x
+    return (
+        torch.unsqueeze(x, dim=3)
+        .expand(bs, slen, n_kv_heads, n_rep, head_dim)
+        .reshape(bs, slen, n_kv_heads * n_rep, head_dim)
+    )
+
+
+class LaneLinear(nn.Module):
+    def __init__(self, in_features: int, out_features: int, num_lanes: int):
+        super().__init__()
+        self.num_lanes = num_lanes
+        self.in_features = in_features
+        self.out_features = out_features
+        self.weight = torch.empty(self.get_shape())
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        raise NotImplementedError
+
+    def get_shape(self) -> Tuple[int, int, int]:
+        raise NotImplementedError
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return torch.einsum(
+            "bsli,loi->bslo", x, self.weight
+        )  # [L, out_features, 1] @ [B*S, L, in_features // L, 1] -> [B*S, L, out_features, 1]
+
+
+class LaneColumnLinear(LaneLinear):
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        num_lanes: int,
+    ):
+        super().__init__(
+            in_features=in_features, out_features=out_features, num_lanes=num_lanes
+        )
+
+    def get_shape(self) -> Tuple[int, int, int]:
+        return (self.num_lanes, self.out_features, self.in_features // self.num_lanes)
+
+    def reset_parameters(self):
+        w = torch.empty(self.out_features, self.in_features)
+        nn.init.kaiming_uniform_(w, a=math.sqrt(5))
+        self.weight = nn.Parameter(
+            w.reshape(
+                self.out_features, self.num_lanes, self.in_features // self.num_lanes
+            )
+            .permute(1, 0, 2)
+            .contiguous()
+        )
+
+
+class LaneRowLinear(LaneLinear):
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        num_lanes: int,
+    ):
+        super().__init__(
+            in_features=in_features, out_features=out_features, num_lanes=num_lanes
+        )
+
+    def get_shape(self) -> Tuple[int, int, int]:
+        return (self.num_lanes, self.out_features // self.num_lanes, self.in_features)
+
+    def reset_parameters(self):
+        w = torch.empty(self.out_features, self.in_features)
+        nn.init.kaiming_uniform_(w, a=math.sqrt(5))
+        self.weight = nn.Parameter(w.reshape(self.weight.shape).contiguous())
+
+
+class LaneLMHead(nn.Module):
+    def __init__(self, in_features: int, out_features: int, num_lanes: int):
+        super().__init__()
+        self.num_lanes = num_lanes
+        self.linear = nn.Linear(
+            in_features=num_lanes * in_features, out_features=out_features, bias=False
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        B, S, L, d = x.shape
+        return self.linear(x.view(B, S, L * d))
+
+
+class LaneRMSNorm(nn.Module):
+    def __init__(self, dim: int, num_lanes: int, eps: float = 1e-8):
+        super().__init__()
+        self.rms_norm = nn.RMSNorm(dim, eps=eps, elementwise_affine=False)
+        self.dim = dim
+        self.eps = eps
+        self.scale = nn.Parameter(torch.ones(num_lanes, dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.rms_norm(x) * self.scale
+
+    def reset_parameters(self):
+        nn.init.ones_(self.scale)
+
+
+def eager_attention_forward(
+    module: nn.Module,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attention_mask: Optional[torch.Tensor],
+    scaling: Optional[float] = None,
+    dropout: float = 0.0,
+    **kwargs,
+):
+    num_key_value_groups = getattr(module, "num_key_value_groups", 1)
+    key_states = repeat_kv(key.transpose(1, 2), num_key_value_groups).transpose(1, 2)
+    value_states = repeat_kv(value.transpose(1, 2), num_key_value_groups).transpose(
+        1, 2
+    )
+
+    if scaling is None:
+        scaling = query.shape[-1] ** -0.5
+
+    attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
+    if attention_mask is not None:
+        causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
+        attn_weights = attn_weights + causal_mask
+
+    attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(
+        query.dtype
+    )
+    attn_weights = nn.functional.dropout(
+        attn_weights, p=dropout, training=module.training
+    )
+    attn_output = torch.matmul(attn_weights, value_states)
+    attn_output = attn_output.transpose(1, 2).contiguous()
+
+    return attn_output, attn_weights
+
+
+def _attention_interface(config: LaneformerTPConfig):
+    attn_implementation = getattr(config, "_attn_implementation", "eager")
+    if attn_implementation == "eager":
+        return eager_attention_forward
+
+    get_interface = getattr(ALL_ATTENTION_FUNCTIONS, "get_interface", None)
+    if get_interface is not None:
+        return get_interface(attn_implementation, eager_attention_forward)
+
+    return ALL_ATTENTION_FUNCTIONS[attn_implementation]
+
+
+_CAUSAL_MASK_USES_INPUTS_EMBEDS = (
+    "inputs_embeds" in inspect.signature(create_causal_mask).parameters
+)
+_SLIDING_MASK_USES_INPUTS_EMBEDS = (
+    "inputs_embeds" in inspect.signature(create_sliding_window_causal_mask).parameters
+)
+
+
+def _call_attention_mask_factory(
+    mask_factory,
+    *,
+    uses_inputs_embeds: bool,
+    config: LaneformerTPConfig,
+    inputs_embeds: torch.Tensor,
+    attention_mask: Optional[torch.Tensor],
+    cache_position: torch.Tensor,
+    past_key_values: Optional[Cache],
+    position_ids: Optional[torch.Tensor],
+):
+    if uses_inputs_embeds:
+        return mask_factory(
+            config=config,
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            past_key_values=past_key_values,
+            position_ids=position_ids,
+        )
+
+    return mask_factory(
+        config=config,
+        input_embeds=inputs_embeds,
+        attention_mask=attention_mask,
+        cache_position=cache_position,
+        past_key_values=past_key_values,
+        position_ids=position_ids,
+    )
+
+
+def _attention_mask_needs_sliding_window_overlay(
+    attention_mask: Optional[torch.Tensor],
+) -> bool:
+    return (
+        isinstance(attention_mask, torch.Tensor)
+        and attention_mask.ndim == 4
+        and attention_mask.shape[-1] == attention_mask.shape[-2]
+    )
+
+
+def _apply_sliding_window_to_attention_mask(
+    attention_mask: Optional[torch.Tensor],
+    *,
+    sliding_window: int,
+) -> Optional[torch.Tensor]:
+    if not _attention_mask_needs_sliding_window_overlay(attention_mask):
+        return attention_mask
+    sequence_length = int(attention_mask.shape[-1])
+    if sequence_length <= sliding_window:
+        return attention_mask
+
+    positions = torch.arange(sequence_length, device=attention_mask.device)
+    outside_window = positions[None, :] <= positions[:, None] - int(sliding_window)
+    if not bool(outside_window.any()):
+        return attention_mask
+
+    sliding_window_attention_mask = attention_mask.clone()
+    if attention_mask.dtype == torch.bool:
+        fill_value = False
+    elif torch.is_floating_point(attention_mask):
+        fill_value = torch.finfo(attention_mask.dtype).min
+    else:
+        fill_value = 0
+    sliding_window_attention_mask.masked_fill_(
+        outside_window[None, None, :, :],
+        fill_value,
+    )
+    return sliding_window_attention_mask
+
+
+def _layer_uses_sliding_window_attention(
+    config: LaneformerTPConfig,
+    layer_index: int,
+) -> bool:
+    sliding_window = getattr(config, "sliding_window", None)
+    if sliding_window is None or int(sliding_window) <= 0:
+        return False
+
+    swa_layers = getattr(config, "swa_layers", None)
+    if swa_layers is not None:
+        return layer_index in {int(layer) for layer in swa_layers}
+
+    layer_types = getattr(config, "layer_types", None)
+    if isinstance(layer_types, (list, tuple)) and layer_index < len(layer_types):
+        return layer_types[layer_index] == "sliding_attention"
+
+    sliding_window_n_layers = getattr(config, "sliding_window_n_layers", None)
+    if sliding_window_n_layers is not None:
+        return layer_index < int(sliding_window_n_layers)
+
+    return True
+
+
+def _layer_attention_masks_from_provided_mask(
+    config: LaneformerTPConfig,
+    attention_mask: Optional[torch.Tensor],
+) -> Optional[list[Optional[torch.Tensor]]]:
+    if not _attention_mask_needs_sliding_window_overlay(attention_mask):
+        return None
+
+    sliding_window = getattr(config, "sliding_window", None)
+    if sliding_window is None or int(sliding_window) <= 0:
+        return [attention_mask] * config.num_hidden_layers
+
+    sliding_window_attention_mask = _apply_sliding_window_to_attention_mask(
+        attention_mask,
+        sliding_window=int(sliding_window),
+    )
+    return [
+        sliding_window_attention_mask
+        if _layer_uses_sliding_window_attention(config, layer_id)
+        else attention_mask
+        for layer_id in range(config.num_hidden_layers)
+    ]
+
+
+class ReduceMode(Enum):
+    NO_REDUCE = "no_reduce"
+    PRESENT = "present"
+    PAST = "past"
+
+
+class LaneModule:
+    _NO_REDUCE_MODE = 0
+    _PRESENT_MODE = 1
+    _PAST_MODE = 2
+
+    def __init__(self, no_reduce_scale: float, reduce_mode: ReduceMode):
+        super().__init__()
+        self.no_reduce_scale = no_reduce_scale
+        self.reduce_mode = reduce_mode
+        if reduce_mode is ReduceMode.NO_REDUCE:
+            self.reduce_mode_code = self._NO_REDUCE_MODE
+        elif reduce_mode is ReduceMode.PRESENT:
+            self.reduce_mode_code = self._PRESENT_MODE
+        elif reduce_mode is ReduceMode.PAST:
+            self.reduce_mode_code = self._PAST_MODE
+        else:
+            raise ValueError(f"Unknown reduce mode: {reduce_mode}")
+
+    def reduce_lanes(self, x: torch.Tensor, past: torch.Tensor | None) -> torch.Tensor:
+        if self.reduce_mode_code == self._NO_REDUCE_MODE:
+            return x * self.no_reduce_scale
+        elif self.reduce_mode_code == self._PRESENT_MODE:
+            return x.sum(dim=2, keepdim=True)
+        elif self.reduce_mode_code == self._PAST_MODE:
+            assert past is not None
+            sum_past = torch.sum(past, dim=-2, keepdim=True) - past
+            return x + sum_past
+        else:
+            raise ValueError(f"Unknown reduce mode code: {self.reduce_mode_code}")
+
+
+class LaneformerAttention(LaneModule, nn.Module):
+    def __init__(
+        self,
+        config: LaneformerTPConfig,
+        layer_idx: int,
+        reduce_mode: ReduceMode = ReduceMode.PRESENT,
+    ):
+        super().__init__(
+            no_reduce_scale=math.sqrt(config.num_lanes), reduce_mode=reduce_mode
+        )
+        self.config = config
+        self.layer_idx = layer_idx
+        self.n_heads = config.num_attention_heads
+        self.num_lanes = config.num_lanes
+        self.n_kv_heads = (
+            config.num_attention_heads
+            if config.num_key_value_heads is None
+            else config.num_key_value_heads
+        )
+        self.n_rep = config.num_attention_heads // self.n_kv_heads
+        self.head_dim = config.hidden_size // config.num_attention_heads
+        # Scaling & causal flag for backend attention APIs
+        self.scaling = self.head_dim**-0.5
+        self.is_causal = True  # decoder-only model
+        self.num_key_value_groups = self.n_rep
+        self.attention_dropout = getattr(config, "attention_dropout", 0.0)
+        self.wq = LaneRowLinear(
+            in_features=config.hidden_size,
+            out_features=self.n_heads * self.head_dim,
+            num_lanes=config.num_lanes,
+        )
+        self.wk = LaneRowLinear(
+            in_features=config.hidden_size,
+            out_features=self.n_kv_heads * self.head_dim,
+            num_lanes=config.num_lanes,
+        )
+        self.wv = LaneRowLinear(
+            in_features=config.hidden_size,
+            out_features=self.n_kv_heads * self.head_dim,
+            num_lanes=config.num_lanes,
+        )
+        self.wo = LaneColumnLinear(
+            in_features=self.n_heads * self.head_dim,
+            out_features=config.hidden_size,
+            num_lanes=config.num_lanes,
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        freqs_cis: torch.Tensor,
+        attention_mask: Optional[torch.Tensor],
+        *,
+        past_key_values: Optional[Cache] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs,
+    ):
+        """
+        Forward pass of the attention module.
+
+        Args:
+            x (torch.Tensor): Input tensor (B, S, L, D)
+            freqs_cis (torch.Tensor): Precomputed frequency tensor.
+
+        Returns:
+            torch.Tensor: Output tensor after attention (B, S, L, D)
+
+        """
+        bs, seqlen, _, _ = x.shape
+
+        xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)
+        xq = xq.reshape(bs, seqlen, -1, self.head_dim)
+        # (B, T, L, n_kv_heads / L, head_dim)
+        xk = xk.reshape(bs, seqlen, -1, self.head_dim)
+        # (B, T, L, n_kv_heads / L, head_dim)
+        xv = xv.reshape(bs, seqlen, -1, self.head_dim)
+
+        xq, xk = apply_rotary_emb(xq, xk, freqs_cis=freqs_cis)
+        xq = xq.transpose(1, 2)  # (B, n_heads, T, head_dim)
+        xk = xk.transpose(1, 2)  # (B, n_kv_heads, T, head_dim)
+        xv = xv.transpose(1, 2)  # (B, n_kv_heads, T, head_dim)
+
+        # Cache update BEFORE replication
+        if past_key_values is not None:
+            # sin and cos are specific to RoPE models; cache_position needed for the static cache
+            cache_kwargs = {"sin": None, "cos": None, "cache_position": cache_position}
+            xk, xv = past_key_values.update(xk, xv, self.layer_idx, cache_kwargs)
+
+        attention_interface = _attention_interface(self.config)
+        output, _ = attention_interface(
+            self,
+            query=xq,
+            key=xk,
+            value=xv,
+            attention_mask=attention_mask,
+            scaling=self.scaling,
+            dropout=0.0 if not self.training else self.attention_dropout,
+            **kwargs,
+        )
+
+        output = output.contiguous().view(
+            bs,
+            seqlen,
+            self.num_lanes,
+            self.n_heads * self.head_dim // self.num_lanes,
+        )
+        output = self.wo(output)
+        return output
+
+
+class LaneformerMLP(LaneModule, nn.Module):
+    def __init__(
+        self, config: LaneformerTPConfig, reduce_mode: ReduceMode = ReduceMode.PRESENT
+    ):
+        super().__init__(
+            no_reduce_scale=math.sqrt(config.num_lanes), reduce_mode=reduce_mode
+        )
+        ffn_dim = config.intermediate_size
+        dim = config.hidden_size
+        num_lanes = config.num_lanes
+
+        self.w1 = LaneRowLinear(dim, ffn_dim, num_lanes=num_lanes)
+        self.w2 = LaneColumnLinear(ffn_dim, dim, num_lanes=num_lanes)
+        self.w3 = LaneRowLinear(dim, ffn_dim, num_lanes=num_lanes)
+
+    def forward(self, x):
+        return self.w2(F.silu(self.w1(x)) * self.w3(x))
+
+
+class LaneformerDecoderLayer(nn.Module):
+    def __init__(
+        self,
+        config: LaneformerTPConfig,
+        layer_idx: int,
+        attention_reduce_mode: ReduceMode,
+        mlp_reduce_mode: ReduceMode,
+        broadcast_attention_to_future: bool,
+        broadcast_mlp_to_future: bool,
+    ):
+        super().__init__()
+        self.n_heads = config.num_attention_heads
+        self.dim = config.hidden_size
+        self.attention = LaneformerAttention(
+            config, layer_idx=layer_idx, reduce_mode=attention_reduce_mode
+        )
+        self.feed_forward = LaneformerMLP(config, reduce_mode=mlp_reduce_mode)
+        norm_fun = (
+            nn.RMSNorm
+            if config.replicated_rmsn_scale
+            else partial(LaneRMSNorm, num_lanes=config.num_lanes)
+        )
+        self.attention_norm = norm_fun(config.hidden_size, eps=config.norm_eps)
+        self.ffn_norm = norm_fun(config.hidden_size, eps=config.norm_eps)
+        self.num_lanes = config.num_lanes
+        self.broadcast_attention_to_future = broadcast_attention_to_future
+        self.broadcast_mlp_to_future = broadcast_mlp_to_future
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        past_attention: Optional[torch.Tensor] = None,
+        past_mlp: Optional[torch.Tensor] = None,
+        use_cache: Optional[bool] = False,
+        cache_position: Optional[torch.LongTensor] = None,
+        freqs_cis: Optional[torch.Tensor] = None,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        residual = hidden_states
+        hidden_states = self.attention_norm(hidden_states)
+        hidden_states = self.attention(
+            x=hidden_states,
+            freqs_cis=freqs_cis,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            **kwargs,
+        )
+        future_attention = hidden_states if self.broadcast_attention_to_future else None
+        hidden_states = self.attention.reduce_lanes(hidden_states, past=past_attention)
+        hidden_states = residual + hidden_states
+
+        residual = hidden_states
+        hidden_states = self.ffn_norm(hidden_states)
+        hidden_states = self.feed_forward(hidden_states)
+        future_mlp = hidden_states if self.broadcast_mlp_to_future else None
+        hidden_states = self.feed_forward.reduce_lanes(hidden_states, past=past_mlp)
+        hidden_states = residual + hidden_states
+        return hidden_states, future_attention, future_mlp
+
+
+# ============================================================
+# Backbone & HF wrappers
+# ============================================================
+class LaneformerPreTrainedModel(PreTrainedModel):
+    config_class = LaneformerTPConfig
+    config: LaneformerTPConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = False
+    _no_split_modules = [
+        "LaneformerTPModel",
+        "LaneformerDecoderLayer",
+        "LaneformerAttention",
+        "LaneformerMLP",
+    ]
+    _skip_keys_device_placement = ["past_key_values"]
+    # Advertise support for Transformers attention backends (vLLM / flash / sdpa routing)
+    _supports_attention_backend = True
+    _supports_sdpa = True
+    _supports_flex_attn = True
+
+
+class LaneformerTPModel(LaneformerPreTrainedModel):
+    def __init__(self, config: LaneformerTPConfig):
+        super().__init__(config)
+        self.padding_idx = config.pad_token_id
+        self.config = config
+        self.vocab_size = config.vocab_size
+        self.n_layers = config.num_hidden_layers
+        self.eos_id = config.eos_token_id
+
+        self.num_lanes = config.num_lanes
+        self.broadcast_delay = config.broadcast_delay
+        self.use_attention_comm = getattr(config, "use_attention_comm", config.use_comm)
+        self.use_mlp_comm = getattr(config, "use_mlp_comm", config.use_comm)
+        self.use_early_comm = config.use_early_comm
+        self.pre_norm_lane_agg = config.pre_norm_lane_agg
+        self.lm_head_type = config.lm_head_type
+        self.replicated_rmsn_scale = config.replicated_rmsn_scale
+        self.swa_layers = set(config.swa_layers)
+        self.rope_dim = config.hidden_size // config.num_attention_heads
+        self.rope_theta = config.rope_theta
+        self.rope_scaling_args = _rope_scaling_from_config(config)
+
+        self.tok_embeddings = nn.Embedding(
+            config.vocab_size, config.hidden_size, self.padding_idx
+        )
+
+        # Build lane blocks
+        self.layers = nn.ModuleList()
+        for layer_id in range(config.num_hidden_layers):
+            if self.broadcast_delay == 0 or layer_id < self.broadcast_delay:
+                attention_reduce_mode = (
+                    ReduceMode.PRESENT if self.use_early_comm else ReduceMode.NO_REDUCE
+                )
+                mlp_reduce_mode = (
+                    ReduceMode.PRESENT if self.use_early_comm else ReduceMode.NO_REDUCE
+                )
+            else:
+                attention_reduce_mode = (
+                    ReduceMode.NO_REDUCE
+                    if not (self.use_attention_comm)
+                    else ReduceMode.PAST
+                )
+                mlp_reduce_mode = (
+                    ReduceMode.NO_REDUCE if not (self.use_mlp_comm) else ReduceMode.PAST
+                )
+
+            broadcast_attention_to_future = (
+                config.num_hidden_layers - self.broadcast_delay > layer_id
+            ) and self.use_attention_comm
+            broadcast_mlp_to_future = (
+                config.num_hidden_layers - self.broadcast_delay > layer_id
+            ) and self.use_mlp_comm
+
+            self.layers.append(
+                LaneformerDecoderLayer(
+                    config=config,
+                    layer_idx=layer_id,
+                    attention_reduce_mode=attention_reduce_mode,
+                    mlp_reduce_mode=mlp_reduce_mode,
+                    broadcast_attention_to_future=broadcast_attention_to_future,
+                    broadcast_mlp_to_future=broadcast_mlp_to_future,
+                )
+            )
+        if self.pre_norm_lane_agg and not self.lm_head_type == "replicate":
+            raise ValueError(
+                "pre_norm_lane_agg is only supported with lm_head_type='replicate'"
+            )
+
+        norm_fun = (
+            nn.RMSNorm
+            if config.replicated_rmsn_scale or config.pre_norm_lane_agg
+            else partial(LaneRMSNorm, num_lanes=self.num_lanes)
+        )
+
+        self.norm = norm_fun(config.hidden_size, eps=config.norm_eps)
+        self.gradient_checkpointing = False
+        self.post_init()
+
+    # Required by vllm integration
+    def get_input_embeddings(self):
+        return self.tok_embeddings
+
+    # Required by vllm integration
+    def set_input_embeddings(self, new_emb):
+        self.tok_embeddings = new_emb
+
+    def _compute_freqs_cis(self, position_ids: torch.Tensor) -> torch.Tensor:
+        return freqs_cis_for_position_ids(
+            position_ids,
+            self.rope_dim,
+            self.rope_theta,
+            self.rope_scaling_args,
+        )
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        **kwargs,
+    ):
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError(
+                "You must specify exactly one of input_ids or inputs_embeds"
+            )
+
+        if inputs_embeds is None:
+            inputs_embeds: torch.Tensor = self.tok_embeddings(input_ids)
+
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+        past_key_values_for_forward = past_key_values if use_cache else None
+
+        if use_cache and past_key_values_for_forward is None:
+            past_key_values_for_forward = DynamicCache(config=self.config)
+
+        if cache_position is None:
+            past_seen_tokens = (
+                past_key_values_for_forward.get_seq_length()
+                if past_key_values_for_forward is not None
+                else 0
+            )
+            cache_position: torch.Tensor = torch.arange(
+                past_seen_tokens,
+                past_seen_tokens + inputs_embeds.shape[1],
+                device=inputs_embeds.device,
+            )
+
+        if position_ids is None:
+            position_ids = cache_position.unsqueeze(0)
+
+        layer_types = list(getattr(self.config, "layer_types", []))
+        if len(layer_types) != self.config.num_hidden_layers:
+            layer_types = [
+                "sliding_attention" if layer_id in self.swa_layers else "full_attention"
+                for layer_id in range(self.config.num_hidden_layers)
+            ]
+        unknown_layer_types = set(layer_types) - {"full_attention", "sliding_attention"}
+        if unknown_layer_types:
+            raise ValueError(
+                f"Unknown attention layer types: {sorted(unknown_layer_types)}"
+            )
+
+        has_sliding_layers = "sliding_attention" in layer_types
+        if self.config.sliding_window is None and (
+            self.swa_layers or has_sliding_layers
+        ):
+            raise ValueError("sliding_window must be set for layers in swa_layers")
+
+        causal_mask_mapping = {}
+        if "full_attention" in layer_types:
+            causal_mask_mapping["full_attention"] = _call_attention_mask_factory(
+                create_causal_mask,
+                uses_inputs_embeds=_CAUSAL_MASK_USES_INPUTS_EMBEDS,
+                config=self.config,
+                inputs_embeds=inputs_embeds,
+                attention_mask=attention_mask,
+                cache_position=cache_position,
+                past_key_values=past_key_values_for_forward,
+                position_ids=position_ids,
+            )
+
+        if has_sliding_layers:
+            sliding_attention_mask = _apply_sliding_window_to_attention_mask(
+                attention_mask,
+                sliding_window=int(self.config.sliding_window),
+            )
+            causal_mask_mapping["sliding_attention"] = _call_attention_mask_factory(
+                create_sliding_window_causal_mask,
+                uses_inputs_embeds=_SLIDING_MASK_USES_INPUTS_EMBEDS,
+                config=self.config,
+                inputs_embeds=inputs_embeds,
+                attention_mask=sliding_attention_mask,
+                cache_position=cache_position,
+                past_key_values=past_key_values_for_forward,
+                position_ids=position_ids,
+            )
+        freqs_cis = self._compute_freqs_cis(position_ids).unsqueeze(-2)
+
+        hidden_states = inputs_embeds
+        hidden_states = hidden_states[:, :, None, :].expand(-1, -1, self.num_lanes, -1)
+        past_attentions = []
+        past_mlps = []
+        for i in range(self.config.num_hidden_layers):
+            if self.broadcast_delay == 0 or i < self.broadcast_delay:
+                past_attention = None
+                past_mlp = None
+            else:
+                past_attention = past_attentions[i - self.broadcast_delay]
+                past_mlp = past_mlps[i - self.broadcast_delay]
+            hidden_states, future_attention, future_mlp = self.layers[i](
+                hidden_states=hidden_states,
+                attention_mask=causal_mask_mapping[layer_types[i]],
+                position_ids=position_ids,
+                past_key_values=past_key_values_for_forward,
+                past_attention=past_attention,
+                past_mlp=past_mlp,
+                use_cache=use_cache,
+                cache_position=cache_position,
+                freqs_cis=freqs_cis,
+                **kwargs,
+            )
+            past_attentions.append(future_attention)
+            past_mlps.append(future_mlp)
+
+        if self.pre_norm_lane_agg:
+            # (B, S, L, D) -> (B, S, D)
+            hidden_states = hidden_states.sum(dim=-2)
+            hidden_states = self.norm(hidden_states)
+        elif self.lm_head_type == "replicate":  # and naturally not pre_lane_agg
+            # (B, S, L, D) -> (B, S, D)
+            hidden_states = self.norm(hidden_states)
+            hidden_states = hidden_states.mean(dim=-2)
+        else:
+            # (B, S, L, D)
+            hidden_states = self.norm(hidden_states)
+
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values_for_forward if use_cache else None,
+        )
+
+
+@auto_docstring
+class LaneformerTPForCausalLM(LaneformerPreTrainedModel, GenerationMixin):
+    _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
+
+    def __init__(self, config: LaneformerTPConfig):
+        super().__init__(config)
+        if config.tie_word_embeddings:
+            raise ValueError("tie_word_embeddings is not supported for LaneformerTP")
+        self.model = LaneformerTPModel(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head_type = config.lm_head_type
+        if self.lm_head_type == "replicate":
+            self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        elif self.lm_head_type == "lane":
+            self.lm_head = LaneLMHead(
+                in_features=config.hidden_size,
+                out_features=config.vocab_size,
+                num_lanes=config.num_lanes,
+            )
+        elif self.lm_head_type == "vocab_parallel":
+            self.lm_head = LaneRowLinear(
+                in_features=config.hidden_size,
+                out_features=config.vocab_size,
+                num_lanes=config.num_lanes,
+            )
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def resize_token_embeddings(
+        self,
+        new_num_tokens: Optional[int] = None,
+        pad_to_multiple_of: Optional[int] = None,
+        mean_resizing: bool = True,
+    ) -> nn.Embedding:
+        if self.lm_head_type != "replicate" and (
+            new_num_tokens is not None or pad_to_multiple_of is not None
+        ):
+            raise NotImplementedError(
+                "resize_token_embeddings is only supported for "
+                "lm_head_type='replicate'"
+            )
+        return super().resize_token_embeddings(
+            new_num_tokens=new_num_tokens,
+            pad_to_multiple_of=pad_to_multiple_of,
+            mean_resizing=mean_resizing,
+        )
+
+    @can_return_tuple
+    @auto_docstring
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        logits_to_keep: Union[int, torch.Tensor] = 0,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> CausalLMOutputWithPast:
+        r"""
+        Example with a tiny randomly initialized LaneformerTP model:
+
+        ```python
+        >>> import torch
+        >>> from torchtitan.models.laneformer_tp.hf import (
+        ...     LaneformerTPConfig,
+        ...     LaneformerTPForCausalLM,
+        ... )
+        >>> config = LaneformerTPConfig(
+        ...     hidden_size=16,
+        ...     num_hidden_layers=2,
+        ...     num_attention_heads=4,
+        ...     num_key_value_heads=2,
+        ...     intermediate_size=32,
+        ...     vocab_size=64,
+        ...     num_lanes=2,
+        ... )
+        >>> model = LaneformerTPForCausalLM(config)
+        >>> input_ids = torch.tensor([[1, 2, 3]])
+        >>> logits = model(input_ids=input_ids, use_cache=False).logits
+        >>> logits.shape
+        torch.Size([1, 3, 64])
+        ```"""
+        outputs: BaseModelOutputWithPast = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            **kwargs,
+        )
+
+        hidden_states = outputs.last_hidden_state
+        # Only compute necessary logits, and do not upcast them to float if
+        # we are not computing the loss.
+        slice_indices = (
+            slice(-logits_to_keep, None)
+            if isinstance(logits_to_keep, int)
+            else logits_to_keep
+        )
+        if self.lm_head_type == "replicate":
+            logits = self.lm_head(hidden_states[:, slice_indices, :])
+        else:
+            logits = self.lm_head(hidden_states[:, slice_indices, :, :])
+            if self.lm_head_type == "vocab_parallel":
+                # (B, S, L, D) -> (B, S, D)
+                B, S, L, D = logits.shape
+                logits = logits.reshape(B, S, L * D)
+
+        loss = None
+        if labels is not None:
+            loss = self.loss_function(
+                logits=logits,
+                labels=labels,
+                vocab_size=self.config.vocab_size,
+                **kwargs,
+            )
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )

special_tokens_map.json

@@ -0,0 +1,24 @@
+{
+  "bos_token": {
+    "content": "<s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "</s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": "</s>",
+  "unk_token": {
+    "content": "<unk>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer_config.json

@@ -0,0 +1,43 @@
+{
+  "add_bos_token": true,
+  "add_eos_token": false,
+  "add_prefix_space": null,
+  "added_tokens_decoder": {
+    "0": {
+      "content": "<unk>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "<s>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "</s>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "bos_token": "<s>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "</s>",
+  "extra_special_tokens": {},
+  "legacy": false,
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "</s>",
+  "padding_side": "right",
+  "sp_model_kwargs": {},
+  "tokenizer_class": "LlamaTokenizer",
+  "unk_token": "<unk>",
+  "use_default_system_prompt": false
+}

torchtitan_metadata.json

@@ -0,0 +1,15 @@
+{
+  "format": "torchtitan_laneformer_tp_metadata",
+  "version": 1,
+  "model_type": "laneformer_tp",
+  "torchtitan_model_args": {
+    "attn_mask_type": "block_causal",
+    "depth_init": true,
+    "multiple_of": 32,
+    "attn_type": "flex"
+  },
+  "deprecated_config_fields": {
+    "sliding_window_n_layers": 10,
+    "use_comm": true
+  }
+}