unsloth/studio/backend/tests/test_mtp_mla_target_ctx.py

# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0

"""MTP draft reserve for MLA models keeps a duplicated target KV context.

llama.cpp's MTP speculative decoding allocates a second full copy of the target
model's KV context (``ctx_tgt=yes``) for draft verification, at f16. On MLA
models (GLM-5.x, DeepSeek, Kimi-K2) that copy is ~the main KV again and dwarfs
the tiny embedded draft head, so omitting it let auto-fit pick a context that
fit on paper but OOMed ``cublasCreate`` at the first decode (e.g. GLM-5.2
UD-IQ1_S advertised the native 1M context on 2x B200, then crashed on the first
generation). Non-MLA MTP (Qwen/Gemma) keeps no such copy and must stay exactly
as #6312 tuned it.
"""

import sys
import types as _types
from pathlib import Path

import pytest

# ---------------------------------------------------------------------------
# Stub heavy/unavailable deps before importing the module under test, so this
# file is order-independent (importing core.inference pulls in orchestrator ->
# structlog, absent in the lightweight test env). Mirrors test_mtp_vram_budget.
# ---------------------------------------------------------------------------

_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
if _BACKEND_DIR not in sys.path:
    sys.path.insert(0, _BACKEND_DIR)

_loggers_stub = _types.ModuleType("loggers")
_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name)
sys.modules.setdefault("loggers", _loggers_stub)

sys.modules.setdefault("structlog", _types.ModuleType("structlog"))

# httpx -- only stub when the real library is missing. Unconditional stubbing
# shadows HTTPError/Response that huggingface_hub.errors imports at load time.
try:
    import httpx as _httpx_real  # noqa: F401
except ImportError:
    _httpx_stub = _types.ModuleType("httpx")
    for _exc_name in (
        "ConnectError",
        "TimeoutException",
        "ReadTimeout",
        "ReadError",
        "RemoteProtocolError",
        "CloseError",
        "HTTPError",
        "RequestError",
    ):
        setattr(_httpx_stub, _exc_name, type(_exc_name, (Exception,), {}))
    _httpx_stub.Timeout = type("Timeout", (), {"__init__": lambda self, *a, **kw: None})
    _httpx_stub.Response = type("Response", (), {})
    _httpx_stub.Client = type(
        "Client",
        (),
        {
            "__init__": lambda self, **kw: None,
            "__enter__": lambda self: self,
            "__exit__": lambda self, *a: None,
        },
    )
    sys.modules["httpx"] = _httpx_stub

from core.inference.llama_cpp import (  # noqa: E402
    LlamaCppBackend,
    _kv_bytes_per_elem,
)

GIB = 1024**3


def _make_mla_backend(
    *,
    n_layers = 79,
    n_kv_heads = 1,
    n_heads = 64,
    kv_key_length = 576,
    kv_value_length = 512,
    kv_lora_rank = 512,
    key_length_mla = 256,
    nextn = 1,
    embedding_length = 6144,
    vocab = 154880,
    native_ctx = 1048576,
):
    """GLM-5.2-class backend: MLA attention + an embedded MTP head."""
    b = LlamaCppBackend.__new__(LlamaCppBackend)
    b._nextn_predict_layers = nextn
    b._n_kv_heads = n_kv_heads
    b._n_heads = n_heads
    b._kv_key_length = kv_key_length
    b._kv_value_length = kv_value_length
    b._embedding_length = embedding_length
    b._n_layers = n_layers
    b._context_length = native_ctx
    b._shared_kv_layers = 0
    b._kv_lora_rank = kv_lora_rank
    b._sliding_window = None
    b._sliding_window_pattern = None
    b._ssm_inner_size = None
    b._full_attention_interval = None
    b._key_length_mla = key_length_mla
    b._n_kv_heads_by_layer = None
    b._kv_key_length_swa = None
    b._kv_value_length_swa = None
    b._draft_backend_cache = None
    b._vocab_size = vocab
    return b


def _make_non_mla_backend(**kw):
    """Qwen3.6-MTP-class embedded head: no MLA (kv_lora_rank is None)."""
    b = _make_mla_backend(
        n_kv_heads = 4,
        n_heads = 24,
        kv_key_length = 256,
        kv_value_length = 256,
        embedding_length = 5120,
        n_layers = 65,
        native_ctx = 262144,
        **kw,
    )
    b._kv_lora_rank = None
    b._key_length_mla = None
    return b


class TestMlaTargetCtxReserve:
    def test_mla_reserve_includes_target_ctx_copy(self):
        b = _make_mla_backend()
        ctx = 1048576
        draft = b._mtp_draft_kv_bytes(ctx)
        overhead = b._estimate_mtp_overhead_bytes(ctx)
        main_kv_f16 = b._estimate_kv_cache_bytes(ctx, "f16")
        # Overhead = embedded draft head + a full f16 copy of the target KV.
        assert overhead == draft + main_kv_f16
        # The copy dominates: GLM-5.2 @1M is a ~2 GiB head next to a ~89 GiB copy.
        assert overhead / GIB > 80
        assert main_kv_f16 > 30 * draft

    def test_target_copy_is_f16_regardless_of_main_cache_type(self):
        # The MTP target context is always f16 in llama.cpp; the reserve must not
        # shrink when the user runs a quantized main KV.
        b = _make_mla_backend()
        ctx = 262144
        f16 = _kv_bytes_per_elem("f16")
        expected_copy = b._estimate_kv_cache_bytes(ctx, "f16")
        assert b._estimate_mtp_overhead_bytes(ctx) == (b._mtp_draft_kv_bytes(ctx) + expected_copy)
        assert f16 == 2.0  # sanity: f16 is 2 bytes/elem

    def test_target_copy_scales_linearly_with_context(self):
        b = _make_mla_backend()
        o_64k = b._estimate_mtp_overhead_bytes(65536)
        o_128k = b._estimate_mtp_overhead_bytes(131072)
        assert o_128k == pytest.approx(2 * o_64k)

    def test_non_mla_embedded_head_unchanged(self):
        # Qwen-class MTP keeps no target copy: overhead == draft KV exactly.
        b = _make_non_mla_backend()
        for ctx in (16384, 131072):
            assert b._estimate_mtp_overhead_bytes(ctx) == b._mtp_draft_kv_bytes(ctx)

    def test_mla_reserve_strictly_larger_than_non_mla_shape(self):
        # Same embedded-head dims, MLA toggled on/off: only MLA adds the copy.
        mla = _make_mla_backend()
        non = _make_mla_backend()
        non._kv_lora_rank = None  # flip MLA off, keep every other dim identical
        ctx = 131072
        assert mla._estimate_mtp_overhead_bytes(ctx) > non._estimate_mtp_overhead_bytes(ctx)


class TestMlaFitPreventsOom:
    """The corrected reserve must actually lower the auto-fit context so the
    config holds at runtime instead of OOMing on the first decode."""

    # 2x B200, mirroring the GLM-5.2 UD-IQ1_S crash (only 2 GPUs were selected).
    AVAIL_MIB = 2 * 182010
    TOTAL_MIB = 2 * 182633
    MODEL_BYTES = 200 * GIB  # ~UD-IQ1_S weight footprint
    REQ_CTX = 1048576

    def test_target_copy_lowers_chosen_context(self):
        b = _make_mla_backend()
        with_copy = b._fit_context_to_vram(
            self.REQ_CTX,
            self.AVAIL_MIB,
            self.MODEL_BYTES,
            mtp_engaged = True,
            total_mib = self.TOTAL_MIB,
            mtp_overhead_fn = lambda c: b._estimate_mtp_overhead_bytes(c) or 0,
        )
        # The old behaviour (draft head only, no target copy) kept the full ctx.
        draft_only = b._fit_context_to_vram(
            self.REQ_CTX,
            self.AVAIL_MIB,
            self.MODEL_BYTES,
            mtp_engaged = True,
            total_mib = self.TOTAL_MIB,
            mtp_overhead_fn = lambda c: (b._mtp_draft_kv_bytes(c) or 0),
        )
        assert draft_only == self.REQ_CTX  # reproduces the over-advertised context
        assert with_copy < self.REQ_CTX  # corrected reserve backs the context off