import math
from dataclasses import dataclass
from typing import Tuple, Optional, Literal
from functools import lru_cache
import torch
from torch import nn
import torch.nn.functional as F
import torch.distributed as dist
import torch_npu
from scipy.linalg import hadamard
import numpy as np
from tilelang_kernels.sparse_attn_kernel import sparse_attn
from tilelang_kernels.hc_split_sinkhorn_kernel import hc_split_sinkhorn as tl_hc_split_sinkhorn
tllib = torch.library.Library("tl", "FRAGMENT")
tllib.define("hc_split_sinkhorn(Tensor mixes, Tensor hc_scale, Tensor hc_base, int hc_mult=4, \
int sinkhorn_iters = 20, float eps=1e-6) -> (Tensor, Tensor, Tensor)")
@torch.library.impl(tllib, "hc_split_sinkhorn", "Meta")
def hc_split_sinkhorn_meta(mixes, hc_scale, hc_base, hc_mult=4, sinkhorn_iters=20, eps=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)
return pre, post, comb
@torch.library.impl(tllib, "hc_split_sinkhorn", "NPU")
def hc_split_sinkhorn_npu_impl(mixes, hc_scale, hc_base, hc_mult=4, sinkhorn_iters=20, eps=1e-6):
return tl_hc_split_sinkhorn(mixes, hc_scale, hc_base, hc_mult, sinkhorn_iters, eps)
hc_split_sinkhorn = torch.ops.tl.hc_split_sinkhorn
world_size = 1
rank = 0
block_size = 128
attn_tp_size = 4
@dataclass
class ModelArgs:
"""
Data class for defining model arguments and hyperparameters.
Attributes:
max_batch_size (int): Maximum batch size.
max_seq_len (int): Maximum sequence length.
dtype (Literal["bf16", "fp8"]): Data type for computations.
scale_fmt (Optional[str]): Format for quantization scale.
vocab_size (int): Vocabulary size.
dim (int): Model dimension.
moe_inter_dim (int): Intermediate dimension for MoE layers.
n_layers (int): Number of transformer layers.
n_hash_layers (int): Number of hash MoE layers in the model.
n_heads (int): Number of attention heads.
n_routed_experts (int): Number of routed experts for MoE layers.
n_shared_experts (int): Number of shared experts for MoE layers.
n_activated_experts (int): Number of activated experts in MoE layers.
score_func (Literal["softmax", "sigmoid"]): Scoring function for MoE routing.
route_scale (float): Scaling factor for routing scores.
q_lora_rank (int): LoRA rank for query projections.
head_dim (int): Dimension for attention.
rope_head_dim (int): Dimension for rotary embeddings.
o_groups: (int): Number of groups of out projections.
o_lora_rank (int): Dimension for out projections.
window_size (int): Window size of sliding window attention.
compress_ratios (list[int]): Each layer's compress_ratio.
original_seq_len (int): Original sequence length.
rope_theta (float): Base for rotary positional encoding.
rope_factor (float): Scaling factor for extended sequence lengths.
beta_fast (int): Fast beta correction factor.
beta_slow (int): Slow beta correction factor.
index_n_heads (int): Number of index heads.
index_head_dim (int): Dimension for index head.
index_topk (int): Top-k for index head.
hc_mult (int): HC hidden size multiplier.
sinkhorn_iters (int): Number of sinkhorn iterations for HC.
"""
max_batch_size: int = 4
max_seq_len: int = 4096
dtype: Literal["bf16", "fp8"] = "bf16"
scale_fmt: Optional[str] = "ue8m0"
vocab_size: int = 129280
dim: int = 4096
moe_inter_dim: int = 4096
n_layers: int = 7
n_hash_layers: int = 0
n_heads: int = 64
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.
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] = (1, 1, 4, 128, 4, 128, 4)
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_n_heads: int = 32
index_head_dim: int = 128
index_topk: int = 512
hc_mult: int = 4
hc_sinkhorn_iters: int = 20
hc_eps: float = 1e-6
class ParallelEmbedding(nn.Module):
"""
Embedding layer with parallelism support across distributed processes.
Args:
vocab_size (int): Vocabulary size.
dim (int): Embedding dimension.
"""
def __init__(self, vocab_size: int, dim: int):
super().__init__()
self.vocab_size = vocab_size
self.dim = dim
if vocab_size % world_size != 0:
raise ValueError(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:
"""
Forward pass for parallel embedding layer.
Args:
x (torch.Tensor): Input tensor containing token indices.
Returns:
torch.Tensor: Embedded representations.
Raises:
ValueError: If `world_size` is not defined.
"""
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,
scale_fmt: Optional[str] = None) -> torch.Tensor:
"""
Applies a linear transformation to the incoming data: y = xA^T + b.
This function supports specialized implementations based on quantization
and tensor formats.
Args:
x (torch.Tensor): The input tensor.
weight (torch.Tensor): The weight tensor.
bias (Optional[torch.Tensor]): The bias tensor to be added. Default is None.
scale_fmt (Optional[str]): The format of scaling factors.
Returns:
torch.Tensor: The result of the linear transformation, which may involve
quantization-aware computations depending on the input parameters.
Notes:
- If `weight` is not quantized, a normal version is used for computation.
- For other cases, the function applies quantization to `x` and uses `fp8_gemm` for computation.
"""
if bias is not None:
raise ValueError("bias must be None")
x = x.to(torch.bfloat16)
weight = weight.to(torch.bfloat16)
return F.linear(x, weight)
class Linear(nn.Module):
"""
Custom linear layer with support for quantized weights and optional bias.
Args:
in_features (int): Number of input features.
out_features (int): Number of output features.
bias (bool): Whether to include a bias term. Defaults to False.
dtype (optional): Data type for the layer. Defaults to `torch.bfloat16`.
"""
dtype = torch.bfloat16
scale_fmt: Optional[str] = None
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
self.weight = nn.Parameter(torch.empty(out_features, in_features, dtype=dtype or torch.bfloat16))
if self.weight.element_size() == 1:
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.bfloat16))
else:
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:
"""
Forward pass for the custom linear layer.
Args:
x (torch.Tensor): Input tensor.
Returns:
torch.Tensor: Transformed tensor after linear computation.
"""
self.weight.data = self.weight.data.to(torch.bfloat16)
return linear(x, self.weight, self.bias, self.scale_fmt)
class ColumnParallelLinear(Linear):
"""
Linear layer with column parallelism, splitting output features across distributed processes.
Args:
in_features (int): Number of input features.
out_features (int): Total number of output features.
bias (bool): Whether to include a bias term. Defaults to False.
dtype (optional): Data type for the layer. Defaults to `torch.bfloat16`.
"""
def __init__(
self, in_features: int, out_features: int, bias: bool = False, dtype=None
):
if out_features % world_size != 0:
raise ValueError(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:
"""
Forward pass for column parallel linear layer.
Args:
x (torch.Tensor): Input tensor.
Returns:
torch.Tensor: Transformed tensor with column-parallel computation.
"""
y = linear(x, self.weight, self.bias, self.scale_fmt)
return y
class RowParallelLinear(Linear):
"""
Linear layer with row parallelism, splitting input features across distributed processes.
Args:
in_features (int): Total number of input features.
out_features (int): Number of output features.
bias (bool): Whether to include a bias term. Defaults to False.
dtype (optional): Data type for the layer. Defaults to `torch.bfloat16`.
"""
def __init__(
self, in_features: int, out_features: int, bias: bool = False, dtype=None
):
if in_features % world_size != 0:
raise ValueError(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, attn_tp_dim: int) -> torch.Tensor:
"""
Forward pass for row parallel linear layer.
Args:
x (torch.Tensor): Input tensor.
Returns:
torch.Tensor: Transformed tensor with row-parallel computation.
"""
y = linear(x, self.weight, None, self.scale_fmt)
if world_size > 1:
y = y.float()
dist.all_reduce(y)
y /= attn_tp_dim
if self.bias is not None:
y += self.bias
return y.type_as(x)
class RMSNorm(nn.Module):
"""
Root Mean Square Layer Normalization (RMSNorm).
Args:
dim (int): Dimension of the input tensor.
eps (float): Epsilon value for numerical stability. Defaults to 1e-6.
"""
def __init__(self, dim: int, eps: float = 1e-6):
super().__init__()
self.dim = dim
self.eps = eps
self.weight = nn.Parameter(torch.ones(dim, dtype=torch.bfloat16))
@torch.compile
def forward(self, x: torch.Tensor):
"""
Forward pass for RMSNorm.
Args:
x (torch.Tensor): Input tensor.
Returns:
torch.Tensor: Normalized tensor with the same shape as input.
"""
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 frequency-based complex exponential values for rotary positional embeddings.
Args:
args (ModelArgs): Model arguments containing positional embedding parameters.
Returns:
torch.Tensor: Precomputed complex exponential values for positional embeddings.
"""
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 to the input tensor.
Args:
x (torch.Tensor): Input tensor with positional embeddings to be applied.
freqs_cis (torch.Tensor): Precomputed complex exponential values for positional embeddings.
Returns:
torch.Tensor: Tensor with rotary embeddings applied.
"""
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 hadamard_transform_fix(x):
dtype = x.dtype
device = x.device
hidden_size = x.size(-1)
H_m = torch.tensor(hadamard(hidden_size, dtype=np.float32) / (hidden_size ** 0.5)).to(device)
x = (x.to(torch.float32) @ H_m).to(dtype)
return x
def rotate_activation(x: torch.Tensor) -> torch.Tensor:
assert x.dtype == torch.bfloat16
return hadamard_transform_fix(x)
@lru_cache(1)
def get_window_topk_idxs(window_size: int, bsz: int, seqlen: int, start_pos: int):
def _get_window_topk_idxs():
if start_pos >= window_size - 1:
return torch.arange(window_size)
elif start_pos > 0:
return 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
return _get_window_topk_idxs().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):
def _get_compress_topk_idxs():
if start_pos > 0:
return torch.arange(0, start_pos // 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
return _get_compress_topk_idxs().unsqueeze(0).expand(bsz, -1, -1)
class Compressor(nn.Module):
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))
self.wkv = Linear(self.dim, coff * self.head_dim, dtype=torch.bfloat16)
self.wgate = Linear(self.dim, coff * self.head_dim, dtype=torch.bfloat16)
self.norm = RMSNorm(self.head_dim, args.norm_eps)
self.kv_cache = None
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)
def overlap_transform(self, tensor: torch.Tensor, value=0):
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, freqs_cis: torch.Tensor):
assert self.kv_cache is not None
bsz, seqlen, _ = x.size()
ratio, overlap, d = self.compress_ratio, self.overlap, self.head_dim
dtype = x.dtype
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
freqs_cis = freqs_cis[:cutoff:ratio]
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
dtype = torch.bfloat16
kv = self.norm(kv.to(dtype))
apply_rotary_emb(kv[..., -self.rope_head_dim:], freqs_cis)
if self.rotate:
kv = rotate_activation(kv)
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):
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.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.attn_tp_dim = attn_tp_size
self.wq_b = ColumnParallelLinear(self.q_lora_rank, self.n_heads * self.head_dim * self.attn_tp_dim)
self.weights_proj = ColumnParallelLinear(self.dim, self.n_heads * self.attn_tp_dim, 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)
def forward(self, x: torch.Tensor, qr: torch.Tensor, start_pos: int, freqs_cis: torch.Tensor, offset: int):
bsz, seqlen, _ = x.size()
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
q = self.wq_b(qr).unflatten(-1, (self.n_local_heads * self.attn_tp_dim, self.head_dim))
apply_rotary_emb(q[..., -rd:], freqs_cis)
q = rotate_activation(q)
self.compressor(x, start_pos, freqs_cis)
weights = self.weights_proj(x) * (self.softmax_scale * self.n_heads ** -0.5)
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)
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-Query Attention (MQA) Layer."""
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.attn_tp_dim = attn_tp_size
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 // self.attn_tp_dim
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 * self.attn_tp_dim, 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.attn_tp_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 * self.attn_tp_dim, dtype=torch.bfloat16)
self.wo_b = RowParallelLinear(self.n_groups * args.o_lora_rank * self.attn_tp_dim, self.dim)
self.softmax_scale = self.head_dim ** -0.5
self.scale_fmt = args.scale_fmt
if self.compress_ratio > 1:
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
self.register_buffer("kv_cache", torch.zeros(args.max_batch_size, args.window_size + args.max_seq_len // self.compress_ratio, self.head_dim), persistent=False)
freqs_cis = precompute_freqs_cis(self.rope_head_dim, args.max_seq_len, args.original_seq_len,
args.compress_rope_theta if self.compress_ratio > 1 else args.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):
"""
Forward pass for the Multi-Head Latent Attention (MLA) Layer.
Args:
x (torch.Tensor): Input tensor of shape (batch_size, seq_len, dim).
start_pos (int): Starting position in the sequence for caching.
Returns:
torch.Tensor: Output tensor with the same shape as the input.
"""
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 > 1 and self.compressor.kv_cache is None:
self.compressor.kv_cache = self.kv_cache[:, win:]
qr = q = self.q_norm(self.wq_a(x))
q = self.wq_b(q)
q = q.unflatten(-1, (self.n_local_heads * self.attn_tp_dim, self.head_dim))
q *= torch.rsqrt(q.square().mean(-1, keepdim=True) + self.eps)
apply_rotary_emb(q[..., -rd:], freqs_cis)
kv = self.wkv(x)
kv = self.kv_norm(kv)
apply_rotary_emb(kv[..., -rd:], freqs_cis)
topk_idxs = get_window_topk_idxs(win, bsz, seqlen, start_pos)
if self.compress_ratio > 1:
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, freqs_cis, 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()
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 > 1:
if (kv_compress := self.compressor(x, start_pos, freqs_cis)) is not None:
kv = torch.cat([kv, kv_compress], dim=1)
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 > 1:
self.compressor(x, start_pos, freqs_cis)
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.view(bsz, seqlen, self.n_local_groups, -1)
wo_a = self.wo_a.weight.view(self.n_local_groups, self.o_lora_rank, -1)
o = torch.einsum("bsgd,grd->bsgr", o, wo_a)
x = self.wo_b(o.flatten(2), self.attn_tp_dim)
return x
class Gate(nn.Module):
"""
Gating mechanism for routing inputs in a mixture-of-experts (MoE) model.
Attributes:
dim (int): Dimensionality of input features.
topk (int): Number of top experts activated for each input.
n_groups (int): Number of groups for routing.
topk_groups (int): Number of groups to route inputs to.
score_func (str): Scoring function ('softmax' or 'sigmoid').
route_scale (float): Scaling factor for routing weights.
weight (torch.nn.Parameter): Learnable weights for the gate.
bias (Optional[torch.nn.Parameter]): Optional bias term for the gate.
"""
def __init__(self, layer_id: int, args: ModelArgs):
"""
Initializes the Gate module.
Args:
args (ModelArgs): Model arguments containing gating parameters.
"""
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.int64), 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]:
"""
Forward pass for the gating mechanism.
Args:
x (torch.Tensor): Input tensor.
input_ids (torch.Tensor): Token IDs tensor.
Returns:
Tuple[torch.Tensor, torch.Tensor]: Routing weights and selected expert indices.
"""
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
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):
"""
Expert layer for Mixture-of-Experts (MoE) models.
Attributes:
w1 (nn.Module): Linear layer for input-to-hidden transformation.
w2 (nn.Module): Linear layer for hidden-to-output transformation.
w3 (nn.Module): Additional linear layer for feature transformation.
"""
def __init__(self, dim: int, inter_dim: int):
"""
Initializes the Expert layer.
Args:
dim (int): Input and output dimensionality.
inter_dim (int): Hidden layer dimensionality.
"""
super().__init__()
self.w1 = Linear(dim, inter_dim)
self.w2 = Linear(inter_dim, dim)
self.w3 = Linear(dim, inter_dim)
@torch.compile
def forward(self, x: torch.Tensor, weights: Optional[torch.Tensor] = None) -> torch.Tensor:
"""
Forward pass for the Expert layer.
Args:
x (torch.Tensor): Input tensor.
Returns:
torch.Tensor: Output tensor after expert computation.
"""
dtype = x.dtype
x = F.silu(self.w1(x).float()) * self.w3(x).float()
if weights is not None:
x = weights * x
return self.w2(x.to(dtype))
class MoE(nn.Module):
"""
Mixture-of-Experts (MoE) module.
Attributes:
dim (int): Dimensionality of input features.
n_routed_experts (int): Total number of experts in the model.
n_local_experts (int): Number of experts handled locally in distributed systems.
n_activated_experts (int): Number of experts activated for each input.
gate (nn.Module): Gating mechanism to route inputs to experts.
experts (nn.ModuleList): List of expert modules.
shared_experts (nn.Module): Shared experts applied to all inputs.
"""
def __init__(self, layer_id: int, args: ModelArgs):
"""
Initializes the MoE module.
Args:
args (ModelArgs): Model arguments containing MoE parameters.
"""
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)
self.experts = nn.ModuleList([Expert(args.dim, args.moe_inter_dim) 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)
def run_gate(self, x, input_ids: Optional[torch.Tensor] = None):
return self.gate(x, input_ids)
def forward(self, x: torch.Tensor, input_ids: torch.Tensor) -> torch.Tensor:
"""
Forward pass for the MoE module.
Args:
x (torch.Tensor): Input tensor.
Returns:
torch.Tensor: Output tensor after expert routing and computation.
"""
shape = x.size()
x = x.view(-1, self.dim)
weights, indices = self.run_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 combining attention and feed-forward layers.
Attributes:
attn (nn.Module): Attention layer (MLA).
ffn (nn.Module): Feed-forward network (MLP or MoE).
attn_norm (nn.Module): Layer normalization for attention.
ffn_norm (nn.Module): Layer normalization for feed-forward network.
"""
def __init__(self, layer_id: int, args: ModelArgs):
"""
Initializes the Transformer block.
Args:
layer_id (int): Layer index in the transformer.
args (ModelArgs): Model arguments containing block parameters.
"""
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
origin_dtype = torch.get_default_dtype()
torch.set_default_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))
torch.set_default_dtype(origin_dtype)
@torch.compile
def hc_pre(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, 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
@torch.compile
def hc_post(self, x: torch.Tensor, residual: torch.Tensor, post: torch.Tensor, comb: torch.Tensor):
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:
"""
Forward pass for the Transformer block.
Args:
x (torch.Tensor): Input tensor.
start_pos (int): Starting position in the sequence.
Returns:
torch.Tensor: Output tensor after block computation.
"""
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 Transformer(nn.Module):
"""
Transformer model with positional embeddings, multiple layers, and output projection.
Attributes:
max_seq_len (int): Maximum sequence length for the transformer.
embed (nn.Module): Embedding layer for input tokens.
layers (torch.nn.ModuleList): List of transformer blocks.
norm (nn.Module): Layer normalization applied after all blocks.
head (nn.Module): Output projection layer mapping to vocabulary size.
"""
def __init__(self, args: ModelArgs):
"""
Initializes the Transformer model.
Args:
args (ModelArgs): Model arguments containing transformer parameters.
"""
global world_size, rank
world_size = dist.get_world_size() if dist.is_initialized() else 1
rank = dist.get_rank() if dist.is_initialized() else 0
Linear.dtype = torch.float8_e4m3fn if args.dtype == "fp8" else torch.bfloat16
Linear.scale_fmt = args.scale_fmt
super().__init__()
self.max_seq_len = args.max_seq_len
self.norm_eps = args.norm_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 = ColumnParallelLinear(args.dim, args.vocab_size, dtype=torch.bfloat16)
self.hc_eps = args.hc_eps
self.hc_mult = hc_mult = args.hc_mult
hc_dim = hc_mult * args.dim
origin_dtype = torch.get_default_dtype()
torch.set_default_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.set_default_dtype(origin_dtype)
@torch.compile
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)
@torch.inference_mode()
def forward(self, input_ids: torch.Tensor, start_pos: int = 0):
"""
Forward pass for the Transformer model.
Args:
input_ids (torch.Tensor): Input tensor of token IDs with shape (batch_size, seq_len).
start_pos (int, optional): Starting position in the sequence for rotary embeddings. Defaults to 0.
Returns:
torch.Tensor: Logits tensor of shape (batch_size, vocab_size).
"""
h = self.embed(input_ids)
h = h.unsqueeze(2).repeat(1, 1, self.hc_mult, 1)
for layer in self.layers:
h = layer(h, start_pos, input_ids)
h = self.hc_head(h, self.hc_head_fn, self.hc_head_scale, self.hc_head_base)
h = self.norm(h)
logits = self.head(h[:, -1].float())
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
