initial prototype

app.py

@@ -0,0 +1,61 @@
+import gradio as gr
+from transformers import AutoConfig
+
+
+# TODO: access token for gated repo
+# TODO: add note about sliding window attention
+def calculate(name: str, ctx_len: int, num_users: int, dtype: str):
+    cfg = AutoConfig.from_pretrained(name, trust_remote_code=True)
+    use_mla = cfg.architectures[0].startswith(("DeepseekV2", "DeepseekV3"))
+
+    if hasattr(cfg, "text_config"):
+        cfg = cfg.text_config
+
+    num_layers = cfg.num_hidden_layers
+    model_config = [
+        ["num_layers", num_layers],
+        ["max_ctx_len", cfg.max_position_embeddings],
+    ]
+
+    # TODO: show attention type, show calculation
+    if use_mla:
+        kv_lora_rank = cfg.kv_lora_rank
+        qk_rope_head_dim = cfg.qk_rope_head_dim
+        nelems_per_token = num_layers * (kv_lora_rank + qk_rope_head_dim)
+
+        model_config.append(["kv_lora_rank", kv_lora_rank])
+        model_config.append(["qk_rope_head_dim", qk_rope_head_dim])
+
+    else:
+        num_kv_heads = cfg.num_key_value_heads
+        head_dim = getattr(cfg, "head_dim", cfg.hidden_size // cfg.num_attention_heads)
+        nelems_per_token = num_layers * num_kv_heads * head_dim * 2
+
+        model_config.append(["num_kv_heads", num_kv_heads])
+        model_config.append(["head_dim", head_dim])
+
+    if dtype == "fp16/bf16":
+        nbytes_per_elem = 2
+    elif dtype == "fp8":
+        nbytes_per_elem = 1 + 2 / cfg.hidden_size  # assume per-token scaling
+
+    kv_cache_size = nelems_per_token * ctx_len * num_users * nbytes_per_elem / 1e9
+    return kv_cache_size, model_config
+
+
+demo = gr.Interface(
+    fn=calculate,
+    inputs=[
+        gr.Textbox(label="model_id", value="google/gemma-3-1b-it"),
+        gr.Number(label="Context length", value=128_000),
+        gr.Number(label="No. of users", value=1),
+        gr.Dropdown(label="KV cache dtype", choices=["fp16/bf16", "fp8"]),
+    ],
+    outputs=[
+        gr.Number(label="KV cache size (GB)", precision=2),
+        gr.Dataframe(
+            label="Model config", headers=["Key", "Value"], datatype=["str", "int"]
+        ),
+    ],
+)
+demo.launch()