Instructions to use JaydenTeoh/soft-mtp-167m-model with libraries, inference providers, notebooks, and local apps. Follow these links to get started.

Libraries

How to use JaydenTeoh/soft-mtp-167m-model with Transformers:

# Use a pipeline as a high-level helper
from transformers import pipeline

pipe = pipeline("text-generation", model="JaydenTeoh/soft-mtp-167m-model", trust_remote_code=True)

# pip install -U transformers accelerate
# Load model directly
from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained("JaydenTeoh/soft-mtp-167m-model", trust_remote_code=True, device_map="auto")

Notebooks
Google Colab
Kaggle
Local Apps Settings

vLLM

How to use JaydenTeoh/soft-mtp-167m-model with vLLM:

Install from pip and serve model

# Install vLLM from pip:
pip install vllm
# Start the vLLM server:
vllm serve "JaydenTeoh/soft-mtp-167m-model"
# Call the server using curl (OpenAI-compatible API):
curl -X POST "http://localhost:8000/v1/completions" \
	-H "Content-Type: application/json" \
	--data '{
		"model": "JaydenTeoh/soft-mtp-167m-model",
		"prompt": "Once upon a time,",
		"max_tokens": 512,
		"temperature": 0.5
	}'

Use Docker

docker model run hf.co/JaydenTeoh/soft-mtp-167m-model

SGLang

How to use JaydenTeoh/soft-mtp-167m-model with SGLang:

Install from pip and serve model

# Install SGLang from pip:
pip install sglang
# Start the SGLang server:
python3 -m sglang.launch_server \
    --model-path "JaydenTeoh/soft-mtp-167m-model" \
    --host 0.0.0.0 \
    --port 30000
# Call the server using curl (OpenAI-compatible API):
curl -X POST "http://localhost:30000/v1/completions" \
	-H "Content-Type: application/json" \
	--data '{
		"model": "JaydenTeoh/soft-mtp-167m-model",
		"prompt": "Once upon a time,",
		"max_tokens": 512,
		"temperature": 0.5
	}'

Use Docker images

docker run --gpus all \
    --shm-size 32g \
    -p 30000:30000 \
    -v ~/.cache/huggingface:/root/.cache/huggingface \
    --env "HF_TOKEN=<secret>" \
    --ipc=host \
    lmsysorg/sglang:latest \
    python3 -m sglang.launch_server \
        --model-path "JaydenTeoh/soft-mtp-167m-model" \
        --host 0.0.0.0 \
        --port 30000
# Call the server using curl (OpenAI-compatible API):
curl -X POST "http://localhost:30000/v1/completions" \
	-H "Content-Type: application/json" \
	--data '{
		"model": "JaydenTeoh/soft-mtp-167m-model",
		"prompt": "Once upon a time,",
		"max_tokens": 512,
		"temperature": 0.5
	}'

Docker Model Runner
How to use JaydenTeoh/soft-mtp-167m-model with Docker Model Runner:
```
docker model run hf.co/JaydenTeoh/soft-mtp-167m-model
```

soft-mtp-167m-model / blocks.py

JaydenTeoh

Sync modelling files with the released source (cascade -> sequential)

b06441f verified 4 days ago

Raw History Blame Contribute Delete

9.27 kB

"""Pure-PyTorch Llama-style transformer block (no FLA/triton; runs on CPU + GPU).
RMSNorm (fp32-upcast, matches Llama) + RoPE (rotate_half convention) + GQA attention via
`F.scaled_dot_product_attention` (uses FlashAttention on CUDA) + SwiGLU. Verified equal to
HF `LlamaForCausalLM` to 1e-4. One `Block` is reused
for trunk layers and MTP heads, and `torch.compile`s cleanly (regional compile).
"""

from contextlib import nullcontext

import torch
import torch.nn as nn
import torch.nn.functional as F


class RMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-5):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(dim))
        self.eps = eps

    def forward(self, x):
        in_dtype = x.dtype  # upcast LOW precision to fp32 (matches Llama; stable under bf16),
        h = x.float() if in_dtype in (torch.float16, torch.bfloat16) else x  # but never downcast fp64
        h = h * torch.rsqrt(h.pow(2).mean(-1, keepdim=True) + self.eps)
        return (self.weight * h).to(in_dtype)


def build_rope_cache(seq_len, head_dim, theta, device, dtype=torch.float32):
    """cos/sin table [T, head_dim]. Trig computed in fp32, cast to `dtype` last."""
    dev = device if isinstance(device, torch.device) else torch.device(device)
    # Disable any outside autocast context if any, to really force fp32 (autocast has no policy for
    # `meta`, which the param-count probes build on, so the context is guarded rather than assumed).
    ctx = (torch.autocast(device_type=dev.type, enabled=False)
           if dev.type in ("cpu", "cuda", "xpu") else nullcontext())
    with ctx:
        inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim))
        freqs = torch.outer(torch.arange(seq_len, device=device).float(), inv_freq)
        emb = torch.cat((freqs, freqs), dim=-1)  # [T, head_dim]
        return emb.cos().to(dtype), emb.sin().to(dtype)


def _rotate_half(x):
    x1, x2 = x.chunk(2, dim=-1)
    return torch.cat((-x2, x1), dim=-1)


def apply_rope(q, k, cos, sin):
    """cos/sin are [T, D] (batch-shared positions) or [B, T, D] (per-row `position_ids`, the HF
    convention — `generate()` derives them as `attention_mask.cumsum(-1) - 1`, so left-padded rows
    get 0..n-1 on their REAL tokens). Both broadcast to q/k's [B, H, T, D]."""
    if cos.dim() == 2:
        cos, sin = cos[None, None], sin[None, None]  # [1, 1, T, D]
    else:
        cos, sin = cos[:, None], sin[:, None]        # [B, 1, T, D]
    return q * cos + _rotate_half(q) * sin, k * cos + _rotate_half(k) * sin


class Attention(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.n_heads, self.n_kv = cfg.num_heads, cfg.num_kv_heads
        self.hd = cfg.hidden_size // cfg.num_heads
        self.q_proj = nn.Linear(cfg.hidden_size, self.n_heads * self.hd, bias=False)
        self.k_proj = nn.Linear(cfg.hidden_size, self.n_kv * self.hd, bias=False)
        self.v_proj = nn.Linear(cfg.hidden_size, self.n_kv * self.hd, bias=False)
        self.o_proj = nn.Linear(self.n_heads * self.hd, cfg.hidden_size, bias=False)

    def forward(self, x, cos, sin, past_key_values=None, layer_idx=0, cache_position=None,
                cache_readonly=False, attention_mask=None):
        """`attention_mask` is the pre-built ADDITIVE 4D mask [B, 1, T, k_len] (causal + padding), or
        None for the training/cache-free path which uses the compile-clean `is_causal` kernel."""
        B, T, _ = x.shape
        q = self.q_proj(x).view(B, T, self.n_heads, self.hd).transpose(1, 2)
        k = self.k_proj(x).view(B, T, self.n_kv, self.hd).transpose(1, 2)
        v = self.v_proj(x).view(B, T, self.n_kv, self.hd).transpose(1, 2)
        q, k = apply_rope(q, k, cos, sin)
        if past_key_values is not None:
            if cache_readonly:
                # READ-ONLY: attend against the cache without writing to it. transformers'
                # `update()` is unconditionally append-in-place, so this is hand-rolled. The sequential
                # drafter needs it — its last positions consume unconfirmed DRAFT tokens, which must
                # attend but must never enter the cache, or a rejected draft poisons later drafts.
                layer = past_key_values.layers[layer_idx]
                pk, pv = getattr(layer, "keys", None), getattr(layer, "values", None)
                if pk is not None and pk.numel():
                    k = torch.cat([pk, k], dim=-2)
                    v = torch.cat([pv, v], dim=-2)
            else:                                        # inference KV cache (eager): append + read full K/V
                k, v = past_key_values.update(k, v, layer_idx)
        if self.n_kv != self.n_heads:  # GQA: expand kv heads
            rep = self.n_heads // self.n_kv
            k, v = k.repeat_interleave(rep, dim=1), v.repeat_interleave(rep, dim=1)
        if attention_mask is None and past_key_values is not None and cache_position is not None:
            # SAFETY NET: a cached call whose caller forgot the pre-built mask would otherwise run
            # BIDIRECTIONAL attention with no error — exactly the f0816ed regression that left
            # `_dsmtp_layer` non-causal on multi-position chunks while every decode gate kept
            # passing (the chunk's LAST position is causal by accident, and that is the only one
            # sampled from). Model-level callers still pass the shared mask, so this never fires on
            # the trunk/sequential paths and their build-once economics are unchanged.
            attention_mask = build_causal_mask(cache_position, k.size(2), None, x.device, q.dtype)
        if past_key_values is None and attention_mask is None:
            # training / cache-free without padding: clean compiled is_causal kernel. This mirrors
            # HF's contract exactly — `create_causal_mask` returns None when the sdpa backend can
            # use `is_causal` (no padding, no cache), and the attention interface falls back to it.
            out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        else:                                            # cached, and/or padded cache-free: additive mask
            # `attention_mask` here is the FINAL additive mask, already combining causal + padding,
            # built ONCE per forward by `build_causal_mask` and shared by every layer — the same split
            # HF uses (DeepseekV3Model calls `create_causal_mask` and DeepseekV3Attention only passes
            # the result to the attention interface). Building it here instead would reallocate the
            # same [B,1,T,k_len] tensor once per layer.
            out = F.scaled_dot_product_attention(q, k, v, attn_mask=attention_mask)
        return self.o_proj(out.transpose(1, 2).reshape(B, T, -1))


def build_causal_mask(cache_position, k_len, attention_mask, device, dtype):
    """The additive [B, 1, T, k_len] mask for the CACHED path: causal AND padding, built once per
    forward and shared by every layer (HF's `create_causal_mask` plays this role for DeepSeek-V3 /
    Llama; their attention modules never construct a mask themselves).
    causal : key position <= query's absolute `cache_position`.
    padding: `attention_mask` is the HF convention [B, k_len], 1 = real token, 0 = pad. Without this
             term a query attends PAD keys, which silently degrades ANY batched generation with mixed
             prompt lengths — lm-eval left-pads its generate_until batches.
    Returns a 2D [T, k_len] mask when there is no padding to apply, so the common single-sequence case
    does not pay for a [B,1,T,k_len] broadcast.
    """
    T = cache_position.numel()
    kpos = torch.arange(k_len, device=device)
    allow = kpos[None, :] <= cache_position[:, None]
    mask = torch.zeros(T, k_len, device=device, dtype=dtype).masked_fill(~allow, float("-inf"))
    if attention_mask is None:
        return mask
    am = attention_mask[:, :k_len]
    mask = mask[None, None].expand(am.size(0), 1, T, k_len).clone()
    return mask.masked_fill_(am[:, None, None, :] == 0, float("-inf"))


class SwiGLU(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.gate_proj = nn.Linear(cfg.hidden_size, cfg.intermediate_size, bias=False)
        self.up_proj = nn.Linear(cfg.hidden_size, cfg.intermediate_size, bias=False)
        self.down_proj = nn.Linear(cfg.intermediate_size, cfg.hidden_size, bias=False)

    def forward(self, x):
        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))


class Block(nn.Module):
    """Pre-norm transformer block; reused for trunk layers and MTP heads."""

    def __init__(self, cfg):
        super().__init__()
        self.attn_norm = RMSNorm(cfg.hidden_size, cfg.norm_eps)
        self.attn = Attention(cfg)
        self.mlp_norm = RMSNorm(cfg.hidden_size, cfg.norm_eps)
        self.mlp = SwiGLU(cfg)

    def forward(self, x, cos, sin, past_key_values=None, layer_idx=0, cache_position=None,
                cache_readonly=False, attention_mask=None):
        x = x + self.attn(self.attn_norm(x), cos, sin, past_key_values, layer_idx, cache_position,
                          cache_readonly, attention_mask)
        x = x + self.mlp(self.mlp_norm(x))
        return x

	"""Pure-PyTorch Llama-style transformer block (no FLA/triton; runs on CPU + GPU).

	RMSNorm (fp32-upcast, matches Llama) + RoPE (rotate_half convention) + GQA attention via
	`F.scaled_dot_product_attention` (uses FlashAttention on CUDA) + SwiGLU. Verified equal to
	HF `LlamaForCausalLM` to 1e-4. One `Block` is reused
	for trunk layers and MTP heads, and `torch.compile`s cleanly (regional compile).
	"""

	from contextlib import nullcontext

	import torch
	import torch.nn as nn
	import torch.nn.functional as F


	class RMSNorm(nn.Module):
	def __init__(self, dim, eps=1e-5):
	super().__init__()
	self.weight = nn.Parameter(torch.ones(dim))
	self.eps = eps

	def forward(self, x):
	in_dtype = x.dtype # upcast LOW precision to fp32 (matches Llama; stable under bf16),
	h = x.float() if in_dtype in (torch.float16, torch.bfloat16) else x # but never downcast fp64
	h = h * torch.rsqrt(h.pow(2).mean(-1, keepdim=True) + self.eps)
	return (self.weight * h).to(in_dtype)


	def build_rope_cache(seq_len, head_dim, theta, device, dtype=torch.float32):
	"""cos/sin table [T, head_dim]. Trig computed in fp32, cast to `dtype` last."""
	dev = device if isinstance(device, torch.device) else torch.device(device)
	# Disable any outside autocast context if any, to really force fp32 (autocast has no policy for
	# `meta`, which the param-count probes build on, so the context is guarded rather than assumed).
	ctx = (torch.autocast(device_type=dev.type, enabled=False)
	if dev.type in ("cpu", "cuda", "xpu") else nullcontext())
	with ctx:
	inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim))
	freqs = torch.outer(torch.arange(seq_len, device=device).float(), inv_freq)
	emb = torch.cat((freqs, freqs), dim=-1) # [T, head_dim]
	return emb.cos().to(dtype), emb.sin().to(dtype)


	def _rotate_half(x):
	x1, x2 = x.chunk(2, dim=-1)
	return torch.cat((-x2, x1), dim=-1)


	def apply_rope(q, k, cos, sin):
	"""cos/sin are [T, D] (batch-shared positions) or [B, T, D] (per-row `position_ids`, the HF
	convention — `generate()` derives them as `attention_mask.cumsum(-1) - 1`, so left-padded rows
	get 0..n-1 on their REAL tokens). Both broadcast to q/k's [B, H, T, D]."""
	if cos.dim() == 2:
	cos, sin = cos[None, None], sin[None, None] # [1, 1, T, D]
	else:
	cos, sin = cos[:, None], sin[:, None] # [B, 1, T, D]
	return q * cos + _rotate_half(q) * sin, k * cos + _rotate_half(k) * sin


	class Attention(nn.Module):
	def __init__(self, cfg):
	super().__init__()
	self.n_heads, self.n_kv = cfg.num_heads, cfg.num_kv_heads
	self.hd = cfg.hidden_size // cfg.num_heads
	self.q_proj = nn.Linear(cfg.hidden_size, self.n_heads * self.hd, bias=False)
	self.k_proj = nn.Linear(cfg.hidden_size, self.n_kv * self.hd, bias=False)
	self.v_proj = nn.Linear(cfg.hidden_size, self.n_kv * self.hd, bias=False)
	self.o_proj = nn.Linear(self.n_heads * self.hd, cfg.hidden_size, bias=False)

	def forward(self, x, cos, sin, past_key_values=None, layer_idx=0, cache_position=None,
	cache_readonly=False, attention_mask=None):
	"""`attention_mask` is the pre-built ADDITIVE 4D mask [B, 1, T, k_len] (causal + padding), or
	None for the training/cache-free path which uses the compile-clean `is_causal` kernel."""
	B, T, _ = x.shape
	q = self.q_proj(x).view(B, T, self.n_heads, self.hd).transpose(1, 2)
	k = self.k_proj(x).view(B, T, self.n_kv, self.hd).transpose(1, 2)
	v = self.v_proj(x).view(B, T, self.n_kv, self.hd).transpose(1, 2)
	q, k = apply_rope(q, k, cos, sin)
	if past_key_values is not None:
	if cache_readonly:
	# READ-ONLY: attend against the cache without writing to it. transformers'
	# `update()` is unconditionally append-in-place, so this is hand-rolled. The sequential
	# drafter needs it — its last positions consume unconfirmed DRAFT tokens, which must
	# attend but must never enter the cache, or a rejected draft poisons later drafts.
	layer = past_key_values.layers[layer_idx]
	pk, pv = getattr(layer, "keys", None), getattr(layer, "values", None)
	if pk is not None and pk.numel():
	k = torch.cat([pk, k], dim=-2)
	v = torch.cat([pv, v], dim=-2)
	else: # inference KV cache (eager): append + read full K/V
	k, v = past_key_values.update(k, v, layer_idx)
	if self.n_kv != self.n_heads: # GQA: expand kv heads
	rep = self.n_heads // self.n_kv
	k, v = k.repeat_interleave(rep, dim=1), v.repeat_interleave(rep, dim=1)
	if attention_mask is None and past_key_values is not None and cache_position is not None:
	# SAFETY NET: a cached call whose caller forgot the pre-built mask would otherwise run
	# BIDIRECTIONAL attention with no error — exactly the f0816ed regression that left
	# `_dsmtp_layer` non-causal on multi-position chunks while every decode gate kept
	# passing (the chunk's LAST position is causal by accident, and that is the only one
	# sampled from). Model-level callers still pass the shared mask, so this never fires on
	# the trunk/sequential paths and their build-once economics are unchanged.
	attention_mask = build_causal_mask(cache_position, k.size(2), None, x.device, q.dtype)
	if past_key_values is None and attention_mask is None:
	# training / cache-free without padding: clean compiled is_causal kernel. This mirrors
	# HF's contract exactly — `create_causal_mask` returns None when the sdpa backend can
	# use `is_causal` (no padding, no cache), and the attention interface falls back to it.
	out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
	else: # cached, and/or padded cache-free: additive mask
	# `attention_mask` here is the FINAL additive mask, already combining causal + padding,
	# built ONCE per forward by `build_causal_mask` and shared by every layer — the same split
	# HF uses (DeepseekV3Model calls `create_causal_mask` and DeepseekV3Attention only passes
	# the result to the attention interface). Building it here instead would reallocate the
	# same [B,1,T,k_len] tensor once per layer.
	out = F.scaled_dot_product_attention(q, k, v, attn_mask=attention_mask)
	return self.o_proj(out.transpose(1, 2).reshape(B, T, -1))


	def build_causal_mask(cache_position, k_len, attention_mask, device, dtype):
	"""The additive [B, 1, T, k_len] mask for the CACHED path: causal AND padding, built once per
	forward and shared by every layer (HF's `create_causal_mask` plays this role for DeepSeek-V3 /
	Llama; their attention modules never construct a mask themselves).

	causal : key position <= query's absolute `cache_position`.
	padding: `attention_mask` is the HF convention [B, k_len], 1 = real token, 0 = pad. Without this
	term a query attends PAD keys, which silently degrades ANY batched generation with mixed
	prompt lengths — lm-eval left-pads its generate_until batches.

	Returns a 2D [T, k_len] mask when there is no padding to apply, so the common single-sequence case
	does not pay for a [B,1,T,k_len] broadcast.
	"""
	T = cache_position.numel()
	kpos = torch.arange(k_len, device=device)
	allow = kpos[None, :] <= cache_position[:, None]
	mask = torch.zeros(T, k_len, device=device, dtype=dtype).masked_fill(~allow, float("-inf"))
	if attention_mask is None:
	return mask
	am = attention_mask[:, :k_len]
	mask = mask[None, None].expand(am.size(0), 1, T, k_len).clone()
	return mask.masked_fill_(am[:, None, None, :] == 0, float("-inf"))


	class SwiGLU(nn.Module):
	def __init__(self, cfg):
	super().__init__()
	self.gate_proj = nn.Linear(cfg.hidden_size, cfg.intermediate_size, bias=False)
	self.up_proj = nn.Linear(cfg.hidden_size, cfg.intermediate_size, bias=False)
	self.down_proj = nn.Linear(cfg.intermediate_size, cfg.hidden_size, bias=False)

	def forward(self, x):
	return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))


	class Block(nn.Module):
	"""Pre-norm transformer block; reused for trunk layers and MTP heads."""

	def __init__(self, cfg):
	super().__init__()
	self.attn_norm = RMSNorm(cfg.hidden_size, cfg.norm_eps)
	self.attn = Attention(cfg)
	self.mlp_norm = RMSNorm(cfg.hidden_size, cfg.norm_eps)
	self.mlp = SwiGLU(cfg)

	def forward(self, x, cos, sin, past_key_values=None, layer_idx=0, cache_position=None,
	cache_readonly=False, attention_mask=None):
	x = x + self.attn(self.attn_norm(x), cos, sin, past_key_values, layer_idx, cache_position,
	cache_readonly, attention_mask)
	x = x + self.mlp(self.mlp_norm(x))
	return x