Skip to content

Commit f01113a

Browse files
felipedilhofelipedilho
committed
Implementation of Phi-4-Mini-Flash-Reasoning inference, conversion from HF weights format into .gguf format and quantization of the same. The model, which is of a hybrid architecture, uses a combination of many different components, namely differential attention on all attention layers, SWA, Mamba-1, YOCO for a full attention layer, and a YOCO for the last Mamba-1 layer, cross attention layers, and GMUs (which could be seen as a "cross attention" analogue for Mamba layers).
1 parent 49bfdde commit f01113a

14 files changed

Lines changed: 1102 additions & 3 deletions

‎convert_hf_to_gguf.py‎

Lines changed: 273 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5453,6 +5453,279 @@ def prepare_tensors(self):
54535453
raise ValueError(f"Unprocessed experts: {experts}")
54545454

54555455

5456+
@ModelBase.register("Phi4FlashForCausalLM")
5457+
class Phi4FlashModel(TextModel):
5458+
model_arch = gguf.MODEL_ARCH.PHI4FLASH
5459+
5460+
def __init__(self, *args, **kwargs):
5461+
super().__init__(*args, **kwargs)
5462+
self.mb_per_layer = self.hparams.get('mb_per_layer', 2)
5463+
self.sliding_window = self.hparams.get('sliding_window', 512)
5464+
5465+
def set_gguf_parameters(self):
5466+
# 1. Standard Transformer Parameters
5467+
n_embd = self.find_hparam(["hidden_size", "n_embd"])
5468+
n_head = self.find_hparam(["num_attention_heads", "n_head"])
5469+
n_head_kv = self.find_hparam(["num_key_value_heads", "n_head_kv", "num_attention_heads"])
5470+
rms_eps = self.find_hparam(["layer_norm_eps", "rms_norm_eps"])
5471+
max_pos_embds = self.find_hparam(["max_position_embeddings", "n_positions"])
5472+
5473+
# FIX: Extract intermediate_size (n_ff)
5474+
n_ff = self.find_hparam(["intermediate_size"])
5475+
5476+
self.gguf_writer.add_embedding_length(n_embd)
5477+
5478+
# FIX: Write the feed-forward length to the GGUF so llama.cpp knows it's 20480 and not 0
5479+
self.gguf_writer.add_feed_forward_length(n_ff)
5480+
5481+
self.gguf_writer.add_block_count(self.block_count)
5482+
self.gguf_writer.add_head_count(n_head)
5483+
self.gguf_writer.add_head_count_kv(n_head_kv)
5484+
self.gguf_writer.add_layer_norm_rms_eps(rms_eps)
5485+
self.gguf_writer.add_context_length(max_pos_embds)
5486+
self.gguf_writer.add_file_type(self.ftype)
5487+
5488+
# 2. Phi-4-Flash Specific Parameters
5489+
layer_types = []
5490+
for i in range(self.block_count):
5491+
is_mamba = (i % self.mb_per_layer == 0)
5492+
if i < self.block_count // 2:
5493+
layer_types.append(0 if is_mamba else 1)
5494+
elif i == self.block_count // 2:
5495+
layer_types.append(0)
5496+
elif i == self.block_count // 2 + 1:
5497+
layer_types.append(2)
5498+
else:
5499+
layer_types.append(3 if is_mamba else 4)
5500+
5501+
self.gguf_writer.add_array("phi4flash.layer_types", layer_types)
5502+
self.gguf_writer.add_uint32("phi4flash.mb_per_layer", self.mb_per_layer)
5503+
self.gguf_writer.add_uint32("phi4flash.sliding_window", self.sliding_window)
5504+
self.gguf_writer.add_uint32("phi4flash.pivot_layer", 17)
5505+
self.gguf_writer.add_uint32("phi4flash.ssm_cache_layer", 16)
5506+
5507+
# Mamba-1 intrinsic dimensions
5508+
self.gguf_writer.add_uint32("phi4flash.ssm_d_conv", 4)
5509+
self.gguf_writer.add_uint32("phi4flash.ssm_d_state", 16)
5510+
self.gguf_writer.add_uint32("phi4flash.ssm_d_inner", 5120)
5511+
self.gguf_writer.add_uint32("phi4flash.ssm_dt_rank", 160)
5512+
self.gguf_writer.add_uint32("phi4flash.ssm_expand", 2)
5513+
5514+
def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
5515+
import torch
5516+
import numpy as np
5517+
5518+
# Helper to ensure tensor is transposed in actual memory layout
5519+
def force_transpose_2d(tensor):
5520+
"""Force transpose with contiguous memory layout via numpy"""
5521+
# Convert to numpy (forces materialization)
5522+
np_array = tensor.cpu().numpy()
5523+
# Transpose in numpy (creates new array with transposed data)
5524+
np_transposed = np.ascontiguousarray(np_array.T)
5525+
# Convert back to torch
5526+
return torch.from_numpy(np_transposed)
5527+
5528+
# Map general tensors (NO transpose for embeddings/norms)
5529+
if name == "model.embed_tokens.weight":
5530+
return [("token_embd.weight", data_torch)]
5531+
if name == "model.final_layernorm.weight":
5532+
return [("output_norm.weight", data_torch)]
5533+
if name == "model.final_layernorm.bias":
5534+
return [("output_norm.bias", data_torch)]
5535+
if name == "lm_head.weight":
5536+
return [] # Skip, tied to embed_tokens
5537+
5538+
if not name.startswith("model.layers.") or bid is None:
5539+
return super().modify_tensors(data_torch, name, bid)
5540+
5541+
rest = name.split(f"model.layers.{bid}.")[-1]
5542+
5543+
# Determine layer type
5544+
is_mamba = (bid % self.mb_per_layer == 0)
5545+
if bid < self.block_count // 2:
5546+
layer_type = 0 if is_mamba else 1
5547+
elif bid == self.block_count // 2:
5548+
layer_type = 0
5549+
elif bid == self.block_count // 2 + 1:
5550+
layer_type = 2
5551+
else:
5552+
layer_type = 3 if is_mamba else 4
5553+
5554+
# Shared block normalizations (NO transpose for norms/biases)
5555+
if rest == "input_layernorm.weight":
5556+
return [(f"blk.{bid}.attn_norm.weight", data_torch)]
5557+
if rest == "input_layernorm.bias":
5558+
return [(f"blk.{bid}.attn_norm.bias", data_torch)]
5559+
if rest == "post_attention_layernorm.weight":
5560+
return [(f"blk.{bid}.ffn_norm.weight", data_torch)]
5561+
if rest == "post_attention_layernorm.bias":
5562+
return [(f"blk.{bid}.ffn_norm.bias", data_torch)]
5563+
5564+
# MLP (TRANSPOSE all weight matrices!)
5565+
# Check what your HF model actually has - could be gate_proj/up_proj/down_proj
5566+
# or fc1/fc2/fc3, etc.
5567+
if rest == "mlp.gate_up_proj.weight":
5568+
# Combined gate+up projection [2*n_ff, n_embd]
5569+
return [(f"blk.{bid}.ffn_up.weight", force_transpose_2d(data_torch))]
5570+
elif rest == "mlp.up_proj.weight":
5571+
return [(f"blk.{bid}.ffn_up.weight", force_transpose_2d(data_torch))]
5572+
elif rest == "mlp.down_proj.weight":
5573+
return [(f"blk.{bid}.ffn_down.weight", force_transpose_2d(data_torch))]
5574+
# Alternative naming
5575+
elif rest == "mlp.fc1.weight":
5576+
# Keep combined [2*n_ff, n_embd] — LLM_FFN_SWIGLU splits it internally
5577+
# No force_transpose_2d — ggml reversal gives ne=[n_embd, 2*n_ff]
5578+
return [(f"blk.{bid}.ffn_up.weight", data_torch)]
5579+
elif rest == "mlp.fc2.weight":
5580+
# FIX: Remove force_transpose_2d!
5581+
return [(f"blk.{bid}.ffn_down.weight", data_torch)]
5582+
5583+
# ==================================================
5584+
# MAMBA LAYERS (type 0)
5585+
# ==================================================
5586+
if layer_type == 0:
5587+
if rest == "attn.in_proj.weight":
5588+
# Remove force_transpose_2d — ggml dimension reversal gives correct ne=[2560,10240]
5589+
return [(f"blk.{bid}.ssm_in.weight", data_torch)]
5590+
5591+
if rest == "attn.conv1d.weight":
5592+
# print(f"Layer {bid}: conv1d.weight input shape: {data_torch.shape}")
5593+
5594+
# PyTorch shape is [5120, 1, 4].
5595+
# Remove the channel dimension (dim 1) to get [5120, 4]
5596+
tensor = data_torch.squeeze(1)
5597+
5598+
# DO NOT transpose! ggml automatically reverses dimensions when loading.
5599+
# A PyTorch tensor of shape [5120, 4] will correctly become a ggml tensor
5600+
# with dimensions ne = [4, 5120], which is exactly what llama.cpp expects.
5601+
5602+
# Just ensure it's contiguous in memory before returning
5603+
if not tensor.is_contiguous():
5604+
tensor = tensor.contiguous()
5605+
5606+
# print(f"Layer {bid}: final shape: {tensor.shape}")
5607+
return [(f"blk.{bid}.ssm_conv1d.weight", tensor)]
5608+
5609+
if rest == "attn.conv1d.bias":
5610+
# Bias stays 1D
5611+
return [(f"blk.{bid}.ssm_conv1d.bias", data_torch)]
5612+
5613+
if rest == "attn.x_proj.weight":
5614+
# Remove force_transpose_2d — gives ne=[5120,192]
5615+
return [(f"blk.{bid}.ssm_x.weight", data_torch)]
5616+
5617+
if rest == "attn.dt_proj.weight":
5618+
# Remove force_transpose_2d — gives ne=[160,5120]
5619+
return [(f"blk.{bid}.ssm_dt.weight", data_torch)]
5620+
5621+
if rest == "attn.dt_proj.bias":
5622+
return [(f"blk.{bid}.ssm_dt.bias", data_torch)]
5623+
5624+
if rest == "attn.A_log":
5625+
# Convert A_log -> A = -exp(A_log)
5626+
# HF: [d_inner, d_state] -> llama.cpp: [d_state, d_inner]
5627+
A = -torch.exp(data_torch.float())
5628+
return [(f"blk.{bid}.ssm_a", A)]
5629+
5630+
if rest == "attn.D":
5631+
return [(f"blk.{bid}.ssm_d", data_torch)]
5632+
5633+
if rest == "attn.out_proj.weight":
5634+
# Remove force_transpose_2d — gives ne=[5120,2560]
5635+
return [(f"blk.{bid}.ssm_out.weight", data_torch)]
5636+
5637+
# ==================================================
5638+
# GMU LAYERS (type 3)
5639+
# ==================================================
5640+
elif layer_type == 3:
5641+
if rest == "attn.in_proj.weight":
5642+
# HF: [d_inner, n_embd] -> llama.cpp: [n_embd, d_inner]
5643+
return [(f"blk.{bid}.gmu_in.weight", data_torch)]
5644+
5645+
if rest == "attn.out_proj.weight":
5646+
# HF: [n_embd, d_inner] -> llama.cpp: [d_inner, n_embd]
5647+
return [(f"blk.{bid}.gmu_out.weight", data_torch)]
5648+
5649+
# ==================================================
5650+
# ATTENTION LAYERS (types 1, 2, 4)
5651+
# ==================================================
5652+
elif layer_type in [1, 2, 4]:
5653+
if rest == "attn.Wqkv.weight":
5654+
w = data_torch
5655+
head_dim = 64
5656+
n_q_heads = 40
5657+
n_kv_heads = 20
5658+
q_size = n_q_heads * head_dim # 2560
5659+
k_size = n_kv_heads * head_dim # 1280
5660+
5661+
def reorder(block, n_heads):
5662+
h = block.reshape(n_heads, head_dim, block.shape[-1])
5663+
return torch.cat([h[0::2], h[1::2]], dim=0).reshape(n_heads * head_dim, block.shape[-1]).contiguous()
5664+
5665+
if w.shape[0] == q_size:
5666+
# Cross-attention layers (type 4): Q only, no K/V block
5667+
result = reorder(w, n_q_heads)
5668+
else:
5669+
# SWA (type 1) and full-attn (type 2): Q | K | V
5670+
q_block = w[:q_size]
5671+
k_block = w[q_size : q_size + k_size]
5672+
v_block = w[q_size + k_size:]
5673+
result = torch.cat([
5674+
reorder(q_block, n_q_heads),
5675+
reorder(k_block, n_kv_heads),
5676+
v_block, # V is NOT reordered
5677+
], dim=0).contiguous()
5678+
5679+
return [(f"blk.{bid}.attn_qkv.weight", result)]
5680+
5681+
if rest == "attn.Wqkv.bias":
5682+
b = data_torch
5683+
head_dim = 64
5684+
n_q_heads = 40
5685+
n_kv_heads = 20
5686+
q_size = n_q_heads * head_dim # 2560
5687+
k_size = n_kv_heads * head_dim # 1280
5688+
5689+
def reorder_bias(block, n_heads):
5690+
h = block.reshape(n_heads, head_dim)
5691+
return torch.cat([h[0::2], h[1::2]], dim=0).reshape(-1).contiguous()
5692+
5693+
if b.shape[0] == q_size:
5694+
result = reorder_bias(b, n_q_heads)
5695+
else:
5696+
result = torch.cat([
5697+
reorder_bias(b[:q_size], n_q_heads),
5698+
reorder_bias(b[q_size : q_size + k_size], n_kv_heads),
5699+
b[q_size + k_size:], # V bias NOT reordered
5700+
], dim=0).contiguous()
5701+
5702+
return [(f"blk.{bid}.attn_qkv.bias", result)]
5703+
5704+
# Lambda parameters (1D vectors, no transpose)
5705+
if rest == "attn.inner_cross_attn.lambda_q1":
5706+
return [(f"blk.{bid}.attn_lambda_q1.weight", data_torch)]
5707+
if rest == "attn.inner_cross_attn.lambda_q2":
5708+
return [(f"blk.{bid}.attn_lambda_q2.weight", data_torch)]
5709+
if rest == "attn.inner_cross_attn.lambda_k1":
5710+
return [(f"blk.{bid}.attn_lambda_k1.weight", data_torch)]
5711+
if rest == "attn.inner_cross_attn.lambda_k2":
5712+
return [(f"blk.{bid}.attn_lambda_k2.weight", data_torch)]
5713+
5714+
# SubLN (1D, no transpose)
5715+
if rest == "attn.inner_cross_attn.subln.weight":
5716+
return [(f"blk.{bid}.attn_subln.weight", data_torch)]
5717+
5718+
# Output projection
5719+
if rest == "attn.out_proj.weight":
5720+
# HF: [n_embd, n_embd] -> llama.cpp: [n_embd, n_embd] (square, but still transpose!)
5721+
return [(f"blk.{bid}.attn_output.weight", data_torch)]
5722+
5723+
if rest == "attn.out_proj.bias":
5724+
return [(f"blk.{bid}.attn_output.bias", data_torch)]
5725+
5726+
return super().modify_tensors(data_torch, name, bid)
5727+
5728+
54565729
@ModelBase.register("PlamoForCausalLM")
54575730
class PlamoModel(TextModel):
54585731
model_arch = gguf.MODEL_ARCH.PLAMO

‎gguf-py/gguf/constants.py‎

Lines changed: 35 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -400,6 +400,7 @@ class MODEL_ARCH(IntEnum):
400400
PHI2 = auto()
401401
PHI3 = auto()
402402
PHIMOE = auto()
403+
PHI4FLASH = auto()
403404
PLAMO = auto()
404405
PLAMO2 = auto()
405406
PLAMO3 = auto()
@@ -806,7 +807,13 @@ class MODEL_TENSOR(IntEnum):
806807
A_ENC_CONV_NORM = auto() # SSM conv
807808
A_ENC_CONV_PW1 = auto()
808809
A_ENC_CONV_PW2 = auto()
809-
810+
GMU_IN = auto()
811+
GMU_OUT = auto()
812+
ATTN_LAMBDA_Q1 = auto()
813+
ATTN_LAMBDA_Q2 = auto()
814+
ATTN_LAMBDA_K1 = auto()
815+
ATTN_LAMBDA_K2 = auto()
816+
ATTN_SUBLN = auto()
810817

811818
MODEL_ARCH_NAMES: dict[MODEL_ARCH, str] = {
812819
MODEL_ARCH.MMPROJ: "clip", # dummy arch for clip.cpp
@@ -846,6 +853,7 @@ class MODEL_TENSOR(IntEnum):
846853
MODEL_ARCH.PHI2: "phi2",
847854
MODEL_ARCH.PHI3: "phi3",
848855
MODEL_ARCH.PHIMOE: "phimoe",
856+
MODEL_ARCH.PHI4FLASH: "phi4flash",
849857
MODEL_ARCH.PLAMO: "plamo",
850858
MODEL_ARCH.PLAMO2: "plamo2",
851859
MODEL_ARCH.PLAMO3: "plamo3",
@@ -2051,6 +2059,32 @@ class MODEL_TENSOR(IntEnum):
20512059
MODEL_TENSOR.FFN_DOWN_EXP,
20522060
MODEL_TENSOR.FFN_UP_EXP,
20532061
],
2062+
MODEL_ARCH.PHI4FLASH: [
2063+
MODEL_TENSOR.TOKEN_EMBD,
2064+
MODEL_TENSOR.OUTPUT_NORM,
2065+
MODEL_TENSOR.OUTPUT,
2066+
MODEL_TENSOR.ATTN_NORM,
2067+
MODEL_TENSOR.FFN_NORM,
2068+
MODEL_TENSOR.ATTN_QKV,
2069+
MODEL_TENSOR.ATTN_OUT,
2070+
MODEL_TENSOR.FFN_UP,
2071+
MODEL_TENSOR.FFN_GATE,
2072+
MODEL_TENSOR.FFN_DOWN,
2073+
MODEL_TENSOR.SSM_IN,
2074+
MODEL_TENSOR.SSM_CONV1D,
2075+
MODEL_TENSOR.SSM_X,
2076+
MODEL_TENSOR.SSM_DT,
2077+
MODEL_TENSOR.SSM_A,
2078+
MODEL_TENSOR.SSM_D,
2079+
MODEL_TENSOR.SSM_OUT,
2080+
MODEL_TENSOR.GMU_IN,
2081+
MODEL_TENSOR.GMU_OUT,
2082+
MODEL_TENSOR.ATTN_LAMBDA_Q1,
2083+
MODEL_TENSOR.ATTN_LAMBDA_Q2,
2084+
MODEL_TENSOR.ATTN_LAMBDA_K1,
2085+
MODEL_TENSOR.ATTN_LAMBDA_K2,
2086+
MODEL_TENSOR.ATTN_SUBLN,
2087+
],
20542088
MODEL_ARCH.CODESHELL: [
20552089
MODEL_TENSOR.TOKEN_EMBD,
20562090
MODEL_TENSOR.POS_EMBD,

‎src/CMakeLists.txt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -115,6 +115,7 @@ add_library(llama
115115
models/pangu-embedded.cpp
116116
models/phi2.cpp
117117
models/phi3.cpp
118+
models/phi4flash.cpp
118119
models/plamo.cpp
119120
models/plamo2.cpp
120121
models/plamo3.cpp

0 commit comments

Comments
 (0)