overlay.patch

diff --git a/deploy/functional/__init__.py b/deploy/functional/__init__.py
index 29951ec..2d74dcc 100644
--- a/deploy/functional/__init__.py
+++ b/deploy/functional/__init__.py
@@ -1,3 +1,18 @@
 from .quantization import pack_i4, unpack_i4, asym_quant_dequant, sym_quant_dequant
-from .online_trans import (
-    matmul_hadU_cuda, get_decompose_dim)
+
+try:
+    from .online_trans import matmul_hadU_cuda, get_decompose_dim
+except ImportError:  # Hadamard extra is optional; packed W4A4 uses CUDA matmul
+    matmul_hadU_cuda = None
+
+    def get_decompose_dim(n):
+        a = int(n ** 0.5)
+        if a * a < n:
+            a += 1
+        while True:
+            tmp = a * a - n
+            b = int(tmp ** 0.5)
+            if b * b == tmp:
+                break
+            a += 1
+        return a - b, a + b
diff --git a/deploy/nn/linear.py b/deploy/nn/linear.py
index 026cb5d..cb7a16b 100644
--- a/deploy/nn/linear.py
+++ b/deploy/nn/linear.py
@@ -27,8 +27,9 @@ class Linear4bit(torch.nn.Module):
         super().__init__()
         self.in_features = in_features
         self.out_features = out_features
+        # CUDA sym_dequant takes half* scale_col; default zeros() is float32.
         self.register_buffer('weight_scales',
-                             torch.zeros((self.out_features, 1), requires_grad=False))
+                             torch.zeros((self.out_features, 1), dtype=torch.float16, requires_grad=False))
         self.register_buffer('weight', (torch.randint(1, 7, (self.out_features, self.in_features // 2),
                                                              # SubByte weight
                                                              dtype=torch.uint8, requires_grad=False)))
@@ -46,12 +47,12 @@ class Linear4bit(torch.nn.Module):
         #shape_handler = ShapeHandler(quantized_x)
         #quantized_x = shape_handler.flatten(quantized_x)
         x = deploy.matmul(x, self.weight)
-        #out = shape_handler.unflatten(
-        #    deploy.sym_dequant(int_result, scales_x, self.weight_scales))
+        scales_x = scales_x.to(dtype=torch.float16)
+        weight_scales = self.weight_scales.to(dtype=torch.float16)
         if self.bias is not None:
-            return deploy.sym_dequant(x, scales_x, self.weight_scales) + self.bias
+            return deploy.sym_dequant(x, scales_x, weight_scales) + self.bias
         else:
-            return deploy.sym_dequant(x, scales_x, self.weight_scales)
+            return deploy.sym_dequant(x, scales_x, weight_scales)
 
     @staticmethod
     def from_float(module: torch.nn.Linear, weight_scales=None,):
@@ -63,15 +64,16 @@ class Linear4bit(torch.nn.Module):
         weight_matrix = module.weight.data
         
         
-        int_module = Linear4bit(module.in_features, module.out_features, bias=module.bias is not None, dtype=weight_matrix.dtype).to(weight_matrix.dtype)
+        int_module = Linear4bit(module.in_features, module.out_features, bias=module.bias is not None, dtype=torch.float16)
         if weight_scales is not None:
             assert weight_scales.shape == (module.out_features, 1), 'weight_scales should have shape (out_features, 1)'
             weight_matrix = weight_matrix.cuda()
-            int_module.weight_scales.copy_(weight_scales.to(weight_matrix.dtype))
-            int_rounded_weight = (weight_matrix/weight_scales.cuda()).round()
+            scales = weight_scales.detach().to(device=weight_matrix.device, dtype=torch.float16)
+            int_module.weight_scales.copy_(scales.cpu() if int_module.weight_scales.device.type == "cpu" else scales)
+            int_rounded_weight = (weight_matrix / scales).round()
             int_module.weight.copy_(deploy.functional.pack_i4(int_rounded_weight.to(torch.int8)).cpu())
         
             if module.bias is not None:
-                int_module.bias.copy_(module.bias)
+                int_module.bias.copy_(module.bias.to(dtype=torch.float16))
         
         return int_module
diff --git a/deploy/nn/quantization.py b/deploy/nn/quantization.py
index 31e6a49..91dedd5 100644
--- a/deploy/nn/quantization.py
+++ b/deploy/nn/quantization.py
@@ -12,6 +12,7 @@ class Quantizer(torch.nn.Module):
 
     def forward(self, x):
         if not isinstance(x, deploy.PackedQuantizedTensor):
+            x = x.to(dtype=torch.float16)
             if self.lac:
                 reshaped_x = x.reshape((-1, x.shape[-1]))
                 xmax, xmin = reshaped_x.amax(1, keepdim=True), reshaped_x.amin(1, keepdim=True)
diff --git a/flatquant/args_utils.py b/flatquant/args_utils.py
index 1cda369..ee49361 100644
--- a/flatquant/args_utils.py
+++ b/flatquant/args_utils.py
@@ -21,6 +21,8 @@ supported_models = [
             './modelzoo/meta-llama/Llama-3.3-70B-Instruct', 
             './modelzoo/Qwen/Qwen2.5-7B-Instruct', 
             './modelzoo/Qwen/Qwen2.5-32B-Instruct', 
+            './modelzoo/microsoft/Phi-3-mini-4k-instruct',
+            'microsoft/Phi-3-mini-4k-instruct',
             ]
 supported_datasets = ['wikitext2', 'c4', 'pile']
 
@@ -30,7 +32,7 @@ def parser_gen():
 
     # General Arguments
     parser.add_argument('--model', type=str, default='meta-llama/Llama-2-7b-hf',
-                        help='Model to load.', choices=supported_models)
+                        help='Model to load.')
     parser.add_argument('--seed', type=int, default=0, help='Random seed for HuggingFace and PyTorch.')
     parser.add_argument('--hf_token', type=str, default=None, help='HuggingFace token for model access.')
 
diff --git a/flatquant/data_utils.py b/flatquant/data_utils.py
index 5ddc680..ffc53dc 100644
--- a/flatquant/data_utils.py
+++ b/flatquant/data_utils.py
@@ -9,13 +9,20 @@ class TokenizerWrapper:
         self.input_ids = input_ids
 
 
+def _load_wikitext(split):
+    local_dir = './datasets/wikitext'
+    if os.path.isdir(local_dir):
+        return datasets.load_dataset(local_dir, 'wikitext-2-raw-v1', split=split)
+    return datasets.load_dataset('wikitext', 'wikitext-2-raw-v1', split=split)
+
+
 def get_wikitext2(nsamples, seqlen, tokenizer, eval_mode=False):
     if eval_mode:
-        testdata = datasets.load_dataset('./datasets/wikitext', 'wikitext-2-raw-v1', split='test')
+        testdata = _load_wikitext('test')
         testenc = tokenizer("\n\n".join(testdata['text']), return_tensors='pt')
         return testenc
     else:
-        traindata = datasets.load_dataset('./datasets/wikitext', 'wikitext-2-raw-v1', split='train')
+        traindata = _load_wikitext('train')
         traindata = traindata.filter(lambda x: len(x) > 0)
         traindata = traindata.map(lambda x : {'text': x['text'].strip()})
         trainenc = tokenizer("\n\n".join(traindata['text']), return_tensors='pt')    
diff --git a/flatquant/flat_linear.py b/flatquant/flat_linear.py
index 6e142de..45172d9 100644
--- a/flatquant/flat_linear.py
+++ b/flatquant/flat_linear.py
@@ -24,6 +24,7 @@ class FlatQuantizedLinear(nn.Module):
             self.sigmoid = nn.Sigmoid()
 
         self._eval_mode = False
+        self._packed = False
 
     def apply_wclip(self, weight):
         wmin, wmax = weight.min(1, keepdim=True)[0], weight.max(1, keepdim=True)[0]
@@ -73,6 +74,8 @@ class FlatQuantizedLinear(nn.Module):
             return self._eval_forward(hidden_states)
 
     def _eval_forward(self, hidden_states):
+        if getattr(self, "_packed", False):
+            return self.linear(self.a4(hidden_states))
         x_dtype = hidden_states.dtype
         hidden_states = self.act_quantizer(hidden_states).to(x_dtype)
 
diff --git a/flatquant/model_tools/phi3_utils.py b/flatquant/model_tools/phi3_utils.py
new file mode 100644
index 0000000..cc62529
--- /dev/null
+++ b/flatquant/model_tools/phi3_utils.py
@@ -0,0 +1,423 @@
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from flatquant.quant_utils import ActivationQuantizer
+from flatquant.utils import skip_initialization
+from flatquant.function_utils import get_init_scale, get_decompose_dim
+from flatquant.trans_utils import SVDSingleTransMatrix, SVDDecomposeTransMatrix
+from flatquant.trans_utils import InvSingleTransMatrix, InvDecomposeTransMatrix
+from flatquant.flat_linear import FlatQuantizedLinear
+
+from transformers.models.phi3.modeling_phi3 import apply_rotary_pos_emb, repeat_kv
+
+
+def _clip_scalar(value):
+    """Triton kron_matmul needs a Python float; 0-d tensors compile as pointers."""
+    if value is None:
+        return 1.0
+    if torch.is_tensor(value):
+        return float(value.detach().reshape(-1)[0].cpu())
+    return float(value)
+
+
+def _eval_proj(module, hidden_states, trans=None):
+    """Eval projection: fuse decompose T + act quant into repo kron_matmul when packed.
+
+    Import deploy.kernels.kron_matmul, not deploy.functional.online_trans (that
+    module imports optional fast_hadamard_transform at load).
+    """
+    packed = getattr(module, "_packed", False)
+    if (
+        packed
+        and trans is not None
+        and getattr(trans, "_eval_mode", False)
+        and hasattr(trans, "matrix_left")
+        and hasattr(trans, "matrix_right")
+        and hidden_states.dim() == 3
+    ):
+        from deploy.kernels.kron_matmul import kron_matmul
+
+        a4 = module.a4
+        bsz, seq_len, _ = hidden_states.shape
+        left = trans.matrix_left.detach().to(dtype=hidden_states.dtype).T.contiguous()
+        right = trans.matrix_right.detach().to(dtype=hidden_states.dtype).contiguous()
+        x = hidden_states.contiguous().reshape(-1, left.shape[0], right.shape[0])
+        packed_x = kron_matmul(
+            left,
+            x,
+            right,
+            seq_len,
+            _clip_scalar(getattr(a4, "clip_factor_a_max", 1.0)),
+            _clip_scalar(getattr(a4, "clip_factor_a_min", 1.0)),
+        )
+        packed_x.quantized_x = packed_x.quantized_x.reshape(bsz, seq_len, -1)
+        packed_x.scales_x = packed_x.scales_x.reshape(bsz, 1, seq_len)
+        return module.linear(packed_x)
+    if trans is not None:
+        hidden_states = trans(hidden_states)
+    if packed:
+        return module.linear(module.a4(hidden_states))
+    return module(hidden_states)
+
+
+class FlatQuantPhi3MLP(nn.Module):
+    """FlatQuant wrapper for Phi-3 fused gate_up_proj + down_proj."""
+
+    def __init__(self, args, module):
+        super().__init__()
+        self.args = args
+        self.config = getattr(module, "config", None)
+        self.hidden_size = getattr(module, "hidden_size", module.gate_up_proj.in_features)
+        self.intermediate_size = module.down_proj.in_features
+        self.act_fn = module.activation_fn
+        self.gate_up_proj = FlatQuantizedLinear(args, module.gate_up_proj)
+        self.down_proj = FlatQuantizedLinear(args, module.down_proj)
+        self.add_fq_trans()
+
+        self._ori_mode = False
+        self.diag_init = args.diag_init
+        if self.diag_init == "sq_style":
+            self.up_smax = torch.ones_like(self.gate_up_proj.linear.weight.abs().max(dim=0)[0]).cuda() * 1e-5
+            self.down_smax = torch.ones_like(self.down_proj.linear.weight.abs().max(dim=0)[0]).cuda() * 1e-5
+
+    def add_fq_trans(self):
+        if self.args.direct_inv:
+            DecomposeTransMatrix = InvDecomposeTransMatrix
+        else:
+            DecomposeTransMatrix = SVDDecomposeTransMatrix
+        if self.args.w_bits < 16 or self.args.a_bits < 16:
+            up_dim_left, up_dim_right = get_decompose_dim(self.gate_up_proj.linear.weight.shape[1])
+            self.up_gate_trans = DecomposeTransMatrix(up_dim_left, up_dim_right, add_diag=self.args.add_diag)
+            down_dim_left, down_dim_right = get_decompose_dim(self.down_proj.linear.weight.shape[1])
+            self.down_trans = DecomposeTransMatrix(down_dim_left, down_dim_right, add_diag=self.args.add_diag)
+        else:
+            self.up_gate_trans, self.down_trans = None, None
+
+    def _trans_forward(self, x):
+        up_states = _eval_proj(self.gate_up_proj, x, self.up_gate_trans)
+        gate, up_states = up_states.chunk(2, dim=-1)
+        x_act_fn = self.act_fn(gate) * up_states
+        return _eval_proj(self.down_proj, x_act_fn, self.down_trans)
+
+    def _ori_forward(self, x):
+        if self.diag_init == "sq_style":
+            self.up_smax = torch.maximum(
+                self.up_smax, x.reshape(-1, x.shape[-1]).abs().max(0)[0].clone().detach()
+            )
+        up_states = self.gate_up_proj._ori_forward(x)
+        gate, up_states = up_states.chunk(2, dim=-1)
+        x = self.act_fn(gate) * up_states
+        if self.diag_init == "sq_style":
+            self.down_smax = torch.maximum(
+                self.down_smax, x.reshape(-1, x.shape[-1]).abs().max(0)[0].clone().detach()
+            )
+        return self.down_proj._ori_forward(x)
+
+    def forward(self, x):
+        if self._ori_mode:
+            return self._ori_forward(x)
+        return self._trans_forward(x)
+
+    def reparameterize(self):
+        if self.up_gate_trans is not None:
+            self.up_gate_trans.to_eval_mode()
+            self.down_trans.to_eval_mode()
+        self.gate_up_proj.reparameterize(qa_trans=self.up_gate_trans)
+        self.down_proj.reparameterize(qa_trans=self.down_trans)
+        if self.up_gate_trans is not None:
+            self.up_gate_trans.use_diag = False
+        if self.down_trans is not None and self.down_trans.add_diag:
+            up_weight = self.gate_up_proj.linear.weight
+            ori_dtype = up_weight.dtype
+            # Fuse down-channel scales into the up half of the fused gate_up matrix.
+            up_half = up_weight[self.intermediate_size :].to(torch.float64)
+            up_half = up_half.T.mul(self.down_trans.diag_scale.to(torch.float64)).T
+            up_weight = up_weight.clone()
+            up_weight[self.intermediate_size :] = up_half.to(ori_dtype)
+            self.gate_up_proj.linear.weight.data = up_weight
+            self.down_trans.use_diag = False
+
+    def init_diag_scale(self, alpha=0.5):
+        assert hasattr(self, "up_smax") and hasattr(self, "down_smax")
+        upw_smax = self.gate_up_proj.linear.weight.abs().max(dim=0)[0]
+        downw_smax = self.down_proj.linear.weight.abs().max(dim=0)[0]
+        if self.up_gate_trans is not None:
+            self.up_gate_trans.diag_scale.data = get_init_scale(upw_smax, self.up_smax, alpha)
+        if self.down_trans is not None:
+            self.down_trans.diag_scale.data = get_init_scale(downw_smax, self.down_smax, alpha)
+        del self.up_smax, self.down_smax
+        self.diag_init = None
+
+    def rep_matrix_only(self):
+        if self.up_gate_trans is not None:
+            self.up_gate_trans.to_eval_mode()
+            self.down_trans.to_eval_mode()
+
+
+class FlatQuantPhi3Attention(nn.Module):
+    """FlatQuant wrapper for Phi-3 fused qkv_proj + o_proj."""
+
+    def __init__(self, args, module):
+        super().__init__()
+        self.args = args
+        self.config = module.config
+        self.layer_idx = module.layer_idx
+        self.hidden_size = module.hidden_size
+        self.num_heads = module.num_heads
+        self.head_dim = module.head_dim
+        self.num_key_value_heads = module.num_key_value_heads
+        self.num_key_value_groups = module.num_key_value_groups
+        self.attention_dropout = module.attention_dropout
+        self.rotary_emb = module.rotary_emb
+        self.query_pos = self.num_heads * self.head_dim
+        self.kv_dim = self.num_key_value_heads * self.head_dim
+
+        self.qkv_proj = FlatQuantizedLinear(args, module.qkv_proj)
+        self.o_proj = FlatQuantizedLinear(args, module.o_proj)
+        self.add_fq_trans()
+
+        if args.q_bits < 16:
+            self.q_cache_quantizer = ActivationQuantizer(
+                bits=args.q_bits, sym=not (args.q_asym), lac=args.lac, groupsize=-1
+            )
+        if args.k_bits < 16:
+            self.k_cache_quantizer = ActivationQuantizer(
+                bits=args.k_bits, sym=not (args.k_asym), lac=args.lac, groupsize=-1
+            )
+        if args.v_bits < 16:
+            self.v_cache_quantizer = ActivationQuantizer(
+                bits=args.v_bits, sym=not (args.v_asym), lac=args.lac, groupsize=-1
+            )
+
+        self._ori_mode = False
+        self._eval_mode = False
+        self.diag_init = args.diag_init
+        if self.diag_init == "sq_style":
+            self.ln_smax = torch.ones_like(self.qkv_proj.linear.weight.abs().max(dim=0)[0]).cuda() * 1e-5
+
+    def add_fq_trans(self):
+        if self.args.direct_inv:
+            SingleTransMatrix, DecomposeTransMatrix = InvSingleTransMatrix, InvDecomposeTransMatrix
+        else:
+            SingleTransMatrix, DecomposeTransMatrix = SVDSingleTransMatrix, SVDDecomposeTransMatrix
+        if self.args.w_bits < 16 or self.args.a_bits < 16:
+            ln_dim_left, ln_dim_right = get_decompose_dim(self.qkv_proj.linear.weight.shape[1])
+            self.ln_trans = DecomposeTransMatrix(ln_dim_left, ln_dim_right, add_diag=self.args.add_diag)
+            self.o_trans = SingleTransMatrix(self.config.num_attention_heads)
+        else:
+            self.ln_trans, self.o_trans = None, None
+
+        if self.args.k_bits < 16 or self.args.q_bits < 16:
+            self.kcache_trans = SingleTransMatrix(self.head_dim)
+        else:
+            self.kcache_trans = None
+        if self.args.v_bits < 16 or self.args.w_bits < 16 or self.args.a_bits < 16:
+            self.vcache_trans = SingleTransMatrix(self.head_dim)
+        else:
+            self.vcache_trans = None
+
+    def _split_qkv(self, qkv):
+        query_states = qkv[..., : self.query_pos]
+        key_states = qkv[..., self.query_pos : self.query_pos + self.kv_dim]
+        value_states = qkv[..., self.query_pos + self.kv_dim :]
+        return query_states, key_states, value_states
+
+    def _trans_forward_after_ln(self, hidden_states):
+        qkv = _eval_proj(self.qkv_proj, hidden_states, self.ln_trans)
+        query_states, key_states, value_states = self._split_qkv(qkv)
+        # Official llama applies T_v as out_trans on V weights, then fuses at
+        # reparameterize. After _eval_mode, do not transform V activations again.
+        if (
+            self.vcache_trans is not None
+            and not self.args.separate_vtrans
+            and not self._eval_mode
+        ):
+            value_states = self.vcache_trans(value_states)
+        return query_states, key_states, value_states
+
+    def _ori_forward_after_ln(self, hidden_states):
+        if self.diag_init == "sq_style" and hasattr(self, "ln_smax"):
+            self.ln_smax = torch.maximum(
+                self.ln_smax,
+                hidden_states.reshape(-1, hidden_states.shape[-1]).abs().max(0)[0].clone().detach(),
+            )
+        qkv = self.qkv_proj._ori_forward(hidden_states)
+        return self._split_qkv(qkv)
+
+    def quant_vcache(self, value_states):
+        if self.args.separate_vtrans and self.vcache_trans is not None:
+            value_states = self.vcache_trans(value_states)
+        if self.args.v_bits < 16:
+            value_states = self.v_cache_quantizer(value_states)
+        return value_states
+
+    def quant_kcache(self, q, k):
+        if not (self.args.k_bits < 16 or self.args.q_bits < 16):
+            return q, k
+        if self.kcache_trans is not None:
+            q = self.kcache_trans(q, inv_t=True)
+            k = self.kcache_trans(k)
+        if self.args.q_bits < 16:
+            q = self.q_cache_quantizer(q).to(q)
+        if self.args.k_bits < 16:
+            k = self.k_cache_quantizer(k).to(q)
+        return q, k
+
+    def forward(
+        self,
+        hidden_states,
+        attention_mask=None,
+        position_ids=None,
+        past_key_value=None,
+        output_attentions=False,
+        use_cache=False,
+        cache_position=None,
+        **kwargs,
+    ):
+        bsz, q_len, _ = hidden_states.size()
+        if self._ori_mode:
+            query_states, key_states, value_states = self._ori_forward_after_ln(hidden_states)
+        else:
+            query_states, key_states, value_states = self._trans_forward_after_ln(hidden_states)
+
+        query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+        key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+        value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+
+        kv_seq_len = key_states.shape[-2]
+        if past_key_value is not None:
+            if self.layer_idx is None:
+                raise ValueError(
+                    f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} "
+                    "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class "
+                    "with a layer index."
+                )
+            kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
+        cos, sin = self.rotary_emb(value_states, position_ids, seq_len=kv_seq_len)
+        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
+
+        if not self._ori_mode:
+            query_states, key_states = self.quant_kcache(query_states, key_states)
+            value_states = self.quant_vcache(value_states)
+
+        if past_key_value is not None:
+            cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
+            key_states, value_states = past_key_value.update(
+                key_states, value_states, self.layer_idx, cache_kwargs
+            )
+
+        key_states = repeat_kv(key_states, self.num_key_value_groups)
+        value_states = repeat_kv(value_states, self.num_key_value_groups)
+        dropout_p = self.attention_dropout if self.training else 0.0
+        kv_len = key_states.shape[-2]
+        if attention_mask is not None:
+            causal_mask = attention_mask[:, :, :, :kv_len]
+            attn_output = F.scaled_dot_product_attention(
+                query_states,
+                key_states,
+                value_states,
+                attn_mask=causal_mask,
+                dropout_p=dropout_p,
+                is_causal=False,
+            )
+        else:
+            # is_causal requires q_len == kv_len. Decode with KV cache is q_len=1.
+            attn_output = F.scaled_dot_product_attention(
+                query_states,
+                key_states,
+                value_states,
+                dropout_p=dropout_p,
+                is_causal=(q_len == kv_len),
+            )
+        attn_output = attn_output.transpose(1, 2).contiguous()
+        attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
+
+        if self._ori_mode:
+            attn_output = self.o_proj._ori_forward(attn_output)
+        else:
+            if self.o_trans is None and self.vcache_trans is not None:
+                init_shape = attn_output.shape
+                attn_output = attn_output.reshape(
+                    -1, self.config.num_attention_heads, self.config.hidden_size // self.config.num_attention_heads
+                )
+                attn_output = torch.matmul(
+                    attn_output, self.vcache_trans.get_matrix(inv_t=True).T.to(attn_output)
+                ).reshape(init_shape)
+                attn_output = self.o_proj(attn_output)
+            else:
+                init_shape = attn_output.shape
+                attn_output = attn_output.reshape(
+                    -1, self.config.num_attention_heads, self.config.hidden_size // self.config.num_attention_heads
+                )
+                attn_output = torch.matmul(
+                    self.o_trans.get_matrix().T.to(attn_output), attn_output
+                ).reshape(init_shape)
+                if not self._eval_mode:
+                    attn_o_og_it = self.o_trans.get_matrix(inv_t=True)
+                    attn_v_og_it = self.vcache_trans.get_matrix(inv_t=True)
+                    attn_output = self.o_proj(attn_output, qa_trans=[attn_o_og_it, attn_v_og_it])
+                else:
+                    attn_output = self.o_proj(attn_output)
+
+        if not output_attentions:
+            attn_weights = None
+        else:
+            attn_weights = None
+        return attn_output, attn_weights, past_key_value
+
+    def reparameterize(self):
+        if self.ln_trans is not None:
+            self.ln_trans.to_eval_mode()
+        if self.kcache_trans is not None:
+            self.kcache_trans.to_eval_mode()
+        if self.vcache_trans is not None:
+            self.vcache_trans.to_eval_mode()
+        if self.o_trans is not None:
+            self.o_trans.to_eval_mode()
+        self.qkv_proj.reparameterize(qa_trans=self.ln_trans)
+        if self.vcache_trans is not None and not self.args.separate_vtrans:
+            weight = self.qkv_proj.linear.weight.data
+            ori_dtype = weight.dtype
+            v_weight = weight[self.query_pos + self.kv_dim :].to(torch.float64)
+            v_weight = self.vcache_trans(v_weight.T).T
+            weight = weight.clone()
+            weight[self.query_pos + self.kv_dim :] = v_weight.to(ori_dtype)
+            self.qkv_proj.linear.weight.data = weight
+        if self.o_trans is not None and self.vcache_trans is not None:
+            attn_o_og_it = self.o_trans.get_matrix(inv_t=True)
+            attn_v_og_it = self.vcache_trans.get_matrix(inv_t=True)
+            self.o_proj.reparameterize(qa_trans=[attn_o_og_it, attn_v_og_it])
+        else:
+            self.o_proj.reparameterize()
+        self._eval_mode = True
+
+    def init_diag_scale(self, alpha=0.5):
+        assert hasattr(self, "ln_smax")
+        qkvw_smax = self.qkv_proj.linear.weight.abs().max(dim=0)[0]
+        if self.ln_trans is not None:
+            self.ln_trans.diag_scale.data = get_init_scale(qkvw_smax, self.ln_smax, alpha)
+        del self.ln_smax
+        self.diag_init = None
+
+    def rep_matrix_only(self):
+        if self.ln_trans is not None:
+            self.ln_trans.to_eval_mode()
+        if self.kcache_trans is not None:
+            self.kcache_trans.to_eval_mode()
+        if self.vcache_trans is not None:
+            self.vcache_trans.to_eval_mode()
+        if self.o_trans is not None:
+            self.o_trans.to_eval_mode()
+
+
+def apply_flatquant_to_phi3(args, model):
+    skip_initialization()
+    for layer in range(model.config.num_hidden_layers):
+        model.model.layers[layer].self_attn = FlatQuantPhi3Attention(
+            args, model.model.layers[layer].self_attn
+        )
+        model.model.layers[layer].mlp = FlatQuantPhi3MLP(args, model.model.layers[layer].mlp)
+    return model
diff --git a/flatquant/model_utils.py b/flatquant/model_utils.py
index 2d64c98..b181a02 100644
--- a/flatquant/model_utils.py
+++ b/flatquant/model_utils.py
@@ -1,3 +1,4 @@
+import os
 import torch
 import transformers
 import logging
@@ -62,6 +63,31 @@ def get_qwen2(model_name, hf_token):
     return model, apply_flatquant_to_qwen
 
 
+def get_phi3(model_name, hf_token):
+    skip_initialization()
+    from transformers import Phi3Config, Phi3ForCausalLM
+    from flatquant.model_tools.phi3_utils import apply_flatquant_to_phi3
+
+    is_local = os.path.isdir(model_name)
+    config = Phi3Config.from_pretrained(
+        model_name, use_auth_token=hf_token, local_files_only=is_local
+    )
+    config._attn_implementation = "eager"
+    config._attn_implementation_internal = "eager"
+    model = Phi3ForCausalLM.from_pretrained(
+        model_name,
+        torch_dtype="auto",
+        config=config,
+        use_auth_token=hf_token,
+        low_cpu_mem_usage=True,
+        attn_implementation="eager",
+        local_files_only=is_local,
+    )
+    model.seqlen = 2048
+    logging.info(f"---> Loading {model_name} Model with seq_len: {model.seqlen}")
+    return model, apply_flatquant_to_phi3
+
+
 def get_opt(model_name):
     skip_initialization()
     model = transformers.OPTForCausalLM.from_pretrained(model_name,
@@ -80,6 +106,8 @@ def get_model(model_name, hf_token=None):
         return get_llama(model_name, hf_token)
     elif 'qwen-2.5' in model_name:
         return get_qwen2(model_name, hf_token)
+    elif 'phi-3' in model_name.lower() or 'phi3' in model_name.lower():
+        return get_phi3(model_name, hf_token)
     else:
         raise ValueError(f'Unknown model {model_name}')
 
diff --git a/flatquant/pack_utils.py b/flatquant/pack_utils.py
new file mode 100644
index 0000000..ff8be0a
--- /dev/null
+++ b/flatquant/pack_utils.py
@@ -0,0 +1,127 @@
+"""Convert a reparameterized FlatQuant model onto official packed W4A4 Linear4bit."""
+from __future__ import annotations
+
+
+def _kernel_scale_dtype():
+    import torch
+
+    return torch.float16
+
+
+def cast_linear4bit_kernel_dtypes(model) -> int:
+    """Match CUDA sym_quant/sym_dequant: scale_row/scale_col must be float16.
+
+    Empty Linear4bit buffers default to float32, so load_state_dict keeps
+    float32 even when the packed file stored fp16 scales.
+    """
+    import torch
+    from deploy.nn.linear import Linear4bit
+
+    dtype = _kernel_scale_dtype()
+    converted = 0
+    for module in model.modules():
+        if not isinstance(module, Linear4bit):
+            continue
+        scales = module.weight_scales.detach().to(dtype=dtype)
+        module.register_buffer("weight_scales", scales)
+        if module.bias is not None:
+            module.register_buffer("bias", module.bias.detach().to(dtype=dtype))
+        converted += 1
+    return converted
+
+
+def prepare_packed_slots(model) -> None:
+    """Replace inner nn.Linear with empty Linear4bit so a packed state_dict can load."""
+    import torch
+    from deploy.nn.linear import Linear4bit
+    from deploy.nn.quantization import Quantizer
+    from flatquant.flat_linear import FlatQuantizedLinear
+
+    dtype = _kernel_scale_dtype()
+    for module in model.modules():
+        if not isinstance(module, FlatQuantizedLinear):
+            continue
+        if getattr(module, "_packed", False) and isinstance(module.linear, Linear4bit):
+            continue
+        linear = module.linear
+        weight_dtype = linear.weight.dtype if linear.weight is not None else dtype
+        packed = Linear4bit(
+            linear.in_features,
+            linear.out_features,
+            bias=linear.bias is not None,
+            dtype=weight_dtype,
+        )
+        module.linear = packed
+        if not hasattr(module, "a4"):
+            module.add_module("a4", Quantizer(lac=True))
+        module._packed = True
+        module._eval_mode = True
+    cast_linear4bit_kernel_dtypes(model)
+
+
+def restore_packed_eval_runtime(model) -> None:
+    """Restore eval-only Python flags that are not stored in state_dict.
+
+    Packed / reparameterized weights already fused T and diag(T). The adapter
+    must call transform.to_eval_mode() *before* load so matrix buffers exist,
+    then set these flags *after* load. Do not run reparameterize_model again
+    on Linear4bit weights.
+    """
+    from flatquant.flat_linear import FlatQuantizedLinear
+
+    for layer in model.model.layers:
+        attn = layer.self_attn
+        mlp = layer.mlp
+        attn._eval_mode = True
+        for trans in (
+            getattr(attn, "ln_trans", None),
+            getattr(attn, "o_trans", None),
+            getattr(attn, "kcache_trans", None),
+            getattr(attn, "vcache_trans", None),
+            getattr(mlp, "up_gate_trans", None),
+            getattr(mlp, "down_trans", None),
+        ):
+            if trans is None:
+                continue
+            if hasattr(trans, "use_diag"):
+                trans.use_diag = False
+            trans._eval_mode = True
+        for module in layer.modules():
+            if isinstance(module, FlatQuantizedLinear):
+                module._eval_mode = True
+
+
+def convert_model_to_w4a4(model, quantizers: dict) -> int:
+    """Pack RTN-quantized inner linears with official pack_i4 / Linear4bit."""
+    import torch
+    from deploy.nn.linear import Linear4bit
+    from deploy.nn.quantization import Quantizer
+    from flatquant.flat_linear import FlatQuantizedLinear
+
+    dtype = _kernel_scale_dtype()
+    converted = 0
+    named = dict(model.named_modules())
+    for name, module in named.items():
+        if not isinstance(module, FlatQuantizedLinear):
+            continue
+        key = f"{name}.linear"
+        quantizer = quantizers.get(key)
+        if quantizer is None:
+            raise RuntimeError(f"missing RTN quantizer for {key}")
+        scale = quantizer.scale.detach().reshape(-1, 1).to(dtype=dtype)
+        packed = Linear4bit.from_float(module.linear, weight_scales=scale)
+        module.linear = packed
+        a4 = Quantizer(lac=True)
+        act = module.act_quantizer
+        if hasattr(act, "clip_factor_a_max"):
+            with torch.no_grad():
+                a4.clip_factor_a_max.fill_(float(act.clip_factor_a_max.detach().reshape(-1)[0]))
+                a4.clip_factor_a_min.fill_(float(act.clip_factor_a_min.detach().reshape(-1)[0]))
+        module.add_module("a4", a4)
+        module._packed = True
+        module._eval_mode = True
+        converted += 1
+    if converted == 0:
+        raise RuntimeError("no FlatQuantizedLinear modules to pack")
+    cast_linear4bit_kernel_dtypes(model)
+    return converted
diff --git a/flatquant/utils.py b/flatquant/utils.py
index 6de9dbe..b23faac 100644
--- a/flatquant/utils.py
+++ b/flatquant/utils.py
@@ -52,7 +52,7 @@ def cleanup_memory(verbose=True) -> None:
 
 def distribute_model(model) -> None:
     """Distribute the model across available GPUs. NB: only implemented for Llama-2/3/Qwen-2."""
-    no_split_module_classes = ['LlamaDecoderLayer', 'Qwen2DecoderLayer']
+    no_split_module_classes = ['LlamaDecoderLayer', 'Qwen2DecoderLayer', 'Phi3DecoderLayer']
     max_memory = get_balanced_memory(model, no_split_module_classes=no_split_module_classes)
 
     device_map = infer_auto_device_map(model, max_memory=max_memory, no_split_module_classes=no_split_module_classes)
diff --git a/gptq_utils.py b/gptq_utils.py
index 714849f..c28f78d 100644
--- a/gptq_utils.py
+++ b/gptq_utils.py
@@ -200,12 +200,21 @@ def gptq_fwrd(model, dataloader, dev, args):
     position_ids = cache['position_ids']
 
     quantizers = {}
-    sequential = [
-                ['self_attn.k_proj.linear', 'self_attn.v_proj.linear', 'self_attn.q_proj.linear'],
-                ['self_attn.o_proj.linear'],
-                ['mlp.up_proj.linear', 'mlp.gate_proj.linear'],
-                ['mlp.down_proj.linear']
-            ]
+    first_attn = layers[0].self_attn
+    if hasattr(first_attn, "qkv_proj"):
+        sequential = [
+                    ['self_attn.qkv_proj.linear'],
+                    ['self_attn.o_proj.linear'],
+                    ['mlp.gate_up_proj.linear'],
+                    ['mlp.down_proj.linear']
+                ]
+    else:
+        sequential = [
+                    ['self_attn.k_proj.linear', 'self_attn.v_proj.linear', 'self_attn.q_proj.linear'],
+                    ['self_attn.o_proj.linear'],
+                    ['mlp.up_proj.linear', 'mlp.gate_proj.linear'],
+                    ['mlp.down_proj.linear']
+                ]
     # sequential = [
     #             ['self_attn.k_proj', 'self_attn.v_proj', 'self_attn.q_proj'],
     #             ['self_attn.o_proj'],
diff --git a/quant_agent_inference_adapter.py b/quant_agent_inference_adapter.py
new file mode 100644
index 0000000..7326b82
--- /dev/null
+++ b/quant_agent_inference_adapter.py
@@ -0,0 +1,66 @@
+"""Reload packed FlatQuant W4A4 weights (Linear4bit + transforms) from the artifact."""
+from __future__ import annotations
+
+import json
+import os
+import sys
+from argparse import Namespace
+
+import torch
+from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
+
+QUANT_AGENT_ADAPTER_API = 1
+
+
+def load_model_and_tokenizer(model_path, model_id, dtype, device, trust_remote_code):
+    method_repo = os.environ.get("QUANT_AGENT_METHOD_REPO")
+    if method_repo:
+        sys.path.insert(0, method_repo)
+    import flatquant.utils as utils
+    from flatquant.model_tools.phi3_utils import apply_flatquant_to_phi3
+    import flatquant.pack_utils as pack_utils
+
+    packed_path = os.path.join(model_path, "packed_w4a4.pt")
+    if not os.path.isfile(packed_path):
+        raise FileNotFoundError("packed W4A4 artifact missing: %s" % packed_path)
+    dtypes = {"float16": torch.float16, "bfloat16": torch.bfloat16, "auto": "auto"}
+    torch_dtype = dtypes.get(dtype, torch.float16)
+    tokenizer = AutoTokenizer.from_pretrained(
+        model_path, trust_remote_code=trust_remote_code, local_files_only=True
+    )
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+    config = AutoConfig.from_pretrained(
+        model_path, trust_remote_code=trust_remote_code, local_files_only=True
+    )
+    config._attn_implementation = "eager"
+    config._attn_implementation_internal = "eager"
+    utils.skip_initialization()
+    model = AutoModelForCausalLM.from_config(
+        config,
+        torch_dtype=torch_dtype if device == "cuda" else torch.float32,
+        trust_remote_code=trust_remote_code,
+        attn_implementation="eager",
+    )
+    model.seqlen = 2048
+    with open(os.path.join(model_path, "flatquant_args.json"), "r", encoding="utf-8") as handle:
+        args = Namespace(**json.load(handle))
+    model = apply_flatquant_to_phi3(args, model)
+    for layer in model.model.layers:
+        layer.self_attn.rep_matrix_only()
+        layer.mlp.rep_matrix_only()
+    pack_utils.prepare_packed_slots(model)
+    payload = torch.load(packed_path, map_location="cpu", weights_only=False)
+    missing, unexpected = model.load_state_dict(payload, strict=False)
+    packed_missing = [key for key in missing if "linear.weight" in key or key.endswith("a4.clip_factor_a_max")]
+    if packed_missing:
+        raise RuntimeError(
+            "packed_w4a4.pt did not restore Linear4bit weights: missing=%s unexpected=%s"
+            % (packed_missing[:12], list(unexpected)[:12])
+        )
+    pack_utils.cast_linear4bit_kernel_dtypes(model)
+    pack_utils.restore_packed_eval_runtime(model)
+    model.to(device)
+    pack_utils.cast_linear4bit_kernel_dtypes(model)
+    model.eval()
+    return model, tokenizer