diff --git a/reference/README.md b/reference/README.md
new file mode 100644
index 00000000..0d3b9037
--- /dev/null
+++ b/reference/README.md
@@ -0,0 +1,49 @@
+# Reference Implementations
+
+This directory contains **read-only** reference implementations from official sources.
+Do not modify these files — they exist to cross-check our production pipeline.
+
+## Directory Layout
+
+```
+reference/
+├── vllm/ # vLLM project reference (Apache-2.0)
+│ ├── tokenizers/
+│ │ ├── deepseek_v4.py # Tokenizer wrapper — apply_chat_template for DSV4
+│ │ └── deepseek_v4_encoding.py # Official prompt encoder (canonical source)
+│ ├── reasoning/
+│ │ ├── deepseek_v3_reasoning_parser.py # Thinking-mode dispatcher
+│ │ └── deepseek_r1_reasoning_parser.py # ）/） reasoning token parser
+│ └── tool_parsers/
+│ ├── deepseekv4_tool_parser.py # DSML tool call parser (V4)
+│ └── deepseekv32_tool_parser.py # DSML tool call parser (V3.2 base)
+│
+└── official_inference/ # Original weight's reference inference code
+ ├── generate.py # Official generate loop + encode_messages usage
+ ├── model.py # BF16/FP8 model implementation
+ ├── kernel.py # Reference CUDA kernels
+ ├── convert.py # Weight conversion
+ └── config.json # Model config (small variant)
+```
+
+## Key Files for Our Pipeline
+
+1. **`vllm/tokenizers/deepseek_v4_encoding.py`** — Canonical prompt encoder.
+ Already copied to `encoding/deepseek_v4_encoding.py` in the repo root (our live import).
+ If vLLM updates this file, diff and sync.
+
+2. **`vllm/tokenizers/deepseek_v4.py`** — Shows how vLLM wraps the tokenizer
+ to add `apply_chat_template` support. Key insight: it calls
+ `encode_messages(messages, thinking_mode=..., ...)` then
+ `tokenizer.encode(prompt_str, add_special_tokens=False)`.
+ This is exactly what our single_shot does.
+
+3. **`official_inference/generate.py`** — The original weight's inference entry point.
+ Uses `tokenizer.encode(encode_messages(messages, thinking_mode="chat"))`
+ (default `add_special_tokens=True`) and `parse_message_from_completion_text()`
+ for output parsing.
+
+4. **`vllm/reasoning/`** — How vLLM detects thinking mode boundaries
+ (`）、` start, `）/）` end). Useful when we integrate streaming.
+
+5. **`vllm/tool_parsers/`** — DSML tool call parsing for future tool-use support.
diff --git a/reference/official_inference/__init__.py b/reference/official_inference/__init__.py
new file mode 100644
index 00000000..8fdcc9db
--- /dev/null
+++ b/reference/official_inference/__init__.py
@@ -0,0 +1 @@
+# Official inference reference — read only, do not modify
diff --git a/reference/official_inference/config.json b/reference/official_inference/config.json
new file mode 100644
index 00000000..e2196678
--- /dev/null
+++ b/reference/official_inference/config.json
@@ -0,0 +1,35 @@
+{
+ "vocab_size": 129280,
+ "dim": 7168,
+ "moe_inter_dim": 3072,
+ "n_layers": 61,
+ "n_hash_layers": 3,
+ "n_heads": 128,
+ "n_routed_experts": 384,
+ "n_shared_experts": 1,
+ "n_activated_experts": 6,
+ "score_func": "sqrtsoftplus",
+ "route_scale": 2.5,
+ "swiglu_limit": 10.0,
+ "q_lora_rank": 1536,
+ "head_dim": 512,
+ "rope_head_dim": 64,
+ "o_groups": 16,
+ "o_lora_rank": 1024,
+ "window_size": 128,
+ "original_seq_len": 65536,
+ "rope_theta": 10000,
+ "rope_factor": 16,
+ "beta_fast": 32,
+ "beta_slow": 1,
+ "index_n_heads": 64,
+ "index_head_dim": 128,
+ "index_topk": 1024,
+ "hc_mult": 4,
+ "hc_sinkhorn_iters": 20,
+ "dtype": "fp8",
+ "scale_fmt": "ue8m0",
+ "expert_dtype": "fp4",
+ "compress_rope_theta": 160000,
+ "compress_ratios": [128, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 0]
+}
\ No newline at end of file
diff --git a/reference/official_inference/convert.py b/reference/official_inference/convert.py
new file mode 100644
index 00000000..77748646
--- /dev/null
+++ b/reference/official_inference/convert.py
@@ -0,0 +1,168 @@
+import os
+import shutil
+from argparse import ArgumentParser
+from glob import glob
+from tqdm import tqdm, trange
+
+import torch
+from safetensors.torch import safe_open, save_file
+
+
+FP4_TABLE = torch.tensor([
+ 0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0,
+ 0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0
+], dtype=torch.float32)
+
+
+def cast_e2m1fn_to_e4m3fn(x: torch.Tensor, scale: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+ """
+ Casts a tensor from e2m1fn to e4m3fn losslessly.
+ """
+ assert x.dtype == torch.int8
+ assert x.ndim == 2
+ out_dim, in_dim = x.size()
+ in_dim *= 2
+ fp8_block_size = 128
+ fp4_block_size = 32
+ assert in_dim % fp8_block_size == 0 and out_dim % fp8_block_size == 0
+ assert scale.size(0) == out_dim and scale.size(1) == in_dim // fp4_block_size
+
+ x = x.view(torch.uint8)
+ low = x & 0x0F
+ high = (x >> 4) & 0x0F
+ x = torch.stack([FP4_TABLE[low.long()], FP4_TABLE[high.long()]], dim=-1).flatten(2)
+
+ # max_fp4 (6.0) * MAX_OFFSET must fit in e4m3fn (max 448)
+ # 6.0 * 2^6 = 384 < 448; 6.0 * 2^7 = 768 > 448; so MAX_OFFSET_BITS = 6
+ MAX_OFFSET_BITS = 6
+
+ bOut = out_dim // fp8_block_size
+ bIn = in_dim // fp8_block_size
+ # bOut, bIn, 128, 128
+ x = x.view(bOut, fp8_block_size, bIn, fp8_block_size).transpose(1, 2)
+ # bOut, bIn, 128*4
+ scale = scale.float().view(bOut, fp8_block_size, bIn, -1).transpose(1, 2).flatten(2)
+ ## bOut, bIn, 1
+ scale_max_offset_bits = scale.amax(dim=-1, keepdim=True) / (2**MAX_OFFSET_BITS)
+ # bOut, bIn, 128*4
+ offset = scale / scale_max_offset_bits
+ # bOut, bIn, 128, 128
+ offset = offset.unflatten(-1, (fp8_block_size, -1)).repeat_interleave(fp4_block_size, dim=-1)
+ x = (x * offset).transpose(1, 2).reshape(out_dim, in_dim)
+ return x.to(torch.float8_e4m3fn), scale_max_offset_bits.squeeze(-1).to(torch.float8_e8m0fnu)
+
+
+mapping = {
+ "embed_tokens": ("embed", 0),
+ "input_layernorm": ("attn_norm", None),
+ "post_attention_layernorm": ("ffn_norm", None),
+ "q_proj": ("wq", 0),
+ "q_a_proj": ("wq_a", None),
+ "q_a_layernorm": ("q_norm", None),
+ "q_b_proj": ("wq_b", 0),
+ "kv_a_proj_with_mqa": ("wkv_a", None),
+ "kv_a_layernorm": ("kv_norm", None),
+ "kv_b_proj": ("wkv_b", 0),
+ "o_proj": ("wo", 1),
+ "gate_proj": ("w1", 0),
+ "down_proj": ("w2", 1),
+ "up_proj": ("w3", 0),
+ "lm_head": ("head", 0),
+
+ "embed": ("embed", 0),
+ "wq_b": ("wq_b", 0),
+ "wo_a": ("wo_a", 0),
+ "wo_b": ("wo_b", 1),
+ "head": ("head", 0),
+ "attn_sink": ("attn_sink", 0),
+ "weights_proj": ("weights_proj", 0),
+}
+
+
+def main(hf_ckpt_path, save_path, n_experts, mp, expert_dtype):
+ """
+ Converts and saves model checkpoint files into a specified format.
+
+ Args:
+ hf_ckpt_path (str): Path to the directory containing the input checkpoint files.
+ save_path (str): Path to the directory where the converted checkpoint files will be saved.
+ n_experts (int): Total number of experts in the model.
+ mp (int): Model parallelism factor.
+
+ Returns:
+ None
+ """
+ torch.set_num_threads(8)
+ n_local_experts = n_experts // mp
+ state_dicts = [{} for _ in range(mp)]
+
+ for file_path in tqdm(glob(os.path.join(hf_ckpt_path, "*.safetensors"))):
+ with safe_open(file_path, framework="pt", device="cpu") as f:
+ for name in f.keys():
+ param: torch.Tensor = f.get_tensor(name)
+ if name.startswith("model."):
+ name = name[len("model."):]
+ if name.startswith("mtp.") and ("emb" in name or name.endswith("head.weight")):
+ continue
+ name = name.replace("self_attn", "attn")
+ name = name.replace("mlp", "ffn")
+ name = name.replace("weight_scale_inv", "scale")
+ name = name.replace("e_score_correction_bias", "bias")
+ if any(x in name for x in ["hc", "attn_sink", "tie2eid", "ape"]): # without .weight
+ key = name.split(".")[-1]
+ else:
+ key = name.split(".")[-2]
+ if key in mapping:
+ new_key, dim = mapping[key]
+ else:
+ new_key, dim = key, None
+ name = name.replace(key, new_key)
+ for i in range(mp):
+ new_param = param
+ if "experts" in name and "shared_experts" not in name:
+ idx = int(name.split(".")[-3])
+ if idx < i * n_local_experts or idx >= (i + 1) * n_local_experts:
+ continue
+ elif dim is not None:
+ assert param.size(dim) % mp == 0, f"Dimension {dim} must be divisible by {mp}"
+ shard_size = param.size(dim) // mp
+ new_param = param.narrow(dim, i * shard_size, shard_size).contiguous()
+ state_dicts[i][name] = new_param
+
+ os.makedirs(save_path, exist_ok=True)
+
+ for i in trange(mp):
+ names = list(state_dicts[i].keys())
+ for name in names:
+ if name.endswith("wo_a.weight"):
+ weight = state_dicts[i][name]
+ scale = state_dicts[i].pop(name.replace("weight", "scale"))
+ weight = weight.unflatten(0, (-1, 128)).unflatten(-1, (-1, 128)).float() * scale[:, None, :, None].float()
+ state_dicts[i][name] = weight.flatten(2, 3).flatten(0, 1).bfloat16()
+ elif "experts" in name and state_dicts[i][name].dtype == torch.int8:
+ if expert_dtype == "fp8":
+ scale_name = name.replace("weight", "scale")
+ weight = state_dicts[i].pop(name)
+ scale = state_dicts[i].pop(scale_name)
+ state_dicts[i][name], state_dicts[i][scale_name] = cast_e2m1fn_to_e4m3fn(weight, scale)
+ else:
+ state_dicts[i][name] = state_dicts[i][name].view(torch.float4_e2m1fn_x2)
+ save_file(state_dicts[i], os.path.join(save_path, f"model{i}-mp{mp}.safetensors"))
+
+ for file in ["tokenizer.json", "tokenizer_config.json"]:
+ old_file_path = os.path.join(hf_ckpt_path, file)
+ new_file_path = os.path.join(save_path, file)
+ if os.path.exists(old_file_path):
+ shutil.copyfile(old_file_path, new_file_path)
+
+
+if __name__ == "__main__":
+ parser = ArgumentParser()
+ parser.add_argument("--hf-ckpt-path", type=str, required=True)
+ parser.add_argument("--save-path", type=str, required=True)
+ parser.add_argument("--n-experts", type=int, required=True)
+ parser.add_argument("--model-parallel", type=int, required=True)
+ parser.add_argument("--expert-dtype", type=str, choices=["fp8", "fp4"], required=False, default=None)
+ args = parser.parse_args()
+ assert args.n_experts % args.model_parallel == 0, "Number of experts must be divisible by model parallelism"
+ main(args.hf_ckpt_path, args.save_path, args.n_experts, args.model_parallel, args.expert_dtype)
diff --git a/reference/official_inference/generate.py b/reference/official_inference/generate.py
new file mode 100644
index 00000000..c35c8030
--- /dev/null
+++ b/reference/official_inference/generate.py
@@ -0,0 +1,155 @@
+import os
+import json
+import sys
+from argparse import ArgumentParser
+from typing import List
+
+import torch
+import torch.distributed as dist
+from transformers import AutoTokenizer
+from safetensors.torch import load_model
+
+from model import Transformer, ModelArgs
+current_dir = os.path.dirname(os.path.abspath(__file__))
+encoding_dir = os.path.join(current_dir, '../encoding')
+sys.path.insert(0, os.path.abspath(encoding_dir))
+from encoding_dsv4 import encode_messages, parse_message_from_completion_text
+
+
+def sample(logits, temperature: float = 1.0):
+ """Gumbel-max trick: equivalent to multinomial sampling but faster on GPU,
+ since it avoids the GPU-to-CPU sync in torch.multinomial."""
+ logits = logits / max(temperature, 1e-5)
+ probs = torch.softmax(logits, dim=-1, dtype=torch.float32)
+ return probs.div_(torch.empty_like(probs).exponential_(1)).argmax(dim=-1)
+
+
+@torch.inference_mode()
+def generate(
+ model: Transformer,
+ prompt_tokens: List[List[int]],
+ max_new_tokens: int,
+ eos_id: int,
+ temperature: float = 1.0
+) -> List[List[int]]:
+ """Batch generation with left-padded prompts.
+
+ The first forward pass processes [min_prompt_len:] tokens (prefill phase).
+ Subsequent passes generate one token at a time (decode phase). For positions
+ still within a prompt, the ground-truth token overrides the model's prediction.
+ """
+ prompt_lens = [len(t) for t in prompt_tokens]
+ assert max(prompt_lens) <= model.max_seq_len, f"Prompt length exceeds model maximum sequence length (max_seq_len={model.max_seq_len})"
+ total_len = min(model.max_seq_len, max_new_tokens + max(prompt_lens))
+ tokens = torch.full((len(prompt_tokens), total_len), -1, dtype=torch.long)
+ for i, t in enumerate(prompt_tokens):
+ tokens[i, :len(t)] = torch.tensor(t, dtype=torch.long)
+ prev_pos = 0
+ finished = torch.tensor([False] * len(prompt_tokens))
+ prompt_mask = tokens != -1
+ for cur_pos in range(min(prompt_lens), total_len):
+ logits = model.forward(tokens[:, prev_pos:cur_pos], prev_pos)
+ if temperature > 0:
+ next_token = sample(logits, temperature)
+ else:
+ next_token = logits.argmax(dim=-1)
+ next_token = torch.where(prompt_mask[:, cur_pos], tokens[:, cur_pos], next_token)
+ tokens[:, cur_pos] = next_token
+ finished |= torch.logical_and(~prompt_mask[:, cur_pos], next_token == eos_id)
+ prev_pos = cur_pos
+ if finished.all():
+ break
+ completion_tokens = []
+ for i, toks in enumerate(tokens.tolist()):
+ toks = toks[prompt_lens[i]:prompt_lens[i]+max_new_tokens]
+ if eos_id in toks:
+ toks = toks[:toks.index(eos_id)]
+ toks.append(eos_id)
+ completion_tokens.append(toks)
+ return completion_tokens
+
+
+def main(
+ ckpt_path: str,
+ config: str,
+ input_file: str = "",
+ interactive: bool = True,
+ max_new_tokens: int = 100,
+ temperature: float = 1.0,
+) -> None:
+ world_size = int(os.getenv("WORLD_SIZE", "1"))
+ rank = int(os.getenv("RANK", "0"))
+ local_rank = int(os.getenv("LOCAL_RANK", "0"))
+ if world_size > 1:
+ dist.init_process_group("nccl")
+ global print
+ if rank != 0:
+ print = lambda *_, **__: None
+ torch.cuda.set_device(local_rank)
+ torch.cuda.memory._set_allocator_settings("expandable_segments:True")
+ torch.set_default_dtype(torch.bfloat16)
+ torch.set_num_threads(8)
+ torch.manual_seed(33377335)
+ with open(config) as f:
+ args = ModelArgs(**json.load(f))
+ if interactive:
+ args.max_batch_size = 1
+ print(args)
+ with torch.device("cuda"):
+ model = Transformer(args)
+ tokenizer = AutoTokenizer.from_pretrained(ckpt_path)
+ print("load model")
+ load_model(model, os.path.join(ckpt_path, f"model{rank}-mp{world_size}.safetensors"), strict=False)
+ torch.set_default_device("cuda")
+ print("I'm DeepSeek 👋")
+
+ if interactive:
+ messages = []
+ while True:
+ if world_size == 1:
+ prompt = input(">>> ")
+ elif rank == 0:
+ prompt = input(">>> ")
+ objects = [prompt]
+ dist.broadcast_object_list(objects, 0)
+ else:
+ objects = [None]
+ dist.broadcast_object_list(objects, 0)
+ prompt = objects[0]
+ if prompt == "/exit":
+ break
+ elif prompt == "/clear":
+ messages.clear()
+ continue
+ messages.append({"role": "user", "content": prompt})
+ prompt_tokens = tokenizer.encode(encode_messages(messages, thinking_mode="chat"))
+ completion_tokens = generate(model, [prompt_tokens], max_new_tokens, tokenizer.eos_token_id, temperature)
+ completion = tokenizer.decode(completion_tokens[0])
+ print(completion)
+ messages.append(parse_message_from_completion_text(completion, thinking_mode="chat"))
+ else:
+ with open(input_file) as f:
+ prompts = f.read().split("\n\n")
+ prompt_tokens = [tokenizer.encode(encode_messages([{"role": "user", "content": prompt}], thinking_mode="chat")) for prompt in prompts]
+ completion_tokens = generate(model, prompt_tokens, max_new_tokens, tokenizer.eos_token_id, temperature)
+ completions = tokenizer.batch_decode(completion_tokens)
+ for prompt, completion in zip(prompts, completions):
+ print("Prompt:", prompt)
+ print("Completion:", completion)
+ print()
+
+ if world_size > 1:
+ dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+ parser = ArgumentParser()
+ parser.add_argument("--ckpt-path", type=str, required=True)
+ parser.add_argument("--config", type=str, required=True)
+ parser.add_argument("--input-file", type=str, default="")
+ parser.add_argument("--interactive", action="store_true")
+ parser.add_argument("--max-new-tokens", type=int, default=300)
+ parser.add_argument("--temperature", type=float, default=0.6)
+ args = parser.parse_args()
+ assert args.input_file or args.interactive, "Either input-file or interactive mode must be specified"
+ main(args.ckpt_path, args.config, args.input_file, args.interactive, args.max_new_tokens, args.temperature)
diff --git a/reference/official_inference/kernel.py b/reference/official_inference/kernel.py
new file mode 100644
index 00000000..ea7976fa
--- /dev/null
+++ b/reference/official_inference/kernel.py
@@ -0,0 +1,536 @@
+import torch
+import tilelang
+import tilelang.language as T
+from typing import Tuple, Optional
+
+
+tilelang.set_log_level("WARNING")
+
+pass_configs = {
+ tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
+ tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
+}
+
+FP8 = "float8_e4m3"
+FP4 = "float4_e2m1fn"
+FE8M0 = "float8_e8m0fnu"
+BF16 = "bfloat16"
+FP32 = "float32"
+INT32 = "int32"
+
+
+def fast_log2_ceil(x):
+ """Compute ceil(log2(x)) via IEEE 754 bit manipulation. Avoids slow log/ceil intrinsics."""
+ bits_x = T.reinterpret("uint32", x)
+ exp_x = (bits_x >> 23) & 0xFF
+ man_bits = bits_x & ((1 << 23) - 1)
+ return T.Cast("int32", exp_x - 127 + T.if_then_else(man_bits != 0, 1, 0))
+
+
+def fast_pow2(x):
+ """Compute 2^x for integer x via IEEE 754 bit manipulation."""
+ bits_x = (x + 127) << 23
+ return T.reinterpret("float32", bits_x)
+
+
+def fast_round_scale(amax, fp8_max_inv):
+ return fast_pow2(fast_log2_ceil(amax * fp8_max_inv))
+
+
+@tilelang.jit(pass_configs=pass_configs)
+def act_quant_kernel(
+ N, block_size=128, in_dtype=BF16, out_dtype=FP8, scale_dtype=FP32,
+ round_scale=False, inplace=False
+):
+ """Block-wise FP8 quantization. inplace=True does fused quant+dequant back to BF16."""
+ M = T.symbolic("M")
+ fp8_min = -448.0
+ fp8_max = 448.0
+ fp8_max_inv = 1 / fp8_max
+ num_stages = 0 if round_scale or inplace else 2
+ blk_m = 32
+ group_size = block_size
+ # Internal computation in FP32; scale_dtype controls output storage format.
+ compute_dtype = FP32
+ out_dtype = in_dtype if inplace else out_dtype
+
+ @T.prim_func
+ def act_quant_kernel_(
+ X: T.Tensor[(M, N), in_dtype],
+ Y: T.Tensor[(M, N), out_dtype],
+ S: T.Tensor[(M, T.ceildiv(N, group_size)), scale_dtype],
+ ):
+ with T.Kernel(T.ceildiv(M, blk_m), T.ceildiv(N, group_size), threads=128) as (
+ pid_m,
+ pid_n,
+ ):
+ x_shared = T.alloc_shared((blk_m, group_size), in_dtype)
+ x_local = T.alloc_fragment((blk_m, group_size), in_dtype)
+ amax_local = T.alloc_fragment((blk_m,), compute_dtype)
+ s_local = T.alloc_fragment((blk_m,), compute_dtype)
+ y_local = T.alloc_fragment((blk_m, group_size), out_dtype)
+ y_shared = T.alloc_shared((blk_m, group_size), out_dtype)
+
+ for _ in T.Pipelined(1, num_stages=num_stages):
+ T.copy(X[pid_m * blk_m, pid_n * group_size], x_shared)
+ T.copy(x_shared, x_local)
+ T.reduce_absmax(x_local, amax_local, dim=1)
+ for i in T.Parallel(blk_m):
+ amax_local[i] = T.max(amax_local[i], 1e-4)
+ if round_scale:
+ s_local[i] = fast_round_scale(amax_local[i], fp8_max_inv)
+ else:
+ s_local[i] = amax_local[i] * fp8_max_inv
+ if inplace:
+ for i, j in T.Parallel(blk_m, group_size):
+ y_local[i, j] = T.Cast(
+ out_dtype,
+ T.Cast(compute_dtype, T.Cast(out_dtype, T.clamp(
+ x_local[i, j] / s_local[i], fp8_min, fp8_max
+ ))) * s_local[i],
+ )
+ else:
+ for i, j in T.Parallel(blk_m, group_size):
+ y_local[i, j] = T.clamp(
+ x_local[i, j] / s_local[i], fp8_min, fp8_max
+ )
+ for i in T.Parallel(blk_m):
+ S[pid_m * blk_m + i, pid_n] = T.Cast(scale_dtype, s_local[i])
+ T.copy(y_local, y_shared)
+ T.copy(y_shared, Y[pid_m * blk_m, pid_n * group_size])
+
+ return act_quant_kernel_
+
+
+def act_quant(
+ x: torch.Tensor, block_size: int = 128, scale_fmt: Optional[str] = None,
+ scale_dtype: torch.dtype = torch.float32, inplace: bool = False,
+) -> torch.Tensor:
+ """Block-wise FP8 quantization. inplace=True does fused quant+dequant back to BF16.
+ When scale_fmt is set, scales are rounded to power-of-2 (MXFP)."""
+ N = x.size(-1)
+ assert N % block_size == 0
+ tl_dtype = FE8M0 if scale_dtype == torch.float8_e8m0fnu else FP32
+ z = x.contiguous()
+ y = torch.empty_like(z) if inplace else torch.empty_like(z, dtype=torch.float8_e4m3fn)
+ s = z.new_empty(*z.size()[:-1], N // block_size, dtype=scale_dtype)
+ kernel = act_quant_kernel(
+ N, block_size, scale_dtype=tl_dtype,
+ round_scale=scale_fmt is not None, inplace=inplace,
+ )
+ kernel(z.view(-1, N), y.view(-1, N), s.view(-1, N // block_size))
+ if inplace:
+ x.copy_(y)
+ return x
+ return y, s
+
+
+@tilelang.jit(pass_configs=pass_configs)
+def fp4_quant_kernel(
+ N, block_size=32, in_dtype=BF16, scale_dtype=FE8M0, inplace=False
+):
+ """Block-wise FP4 quantization. Power-of-2 scale via bit ops. inplace=True does fused quant+dequant."""
+ M = T.symbolic("M")
+ fp4_max = 6.0
+ fp4_max_inv = 1.0 / fp4_max
+ blk_m = 32
+ group_size = block_size
+ compute_dtype = FP32
+ out_dtype = in_dtype if inplace else FP4
+
+ @T.prim_func
+ def fp4_quant_kernel_(
+ X: T.Tensor[(M, N), in_dtype],
+ Y: T.Tensor[(M, N), out_dtype],
+ S: T.Tensor[(M, T.ceildiv(N, group_size)), scale_dtype],
+ ):
+ with T.Kernel(T.ceildiv(M, blk_m), T.ceildiv(N, group_size), threads=128) as (
+ pid_m,
+ pid_n,
+ ):
+ x_shared = T.alloc_shared((blk_m, group_size), in_dtype)
+ x_local = T.alloc_fragment((blk_m, group_size), in_dtype)
+ amax_local = T.alloc_fragment((blk_m,), compute_dtype)
+ s_local = T.alloc_fragment((blk_m,), compute_dtype)
+ y_local = T.alloc_fragment((blk_m, group_size), out_dtype)
+ y_shared = T.alloc_shared((blk_m, group_size), out_dtype)
+
+ for _ in T.Pipelined(1, num_stages=2):
+ T.copy(X[pid_m * blk_m, pid_n * group_size], x_shared)
+ T.copy(x_shared, x_local)
+ T.reduce_absmax(x_local, amax_local, dim=1)
+ for i in T.Parallel(blk_m):
+ amax_local[i] = T.max(amax_local[i], 6 * (2**-126))
+ s_local[i] = fast_round_scale(amax_local[i], fp4_max_inv)
+ if inplace:
+ for i, j in T.Parallel(blk_m, group_size):
+ y_local[i, j] = T.Cast(
+ out_dtype,
+ T.Cast(compute_dtype, T.Cast(FP4, T.clamp(
+ x_local[i, j] / s_local[i], -fp4_max, fp4_max
+ ))) * s_local[i],
+ )
+ else:
+ for i, j in T.Parallel(blk_m, group_size):
+ y_local[i, j] = T.clamp(
+ x_local[i, j] / s_local[i], -fp4_max, fp4_max
+ )
+ for i in T.Parallel(blk_m):
+ S[pid_m * blk_m + i, pid_n] = T.Cast(scale_dtype, s_local[i])
+ T.copy(y_local, y_shared)
+ T.copy(y_shared, Y[pid_m * blk_m, pid_n * group_size])
+
+ return fp4_quant_kernel_
+
+
+def fp4_act_quant(
+ x: torch.Tensor, block_size: int = 32, inplace: bool = False,
+) -> torch.Tensor:
+ """Block-wise FP4 quantization. inplace=True does fused quant+dequant back to BF16."""
+ N = x.size(-1)
+ assert N % block_size == 0
+ z = x.contiguous()
+ y = torch.empty_like(z) if inplace else z.new_empty(*z.shape[:-1], N // 2, dtype=torch.float4_e2m1fn_x2)
+ s = z.new_empty(*z.size()[:-1], N // block_size, dtype=torch.float8_e8m0fnu)
+ kernel = fp4_quant_kernel(N, block_size, inplace=inplace)
+ kernel(z.view(-1, N), y.view(-1, y.size(-1)), s.view(-1, N // block_size))
+ if inplace:
+ x.copy_(y)
+ return x
+ return y, s
+
+
+@tilelang.jit(pass_configs=pass_configs)
+def fp8_gemm_kernel(N, K, out_dtype=BF16, accum_dtype=FP32, scale_dtype=FP32):
+ assert out_dtype in [BF16, FP32]
+
+ M = T.symbolic("M")
+ group_size = 128
+ block_M = 32
+ block_N = 128
+ block_K = 128
+
+ @T.prim_func
+ def fp8_gemm_kernel_(
+ A: T.Tensor[(M, K), FP8],
+ B: T.Tensor[(N, K), FP8],
+ C: T.Tensor[(M, N), out_dtype],
+ scales_a: T.Tensor[(M, T.ceildiv(K, group_size)), scale_dtype],
+ scales_b: T.Tensor[(T.ceildiv(N, group_size), T.ceildiv(K, group_size)), scale_dtype],
+ ):
+ with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) as (
+ bx,
+ by,
+ ):
+ A_shared = T.alloc_shared((block_M, block_K), FP8)
+ B_shared = T.alloc_shared((block_N, block_K), FP8)
+ C_shared = T.alloc_shared((block_M, block_N), out_dtype)
+ Scale_C_shared = T.alloc_shared((block_M), FP32)
+ C_local = T.alloc_fragment((block_M, block_N), accum_dtype)
+ C_local_accum = T.alloc_fragment((block_M, block_N), accum_dtype)
+
+ # Improve L2 Cache
+ T.use_swizzle(panel_size=10)
+ T.clear(C_local)
+ T.clear(C_local_accum)
+
+ K_iters = T.ceildiv(K, block_K)
+ for k in T.Pipelined(K_iters, num_stages=4):
+ T.copy(A[by * block_M, k * block_K], A_shared)
+ T.copy(B[bx * block_N, k * block_K], B_shared)
+ # Cast scales to FP32 for computation; scales_b has one value per block_N group
+ Scale_B = T.Cast(FP32, scales_b[bx * block_N // group_size, k])
+ for i in T.Parallel(block_M):
+ Scale_C_shared[i] = T.Cast(FP32, scales_a[by * block_M + i, k]) * Scale_B
+
+ T.gemm(A_shared, B_shared, C_local, transpose_B=True)
+ # Separate accumulator for scale-corrected results (2x accumulation precision)
+ for i, j in T.Parallel(block_M, block_N):
+ C_local_accum[i, j] += C_local[i, j] * Scale_C_shared[i]
+ T.clear(C_local)
+ T.copy(C_local_accum, C_shared)
+ T.copy(C_shared, C[by * block_M, bx * block_N])
+
+ return fp8_gemm_kernel_
+
+
+def fp8_gemm(
+ a: torch.Tensor, a_s: torch.Tensor, b: torch.Tensor, b_s: torch.Tensor,
+ scale_dtype: torch.dtype = torch.float32,
+) -> torch.Tensor:
+ """C[M,N] = A[M,K] @ B[N,K]^T with per-128 block FP8 scaling on both A and B."""
+ assert a.is_contiguous() and b.is_contiguous(), "Input tensors must be contiguous"
+ assert a_s.is_contiguous() and b_s.is_contiguous(), (
+ "Scaling factor tensors must be contiguous"
+ )
+ tl_dtype = FE8M0 if scale_dtype == torch.float8_e8m0fnu else FP32
+ K = a.size(-1)
+ M = a.numel() // K
+ N = b.size(0)
+ c = a.new_empty(*a.size()[:-1], N, dtype=torch.get_default_dtype())
+ kernel = fp8_gemm_kernel(N, K, scale_dtype=tl_dtype)
+ kernel(a.view(M, K), b, c.view(M, N), a_s.view(M, -1), b_s)
+ return c
+
+
+@tilelang.jit(pass_configs=pass_configs)
+def sparse_attn_kernel(h: int, d: int, scale=None):
+ """Sparse multi-head attention via index gathering + online softmax (FlashAttention-style).
+ For each (batch, seq_pos), gathers top-k KV positions by index, computes attention
+ with numerically stable running max/sum, and includes a learnable attn_sink bias."""
+ b = T.symbolic("b")
+ m = T.symbolic("m")
+ n = T.symbolic("n")
+ topk = T.symbolic("topk")
+ if scale is None:
+ scale = (1.0 / d) ** 0.5
+
+ num_stages = 2
+ threads = 256
+ block = 64
+ num_blocks = tilelang.cdiv(topk, block)
+
+ @T.prim_func
+ def sparse_attn_kernel_(
+ q: T.Tensor[(b, m, h, d), BF16],
+ kv: T.Tensor[(b, n, d), BF16],
+ o: T.Tensor[(b, m, h, d), BF16],
+ attn_sink: T.Tensor[(h,), FP32],
+ topk_idxs: T.Tensor[(b, m, topk), INT32],
+ ):
+ with T.Kernel(m, b, threads=threads) as (bx, by):
+ q_shared = T.alloc_shared((h, d), BF16)
+ kv_shared = T.alloc_shared((block, d), BF16)
+ o_shared = T.alloc_shared((h, d), BF16)
+ acc_s_cast = T.alloc_shared((h, block), BF16)
+
+ idxs = T.alloc_fragment(block, INT32)
+ acc_s = T.alloc_fragment((h, block), FP32)
+ acc_o = T.alloc_fragment((h, d), FP32)
+ scores_max = T.alloc_fragment(h, FP32)
+ scores_max_prev = T.alloc_fragment(h, FP32)
+ scores_scale = T.alloc_fragment(h, FP32)
+ scores_sum = T.alloc_fragment(h, FP32)
+ sum_exp = T.alloc_fragment(h, FP32)
+
+ T.clear(acc_o)
+ T.clear(sum_exp)
+ T.fill(scores_max, -T.infinity(FP32))
+ T.copy(q[by, bx, :, :], q_shared)
+
+ for t in T.Pipelined(num_blocks, num_stages=num_stages):
+ for i in T.Parallel(block):
+ idxs[i] = T.if_then_else(t * block + i < topk, topk_idxs[by, bx, t * block + i], -1)
+ for i, j in T.Parallel(block, d):
+ kv_shared[i, j] = T.if_then_else(idxs[i] != -1, kv[by, idxs[i], j], 0)
+ for i, j in T.Parallel(h, block):
+ acc_s[i, j] = T.if_then_else(idxs[j] != -1, 0, -T.infinity(FP32))
+ T.gemm(q_shared, kv_shared, acc_s, transpose_B=True, policy=T.GemmWarpPolicy.FullRow)
+ for i, j in T.Parallel(h, block):
+ acc_s[i, j] *= scale
+ T.copy(scores_max, scores_max_prev)
+ T.reduce_max(acc_s, scores_max, dim=1, clear=False)
+ for i in T.Parallel(h):
+ scores_scale[i] = T.exp(scores_max_prev[i] - scores_max[i])
+ for i, j in T.Parallel(h, block):
+ acc_s[i, j] = T.exp(acc_s[i, j] - scores_max[i])
+ T.reduce_sum(acc_s, scores_sum, dim=1)
+ for i in T.Parallel(h):
+ sum_exp[i] = sum_exp[i] * scores_scale[i] + scores_sum[i]
+ T.copy(acc_s, acc_s_cast)
+ for i, j in T.Parallel(h, d):
+ acc_o[i, j] *= scores_scale[i]
+ T.gemm(acc_s_cast, kv_shared, acc_o, policy=T.GemmWarpPolicy.FullRow)
+
+ for i in T.Parallel(h):
+ sum_exp[i] += T.exp(attn_sink[i] - scores_max[i])
+ for i, j in T.Parallel(h, d):
+ acc_o[i, j] /= sum_exp[i]
+ T.copy(acc_o, o_shared)
+ T.copy(o_shared, o[by, bx, :, :])
+
+ return sparse_attn_kernel_
+
+
+def sparse_attn(
+ q: torch.Tensor, kv: torch.Tensor, attn_sink: torch.Tensor, topk_idxs: torch.Tensor, softmax_scale: float
+) -> torch.Tensor:
+ b, s, h, d = q.size()
+ # Pad heads to 16 for kernel efficiency (stripped after)
+ if h < 16:
+ q = torch.cat([q, q.new_zeros(b, s, 16 - h, d)], dim=2)
+ attn_sink = torch.cat([attn_sink, attn_sink.new_zeros(16 - h)])
+ o = torch.empty_like(q)
+ kernel = sparse_attn_kernel(q.size(2), d, softmax_scale)
+ kernel(q, kv, o, attn_sink, topk_idxs)
+ if h < 16:
+ o = o.narrow(2, 0, h).contiguous()
+ return o
+
+
+@tilelang.jit(pass_configs=pass_configs)
+def hc_split_sinkhorn_kernel(hc: int, sinkhorn_iters: int, eps: float):
+ n = T.symbolic("n")
+ mix_hc = (2 + hc) * hc
+ threads = 64
+
+ @T.prim_func
+ def hc_split_sinkhorn_kernel_(
+ mixes: T.Tensor[(n, mix_hc), FP32],
+ hc_scale: T.Tensor[(3,), FP32],
+ hc_base: T.Tensor[(mix_hc,), FP32],
+ pre: T.Tensor[(n, hc), FP32],
+ post: T.Tensor[(n, hc), FP32],
+ comb: T.Tensor[(n, hc, hc), FP32],
+ ):
+ with T.Kernel(n, threads=threads) as i:
+ mixes_shared = T.alloc_shared(mix_hc, FP32)
+ comb_frag = T.alloc_fragment((hc, hc), FP32)
+ T.copy(mixes[i, :], mixes_shared)
+
+ for j in T.Parallel(hc):
+ pre[i, j] = T.sigmoid(mixes_shared[j] * hc_scale[0] + hc_base[j]) + eps
+ for j in T.Parallel(hc):
+ post[i, j] = 2 * T.sigmoid(mixes_shared[j + hc] * hc_scale[1] + hc_base[j + hc])
+ for j, k in T.Parallel(hc, hc):
+ comb_frag[j, k] = mixes_shared[j * hc + k + hc * 2] * hc_scale[2] + hc_base[j * hc + k + hc * 2]
+
+ row_sum = T.alloc_fragment(hc, FP32)
+ col_sum = T.alloc_fragment(hc, FP32)
+
+ # comb = comb.softmax(-1) + eps
+ row_max = T.alloc_fragment(hc, FP32)
+ T.reduce_max(comb_frag, row_max, dim=1)
+ for j, k in T.Parallel(hc, hc):
+ comb_frag[j, k] = T.exp(comb_frag[j, k] - row_max[j])
+ T.reduce_sum(comb_frag, row_sum, dim=1)
+ for j, k in T.Parallel(hc, hc):
+ comb_frag[j, k] = comb_frag[j, k] / row_sum[j] + eps
+
+ # comb = comb / (comb.sum(-2) + eps)
+ T.reduce_sum(comb_frag, col_sum, dim=0)
+ for j, k in T.Parallel(hc, hc):
+ comb_frag[j, k] = comb_frag[j, k] / (col_sum[k] + eps)
+
+ for _ in T.serial(sinkhorn_iters - 1):
+ # comb = comb / (comb.sum(-1) + eps)
+ T.reduce_sum(comb_frag, row_sum, dim=1)
+ for j, k in T.Parallel(hc, hc):
+ comb_frag[j, k] = comb_frag[j, k] / (row_sum[j] + eps)
+ # comb = comb / (comb.sum(-2) + eps)
+ T.reduce_sum(comb_frag, col_sum, dim=0)
+ for j, k in T.Parallel(hc, hc):
+ comb_frag[j, k] = comb_frag[j, k] / (col_sum[k] + eps)
+
+ T.copy(comb_frag, comb[i, :, :])
+
+ return hc_split_sinkhorn_kernel_
+
+
+def hc_split_sinkhorn(mixes: torch.Tensor, hc_scale: torch.Tensor, hc_base: torch.Tensor, hc_mult: int = 4, sinkhorn_iters: int = 20, eps: float = 1e-6):
+ b, s, _ = mixes.size()
+ pre = mixes.new_empty(b, s, hc_mult)
+ post = mixes.new_empty(b, s, hc_mult)
+ comb = mixes.new_empty(b, s, hc_mult, hc_mult)
+ kernel = hc_split_sinkhorn_kernel(hc_mult, sinkhorn_iters, eps)
+ kernel(mixes.view(-1, (2 + hc_mult) * hc_mult), hc_scale, hc_base,
+ pre.view(-1, hc_mult), post.view(-1, hc_mult), comb.view(-1, hc_mult, hc_mult))
+ return pre, post, comb
+
+
+@tilelang.jit(pass_configs=pass_configs)
+def fp4_gemm_kernel(N, K, out_dtype=BF16, accum_dtype=FP32, scale_dtype=FP32):
+ """FP8 act x FP4 weight GEMM kernel.
+
+ C[M, N] = A_fp8[M, K] @ B_fp4[N, K]^T
+
+ Act: 1x128 quant on K (reduce dim), FP8 with configurable scale dtype
+ Weight: 1x32 quant on K (reduce dim), FP4 with E8M0 scale
+
+ B is stored as [N, K//2] in float4_e2m1fn_x2, logical [N, K] in fp4.
+ The FP4 values are packed along the K (last) dimension.
+
+ Strategy: load FP4 sub-blocks of size [block_N, sub_K] (sub_K=32),
+ cast FP4 to FP8 via float, then do FP8xFP8 GEMM.
+ Apply act scale (per 128 on K) and weight scale (per 32 on K) to the accumulator.
+ """
+ M = T.symbolic("M")
+ act_group_size = 128
+ weight_group_size = 32
+ block_M = 32
+ block_N = 128
+ block_K = 32 # matches weight_group_size for simple scale handling
+ n_sub = act_group_size // block_K # 4 sub-blocks per act scale group
+
+ @T.prim_func
+ def fp4_gemm_kernel_(
+ A: T.Tensor[(M, K), FP8],
+ B: T.Tensor[(N, K), FP4],
+ C: T.Tensor[(M, N), out_dtype],
+ scales_a: T.Tensor[(M, T.ceildiv(K, act_group_size)), scale_dtype],
+ scales_b: T.Tensor[(N, T.ceildiv(K, weight_group_size)), scale_dtype],
+ ):
+ with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) as (
+ bx,
+ by,
+ ):
+ A_shared = T.alloc_shared((block_M, block_K), FP8)
+ B_fp4_shared = T.alloc_shared((block_N, block_K), FP4)
+ B_shared = T.alloc_shared((block_N, block_K), FP8)
+ C_shared = T.alloc_shared((block_M, block_N), out_dtype)
+ C_local = T.alloc_fragment((block_M, block_N), accum_dtype)
+ C_local_accum = T.alloc_fragment((block_M, block_N), accum_dtype)
+ scale_a_frag = T.alloc_fragment((block_M,), FP32)
+ scale_b_frag = T.alloc_fragment((block_N,), FP32)
+
+ T.use_swizzle(panel_size=10)
+ T.clear(C_local)
+ T.clear(C_local_accum)
+
+ K_iters = T.ceildiv(K, block_K)
+ for k in T.Pipelined(K_iters, num_stages=2):
+ T.copy(A[by * block_M, k * block_K], A_shared)
+ T.copy(B[bx * block_N, k * block_K], B_fp4_shared)
+ # FP4->FP8 cast must go through FP32 to avoid ambiguous C++ overload
+ for i, j in T.Parallel(block_N, block_K):
+ B_shared[i, j] = T.Cast(FP8, T.Cast(FP32, B_fp4_shared[i, j]))
+
+ # Weight scale: per 32 on K, indexed by k (each k is one block_K=32)
+ for i in T.Parallel(block_N):
+ scale_b_frag[i] = T.Cast(FP32, scales_b[bx * block_N + i, k])
+
+ # Act scale: per 128 on K, indexed by k // 4
+ for i in T.Parallel(block_M):
+ scale_a_frag[i] = T.Cast(FP32, scales_a[by * block_M + i, k // n_sub])
+
+ T.gemm(A_shared, B_shared, C_local, transpose_B=True)
+
+ for i, j in T.Parallel(block_M, block_N):
+ C_local_accum[i, j] += C_local[i, j] * scale_a_frag[i] * scale_b_frag[j]
+ T.clear(C_local)
+
+ T.copy(C_local_accum, C_shared)
+ T.copy(C_shared, C[by * block_M, bx * block_N])
+
+ return fp4_gemm_kernel_
+
+
+def fp4_gemm(
+ a: torch.Tensor, a_s: torch.Tensor, b: torch.Tensor, b_s: torch.Tensor,
+ scale_dtype: torch.dtype = torch.float32,
+) -> torch.Tensor:
+ """C[M,N] = A_fp8[M,K] @ B_fp4[N,K]^T.
+ A has per-128 act scale; B has per-32 E8M0 weight scale.
+ B is stored as [N, K//2] in float4_e2m1fn_x2 (2 FP4 values per byte, packed along K)."""
+ assert a.is_contiguous() and b.is_contiguous(), "Input tensors must be contiguous"
+ assert a_s.is_contiguous() and b_s.is_contiguous(), (
+ "Scaling factor tensors must be contiguous"
+ )
+ tl_dtype = FE8M0 if scale_dtype == torch.float8_e8m0fnu else FP32
+ K = a.size(-1)
+ M = a.numel() // K
+ N = b.size(0)
+ c = a.new_empty(*a.size()[:-1], N, dtype=torch.get_default_dtype())
+ kernel = fp4_gemm_kernel(N, K, scale_dtype=tl_dtype)
+ kernel(a.view(M, K), b, c.view(M, N), a_s.view(M, -1), b_s)
+ return c
diff --git a/reference/official_inference/model.py b/reference/official_inference/model.py
new file mode 100644
index 00000000..167ade8f
--- /dev/null
+++ b/reference/official_inference/model.py
@@ -0,0 +1,827 @@
+import math
+from dataclasses import dataclass
+from typing import Tuple, Optional, Literal
+from functools import lru_cache
+from contextlib import contextmanager
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+import torch.distributed as dist
+
+from kernel import act_quant, fp4_act_quant, fp8_gemm, fp4_gemm, sparse_attn, hc_split_sinkhorn
+
+
+world_size = 1
+rank = 0
+block_size = 128
+fp4_block_size = 32
+default_dtype = torch.bfloat16
+scale_fmt = None
+scale_dtype = torch.float32
+
+
+@contextmanager
+def set_dtype(dtype):
+ """Temporarily override torch default dtype, restoring it on exit (even if an exception occurs)."""
+ prev = torch.get_default_dtype()
+ torch.set_default_dtype(dtype)
+ try:
+ yield
+ finally:
+ torch.set_default_dtype(prev)
+
+@dataclass
+class ModelArgs:
+ """Model hyperparameters. Field names match the config JSON keys."""
+ max_batch_size: int = 4
+ max_seq_len: int = 4096
+ dtype: Literal["bf16", "fp8"] = "fp8"
+ scale_fmt: Literal[None, "ue8m0"] = "ue8m0"
+ expert_dtype: Literal[None, "fp4"] = None
+ scale_dtype: Literal["fp32", "fp8"] = "fp8"
+ vocab_size: int = 129280
+ dim: int = 4096
+ moe_inter_dim: int = 4096
+ n_layers: int = 7
+ n_hash_layers: int = 0
+ n_mtp_layers: int = 1
+ n_heads: int = 64
+ # moe
+ n_routed_experts: int = 8
+ n_shared_experts: int = 1
+ n_activated_experts: int = 2
+ score_func: Literal["softmax", "sigmoid", "sqrtsoftplus"] = "sqrtsoftplus"
+ route_scale: float = 1.
+ swiglu_limit: float = 0.
+ # mqa
+ q_lora_rank: int = 1024
+ head_dim: int = 512
+ rope_head_dim: int = 64
+ norm_eps: float = 1e-6
+ o_groups: int = 8
+ o_lora_rank: int = 1024
+ window_size: int = 128
+ compress_ratios: Tuple[int] = (0, 0, 4, 128, 4, 128, 4, 0)
+ # yarn
+ compress_rope_theta: float = 40000.0
+ original_seq_len: int = 0
+ rope_theta: float = 10000.0
+ rope_factor: float = 40
+ beta_fast: int = 32
+ beta_slow: int = 1
+ # index
+ index_n_heads: int = 64
+ index_head_dim: int = 128
+ index_topk: int = 512
+ # hc
+ hc_mult: int = 4
+ hc_sinkhorn_iters: int = 20
+ hc_eps: float = 1e-6
+
+
+class ParallelEmbedding(nn.Module):
+ """Embedding sharded along the vocab dimension. Each rank holds vocab_size // world_size rows.
+ Out-of-range indices are zero-masked before all_reduce to combine partial embeddings."""
+ def __init__(self, vocab_size: int, dim: int):
+ super().__init__()
+ self.vocab_size = vocab_size
+ self.dim = dim
+ assert vocab_size % world_size == 0, f"Vocabulary size must be divisible by world size (world_size={world_size})"
+ self.part_vocab_size = (vocab_size // world_size)
+ self.vocab_start_idx = rank * self.part_vocab_size
+ self.vocab_end_idx = self.vocab_start_idx + self.part_vocab_size
+ self.weight = nn.Parameter(torch.empty(self.part_vocab_size, self.dim))
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ if world_size > 1:
+ mask = (x < self.vocab_start_idx) | (x >= self.vocab_end_idx)
+ x = x - self.vocab_start_idx
+ x[mask] = 0
+ y = F.embedding(x, self.weight)
+ if world_size > 1:
+ y[mask] = 0
+ dist.all_reduce(y)
+ return y
+
+
+def linear(x: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor] = None) -> torch.Tensor:
+ """Dispatches to fp4_gemm / fp8_gemm / F.linear based on weight dtype.
+ For quantized weights, x is first quantized to FP8 via act_quant."""
+ assert bias is None
+
+ if weight.dtype == torch.float4_e2m1fn_x2:
+ x, s = act_quant(x, block_size, scale_fmt, scale_dtype)
+ return fp4_gemm(x, s, weight, weight.scale, scale_dtype)
+ elif weight.dtype == torch.float8_e4m3fn:
+ x, s = act_quant(x, block_size, scale_fmt, scale_dtype)
+ return fp8_gemm(x, s, weight, weight.scale, scale_dtype)
+ else:
+ return F.linear(x, weight)
+
+
+class Linear(nn.Module):
+ """Linear layer supporting BF16, FP8, and FP4 weight formats with per-block scaling."""
+
+ def __init__(self, in_features: int, out_features: int, bias: bool = False, dtype = None):
+ super().__init__()
+ self.in_features = in_features
+ self.out_features = out_features
+ dtype = dtype or default_dtype
+ if dtype == torch.float4_e2m1fn_x2:
+ # FP4: weight is [out, in//2] in float4_e2m1fn_x2, logically [out, in] in fp4
+ # Scale is [out, in//32] in float8_e8m0fnu (1 scale per 32 fp4 elements along K)
+ self.weight = nn.Parameter(torch.empty(out_features, in_features // 2, dtype=torch.float4_e2m1fn_x2))
+ scale_out_features = out_features
+ scale_in_features = in_features // fp4_block_size
+ self.weight.scale = self.scale = nn.Parameter(torch.empty(scale_out_features, scale_in_features, dtype=torch.float8_e8m0fnu))
+ elif dtype == torch.float8_e4m3fn:
+ self.weight = nn.Parameter(torch.empty(out_features, in_features, dtype=dtype))
+ scale_out_features = (out_features + block_size - 1) // block_size
+ scale_in_features = (in_features + block_size - 1) // block_size
+ self.weight.scale = self.scale = nn.Parameter(torch.empty(scale_out_features, scale_in_features, dtype=torch.float8_e8m0fnu))
+ else:
+ self.weight = nn.Parameter(torch.empty(out_features, in_features, dtype=dtype))
+ self.register_parameter("scale", None)
+ if bias:
+ self.bias = nn.Parameter(torch.empty(out_features))
+ else:
+ self.register_parameter("bias", None)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ return linear(x, self.weight, self.bias)
+
+
+class ColumnParallelLinear(Linear):
+ """Shards output dim across TP ranks. No all-reduce needed on output."""
+ def __init__(self, in_features: int, out_features: int, bias: bool = False, dtype = None):
+ assert out_features % world_size == 0, f"Output features must be divisible by world size (world_size={world_size})"
+ self.part_out_features = out_features // world_size
+ super().__init__(in_features, self.part_out_features, bias, dtype)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ return linear(x, self.weight, self.bias)
+
+
+class RowParallelLinear(Linear):
+ """Shards input dim across TP ranks. All-reduce on output to sum partial results."""
+ def __init__(self, in_features: int, out_features: int, bias: bool = False, dtype = None):
+ assert in_features % world_size == 0, f"Input features must be divisible by world size (world_size={world_size})"
+ self.part_in_features = in_features // world_size
+ super().__init__(self.part_in_features, out_features, bias, dtype)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ y = linear(x, self.weight, None)
+ if world_size > 1:
+ y = y.float()
+ dist.all_reduce(y)
+ if self.bias is not None:
+ y += self.bias
+ return y.type_as(x)
+
+
+class RMSNorm(nn.Module):
+ def __init__(self, dim: int, eps: float = 1e-6):
+ super().__init__()
+ self.dim = dim
+ self.eps = eps
+ # rmsnorm in the checkpoint is stored in bf16, while the parameter here is stored in fp32 for convenient.
+ self.weight = nn.Parameter(torch.ones(dim, dtype=torch.float32))
+
+ def forward(self, x: torch.Tensor):
+ dtype = x.dtype
+ x = x.float()
+ var = x.square().mean(-1, keepdim=True)
+ x = x * torch.rsqrt(var + self.eps)
+ return (self.weight * x).to(dtype)
+
+
+@lru_cache(2)
+def precompute_freqs_cis(dim, seqlen, original_seq_len, base, factor, beta_fast, beta_slow) -> torch.Tensor:
+ """Precomputes complex exponentials for rotary embeddings with YaRN scaling.
+ When original_seq_len > 0, applies frequency interpolation with a smooth
+ linear ramp between beta_fast and beta_slow correction ranges."""
+
+ def find_correction_dim(num_rotations, dim, base, max_seq_len):
+ return dim * math.log(max_seq_len / (num_rotations * 2 * math.pi)) / (2 * math.log(base))
+
+ def find_correction_range(low_rot, high_rot, dim, base, max_seq_len):
+ low = math.floor(find_correction_dim(low_rot, dim, base, max_seq_len))
+ high = math.ceil(find_correction_dim(high_rot, dim, base, max_seq_len))
+ return max(low, 0), min(high, dim-1)
+
+ def linear_ramp_factor(min, max, dim):
+ if min == max:
+ max += 0.001
+ linear_func = (torch.arange(dim, dtype=torch.float32) - min) / (max - min)
+ ramp_func = torch.clamp(linear_func, 0, 1)
+ return ramp_func
+
+ freqs = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim))
+ if original_seq_len > 0:
+ low, high = find_correction_range(beta_fast, beta_slow, dim, base, original_seq_len)
+ smooth = 1 - linear_ramp_factor(low, high, dim // 2)
+ freqs = freqs / factor * (1 - smooth) + freqs * smooth
+
+ t = torch.arange(seqlen)
+ freqs = torch.outer(t, freqs)
+ freqs_cis = torch.polar(torch.ones_like(freqs), freqs)
+ return freqs_cis
+
+
+def apply_rotary_emb(x: torch.Tensor, freqs_cis: torch.Tensor, inverse: bool = False) -> torch.Tensor:
+ """Applies rotary positional embeddings in-place. Uses conjugate for inverse (de-rotation)."""
+ y = x
+ x = torch.view_as_complex(x.float().unflatten(-1, (-1, 2)))
+ if inverse:
+ freqs_cis = freqs_cis.conj()
+ if x.ndim == 3:
+ freqs_cis = freqs_cis.view(1, x.size(1), x.size(-1))
+ else:
+ freqs_cis = freqs_cis.view(1, x.size(1), 1, x.size(-1))
+ x = torch.view_as_real(x * freqs_cis).flatten(-2)
+ y.copy_(x)
+ return y
+
+
+def rotate_activation(x: torch.Tensor) -> torch.Tensor:
+ """Applies randomized Hadamard rotation to spread information across dims before FP8 quant."""
+ assert x.dtype == torch.bfloat16
+ from fast_hadamard_transform import hadamard_transform
+ return hadamard_transform(x, scale=x.size(-1) ** -0.5)
+
+
+@lru_cache(1)
+def get_window_topk_idxs(window_size: int, bsz: int, seqlen: int, start_pos: int):
+ if start_pos >= window_size - 1:
+ start_pos %= window_size
+ matrix = torch.cat([torch.arange(start_pos + 1, window_size), torch.arange(0, start_pos + 1)], dim=0)
+ elif start_pos > 0:
+ matrix = F.pad(torch.arange(start_pos + 1), (0, window_size - start_pos - 1), value=-1)
+ else:
+ base = torch.arange(seqlen).unsqueeze(1)
+ matrix = (base - window_size + 1).clamp(0) + torch.arange(min(seqlen, window_size))
+ matrix = torch.where(matrix > base, -1, matrix)
+ return matrix.unsqueeze(0).expand(bsz, -1, -1)
+
+
+@lru_cache(2)
+def get_compress_topk_idxs(ratio: int, bsz: int, seqlen: int, start_pos: int, offset: int):
+ if start_pos > 0:
+ matrix = torch.arange(0, (start_pos + 1) // ratio) + offset
+ else:
+ matrix = torch.arange(seqlen // ratio).repeat(seqlen, 1)
+ mask = matrix >= torch.arange(1, seqlen + 1).unsqueeze(1) // ratio
+ matrix = torch.where(mask, -1, matrix + offset)
+ return matrix.unsqueeze(0).expand(bsz, -1, -1)
+
+
+class Compressor(nn.Module):
+ """Compresses KV cache via learned gated pooling over `compress_ratio` consecutive tokens.
+ When overlap=True (ratio==4), uses overlapping windows for smoother compression boundaries."""
+
+ def __init__(self, args: ModelArgs, compress_ratio: int = 4, head_dim: int = 512, rotate: bool = False):
+ super().__init__()
+ self.dim = args.dim
+ self.head_dim = head_dim
+ self.rope_head_dim = args.rope_head_dim
+ self.nope_head_dim = head_dim - args.rope_head_dim
+ self.compress_ratio = compress_ratio
+ self.overlap = compress_ratio == 4
+ self.rotate = rotate
+ coff = 1 + self.overlap
+
+ self.ape = nn.Parameter(torch.empty(compress_ratio, coff * self.head_dim, dtype=torch.float32))
+ # wkv and wgate in the checkpoint is stored in bf16, while the parameter here is stored in fp32 for convenient.
+ # When overlap, the first half of dims is for overlapping compression, second half for normal.
+ self.wkv = Linear(self.dim, coff * self.head_dim, dtype=torch.float32)
+ self.wgate = Linear(self.dim, coff * self.head_dim, dtype=torch.float32)
+ self.norm = RMSNorm(self.head_dim, args.norm_eps)
+ self.kv_cache: torch.Tensor = None # assigned lazily from Attention.kv_cache
+ # State buffers for decode-phase incremental compression.
+ # With overlap: state[:, :ratio] = overlapping window, state[:, ratio:] = current window.
+ self.register_buffer("kv_state", torch.zeros(args.max_batch_size, coff * compress_ratio, coff * self.head_dim, dtype=torch.float32), persistent=False)
+ self.register_buffer("score_state", torch.full((args.max_batch_size, coff * compress_ratio, coff * self.head_dim), float("-inf"), dtype=torch.float32), persistent=False)
+ self.freqs_cis: torch.Tensor = None
+
+ def overlap_transform(self, tensor: torch.Tensor, value=0):
+ # tensor: [b,s,r,2d]
+ b, s, _, _ = tensor.size()
+ ratio, d = self.compress_ratio, self.head_dim
+ new_tensor = tensor.new_full((b, s, 2 * ratio, d), value)
+ new_tensor[:, :, ratio:] = tensor[:, :, :, d:]
+ new_tensor[:, 1:, :ratio] = tensor[:, :-1, :, :d]
+ return new_tensor
+
+ def forward(self, x: torch.Tensor, start_pos: int):
+ assert self.kv_cache is not None
+ bsz, seqlen, _ = x.size()
+ ratio, overlap, d, rd = self.compress_ratio, self.overlap, self.head_dim, self.rope_head_dim
+ dtype = x.dtype
+ # compression need fp32
+ x = x.float()
+ kv = self.wkv(x)
+ score = self.wgate(x)
+ if start_pos == 0:
+ should_compress = seqlen >= ratio
+ remainder = seqlen % ratio
+ cutoff = seqlen - remainder
+ offset = ratio if overlap else 0
+ if overlap and cutoff >= ratio:
+ self.kv_state[:bsz, :ratio] = kv[:, cutoff-ratio : cutoff]
+ self.score_state[:bsz, :ratio] = score[:, cutoff-ratio : cutoff] + self.ape
+ if remainder > 0:
+ kv, self.kv_state[:bsz, offset : offset+remainder] = kv.split([cutoff, remainder], dim=1)
+ self.score_state[:bsz, offset : offset+remainder] = score[:, cutoff:] + self.ape[:remainder]
+ score = score[:, :cutoff]
+ kv = kv.unflatten(1, (-1, ratio))
+ score = score.unflatten(1, (-1, ratio)) + self.ape
+ if overlap:
+ kv = self.overlap_transform(kv, 0)
+ score = self.overlap_transform(score, float("-inf"))
+ kv = (kv * score.softmax(dim=2)).sum(dim=2)
+ else:
+ should_compress = (start_pos + 1) % self.compress_ratio == 0
+ score += self.ape[start_pos % ratio]
+ if overlap:
+ self.kv_state[:bsz, ratio + start_pos % ratio] = kv.squeeze(1)
+ self.score_state[:bsz, ratio + start_pos % ratio] = score.squeeze(1)
+ if should_compress:
+ kv_state = torch.cat([self.kv_state[:bsz, :ratio, :d], self.kv_state[:bsz, ratio:, d:]], dim=1)
+ score_state = torch.cat([self.score_state[:bsz, :ratio, :d], self.score_state[:bsz, ratio:, d:]], dim=1)
+ kv = (kv_state * score_state.softmax(dim=1)).sum(dim=1, keepdim=True)
+ self.kv_state[:bsz, :ratio] = self.kv_state[:bsz, ratio:]
+ self.score_state[:bsz, :ratio] = self.score_state[:bsz, ratio:]
+ else:
+ self.kv_state[:bsz, start_pos % ratio] = kv.squeeze(1)
+ self.score_state[:bsz, start_pos % ratio] = score.squeeze(1)
+ if should_compress:
+ kv = (self.kv_state[:bsz] * self.score_state[:bsz].softmax(dim=1)).sum(dim=1, keepdim=True)
+ if not should_compress:
+ return
+ kv = self.norm(kv.to(dtype))
+ if start_pos == 0:
+ freqs_cis = self.freqs_cis[:cutoff:ratio]
+ else:
+ freqs_cis = self.freqs_cis[start_pos + 1 - self.compress_ratio].unsqueeze(0)
+ apply_rotary_emb(kv[..., -rd:], freqs_cis)
+ if self.rotate:
+ kv = rotate_activation(kv)
+ fp4_act_quant(kv, fp4_block_size, True)
+ else:
+ act_quant(kv[..., :-rd], 64, scale_fmt, scale_dtype, True)
+ if start_pos == 0:
+ self.kv_cache[:bsz, :seqlen // ratio] = kv
+ else:
+ self.kv_cache[:bsz, start_pos // ratio] = kv.squeeze(1)
+ return kv
+
+
+class Indexer(torch.nn.Module):
+ """Selects top-k compressed KV positions for sparse attention via learned scoring.
+ Has its own Compressor (with Hadamard rotation) to build compressed KV for scoring."""
+
+ def __init__(self, args: ModelArgs, compress_ratio: int = 4):
+ super().__init__()
+ self.dim = args.dim
+ self.n_heads = args.index_n_heads
+ self.n_local_heads = args.index_n_heads // world_size
+ self.head_dim = args.index_head_dim
+ self.rope_head_dim = args.rope_head_dim
+ self.index_topk = args.index_topk
+ self.q_lora_rank = args.q_lora_rank
+ self.wq_b = ColumnParallelLinear(self.q_lora_rank, self.n_heads * self.head_dim)
+ self.weights_proj = ColumnParallelLinear(self.dim, self.n_heads, dtype=torch.bfloat16)
+ self.softmax_scale = self.head_dim ** -0.5
+ self.compress_ratio = compress_ratio
+
+ self.compressor = Compressor(args, compress_ratio, self.head_dim, True)
+ self.register_buffer("kv_cache", torch.zeros(args.max_batch_size, args.max_seq_len // compress_ratio, self.head_dim), persistent=False)
+ self.freqs_cis = None
+
+ def forward(self, x: torch.Tensor, qr: torch.Tensor, start_pos: int, offset: int):
+ bsz, seqlen, _ = x.size()
+ freqs_cis = self.freqs_cis[start_pos:start_pos+seqlen]
+ ratio = self.compress_ratio
+ rd = self.rope_head_dim
+ end_pos = start_pos + seqlen
+ if self.compressor.kv_cache is None:
+ self.compressor.kv_cache = self.kv_cache
+ self.compressor.freqs_cis = self.freqs_cis
+ q = self.wq_b(qr)
+ q = q.unflatten(-1, (self.n_local_heads, self.head_dim))
+ apply_rotary_emb(q[..., -rd:], freqs_cis)
+ q = rotate_activation(q)
+ # use fp4 simulation for q and kv in indexer
+ fp4_act_quant(q, fp4_block_size, True)
+ self.compressor(x, start_pos)
+ weights = self.weights_proj(x) * (self.softmax_scale * self.n_heads ** -0.5)
+ # We performed QAT here, kv could also use fp8 format, though current implementation uses bf16
+ index_score = torch.einsum("bshd,btd->bsht", q, self.kv_cache[:bsz, :end_pos // ratio])
+ index_score = (index_score.relu_() * weights.unsqueeze(-1)).sum(dim=2)
+ if world_size > 1:
+ dist.all_reduce(index_score)
+ if start_pos == 0:
+ mask = torch.arange(seqlen // ratio).repeat(seqlen, 1) >= torch.arange(1, seqlen + 1).unsqueeze(1) // ratio
+ index_score += torch.where(mask, float("-inf"), 0)
+ topk_idxs = index_score.topk(min(self.index_topk, end_pos // ratio), dim=-1)[1]
+ if start_pos == 0:
+ mask = topk_idxs >= torch.arange(1, seqlen + 1).unsqueeze(1) // ratio
+ topk_idxs = torch.where(mask, -1, topk_idxs + offset)
+ else:
+ topk_idxs += offset
+ return topk_idxs
+
+
+class Attention(nn.Module):
+ """Multi-head Latent Attention (MLA) with sliding window + optional KV compression.
+ Uses low-rank Q projection (wq_a -> q_norm -> wq_b) and grouped low-rank O projection."""
+ def __init__(self, layer_id: int, args: ModelArgs):
+ super().__init__()
+ self.layer_id = layer_id
+ self.dim = args.dim
+ self.n_heads = args.n_heads
+ self.n_local_heads = args.n_heads // world_size
+ self.q_lora_rank = args.q_lora_rank
+ self.o_lora_rank = args.o_lora_rank
+ self.head_dim = args.head_dim
+ self.rope_head_dim = args.rope_head_dim
+ self.nope_head_dim = args.head_dim - args.rope_head_dim
+ self.n_groups = args.o_groups
+ self.n_local_groups = self.n_groups // world_size
+ self.window_size = args.window_size
+ self.compress_ratio = args.compress_ratios[layer_id]
+ self.eps = args.norm_eps
+
+ self.attn_sink = nn.Parameter(torch.empty(self.n_local_heads, dtype=torch.float32))
+ self.wq_a = Linear(self.dim, self.q_lora_rank)
+ self.q_norm = RMSNorm(self.q_lora_rank, self.eps)
+ self.wq_b = ColumnParallelLinear(self.q_lora_rank, self.n_heads * self.head_dim)
+ self.wkv = Linear(self.dim, self.head_dim)
+ self.kv_norm = RMSNorm(self.head_dim, self.eps)
+ self.wo_a = ColumnParallelLinear(self.n_heads * self.head_dim // self.n_groups, self.n_groups * args.o_lora_rank, dtype=torch.bfloat16)
+ self.wo_b = RowParallelLinear(self.n_groups * args.o_lora_rank, self.dim)
+ self.softmax_scale = self.head_dim ** -0.5
+
+ if self.compress_ratio:
+ self.compressor = Compressor(args, self.compress_ratio, self.head_dim)
+ if self.compress_ratio == 4:
+ self.indexer = Indexer(args, self.compress_ratio)
+ else:
+ self.indexer = None
+
+ kv_cache_size = args.window_size + (args.max_seq_len // self.compress_ratio if self.compress_ratio else 0)
+ self.register_buffer("kv_cache", torch.zeros(args.max_batch_size, kv_cache_size, self.head_dim), persistent=False)
+ if self.compress_ratio:
+ original_seq_len, rope_theta = args.original_seq_len, args.compress_rope_theta
+ else:
+ # disable YaRN and use base rope_theta in pure sliding-window attention
+ original_seq_len, rope_theta = 0, args.rope_theta
+ freqs_cis = precompute_freqs_cis(self.rope_head_dim, args.max_seq_len, original_seq_len,
+ rope_theta, args.rope_factor, args.beta_fast, args.beta_slow)
+ self.register_buffer("freqs_cis", freqs_cis, persistent=False)
+
+ def forward(self, x: torch.Tensor, start_pos: int):
+ bsz, seqlen, _ = x.size()
+ freqs_cis = self.freqs_cis[start_pos:start_pos+seqlen]
+ win = self.window_size
+ ratio = self.compress_ratio
+ rd = self.rope_head_dim
+ if self.compress_ratio and self.compressor.kv_cache is None:
+ self.compressor.kv_cache = self.kv_cache[:, win:]
+ self.compressor.freqs_cis = self.freqs_cis
+ if self.indexer is not None:
+ self.indexer.freqs_cis = self.freqs_cis
+ # q
+ qr = q = self.q_norm(self.wq_a(x))
+ q = self.wq_b(q).unflatten(-1, (self.n_local_heads, self.head_dim))
+ q *= torch.rsqrt(q.square().mean(-1, keepdim=True) + self.eps)
+ apply_rotary_emb(q[..., -rd:], freqs_cis)
+
+ # win kv & topk_idxs
+ kv = self.wkv(x)
+ kv = self.kv_norm(kv)
+ apply_rotary_emb(kv[..., -rd:], freqs_cis)
+ # FP8-simulate non-rope dims to match QAT; rope dims stay bf16 for positional precision
+ act_quant(kv[..., :-rd], 64, scale_fmt, scale_dtype, True)
+ topk_idxs = get_window_topk_idxs(win, bsz, seqlen, start_pos)
+ if self.compress_ratio:
+ offset = kv.size(1) if start_pos == 0 else win
+ if self.indexer is not None:
+ compress_topk_idxs = self.indexer(x, qr, start_pos, offset)
+ else:
+ compress_topk_idxs = get_compress_topk_idxs(ratio, bsz, seqlen, start_pos, offset)
+ topk_idxs = torch.cat([topk_idxs, compress_topk_idxs], dim=-1)
+ topk_idxs = topk_idxs.int()
+
+ # compress kv & attn
+ if start_pos == 0:
+ if seqlen <= win:
+ self.kv_cache[:bsz, :seqlen] = kv
+ else:
+ cutoff = seqlen % win
+ self.kv_cache[:bsz, cutoff: win], self.kv_cache[:bsz, :cutoff] = kv[:, -win:].split([win - cutoff, cutoff], dim=1)
+ if self.compress_ratio:
+ if (kv_compress := self.compressor(x, start_pos)) is not None:
+ kv = torch.cat([kv, kv_compress], dim=1)
+ # We performed QAT here, kv could also use fp8 format, though current implementation uses bf16
+ o = sparse_attn(q, kv, self.attn_sink, topk_idxs, self.softmax_scale)
+ else:
+ self.kv_cache[:bsz, start_pos % win] = kv.squeeze(1)
+ if self.compress_ratio:
+ self.compressor(x, start_pos)
+ o = sparse_attn(q, self.kv_cache[:bsz], self.attn_sink, topk_idxs, self.softmax_scale)
+ apply_rotary_emb(o[..., -rd:], freqs_cis, True)
+
+ # o
+ o = o.view(bsz, seqlen, self.n_local_groups, -1)
+ wo_a = self.wo_a.weight.view(self.n_local_groups, self.o_lora_rank, -1)
+ # NOTE: wo_a is FP8 in checkpoint; could do FP8 einsum here for better perf,
+ # but using BF16 for simplicity.
+ o = torch.einsum("bsgd,grd->bsgr", o, wo_a)
+ x = self.wo_b(o.flatten(2))
+ return x
+
+
+class Gate(nn.Module):
+ """MoE gating: computes expert routing scores and selects top-k experts.
+ Supports hash-based routing (first n_hash_layers) where expert indices are
+ predetermined per token ID, and score-based routing (remaining layers)."""
+ def __init__(self, layer_id: int, args: ModelArgs):
+ super().__init__()
+ self.dim = args.dim
+ self.topk = args.n_activated_experts
+ self.score_func = args.score_func
+ self.route_scale = args.route_scale
+ self.hash = layer_id < args.n_hash_layers
+ self.weight = nn.Parameter(torch.empty(args.n_routed_experts, args.dim))
+ if self.hash:
+ self.tid2eid = nn.Parameter(torch.empty(args.vocab_size, args.n_activated_experts, dtype=torch.int32), requires_grad=False)
+ self.bias = None
+ else:
+ self.bias = nn.Parameter(torch.empty(args.n_routed_experts, dtype=torch.float32))
+
+ def forward(self, x: torch.Tensor, input_ids: Optional[torch.Tensor] = None) -> Tuple[torch.Tensor, torch.Tensor]:
+ scores = linear(x.float(), self.weight.float())
+ if self.score_func == "softmax":
+ scores = scores.softmax(dim=-1)
+ elif self.score_func == "sigmoid":
+ scores = scores.sigmoid()
+ else:
+ scores = F.softplus(scores).sqrt()
+ original_scores = scores
+ # Bias shifts scores for expert selection (topk) but does not affect routing weights.
+ if self.bias is not None:
+ scores = scores + self.bias
+ if self.hash:
+ indices = self.tid2eid[input_ids]
+ else:
+ indices = scores.topk(self.topk, dim=-1)[1]
+ weights = original_scores.gather(1, indices)
+ if self.score_func != "softmax":
+ weights /= weights.sum(dim=-1, keepdim=True)
+ weights *= self.route_scale
+ return weights, indices
+
+
+class Expert(nn.Module):
+ """Single MoE expert: SwiGLU FFN (w1, w2, w3). Computation in float32 for stability."""
+ def __init__(self, dim: int, inter_dim: int, dtype=None, swiglu_limit=0):
+ super().__init__()
+ self.w1 = Linear(dim, inter_dim, dtype=dtype)
+ self.w2 = Linear(inter_dim, dim, dtype=dtype)
+ self.w3 = Linear(dim, inter_dim, dtype=dtype)
+ self.swiglu_limit = swiglu_limit
+
+ def forward(self, x: torch.Tensor, weights: Optional[torch.Tensor] = None) -> torch.Tensor:
+ dtype = x.dtype
+ gate = self.w1(x).float()
+ up = self.w3(x).float()
+ if self.swiglu_limit > 0:
+ up = torch.clamp(up, min=-self.swiglu_limit, max=self.swiglu_limit)
+ gate = torch.clamp(gate, max=self.swiglu_limit)
+ x = F.silu(gate) * up
+ if weights is not None:
+ x = weights * x
+ return self.w2(x.to(dtype))
+
+
+class MoE(nn.Module):
+ """Mixture-of-Experts: gate routes each token to top-k routed experts + 1 shared expert.
+ Experts are sharded across TP ranks; each rank handles n_routed_experts // world_size experts."""
+ def __init__(self, layer_id: int, args: ModelArgs):
+ super().__init__()
+ self.layer_id = layer_id
+ self.dim = args.dim
+ assert args.n_routed_experts % world_size == 0, f"Number of experts must be divisible by world size (world_size={world_size})"
+ self.n_routed_experts = args.n_routed_experts
+ self.n_local_experts = args.n_routed_experts // world_size
+ self.n_activated_experts = args.n_activated_experts
+ self.experts_start_idx = rank * self.n_local_experts
+ self.experts_end_idx = self.experts_start_idx + self.n_local_experts
+ self.gate = Gate(layer_id, args)
+ expert_dtype = torch.float4_e2m1fn_x2 if args.expert_dtype == "fp4" else None
+ self.experts = nn.ModuleList([Expert(args.dim, args.moe_inter_dim, dtype=expert_dtype, swiglu_limit=args.swiglu_limit) if self.experts_start_idx <= i < self.experts_end_idx else None
+ for i in range(self.n_routed_experts)])
+ assert args.n_shared_experts == 1
+ self.shared_experts = Expert(args.dim, args.moe_inter_dim, swiglu_limit=args.swiglu_limit)
+
+ def forward(self, x: torch.Tensor, input_ids: torch.Tensor) -> torch.Tensor:
+ shape = x.size()
+ x = x.view(-1, self.dim)
+ weights, indices = self.gate(x, input_ids.flatten())
+ y = torch.zeros_like(x, dtype=torch.float32)
+ counts = torch.bincount(indices.flatten(), minlength=self.n_routed_experts).tolist()
+ for i in range(self.experts_start_idx, self.experts_end_idx):
+ if counts[i] == 0:
+ continue
+ expert = self.experts[i]
+ idx, top = torch.where(indices == i)
+ y[idx] += expert(x[idx], weights[idx, top, None])
+ if world_size > 1:
+ dist.all_reduce(y)
+ y += self.shared_experts(x)
+ return y.type_as(x).view(shape)
+
+
+class Block(nn.Module):
+ """Transformer block with Hyper-Connections (HC) mixing.
+ Instead of a simple residual, HC maintains `hc_mult` copies of the hidden state.
+ hc_pre: reduces hc copies -> 1 via learned weighted sum (pre-weights from Sinkhorn).
+ hc_post: expands 1 -> hc copies via learned post-weights + combination matrix."""
+ def __init__(self, layer_id: int, args: ModelArgs):
+ super().__init__()
+ self.layer_id = layer_id
+ self.norm_eps = args.norm_eps
+ self.attn = Attention(layer_id, args)
+ self.ffn = MoE(layer_id, args)
+ self.attn_norm = RMSNorm(args.dim, self.norm_eps)
+ self.ffn_norm = RMSNorm(args.dim, self.norm_eps)
+ self.hc_mult = hc_mult = args.hc_mult
+ self.hc_sinkhorn_iters = args.hc_sinkhorn_iters
+ self.hc_eps = args.hc_eps
+ mix_hc = (2 + hc_mult) * hc_mult
+ hc_dim = hc_mult * args.dim
+ with set_dtype(torch.float32):
+ self.hc_attn_fn = nn.Parameter(torch.empty(mix_hc, hc_dim))
+ self.hc_ffn_fn = nn.Parameter(torch.empty(mix_hc, hc_dim))
+ self.hc_attn_base = nn.Parameter(torch.empty(mix_hc))
+ self.hc_ffn_base = nn.Parameter(torch.empty(mix_hc))
+ self.hc_attn_scale = nn.Parameter(torch.empty(3))
+ self.hc_ffn_scale = nn.Parameter(torch.empty(3))
+
+ def hc_pre(self, x: torch.Tensor, hc_fn: torch.Tensor, hc_scale: torch.Tensor, hc_base: torch.Tensor):
+ # x: [b,s,hc,d], hc_fn: [mix_hc,hc*d], hc_scale: [3], hc_base: [mix_hc], y: [b,s,hc,d]
+ shape, dtype = x.size(), x.dtype
+ x = x.flatten(2).float()
+ rsqrt = torch.rsqrt(x.square().mean(-1, keepdim=True) + self.norm_eps)
+ mixes = F.linear(x, hc_fn) * rsqrt
+ pre, post, comb = hc_split_sinkhorn(mixes, hc_scale, hc_base, self.hc_mult, self.hc_sinkhorn_iters, self.hc_eps)
+ y = torch.sum(pre.unsqueeze(-1) * x.view(shape), dim=2)
+ return y.to(dtype), post, comb
+
+ def hc_post(self, x: torch.Tensor, residual: torch.Tensor, post: torch.Tensor, comb: torch.Tensor):
+ # x: [b,s,d], residual: [b,s,hc,d], post: [b,s,hc], comb: [b,s,hc,hc], y: [b,s,hc,d]
+ y = post.unsqueeze(-1) * x.unsqueeze(-2) + torch.sum(comb.unsqueeze(-1) * residual.unsqueeze(-2), dim=2)
+ return y.type_as(x)
+
+ def forward(self, x: torch.Tensor, start_pos: int, input_ids: Optional[torch.Tensor]) -> torch.Tensor:
+ residual = x
+ x, post, comb = self.hc_pre(x, self.hc_attn_fn, self.hc_attn_scale, self.hc_attn_base)
+ x = self.attn_norm(x)
+ x = self.attn(x, start_pos)
+ x = self.hc_post(x, residual, post, comb)
+
+ residual = x
+ x, post, comb = self.hc_pre(x, self.hc_ffn_fn, self.hc_ffn_scale, self.hc_ffn_base)
+ x = self.ffn_norm(x)
+ x = self.ffn(x, input_ids)
+ x = self.hc_post(x, residual, post, comb)
+ return x
+
+
+class ParallelHead(nn.Module):
+
+ def __init__(self, vocab_size: int, dim: int, norm_eps: float = 1e-6, hc_eps: float = 1e-6):
+ super().__init__()
+ self.vocab_size = vocab_size
+ self.dim = dim
+ self.norm_eps = norm_eps
+ self.hc_eps = hc_eps
+ self.part_vocab_size = (vocab_size // world_size)
+ # lm_head in the checkpoint is stored in bf16, while the parameter here is stored in fp32 for easier computation of logits later.
+ self.weight = nn.Parameter(torch.empty(self.part_vocab_size, self.dim, dtype=torch.float32))
+
+ def get_logits(self, x):
+ return F.linear(x[:, -1].float(), self.weight)
+
+ def forward(self, x: torch.Tensor, hc_fn: torch.Tensor, hc_scale: torch.Tensor, hc_base: torch.Tensor, norm: RMSNorm):
+ # x: [b,s,hc,d]
+ x = self.hc_head(x, hc_fn, hc_scale, hc_base)
+ logits = self.get_logits(norm(x))
+ if world_size > 1:
+ all_logits = [torch.empty_like(logits) for _ in range(world_size)]
+ dist.all_gather(all_logits, logits)
+ logits = torch.cat(all_logits, dim=-1)
+ return logits
+
+ def hc_head(self, x: torch.Tensor, hc_fn: torch.Tensor, hc_scale: torch.Tensor, hc_base: torch.Tensor):
+ shape, dtype = x.size(), x.dtype
+ x = x.flatten(2).float()
+ rsqrt = torch.rsqrt(x.square().mean(-1, keepdim=True) + self.norm_eps)
+ mixes = F.linear(x, hc_fn) * rsqrt
+ pre = torch.sigmoid(mixes * hc_scale + hc_base) + self.hc_eps
+ y = torch.sum(pre.unsqueeze(-1) * x.view(shape), dim=2)
+ return y.to(dtype)
+
+
+class MTPBlock(Block):
+
+ def __init__(self, layer_id: int, args: ModelArgs):
+ super().__init__(layer_id, args)
+ self.e_proj = Linear(args.dim, args.dim)
+ self.h_proj = Linear(args.dim, args.dim)
+ self.enorm = RMSNorm(args.dim, args.norm_eps)
+ self.hnorm = RMSNorm(args.dim, args.norm_eps)
+ self.norm = RMSNorm(args.dim, args.norm_eps)
+ self.hc_mult = hc_mult = args.hc_mult
+ hc_dim = hc_mult * args.dim
+ with set_dtype(torch.float32):
+ self.hc_head_fn = nn.Parameter(torch.empty(hc_mult, hc_dim))
+ self.hc_head_base = nn.Parameter(torch.empty(hc_mult))
+ self.hc_head_scale = nn.Parameter(torch.empty(1))
+ self.embed: ParallelEmbedding = None
+ self.head: ParallelHead = None
+
+ @torch.inference_mode()
+ def forward(self, x: torch.Tensor, start_pos: int, input_ids: torch.Tensor) -> torch.Tensor:
+ # x: [b,s,hc,d]
+ assert self.embed is not None and self.head is not None
+ e = self.embed(input_ids)
+ e = self.enorm(e)
+ x = self.hnorm(x)
+ x = self.e_proj(e).unsqueeze(2) + self.h_proj(x)
+ x = super().forward(x, start_pos, input_ids)
+ logits = self.head(x, self.hc_head_fn, self.hc_head_scale, self.hc_head_base, self.norm)
+ return logits
+
+
+class Transformer(nn.Module):
+ """Full DeepSeek-V4 model: embed -> HC-expand -> N blocks -> HC-head -> logits.
+ Sets global state (world_size, rank, default_dtype, scale_fmt, scale_dtype) in __init__."""
+ def __init__(self, args: ModelArgs):
+ global world_size, rank, default_dtype, scale_fmt, scale_dtype
+ world_size = dist.get_world_size() if dist.is_initialized() else 1
+ rank = dist.get_rank() if dist.is_initialized() else 0
+ default_dtype = torch.float8_e4m3fn if args.dtype == "fp8" else torch.bfloat16
+ scale_fmt = "ue8m0" if args.scale_dtype == "fp8" else args.scale_fmt
+ scale_dtype = torch.float8_e8m0fnu if args.scale_dtype == "fp8" else torch.float32
+ super().__init__()
+ self.max_seq_len = args.max_seq_len
+ self.norm_eps = args.norm_eps
+ self.hc_eps = args.hc_eps
+ self.embed = ParallelEmbedding(args.vocab_size, args.dim)
+ self.layers = torch.nn.ModuleList()
+ for layer_id in range(args.n_layers):
+ self.layers.append(Block(layer_id, args))
+ self.norm = RMSNorm(args.dim, self.norm_eps)
+ self.head = ParallelHead(args.vocab_size, args.dim, self.norm_eps, self.hc_eps)
+ self.mtp = torch.nn.ModuleList()
+ for layer_id in range(args.n_mtp_layers):
+ self.mtp.append(MTPBlock(args.n_layers + layer_id, args))
+ self.mtp[-1].embed = self.embed
+ self.mtp[-1].head = self.head
+ self.hc_mult = hc_mult = args.hc_mult
+ hc_dim = hc_mult * args.dim
+ with set_dtype(torch.float32):
+ self.hc_head_fn = nn.Parameter(torch.empty(hc_mult, hc_dim))
+ self.hc_head_base = nn.Parameter(torch.empty(hc_mult))
+ self.hc_head_scale = nn.Parameter(torch.empty(1))
+
+ @torch.inference_mode()
+ def forward(self, input_ids: torch.Tensor, start_pos: int = 0):
+ h = self.embed(input_ids)
+ # Expand to hc_mult copies for Hyper-Connections
+ h = h.unsqueeze(2).repeat(1, 1, self.hc_mult, 1)
+ for layer in self.layers:
+ h = layer(h, start_pos, input_ids)
+ logits = self.head(h, self.hc_head_fn, self.hc_head_scale, self.hc_head_base, self.norm)
+ return logits
+
+
+if __name__ == "__main__":
+ torch.set_default_dtype(torch.bfloat16)
+ torch.set_default_device("cuda")
+ torch.manual_seed(0)
+ args = ModelArgs(n_hash_layers=0)
+ x = torch.randint(0, args.vocab_size, (2, 128))
+ model = Transformer(args)
+
+ print(model(x).size())
+ for i in range(128, 150):
+ print(i, model(x[:, 0:1], i).size())
+
+ h = torch.randn(2, 128, args.hc_mult, args.dim)
+ mtp = model.mtp[0]
+ print(mtp(h, 0, x).size())
+ print(mtp(h[:, 0:1], 1, x[:, 0:1]).size())
diff --git a/reference/vllm/__init__.py b/reference/vllm/__init__.py
new file mode 100644
index 00000000..b0af980a
--- /dev/null
+++ b/reference/vllm/__init__.py
@@ -0,0 +1 @@
+# vLLM reference — read only, do not modify
diff --git a/reference/vllm/reasoning/deepseek_r1_reasoning_parser.py b/reference/vllm/reasoning/deepseek_r1_reasoning_parser.py
new file mode 100644
index 00000000..d24e331b
--- /dev/null
+++ b/reference/vllm/reasoning/deepseek_r1_reasoning_parser.py
@@ -0,0 +1,67 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from collections.abc import Sequence
+
+from vllm.entrypoints.openai.engine.protocol import DeltaMessage
+from vllm.reasoning.basic_parsers import BaseThinkingReasoningParser
+
+
+class DeepSeekR1ReasoningParser(BaseThinkingReasoningParser):
+ """
+ Reasoning parser for DeepSeek R1 model.
+
+ The DeepSeek R1 model uses ... tokens to denote reasoning
+ text. This parser extracts the reasoning content from the model output.
+ """
+
+ @property
+ def start_token(self) -> str:
+ """The token that starts reasoning content."""
+ return ""
+
+ @property
+ def end_token(self) -> str:
+ """The token that ends reasoning content."""
+ return ""
+
+ def extract_reasoning_streaming(
+ self,
+ previous_text: str,
+ current_text: str,
+ delta_text: str,
+ previous_token_ids: Sequence[int],
+ current_token_ids: Sequence[int],
+ delta_token_ids: Sequence[int],
+ ) -> DeltaMessage | None:
+ ret = super().extract_reasoning_streaming(
+ previous_text,
+ current_text,
+ delta_text,
+ previous_token_ids,
+ current_token_ids,
+ delta_token_ids,
+ )
+ if (
+ ret is not None
+ and self.start_token_id not in previous_token_ids
+ and self.start_token_id not in delta_token_ids
+ ):
+ if self.end_token_id in delta_token_ids:
+ # end token in delta with more tokens,
+ # extract reasoning content and content
+ end_index = delta_text.find(self.end_token)
+ reasoning = delta_text[:end_index]
+ content = delta_text[end_index + len(self.end_token) :]
+ return DeltaMessage(
+ reasoning=reasoning,
+ content=content if content else None,
+ )
+ elif self.end_token_id in previous_token_ids:
+ # end token in previous, thinking content ends
+ return DeltaMessage(content=delta_text)
+ else:
+ # no end token in previous or delta, reasoning content continues
+ return DeltaMessage(reasoning=delta_text)
+
+ return ret
diff --git a/reference/vllm/reasoning/deepseek_v3_reasoning_parser.py b/reference/vllm/reasoning/deepseek_v3_reasoning_parser.py
new file mode 100644
index 00000000..bb79afd8
--- /dev/null
+++ b/reference/vllm/reasoning/deepseek_v3_reasoning_parser.py
@@ -0,0 +1,99 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from collections.abc import Iterable, Sequence
+from typing import TYPE_CHECKING
+
+from transformers import PreTrainedTokenizerBase
+
+from vllm.logger import init_logger
+from vllm.reasoning import ReasoningParser
+from vllm.reasoning.deepseek_r1_reasoning_parser import DeepSeekR1ReasoningParser
+
+from .identity_reasoning_parser import IdentityReasoningParser
+
+if TYPE_CHECKING:
+ from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
+ from vllm.entrypoints.openai.engine.protocol import DeltaMessage
+ from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
+
+logger = init_logger(__name__)
+
+
+class DeepSeekV3ReasoningParser(ReasoningParser):
+ """
+ V3 parser that delegates to either DeepSeekR1ReasoningParser or
+ IdentityReasoningParser based on `thinking` and `separate_reasoning`.
+ """
+
+ def __init__(self, tokenizer: PreTrainedTokenizerBase, *args, **kwargs):
+ super().__init__(tokenizer, *args, **kwargs)
+
+ chat_kwargs = kwargs.get("chat_template_kwargs", {}) or {}
+ thinking = bool(chat_kwargs.get("thinking", False))
+ enable_thinking = bool(chat_kwargs.get("enable_thinking", False))
+ thinking = thinking or enable_thinking
+
+ self._parser: ReasoningParser
+ if thinking:
+ self._parser = DeepSeekR1ReasoningParser(tokenizer, *args, **kwargs)
+ else:
+ self._parser = IdentityReasoningParser(tokenizer, *args, **kwargs)
+
+ @property
+ def reasoning_start_str(self) -> str | None:
+ return self._parser.reasoning_start_str
+
+ @property
+ def reasoning_end_str(self) -> str | None:
+ return self._parser.reasoning_end_str
+
+ def is_reasoning_end(self, input_ids: Sequence[int]) -> bool:
+ return self._parser.is_reasoning_end(input_ids)
+
+ def is_reasoning_end_streaming(
+ self, input_ids: Sequence[int], delta_ids: Iterable[int]
+ ) -> bool:
+ return self._parser.is_reasoning_end_streaming(input_ids, delta_ids)
+
+ def extract_content_ids(self, input_ids: list[int]) -> list[int]:
+ return self._parser.extract_content_ids(input_ids)
+
+ def extract_reasoning(
+ self, model_output: str, request: "ChatCompletionRequest | ResponsesRequest"
+ ) -> tuple[str | None, str | None]:
+ return self._parser.extract_reasoning(model_output, request)
+
+ def extract_reasoning_streaming(
+ self,
+ previous_text: str,
+ current_text: str,
+ delta_text: str,
+ previous_token_ids: Sequence[int],
+ current_token_ids: Sequence[int],
+ delta_token_ids: Sequence[int],
+ ) -> "DeltaMessage | None":
+ return self._parser.extract_reasoning_streaming(
+ previous_text,
+ current_text,
+ delta_text,
+ previous_token_ids,
+ current_token_ids,
+ delta_token_ids,
+ )
+
+
+class DeepSeekV3ReasoningWithThinkingParser(DeepSeekV3ReasoningParser):
+ """
+ DeepSeekV3ReasoningParser that defaults to thinking mode.
+ """
+
+ def __init__(self, tokenizer: PreTrainedTokenizerBase, *args, **kwargs):
+ chat_kwargs = kwargs.get("chat_template_kwargs", {}) or {}
+ thinking = chat_kwargs.get("thinking", None)
+ enable_thinking = chat_kwargs.get("enable_thinking", None)
+ if thinking is None and enable_thinking is None:
+ chat_kwargs["thinking"] = True
+ chat_kwargs["enable_thinking"] = True
+ kwargs["chat_template_kwargs"] = chat_kwargs
+ super().__init__(tokenizer, *args, **kwargs)
diff --git a/reference/vllm/tokenizers/deepseek_v4.py b/reference/vllm/tokenizers/deepseek_v4.py
new file mode 100644
index 00000000..2a6aaaf7
--- /dev/null
+++ b/reference/vllm/tokenizers/deepseek_v4.py
@@ -0,0 +1,96 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import copy
+from typing import Any
+
+from transformers import PreTrainedTokenizerFast
+
+from vllm.entrypoints.chat_utils import ChatCompletionMessageParam
+
+from .deepseek_v4_encoding import encode_messages
+from .hf import HfTokenizer, get_cached_tokenizer
+from .protocol import TokenizerLike
+
+
+def get_deepseek_v4_tokenizer(tokenizer: HfTokenizer) -> HfTokenizer:
+ """
+ Wraps a tokenizer to use the custom DeepSeek V4 chat template encoding.
+ """
+ dsv4_tokenizer = copy.copy(tokenizer)
+
+ added_vocab = tokenizer.get_added_vocab()
+ added_vocab_size = len(added_vocab)
+ tokenizer_vocab_size = tokenizer.vocab_size
+
+ class _DeepseekV4Tokenizer(tokenizer.__class__): # type: ignore
+ def apply_chat_template(
+ self,
+ messages: list["ChatCompletionMessageParam"],
+ tools: list[dict[str, Any]] | None = None,
+ **kwargs,
+ ) -> str | list[int]:
+ thinking = kwargs.get("thinking", False)
+ enable_thinking = kwargs.get("enable_thinking", False)
+ thinking = thinking or enable_thinking
+ thinking_mode = "thinking" if thinking else "chat"
+
+ conversation = kwargs.get("conversation", messages)
+ messages = conversation.copy()
+ if tools is not None and len(tools) > 0:
+ messages.insert(0, {"role": "system"})
+ messages[0]["tools"] = tools # type: ignore[typeddict-unknown-key]
+
+ reasoning_effort = kwargs.get("reasoning_effort")
+ if not isinstance(reasoning_effort, str):
+ reasoning_effort = None
+ elif reasoning_effort == "none":
+ thinking_mode = "chat"
+ reasoning_effort = None
+ elif reasoning_effort in ("max", "xhigh"):
+ reasoning_effort = "max"
+ else:
+ reasoning_effort = "high"
+
+ encode_config = dict(
+ thinking_mode=thinking_mode,
+ drop_thinking=kwargs.get("drop_thinking", True),
+ reasoning_effort=reasoning_effort,
+ )
+
+ prompt_str = encode_messages(messages, **encode_config) # type: ignore
+
+ if kwargs.get("tokenize", True):
+ tokenizer_kwargs = {
+ k: kwargs[k] for k in ("truncation", "max_length") if k in kwargs
+ }
+ return self.encode(
+ prompt_str,
+ add_special_tokens=False,
+ **tokenizer_kwargs,
+ )
+
+ return prompt_str
+
+ def num_special_tokens_to_add(self) -> int:
+ return len(self.encode(""))
+
+ def __len__(self) -> int:
+ return tokenizer_vocab_size + added_vocab_size
+
+ def get_added_vocab(self) -> dict[str, int]:
+ return added_vocab.copy()
+
+ def __reduce__(self):
+ return get_deepseek_v4_tokenizer, (tokenizer,)
+
+ _DeepseekV4Tokenizer.__name__ = f"DSV4{tokenizer.__class__.__name__}"
+
+ dsv4_tokenizer.__class__ = _DeepseekV4Tokenizer
+ return dsv4_tokenizer
+
+
+class DeepseekV4Tokenizer(TokenizerLike):
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs) -> HfTokenizer:
+ tokenizer = PreTrainedTokenizerFast.from_pretrained(*args, **kwargs)
+ return get_cached_tokenizer(get_deepseek_v4_tokenizer(tokenizer))
diff --git a/reference/vllm/tokenizers/deepseek_v4_encoding.py b/reference/vllm/tokenizers/deepseek_v4_encoding.py
new file mode 100644
index 00000000..6895771e
--- /dev/null
+++ b/reference/vllm/tokenizers/deepseek_v4_encoding.py
@@ -0,0 +1,757 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# ruff: noqa
+# fmt: off
+
+"""
+DeepSeek-V4 Encoding
+
+A self-contained implementation for encoding/decoding DeepSeek-V4 chat messages
+with tool calling, thinking mode, and quick instruction task support.
+"""
+
+from typing import Any, Dict, List, Union, Optional, Tuple
+import copy
+import json
+
+import regex as re
+
+# ============================================================
+# Special Tokens
+# ============================================================
+
+bos_token: str = "<｜begin▁of▁sentence｜>"
+eos_token: str = "<｜end▁of▁sentence｜>"
+thinking_start_token: str = ""
+thinking_end_token: str = ""
+dsml_token: str = "｜DSML｜"
+
+USER_SP_TOKEN = "<｜User｜>"
+ASSISTANT_SP_TOKEN = "<｜Assistant｜>"
+LATEST_REMINDER_SP_TOKEN = "<｜latest_reminder｜>"
+
+# Task special tokens for internal classification tasks
+DS_TASK_SP_TOKENS = {
+ "action": "<｜action｜>",
+ "query": "<｜query｜>",
+ "authority": "<｜authority｜>",
+ "domain": "<｜domain｜>",
+ "title": "<｜title｜>",
+ "read_url": "<｜read_url｜>",
+}
+VALID_TASKS = set(DS_TASK_SP_TOKENS.keys())
+
+# ============================================================
+# Templates
+# ============================================================
+
+system_msg_template: str = "{content}"
+user_msg_template: str = "{content}"
+latest_reminder_msg_template: str = "{content}"
+assistant_msg_template: str = "{reasoning}{content}{tool_calls}" + eos_token
+assistant_msg_wo_eos_template: str = "{reasoning}{content}{tool_calls}"
+thinking_template: str = "{reasoning}"
+
+response_format_template: str = (
+ "## Response Format:\n\nYou MUST strictly adhere to the following schema to reply:\n{schema}"
+)
+tool_call_template: str = (
+ "<{dsml_token}invoke name=\"{name}\">\n{arguments}\n"
+)
+tool_calls_template = (
+ "<{dsml_token}{tc_block_name}>\n{tool_calls}\n"
+)
+tool_calls_block_name: str = "tool_calls"
+
+tool_output_template: str = (
+ "{content}"
+)
+
+REASONING_EFFORT_MAX = (
+ "Reasoning Effort: Absolute maximum with no shortcuts permitted.\n"
+ "You MUST be very thorough in your thinking and comprehensively decompose the problem to resolve the root cause, rigorously stress-testing your logic against all potential paths, edge cases, and adversarial scenarios.\n"
+ "Explicitly write out your entire deliberation process, documenting every intermediate step, considered alternative, and rejected hypothesis to ensure absolutely no assumption is left unchecked.\n\n"
+)
+
+TOOLS_TEMPLATE = """## Tools
+
+You have access to a set of tools to help answer the user's question. You can invoke tools by writing a "<{dsml_token}tool_calls>" block like the following:
+
+<{dsml_token}tool_calls>
+<{dsml_token}invoke name="$TOOL_NAME">
+<{dsml_token}parameter name="$PARAMETER_NAME" string="true|false">$PARAMETER_VALUE
+...
+
+<{dsml_token}invoke name="$TOOL_NAME2">
+...
+
+
+
+String parameters should be specified as is and set `string="true"`. For all other types (numbers, booleans, arrays, objects), pass the value in JSON format and set `string="false"`.
+
+If thinking_mode is enabled (triggered by {thinking_start_token}), you MUST output your complete reasoning inside {thinking_start_token}...{thinking_end_token} BEFORE any tool calls or final response.
+
+Otherwise, output directly after {thinking_end_token} with tool calls or final response.
+
+### Available Tool Schemas
+
+{tool_schemas}
+
+You MUST strictly follow the above defined tool name and parameter schemas to invoke tool calls.
+"""
+
+# ============================================================
+# Utility Functions
+# ============================================================
+
+def to_json(value: Any) -> str:
+ """Serialize a value to JSON string."""
+ try:
+ return json.dumps(value, ensure_ascii=False)
+ except Exception:
+ return json.dumps(value, ensure_ascii=True)
+
+
+def tools_from_openai_format(tools):
+ """Extract function definitions from OpenAI-format tool list."""
+ return [tool["function"] for tool in tools]
+
+
+def tool_calls_from_openai_format(tool_calls):
+ """Convert OpenAI-format tool calls to internal format."""
+ return [
+ {
+ "name": tool_call["function"]["name"],
+ "arguments": tool_call["function"]["arguments"],
+ }
+ for tool_call in tool_calls
+ ]
+
+
+def tool_calls_to_openai_format(tool_calls):
+ """Convert internal tool calls to OpenAI format."""
+ return [
+ {
+ "type": "function",
+ "function": {
+ "name": tool_call["name"],
+ "arguments": tool_call["arguments"],
+ }
+ }
+ for tool_call in tool_calls
+ ]
+
+
+def encode_arguments_to_dsml(tool_call: Dict[str, Any]) -> str:
+ """
+ Encode tool call arguments into DSML parameter format.
+
+ Args:
+ tool_call: Dict with "name" and "arguments" keys.
+
+ Returns:
+ DSML-formatted parameter string.
+ """
+ p_dsml_template = '<{dsml_token}parameter name="{key}" string="{is_str}">{value}'
+ P_dsml_strs = []
+
+ if isinstance(tool_call["arguments"], str):
+ arguments = json.loads(tool_call["arguments"])
+ else:
+ arguments = tool_call["arguments"]
+
+ for k, v in arguments.items():
+ p_dsml_str = p_dsml_template.format(
+ dsml_token=dsml_token,
+ key=k,
+ is_str="true" if isinstance(v, str) else "false",
+ value=v if isinstance(v, str) else to_json(v),
+ )
+ P_dsml_strs.append(p_dsml_str)
+
+ return "\n".join(P_dsml_strs)
+
+
+def decode_dsml_to_arguments(tool_name: str, tool_args: Dict[str, Tuple[str, str]]) -> Dict[str, str]:
+ """
+ Decode DSML parameters back to a tool call dict.
+
+ Args:
+ tool_name: Name of the tool.
+ tool_args: Dict mapping param_name -> (value, is_string_flag).
+
+ Returns:
+ Dict with "name" and "arguments" (JSON string) keys.
+ """
+ def _decode_value(key: str, value: str, string: str):
+ if string == "true":
+ value = to_json(value)
+ return f"{to_json(key)}: {value}"
+
+ tool_args_json = "{" + ", ".join([_decode_value(k, v, string=is_str) for k, (v, is_str) in tool_args.items()]) + "}"
+ return dict(name=tool_name, arguments=tool_args_json)
+
+
+def render_tools(tools: List[Dict[str, Union[str, Dict[str, Any]]]]) -> str:
+ """
+ Render tool schemas into the system prompt format.
+
+ Args:
+ tools: List of tool schema dicts (each with name, description, parameters).
+
+ Returns:
+ Formatted tools section string.
+ """
+ tools_json = [to_json(t) for t in tools]
+
+ return TOOLS_TEMPLATE.format(
+ tool_schemas="\n".join(tools_json),
+ dsml_token=dsml_token,
+ thinking_start_token=thinking_start_token,
+ thinking_end_token=thinking_end_token,
+ )
+
+
+def find_last_user_index(messages: List[Dict[str, Any]]) -> int:
+ """Find the index of the last user/developer message."""
+ last_user_index = -1
+ for idx in range(len(messages) - 1, -1, -1):
+ if messages[idx].get("role") in ["user", "developer"]:
+ last_user_index = idx
+ break
+ return last_user_index
+
+
+# ============================================================
+# Message Rendering
+# ============================================================
+
+def render_message(index: int, messages: List[Dict[str, Any]], thinking_mode: str, drop_thinking: bool = True, reasoning_effort: Optional[str] = None) -> str:
+ """
+ Render a single message at the given index into its encoded string form.
+
+ This is the core function that converts each message in the conversation
+ into the DeepSeek-V4 format.
+
+ Args:
+ index: Index of the message to render.
+ messages: Full list of messages in the conversation.
+ thinking_mode: Either "chat" or "thinking".
+ drop_thinking: Whether to drop reasoning content from earlier turns.
+ reasoning_effort: Optional reasoning effort level ("max", "high", or None).
+
+ Returns:
+ Encoded string for this message.
+ """
+ assert 0 <= index < len(messages)
+ assert thinking_mode in ["chat", "thinking"], f"Invalid thinking_mode `{thinking_mode}`"
+
+ prompt = ""
+ msg = messages[index]
+ last_user_idx = find_last_user_index(messages)
+
+ role = msg.get("role")
+ content = msg.get("content")
+ tools = msg.get("tools")
+ response_format = msg.get("response_format")
+ tool_calls = msg.get("tool_calls")
+ reasoning = msg.get("reasoning")
+ wo_eos = msg.get("wo_eos", False)
+
+ if tools:
+ tools = tools_from_openai_format(tools)
+ if tool_calls:
+ tool_calls = tool_calls_from_openai_format(tool_calls)
+
+ # Reasoning effort prefix (only at index 0 in thinking mode with max effort)
+ assert reasoning_effort in ['max', None, 'high'], f"Invalid reasoning effort: {reasoning_effort}"
+ if index == 0 and thinking_mode == "thinking" and reasoning_effort == 'max':
+ prompt += REASONING_EFFORT_MAX
+
+ if role == "system":
+ prompt += system_msg_template.format(content=content or "")
+ if tools:
+ prompt += "\n\n" + render_tools(tools)
+ if response_format:
+ prompt += "\n\n" + response_format_template.format(schema=to_json(response_format))
+
+ elif role == "developer":
+ assert content, f"Invalid message for role `{role}`: {msg}"
+
+ content_developer = USER_SP_TOKEN
+ content_developer += content
+
+ if tools:
+ content_developer += "\n\n" + render_tools(tools)
+ if response_format:
+ content_developer += "\n\n" + response_format_template.format(schema=to_json(response_format))
+
+ prompt += user_msg_template.format(content=content_developer)
+
+ elif role == "user":
+ prompt += USER_SP_TOKEN
+
+ # Handle content blocks (tool results mixed with text)
+ content_blocks = msg.get("content_blocks")
+ if content_blocks:
+ parts = []
+ for block in content_blocks:
+ block_type = block.get("type")
+ if block_type == "text":
+ parts.append(block.get("text", ""))
+ elif block_type == "tool_result":
+ tool_content = block.get("content", "")
+ if isinstance(tool_content, list):
+ text_parts = []
+ for b in tool_content:
+ if b.get("type") == "text":
+ text_parts.append(b.get("text", ""))
+ else:
+ text_parts.append(f"[Unsupported {b.get('type')}]")
+ tool_content = "\n\n".join(text_parts)
+ parts.append(tool_output_template.format(content=tool_content))
+ else:
+ parts.append(f"[Unsupported {block_type}]")
+ prompt += "\n\n".join(parts)
+ else:
+ prompt += content or ""
+
+ elif role == "latest_reminder":
+ prompt += LATEST_REMINDER_SP_TOKEN + latest_reminder_msg_template.format(content=content)
+
+ elif role == "tool":
+ raise NotImplementedError("deepseek_v4 merges tool messages into user; please preprocess with merge_tool_messages()")
+
+ elif role == "assistant":
+ thinking_part = ""
+ tc_content = ""
+
+ if tool_calls:
+ tc_list = [
+ tool_call_template.format(
+ dsml_token=dsml_token,
+ name=tc.get("name"),
+ arguments=encode_arguments_to_dsml(tc)
+ )
+ for tc in tool_calls
+ ]
+ tc_content += '\n\n' + tool_calls_template.format(
+ dsml_token=dsml_token,
+ tool_calls="\n".join(tc_list),
+ tc_block_name=tool_calls_block_name,
+ )
+
+ summary_content = content or ""
+ reasoning = reasoning or ""
+
+ # Check if previous message has a task - if so, this is a task output (no thinking)
+ prev_has_task = index - 1 >= 0 and messages[index - 1].get("task") is not None
+
+ if thinking_mode == "thinking" and not prev_has_task:
+ if not drop_thinking or index > last_user_idx:
+ thinking_part = thinking_template.format(reasoning=reasoning) + thinking_end_token
+ else:
+ thinking_part = ""
+
+ if wo_eos:
+ prompt += assistant_msg_wo_eos_template.format(
+ reasoning=thinking_part,
+ content=summary_content,
+ tool_calls=tc_content,
+ )
+ else:
+ prompt += assistant_msg_template.format(
+ reasoning=thinking_part,
+ content=summary_content,
+ tool_calls=tc_content,
+ )
+ else:
+ raise NotImplementedError(f"Unknown role: {role}")
+
+ # Append transition tokens based on what follows
+ if index + 1 < len(messages) and messages[index + 1].get("role") not in ["assistant", "latest_reminder"]:
+ return prompt
+
+ task = messages[index].get("task")
+ if task is not None:
+ # Task special token for internal classification tasks
+ assert task in VALID_TASKS, f"Invalid task: '{task}'. Valid tasks are: {list(VALID_TASKS)}"
+ task_sp_token = DS_TASK_SP_TOKENS[task]
+
+ if task != "action":
+ # Non-action tasks: append task sp token directly after the message
+ prompt += task_sp_token
+ else:
+ # Action task: append Assistant + thinking token + action sp token
+ prompt += ASSISTANT_SP_TOKEN
+ prompt += thinking_end_token if thinking_mode != "thinking" else thinking_start_token
+ prompt += task_sp_token
+
+ elif messages[index].get("role") in ["user", "developer"]:
+ # Normal generation: append Assistant + thinking token
+ prompt += ASSISTANT_SP_TOKEN
+ if not drop_thinking and thinking_mode == "thinking":
+ prompt += thinking_start_token
+ elif drop_thinking and thinking_mode == "thinking" and index >= last_user_idx:
+ prompt += thinking_start_token
+ else:
+ prompt += thinking_end_token
+
+ return prompt
+
+
+# ============================================================
+# Preprocessing
+# ============================================================
+
+def merge_tool_messages(messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+ """
+ Merge tool messages into the preceding user message using content_blocks format.
+
+ DeepSeek-V4 does not have a standalone "tool" role; instead, tool results
+ are encoded as blocks within user messages.
+
+ This function converts a standard OpenAI-format conversation (with separate
+ "tool" role messages) into V4 format where tool results are merged into
+ user messages.
+
+ Args:
+ messages: List of message dicts in OpenAI format.
+
+ Returns:
+ Processed message list with tool messages merged into user messages.
+ """
+ merged: List[Dict[str, Any]] = []
+
+ for msg in messages:
+ msg = copy.deepcopy(msg)
+ role = msg.get("role")
+
+ if role == "tool":
+ # Convert tool message to a user message with tool_result block
+ tool_block = {
+ "type": "tool_result",
+ "tool_use_id": msg.get("tool_call_id", ""),
+ "content": msg.get("content", ""),
+ }
+ # Merge into previous message if it's already a user (merged tool)
+ if merged and merged[-1].get("role") == "user" and "content_blocks" in merged[-1]:
+ merged[-1]["content_blocks"].append(tool_block)
+ else:
+ merged.append({
+ "role": "user",
+ "content_blocks": [tool_block],
+ })
+ elif role == "user":
+ text_block = {"type": "text", "text": msg.get("content", "")}
+ if merged and merged[-1].get("role") == "user" and "content_blocks" in merged[-1] and merged[-1].get("task") is None:
+ merged[-1]["content_blocks"].append(text_block)
+ else:
+ new_msg = {
+ "role": "user",
+ "content": msg.get("content", ""),
+ "content_blocks": [text_block],
+ }
+ # Preserve extra fields (task, wo_eos, mask, etc.)
+ for key in ("task", "wo_eos", "mask"):
+ if key in msg:
+ new_msg[key] = msg[key]
+ merged.append(new_msg)
+ else:
+ merged.append(msg)
+
+ return merged
+
+
+def sort_tool_results_by_call_order(messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+ """
+ Sort tool_result blocks within user messages by the order of tool_calls
+ in the preceding assistant message.
+
+ Args:
+ messages: Preprocessed message list (after merge_tool_messages).
+
+ Returns:
+ Message list with sorted tool result blocks.
+ """
+ last_tool_call_order: Dict[str, int] = {}
+
+ for msg in messages:
+ role = msg.get("role")
+ if role == "assistant" and msg.get("tool_calls"):
+ last_tool_call_order = {}
+ for idx, tc in enumerate(msg["tool_calls"]):
+ tc_id = tc.get("id") or tc.get("function", {}).get("id", "")
+ if tc_id:
+ last_tool_call_order[tc_id] = idx
+
+ elif role == "user" and msg.get("content_blocks"):
+ tool_blocks = [b for b in msg["content_blocks"] if b.get("type") == "tool_result"]
+ if len(tool_blocks) > 1 and last_tool_call_order:
+ sorted_blocks = sorted(
+ tool_blocks,
+ key=lambda b: last_tool_call_order.get(b.get("tool_use_id", ""), 0)
+ )
+ sorted_idx = 0
+ new_blocks = []
+ for block in msg["content_blocks"]:
+ if block.get("type") == "tool_result":
+ new_blocks.append(sorted_blocks[sorted_idx])
+ sorted_idx += 1
+ else:
+ new_blocks.append(block)
+ msg["content_blocks"] = new_blocks
+
+ return messages
+
+
+# ============================================================
+# Main Encoding Function
+# ============================================================
+
+def encode_messages(
+ messages: List[Dict[str, Any]],
+ thinking_mode: str,
+ context: Optional[List[Dict[str, Any]]] = None,
+ drop_thinking: bool = True,
+ add_default_bos_token: bool = True,
+ reasoning_effort: Optional[str] = None,
+) -> str:
+ """
+ Encode a list of messages into the DeepSeek-V4 prompt format.
+
+ This is the main entry point for encoding conversations. It handles:
+ - BOS token insertion
+ - Thinking mode with optional reasoning content dropping
+ - Tool message merging into user messages
+ - Multi-turn conversation context
+
+ Args:
+ messages: List of message dicts to encode.
+ thinking_mode: Either "chat" or "thinking".
+ context: Optional preceding context messages (already encoded prefix).
+ drop_thinking: If True, drop reasoning from earlier assistant turns
+ (only keep reasoning for messages after the last user message).
+ add_default_bos_token: Whether to prepend BOS token at conversation start.
+ reasoning_effort: Optional reasoning effort level ("max", "high", or None).
+
+ Returns:
+ The encoded prompt string.
+ """
+ context = context if context else []
+
+ # Preprocess: merge tool messages and sort tool results
+ messages = merge_tool_messages(messages)
+ messages = sort_tool_results_by_call_order(context + messages)[len(context):]
+ if context:
+ context = merge_tool_messages(context)
+ context = sort_tool_results_by_call_order(context)
+
+ full_messages = context + messages
+
+ prompt = bos_token if add_default_bos_token and len(context) == 0 else ""
+
+ # Resolve drop_thinking: if any message has tools defined, don't drop thinking
+ effective_drop_thinking = drop_thinking
+ if any(m.get("tools") for m in full_messages):
+ effective_drop_thinking = False
+
+ if thinking_mode == "thinking" and effective_drop_thinking:
+ full_messages = _drop_thinking_messages(full_messages)
+ # After dropping, recalculate how many messages to render
+ # (context may have shrunk too)
+ num_to_render = len(full_messages) - len(_drop_thinking_messages(context))
+ context_len = len(full_messages) - num_to_render
+ else:
+ num_to_render = len(messages)
+ context_len = len(context)
+
+ for idx in range(num_to_render):
+ prompt += render_message(
+ idx + context_len,
+ full_messages,
+ thinking_mode=thinking_mode,
+ drop_thinking=effective_drop_thinking,
+ reasoning_effort=reasoning_effort,
+ )
+
+ return prompt
+
+
+def _drop_thinking_messages(messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+ """
+ Drop reasoning and non-essential messages before the last user message.
+
+ Behavior:
+ - Messages with role in ["user", "system", "tool", "latest_reminder"] are always kept.
+ - Messages at or after the last user index are always kept.
+ - Assistant messages before the last user get reasoning removed.
+ - Developer messages before the last user are dropped entirely.
+ """
+ last_user_idx = find_last_user_index(messages)
+ result = []
+ keep_roles = {"user", "system", "tool", "latest_reminder", "direct_search_results"}
+
+ for idx, msg in enumerate(messages):
+ role = msg.get("role")
+ if role in keep_roles or idx >= last_user_idx:
+ result.append(msg)
+ elif role == "assistant":
+ msg = copy.copy(msg)
+ msg.pop("reasoning", None)
+ result.append(msg)
+ # developer and other roles before last_user_idx are dropped
+
+ return result
+
+
+# ============================================================
+# Parsing (Decoding model output)
+# ============================================================
+
+def _read_until_stop(index: int, text: str, stop: List[str]) -> Tuple[int, str, Optional[str]]:
+ """
+ Read text from index until one of the stop strings is found.
+
+ Returns:
+ Tuple of (new_index, content_before_stop, matched_stop_string_or_None).
+ """
+ min_pos = len(text)
+ matched_stop = None
+
+ for s in stop:
+ pos = text.find(s, index)
+ if pos != -1 and pos < min_pos:
+ min_pos = pos
+ matched_stop = s
+
+ if matched_stop:
+ content = text[index:min_pos]
+ return min_pos + len(matched_stop), content, matched_stop
+ else:
+ content = text[index:]
+ return len(text), content, None
+
+
+def parse_tool_calls(index: int, text: str) -> Tuple[int, Optional[str], List[Dict[str, str]]]:
+ """
+ Parse DSML tool calls from text starting at the given index.
+
+ Args:
+ index: Starting position in text.
+ text: The full text to parse.
+
+ Returns:
+ Tuple of (new_index, last_stop_token, list_of_tool_call_dicts).
+ Each tool call dict has "name" and "arguments" keys.
+ """
+ tool_calls: List[Dict[str, Any]] = []
+ stop_token = None
+ tool_calls_end_token = f""
+
+ while index < len(text):
+ index, content_before, stop_token = _read_until_stop(index, text, [f"<{dsml_token}invoke", tool_calls_end_token])
+ if content_before != ">\n":
+ raise ValueError(f"Tool call format error: expected '>\\n' but got '{content_before}'")
+
+ if stop_token == tool_calls_end_token:
+ break
+
+ if stop_token is None:
+ raise ValueError("Missing special token in tool calls")
+
+ index, tool_name_content, stop_token = _read_until_stop(index, text, [f"<{dsml_token}parameter", f"\n$', tool_name_content, flags=re.DOTALL)
+ if len(p_tool_name) != 1:
+ raise ValueError(f"Tool name format error: '{tool_name_content}'")
+ tool_name = p_tool_name[0]
+
+ tool_args: Dict[str, Tuple[str, str]] = {}
+ while stop_token == f"<{dsml_token}parameter":
+ index, param_content, stop_token = _read_until_stop(index, text, [f"/{dsml_token}parameter"])
+
+ param_kv = re.findall(r'^ name="(.*?)" string="(true|false)">(.*?)<$', param_content, flags=re.DOTALL)
+ if len(param_kv) != 1:
+ raise ValueError(f"Parameter format error: '{param_content}'")
+ param_name, string, param_value = param_kv[0]
+
+ if param_name in tool_args:
+ raise ValueError(f"Duplicate parameter name: '{param_name}'")
+ tool_args[param_name] = (param_value, string)
+
+ index, content, stop_token = _read_until_stop(index, text, [f"<{dsml_token}parameter", f"\n":
+ raise ValueError(f"Parameter format error: expected '>\\n' but got '{content}'")
+
+ tool_call = decode_dsml_to_arguments(tool_name=tool_name, tool_args=tool_args)
+ tool_calls.append(tool_call)
+
+ return index, stop_token, tool_calls
+
+
+def parse_message_from_completion_text(text: str, thinking_mode: str) -> Dict[str, Any]:
+ """
+ Parse a model completion text into a structured assistant message.
+
+ This function takes the raw text output from the model (a single assistant turn)
+ and extracts:
+ - reasoning (thinking block)
+ - content (summary/response)
+ - tool_calls (if any)
+
+ NOTE: This function is designed to parse only correctly formatted strings and
+ will raise ValueError for malformed output.
+
+ Args:
+ text: The raw completion text (including EOS token).
+ thinking_mode: Either "chat" or "thinking".
+
+ Returns:
+ Dict with keys: "role", "content", "reasoning", "tool_calls".
+ tool_calls are in OpenAI format.
+ """
+ summary_content, reasoning = "", ""
+ tool_calls: List[Dict[str, str]] = []
+ index, stop_token = 0, None
+ tool_calls_start_token = f"\n\n<{dsml_token}{tool_calls_block_name}"
+
+ is_thinking = thinking_mode == "thinking"
+ is_tool_calling = False
+
+ if is_thinking:
+ index, content_delta, stop_token = _read_until_stop(index, text, [thinking_end_token, tool_calls_start_token])
+ reasoning = content_delta
+ if stop_token != thinking_end_token:
+ raise ValueError("Invalid thinking format: missing ")
+
+ index, content_delta, stop_token = _read_until_stop(index, text, [eos_token, tool_calls_start_token])
+ summary_content = content_delta
+ if stop_token == tool_calls_start_token:
+ is_tool_calling = True
+ else:
+ if stop_token != eos_token:
+ raise ValueError("Invalid format: missing EOS token")
+
+ if is_tool_calling:
+ index, stop_token, tool_calls = parse_tool_calls(index, text)
+
+ index, tool_ends_text, stop_token = _read_until_stop(index, text, [eos_token])
+ if tool_ends_text:
+ raise ValueError("Unexpected content after tool calls")
+
+ if len(text) != index or stop_token not in [eos_token, None]:
+ raise ValueError("Unexpected content at end")
+
+ for sp_token in [bos_token, eos_token, thinking_start_token, thinking_end_token, dsml_token]:
+ if sp_token in summary_content or sp_token in reasoning:
+ raise ValueError(f"Unexpected special token '{sp_token}' in content")
+
+ return {
+ "role": "assistant",
+ "content": summary_content,
+ "reasoning": reasoning,
+ "tool_calls": tool_calls_to_openai_format(tool_calls)
+ }
+
+# fmt: on
diff --git a/reference/vllm/tool_parsers/deepseekv32_tool_parser.py b/reference/vllm/tool_parsers/deepseekv32_tool_parser.py
new file mode 100644
index 00000000..02182e22
--- /dev/null
+++ b/reference/vllm/tool_parsers/deepseekv32_tool_parser.py
@@ -0,0 +1,322 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import json
+import uuid
+from collections.abc import Sequence
+from typing import Any
+
+import regex as re
+
+from vllm.entrypoints.openai.chat_completion.protocol import (
+ ChatCompletionRequest,
+)
+from vllm.entrypoints.openai.engine.protocol import (
+ DeltaFunctionCall,
+ DeltaMessage,
+ DeltaToolCall,
+ ExtractedToolCallInformation,
+ FunctionCall,
+ ToolCall,
+)
+from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
+from vllm.logger import init_logger
+from vllm.tokenizers import TokenizerLike
+from vllm.tool_parsers.abstract_tool_parser import (
+ Tool,
+ ToolParser,
+)
+from vllm.tool_parsers.utils import partial_tag_overlap
+
+logger = init_logger(__name__)
+
+
+class DeepSeekV32ToolParser(ToolParser):
+ """
+ example tool call content:
+ <｜DSML｜function_calls>
+ <｜DSML｜invoke name="get_weather">
+ <｜DSML｜parameter name="location" string="true">杭州
+ <｜DSML｜parameter name="date" string="true">2024-01-16
+
+ <｜DSML｜invoke name="get_weather">
+ <｜DSML｜parameter name="location" string="true">北京
+ <｜DSML｜parameter name="date" string="true">2024-01-16
+
+
+ """
+
+ tool_call_start_token: str = "<｜DSML｜function_calls>"
+ tool_call_end_token: str = ""
+
+ def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None):
+ super().__init__(tokenizer, tools)
+
+ self.prev_tool_call_arr: list[dict] = []
+
+ # Streaming state
+ self.current_tool_index: int = 0
+ self._sent_content_idx: int = 0
+
+ # Regex patterns for complete parsing
+ self.tool_call_complete_regex = re.compile(
+ re.escape(self.tool_call_start_token)
+ + r"(.*?)"
+ + re.escape(self.tool_call_end_token),
+ re.DOTALL,
+ )
+ self.invoke_complete_regex = re.compile(
+ r'<｜DSML｜invoke\s+name="([^"]+)"\s*>(.*?)', re.DOTALL
+ )
+ self.parameter_complete_regex = re.compile(
+ r'<｜DSML｜parameter\s+name="([^"]+)"\s+string="(?:true|false)"\s*>(.*?)',
+ re.DOTALL,
+ )
+
+ if not self.model_tokenizer:
+ raise ValueError(
+ "The model tokenizer must be passed to the ToolParser "
+ "constructor during construction."
+ )
+
+ logger.debug(
+ "vLLM Successfully import tool parser %s !", self.__class__.__name__
+ )
+
+ def adjust_request(
+ self, request: ChatCompletionRequest | ResponsesRequest
+ ) -> ChatCompletionRequest | ResponsesRequest:
+ request = super().adjust_request(request)
+ if request.tools and request.tool_choice != "none":
+ # Ensure tool call tokens
+ # (e.g. <｜DSML｜function_calls>, )
+ # are not skippedduring decoding.
+ # Even though they are not marked as special tokens,
+ # setting skip_special_tokens=False ensures proper handling in
+ # transformers 5.x where decoding behavior may have changed.
+ request.skip_special_tokens = False
+ return request
+
+ def _generate_tool_call_id(self) -> str:
+ """Generate a unique tool call ID."""
+ return f"call_{uuid.uuid4().hex[:24]}"
+
+ def _parse_invoke_params(self, invoke_str: str) -> dict:
+ param_dict = dict()
+ for param_name, param_val in self.parameter_complete_regex.findall(invoke_str):
+ param_dict[param_name] = param_val
+ return param_dict
+
+ def _convert_param_value_checked(self, value: str, param_type: str) -> Any:
+ """Convert parameter value to the correct type."""
+ if value.lower() == "null":
+ return None
+
+ param_type = param_type.lower()
+ if param_type in ["string", "str", "text"]:
+ return value
+ elif param_type in ["integer", "int"]:
+ return int(value)
+ elif param_type in ["number", "float"]:
+ val = float(value)
+ return val if val != int(val) else int(val)
+ elif param_type in ["boolean", "bool"]:
+ value = value.strip()
+ if value.lower() not in ["false", "0", "true", "1"]:
+ raise ValueError("Invalid boolean value")
+ return value.lower() in ["true", "1"]
+ elif param_type in ["object", "array"]:
+ return json.loads(value)
+ else:
+ return json.loads(value)
+
+ def _convert_param_value(self, value: str, param_type: str | list[str]) -> Any:
+ """Convert parameter value to the correct type."""
+ if not isinstance(param_type, list):
+ param_type = [param_type]
+ for current_type in param_type:
+ try:
+ return self._convert_param_value_checked(value, current_type)
+ except Exception:
+ continue
+ # return value as fallback
+ return value
+
+ def _convert_params_with_schema(
+ self,
+ function_name: str,
+ param_dict: dict[str, str],
+ ) -> dict[str, Any]:
+ """Convert raw string param values using the tool schema types."""
+ param_config: dict = {}
+ if self.tools:
+ for tool in self.tools:
+ if (
+ hasattr(tool, "function")
+ and tool.function.name == function_name
+ and hasattr(tool.function, "parameters")
+ ):
+ schema = tool.function.parameters
+ if isinstance(schema, dict) and "properties" in schema:
+ param_config = schema["properties"]
+ break
+
+ converted: dict[str, Any] = {}
+ for name, value in param_dict.items():
+ param_type = "string"
+ if name in param_config and isinstance(param_config[name], dict):
+ param_type = param_config[name].get("type", "string")
+ converted[name] = self._convert_param_value(value, param_type)
+ return converted
+
+ def extract_tool_calls(
+ self,
+ model_output: str,
+ request: ChatCompletionRequest,
+ ) -> ExtractedToolCallInformation:
+ """Extract tool calls from complete model output (non-streaming)."""
+ # Quick check
+ if self.tool_call_start_token not in model_output:
+ return ExtractedToolCallInformation(
+ tools_called=False, tool_calls=[], content=model_output
+ )
+
+ try:
+ tool_calls = []
+
+ # Find all complete tool_call blocks
+ for tool_call_match in self.tool_call_complete_regex.findall(model_output):
+ # Find all invokes within this tool_call
+ for invoke_name, invoke_content in self.invoke_complete_regex.findall(
+ tool_call_match
+ ):
+ param_dict = self._parse_invoke_params(invoke_content)
+ params = self._convert_params_with_schema(invoke_name, param_dict)
+ tool_calls.append(
+ ToolCall(
+ type="function",
+ function=FunctionCall(
+ name=invoke_name,
+ arguments=json.dumps(params, ensure_ascii=False),
+ ),
+ )
+ )
+
+ if not tool_calls:
+ return ExtractedToolCallInformation(
+ tools_called=False, tool_calls=[], content=model_output
+ )
+
+ # Extract content before first tool call
+ first_tool_idx = model_output.find(self.tool_call_start_token)
+ content = model_output[:first_tool_idx] if first_tool_idx > 0 else None
+
+ return ExtractedToolCallInformation(
+ tools_called=True, tool_calls=tool_calls, content=content
+ )
+
+ except Exception:
+ logger.exception("Error extracting tool calls")
+ return ExtractedToolCallInformation(
+ tools_called=False, tool_calls=[], content=model_output
+ )
+
+ def _reset_streaming_state(self):
+ """Reset all streaming state."""
+ self.current_tool_index = 0
+ self._sent_content_idx = 0
+ self.prev_tool_call_arr.clear()
+ self.streamed_args_for_tool.clear()
+
+ def _extract_delta_tool_calls(
+ self,
+ current_text: str,
+ request: ChatCompletionRequest | None,
+ ) -> list[DeltaToolCall]:
+ """Extract DeltaToolCalls from newly completed blocks.
+
+ Tracks progress via ``current_tool_index`` so each block is
+ extracted exactly once across successive streaming calls.
+ """
+ complete_invokes = self.invoke_complete_regex.findall(current_text)
+ delta_tool_calls: list[DeltaToolCall] = []
+
+ while len(complete_invokes) > self.current_tool_index:
+ invoke_name, invoke_body = complete_invokes[self.current_tool_index]
+ param_dict = self._parse_invoke_params(invoke_body)
+
+ converted = self._convert_params_with_schema(invoke_name, param_dict)
+ args_json = json.dumps(converted, ensure_ascii=False)
+ idx = self.current_tool_index
+ self.current_tool_index += 1
+
+ self.prev_tool_call_arr.append(
+ {"name": invoke_name, "arguments": converted}
+ )
+ self.streamed_args_for_tool.append(args_json)
+
+ delta_tool_calls.append(
+ DeltaToolCall(
+ index=idx,
+ id=self._generate_tool_call_id(),
+ function=DeltaFunctionCall(
+ name=invoke_name,
+ arguments=args_json,
+ ),
+ type="function",
+ )
+ )
+
+ return delta_tool_calls
+
+ def _extract_content(self, current_text: str) -> str | None:
+ """Return unsent non-tool-call text, or None.
+
+ Holds back any suffix that could be a partial start marker
+ so that split markers are never leaked as content.
+ """
+ if self.tool_call_start_token not in current_text:
+ overlap = partial_tag_overlap(current_text, self.tool_call_start_token)
+ sendable_idx = len(current_text) - overlap
+ else:
+ sendable_idx = current_text.index(self.tool_call_start_token)
+
+ if sendable_idx > self._sent_content_idx:
+ content = current_text[self._sent_content_idx : sendable_idx]
+ self._sent_content_idx = sendable_idx
+ return content
+ return None
+
+ def extract_tool_calls_streaming(
+ self,
+ previous_text: str,
+ current_text: str,
+ delta_text: str,
+ previous_token_ids: Sequence[int], # pylint: disable=unused-argument
+ current_token_ids: Sequence[int], # pylint: disable=unused-argument
+ delta_token_ids: Sequence[int],
+ request: ChatCompletionRequest,
+ ) -> DeltaMessage | None:
+ """Extract tool calls from streaming model output.
+
+ Uses a buffer-until-complete-invoke strategy: tokens are buffered
+ until a complete invoke block is available, then parsed and emitted
+ in one shot.
+ """
+
+ # First chunk of a new stream — reset state from prior request.
+ if not previous_text:
+ self._reset_streaming_state()
+
+ content = self._extract_content(current_text)
+ delta_tool_calls = self._extract_delta_tool_calls(current_text, request)
+
+ if delta_tool_calls or content:
+ return DeltaMessage(content=content, tool_calls=delta_tool_calls)
+
+ # Empty delta with token ids means EOS or closing tag; return
+ # non-None so the serving framework can finalize finish_reason.
+ if not delta_text and delta_token_ids and self.prev_tool_call_arr:
+ return DeltaMessage(content="")
+
+ return None
diff --git a/reference/vllm/tool_parsers/deepseekv4_tool_parser.py b/reference/vllm/tool_parsers/deepseekv4_tool_parser.py
new file mode 100644
index 00000000..e32451cd
--- /dev/null
+++ b/reference/vllm/tool_parsers/deepseekv4_tool_parser.py
@@ -0,0 +1,31 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from vllm.entrypoints.openai.chat_completion.protocol import (
+ ChatCompletionRequest,
+)
+from vllm.tool_parsers.deepseekv32_tool_parser import DeepSeekV32ToolParser
+from vllm.tool_parsers.structural_tag_registry import (
+ get_enable_structured_outputs_in_reasoning,
+ get_model_structural_tag,
+)
+
+
+class DeepSeekV4ToolParser(DeepSeekV32ToolParser):
+ """
+ DeepSeek V4 DSML tool parser.
+
+ V4 keeps the V3.2 DSML invoke/parameter grammar, but wraps tool calls in
+ ``<｜DSML｜tool_calls>`` instead of ``<｜DSML｜function_calls>``.
+ """
+
+ tool_call_start_token: str = "<｜DSML｜tool_calls>"
+ tool_call_end_token: str = ""
+
+ def get_structural_tag(self, request: ChatCompletionRequest):
+ return get_model_structural_tag(
+ model="deepseek_v4",
+ tools=request.tools,
+ tool_choice=request.tool_choice,
+ reasoning=get_enable_structured_outputs_in_reasoning(),
+ )