Coverage for src/lilbee/providers/fleet/ctx.py: 100%

1"""Per-device context sizing for a tensor-split chat model in the fleet. 

2 

3Single-GPU chat ctx lives in ``engine_params.resolve_chat_ctx``; this is its 

4multi-GPU counterpart, sizing the per-slot context against the busiest card's 

5headroom rather than the summed pool. See docs/architecture.md (VRAM estimation). 

6""" 

7 

8from __future__ import annotations 

9 

10from collections.abc import Sequence 

11from pathlib import Path 

12 

13from lilbee.core.config.enums import KvCacheType 

14from lilbee.providers.engine_params import chat_ctx_ceiling 

15from lilbee.providers.fleet.placement import _vram_proportional_split 

16from lilbee.providers.fleet.vram import estimate_instance_footprint, usable_vram_fraction 

17from lilbee.providers.model_cache import _DYNAMIC_CTX_FLOOR, _DYNAMIC_CTX_QUANTUM 

18 

19# Extra VRAM held back on the busiest card on top of the gguf-parser estimate. 

20# In --split-mode layer, gguf-parser ignores --main-gpu and so under-models the 

21# compute graph + logits that real llama.cpp concentrates on the main device. 

22# Measured on 2x4090 (70B Q4_K_M, ctx 5888): the main device lands 0.26 GiB over 

23# its estimate, an order of magnitude inside the usable_vram_fraction already held 

24# back per card, so the reserve stays 0. 

25_MAIN_GPU_SKEW_RESERVE_BYTES = 0 

26 

27 

28def fit_split_ctx( 

29 model_path: Path, 

30 *, 

31 meta: dict[str, str] | None, 

32 slots: int, 

33 ratio: tuple[int, ...], 

34 per_device_free_bytes: Sequence[int], 

35 gpu_layers: int, 

36 flash_attn: bool, 

37 kv_cache_type: KvCacheType, 

38 kv_cache_type_v: KvCacheType, 

39 ctx_ceiling: int, 

40) -> int: 

41 """Largest quantized per-slot n_ctx that fits every card, capped at *ctx_ceiling*. 

42 

43 Binary-searches the gguf-parser estimate at the launch tensor-split *ratio*: 

44 each probe passes a per-slot value, which the estimator charges across the 

45 slot count as the server does, and accepts it when every device's own share 

46 stays under that device's usable headroom. *ctx_ceiling* is the working context the 

47 caller planned for (``planning._placement_estimate_ctx``: a ``cfg.num_ctx`` pin, 

48 else ``cfg.chat_n_ctx_target``); the search never exceeds it, nor the model's 

49 trained context. Note the ceiling bounds the PER-SLOT window, not the total: 

50 placement reserves KV for one full window, while a split whose cards hold 

51 several may serve up to ``_CHAT_SLOTS`` of them. What keeps that honest is 

52 the per-device check below against real free bytes, not the placement 

53 reserve. Falls to the floor when even the floor overflows; plan_launches 

54 refuses a chat launch left at a window below the minimum grounded prompt. 

55 

56 An empty *ratio* is the tight placement, which launches without one so the 

57 engine runs its own fit pass. It is also the estimator's only device-count 

58 signal, so size against a headroom-proportional one here: without it 

59 gguf-parser reports the whole model as a single card and nothing fits. 

60 """ 

61 headrooms = [ 

62 int(free * usable_vram_fraction()) - _MAIN_GPU_SKEW_RESERVE_BYTES 

63 for free in per_device_free_bytes 

64 ] 

65 if min(headrooms) <= 0: 

66 return _DYNAMIC_CTX_FLOOR 

67 if not ratio and len(per_device_free_bytes) > 1: 

68 by_position = {i: float(free) for i, free in enumerate(per_device_free_bytes)} 

69 ratio = _vram_proportional_split(list(by_position), by_position) 

70 # Bound the per-slot search by the planned working context, not just the model's 

71 # trained max: filling VRAM to that max OOM'd large tensor-split models under 

72 # load (a 235B took the full 262144-token ctx and crashed). The caller passes the 

73 # target placement sized its reserve against, so no single sequence exceeds the 

74 # plan; the total across slots can, and is held instead by the per-device 

75 # headroom test, which measures each card's real free bytes at launch. 

76 upper = min(chat_ctx_ceiling(meta, model_path), ctx_ceiling) 

77 

78 def _peak_fits(per_slot: int) -> bool: 

79 est = estimate_instance_footprint( 

80 model_path, 

81 ctx=per_slot, 

82 slots=slots, 

83 gpu_layers=gpu_layers, 

84 flash_attn=flash_attn, 

85 kv_cache_type=kv_cache_type, 

86 kv_cache_type_v=kv_cache_type_v, 

87 tensor_split=ratio, 

88 ) 

89 shares = est.per_device_vram 

90 if len(shares) != len(headrooms): 

91 # No usable per-device breakdown: fall back to peak vs the tightest card. 

92 return est.peak_footprint(unified=False) <= min(headrooms) 

93 return all(share <= room for share, room in zip(shares, headrooms, strict=True)) 

94 

95 if not _peak_fits(_DYNAMIC_CTX_FLOOR): 

96 return _DYNAMIC_CTX_FLOOR 

97 steps = max(0, (upper - _DYNAMIC_CTX_FLOOR) // _DYNAMIC_CTX_QUANTUM) 

98 lo, hi, best = 0, steps, 0 

99 while lo <= hi: 

100 mid = (lo + hi) // 2 

101 if _peak_fits(_DYNAMIC_CTX_FLOOR + mid * _DYNAMIC_CTX_QUANTUM): 

102 best, lo = mid, mid + 1 

103 else: 

104 hi = mid - 1 

105 return _DYNAMIC_CTX_FLOOR + best * _DYNAMIC_CTX_QUANTUM