diff --git a/c/resource_plan.py b/c/resource_plan.py index fbd8604df..955f29a77 100644 --- a/c/resource_plan.py +++ b/c/resource_plan.py @@ -963,6 +963,26 @@ def _next_actions(bottleneck_class, projected_hit, probe_state, probe_gbs, } +def _family_expert_cache_knob(family_id, cache_bytes): + """Env var that actually sizes the expert LRU, when it is not RAM_GB. + + Kimi K3 reads K3_EXPERT_GB (default 8) and treats RAM_GB as a ceiling; + main() never takes argv as a cap. GLM-5.3 reads GLM53_EXPERT_GB and never + RAM_GB. Exporting only RAM_GB therefore left both engines on their own + default cache under --auto-tier. + """ + if (not isinstance(cache_bytes, int) or isinstance(cache_bytes, bool) + or cache_bytes <= 0): + return None + if family_id == "kimi": + return ("K3_EXPERT_GB", + "Kimi K3 sizes the expert LRU from K3_EXPERT_GB; RAM_GB is only a ceiling") + if family_id == "glm53": + return ("GLM53_EXPERT_GB", + "GLM-5.3 sizes the expert LRU from GLM53_EXPERT_GB, not RAM_GB") + return None + + def build_plan(model, ram_gb=0, context=4096, gpu_indices=None, vram_gb=0, available_memory=None, available_disk=None, gpus=None, policy="quality", physical_cpus=None, cpu_sockets=None, @@ -1129,6 +1149,13 @@ def build_plan(model, ram_gb=0, context=4096, gpu_indices=None, vram_gb=0, tune = _auto_tune(bottleneck_class, projected_hit, planning_gpus, cpu_sockets, plan_has_metal=False, engine_group=resolved.descriptor.engine_group) + # DRAFT/PIPE/PIN/NUMA stay colibri.c-only (#1585). These two are the + # opposite leftover: the engines that size the expert LRU from their own + # *EXPERT_GB variable, which RAM_GB does not set. + knob = _family_expert_cache_knob(resolved.descriptor.id, cache_bytes) + if knob: + name, reason = knob + tune[name] = {"value": f"{cache_bytes / GB:.3f}", "reason": reason} probe_state, probe_gbs = ssd_probe_state(info["path"]) actions = _next_actions(bottleneck_class, projected_hit, probe_state, probe_gbs, planning_gpus) @@ -1213,6 +1240,10 @@ def environment_for_plan(plan, env=None, cuda_enabled=True): result.setdefault("REPIN", "64") ram = plan["tiers"]["ram"] result.setdefault("RAM_GB", f"{ram['budget_bytes'] / GB:.3f}") + knob = _family_expert_cache_knob(plan.get("model", {}).get("family_id"), + ram.get("expert_cache_bytes")) + if knob: + result.setdefault(knob[0], f"{ram['expert_cache_bytes'] / GB:.3f}") planned_cap = ram.get("cache_slots_per_layer") if (plan.get("model", {}).get("family_id") == "qwen38" and (not isinstance(planned_cap, int) or diff --git a/c/tests/test_resource_plan.py b/c/tests/test_resource_plan.py index 03a7e8373..539fb484e 100644 --- a/c/tests/test_resource_plan.py +++ b/c/tests/test_resource_plan.py @@ -220,6 +220,105 @@ def test_sibling_plan_advises_no_colibri_knob(self): for key in ("DRAFT", "PIPE", "COLI_CUDA_PIPE", "COLI_NUMA", "PIN_GB"): self.assertNotIn(key, env) + def _kimi_model(self): + other = tempfile.TemporaryDirectory() + self.addCleanup(other.cleanup) + model = Path(other.name) + (model / "config.json").write_text(json.dumps({ + "model_type": "kimi_k3", + "hidden_size": 64, + "num_hidden_layers": 2, + "num_attention_heads": 4, + "q_lora_rank": 8, + "kv_lora_rank": 8, + "qk_nope_head_dim": 8, + "qk_rope_head_dim": 4, + "v_head_dim": 8, + "num_experts": 2, + "linear_attn_config": { + "num_heads": 2, + "head_dim": 8, + "kda_layers": [1], + }, + })) + write_shard(model / "model.safetensors", [ + ("model.embed_tokens.weight", 100), + ("model.layers.0.block_sparse_moe.experts.0.w1.weight", 80), + ("model.layers.0.block_sparse_moe.experts.1.w1.weight", 80), + ("model.layers.1.block_sparse_moe.experts.0.w1.weight", 80), + ("model.layers.1.block_sparse_moe.experts.1.w1.weight", 80), + ]) + return model + + def _glm53_model(self): + other = tempfile.TemporaryDirectory() + self.addCleanup(other.cleanup) + model = Path(other.name) + (model / "config.json").write_text(json.dumps({ + "model_type": "glm5_next", + "text_config": { + "num_hidden_layers": 2, + "n_routed_experts": 2, + "hidden_size": 32, + "num_attention_heads": 4, + "q_lora_rank": 8, + "kv_lora_rank": 8, + "qk_nope_head_dim": 8, + "v_head_dim": 8, + "index_head_dim": 8, + "hc_mult": 2, + "layer_types": ["linear", "full"], + "linear_attn_config": { + "num_heads": 2, + "head_dim": 8, + "short_conv_kernel_size": 2, + }, + }, + })) + write_shard(model / "model.safetensors", [ + ("model.language_model.embed_tokens.weight", 100), + ("model.language_model.layers.1.mlp.experts.0.gate_proj.weight", 80), + ("model.language_model.layers.1.mlp.experts.1.gate_proj.weight", 80), + ]) + return model + + def test_kimi_plan_exports_the_expert_cache_knob_the_engine_reads(self): + """Kimi K3 sizes the expert LRU from K3_EXPERT_GB (default 8 GB). + RAM_GB is only a ceiling, and main() never takes argv as a cap, so + --auto-tier that set RAM_GB and COLI_PLAN_CAP left the 8 GB default + in place on a machine whose plan had hundreds of GB of warm experts.""" + plan = build_plan(self._kimi_model(), context=32, available_memory=32 * GB, + available_disk=1, gpus=[], cpu_sockets=1) + cache = plan["tiers"]["ram"]["expert_cache_bytes"] + self.assertGreater(cache, 0) + expected = f"{cache / GB:.3f}" + self.assertIn("K3_EXPERT_GB", plan["tune"]) + self.assertEqual(plan["tune"]["K3_EXPERT_GB"]["value"], expected) + self.assertNotIn("GLM53_EXPERT_GB", plan["tune"]) + self.assertIn("K3_EXPERT_GB=", format_plan(plan)) + env = environment_for_plan(plan, {}) + self.assertIn("K3_EXPERT_GB", env) + self.assertEqual(env["K3_EXPERT_GB"], expected) + kept = environment_for_plan(plan, {"K3_EXPERT_GB": "3.5"}) + self.assertEqual(kept["K3_EXPERT_GB"], "3.5") + + def test_glm53_plan_exports_the_expert_cache_knob_the_engine_reads(self): + """glm53.c reads GLM53_EXPERT_GB, not RAM_GB. --auto-tier exported + RAM_GB (inert) and relied on COLI_PLAN_CAP as argv, which a direct + engine launch never sees.""" + plan = build_plan(self._glm53_model(), context=32, available_memory=32 * GB, + available_disk=1, gpus=[], cpu_sockets=1) + cache = plan["tiers"]["ram"]["expert_cache_bytes"] + self.assertGreater(cache, 0) + expected = f"{cache / GB:.3f}" + self.assertIn("GLM53_EXPERT_GB", plan["tune"]) + self.assertEqual(plan["tune"]["GLM53_EXPERT_GB"]["value"], expected) + self.assertNotIn("K3_EXPERT_GB", plan["tune"]) + env = environment_for_plan(plan, {}) + self.assertIn("GLM53_EXPERT_GB", env) + self.assertEqual(env["GLM53_EXPERT_GB"], expected) + self.assertNotIn("K3_EXPERT_GB", env) + def test_cpu_socket_count_is_positive(self): self.assertGreaterEqual(cpu_socket_count(), 1)