Add configurable context-window with memory-safety validation

[alytaphoenix]

Summary

Adds a max_context_length option to instance placement, letting a request cap a model's context window below its native max — while making sure a larger context can't silently OOM a node.

• ModelCard gains head_dim (parsed from config, used to size KV cache per token per layer).
• PlaceInstance/PlaceInstanceParams gain max_context_length: int | None (None = use the model's native context_length; never allowed to exceed it, only cap down).
• New kv_cache_bytes_per_layer() estimates KV-cache memory (uses num_key_value_heads * head_dim when known, conservatively falls back to hidden_size otherwise — never under-counts).
• Pipeline sharding: KV-cache cost is now folded into the existing per-node memory check (_allocate_and_validate_layers) alongside layer weights. Because higher-memory nodes already get proportionally more layers via the existing largest-remainder allocation, correctly costing KV cache here means larger-context instances naturally land more KV burden on higher-memory nodes — no new redistribution logic needed.
• Tensor sharding: added per-rank memory validation that didn't exist before (_validate_tensor_parallel_memory) — this also fixes Max available combined RDMA memory needs to account for uneven devices under tensor parallelism #1936, where an uneven-memory cluster could accept a tensor-parallel rank that can't actually hold its equal share of model weights. Now checked for both weights and KV cache.
• On failure, the error names the largest context length that would fit, computed via _largest_fitting_context_length().
• max_context_length is threaded down through MlxBuilder → SequentialGenerator/BatchGenerator → mlx_generate/run_prefill_for_request → make_kv_cache(..., max_kv_size=...), so the cap is (where architecturally possible) enforced at runtime, not just placement time.

Known limitations (flagging for on-device review — I don't have mlx hardware in this environment)

• Runtime enforcement is code-complete but unverified. When max_kv_size is set, make_kv_cache returns a RotatingKVCache. This PR routes that into the existing prefix-cache path (KVPrefixCache). I have not verified that prefix-cache trim/restore stays correct once a rotating cache has wrapped past its capacity — the existing rotating-aware helpers (copy_rotating_kv_cache, etc.) look like they were built for hybrid models' internal caches, not necessarily a top-level rotating cache reaching them via this new path. Would appreciate a maintainer with real hardware checking a long conversation that exceeds the configured cap.
• Hybrid architectures that define their own model.make_cache() (NemotronH, Qwen3-Next) ignore max_kv_size entirely — this is pre-existing mlx_lm/make_kv_cache behavior, not something this PR changes. Placement-time memory protection still applies (the KV estimate is still validated against node memory before load), but the runtime cap isn't enforced during generation for these architectures.
• Tensor-parallel validation runs only on the cycle already selected by filter_cycles_by_memory (chosen by download-score/total-RAM before per-node validation runs). In a topology with multiple tensor-capable cycles, this can reject a selection containing one undersized node even though another all-large cycle would have fit. This mirrors existing pipeline-path behavior (not a regression this PR introduces), but is worth a follow-up if it becomes an issue in practice.
• Per-rank weight estimate (storage_size // world_size) doesn't account for replicated parameters (e.g. embeddings/router/norms in MoE models aren't actually sharded), so it slightly under-counts real per-rank weight memory. The KV-cache estimate's GQA overestimate (falling back to hidden_size when head_dim/num_key_value_heads aren't both known) partially offsets this in practice, but it's an approximation, same basis as the already-shipped pipeline path.

Test plan

• uv run basedpyright — 0 errors (verified identical to main baseline for all touched files)
• uv run ruff check — passes
• uv run ruff format — applied
• uv run pytest src/exo/master — 61 passed (28 in test_placement_utils.py, 27 in test_placement.py, including new tests for kv_cache_bytes_per_layer, tensor-parallel per-rank validation with and without a context cap, the Max available combined RDMA memory needs to account for uneven devices under tensor parallelism #1936 regression, and context-window placement fit/reject/native-max cases)
• Full suite excluding mlx-dependent dirs: 330 passed, 3 skipped
• nix fmt / nix flake check not run in this environment (ruff + basedpyright substituted)
• No manual on-device verification of runtime KV-cache capping — see limitations above

🤖 Generated with Claude Code

https://claude.ai/code/session_011rjSfwDBTkmySmfU6NgHKF