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

1"""Context sizing for a chat model, from gguf-parser estimates. 

2 

3:func:`fit_single_ctx` sizes one device against the budget the planner scaled; 

4:func:`fit_split_ctx` sizes a tensor split against the busiest card's headroom 

5rather than the summed pool. Both bisect the same quantized grid. 

6``engine_params.resolve_chat_ctx`` is the entry point for the single-device fit 

7and holds the header-math fallback. See docs/architecture.md (VRAM estimation). 

8""" 

9 

10from __future__ import annotations 

11 

12from collections.abc import Callable, Sequence 

13from pathlib import Path 

14 

15from lilbee.core.config.enums import KvCacheType 

16from lilbee.providers.engine_params import chat_ctx_ceiling 

17from lilbee.providers.fleet.placement import _vram_proportional_split 

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

19from lilbee.providers.model_cache import _DYNAMIC_CTX_FLOOR, _DYNAMIC_CTX_QUANTUM 

20 

21 

22def _largest_fitting_ctx(upper: int, fits: Callable[[int], bool]) -> int: 

23 """The largest quantized per-slot n_ctx at or below *upper* that *fits*. 

24 

25 The grid is ``_DYNAMIC_CTX_FLOOR`` plus whole ``_DYNAMIC_CTX_QUANTUM`` steps, 

26 so every probe is a context the engine launches with cleanly. Returns the 

27 floor when even the floor overflows, which leaves the caller to refuse a 

28 window too small to serve. 

29 """ 

30 if not fits(_DYNAMIC_CTX_FLOOR): 

31 return _DYNAMIC_CTX_FLOOR 

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

33 lo, hi, best = 0, steps, 0 

34 while lo <= hi: 

35 mid = (lo + hi) // 2 

36 if fits(_DYNAMIC_CTX_FLOOR + mid * _DYNAMIC_CTX_QUANTUM): 

37 best, lo = mid, mid + 1 

38 else: 

39 hi = mid - 1 

40 return _DYNAMIC_CTX_FLOOR + best * _DYNAMIC_CTX_QUANTUM 

41 

42 

43def fit_single_ctx( 

44 model_path: Path, 

45 *, 

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

47 slots: int, 

48 available_bytes: int, 

49 gpu_layers: int, 

50 flash_attn: bool, 

51 kv_cache_type: KvCacheType, 

52 unified: bool, 

53 ctx_ceiling: int, 

54 expert_offload: tuple[str, ...], 

55) -> int: 

56 """Largest quantized n_ctx whose gguf-parser estimate fits *available_bytes*. 

57 

58 The estimator prices the cache each layer of this architecture actually 

59 holds, so a linear-attention, sliding-window or MLA model is granted the 

60 window it can serve rather than one budgeted for dense attention everywhere. 

61 

62 *available_bytes* is the budget the caller already scaled 

63 (``planning.plan_sizing_budget``), so no further fraction applies here. 

64 *unified* charges the shared-memory figure, which is the whole resident 

65 footprint on a host whose GPU memory is the system's memory. 

66 *expert_offload* names the tensors the launch moves to system memory, so a 

67 mixture-of-experts model is not charged VRAM for experts it will not hold. 

68 """ 

69 if available_bytes <= 0: 

70 return _DYNAMIC_CTX_FLOOR 

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

72 

73 def _fits(per_slot: int) -> bool: 

74 est = estimate_instance_footprint( 

75 model_path, 

76 ctx=per_slot, 

77 slots=slots, 

78 gpu_layers=gpu_layers, 

79 flash_attn=flash_attn, 

80 kv_cache_type=kv_cache_type, 

81 expert_offload=expert_offload, 

82 ) 

83 return est.footprint(unified=unified) <= available_bytes 

84 

85 return _largest_fitting_ctx(upper, _fits) 

86 

87 

88def fit_split_ctx( 

89 model_path: Path, 

90 *, 

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

92 slots: int, 

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

94 per_device_free_bytes: Sequence[int], 

95 gpu_layers: int, 

96 flash_attn: bool, 

97 kv_cache_type: KvCacheType, 

98 ctx_ceiling: int, 

99 expert_offload: tuple[str, ...], 

100) -> int: 

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

102 

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

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

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

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

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

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

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

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

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

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

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

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

115 

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

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

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

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

120 """ 

121 headrooms = [int(free * usable_vram_fraction()) for free in per_device_free_bytes] 

122 if min(headrooms) <= 0: 

123 return _DYNAMIC_CTX_FLOOR 

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

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

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

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

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

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

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

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

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

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

134 

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

136 est = estimate_instance_footprint( 

137 model_path, 

138 ctx=per_slot, 

139 slots=slots, 

140 gpu_layers=gpu_layers, 

141 flash_attn=flash_attn, 

142 kv_cache_type=kv_cache_type, 

143 tensor_split=ratio, 

144 expert_offload=expert_offload, 

145 ) 

146 shares = est.per_device_vram 

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

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

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

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

151 

152 return _largest_fitting_ctx(upper, _peak_fits)