update

src/backend/run_eval_suite.py

@@ -17,16 +17,22 @@ def process_results_decorator(func):
         end_to_end_time = sum([r[1] for r in results]) / len(results)
         prefilling_time = sum([r[2] for r in results]) / len(results)
         decoding_throughput = sum([r[3] for r in results]) / len(results)
-        mfu = sum([r[4] for r in results]) / len(results)
-        mbu = sum([r[5] for r in results]) / len(results)
+        decoding_mfu = sum([r[4] for r in results]) / len(results)
+        decoding_mbu = sum([r[5] for r in results]) / len(results)
+        prefill_throughput = sum([r[6] for r in results]) / len(results)
+        prefill_mfu = sum([r[7] for r in results]) / len(results)
+        prefill_mbu = sum([r[8] for r in results]) / len(results)
         # print(f"end_to_end_time: {end_to_end_time}, prefilling_time: {prefilling_time}, decoding_throughput: {decoding_throughput}")
 
         result_dict = func(self, doc, processed_results, *args, **kwargs)
         result_dict["end_to_end_time"] = end_to_end_time
         result_dict["prefilling_time"] = prefilling_time
         result_dict["decoding_throughput"] = decoding_throughput
-        result_dict["mfu"] = mfu
-        result_dict["mbu"] = mbu
+        result_dict["decoding_mfu"] = decoding_mfu
+        result_dict["decoding_mbu"] = decoding_mbu
+        result_dict["prefill_throughput"] = prefill_throughput
+        result_dict["prefill_mfu"] = prefill_mfu
+        result_dict["prefill_mbu"] = prefill_mbu
         return result_dict
     return wrapper
 ConfigurableTask.process_results = process_results_decorator(orig_process_results)
@@ -37,8 +43,11 @@ def aggregation_decorator(func):
         aggregation_list["end_to_end_time"] = mean
         aggregation_list["prefilling_time"] = mean
         aggregation_list["decoding_throughput"] = mean
-        aggregation_list["mfu"] = mean
-        aggregation_list["mbu"] = mean
+        aggregation_list["decoding_mfu"] = mean
+        aggregation_list["decoding_mbu"] = mean
+        aggregation_list["prefill_throughput"] = mean
+        aggregation_list["prefill_mfu"] = mean
+        aggregation_list["prefill_mbu"] = mean
         return aggregation_list
     return wrapper
 ConfigurableTask.aggregation = aggregation_decorator(orig_aggregation)
@@ -49,8 +58,11 @@ def higher_is_better_decorator(func):
         higher_is_better_dict["end_to_end_time"] = False
         higher_is_better_dict["prefilling_time"] = False
         higher_is_better_dict["decoding_throughput"] = True
-        higher_is_better_dict["mfu"] = True
-        higher_is_better_dict["mbu"] = True
+        higher_is_better_dict["decoding_mfu"] = True
+        higher_is_better_dict["decoding_mbu"] = True
+        higher_is_better_dict["prefill_throughput"] = True
+        higher_is_better_dict["prefill_mfu"] = True
+        higher_is_better_dict["prefill_mbu"] = True
         return higher_is_better_dict
     return wrapper
 ConfigurableTask.higher_is_better = higher_is_better_decorator(orig_higher_is_better)
@@ -65,6 +77,8 @@ from src.backend.tasks.selfcheckgpt.task import SelfCheckGPT
 
 from src.backend.huggingface_generate_until import HFLMwithChatTemplate
 from src.backend.moe_infinity import MoEHFLM
+from src.backend.vllm import VLLM_MOE
+from src.backend.sglang import SGLangMoE
 
 def run_evaluation(
     eval_request: EvalRequest,

src/backend/tasks/measurement_task_utils.py

@@ -12,8 +12,12 @@ def process_results_decorator(func):
         end_to_end_time = sum([r[1] for r in results]) / len(results)
         prefilling_time = sum([r[2] for r in results]) / len(results)
         decoding_throughput = sum([r[3] for r in results]) / len(results)
-        mfu = sum([r[4] for r in results]) / len(results)
-        mbu = sum([r[5] for r in results]) / len(results)
+        decoding_mfu = sum([r[4] for r in results]) / len(results)
+        decoding_mbu = sum([r[5] for r in results]) / len(results)
+        prefill_throughput = sum([r[6] for r in results]) / len(results)
+        prefill_mfu = sum([r[7] for r in results]) / len(results)
+        prefill_mbu = sum([r[8] for r in results]) / len(results)
+
         
         # print(f"end_to_end_time: {end_to_end_time}, prefilling_time: {prefilling_time}, decoding_throughput: {decoding_throughput}")
 
@@ -22,8 +26,11 @@ def process_results_decorator(func):
         result_dict["end_to_end_time"] = end_to_end_time
         result_dict["prefilling_time"] = prefilling_time
         result_dict["decoding_throughput"] = decoding_throughput
-        result_dict["mfu"] = mfu
-        result_dict["mbu"] = mbu
+        result_dict["decoding_mfu"] = decoding_mfu
+        result_dict["decoding_mbu"] = decoding_mbu
+        result_dict["prefill_throughput"] = prefill_throughput
+        result_dict["prefill_mfu"] = prefill_mfu
+        result_dict["prefill_mbu"] = prefill_mbu
         return result_dict
     return wrapper
 
@@ -35,8 +42,11 @@ def aggregation_decorator(func):
         aggregation_list["end_to_end_time"] = mean
         aggregation_list["prefilling_time"] = mean
         aggregation_list["decoding_throughput"] = mean
-        aggregation_list["mfu"] = mean
-        aggregation_list["mbu"] = mean
+        aggregation_list["decoding_mfu"] = mean
+        aggregation_list["decoding_mbu"] = mean
+        aggregation_list["prefill_throughput"] = mean
+        aggregation_list["prefill_mfu"] = mean
+        aggregation_list["prefill_mbu"] = mean
         return aggregation_list
     return wrapper
 
@@ -48,8 +58,11 @@ def higher_is_better_decorator(func):
         higher_is_better_dict["end_to_end_time"] = False
         higher_is_better_dict["prefilling_time"] = False
         higher_is_better_dict["decoding_throughput"] = True
-        higher_is_better_dict["mfu"] = True
-        higher_is_better_dict["mbu"] = True
+        higher_is_better_dict["decoding_mfu"] = True
+        higher_is_better_dict["decoding_mbu"] = True
+        higher_is_better_dict["prefill_throughput"] = True
+        higher_is_better_dict["prefill_mfu"] = True
+        higher_is_better_dict["prefill_mbu"] = True
         return higher_is_better_dict
     return wrapper
 

src/display/utils.py

@@ -9,25 +9,32 @@ def fields(raw_class):
 
 E2Es = "E2E(s)" #"End-to-end time (s)"
 PREs = "PRE(s)" #"Prefilling time (s)"
-TS = "T/s" #Decoding throughput (tok/s)
+TS = "Decoding T/s" #Decoding throughput (tok/s)
+PTS = "Prefill T/s" #Prefill throughput (tok/s)
 InFrame = "Method" #"Inference framework"
 MULTIPLE_CHOICEs = ["mmlu"]
 
+
 GPU_TEMP = 'Temp(C)'
 GPU_Power = 'Power(W)'
 GPU_Mem = 'Mem(G)'
 GPU_Name = "GPU"
 GPU_Util = 'Util(%)'
-MFU = 'S-MFU(%)'
-MBU = 'S-MBU(%)'
+DSMFU = 'Decoding S-MFU(%)'
+DSMBU = 'Decoding S-MBU(%)'
+PSMFU = 'Prefill S-MFU(%)'
+PSMBU = 'Prefill S-MBU(%)'
 BATCH_SIZE = 'bs'
 PRECISION = "Precision"
 system_metrics_to_name_map = {
     "end_to_end_time": f"{E2Es}",
     "prefilling_time": f"{PREs}",
     "decoding_throughput": f"{TS}",
-    "mfu": f"{MFU}",
-    "mbu": f"{MBU}"
+    "decoding_mfu": f"{DSMFU}",
+    "decoding_mbu": f"{DSMBU}",
+    "prefill_throughput": f"{PTS}",
+    "prefill_mfu": f"{PSMFU}",
+    "prefill_mbu": f"{PSMBU}",
 }
 
 gpu_metrics_to_name_map = {
@@ -78,10 +85,11 @@ class Tasks(Enum):
 
     # # XXX include me back at some point
     # selfcheck = Task("selfcheckgpt", "max-selfcheckgpt", "SelfCheckGPT")
-    mmlu = Task("mmlu", "acc", "MMLU") #MMLU/Acc (5-shot)
+    # selfcheck = Task("selfcheckgpt", "max-selfcheckgpt", "SelfCheckGPT")
     gsm8k = Task("gsm8k_custom", "em", "GSM8K") #GSM8K/EM (5-shot)
     # gsm8k_cot = Task("gsm8k_cot", "em", "GSM8K COT") #GSM8K COT/EM (5-shot)
     arena_hard = Task("arena_hard", "score", "Arena Hard") #Arena Hard/Score
+    mmlu = Task("mmlu", "acc", "MMLU") #MMLU/Acc (5-shot)
 
 
 # These classes are for user facing column names,
@@ -121,8 +129,14 @@ for task in Tasks:
         continue
     auto_eval_column_dict.append([f"{task.name}_prefilling_time", ColumnContent, ColumnContent(f"{task.value.col_name} {PREs}", "number", False, hidden=True)])
     auto_eval_column_dict.append([f"{task.name}_decoding_throughput", ColumnContent, ColumnContent(f"{task.value.col_name} {TS}", "number", True, hidden=True)])
-    auto_eval_column_dict.append([f"{task.name}_mbu", ColumnContent, ColumnContent(f"{task.value.col_name} {MBU}", "number", True, hidden=True)])
-    auto_eval_column_dict.append([f"{task.name}_mfu", ColumnContent, ColumnContent(f"{task.value.col_name} {MFU}", "number", True, hidden=True)])
+    # if task.value.benchmark != "gsm8k_custom":
+    #     continue
+    auto_eval_column_dict.append([f"{task.name}_decoding_mbu", ColumnContent, ColumnContent(f"{task.value.col_name} {DSMBU}", "number", True, hidden=True)])
+    auto_eval_column_dict.append([f"{task.name}_decoding_mfu", ColumnContent, ColumnContent(f"{task.value.col_name} {DSMFU}", "number", True, hidden=True)])
+    auto_eval_column_dict.append([f"{task.name}_prefill_throughput", ColumnContent, ColumnContent(f"{task.value.col_name} {PTS}", "number", True, hidden=True)])
+    auto_eval_column_dict.append([f"{task.name}_prefill_mbu", ColumnContent, ColumnContent(f"{task.value.col_name} {PSMBU}", "number", True, hidden=True)])
+    auto_eval_column_dict.append([f"{task.name}_prefill_mfu", ColumnContent, ColumnContent(f"{task.value.col_name} {PSMFU}", "number", True, hidden=True)])
+    
 
 
 # Model information
@@ -187,8 +201,9 @@ class InferenceFramework(Enum):
     # MoE_Infinity = ModelDetails("moe-infinity")
     HF_Chat = ModelDetails("hf-chat")
     VLLM = ModelDetails("vllm_moe")
-    TRTLLM = ModelDetails("tensorrt_llm")
     VLLM_FIX = ModelDetails("vllm_moe_fixbs")
+    TRTLLM = ModelDetails("tensorrt_llm")
+    SGLANG = ModelDetails("sglang")
     Unknown = ModelDetails("?")
 
     def to_str(self):
@@ -206,10 +221,13 @@ class InferenceFramework(Enum):
             return InferenceFramework.VLLM
         if inference_framework in ["vllm_moe_fixbs"]:
             return InferenceFramework.VLLM_FIX
+        if inference_framework in ["sglang"]:
+            return InferenceFramework.SGLANG
         return InferenceFramework.Unknown
 
 class GPUType(Enum):
     A100_sxm = ModelDetails("NVIDIA-A100-SXM4-80GB")
+    A100_sxm4 = ModelDetails("NVIDIA-A100-SMX4-80GB")
     A100_pcie = ModelDetails("NVIDIA-A100-PCIe-80GB")
     Unknown = ModelDetails("?")
 

src/utils.py

@@ -4,6 +4,8 @@ import subprocess
 import re
 import os
 import GPUtil
+from transformers import AutoConfig
+from typing import List
 
 try:
     from src.display.utils import GPU_TEMP, GPU_Mem, GPU_Power, GPU_Util, GPU_Name
@@ -12,44 +14,63 @@ except:
     from display.utils import GPU_TEMP, GPU_Mem, GPU_Power, GPU_Util, GPU_Name
     
 MEM_BW_DICT ={
-    "NVIDIA-A100-PCIe-80GB": 1935,
-    "NVIDIA-A100-SXM-80GB": 2039,
-    "NVIDIA-H100-PCIe-80GB": 2039,
-    "NVIDIA-RTX-A5000-24GB": 768
+    "NVIDIA-A100-PCIe-80GB": 1935e9,
+    "NVIDIA-A100-SXM4-80GB": 2039e9,
+    "NVIDIA-H100-PCIe-80GB": 2039e9,
+    "NVIDIA-RTX-A5000-24GB": 768e9,
+    "NVIDIA-RTX-A6000-48GB": 768e9,
 }
 
 PEAK_FLOPS_DICT = {
     "float32":{
         "NVIDIA-A100-PCIe-80GB": 312e12,
-        "NVIDIA-A100-SXM-80GB": 312e12,
+        "NVIDIA-A100-SXM4-80GB": 312e12,
         "NVIDIA-H100-PCIe-80GB": 756e12,
-        "NVIDIA-RTX-A5000-24GB": 222.2e12
+        "NVIDIA-RTX-A5000-24GB": 222.2e12,
+        "NVIDIA-RTX-A6000-48GB": 309.7e12
     },
     "float16":{
         "NVIDIA-A100-PCIe-80GB": 624e12,
-        "NVIDIA-A100-SXM-80GB": 624e12,
+        "NVIDIA-A100-SXM4-80GB": 624e12,
         "NVIDIA-H100-PCIe-80GB": 1513e12,
-        "NVIDIA-RTX-A5000-24GB": 444.4e12
+        "NVIDIA-RTX-A5000-24GB": 222.2e12,
+        "NVIDIA-RTX-A6000-48GB": 309.7e12
     },
     "bfloat16":{
         "NVIDIA-A100-PCIe-80GB": 624e12,
-        "NVIDIA-A100-SXM-80GB": 624e12,
+        "NVIDIA-A100-SXM4-80GB": 624e12,
         "NVIDIA-H100-PCIe-80GB": 1513e12,
-        "NVIDIA-RTX-A5000-24GB": 444.4e12
+        "NVIDIA-RTX-A5000-24GB": 222.2e12,
+        "NVIDIA-RTX-A6000-48GB": 309.7e12
     },
-    "8bit":{
+    "int8":{
         "NVIDIA-A100-PCIe-80GB": 1248e12,
-        "NVIDIA-A100-SXM-80GB": 1248e12,
+        "NVIDIA-A100-SXM4-80GB": 1248e12,
         "NVIDIA-H100-PCIe-80GB": 3026e12,
-        "NVIDIA-RTX-A5000-24GB": 889e12
+        "NVIDIA-RTX-A5000-24GB": 222.2e12,
+        "NVIDIA-RTX-A6000-48GB": 309.7e12
     },
-    "4bit": {
-        "NVIDIA-A100-PCIe-80GB": 2496e12,
-        "NVIDIA-A100-SXM-80GB": 2496e12,
-        "NVIDIA-H100-PCIe-80GB": 6052e12,
-        "NVIDIA-RTX-A5000-24GB": 1778e12
+    "fp8":{
+        "NVIDIA-A100-PCIe-80GB": 1248e12,
+        "NVIDIA-A100-SXM4-80GB": 1248e12,
+        "NVIDIA-H100-PCIe-80GB": 3026e12,
+        "NVIDIA-RTX-A5000-24GB": 0,
+        "NVIDIA-RTX-A6000-48GB": 0
+    },
+    "fp4": {
+        "NVIDIA-A100-PCIe-80GB": 1248e12,
+        "NVIDIA-A100-SXM4-80GB": 1248e12,
+        "NVIDIA-H100-PCIe-80GB": 3026e12,
+        "NVIDIA-RTX-A5000-24GB": 0,
+        "NVIDIA-RTX-A6000-48GB": 0
+    },
+    "int4": {
+        "NVIDIA-A100-PCIe-80GB": 1248e12,
+        "NVIDIA-A100-SXM4-80GB": 1248e12,
+        "NVIDIA-H100-PCIe-80GB": 3026e12,
+        "NVIDIA-RTX-A5000-24GB": 222.2e12,
+        "NVIDIA-RTX-A6000-48GB": 309.7e12
     }
-
 }
 
 def my_snapshot_download(repo_id, revision, local_dir, repo_type, max_workers):
@@ -97,7 +118,7 @@ def parse_nvidia_smi():
     # print(f"gpu_indices: {gpu_indices}")
     gpu_stats = []
 
-    gpu_info_pattern = re.compile(r'(\d+)C\s+P\d+\s+(\d+)W / \d+W\s+\|\s+(\d+)MiB / \d+MiB\s+\|\s+(\d+)%')
+    gpu_info_pattern = re.compile(r'(\d+)C\s+P\d+\s+(\d+)W\s*/\s*\d+W\s*\|\s*(\d+)MiB\s*/\s*\d+MiB\s*\|\s*(\d+)%')
     # gpu_name_pattern = re.compile(r'NVIDIA\s+([\w\s]+\d+(?:\s*GB)?)')
     gpu_name_pattern = re.compile(r'NVIDIA\s+(RTX\s+)?([A-Z0-9]+)')
 
@@ -195,17 +216,790 @@ def get_peak_bw(gpu_name):
 def get_peak_flops(gpu_name, precision):
     return PEAK_FLOPS_DICT[precision][gpu_name]
 
-def transfer_precision2bytes(precision):
-    if precision == "float32":
-        return 4
-    elif precision in ["float16", "bfloat16"]:
-        return 2
-    elif precision == "8bit":
-        return 1
-    elif precision == "4bit":
-        return 0.5
+def _calculate_batch_metrics(outputs, decoding_tp, n_layers, d_model, 
+                                n_attn_heads, d_head, n_kv_heads, n_experts_per_tok, d_ff, 
+                                avg_activated_experts, hf_config, num_gpus, model_name, 
+                                used_dtype, batch_size, precision):
+    """Calculate metrics for a batch of outputs"""
+    gpu_type = get_gpu_details()
+    hardware_specs = {
+        "peak_bandwidth_tb": get_peak_bw(gpu_type) / 1e12,
+        "peak_flops_tf": get_peak_flops(gpu_type, precision=used_dtype) / 1e12,
+    }
+    kvs = []
+    true_kvs = []
+    attn_score = []
+    
+    # Calculate KV sizes
+    per_token_kv_size = 2 * n_layers * d_head * n_kv_heads  # Default calculation
+    
+    if "DeepSeek" in model_name:
+        if hasattr(hf_config, "kv_lora_rank") and hasattr(hf_config, "qk_rope_head_dim"):
+            per_token_kv_size = n_layers * (hf_config.kv_lora_rank + hf_config.qk_rope_head_dim)
+    
+    # Process each output
+    for x in outputs:
+        output_len = len(x.outputs[0].token_ids)
+        context_prefill_size = len(x.prompt_token_ids)
+        
+        # Calculate attention scores
+        if "DeepSeek" in model_name and hasattr(hf_config, "qk_rope_head_dim") and hasattr(hf_config, "qk_nope_head_dim") and hasattr(hf_config, "v_head_dim"):
+            q_head_dim = hf_config.qk_rope_head_dim + hf_config.qk_nope_head_dim
+            origin_per_token_k_state_size = n_layers * n_attn_heads * q_head_dim
+            origin_per_token_v_state_size = n_layers * n_attn_heads * hf_config.v_head_dim
+            attention_score = context_prefill_size * origin_per_token_k_state_size + (output_len - 1) * origin_per_token_k_state_size / 2
+            attention_score += context_prefill_size * origin_per_token_v_state_size + (output_len - 1) * origin_per_token_v_state_size / 2
+            attention_score = attention_score / 1e12
+        else:
+            origin_per_token_kv_states_size = n_layers * n_attn_heads * d_head
+            attention_score = context_prefill_size * origin_per_token_kv_states_size + (output_len - 1) * origin_per_token_kv_states_size / 2
+            attention_score = attention_score * 2 / 1e12
+        
+        # Store attention scores and KV sizes
+        attn_score.append(attention_score)
+        kv_size = context_prefill_size * per_token_kv_size + (output_len - 1) * per_token_kv_size / 2
+        kv_size = kv_size / 1e12
+        true_kv = (context_prefill_size * per_token_kv_size + output_len * per_token_kv_size) / 1e12 * 1e3
+        kvs.append(kv_size)
+        true_kvs.append(true_kv)
+    
+    # Calculate aggregate values
+    kv_size = sum(kvs)
+    true_kv_size = sum(true_kvs) * 1e3
+    attention_score = sum(attn_score) / len(attn_score)
+    
+    # Calculate attention size per token
+    if "DeepSeek" in model_name and hasattr(hf_config, "qk_rope_head_dim") and hasattr(hf_config, "qk_nope_head_dim") and hasattr(hf_config, "v_head_dim") and hasattr(hf_config, "kv_lora_rank"):
+        q_head_dim = hf_config.qk_rope_head_dim + hf_config.qk_nope_head_dim
+        if not hasattr(hf_config, "q_lora_rank") or not hf_config.q_lora_rank:
+            attention_size_per_token = (d_model * n_attn_heads * q_head_dim) + \
+                (d_model * (hf_config.kv_lora_rank + hf_config.qk_rope_head_dim)) + \
+                    (hf_config.kv_lora_rank * n_attn_heads * (q_head_dim - hf_config.qk_rope_head_dim + hf_config.v_head_dim)) + \
+                        (hf_config.v_head_dim * n_attn_heads * d_model)
+            attention_size_per_token = attention_size_per_token / 1e12
+        else:
+            attention_size_per_token = (d_model * hf_config.q_lora_rank) + \
+                (hf_config.q_lora_rank * n_attn_heads * q_head_dim) + \
+                    (d_model * (hf_config.kv_lora_rank + hf_config.qk_rope_head_dim)) + \
+                        (hf_config.kv_lora_rank * n_attn_heads * (q_head_dim - hf_config.qk_rope_head_dim + hf_config.v_head_dim)) + \
+                            (hf_config.v_head_dim * n_attn_heads * d_model)
+            attention_size_per_token = attention_size_per_token / 1e12
+    else:
+        attention_size_per_token = d_model * (n_attn_heads * d_head + n_kv_heads * d_head * 2) + n_attn_heads * d_head * d_model
+        attention_size_per_token = attention_size_per_token / 1e12
+    
+    # Calculate expert sizes
+    expert_size = d_ff * 3 * d_model / 1e12
+    shared_experts_size_total = 0
+    deepseek_dense_ffn_size = 0
+    deepseek_sparse_layer_num = 0
+    
+    if "Qwen" in model_name and hasattr(hf_config, "moe_intermediate_size") and hasattr(hf_config, "shared_expert_intermediate_size"):
+        d_ff = hf_config.moe_intermediate_size
+        d_ff_share = hf_config.shared_expert_intermediate_size
+        shared_experts_size = d_ff_share * 3 * d_model
+        expert_size = d_ff * 3 * d_model
+        shared_experts_size_total = shared_experts_size / 1e12
+        expert_size = expert_size / 1e12
+    elif "Qwen3" in model_name and hasattr(hf_config, "moe_intermediate_size"):
+        d_ff = hf_config.moe_intermediate_size
+        expert_size = d_ff * 3 * d_model
+        expert_size = expert_size / 1e12
+    elif "DeepSeek" in model_name and hasattr(hf_config, "moe_intermediate_size") and hasattr(hf_config, "intermediate_size") and hasattr(hf_config, "first_k_dense_replace"):
+        d_ff = hf_config.moe_intermediate_size
+        d_ff_dense = hf_config.intermediate_size
+        deepseek_num_dense_layer = hf_config.first_k_dense_replace
+        shared_experts_size = d_ff * 3 * d_model
+        expert_size = d_ff * 3 * d_model
+        shared_experts = 2
+        shared_experts_size_total = shared_experts_size * shared_experts / 1e12
+        expert_size = expert_size / 1e12
+        deepseek_sparse_layer_num = n_layers - deepseek_num_dense_layer
+        deepseek_dense_ffn_size = d_ff_dense * 3 * d_model / 1e12
+    
+    # Calculate S-MBU and S-MFU
+    if "Qwen" in model_name and not "Qwen3" in model_name:
+        smbu = ((n_layers*(avg_activated_experts * expert_size + shared_experts_size_total + attention_size_per_token) + 
+                kv_size) * precision/ (batch_size / decoding_tp)) / (num_gpus * hardware_specs['peak_bandwidth_tb'])
+        smfu = (n_layers * (attention_size_per_token + n_experts_per_tok * expert_size + shared_experts_size_total) + attention_score) \
+            * 2 * decoding_tp / (num_gpus * hardware_specs['peak_flops_tf'] / 2)
+    elif "Qwen3" in model_name:
+        smbu = ((n_layers * (avg_activated_experts * expert_size + attention_size_per_token) + 
+                kv_size) * precision/ (batch_size / decoding_tp)) / (num_gpus * hardware_specs['peak_bandwidth_tb'])
+        smfu = (n_layers * (attention_size_per_token + n_experts_per_tok * expert_size) + attention_score) \
+            * 2 * decoding_tp / (num_gpus * hardware_specs['peak_flops_tf'] / 2)
+    elif "DeepSeek" in model_name:
+        smbu = ((n_layers * attention_size_per_token + deepseek_sparse_layer_num * \
+                (avg_activated_experts * expert_size + shared_experts_size_total) + \
+                deepseek_num_dense_layer * deepseek_dense_ffn_size + \
+                kv_size) * precision/ (batch_size / decoding_tp)) / (num_gpus * hardware_specs['peak_bandwidth_tb'])
+        smfu = (n_layers * attention_size_per_token + deepseek_sparse_layer_num * \
+                (n_experts_per_tok * expert_size + shared_experts_size_total) + \
+                deepseek_num_dense_layer * deepseek_dense_ffn_size + attention_score) \
+                * 2 * decoding_tp / (num_gpus * hardware_specs['peak_flops_tf'] / 2)
+    else:
+        smbu = ((n_layers*(avg_activated_experts * expert_size + attention_size_per_token) + 
+                kv_size) * precision/ (batch_size / decoding_tp) ) / (num_gpus * hardware_specs['peak_bandwidth_tb'])
+        smfu = (n_layers * (attention_size_per_token + n_experts_per_tok * expert_size) + attention_score) \
+            * 2 * decoding_tp / (num_gpus * hardware_specs['peak_flops_tf'] / 2)
+    
+    return {
+        'smbu': smbu,
+        'smfu': smfu,
+        'kv_size': true_kv_size,
+        'decoding_throughput': decoding_tp
+    }
+
+def _calculate_batch_metrics_sglang(outputs, decoding_tp, n_layers, d_model, 
+                                n_attn_heads, d_head, n_kv_heads, n_experts_per_tok, d_ff, 
+                                avg_activated_experts, hf_config, num_gpus, model_name, 
+                                used_dtype, batch_size, precision, ttft=None, prefill_tp=None):
+    """Calculate metrics for a batch of outputs"""
+    # Initialize hardware specs and output lists
+    hardware_specs = _get_hardware_specs(used_dtype)
+    output_data = _extract_output_data(outputs)
+    
+    # Calculate model-specific sizes
+    per_token_kv_size = _calculate_kv_size(model_name, hf_config, n_layers, d_head, n_kv_heads)
+    attention_size_per_token = _calculate_attention_size(model_name, hf_config, d_model, n_attn_heads, d_head, n_kv_heads)
+    expert_config = _calculate_expert_config(model_name, hf_config, d_ff, d_model, n_layers)
+    
+    # Process outputs and calculate metrics
+    metrics_data = _process_outputs(output_data, per_token_kv_size, attention_size_per_token, 
+                                  model_name, hf_config, n_layers, n_attn_heads, d_head)
+
+    # Calculate throughput metrics
+    if ttft is None or prefill_tp is None:
+        ttft, prefill_tp = _calculate_throughput_metrics(batch_size, output_data['prefill_lengths'],
+                                                       output_data['max_duration'])
+
+    
+    # Calculate S-MBU and S-MFU
+    smbu_smfu_metrics = _calculate_smbu_smfu(model_name, n_layers, attention_size_per_token,
+                                           expert_config, avg_activated_experts, metrics_data,
+                                           hardware_specs, num_gpus, precision, ttft, prefill_tp,
+                                           batch_size, decoding_tp)
+    
+    return {
+        'prefill_smbu': smbu_smfu_metrics['prefill_smbu'],
+        'prefill_smfu': smbu_smfu_metrics['prefill_smfu'],
+        'decoding_smbu': smbu_smfu_metrics['decoding_smbu'],
+        'decoding_smfu': smbu_smfu_metrics['decoding_smfu'],
+        'kv_size': metrics_data['true_kv_size'],
+        'decoding_throughput': decoding_tp,
+        'prefill_tp': prefill_tp,
+        'ttft': ttft
+    }
+
+
+def _get_hardware_specs(used_dtype):
+    """Get hardware specifications"""
+    gpu_type = get_gpu_details()
+    return {
+        "peak_bandwidth_tb": get_peak_bw(gpu_type) / 1e12,
+        "peak_flops_tf": get_peak_flops(gpu_type, precision=used_dtype) / 1e12,
+    }
+
+
+def _extract_output_data(outputs):
+    """Extract relevant data from outputs"""
+    prefill_lengths = []
+    output_lengths = []
+    max_duration = 0.0
+    
+    for x in outputs:
+        output_lengths.append(x['meta_info']['completion_tokens'])
+        prefill_lengths.append(x['meta_info']['prompt_tokens'])
+        max_duration = max(max_duration, x['meta_info']['e2e_latency'])
+    
+    return {
+        'prefill_lengths': prefill_lengths,
+        'output_lengths': output_lengths,
+        'max_duration': max_duration
+    }
+
+
+def _calculate_kv_size(model_name, hf_config, n_layers, d_head, n_kv_heads):
+    """Calculate per-token KV size based on model type"""
+    if "DeepSeek" in model_name and hasattr(hf_config, "kv_lora_rank") and hasattr(hf_config, "qk_rope_head_dim"):
+        return n_layers * (hf_config.kv_lora_rank + hf_config.qk_rope_head_dim)
+    return 2 * n_layers * d_head * n_kv_heads
+
+
+def _calculate_attention_size(model_name, hf_config, d_model, n_attn_heads, d_head, n_kv_heads):
+    """Calculate attention size per token based on model type"""
+    if ("DeepSeek" in model_name and 
+        hasattr(hf_config, "qk_rope_head_dim") and 
+        hasattr(hf_config, "qk_nope_head_dim") and 
+        hasattr(hf_config, "v_head_dim") and 
+        hasattr(hf_config, "kv_lora_rank")):
+        
+        return _calculate_deepseek_attention_size(hf_config, d_model, n_attn_heads)
+    
+    return (d_model * (n_attn_heads * d_head + n_kv_heads * d_head * 2) + 
+            n_attn_heads * d_head * d_model) / 1e12
+
+
+def _calculate_deepseek_attention_size(hf_config, d_model, n_attn_heads):
+    """Calculate DeepSeek-specific attention size"""
+    q_head_dim = hf_config.qk_rope_head_dim + hf_config.qk_nope_head_dim
+    
+    base_size = ((d_model * (hf_config.kv_lora_rank + hf_config.qk_rope_head_dim)) +
+                (hf_config.kv_lora_rank * n_attn_heads * 
+                 (q_head_dim - hf_config.qk_rope_head_dim + hf_config.v_head_dim)) +
+                (hf_config.v_head_dim * n_attn_heads * d_model))
+    
+    if hasattr(hf_config, "q_lora_rank") and hf_config.q_lora_rank:
+        q_size = (d_model * hf_config.q_lora_rank + 
+                 hf_config.q_lora_rank * n_attn_heads * q_head_dim)
+    else:
+        q_size = d_model * n_attn_heads * q_head_dim
+    
+    return (base_size + q_size) / 1e12
+
+
+def _calculate_expert_config(model_name, hf_config, d_ff, d_model, n_layers):
+    """Calculate expert configuration based on model type"""
+    config = {
+        'expert_size': d_ff * 3 * d_model / 1e12,
+        'shared_experts_size_total': 0,
+        'deepseek_dense_ffn_size': 0,
+        'deepseek_sparse_layer_num': 0,
+        'deepseek_num_dense_layer': 0
+    }
+    
+    if "Qwen" in model_name and not "Qwen3" in model_name:
+        config.update(_get_qwen_expert_config(hf_config, d_model))
+    elif "Qwen3" in model_name:
+        config.update(_get_qwen3_expert_config(hf_config, d_model))
+    elif "DeepSeek" in model_name:
+        config.update(_get_deepseek_expert_config(hf_config, d_model, n_layers))
+    
+    return config
+
+
+def _get_qwen_expert_config(hf_config, d_model):
+    """Get Qwen-specific expert configuration"""
+    if (hasattr(hf_config, "moe_intermediate_size") and 
+        hasattr(hf_config, "shared_expert_intermediate_size")):
+        
+        return {
+            'expert_size': hf_config.moe_intermediate_size * 3 * d_model / 1e12,
+            'shared_experts_size_total': hf_config.shared_expert_intermediate_size * 3 * d_model / 1e12
+        }
+    return {}
+
+
+def _get_qwen3_expert_config(hf_config, d_model):
+    """Get Qwen3-specific expert configuration"""
+    if hasattr(hf_config, "moe_intermediate_size"):
+        return {
+            'expert_size': hf_config.moe_intermediate_size * 3 * d_model / 1e12
+        }
+    return {}
+
+
+def _get_deepseek_expert_config(hf_config, d_model, n_layers):
+    """Get DeepSeek-specific expert configuration"""
+    if (hasattr(hf_config, "moe_intermediate_size") and 
+        hasattr(hf_config, "intermediate_size") and 
+        hasattr(hf_config, "first_k_dense_replace")):
+        
+        deepseek_num_dense_layer = hf_config.first_k_dense_replace
+        return {
+            'expert_size': hf_config.moe_intermediate_size * 3 * d_model / 1e12,
+            'shared_experts_size_total': hf_config.moe_intermediate_size * 3 * d_model * 2 / 1e12,
+            'deepseek_dense_ffn_size': hf_config.intermediate_size * 3 * d_model / 1e12,
+            'deepseek_sparse_layer_num': n_layers - deepseek_num_dense_layer,
+            'deepseek_num_dense_layer': deepseek_num_dense_layer
+        }
+    return {}
+
+
+def _process_outputs(output_data, per_token_kv_size, attention_size_per_token, 
+                    model_name, hf_config, n_layers, n_attn_heads, d_head):
+    """Process outputs to calculate KV sizes and attention scores"""
+    kvs = []
+    true_kvs = []
+    attn_scores = []
+    
+    for prefill_len, output_len in zip(output_data['prefill_lengths'], output_data['output_lengths']):
+        # Calculate attention score
+        attn_score = _calculate_attention_score(model_name, hf_config, prefill_len, output_len,
+                                              n_layers, n_attn_heads, d_head)
+        attn_scores.append(attn_score)
+        
+        # Calculate KV sizes
+        kv_size = (prefill_len * per_token_kv_size + (output_len - 1) * per_token_kv_size / 2) / 1e12
+        true_kv = (prefill_len * per_token_kv_size + output_len * per_token_kv_size) / 1e9
+        
+        kvs.append(kv_size)
+        true_kvs.append(true_kv)
+    
+    return {
+        'kv_size': sum(kvs),
+        'true_kv_size': sum(true_kvs) * 1e3,
+        'attention_score': sum(attn_scores) / len(attn_scores)
+    }
+
+
+def _calculate_attention_score(model_name, hf_config, prefill_len, output_len, 
+                             n_layers, n_attn_heads, d_head):
+    """Calculate attention score for a single output"""
+    if ("DeepSeek" in model_name and 
+        hasattr(hf_config, "qk_rope_head_dim") and 
+        hasattr(hf_config, "qk_nope_head_dim") and 
+        hasattr(hf_config, "v_head_dim")):
+        
+        q_head_dim = hf_config.qk_rope_head_dim + hf_config.qk_nope_head_dim
+        k_size = n_layers * n_attn_heads * q_head_dim
+        v_size = n_layers * n_attn_heads * hf_config.v_head_dim
+        
+        score = (prefill_len * k_size + (output_len - 1) * k_size / 2 +
+                prefill_len * v_size + (output_len - 1) * v_size / 2)
+    else:
+        kv_size = n_layers * n_attn_heads * d_head
+        score = (prefill_len * kv_size + (output_len - 1) * kv_size / 2) * 2
+    
+    return score / 1e12
+
+
+def _calculate_throughput_metrics(batch_size, prefill_lengths, max_duration):
+    """Calculate throughput metrics"""
+    total_prefill = sum(prefill_lengths)
+    prefill_tp = total_prefill / (max_duration)
+    ttft = max_duration / batch_size
+    return ttft, prefill_tp
+
+
+def _calculate_smbu_smfu(model_name, n_layers, attention_size_per_token, expert_config,
+                        avg_activated_experts, metrics_data, hardware_specs, num_gpus,
+                        precision, ttft, prefill_tp, batch_size, decoding_tp):
+    """Calculate S-MBU and S-MFU metrics"""
+    prefill_activation = avg_activated_experts[1]
+    decode_steps_activation = avg_activated_experts[2:]
+    
+    # Calculate prefill metrics
+    prefill_smbu, prefill_smfu = _calculate_prefill_metrics(
+        model_name, n_layers, attention_size_per_token, expert_config,
+        prefill_activation, metrics_data['attention_score'], hardware_specs,
+        num_gpus, precision, ttft, prefill_tp
+    )
+    
+    # Calculate decoding metrics
+    decoding_smbu, decoding_smfu = _calculate_decoding_metrics(
+        model_name, n_layers, attention_size_per_token, expert_config,
+        decode_steps_activation, metrics_data, hardware_specs,
+        num_gpus, precision, batch_size, decoding_tp
+    )
+    
+    return {
+        'prefill_smbu': prefill_smbu,
+        'prefill_smfu': prefill_smfu,
+        'decoding_smbu': decoding_smbu,
+        'decoding_smfu': decoding_smfu
+    }
+
+
+def _calculate_prefill_metrics(model_name, n_layers, attention_size_per_token, expert_config,
+                             prefill_activation, attention_score, hardware_specs,
+                             num_gpus, precision, ttft, prefill_tp):
+    """Calculate prefill S-MBU and S-MFU"""
+    model_calculators = {
+        'Qwen': _calculate_qwen_prefill,
+        'Qwen3': _calculate_qwen3_prefill,
+        'DeepSeek': _calculate_deepseek_prefill
+    }
+    
+    for model_type, calculator in model_calculators.items():
+        if model_type in model_name and (model_type != 'Qwen' or 'Qwen3' not in model_name):
+            return calculator(n_layers, attention_size_per_token, expert_config,
+                            prefill_activation, attention_score, hardware_specs,
+                            num_gpus, precision, ttft, prefill_tp)
+    
+    # Default case
+    return _calculate_default_prefill(n_layers, attention_size_per_token, expert_config,
+                                    prefill_activation, attention_score, hardware_specs,
+                                    num_gpus, precision, ttft, prefill_tp)
+
+
+def _calculate_decoding_metrics(model_name, n_layers, attention_size_per_token, expert_config,
+                              decode_steps_activation, metrics_data, hardware_specs,
+                              num_gpus, precision, batch_size, decoding_tp):
+    """Calculate decoding S-MBU and S-MFU"""
+    decoding_smbus = []
+    
+    for activation in decode_steps_activation:
+        if "Qwen" in model_name and "Qwen3" not in model_name:
+            smbu, smfu = _calculate_qwen_decoding(n_layers, attention_size_per_token, expert_config,
+                                                activation, metrics_data, hardware_specs, num_gpus,
+                                                precision, batch_size, decoding_tp)
+        elif "Qwen3" in model_name:
+            smbu, smfu = _calculate_qwen3_decoding(n_layers, attention_size_per_token, expert_config,
+                                                 activation, metrics_data, hardware_specs, num_gpus,
+                                                 precision, batch_size, decoding_tp)
+        elif "DeepSeek" in model_name:
+            smbu, smfu = _calculate_deepseek_decoding(n_layers, attention_size_per_token, expert_config,
+                                                    activation, metrics_data, hardware_specs, num_gpus,
+                                                    precision, batch_size, decoding_tp)
+        else:
+            smbu, smfu = _calculate_default_decoding(n_layers, attention_size_per_token, expert_config,
+                                                   activation, metrics_data, hardware_specs, num_gpus,
+                                                   precision, batch_size, decoding_tp)
+        decoding_smbus.append(smbu)
+    
+    return sum(decoding_smbus) / len(decoding_smbus), smfu
+
+
+# Helper functions for specific model calculations
+def _calculate_qwen_prefill(n_layers, attention_size_per_token, expert_config, prefill_activation,
+                          attention_score, hardware_specs, num_gpus, precision, ttft, prefill_tp):
+    smbu_numerator = (n_layers * (prefill_activation * expert_config['expert_size'] + 
+                                expert_config['shared_experts_size_total'] + 
+                                attention_size_per_token)) * precision / ttft
+    smbu = smbu_numerator / (num_gpus * hardware_specs['peak_bandwidth_tb'])
+    
+    smfu_numerator = (n_layers * (attention_size_per_token + expert_config['expert_size'] + 
+                                expert_config['shared_experts_size_total']) + attention_score) * 2 * prefill_tp
+    smfu = smfu_numerator / (num_gpus * hardware_specs['peak_flops_tf'] / 2)
+    
+    return smbu, smfu
+
+
+def _calculate_qwen3_prefill(n_layers, attention_size_per_token, expert_config, prefill_activation,
+                           attention_score, hardware_specs, num_gpus, precision, ttft, prefill_tp):
+    smbu_numerator = (n_layers * (prefill_activation * expert_config['expert_size'] + 
+                                attention_size_per_token)) * precision / ttft
+    smbu = smbu_numerator / (num_gpus * hardware_specs['peak_bandwidth_tb'])
+    
+    smfu_numerator = (n_layers * (attention_size_per_token + expert_config['expert_size']) + 
+                     attention_score) * 2 * prefill_tp
+    smfu = smfu_numerator / (num_gpus * hardware_specs['peak_flops_tf'] / 2)
+    
+    return smbu, smfu
+
+
+def _calculate_deepseek_prefill(n_layers, attention_size_per_token, expert_config, prefill_activation,
+                              attention_score, hardware_specs, num_gpus, precision, ttft, prefill_tp):
+    smbu_numerator = ((n_layers * attention_size_per_token + 
+                      expert_config['deepseek_sparse_layer_num'] * 
+                      (prefill_activation * expert_config['expert_size'] + 
+                       expert_config['shared_experts_size_total']) + 
+                      expert_config['deepseek_num_dense_layer'] * 
+                      expert_config['deepseek_dense_ffn_size']) * precision / ttft)
+    smbu = smbu_numerator / (num_gpus * hardware_specs['peak_bandwidth_tb'])
+    
+    smfu_numerator = ((n_layers * attention_size_per_token + 
+                      expert_config['deepseek_sparse_layer_num'] * 
+                      (expert_config['expert_size'] + expert_config['shared_experts_size_total']) + 
+                      expert_config['deepseek_num_dense_layer'] * 
+                      expert_config['deepseek_dense_ffn_size'] + attention_score) * 2 * prefill_tp)
+    smfu = smfu_numerator / (num_gpus * hardware_specs['peak_flops_tf'] / 2)
+    
+    return smbu, smfu
+
+
+def _calculate_default_prefill(n_layers, attention_size_per_token, expert_config, prefill_activation,
+                             attention_score, hardware_specs, num_gpus, precision, ttft, prefill_tp):
+    # Default implementation
+    smbu_numerator = (n_layers * (prefill_activation * expert_config['expert_size'] + 
+                                attention_size_per_token)) * precision / ttft
+    smbu = smbu_numerator / (num_gpus * hardware_specs['peak_bandwidth_tb'])
+    
+    smfu_numerator = (n_layers * (attention_size_per_token + expert_config['expert_size']) + 
+                     attention_score) * 2 * prefill_tp
+    smfu = smfu_numerator / (num_gpus * hardware_specs['peak_flops_tf'] / 2)
+    
+    return smbu, smfu
+
+
+def _calculate_qwen_decoding(n_layers, attention_size_per_token, expert_config, activation,
+                           metrics_data, hardware_specs, num_gpus, precision, batch_size, decoding_tp):
+    smbu_numerator = ((n_layers * (activation * expert_config['expert_size'] + 
+                                 expert_config['shared_experts_size_total'] + 
+                                 attention_size_per_token) + 
+                      metrics_data['kv_size']) * precision / (batch_size / decoding_tp))
+    smbu = smbu_numerator / (num_gpus * hardware_specs['peak_bandwidth_tb'])
+    
+    smfu_numerator = ((n_layers * (attention_size_per_token + expert_config['expert_size'] + 
+                                 expert_config['shared_experts_size_total']) + 
+                      metrics_data['attention_score']) * 2 * decoding_tp)
+    smfu = smfu_numerator / (num_gpus * hardware_specs['peak_flops_tf'] / 2)
+    
+    return smbu, smfu
+
+
+def _calculate_qwen3_decoding(n_layers, attention_size_per_token, expert_config, activation,
+                            metrics_data, hardware_specs, num_gpus, precision, batch_size, decoding_tp):
+    smbu_numerator = ((n_layers * (activation * expert_config['expert_size'] + 
+                                 attention_size_per_token) + 
+                      metrics_data['kv_size']) * precision / (batch_size / decoding_tp))
+    smbu = smbu_numerator / (num_gpus * hardware_specs['peak_bandwidth_tb'])
+    
+    smfu_numerator = ((n_layers * (attention_size_per_token + expert_config['expert_size']) + 
+                      metrics_data['attention_score']) * 2 * decoding_tp)
+    smfu = smfu_numerator / (num_gpus * hardware_specs['peak_flops_tf'] / 2)
+    
+    return smbu, smfu
+
+
+def _calculate_deepseek_decoding(n_layers, attention_size_per_token, expert_config, activation,
+                               metrics_data, hardware_specs, num_gpus, precision, batch_size, decoding_tp):
+    smbu_numerator = ((n_layers * attention_size_per_token + 
+                      expert_config['deepseek_sparse_layer_num'] * 
+                      (activation * expert_config['expert_size'] + 
+                       expert_config['shared_experts_size_total']) + 
+                      expert_config['deepseek_num_dense_layer'] * 
+                      expert_config['deepseek_dense_ffn_size'] + 
+                      metrics_data['kv_size']) * precision / (batch_size / decoding_tp))
+    smbu = smbu_numerator / (num_gpus * hardware_specs['peak_bandwidth_tb'])
+    
+    smfu_numerator = ((n_layers * attention_size_per_token + 
+                      expert_config['deepseek_sparse_layer_num'] * 
+                      (expert_config['expert_size'] + expert_config['shared_experts_size_total']) + 
+                      expert_config['deepseek_num_dense_layer'] * 
+                      expert_config['deepseek_dense_ffn_size'] + 
+                      metrics_data['attention_score']) * 2 * decoding_tp)
+    smfu = smfu_numerator / (num_gpus * hardware_specs['peak_flops_tf'] / 2)
+    
+    return smbu, smfu
+
+
+def _calculate_default_decoding(n_layers, attention_size_per_token, expert_config, activation,
+                              metrics_data, hardware_specs, num_gpus, precision, batch_size, decoding_tp):
+    smbu_numerator = ((n_layers * (activation * expert_config['expert_size'] + 
+                                 attention_size_per_token) + 
+                      metrics_data['kv_size']) * precision / (batch_size / decoding_tp))
+    smbu = smbu_numerator / (num_gpus * hardware_specs['peak_bandwidth_tb'])
+    
+    smfu_numerator = ((n_layers * (attention_size_per_token + expert_config['expert_size']) + 
+                      metrics_data['attention_score']) * 2 * decoding_tp)
+    smfu = smfu_numerator / (num_gpus * hardware_specs['peak_flops_tf'] / 2)
+    
+    return smbu, smfu
+
+def _calculate_batch_metrics_hflm(output_len, context_prefill_size, decoding_tp, n_layers, d_model, 
+                                n_attn_heads, d_head, n_kv_heads, n_experts_per_tok, d_ff, 
+                                avg_activated_experts, hf_config, num_gpus, model_name, 
+                                used_dtype, batch_size, precision):
+    """Calculate metrics for a batch of outputs"""
+    gpu_type = get_gpu_details()
+    hardware_specs = {
+        "peak_bandwidth_tb": get_peak_bw(gpu_type) / 1e12,
+        "peak_flops_tf": get_peak_flops(gpu_type, precision=used_dtype) / 1e12,
+    }
+    
+    # Calculate KV sizes
+    per_token_kv_size = 2 * n_layers * d_head * n_kv_heads  # Default calculation
+    
+    if "DeepSeek" in model_name:
+        if hasattr(hf_config, "kv_lora_rank") and hasattr(hf_config, "qk_rope_head_dim"):
+            per_token_kv_size = n_layers * (hf_config.kv_lora_rank + hf_config.qk_rope_head_dim)
+    
+        
+    # Calculate attention scores
+    if "DeepSeek" in model_name and hasattr(hf_config, "qk_rope_head_dim") and hasattr(hf_config, "qk_nope_head_dim") and hasattr(hf_config, "v_head_dim"):
+        q_head_dim = hf_config.qk_rope_head_dim + hf_config.qk_nope_head_dim
+        origin_per_token_k_state_size = n_layers * n_attn_heads * q_head_dim
+        origin_per_token_v_state_size = n_layers * n_attn_heads * hf_config.v_head_dim
+        attention_score = context_prefill_size * origin_per_token_k_state_size + (output_len - 1) * origin_per_token_k_state_size / 2
+        attention_score += context_prefill_size * origin_per_token_v_state_size + (output_len - 1) * origin_per_token_v_state_size / 2
+        attention_score = attention_score / 1e12
     else:
-        raise ValueError(f"Unsupported precision: {precision}")
+        origin_per_token_kv_states_size = n_layers * n_attn_heads * d_head
+        attention_score = context_prefill_size * origin_per_token_kv_states_size + (output_len - 1) * origin_per_token_kv_states_size / 2
+        attention_score = attention_score * 2 / 1e12
+    
+    # Store attention scores and KV sizes
+    kv_size = context_prefill_size * per_token_kv_size + (output_len - 1) * per_token_kv_size / 2
+    kv_size = kv_size / 1e12
+    true_kv = (context_prefill_size * per_token_kv_size + output_len * per_token_kv_size) / 1e12 * 1e3
+    
+    # Calculate aggregate values
+    kv_size = kv_size * batch_size
+    true_kv_size = true_kv * batch_size * 1e3    
+    # Calculate attention size per token
+    if "DeepSeek" in model_name and hasattr(hf_config, "qk_rope_head_dim") and hasattr(hf_config, "qk_nope_head_dim") and hasattr(hf_config, "v_head_dim") and hasattr(hf_config, "kv_lora_rank"):
+        q_head_dim = hf_config.qk_rope_head_dim + hf_config.qk_nope_head_dim
+        if not hasattr(hf_config, "q_lora_rank") or not hf_config.q_lora_rank:
+            attention_size_per_token = (d_model * n_attn_heads * q_head_dim) + \
+                (d_model * (hf_config.kv_lora_rank + hf_config.qk_rope_head_dim)) + \
+                    (hf_config.kv_lora_rank * n_attn_heads * (q_head_dim - hf_config.qk_rope_head_dim + hf_config.v_head_dim)) + \
+                        (hf_config.v_head_dim * n_attn_heads * d_model)
+            attention_size_per_token = attention_size_per_token / 1e12
+        else:
+            attention_size_per_token = (d_model * hf_config.q_lora_rank) + \
+                (hf_config.q_lora_rank * n_attn_heads * q_head_dim) + \
+                    (d_model * (hf_config.kv_lora_rank + hf_config.qk_rope_head_dim)) + \
+                        (hf_config.kv_lora_rank * n_attn_heads * (q_head_dim - hf_config.qk_rope_head_dim + hf_config.v_head_dim)) + \
+                            (hf_config.v_head_dim * n_attn_heads * d_model)
+            attention_size_per_token = attention_size_per_token / 1e12
+    else:
+        attention_size_per_token = d_model * (n_attn_heads * d_head + n_kv_heads * d_head * 2) + n_attn_heads * d_head * d_model
+        attention_size_per_token = attention_size_per_token / 1e12
+    
+    # Calculate expert sizes
+    expert_size = d_ff * 3 * d_model / 1e12
+    shared_experts_size_total = 0
+    deepseek_dense_ffn_size = 0
+    deepseek_sparse_layer_num = 0
+    
+    if "Qwen" in model_name and hasattr(hf_config, "moe_intermediate_size") and hasattr(hf_config, "shared_expert_intermediate_size"):
+        d_ff = hf_config.moe_intermediate_size
+        d_ff_share = hf_config.shared_expert_intermediate_size
+        shared_experts_size = d_ff_share * 3 * d_model
+        expert_size = d_ff * 3 * d_model
+        shared_experts_size_total = shared_experts_size / 1e12
+        expert_size = expert_size / 1e12
+    elif "Qwen3" in model_name and hasattr(hf_config, "moe_intermediate_size"):
+        d_ff = hf_config.moe_intermediate_size
+        expert_size = d_ff * 3 * d_model
+        expert_size = expert_size / 1e12
+    elif "DeepSeek" in model_name and hasattr(hf_config, "moe_intermediate_size") and hasattr(hf_config, "intermediate_size") and hasattr(hf_config, "first_k_dense_replace"):
+        d_ff = hf_config.moe_intermediate_size
+        d_ff_dense = hf_config.intermediate_size
+        deepseek_num_dense_layer = hf_config.first_k_dense_replace
+        shared_experts_size = d_ff * 3 * d_model
+        expert_size = d_ff * 3 * d_model
+        shared_experts = 2
+        shared_experts_size_total = shared_experts_size * shared_experts / 1e12
+        expert_size = expert_size / 1e12
+        deepseek_sparse_layer_num = n_layers - deepseek_num_dense_layer
+        deepseek_dense_ffn_size = d_ff_dense * 3 * d_model / 1e12
+    
+    # Calculate S-MBU and S-MFU
+    if "Qwen" in model_name:
+        smbu = ((n_layers*(avg_activated_experts * expert_size + shared_experts_size_total + attention_size_per_token) + 
+                kv_size) * precision/(batch_size / decoding_tp) ) / (num_gpus * hardware_specs['peak_bandwidth_tb'])
+        smfu = (n_layers * (attention_size_per_token + n_experts_per_tok * expert_size + shared_experts_size_total) + attention_score) \
+            * 2 * decoding_tp / (num_gpus * hardware_specs['peak_flops_tf'] / 2)
+    elif "Qwen3" in model_name:
+        smbu = ((n_layers * (avg_activated_experts * expert_size + attention_size_per_token) + 
+                kv_size) * precision/(batch_size / decoding_tp) ) / (num_gpus * hardware_specs['peak_bandwidth_tb'])
+        smfu = (n_layers * (attention_size_per_token + n_experts_per_tok * expert_size) + attention_score) \
+            * 2 * decoding_tp / (num_gpus * hardware_specs['peak_flops_tf'] / 2)
+    elif "DeepSeek" in model_name:
+        smbu = ((n_layers * attention_size_per_token + deepseek_sparse_layer_num * \
+                (avg_activated_experts * expert_size + shared_experts_size_total) + \
+                deepseek_num_dense_layer * deepseek_dense_ffn_size + \
+                kv_size) * precision/(batch_size / decoding_tp) ) / (num_gpus * hardware_specs['peak_bandwidth_tb'])
+        smfu = (n_layers * attention_size_per_token + deepseek_sparse_layer_num * \
+                (n_experts_per_tok * expert_size + shared_experts_size_total) + \
+                deepseek_num_dense_layer * deepseek_dense_ffn_size + attention_score) \
+                * 2 * decoding_tp / (num_gpus * hardware_specs['peak_flops_tf'] / 2)
+    else:
+        smbu = ((n_layers*(avg_activated_experts * expert_size + attention_size_per_token) + 
+                kv_size) * precision/(batch_size / decoding_tp) ) / (num_gpus * hardware_specs['peak_bandwidth_tb'])
+        smfu = (n_layers * (attention_size_per_token + n_experts_per_tok * expert_size) + attention_score) \
+            * 2 * decoding_tp / (num_gpus * hardware_specs['peak_flops_tf'] / 2)
+    
+    return {
+        'smbu': smbu,
+        'smfu': smfu,
+        'kv_size': true_kv_size,
+        'decoding_throughput': decoding_tp,
+        'ttft': 0
+    }
+class ModelInfoRetriever:
+    def __init__(self, model_name: str, precision: str = 'float16'):
+        if precision not in ['float32', 'float16', 'bfloat16', 'int8', 'int4', 'awq', 'gptq', 'fp8', 'fp4']:
+            raise ValueError("Precision must be one of ['float32', 'float16', 'bfloat16', 'int8', 'int4', 'awq', 'gptq', 'fp8', 'fp4']")
+        self.model_name = model_name
+        self.precision = precision
+        self.config = AutoConfig.from_pretrained(model_name, trust_remote_code=True)
+        self.model_type = self.config.model_type
+
+    def get_model_precision_bits(self):
+        """Returns bit width used by the given quantization format."""
+        if self.precision == 'float32':
+            return 4
+        if self.precision in ['float16', 'bfloat16']:
+            return 2
+        if self.precision in ['int8', 'fp8']:
+            return 1
+        if self.precision in ['int4', 'fp4', 'gptq', 'awq']:
+            return 0.5
+        raise ValueError(f"Unsupported precision: {self.precision}")
+
+    def get_attention_info(self):
+        """Returns attention-related info"""
+        return {
+            'num_attention_heads': getattr(self.config, "num_attention_heads", None),
+            'num_key_value_heads': getattr(self.config, "num_key_value_heads", getattr(self.config, "num_kv_heads", None)),
+            'head_dim': getattr(self.config, "head_dim", getattr(self.config, "hidden_size", None) // getattr(self.config, "num_attention_heads", 1))
+        }
+
+    def get_rope_info(self):
+        """Returns RoPE (rotary embedding) info if available"""
+        if hasattr(self.config, "rope_scaling"):
+            return {
+                "type": self.config.rope_scaling.get("type"),
+                "factor": self.config.rope_scaling.get("factor")
+            }
+        elif hasattr(self.config, "use_alibi"):
+            return {"type": "alibi", "enabled": self.config.use_alibi}
+        else:
+            return {"type": "none"}
+
+    def get_moe_info(self, d_model=None):
+        """Returns MoE configuration such as number of experts and FFN dim"""
+        if d_model is None:
+            d_model = getattr(self.config, "hidden_size", None)
+
+        num_experts = (
+            getattr(self.config, "num_local_experts", None) or
+            getattr(self.config, "num_experts", None) or
+            getattr(self.config, "n_routed_experts", None) or
+            getattr(getattr(self.config, "ffn_config", {}), "moe_num_experts", None) or
+            1
+        )
+        n_experts_per_tok = (
+            getattr(self.config, "num_experts_per_tok", None) or
+            getattr(self.config, "num_selected_experts", None) or
+            getattr(getattr(self.config, "ffn_config", {}), "moe_top_k", None) or
+            1
+        )
+        d_ff = (
+            getattr(self.config, "ffn_dim", None) or
+            getattr(self.config, "intermediate_size", None) or
+            getattr(self.config, "d_ff", None) or
+            (d_model * getattr(self.config, "ff_ratio", 4)) or
+            getattr(getattr(self.config, "ffn_config", {}), "ffn_hidden_size", None) or
+            (4 * d_model)
+        )
+
+        return {
+            "num_experts": num_experts,
+            "experts_per_token": n_experts_per_tok,
+            "ffn_dim": d_ff
+        }
+
+    def get_architecture_info(self):
+        """Returns model-wide architecture info"""
+        return {
+            "model_type": self.model_type,
+            "hidden_size": getattr(self.config, "hidden_size", None),
+            "num_hidden_layers": getattr(self.config, "num_hidden_layers", None),
+            "max_position_embeddings": getattr(self.config, "max_position_embeddings", None),
+            "vocab_size": getattr(self.config, "vocab_size", None),
+            "architectures": getattr(self.config, "architectures", []),
+        }
+
+    def summarize(self):
+        """Aggregate all extracted info in a dictionary"""
+        d_model = getattr(self.config, "hidden_size", None)
+        return {
+            "model_name": self.model_name,
+            "model_type": self.model_type,
+            "precision_bits": self.get_model_precision_bits(),
+            "architecture": self.get_architecture_info(),
+            "attention": self.get_attention_info(),
+            "rope": self.get_rope_info(),
+            "moe": self.get_moe_info(d_model)
+        }
+    
+
 
-if __name__ == "__main__":
-    print(analyze_gpu_stats(parse_nvidia_smi()))
+# if __name__ == "__main__":
+#     print(analyze_gpu_stats(parse_nvidia_smi()))
+#     print(get_gpu_details())