AGENTS.md

Agent instructions for this project. Read this before answering any user request inside this repo.

Compatible with the cross-tool AGENTS.md convention (Codex CLI, Cursor, GitHub Copilot, Gemini CLI, Aider, Sourcegraph Amp, Windsurf, Zed, Continue, RooCode, Cline, Factory). Claude Code reads CLAUDE.md, which is a symlink to this file.

Project scope

This repo ports new Hugging Face model architectures into MAX as custom architectures. The workhorse is the import-model agent skill, vendored under .agents/skills/import-model/.

A bring-up writes five files under <slug>/: arch.py, model_config.py, model.py, weight_adapters.py, <slug>.py (the graph). The port is done when MAX greedy output matches the HF reference byte-for-byte (or near-identical past BF16 rounding).

How to start a bring-up

Every bring-up is driven by a plan document at .agents/plans/<model>.plan.md. The plan is the spec — donor choice, delta list, per-file edit order, verification gates with pass criteria. The canonical worked example is .agents/plans/mellum2-bring-up.plan.md (the JetBrains Mellum2 port that lives at mellum/).

When the user asks to bring up a new model:

1. Copy the Mellum plan as a template:

   cp .agents/plans/mellum2-bring-up.plan.md .agents/plans/<your-model>.plan.md

2. Run Phase 1 (inspect, list-keys, check-walls) to fill in the new plan's delta list and donor choice.
3. Confirm the plan with the user. Then execute it phase-by-phase, marking todos completed in the plan frontmatter as gates pass.
4. Read .agents/skills/import-model/SKILL.md alongside the plan — it contains the work-order rules (e.g. scaffold before edit, audit before serve) the plan assumes.

Reference docs live in .agents/skills/import-model/references/ (20+ topic-indexed files). Two that get linked from every plan:

• implement-graph.md — edit order + completion criteria for Phase 2
• divergences.md — symptom → cause table for Phase 3 debugging

Run everything through pixi

There are two pixi workspaces. Do not call bare python / max / transformers — always pixi run. List tasks with pixi task list.

Project root (./pixi.toml):

Skill env (./.agents/skills/import-model/pixi.toml) — cd there first. Tasks that take filesystem paths (--output-dir, --port-dir, --custom-architectures) resolve them relative to the skill root, so use ../../../<slug>/ (or an absolute path) to point at a port folder at the repo root.

Reference serve invocation (Mellum, BF16 weights = 24 GB, single GPU):

cd .agents/skills/import-model
pixi run serve \
 --model JetBrains/Mellum2-12B-A2.5B-Thinking \
 --custom-architectures ../../../mellum \
 --quantization-encoding bfloat16 \
 --trust-remote-code \
 --port 8000
  # Add --device-memory-utilization 0.45 ONLY when HF reference must
  # coexist on the same GPU (validation workflow on 80 GB card).
  # On a 40 GB card pass --device-memory-utilization 0.85 instead;
  # default 0.9 leaves too little headroom. ≥48 GB: default works.

_serve_wrapper.py exists because new transformers versions need torch imported and model_rebuild() called on PipelineConfig / MAXModelConfig before max serve boots. Do not bypass it.

Conventions

When implementing or reviewing a port, hold these as hard rules:

• Scaffold ≠ port. scaffold.py only copies the donor's graph. Do not max serve until every row in the plan's delta list is implemented in <slug>.py.
• Docstring audit is mandatory. After scaffold, donor lineage strings (gpt_oss, llama, qwen, single-GPU, etc.) are still in module docstrings and comments. Run the grep recipe from honest-docstrings.md and classify every hit OK / Lie / Stale before declaring Phase 2 done.
• State-dict coverage. Wire the audit from state-dict-audit.md after editing weight_adapters.py. Orphan keys or missing experts will fail at load.
• Preflight before serve. Four local checks (import smoke, graph dry-build, adapter↔graph key diff, weights-format preflight) from serve-and-iterate.md must pass before max serve. Cold compile is 5–25 minutes; do not eat that cost on a broken graph.
• Read HF modeling code, not just config.json. The Hub config's field names are not authoritative. The Mellum plan's delta list said "no QK-norm" — modeling_mellum.py clearly has per-head RMSNorm on Q and K before RoPE. Trust the modeling code.
• --custom-architectures takes the port dir, not its parent. MAX adds dirname(<port_dir>) to sys.path and imports basename(<port_dir>). Pass <repo_root>/<slug>/, not <repo_root>/.

Pitfalls (Mellum bring-up)

These bit the worked-example port and are likely to bite any port that involves MoE, multi-GPU sharding, or recent transformers builds:

• MoE gate sharding. A custom MoEGate.__call__ that adds softmax / norm_topk_prob renormalization must also override MoEGate.shard(). The base implementation returns plain MoEGate instances and silently strips the subclass routing logic. Canonical examples: Qwen3MoEGate.shard(), mellum/mellum.py::MellumMoEGate.shard.
• Override shard() whenever you override __call__ on any shardable MAX module. Same failure mode as MoE gate above.
• Empty token_logprobs from custom archs. The donor PipelineModel doesn't implement compute_log_probabilities, so the OpenAI completion route returns token_logprobs: []. Substitute top-1 token parity at varied prefix lengths — see _compare_top1.py for the pattern.
• GPU memory contention. MAX reserves 90% of the card by default (--device-memory-utilization 0.9), so max serve + an HF reference on the same GPU OOMs unless you either (a) precompute HF answers while MAX is stopped (_hf_precompute.py
  • _max_check.py) — the workflow used for the Mellum port — or (b) pass MAX --device-memory-utilization 0.45 so both fit on an 80 GB card. On smaller GPUs (40 GB), default 0.9 is already too tight; lower to 0.85.
• RoPE max_seq_len is the precompute budget, not a position cap. Sliding-attention layers still see absolute positions up to max_position_embeddings. Do not shrink the sliding RoPE's max_seq_len to sliding_window.
• Pydantic forward-reference. transformers from main triggers PydanticUserError: 'PipelineConfig' is not fully defined until torch is imported and model_rebuild() is called. Always launch max serve via _serve_wrapper.py (pixi run serve).
• NVML driver/library mismatch. nvidia-smi returns "Failed to initialize NVML: Driver/library version mismatch" and MAX cannot use the GPU. Try:

  sudo modprobe -r nvidia_uvm nvidia_drm nvidia_modeset nvidia
  sudo modprobe nvidia

  Both steps can fail: modprobe -r fails when something still holds /dev/nvidia* (kill the offending process first; lsmod | rg nvidia shows the refcount), and modprobe nvidia fails when the running kernel has no matching module in /lib/modules/$(uname -r) (kernel ↔ driver skew on containerized / coder workspaces). When unloading isn't possible, reboot the workspace — that re-aligns the kernel with whatever NVIDIA driver package is installed.

Full pitfall index: pitfalls.md.