Workspace/Mini projects
Loading progress
All mini projects
MODULE 04 · 5 HOUR BUILD

Transformer resource worksheet

Build a reproducible resource report that compares cache layouts, prompt/output reservations, and idealized concurrency under an explicit memory budget.

Build evidence Record your actual checks, results, and limitations.

Build it in stages

  1. Run the seed and verify one cache estimate by hand.
  2. Replace embedded configurations with validated JSON input.
  3. Include separate terms for weight storage, packed KV cache, and configurable headroom.
  4. Reserve expected output tokens and compare against prompt-only estimates.
  5. Compare multi-head and grouped-query layouts while keeping other dimensions fixed.
  6. Write an assumptions section and a measurement plan for a real serving engine.

Your acceptance criteria

Use these as your project review. Record commands, outputs, and failure cases in your repository.

  • Every number has explicit bytes, MiB, or GiB units.
  • Doubling sequence length doubles only the relevant cache term.
  • The report distinguishes estimated memory from measured latency or quality.
  • Invalid and impossible budgets produce explicit results.
  • A source and configuration record accompanies each architecture comparison.

A working starting point

The seed runs as supplied. Extend it to satisfy the full brief. It is a teaching starting point, not a finished portfolio submission.

main.py
python
import json

CONFIGS = [
    {"name": "teaching-mha", "layers": 24, "kv_heads": 8, "head_dim": 64},
    {"name": "teaching-gqa", "layers": 24, "kv_heads": 2, "head_dim": 64},
]

def estimate(config, prompt_tokens, output_reserve, bytes_per_element=2):
    total_tokens = prompt_tokens + output_reserve
    per_token = 2 * config["layers"] * config["kv_heads"] * config["head_dim"] * bytes_per_element
    return per_token * total_tokens

def report(cache_budget_mib=512):
    budget = cache_budget_mib * 2**20
    rows = []
    for config in CONFIGS:
        for prompt in (2048, 4096):
            reserve = 1024
            size = estimate(config, prompt, reserve)
            rows.append({"configuration": config["name"],
                         "prompt_tokens": prompt, "output_reserve": reserve,
                         "cache_bytes_per_request": size,
                         "cache_mib_per_request": size / 2**20,
                         "idealized_concurrent_requests": budget // size})
    return {"cache_budget_mib": cache_budget_mib,
            "assumptions": ["dense unquantized two-byte KV values", "uniform request lengths per row",
                            "cache-only budget excludes weights and temporary memory",
                            "teaching configurations, not named deployed models"],
            "rows": rows}

if __name__ == "__main__":
    result = report()
    assert result["rows"][0]["cache_mib_per_request"] == 144.0
    print(json.dumps(result, indent=2))

Push it further

Simulate an arrival queue with variable output lengths and compare static padded batching against a simplified packed continuous-batching policy, reporting both utilization and waiting time.