diff --git a/src/gpu/modal_distill.py b/src/gpu/modal_distill.py
index fd1ddae..22c0c88 100644
--- a/src/gpu/modal_distill.py
+++ b/src/gpu/modal_distill.py
@@ -659,7 +659,7 @@ def distill_sequence(spec_id: str, epochs: int = 3) -> dict:
             vocab_size=259,
             d_model=cfg.get("d_model", 384),
             d_ff=cfg.get("d_ff", 1536),
-            d_kv=cfg.get("d_model", 384) // cfg.get("num_heads", 6),
+            d_kv=cfg.get("d_kv", cfg.get("d_model", 384) // cfg.get("num_heads", 6)),
             num_layers=cfg.get("enc_layers", 8),
             num_decoder_layers=cfg.get("dec_layers", 8),
             num_heads=cfg.get("num_heads", 6),
@@ -1217,7 +1217,7 @@ def distill_microkimi(spec_id: str, epochs: int = 3, calib_batches: int = 64,
             vocab_size=259,
             d_model=cfg.get("d_model", 384),
             d_ff=cfg.get("d_ff", 1536),
-            d_kv=cfg.get("d_model", 384) // cfg.get("num_heads", 6),
+            d_kv=cfg.get("d_kv", cfg.get("d_model", 384) // cfg.get("num_heads", 6)),
             num_layers=cfg.get("enc_layers", 8),
             num_decoder_layers=cfg.get("dec_layers", 8),
             num_heads=cfg.get("num_heads", 6),