| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from transformers import Qwen3_5ForConditionalGeneration |
|
|
| class FP8RowScaledEmbedding(nn.Module): |
| def __init__(self, num_embeddings, embedding_dim): |
| super().__init__() |
| self.num_embeddings = num_embeddings |
| self.embedding_dim = embedding_dim |
| self.register_buffer("weight", torch.zeros((num_embeddings, embedding_dim), dtype=torch.float8_e4m3fn)) |
| self.register_buffer("weight_scale", torch.zeros((num_embeddings, 1), dtype=torch.bfloat16)) |
|
|
| def forward(self, input_ids): |
| w_sub = self.weight[input_ids] |
| s_sub = self.weight_scale[input_ids].to(torch.float32) |
| deq = w_sub.to(torch.float32) * s_sub |
| return deq.to(torch.bfloat16) |
|
|
| class FlatQuantFP8Linear(nn.Module): |
| def __init__(self, in_features, out_features, bias=False): |
| super().__init__() |
| self.in_features = in_features |
| self.out_features = out_features |
| self.register_buffer("weight", torch.zeros((out_features, in_features), dtype=torch.float8_e4m3fn)) |
| self.register_buffer("weight_scale", torch.zeros((), dtype=torch.bfloat16)) |
| if bias: |
| self.bias = nn.Parameter(torch.zeros(out_features, dtype=torch.bfloat16)) |
| else: |
| self.register_parameter("bias", None) |
|
|
| def forward(self, x): |
| w_deq = (self.weight.to(torch.float32) * self.weight_scale.to(torch.float32)).to(x.dtype) |
| return F.linear(x, w_deq, self.bias.to(x.dtype) if self.bias is not None else None) |
|
|
| class FlatQuantW4A8Linear(nn.Module): |
| def __init__(self, in_features, out_features, bias=False, group_size=16): |
| super().__init__() |
| self.in_features = in_features |
| self.out_features = out_features |
| self.group_size = group_size |
| self.register_buffer("weight", torch.zeros((out_features, in_features // 2), dtype=torch.uint8)) |
| self.register_buffer("weight_scale", torch.zeros((out_features, in_features // group_size), dtype=torch.bfloat16)) |
| if bias: |
| self.bias = nn.Parameter(torch.zeros(out_features, dtype=torch.bfloat16)) |
| else: |
| self.register_parameter("bias", None) |
|
|
| def forward(self, x): |
| low = (self.weight & 0x0F).to(torch.int8) |
| high = (self.weight >> 4).to(torch.int8) |
| low = torch.where(low >= 8, low - 16, low) |
| high = torch.where(high >= 8, high - 16, high) |
| unpacked = torch.empty(self.out_features, self.in_features, device=x.device, dtype=x.dtype) |
| unpacked[:, 0::2] = low.to(x.dtype) |
| unpacked[:, 1::2] = high.to(x.dtype) |
| scales = self.weight_scale.to(x.dtype).repeat_interleave(self.group_size, dim=1) |
| return F.linear(x, unpacked * scales, self.bias.to(x.dtype) if self.bias is not None else None) |
|
|
| class WeMMEmbedding(Qwen3_5ForConditionalGeneration): |
| def __init__(self, config): |
| super().__init__(config) |
| for name, mod in list(self.model.named_modules()): |
| if name.endswith("embed_tokens") and isinstance(mod, nn.Embedding): |
| parent = self.model.get_submodule(name.rsplit(".", 1)[0]) if "." in name else self.model |
| child = name.rsplit(".", 1)[-1] |
| setattr(parent, child, FP8RowScaledEmbedding(mod.num_embeddings, mod.embedding_dim)) |
| elif isinstance(mod, nn.Linear): |
| parent = self.model.get_submodule(name.rsplit(".", 1)[0]) if "." in name else self.model |
| child = name.rsplit(".", 1)[-1] |
| is_attn = ("self_attn" in name) or ("linear_attn" in name) |
| is_down = "down_proj" in name |
| if is_attn or is_down: |
| setattr(parent, child, FlatQuantFP8Linear(mod.in_features, mod.out_features, bias=mod.bias is not None)) |
| else: |
| setattr(parent, child, FlatQuantW4A8Linear(mod.in_features, mod.out_features, bias=mod.bias is not None, group_size=16)) |
|
|
| def embedding(self, input_ids=None, attention_mask=None, **kwargs): |
| self.model.rope_deltas = None |
| outputs = self.model(input_ids=input_ids, attention_mask=attention_mask, **kwargs) |
| last_hidden_state = outputs.last_hidden_state |
| if attention_mask is not None: |
| eos_positions = attention_mask.sum(dim=1) - 1 |
| else: |
| eos_positions = torch.full((last_hidden_state.shape[0],), last_hidden_state.shape[1] - 1, device=last_hidden_state.device) |
| eos_positions = eos_positions.clamp(min=0) |
| batch_indices = torch.arange(last_hidden_state.size(0), device=last_hidden_state.device) |
| embeddings = last_hidden_state[batch_indices, eos_positions] |
| embeddings = F.normalize(embeddings, dim=-1) |
| return embeddings |
|
|
