Spaces:

ecs-lab
/

SIMPLE

Running

SIMPLE / app.js

AndreP-git

added model name in plots

e96a59e 5 days ago

22.3 kB

// =====================================================
// SIMPLE Law analytical model
// =====================================================


// Fitted parameters (thetas.csv)
const THETA0 = -1.86626365e-03;
const THETA1 = 3.30843561e-11;
const THETA2 = 1.04066816e-09;
const THETA3 = 2.02013691e-07;

const THETAS = [THETA0, THETA1, THETA2, THETA3];

const COMPONENT_LABELS =
[
    "Prefill Amortization",
    "Decode Accum. KV-cache",
    "Decode Accum. Next-token"
];

const COMPONENT_COLORS = ["#55A868", "#E2C044", "#DD8452"];
const COMPONENT_HATCHES = ["x", "/", "-"];

// bottom -> top of the stack (indices into the component arrays)
const STACK_ORDER = [2, 1, 0];

const GRID = "#e1e0d9";

// matplotlib RdYlGn_r
const RDYLGN_R =
[
    "#006837", "#1a9850", "#66bd63", "#a6d96a", "#d9ef8b", "#ffffbf",
    "#fee08b", "#fdae61", "#f46d43", "#d73027", "#a50026"
];

// Validated models, from models.json (up to ~9B)
const PRESET_MODELS =
{
    "Llama 3.2 1B":
    {
        params:1.24e+09,
        d:2048,
        L:16,
        nq:32,
        nkv:8
    },

    "Llama 3.2 3B":
    {
        params:3.21e+09,
        d:3072,
        L:28,
        nq:24,
        nkv:8
    },

    "Llama 3.1 8B":
    {
        params:8.03e+09,
        d:4096,
        L:32,
        nq:32,
        nkv:8
    },

    "Gemma 2 2B":
    {
        params:2.61e+09,
        d:2304,
        L:26,
        nq:8,
        nkv:4
    },

    "Gemma 2 9B":
    {
        params:9.24e+09,
        d:3584,
        L:42,
        nq:16,
        nkv:8
    },

    "Granite 3B":
    {
        params:3.48e+09,
        d:2560,
        L:32,
        nq:32,
        nkv:32
    },

    "Granite-8B":
    {
        params:8.05e+09,
        d:4096,
        L:36,
        nq:32,
        nkv:8
    },

    "Qwen2 1.5B":
    {
        params:1.54e+09,
        d:1536,
        L:28,
        nq:12,
        nkv:2
    },

    "Qwen2 7B":
    {
        params:7.62e+09,
        d:3584,
        L:28,
        nq:28,
        nkv:4
    },

    "OPT 1.3B":
    {
        params:1.32e+09,
        d:2048,
        L:24,
        nq:32,
        nkv:32
    },

    "OPT 2.7B":
    {
        params:2.75e+09,
        d:2560,
        L:32,
        nq:32,
        nkv:32
    },

    "OPT 6.7B":
    {
        params:6.71e+09,
        d:4096,
        L:32,
        nq:32,
        nkv:32
    },

    "Falcon 7B":
    {
        params:7.78e+09,
        d:4096,
        L:36,
        nq:64,
        nkv:64
    }
};

// Size classes for the preset dropdown, by parameter count
const PRESET_GROUPS =
[
    {label:"XS · ~1B", max:2e9},
    {label:"S · ~2B–3B", max:5e9},
    {label:"M · ~6B–9B", max:Infinity}
];

window.onload = function()
{

    let selector = document.getElementById("preset");


    let names =
        Object.keys(PRESET_MODELS)
        .sort((a,b)=>PRESET_MODELS[a].params-PRESET_MODELS[b].params);

    let min=0;

    for(let group of PRESET_GROUPS)
    {
        let optgroup = document.createElement("optgroup");

        optgroup.label=group.label;

        for(let name of names.filter(n=>PRESET_MODELS[n].params>=min && PRESET_MODELS[n].params<group.max))
        {
            let option = document.createElement("option");

            option.text=name;

            option.value=name;

            optgroup.appendChild(option);
        }

        selector.appendChild(optgroup);

        min=group.max;
    }

    selector.value = "Llama 3.1 8B";

    renderThetas();

    renderFormula();

    fitFormula();
    document.fonts.ready.then(fitFormula);

    // Refit only when the width changes (fitting itself changes the height)
    let formulaWidth=0;
    new ResizeObserver(entries=>{
        const width=entries[0].contentRect.width;
        if(width!==formulaWidth)
        {
            formulaWidth=width;
            fitFormula();
        }
    }).observe(document.querySelector(".formula"));

    loadPreset();

}

// kind: "ok" | "loading" | "error" | "custom"
function setStatus(kind,html)
{

    let status = document.getElementById("status");

    status.className = "status "+kind;

    status.innerHTML = html;

    // The plots keep showing the previous model while loading or on error
    if(kind==="ok" || kind==="custom")
        document.querySelectorAll(".plotModel").forEach(
            el=>el.innerHTML = html
        );

}

// n_kv / n_q reduced to lowest terms, e.g. 8/32 -> "1/4"
function ratioFraction(nkv,nq)
{

    let gcd=(a,b)=>b ? gcd(b,a%b) : a;

    let g=gcd(nkv,nq) || 1;

    let num=nkv/g, den=nq/g;

    return den===1 ? String(num) : `${num}/${den}`;

}

function attentionLabel(nq,nkv)
{

    if(nq===nkv) return "MHA";

    if(nkv===1) return "MQA";

    return `GQA (${+(nq/nkv).toFixed(1)}x)`;

}

// Derived quantities shown under the architecture fields
function updateSummary(d,L,nq,nkv)
{

    document.getElementById("archSummary").innerHTML =
    `
    <span class="chip">${attentionLabel(nq,nkv)}</span>
    <span class="chip">R = ${ratioFraction(nkv,nq)}</span>
    `;

}

function loadPreset()
{

    let name = document.getElementById("preset").value;


    let model = PRESET_MODELS[name];

    if (!model)
    {
        console.error("Unknown preset:", name);
        return;
    }

    document.getElementById("d").value=model.d;
    document.getElementById("L").value=model.L;
    document.getElementById("nq").value=model.nq;
    document.getElementById("nkv").value=model.nkv;


    setStatus("ok", `<b>${name}</b> · validated`);


    draw();

}

// Returns the three additive terms, as in the notebook's compute_components
function computeComponents(d,L,nq,nkv,nin,nout)
{

    const ratio = nkv/nq;

    let prefill=[];
    let kvCache=[];
    let nextTok=[];


    for(let n of nout)
    {

        prefill.push(
            THETA1*d*d*L*nin/n
        );

        kvCache.push(
            THETA2*d*L*ratio*(nin + n/2)
        );

        nextTok.push(
            THETA0 + THETA3*d*L
        );
    }


    return [
        prefill,kvCache,nextTok
    ];
}

function energyPerToken(d,L,nq,nkv,nin,nout)
{
    return THETA1*d*d*L*nin/nout
        + THETA2*d*L*(nkv/nq)*(nin + nout/2)
        + THETA0 + THETA3*d*L;
}



function logspace2(start,end,points)
{
    let arr=[];

    let a=Math.log2(start);
    let b=Math.log2(end);

    for(let i=0;i<points;i++)
    {
        let x=a+(b-a)*i/(points-1);
        arr.push(2**x);
    }

    return arr;
}



function draw()
{

    let d = Number(document.getElementById("d").value);

    let L = Number(document.getElementById("L").value);

    let nq = Number(document.getElementById("nq").value);

    let nkv = Number(document.getElementById("nkv").value);

    let nin = Number(document.getElementById("nin").value);


    updateSummary(d,L,nq,nkv);

    drawComponents(d,L,nq,nkv,nin);

    drawContour(d,L,nq,nkv,nin);

}



function drawComponents(d,L,nq,nkv,nin)
{

    let nout = logspace2(1,16384,400);


    let components =
        computeComponents(
            d,L,nq,nkv,nin,nout
        );


    let total =
        nout.map(
            (_,i)=>
            components.reduce(
                (a,b)=>a+b[i],
                0
            )
        );


    let traces=[];


    for(let i of STACK_ORDER)
    {

        traces.push({
            x:nout,
            y:components[i],
            mode:"lines",
            name:COMPONENT_LABELS[i],
            legendrank:i+1,
            stackgroup:"one",
            line:{
                color:"rgba(0,0,0,0.2)",
                width:0.5
            },
            fillcolor:COMPONENT_COLORS[i],
            fillpattern:{
                shape:COMPONENT_HATCHES[i],
                bgcolor:COMPONENT_COLORS[i],
                fgcolor:"rgba(0,0,0,0.6)",
                size:14,
                solidity:0.05
            },
            hovertemplate:"%{y:.3g} J/tok<extra>"+COMPONENT_LABELS[i]+"</extra>"
        });

    }

    traces.push({
        x:nout,
        y:total,
        mode:"lines",
        name:"Total E<sub>tok</sub>",
        legendrank:4,
        line:{
            color:"black",
            width:2.5
        },
        hovertemplate:"%{y:.3g} J/tok<extra>Total</extra>"

    });


    // Every other power of two, so the ticks fit the narrower panel
    let ticks=[];
    for(let e=0;e<=14;e+=2) ticks.push(2**e);


    let layout={
        height:450,
        font:{
            family:"sans-serif",
            size:13
        },
        xaxis:{
            title:
            `Output Tokens τ<sub>out</sub> (τ<sub>in</sub>=${nin}) [log<sub>2</sub>]`,
            type:"log",
            range:[0,Math.log10(16384)],
            tickmode:"array",
            tickvals:ticks,
            ticktext:ticks.map(String),
            tickangle:-45,
            gridcolor:GRID,
            showline:true,
            linecolor:"black",
            ticks:"outside"
        },


        yaxis:{
            title:
            "Energy-Per-Token E<sub>tok</sub> [J/tok]",
            rangemode:"tozero",
            gridcolor:GRID,
            showline:true,
            linecolor:"black",
            ticks:"outside",
            zeroline:false
        },

        hovermode:"x unified",

        legend:{
            x:0.98,
            xanchor:"right",
            y:1,
            traceorder:"normal",
            bgcolor:"rgba(255,255,255,0)"
        },

        plot_bgcolor:"white",

        margin:{
            l:70,
            r:20,
            t:20,
            b:80
        }

    };


    Plotly.newPlot(
        "plot",
        traces,
        layout,
        {
            responsive:true
        }
    );

}



// Same grid and level scheme as plot_E_tok_contour in the notebook.
// Axes are plotted in log2 units and E in log10, so evenly spaced
// Plotly contour levels match np.geomspace in the notebook.
const CONTOUR_NIN_RANGE = [16, 65536];
const CONTOUR_NOUT_RANGE = [4, 16384];
const CONTOUR_POINTS = 200;
const CONTOUR_LEVELS = 24;

function drawContour(d,L,nq,nkv,nin)
{

    let xs = logspace2(...CONTOUR_NIN_RANGE, CONTOUR_POINTS).map(Math.log2);

    let ys = logspace2(...CONTOUR_NOUT_RANGE, CONTOUR_POINTS).map(Math.log2);


    let z=[];
    let text=[];

    for(let y of ys)
    {
        let zrow=[];
        let trow=[];

        for(let x of xs)
        {
            let e = energyPerToken(d,L,nq,nkv,2**x,2**y);

            zrow.push(Math.log10(e));

            trow.push(
                `τ<sub>in</sub> = ${Math.round(2**x)}<br>`
                +`τ<sub>out</sub> = ${Math.round(2**y)}<br>`
                +`E<sub>tok</sub> = ${e.toPrecision(3)} J/tok`
            );
        }

        z.push(zrow);
        text.push(trow);
    }


    let flat = z.flat();
    let zmin = Math.min(...flat);
    let zmax = Math.max(...flat);
    let step = (zmax-zmin)/(CONTOUR_LEVELS-1);

    let levels=[];
    for(let i=0;i<CONTOUR_LEVELS;i++) levels.push(zmin+i*step);


    let colorscale =
        RDYLGN_R.map((c,i)=>[i/(RDYLGN_R.length-1),c]);


    let fill = {
        type:"contour",
        x:xs,
        y:ys,
        z:z,
        text:text,
        hovertemplate:"%{text}<extra></extra>",
        colorscale:colorscale,
        zmin:zmin,
        zmax:zmax,
        autocontour:false,
        contours:{
            start:zmin,
            end:zmax,
            size:step,
            coloring:"fill"
        },
        line:{
            width:0
        },
        colorbar:{
            title:{
                text:"Energy-Per-Token E<sub>tok</sub> [J/tok]",
                side:"right"
            },
            tickvals:levels.filter((_,i)=>i%3===0),
            ticktext:levels.filter((_,i)=>i%3===0).map(v=>formatEnergy(10**v)),
            outlinewidth:0.8
        }
    };


    let lines = {
        type:"contour",
        x:xs,
        y:ys,
        z:z,
        hoverinfo:"skip",
        showscale:false,
        autocontour:false,
        contours:{
            start:zmin,
            end:zmax,
            size:2*step,
            coloring:"none"
        },
        line:{
            color:"black",
            width:0.5
        }
    };


    let xticks=[];
    for(let e=Math.log2(CONTOUR_NIN_RANGE[0]);e<=Math.log2(CONTOUR_NIN_RANGE[1]);e++) xticks.push(e);

    let yticks=[];
    for(let e=Math.log2(CONTOUR_NOUT_RANGE[0]);e<=Math.log2(CONTOUR_NOUT_RANGE[1]);e++) yticks.push(e);


    let annotations =
        contourLabels(xs,ys,z,levels.filter((_,i)=>i%2===0),colorscale,zmin,zmax);


    // Mark the input-token count currently selected with the slider
    let shapes=[];
    let ninLog = Math.log2(nin);

    if(ninLog>=xs[0] && ninLog<=xs[xs.length-1])
    {
        shapes.push({
            type:"line",
            x0:ninLog,
            x1:ninLog,
            y0:ys[0],
            y1:ys[ys.length-1],
            line:{
                color:"black",
                width:1.5,
                dash:"dash"
            }
        });
    }


    let layout={
        height:450,
        font:{
            family:"sans-serif",
            size:13
        },
        xaxis:{
            title:"Input Tokens τ<sub>in</sub> [log<sub>2</sub>]",
            range:[xs[0],xs[xs.length-1]],
            tickmode:"array",
            tickvals:xticks,
            ticktext:xticks.map(e=>String(2**e)),
            tickangle:-45,
            showline:true,
            linecolor:"black",
            ticks:"outside",
            showgrid:false,
            zeroline:false
        },
        yaxis:{
            title:"Output Tokens τ<sub>out</sub> [log<sub>2</sub>]",
            range:[ys[0],ys[ys.length-1]],
            tickmode:"array",
            tickvals:yticks,
            ticktext:yticks.map(e=>String(2**e)),
            showline:true,
            linecolor:"black",
            ticks:"outside",
            showgrid:false,
            zeroline:false
        },
        annotations:annotations,
        shapes:shapes,
        margin:{
            l:80,
            r:20,
            t:20,
            b:90
        }
    };


    Plotly.newPlot(
        "contour",
        [fill,lines],
        layout,
        {
            responsive:true
        }
    );

}



// Place one label per contour line where it is farthest from every plot
// edge, skipping lines that never get far enough (as in the notebook).
function contourLabels(xs,ys,z,lineLevels,colorscale,zmin,zmax)
{

    const MARGIN_THRESHOLD = 0.09;

    let xmin=xs[0], xmax=xs[xs.length-1];
    let ymin=ys[0], ymax=ys[ys.length-1];

    let annotations=[];


    for(let level of lineLevels)
    {

        let best=null;
        let bestMargin=-1;

        let consider=(x,y)=>
        {
            let nx=(x-xmin)/(xmax-xmin);
            let ny=(y-ymin)/(ymax-ymin);
            let m=Math.min(nx,1-nx,ny,1-ny);

            if(m>bestMargin)
            {
                bestMargin=m;
                best=[x,y];
            }
        };


        for(let i=0;i<ys.length;i++)
        {
            for(let j=0;j<xs.length;j++)
            {
                let a=z[i][j]-level;

                if(j+1<xs.length)
                {
                    let b=z[i][j+1]-level;

                    if(a*b<0)
                        consider(xs[j]+(xs[j+1]-xs[j])*a/(a-b), ys[i]);
                }

                if(i+1<ys.length)
                {
                    let b=z[i+1][j]-level;

                    if(a*b<0)
                        consider(xs[j], ys[i]+(ys[i+1]-ys[i])*a/(a-b));
                }
            }
        }


        if(best && bestMargin>MARGIN_THRESHOLD)
        {
            annotations.push({
                x:best[0],
                y:best[1],
                text:formatEnergy(10**level),
                showarrow:false,
                font:{
                    size:11,
                    color:"black"
                },
                bgcolor:colorAt(colorscale,(level-zmin)/(zmax-zmin)),
                borderpad:1
            });
        }

    }


    return annotations;

}



function colorAt(colorscale,t)
{

    t=Math.min(1,Math.max(0,t));

    for(let i=1;i<colorscale.length;i++)
    {
        let [t1,c1]=colorscale[i];

        if(t<=t1)
        {
            let [t0,c0]=colorscale[i-1];
            let f=(t-t0)/(t1-t0);

            let rgb0=hexToRgb(c0);
            let rgb1=hexToRgb(c1);

            let mix=rgb0.map((v,k)=>Math.round(v+(rgb1[k]-v)*f));

            return `rgb(${mix.join(",")})`;
        }
    }

    return colorscale[colorscale.length-1][1];

}

function hexToRgb(hex)
{
    return [1,3,5].map(i=>parseInt(hex.slice(i,i+2),16));
}

function formatEnergy(v)
{
    return String(Number(v.toPrecision(2)));
}



// Shrinks the formula so it always fits on one line.
// All its inner sizes are in em, so the width scales linearly with font-size.
function fitFormula()
{

    const formula=document.querySelector(".formula");

    // Back to the CSS size, then scale down only if it overflows
    formula.style.fontSize="";

    const available=formula.clientWidth;
    const needed=formula.scrollWidth;

    if(needed>available)
    {
        const base=parseFloat(getComputedStyle(formula).fontSize);
        formula.style.fontSize=(base*available/needed*0.98)+"px";
    }

}

// Renders every element carrying a data-tex attribute
function renderFormula()
{

    document.querySelectorAll("[data-tex]").forEach(
        el=>katex.render(
            el.dataset.tex,
            el,
            {
                throwOnError:false
            }
        )
    );

}

// Component each theta belongs to (indices into COMPONENT_LABELS)
const THETA_COMPONENT = [2, 0, 1, 2];

const COMPONENT_CLASSES = ["prefill", "kvcache", "nexttok"];

function renderThetas()
{

    let body = document.getElementById("thetaBody");

    body.innerHTML = THETAS.map(
        (t,i)=>
        {
            let [mantissa,exponent] = t.toExponential(4).split("e");
            let c = THETA_COMPONENT[i];

            return `
            <tr class="${COMPONENT_CLASSES[c]}">
                <td><span class="dot"></span><span data-tex="\\theta_${i}"></span></td>
                <td data-tex="${mantissa}\\times10^{${Number(exponent)}}"></td>
                <td class="termName">${COMPONENT_LABELS[c]}</td>
            </tr>`;
        }
    ).join("");

}

// Reads d, L, n_q, n_kv from a config.json.
// Multimodal models (e.g. Gemma 3, Qwen-VL) nest the language model
// settings under text_config, so that is used when present.
function readArchitecture(config)
{

    let cfg = config.text_config || config;

    let d =
        cfg.hidden_size ||
        cfg.d_model ||
        cfg.n_embd;


    let L =
        cfg.num_hidden_layers ||
        cfg.num_layers ||
        cfg.n_layer;


    let nq =
        cfg.num_attention_heads ||
        cfg.n_head;


    let nkv =
        cfg.num_key_value_heads ||
        nq;


    return {cfg,d,L,nq,nkv};

}

function validateModel(config)
{

    let {cfg,d,L} = readArchitecture(config);

    // ------------------------------
    // Reject Mixture-of-Experts models
    // ------------------------------

    const moeKeys =
    [
        "num_experts",
        "num_local_experts",
        "num_experts_per_tok",
        "num_sparse_layers",
        "router_aux_loss_coef",
        "moe_layer_freq"
    ];


    for (let key of moeKeys)
    {
        if (cfg[key] !== undefined)
        {
            throw new Error(
                "MoE models are not supported. "
                +"Please select a dense Transformer model."
            );
        }
    }



    // ------------------------------
    // Reject hybrid / linear-attention models
    // (e.g. Qwen3.8: only 1 layer in 4 keeps a growing KV-cache)
    // ------------------------------

    const attentionTypes =
    [
        "full_attention",
        "sliding_attention"
    ];

    let layerTypes = cfg.layer_types || [];

    if(layerTypes.some(t=>!attentionTypes.includes(t)))
    {
        throw new Error(
            "Hybrid / linear-attention models are not supported. "
            +"Please select a standard Transformer model."
        );
    }



    // ------------------------------
    // Estimate parameter count
    // ------------------------------

    if(!d || !L)
    {
        throw new Error(
            "Cannot estimate model size from config.json"
        );
    }

    // Dense Transformer approximation
    let params = 12 * L * d * d;


    let paramsB = params / 1e9;

    // ------------------------------
    // Size restriction
    // ------------------------------

    const MAX_PARAMS_B = 20;


    if(paramsB > MAX_PARAMS_B)
    {
        throw new Error(
            `Model too large (~${paramsB.toFixed(1)}B parameters). `
            +
            `Only dense models below ${MAX_PARAMS_B}B are allowed.`
        );
    }


    return paramsB;

}

async function loadModel()
{

    let repo =
        document.getElementById("repo").value.trim();


    await loadRepo(repo);

}

async function loadRepo(repo)
{

    if(!repo)
    {
        setStatus("error", "Enter a repository id, e.g. deepseek-ai/DeepSeek-R1-Distill-Llama-8B");
        return;
    }

    let url =
    `https://huggingface.co/${repo}/resolve/main/config.json`;

    setStatus("loading", `Loading <b>${repo}</b>…`);


    try
    {

        let response = await fetch(url);

        if(!response.ok)
            throw Error(
                "Could not load config.json. Some models (e.g. meta-llama models) require licenses."
            );


        let cfg = await response.json();

        validateModel(cfg);

        let {d,L,nq,nkv} = readArchitecture(cfg);


        document.getElementById("d").value=d;
        document.getElementById("L").value=L;
        document.getElementById("nq").value=nq;
        document.getElementById("nkv").value=nkv;



        setStatus("ok", `<b>${repo}</b> · Hugging Face`);


        draw();


    }
    catch(e)
    {

        setStatus("error", e.message);

    }

}

// redraw when the architecture fields change

document
.querySelectorAll("input[type=number]")
.forEach(
    x=>x.addEventListener("change",()=>
    {
        setStatus("custom", "Custom architecture");
        draw();
    })
);

document
.getElementById("repo")
.addEventListener("keydown",e=>
{
    if(e.key==="Enter") loadModel();
});


draw();

	// =====================================================
	// SIMPLE Law analytical model
	// =====================================================


	// Fitted parameters (thetas.csv)
	const THETA0 = -1.86626365e-03;
	const THETA1 = 3.30843561e-11;
	const THETA2 = 1.04066816e-09;
	const THETA3 = 2.02013691e-07;

	const THETAS = [THETA0, THETA1, THETA2, THETA3];

	const COMPONENT_LABELS =
	[
	"Prefill Amortization",
	"Decode Accum. KV-cache",
	"Decode Accum. Next-token"
	];

	const COMPONENT_COLORS = ["#55A868", "#E2C044", "#DD8452"];
	const COMPONENT_HATCHES = ["x", "/", "-"];

	// bottom -> top of the stack (indices into the component arrays)
	const STACK_ORDER = [2, 1, 0];

	const GRID = "#e1e0d9";

	// matplotlib RdYlGn_r
	const RDYLGN_R =
	[
	"#006837", "#1a9850", "#66bd63", "#a6d96a", "#d9ef8b", "#ffffbf",
	"#fee08b", "#fdae61", "#f46d43", "#d73027", "#a50026"
	];

	// Validated models, from models.json (up to ~9B)
	const PRESET_MODELS =
	{
	"Llama 3.2 1B":
	{
	params:1.24e+09,
	d:2048,
	L:16,
	nq:32,
	nkv:8
	},

	"Llama 3.2 3B":
	{
	params:3.21e+09,
	d:3072,
	L:28,
	nq:24,
	nkv:8
	},

	"Llama 3.1 8B":
	{
	params:8.03e+09,
	d:4096,
	L:32,
	nq:32,
	nkv:8
	},

	"Gemma 2 2B":
	{
	params:2.61e+09,
	d:2304,
	L:26,
	nq:8,
	nkv:4
	},

	"Gemma 2 9B":
	{
	params:9.24e+09,
	d:3584,
	L:42,
	nq:16,
	nkv:8
	},

	"Granite 3B":
	{
	params:3.48e+09,
	d:2560,
	L:32,
	nq:32,
	nkv:32
	},

	"Granite-8B":
	{
	params:8.05e+09,
	d:4096,
	L:36,
	nq:32,
	nkv:8
	},

	"Qwen2 1.5B":
	{
	params:1.54e+09,
	d:1536,
	L:28,
	nq:12,
	nkv:2
	},

	"Qwen2 7B":
	{
	params:7.62e+09,
	d:3584,
	L:28,
	nq:28,
	nkv:4
	},

	"OPT 1.3B":
	{
	params:1.32e+09,
	d:2048,
	L:24,
	nq:32,
	nkv:32
	},

	"OPT 2.7B":
	{
	params:2.75e+09,
	d:2560,
	L:32,
	nq:32,
	nkv:32
	},

	"OPT 6.7B":
	{
	params:6.71e+09,
	d:4096,
	L:32,
	nq:32,
	nkv:32
	},

	"Falcon 7B":
	{
	params:7.78e+09,
	d:4096,
	L:36,
	nq:64,
	nkv:64
	}
	};

	// Size classes for the preset dropdown, by parameter count
	const PRESET_GROUPS =
	[
	{label:"XS · ~1B", max:2e9},
	{label:"S · ~2B–3B", max:5e9},
	{label:"M · ~6B–9B", max:Infinity}
	];

	window.onload = function()
	{

	let selector = document.getElementById("preset");


	let names =
	Object.keys(PRESET_MODELS)
	.sort((a,b)=>PRESET_MODELS[a].params-PRESET_MODELS[b].params);

	let min=0;

	for(let group of PRESET_GROUPS)
	{
	let optgroup = document.createElement("optgroup");

	optgroup.label=group.label;

	for(let name of names.filter(n=>PRESET_MODELS[n].params>=min && PRESET_MODELS[n].params<group.max))
	{
	let option = document.createElement("option");

	option.text=name;

	option.value=name;

	optgroup.appendChild(option);
	}

	selector.appendChild(optgroup);

	min=group.max;
	}

	selector.value = "Llama 3.1 8B";

	renderThetas();

	renderFormula();

	fitFormula();
	document.fonts.ready.then(fitFormula);

	// Refit only when the width changes (fitting itself changes the height)
	let formulaWidth=0;
	new ResizeObserver(entries=>{
	const width=entries[0].contentRect.width;
	if(width!==formulaWidth)
	{
	formulaWidth=width;
	fitFormula();
	}
	}).observe(document.querySelector(".formula"));

	loadPreset();

	}

	// kind: "ok" \| "loading" \| "error" \| "custom"
	function setStatus(kind,html)
	{

	let status = document.getElementById("status");

	status.className = "status "+kind;

	status.innerHTML = html;

	// The plots keep showing the previous model while loading or on error
	if(kind==="ok" \|\| kind==="custom")
	document.querySelectorAll(".plotModel").forEach(
	el=>el.innerHTML = html
	);

	}

	// n_kv / n_q reduced to lowest terms, e.g. 8/32 -> "1/4"
	function ratioFraction(nkv,nq)
	{

	let gcd=(a,b)=>b ? gcd(b,a%b) : a;

	let g=gcd(nkv,nq) \|\| 1;

	let num=nkv/g, den=nq/g;

	return den===1 ? String(num) : `${num}/${den}`;

	}

	function attentionLabel(nq,nkv)
	{

	if(nq===nkv) return "MHA";

	if(nkv===1) return "MQA";

	return `GQA (${+(nq/nkv).toFixed(1)}x)`;

	}

	// Derived quantities shown under the architecture fields
	function updateSummary(d,L,nq,nkv)
	{

	document.getElementById("archSummary").innerHTML =
	`
	<span class="chip">${attentionLabel(nq,nkv)}</span>
	<span class="chip">R = ${ratioFraction(nkv,nq)}</span>
	`;

	}

	function loadPreset()
	{

	let name = document.getElementById("preset").value;


	let model = PRESET_MODELS[name];

	if (!model)
	{
	console.error("Unknown preset:", name);
	return;
	}

	document.getElementById("d").value=model.d;
	document.getElementById("L").value=model.L;
	document.getElementById("nq").value=model.nq;
	document.getElementById("nkv").value=model.nkv;


	setStatus("ok", `<b>${name}</b> · validated`);


	draw();

	}

	// Returns the three additive terms, as in the notebook's compute_components
	function computeComponents(d,L,nq,nkv,nin,nout)
	{

	const ratio = nkv/nq;

	let prefill=[];
	let kvCache=[];
	let nextTok=[];


	for(let n of nout)
	{

	prefill.push(
	THETA1ddLnin/n
	);

	kvCache.push(
	THETA2dLratio(nin + n/2)
	);

	nextTok.push(
	THETA0 + THETA3dL
	);
	}


	return [
	prefill,kvCache,nextTok
	];
	}

	function energyPerToken(d,L,nq,nkv,nin,nout)
	{
	return THETA1ddLnin/nout
	+ THETA2dL(nkv/nq)(nin + nout/2)
	+ THETA0 + THETA3dL;
	}



	function logspace2(start,end,points)
	{
	let arr=[];

	let a=Math.log2(start);
	let b=Math.log2(end);

	for(let i=0;i<points;i++)
	{
	let x=a+(b-a)*i/(points-1);
	arr.push(2**x);
	}

	return arr;
	}



	function draw()
	{

	let d = Number(document.getElementById("d").value);

	let L = Number(document.getElementById("L").value);

	let nq = Number(document.getElementById("nq").value);

	let nkv = Number(document.getElementById("nkv").value);

	let nin = Number(document.getElementById("nin").value);


	updateSummary(d,L,nq,nkv);

	drawComponents(d,L,nq,nkv,nin);

	drawContour(d,L,nq,nkv,nin);

	}



	function drawComponents(d,L,nq,nkv,nin)
	{

	let nout = logspace2(1,16384,400);


	let components =
	computeComponents(
	d,L,nq,nkv,nin,nout
	);


	let total =
	nout.map(
	(_,i)=>
	components.reduce(
	(a,b)=>a+b[i],
	0
	)
	);


	let traces=[];


	for(let i of STACK_ORDER)
	{

	traces.push({
	x:nout,
	y:components[i],
	mode:"lines",
	name:COMPONENT_LABELS[i],
	legendrank:i+1,
	stackgroup:"one",
	line:{
	color:"rgba(0,0,0,0.2)",
	width:0.5
	},
	fillcolor:COMPONENT_COLORS[i],
	fillpattern:{
	shape:COMPONENT_HATCHES[i],
	bgcolor:COMPONENT_COLORS[i],
	fgcolor:"rgba(0,0,0,0.6)",
	size:14,
	solidity:0.05
	},
	hovertemplate:"%{y:.3g} J/tok<extra>"+COMPONENT_LABELS[i]+"</extra>"
	});

	}

	traces.push({
	x:nout,
	y:total,
	mode:"lines",
	name:"Total E<sub>tok</sub>",
	legendrank:4,
	line:{
	color:"black",
	width:2.5
	},
	hovertemplate:"%{y:.3g} J/tok<extra>Total</extra>"

	});


	// Every other power of two, so the ticks fit the narrower panel
	let ticks=[];
	for(let e=0;e<=14;e+=2) ticks.push(2**e);


	let layout={
	height:450,
	font:{
	family:"sans-serif",
	size:13
	},
	xaxis:{
	title:
	`Output Tokens τ<sub>out</sub> (τ<sub>in</sub>=${nin}) [log<sub>2</sub>]`,
	type:"log",
	range:[0,Math.log10(16384)],
	tickmode:"array",
	tickvals:ticks,
	ticktext:ticks.map(String),
	tickangle:-45,
	gridcolor:GRID,
	showline:true,
	linecolor:"black",
	ticks:"outside"
	},


	yaxis:{
	title:
	"Energy-Per-Token E<sub>tok</sub> [J/tok]",
	rangemode:"tozero",
	gridcolor:GRID,
	showline:true,
	linecolor:"black",
	ticks:"outside",
	zeroline:false
	},

	hovermode:"x unified",

	legend:{
	x:0.98,
	xanchor:"right",
	y:1,
	traceorder:"normal",
	bgcolor:"rgba(255,255,255,0)"
	},

	plot_bgcolor:"white",

	margin:{
	l:70,
	r:20,
	t:20,
	b:80
	}

	};


	Plotly.newPlot(
	"plot",
	traces,
	layout,
	{
	responsive:true
	}
	);

	}



	// Same grid and level scheme as plot_E_tok_contour in the notebook.
	// Axes are plotted in log2 units and E in log10, so evenly spaced
	// Plotly contour levels match np.geomspace in the notebook.
	const CONTOUR_NIN_RANGE = [16, 65536];
	const CONTOUR_NOUT_RANGE = [4, 16384];
	const CONTOUR_POINTS = 200;
	const CONTOUR_LEVELS = 24;

	function drawContour(d,L,nq,nkv,nin)
	{

	let xs = logspace2(...CONTOUR_NIN_RANGE, CONTOUR_POINTS).map(Math.log2);

	let ys = logspace2(...CONTOUR_NOUT_RANGE, CONTOUR_POINTS).map(Math.log2);


	let z=[];
	let text=[];

	for(let y of ys)
	{
	let zrow=[];
	let trow=[];

	for(let x of xs)
	{
	let e = energyPerToken(d,L,nq,nkv,2x,2y);

	zrow.push(Math.log10(e));

	trow.push(
	`τ<sub>in</sub> = ${Math.round(2**x)}<br>`
	+`τ<sub>out</sub> = ${Math.round(2**y)}<br>`
	+`E<sub>tok</sub> = ${e.toPrecision(3)} J/tok`
	);
	}

	z.push(zrow);
	text.push(trow);
	}


	let flat = z.flat();
	let zmin = Math.min(...flat);
	let zmax = Math.max(...flat);
	let step = (zmax-zmin)/(CONTOUR_LEVELS-1);

	let levels=[];
	for(let i=0;i<CONTOUR_LEVELS;i++) levels.push(zmin+i*step);


	let colorscale =
	RDYLGN_R.map((c,i)=>[i/(RDYLGN_R.length-1),c]);


	let fill = {
	type:"contour",
	x:xs,
	y:ys,
	z:z,
	text:text,
	hovertemplate:"%{text}<extra></extra>",
	colorscale:colorscale,
	zmin:zmin,
	zmax:zmax,
	autocontour:false,
	contours:{
	start:zmin,
	end:zmax,
	size:step,
	coloring:"fill"
	},
	line:{
	width:0
	},
	colorbar:{
	title:{
	text:"Energy-Per-Token E<sub>tok</sub> [J/tok]",
	side:"right"
	},
	tickvals:levels.filter((_,i)=>i%3===0),
	ticktext:levels.filter((_,i)=>i%3===0).map(v=>formatEnergy(10**v)),
	outlinewidth:0.8
	}
	};


	let lines = {
	type:"contour",
	x:xs,
	y:ys,
	z:z,
	hoverinfo:"skip",
	showscale:false,
	autocontour:false,
	contours:{
	start:zmin,
	end:zmax,
	size:2*step,
	coloring:"none"
	},
	line:{
	color:"black",
	width:0.5
	}
	};


	let xticks=[];
	for(let e=Math.log2(CONTOUR_NIN_RANGE[0]);e<=Math.log2(CONTOUR_NIN_RANGE[1]);e++) xticks.push(e);

	let yticks=[];
	for(let e=Math.log2(CONTOUR_NOUT_RANGE[0]);e<=Math.log2(CONTOUR_NOUT_RANGE[1]);e++) yticks.push(e);


	let annotations =
	contourLabels(xs,ys,z,levels.filter((_,i)=>i%2===0),colorscale,zmin,zmax);


	// Mark the input-token count currently selected with the slider
	let shapes=[];
	let ninLog = Math.log2(nin);

	if(ninLog>=xs[0] && ninLog<=xs[xs.length-1])
	{
	shapes.push({
	type:"line",
	x0:ninLog,
	x1:ninLog,
	y0:ys[0],
	y1:ys[ys.length-1],
	line:{
	color:"black",
	width:1.5,
	dash:"dash"
	}
	});
	}


	let layout={
	height:450,
	font:{
	family:"sans-serif",
	size:13
	},
	xaxis:{
	title:"Input Tokens τ<sub>in</sub> [log<sub>2</sub>]",
	range:[xs[0],xs[xs.length-1]],
	tickmode:"array",
	tickvals:xticks,
	ticktext:xticks.map(e=>String(2**e)),
	tickangle:-45,
	showline:true,
	linecolor:"black",
	ticks:"outside",
	showgrid:false,
	zeroline:false
	},
	yaxis:{
	title:"Output Tokens τ<sub>out</sub> [log<sub>2</sub>]",
	range:[ys[0],ys[ys.length-1]],
	tickmode:"array",
	tickvals:yticks,
	ticktext:yticks.map(e=>String(2**e)),
	showline:true,
	linecolor:"black",
	ticks:"outside",
	showgrid:false,
	zeroline:false
	},
	annotations:annotations,
	shapes:shapes,
	margin:{
	l:80,
	r:20,
	t:20,
	b:90
	}
	};


	Plotly.newPlot(
	"contour",
	[fill,lines],
	layout,
	{
	responsive:true
	}
	);

	}



	// Place one label per contour line where it is farthest from every plot
	// edge, skipping lines that never get far enough (as in the notebook).
	function contourLabels(xs,ys,z,lineLevels,colorscale,zmin,zmax)
	{

	const MARGIN_THRESHOLD = 0.09;

	let xmin=xs[0], xmax=xs[xs.length-1];
	let ymin=ys[0], ymax=ys[ys.length-1];

	let annotations=[];


	for(let level of lineLevels)
	{

	let best=null;
	let bestMargin=-1;

	let consider=(x,y)=>
	{
	let nx=(x-xmin)/(xmax-xmin);
	let ny=(y-ymin)/(ymax-ymin);
	let m=Math.min(nx,1-nx,ny,1-ny);

	if(m>bestMargin)
	{
	bestMargin=m;
	best=[x,y];
	}
	};


	for(let i=0;i<ys.length;i++)
	{
	for(let j=0;j<xs.length;j++)
	{
	let a=z[i][j]-level;

	if(j+1<xs.length)
	{
	let b=z[i][j+1]-level;

	if(a*b<0)
	consider(xs[j]+(xs[j+1]-xs[j])*a/(a-b), ys[i]);
	}

	if(i+1<ys.length)
	{
	let b=z[i+1][j]-level;

	if(a*b<0)
	consider(xs[j], ys[i]+(ys[i+1]-ys[i])*a/(a-b));
	}
	}
	}


	if(best && bestMargin>MARGIN_THRESHOLD)
	{
	annotations.push({
	x:best[0],
	y:best[1],
	text:formatEnergy(10**level),
	showarrow:false,
	font:{
	size:11,
	color:"black"
	},
	bgcolor:colorAt(colorscale,(level-zmin)/(zmax-zmin)),
	borderpad:1
	});
	}

	}


	return annotations;

	}



	function colorAt(colorscale,t)
	{

	t=Math.min(1,Math.max(0,t));

	for(let i=1;i<colorscale.length;i++)
	{
	let [t1,c1]=colorscale[i];

	if(t<=t1)
	{
	let [t0,c0]=colorscale[i-1];
	let f=(t-t0)/(t1-t0);

	let rgb0=hexToRgb(c0);
	let rgb1=hexToRgb(c1);

	let mix=rgb0.map((v,k)=>Math.round(v+(rgb1[k]-v)*f));

	return `rgb(${mix.join(",")})`;
	}
	}

	return colorscale[colorscale.length-1][1];

	}

	function hexToRgb(hex)
	{
	return [1,3,5].map(i=>parseInt(hex.slice(i,i+2),16));
	}

	function formatEnergy(v)
	{
	return String(Number(v.toPrecision(2)));
	}



	// Shrinks the formula so it always fits on one line.
	// All its inner sizes are in em, so the width scales linearly with font-size.
	function fitFormula()
	{

	const formula=document.querySelector(".formula");

	// Back to the CSS size, then scale down only if it overflows
	formula.style.fontSize="";

	const available=formula.clientWidth;
	const needed=formula.scrollWidth;

	if(needed>available)
	{
	const base=parseFloat(getComputedStyle(formula).fontSize);
	formula.style.fontSize=(baseavailable/needed0.98)+"px";
	}

	}

	// Renders every element carrying a data-tex attribute
	function renderFormula()
	{

	document.querySelectorAll("[data-tex]").forEach(
	el=>katex.render(
	el.dataset.tex,
	el,
	{
	throwOnError:false
	}
	)
	);

	}

	// Component each theta belongs to (indices into COMPONENT_LABELS)
	const THETA_COMPONENT = [2, 0, 1, 2];

	const COMPONENT_CLASSES = ["prefill", "kvcache", "nexttok"];

	function renderThetas()
	{

	let body = document.getElementById("thetaBody");

	body.innerHTML = THETAS.map(
	(t,i)=>
	{
	let [mantissa,exponent] = t.toExponential(4).split("e");
	let c = THETA_COMPONENT[i];

	return `
	<tr class="${COMPONENT_CLASSES[c]}">
	<td><span class="dot"></span><span data-tex="\\theta_${i}"></span></td>
	<td data-tex="${mantissa}\\times10^{${Number(exponent)}}"></td>
	<td class="termName">${COMPONENT_LABELS[c]}</td>
	</tr>`;
	}
	).join("");

	}

	// Reads d, L, n_q, n_kv from a config.json.
	// Multimodal models (e.g. Gemma 3, Qwen-VL) nest the language model
	// settings under text_config, so that is used when present.
	function readArchitecture(config)
	{

	let cfg = config.text_config \|\| config;

	let d =
	cfg.hidden_size \|\|
	cfg.d_model \|\|
	cfg.n_embd;


	let L =
	cfg.num_hidden_layers \|\|
	cfg.num_layers \|\|
	cfg.n_layer;


	let nq =
	cfg.num_attention_heads \|\|
	cfg.n_head;


	let nkv =
	cfg.num_key_value_heads \|\|
	nq;


	return {cfg,d,L,nq,nkv};

	}

	function validateModel(config)
	{

	let {cfg,d,L} = readArchitecture(config);

	// ------------------------------
	// Reject Mixture-of-Experts models
	// ------------------------------

	const moeKeys =
	[
	"num_experts",
	"num_local_experts",
	"num_experts_per_tok",
	"num_sparse_layers",
	"router_aux_loss_coef",
	"moe_layer_freq"
	];


	for (let key of moeKeys)
	{
	if (cfg[key] !== undefined)
	{
	throw new Error(
	"MoE models are not supported. "
	+"Please select a dense Transformer model."
	);
	}
	}



	// ------------------------------
	// Reject hybrid / linear-attention models
	// (e.g. Qwen3.8: only 1 layer in 4 keeps a growing KV-cache)
	// ------------------------------

	const attentionTypes =
	[
	"full_attention",
	"sliding_attention"
	];

	let layerTypes = cfg.layer_types \|\| [];

	if(layerTypes.some(t=>!attentionTypes.includes(t)))
	{
	throw new Error(
	"Hybrid / linear-attention models are not supported. "
	+"Please select a standard Transformer model."
	);
	}



	// ------------------------------
	// Estimate parameter count
	// ------------------------------

	if(!d \|\| !L)
	{
	throw new Error(
	"Cannot estimate model size from config.json"
	);
	}

	// Dense Transformer approximation
	let params = 12 * L * d * d;


	let paramsB = params / 1e9;

	// ------------------------------
	// Size restriction
	// ------------------------------

	const MAX_PARAMS_B = 20;


	if(paramsB > MAX_PARAMS_B)
	{
	throw new Error(
	`Model too large (~${paramsB.toFixed(1)}B parameters). `
	+
	`Only dense models below ${MAX_PARAMS_B}B are allowed.`
	);
	}


	return paramsB;

	}

	async function loadModel()
	{

	let repo =
	document.getElementById("repo").value.trim();


	await loadRepo(repo);

	}

	async function loadRepo(repo)
	{

	if(!repo)
	{
	setStatus("error", "Enter a repository id, e.g. deepseek-ai/DeepSeek-R1-Distill-Llama-8B");
	return;
	}

	let url =
	`https://huggingface.co/${repo}/resolve/main/config.json`;

	setStatus("loading", `Loading <b>${repo}</b>…`);


	try
	{

	let response = await fetch(url);

	if(!response.ok)
	throw Error(
	"Could not load config.json. Some models (e.g. meta-llama models) require licenses."
	);


	let cfg = await response.json();

	validateModel(cfg);

	let {d,L,nq,nkv} = readArchitecture(cfg);


	document.getElementById("d").value=d;
	document.getElementById("L").value=L;
	document.getElementById("nq").value=nq;
	document.getElementById("nkv").value=nkv;



	setStatus("ok", `<b>${repo}</b> · Hugging Face`);


	draw();


	}
	catch(e)
	{

	setStatus("error", e.message);

	}

	}

	// redraw when the architecture fields change

	document
	.querySelectorAll("input[type=number]")
	.forEach(
	x=>x.addEventListener("change",()=>
	{
	setStatus("custom", "Custom architecture");
	draw();
	})
	);

	document
	.getElementById("repo")
	.addEventListener("keydown",e=>
	{
	if(e.key==="Enter") loadModel();
	});


	draw();