diff --git a/DeepSeek_V3.pdf b/DeepSeek_V3.pdf
deleted file mode 100644
index 29dcd6a..0000000
Binary files a/DeepSeek_V3.pdf and /dev/null differ
diff --git a/figures/benchmark.png b/figures/benchmark.png
deleted file mode 100644
index d7ee0c0..0000000
Binary files a/figures/benchmark.png and /dev/null differ
diff --git a/figures/niah.png b/figures/niah.png
deleted file mode 100644
index 2a2e920..0000000
Binary files a/figures/niah.png and /dev/null differ
diff --git a/inference/models/__init__.py b/inference/models/__init__.py
new file mode 100644
index 0000000..c96277a
--- /dev/null
+++ b/inference/models/__init__.py
@@ -0,0 +1,15 @@
+from .config import ModelArgs
+from .attention import MLA
+from .moe import Gate, Expert, MoE
+from .linear import Linear, ColumnParallelLinear, RowParallelLinear
+
+__all__ = [
+    'ModelArgs',
+    'MLA',
+    'Gate',
+    'Expert', 
+    'MoE',
+    'Linear',
+    'ColumnParallelLinear',
+    'RowParallelLinear'
+]
\ No newline at end of file
diff --git a/inference/models/attention.py b/inference/models/attention.py
new file mode 100644
index 0000000..760e02a
--- /dev/null
+++ b/inference/models/attention.py
@@ -0,0 +1,25 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.distributed as dist
+from .config import ModelArgs
+from ..kernel import act_quant, weight_dequant, fp8_gemm
+
+class MLA(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.dim = args.dim
+        self.n_heads = args.n_heads
+        self.n_local_heads = args.n_heads // dist.get_world_size() if dist.is_initialized() else args.n_heads
+        self.q_lora_rank = args.q_lora_rank
+        self.kv_lora_rank = args.kv_lora_rank
+        self.qk_nope_head_dim = args.qk_nope_head_dim
+        self.qk_rope_head_dim = args.qk_rope_head_dim
+        self.qk_head_dim = args.qk_nope_head_dim + args.qk_rope_head_dim
+        self.v_head_dim = args.v_head_dim
+        
+        # Initialize components (implementation from original MLA class)
+        # ... (rest of the MLA implementation)
+
+    def forward(self, x: torch.Tensor, start_pos: int, freqs_cis: torch.Tensor, mask: Optional[torch.Tensor]):
+        # ... (MLA forward implementation)
\ No newline at end of file
diff --git a/inference/models/config.py b/inference/models/config.py
new file mode 100644
index 0000000..2ca3dbd
--- /dev/null
+++ b/inference/models/config.py
@@ -0,0 +1,36 @@
+from dataclasses import dataclass
+from typing import Literal
+
+@dataclass
+class ModelArgs:
+    max_batch_size: int = 8
+    max_seq_len: int = 4096 * 4
+    dtype: Literal["bf16", "fp8"] = "bf16"
+    vocab_size: int = 102400
+    dim: int = 2048
+    inter_dim: int = 10944
+    moe_inter_dim: int = 1408
+    n_layers: int = 27
+    n_dense_layers: int = 1
+    n_heads: int = 16
+    # moe
+    n_routed_experts: int = 64
+    n_shared_experts: int = 2
+    n_activated_experts: int = 6
+    n_expert_groups: int = 1
+    n_limited_groups: int = 1
+    score_func: Literal["softmax", "sigmoid"] = "softmax"
+    route_scale: float = 1.
+    # mla
+    q_lora_rank: int = 0
+    kv_lora_rank: int = 512
+    qk_nope_head_dim: int = 128
+    qk_rope_head_dim: int = 64
+    v_head_dim: int = 128
+    # yarn
+    original_seq_len: int = 4096
+    rope_theta: float = 10000.0
+    rope_factor: float = 40
+    beta_fast: int = 32
+    beta_slow: int = 1
+    mscale: float = 1.
\ No newline at end of file
diff --git a/inference/models/linear.py b/inference/models/linear.py
new file mode 100644
index 0000000..c7ac7a6
--- /dev/null
+++ b/inference/models/linear.py
@@ -0,0 +1,28 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.distributed as dist
+from ..kernel import act_quant, weight_dequant, fp8_gemm
+
+class Linear(nn.Module):
+    dtype = torch.bfloat16
+
+    def __init__(self, in_features: int, out_features: int, bias: bool = False, dtype = None):
+        # ... (Linear implementation)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # ... (Linear forward implementation)
+
+class ColumnParallelLinear(Linear):
+    def __init__(self, in_features: int, out_features: int, bias: bool = False, dtype = None):
+        # ... (ColumnParallelLinear implementation)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # ... (ColumnParallelLinear forward implementation)
+
+class RowParallelLinear(Linear):
+    def __init__(self, in_features: int, out_features: int, bias: bool = False, dtype = None):
+        # ... (RowParallelLinear implementation)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # ... (RowParallelLinear forward implementation)
\ No newline at end of file
diff --git a/inference/models/moe.py b/inference/models/moe.py
new file mode 100644
index 0000000..c715f3b
--- /dev/null
+++ b/inference/models/moe.py
@@ -0,0 +1,30 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.distributed as dist
+from .config import ModelArgs
+from .linear import Linear, ColumnParallelLinear, RowParallelLinear
+
+class Gate(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        # ... (Gate implementation)
+
+    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        # ... (Gate forward implementation)
+
+class Expert(nn.Module):
+    def __init__(self, dim: int, inter_dim: int):
+        super().__init__()
+        # ... (Expert implementation)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # ... (Expert forward implementation)
+
+class MoE(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        # ... (MoE implementation)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # ... (MoE forward implementation)
\ No newline at end of file
diff --git a/package-lock.json b/package-lock.json
new file mode 100644
index 0000000..af2503e
--- /dev/null
+++ b/package-lock.json
@@ -0,0 +1,6 @@
+{
+  "name": "project",
+  "lockfileVersion": 3,
+  "requires": true,
+  "packages": {}
+}