Publich Laneformer 1B instruct

LICENSE

@@ -1,4 +1,4 @@
-Kog LaneformerTP Repository License Notice
+Kog Laneformer Repository License Notice
 =========================================
 
 This repository contains materials under different licenses.

README.md

@@ -29,7 +29,7 @@ datasets:
   - nvidia/Nemotron-CC-v2
   - nvidia/Nemotron-Pretraining-SFT-v1
   - nvidia/Nemotron-Pretraining-Code-v1
-model_type: laneformer_tp
+model_type: laneformer
 parameters: 2B
 widget:
   - text: "Explain the A* algorithm in details, then write a full Python implementation from scratch."
@@ -43,11 +43,11 @@ widget:
 
 # 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.
+Kog Laneformer 2B is a latency-oriented Transformer variant built using **Delayed Tensor Parallelism (DTP)** 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.
+The model is intended to make our Laneformer architecture available on Hugging Face for research, inspection, and fine-tuning. It is also the architecture family used in our public inference-engine preview, where single-request decoding speeds are **3,000 output tokens/s per request on 8x AMD MI300X** and **2,100 output tokens/s per 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.
+Those figures are **Kog Inference Engine (KIE) benchmark results**, not expected performance from the generic Transformers runtime.
 
 For background, see:
 
@@ -56,59 +56,50 @@ For background, see:
 
 > **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
+### 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.
+Many LLM serving systems optimize aggregate throughput across many concurrent requests. Our Laneformer 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.
+### 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 block typically requires communication before downstream computation can proceed. In **Delayed Tensor Parallelism (DTP)**, local outputs are communicated asynchronously and consumed several modules later, allowing communication latency to be hidden behind subsequent computation and weight streaming.
+
+| Field                         |                                 Value |
+| ----------------------------- | ------------------------------------: |
+| Model family                  |                            Laneformer |
+| HF model type                 |                          `laneformer` |
+| Architecture class            |               `LaneformerForCausalLM` |
+| Task                          | Decoder-only causal language modeling |
+| Parameters                    |                                 ~2.3B |
+| Hidden size                   |                                  3072 |
+| Intermediate size             |                                 12288 |
+| MLP type                      |                                SwiGLU |
+| Decoder layers                |                                    15 |
+| Attention heads               |                                    32 |
+| KV heads                      |                                    16 |
+| Context length                |                           4096 tokens |
+| Sliding window                |                           2048 tokens |
+| Sliding-window layers         |                                   0-9 |
+| Full-attention layers         |                                 10-14 |
+| Vocabulary size               |                                 32000 |
+| Tokenizer                     |                     Llama 2 tokenizer |
+| Number of lanes               |                                     8 |
+| DTP / Broadcast delay         |                                     2 |
+| LM head type                  |                      `vocab_parallel` |
+| RoPE theta                    |                                 10000 |
+| Tied word embeddings          |                                 false |
+| Expected Transformers version |                            `>=4.57.1` |
+| Published weight dtype        |                              BFloat16 |
 
 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
+# Installation
 
-### Nvidia
+## Nvidia
 
 ```bash
 uv venv --python 3.12
@@ -117,7 +108,7 @@ uv pip install torch --index-url https://download.pytorch.org/whl/cu130
 uv pip install "transformers>=4.57.1" accelerate safetensors sentencepiece protobuf
 ```
 
-### AMD
+## AMD
 
 ```bash
 uv venv --python 3.12
@@ -126,19 +117,21 @@ 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
+# Usage
+
+This repository includes `chat_template.jinja`; use `tokenizer.apply_chat_template` for chat and instruction prompts.
 
 ```python
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, set_seed
 
 model_id = "kogai/laneformer-2b-it"
-revision = "main"
+revision = "main" # this model uses remote custom code, pin a reviewed commit hash for reproducible and safer usage
 seed = 42
 temperature = 0.8
 
+print(f"Using random seed {seed}.") 
 set_seed(seed)
-print(f"Using random seed {seed}.")
 
 print(f"Loading tokenizer for {model_id} ({revision=})...")
 tokenizer = AutoTokenizer.from_pretrained(
@@ -146,7 +139,6 @@ tokenizer = AutoTokenizer.from_pretrained(
     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(
@@ -156,6 +148,7 @@ model = AutoModelForCausalLM.from_pretrained(
     dtype="bfloat16",
     device_map="auto",
 )
+model.eval()
 print(f"Model loaded on {model.device}.")
 
 # The Llama 2 tokenizer has no native pad token. If this repo sets PAD to EOS,
@@ -164,19 +157,22 @@ 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.",
+messages = [
+    {
+        "role": "user",
+        "content": "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,
+print(messages[0]["content"])
+print("\nApplying chat template...")
+inputs = tokenizer.apply_chat_template(
+    messages,
+    add_generation_prompt=True,
+    tokenize=True,
     return_tensors="pt",
-    padding=True,
-    truncation=True,
-    max_length=4096,
+    return_dict=True,
 ).to(model.device)
 print(f"Prompt token count: {inputs.input_ids.shape[-1]}")
 
@@ -198,116 +194,144 @@ print("-" * 80)
 print(tokenizer.batch_decode(generated, skip_special_tokens=True)[0])
 ```
 
-### Safer custom-code loading
+## Tokenizer & Chat template
 
-Because this model uses remote custom code, pin a reviewed commit hash for reproducible and safer usage:
+This repository uses the **Llama 2 tokenizer** with a 32,000-token vocabulary. The expected special-token convention is:
 
-```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",
-)
-```
+| Token   |  ID |
+| ------- | --: |
+| `<unk>` |   0 |
+| `<s>`   |   1 |
+| `</s>`  |   2 |
 
-## Tokenizer
+Instruction-tuned models in this family, such as `kogai/laneformer-2b-it`, include a repository-level `chat_template.jinja`.
 
-This repository uses the **Llama 2 tokenizer** with a 32,000-token vocabulary. The expected special-token convention is:
+For the current instruction-tuned template:
+
+- Only `user` and `assistant` roles are supported.
+- System messages are ignored.
+- Each user or assistant message is formatted with a Llama-3-style role header, followed by trimmed message content and `</s>`.
+
+Usage:
+
+```python
+messages = [
+    {
+        "role": "user",
+        "content": "Explain how a binary heap works, then write a complete min-heap implementation from scratch in Python.",
+    },
+]
+
+inputs = tokenizer.apply_chat_template(
+    messages,
+    add_generation_prompt=True,
+    tokenize=True,
+    return_tensors="pt",
+    return_dict=True,
+).to(model.device)
+```
 
-| Token | ID |
-|---|---:|
-| `<unk>` | 0 |
-| `<s>` | 1 |
-| `</s>` | 2 |
+# Training data and procedure
 
-The model config uses `eos_token_id = 2` and `pad_token_id = 2`.
+Laneformer 2B was trained with a staged language-modeling recipe:
 
-## Training data and procedure
+1. Broad pre-training to build general language-modeling capabilities.
+2. Code- and reasoning-heavy continuation training.
+3. Lightweight post-training to produce the released instruction-tuned model.
 
-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.
+The TorchTitan training stages used sequence length 4,096 and global batch size 1,536, or approximately 6.29M tokens per optimizer step.
 
-### Training stages
+## 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. |
+| Stage         | Focus                                    | Framework                 | Tokens |   Steps | Approx. tokens / step | Notes                                                                           |
+| ------------- | ---------------------------------------- | ------------------------- | -----: | ------: | --------------------: | ------------------------------------------------------------------------------- |
+| Pre-training  | Broad language-model pre-training        | TorchTitan                |    ~4T | 620,000 |                ~6.29M | Nemotron-derived broad mixture; final checkpoint selected.                      |
+| Mid-training  | Code- and reasoning-focused continuation | TorchTitan                |    ~2T | 310,000 |                ~6.29M | Continued from the pre-training checkpoint with a code/reasoning-heavy mixture. |
+| Post-training | SFT, instruction tuning, identity tuning | Hugging Face Transformers |  ~210M |     200 |                ~1.05M | Lightweight custom data mixture initialized from the mid-training checkpoint.   |
 
-### Data sources
+## 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
+| Category                    | Pre-training / Phase 1 | Mid-training / Phase 2 | Notes                                                                                |
+| --------------------------- | ---------------------: | ---------------------: | ------------------------------------------------------------------------------------ |
+| General / web crawl         |                  63.2% |                  10.0% | Phase 2 keeps a smaller high-quality and synthetic web component.                    |
+| Code                        |                  21.0% |                 59.01% | Large increase from code metadata, synthetic code data, and Common Crawl code pages. |
+| Math / STEM / reasoning     |                   6.3% |                 30.01% | Large increase from math, RQA, STEM, and textbook-like data.                         |
+| Multilingual                |                   5.0% |                   0.0% | Removed in the current Phase 2 mixture.                                              |
+| QA / instruction / academic |                   4.5% |                   1.0% | Only a small general SFT-style component remains.                                    |
+
+## 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                     | FSDP                                          |
+| 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 |
+| 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
+The following internal Kog evaluations were run in June 2026 against a local Kog serving endpoint for the Laneformer 2B instruction-tuned checkpoint at batch size 1 in FP16. These values evaluate the served preview model; small numerical differences may appear when evaluating the BF16 Hugging Face checkpoint directly through a different runtime.
+
+## Code generation
+
+HumanEval+ and MBPP+ were evaluated with greedy decoding. Generation was performed through the Kog serving endpoint, scoring used EvalPlus.
+
+A custom code-block selection step named `target_function` was applied before scoring. When possible, this step selects the code block containing the target function name before EvalPlus preprocessing.
+
+| Benchmark  | Metric | Value | Samples | Decoding                                 | Scoring  | Postprocessing                    |
+| ---------- | -----: | ----: | ------: | ---------------------------------------- | -------- | --------------------------------- |
+| HumanEval+ | pass@1 |  45.1 |     164 | Greedy, temperature=0, `do_sample=False` | EvalPlus | `target_function` block selection |
+| MBPP+      | pass@1 |  51.6 |     378 | Greedy, temperature=0, `do_sample=False` | EvalPlus | `target_function` block selection |
+
+## General multiple-choice checks
+
+ARC-Challenge and ARC-Easy were evaluated with 0-shot multiple-choice logprobs. Orchestration used `lm-evaluation-harness` 0.4.12 against the Kog serving endpoint.
+
+| Benchmark     |              Metric | Value | Samples | Setting                                    |
+| ------------- | ------------------: | ----: | ------: | ------------------------------------------ |
+| ARC-Challenge | Normalized accuracy | 31.06 |   1,172 | 0-shot multiple-choice logprobs, `lm_eval` |
+| ARC-Easy      | Normalized accuracy | 47.90 |   2,376 | 0-shot multiple-choice logprobs, `lm_eval` |
+
+## Long-context synthetic checks
+
+Long-context checks were also run with RULER-style synthetic tasks at 2,048 and 4,096 tokens using 0-shot greedy chat generation using the LM Evaluation Harness framework. Values below are string-match scores.
+
+| Task                            | 2,048 tokens | 4,096 tokens | Effective samples |         Setting |
+| ------------------------------- | -----------: | -----------: | ----------------: | --------------: |
+| NIAH single 1                   |       100.00 |       100.00 |               500 | using `lm_eval` |
+| NIAH single 2                   |       100.00 |       100.00 |               500 | using `lm_eval` |
+| NIAH single 3                   |        99.80 |        91.60 |               500 | using `lm_eval` |
+| NIAH multikey 1                 |        99.20 |        83.40 |               500 | using `lm_eval` |
+| NIAH multikey 2                 |        97.80 |        60.00 |               500 | using `lm_eval` |
+| NIAH multikey 3                 |        91.80 |        91.20 |               500 | using `lm_eval` |
+| NIAH multiquery                 |        95.40 |        69.70 |               500 | using `lm_eval` |
+| NIAH multivalue                 |        94.60 |        82.20 |               500 | using `lm_eval` |
+| RULER variable tracking         |        77.92 |         3.92 |               500 | using `lm_eval` |
+| RULER common words extraction   |        24.72 |        78.14 |               500 | using `lm_eval` |
+| RULER frequent words extraction |        78.80 |        64.73 |               500 | using `lm_eval` |
+
+# Intended use
 
 This model is intended for:
 
@@ -319,14 +343,16 @@ This model is intended for:
 
 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
+# 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.
+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.
 
-## Limitations and risks
+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 is primarily mid-trained and post-trained for code generation: it is not optimized as a broad general-purpose assistant.
 - 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.
@@ -334,17 +360,17 @@ Do not use this model or tokenizer in ways that violate the repository licenses,
 - 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
+# License
 
 This repository is **multi-license**.
 
-### Kog-owned materials: Apache License 2.0
+## 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
+## 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:
 
@@ -362,19 +388,27 @@ Required Llama 2 attribution notice:
 
 Users and redistributors are responsible for complying with the Llama 2 Community License and its Acceptable Use Policy.
 
-## Citation
+# Citation
 
-If you use this model or architecture, please cite the model repository and the Kog technical blog posts:
+If you use this model or architecture, please cite the model repository and the relevant Kog technical posts:
 
 ```bibtex
-@misc{kog_laneformer_tp_2b_2026,
-  title        = {Kog Laneformer 2B},
-  author       = {Kog AI},
+@misc{kog_laneformer_2b_it_2026,
+  title        = {Kog Laneformer 2B Instruct Model},
+  author       = {Kog Team},
   year         = {2026},
   publisher    = {Hugging Face},
   howpublished = {\url{https://huggingface.co/kogai/laneformer-2b-it}}
 }
 
+@online{kog_laneformer_2b_2026,
+  title        = {{Laneformer 2B: The Latency-First Model Behind Kog Inference Engine}},
+  author       = {Kog Team},
+  year         = {2026},
+  url          = {https://huggingface.co/blog/kogai/kog-laneformer-2b-the-latency-first-model}
+  note         = {HuggingFace blog}
+}
+
 @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},
@@ -390,6 +424,6 @@ If you use this model or architecture, please cite the model repository and the
 }
 ```
 
-## Contact
+# 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/).

__init__.py

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

config.json

@@ -1,20 +1,20 @@
 {
-  "model_type": "laneformer_tp",
+  "model_type": "laneformer",
   "architectures": [
-    "LaneformerTPForCausalLM"
+    "LaneformerForCausalLM"
   ],
   "auto_map": {
-    "AutoConfig": "configuration_laneformer.LaneformerTPConfig",
-    "AutoModel": "modeling_laneformer.LaneformerTPModel",
-    "AutoModelForCausalLM": "modeling_laneformer.LaneformerTPForCausalLM"
+    "AutoConfig": "configuration_laneformer.LaneformerConfig",
+    "AutoModel": "modeling_laneformer.LaneformerModel",
+    "AutoModelForCausalLM": "modeling_laneformer.LaneformerForCausalLM"
   },
   "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,
+  "max_position_embeddings": 4096,
+  "max_seq_len": 4096,
   "eos_token_id": 2,
   "pad_token_id": 2,
   "vocab_size": 32000,
@@ -60,7 +60,7 @@
     "full_attention",
     "full_attention"
   ],
-  "_attn_implementation": "flex_attention",
+  "_attn_implementation": "sdpa",
   "dtype": "bfloat16",
   "transformers_version": "4.57.6"
 }

configuration_laneformer.py

@@ -121,8 +121,8 @@ def _validate_rope_scaling(rope_scaling: Optional[dict[str, Any]]) -> None:
         raise ValueError("rope_scaling.original_max_position_embeddings must be > 0")
 
 
-class LaneformerTPConfig(PretrainedConfig):
-    model_type = "laneformer_tp"
+class LaneformerConfig(PretrainedConfig):
+    model_type = "laneformer"
     keys_to_ignore_at_inference = ["past_key_values"]
 
     def __init__(
@@ -228,7 +228,7 @@ class LaneformerTPConfig(PretrainedConfig):
         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")
+            raise ValueError("tie_word_embeddings is not supported for Laneformer")
         _validate_rope_scaling(rope_scaling)
         if swa_layers is None:
             swa_layers = list(range(sliding_window_n_layers))

modeling_laneformer.py

@@ -2,9 +2,9 @@
 Hugging Face-compatible Laneformer implementation.
 
 This file defines:
-- LaneformerTPPreTrainedModel
-- LaneformerTPModel (backbone)
-- LaneformerTPForCausalLM (head + loss)
+- LaneformerPreTrainedModel
+- LaneformerModel (backbone)
+- LaneformerForCausalLM (head + loss)
 
 It mirrors the TorchTitan Laneformer structure so that state_dict keys largely match,
 which makes converting checkpoints trivial (often identity mapping).
@@ -35,7 +35,7 @@ from transformers.modeling_outputs import (
 from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
 from transformers.utils import auto_docstring, can_return_tuple, TransformersKwargs
 
-from .configuration_laneformer import LaneformerTPConfig
+from .configuration_laneformer import LaneformerConfig
 
 
 # ============================================================
@@ -50,7 +50,7 @@ class RoPEScalingArgs:
 
 
 def _rope_scaling_from_config(
-    config: LaneformerTPConfig,
+    config: LaneformerConfig,
 ) -> RoPEScalingArgs | None:
     if config.rope_scaling is None:
         return None
@@ -295,7 +295,7 @@ def eager_attention_forward(
     return attn_output, attn_weights
 
 
-def _attention_interface(config: LaneformerTPConfig):
+def _attention_interface(config: LaneformerConfig):
     attn_implementation = getattr(config, "_attn_implementation", "eager")
     if attn_implementation == "eager":
         return eager_attention_forward
@@ -319,7 +319,7 @@ def _call_attention_mask_factory(
     mask_factory,
     *,
     uses_inputs_embeds: bool,
-    config: LaneformerTPConfig,
+    config: LaneformerConfig,
     inputs_embeds: torch.Tensor,
     attention_mask: Optional[torch.Tensor],
     cache_position: torch.Tensor,
@@ -386,7 +386,7 @@ def _apply_sliding_window_to_attention_mask(
 
 
 def _layer_uses_sliding_window_attention(
-    config: LaneformerTPConfig,
+    config: LaneformerConfig,
     layer_index: int,
 ) -> bool:
     sliding_window = getattr(config, "sliding_window", None)
@@ -409,7 +409,7 @@ def _layer_uses_sliding_window_attention(
 
 
 def _layer_attention_masks_from_provided_mask(
-    config: LaneformerTPConfig,
+    config: LaneformerConfig,
     attention_mask: Optional[torch.Tensor],
 ) -> Optional[list[Optional[torch.Tensor]]]:
     if not _attention_mask_needs_sliding_window_overlay(attention_mask):
@@ -471,7 +471,7 @@ class LaneModule:
 class LaneformerAttention(LaneModule, nn.Module):
     def __init__(
         self,
-        config: LaneformerTPConfig,
+        config: LaneformerConfig,
         layer_idx: int,
         reduce_mode: ReduceMode = ReduceMode.PRESENT,
     ):
@@ -580,7 +580,7 @@ class LaneformerAttention(LaneModule, nn.Module):
 
 class LaneformerMLP(LaneModule, nn.Module):
     def __init__(
-        self, config: LaneformerTPConfig, reduce_mode: ReduceMode = ReduceMode.PRESENT
+        self, config: LaneformerConfig, reduce_mode: ReduceMode = ReduceMode.PRESENT
     ):
         super().__init__(
             no_reduce_scale=math.sqrt(config.num_lanes), reduce_mode=reduce_mode
@@ -600,7 +600,7 @@ class LaneformerMLP(LaneModule, nn.Module):
 class LaneformerDecoderLayer(nn.Module):
     def __init__(
         self,
-        config: LaneformerTPConfig,
+        config: LaneformerConfig,
         layer_idx: int,
         attention_reduce_mode: ReduceMode,
         mlp_reduce_mode: ReduceMode,
@@ -667,12 +667,12 @@ class LaneformerDecoderLayer(nn.Module):
 # Backbone & HF wrappers
 # ============================================================
 class LaneformerPreTrainedModel(PreTrainedModel):
-    config_class = LaneformerTPConfig
-    config: LaneformerTPConfig
+    config_class = LaneformerConfig
+    config: LaneformerConfig
     base_model_prefix = "model"
     supports_gradient_checkpointing = False
     _no_split_modules = [
-        "LaneformerTPModel",
+        "LaneformerModel",
         "LaneformerDecoderLayer",
         "LaneformerAttention",
         "LaneformerMLP",
@@ -681,11 +681,10 @@ class LaneformerPreTrainedModel(PreTrainedModel):
     # 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):
+class LaneformerModel(LaneformerPreTrainedModel):
+    def __init__(self, config: LaneformerConfig):
         super().__init__(config)
         self.padding_idx = config.pad_token_id
         self.config = config
@@ -911,14 +910,14 @@ class LaneformerTPModel(LaneformerPreTrainedModel):
 
 
 @auto_docstring
-class LaneformerTPForCausalLM(LaneformerPreTrainedModel, GenerationMixin):
+class LaneformerForCausalLM(LaneformerPreTrainedModel, GenerationMixin):
     _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
 
-    def __init__(self, config: LaneformerTPConfig):
+    def __init__(self, config: LaneformerConfig):
         super().__init__(config)
         if config.tie_word_embeddings:
-            raise ValueError("tie_word_embeddings is not supported for LaneformerTP")
-        self.model = LaneformerTPModel(config)
+            raise ValueError("tie_word_embeddings is not supported for Laneformer")
+        self.model = LaneformerModel(config)
         self.vocab_size = config.vocab_size
         self.lm_head_type = config.lm_head_type
         if self.lm_head_type == "replicate":
@@ -979,15 +978,15 @@ class LaneformerTPForCausalLM(LaneformerPreTrainedModel, GenerationMixin):
         **kwargs: Unpack[TransformersKwargs],
     ) -> CausalLMOutputWithPast:
         r"""
-        Example with a tiny randomly initialized LaneformerTP model:
+        Example with a tiny randomly initialized Laneformer model:
 
         ```python
         >>> import torch
-        >>> from torchtitan.models.laneformer_tp.hf import (
-        ...     LaneformerTPConfig,
-        ...     LaneformerTPForCausalLM,
+        >>> from torchtitan.models.laneformer.hf import (
+        ...     LaneformerConfig,
+        ...     LaneformerForCausalLM,
         ... )
-        >>> config = LaneformerTPConfig(
+        >>> config = LaneformerConfig(
         ...     hidden_size=16,
         ...     num_hidden_layers=2,
         ...     num_attention_heads=4,
@@ -996,7 +995,7 @@ class LaneformerTPForCausalLM(LaneformerPreTrainedModel, GenerationMixin):
         ...     vocab_size=64,
         ...     num_lanes=2,
         ... )
-        >>> model = LaneformerTPForCausalLM(config)
+        >>> model = LaneformerForCausalLM(config)
         >>> input_ids = torch.tensor([[1, 2, 3]])
         >>> logits = model(input_ids=input_ids, use_cache=False).logits
         >>> logits.shape

tokenizer_config.json

@@ -32,12 +32,12 @@
   "clean_up_tokenization_spaces": false,
   "eos_token": "</s>",
   "extra_special_tokens": {},
-  "legacy": false,
+  "legacy": true,
   "model_max_length": 1000000000000000019884624838656,
   "pad_token": "</s>",
   "padding_side": "right",
   "sp_model_kwargs": {},
-  "tokenizer_class": "LlamaTokenizer",
+  "tokenizer_class": "LlamaTokenizerFast",
   "unk_token": "<unk>",
   "use_default_system_prompt": false
 }

torchtitan_metadata.json

@@ -1,7 +1,7 @@
 {
-  "format": "torchtitan_laneformer_tp_metadata",
+  "format": "torchtitan_laneformer_metadata",
   "version": 1,
-  "model_type": "laneformer_tp",
+  "model_type": "laneformer",
   "torchtitan_model_args": {
     "attn_mask_type": "block_causal",
     "depth_init": true,