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_fit`` is the entry point for the single-device fit 

7and holds the header-math fallback; ``planning.fit_chat_ctx`` searches the offload 

8around this window search. See docs/architecture.md (VRAM estimation). 

9""" 

10 

11from __future__ import annotations 

12 

13from collections.abc import Callable, Sequence 

14from pathlib import Path 

15 

16from lilbee.core.config.enums import KvCacheType 

17from lilbee.providers.engine_params import chat_ctx_ceiling 

18from lilbee.providers.fleet.placement import _vram_proportional_split 

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

20from lilbee.providers.model_cache import _DYNAMIC_CTX_FLOOR, _DYNAMIC_CTX_QUANTUM 

21 

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

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

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

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

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

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

28_MAIN_GPU_SKEW_RESERVE_BYTES = 0 

29 

30 

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

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

33 

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

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

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

37 window too small to serve. 

38 """ 

39 if not fits(_DYNAMIC_CTX_FLOOR): 

40 return _DYNAMIC_CTX_FLOOR 

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

42 lo, hi, best = 0, steps, 0 

43 while lo <= hi: 

44 mid = (lo + hi) // 2 

45 if fits(_DYNAMIC_CTX_FLOOR + mid * _DYNAMIC_CTX_QUANTUM): 

46 best, lo = mid, mid + 1 

47 else: 

48 hi = mid - 1 

49 return _DYNAMIC_CTX_FLOOR + best * _DYNAMIC_CTX_QUANTUM 

50 

51 

52def fit_single_ctx( 

53 model_path: Path, 

54 *, 

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

56 slots: int, 

57 available_bytes: int, 

58 gpu_layers: int, 

59 flash_attn: bool, 

60 kv_cache_type: KvCacheType, 

61 kv_cache_type_v: KvCacheType, 

62 unified: bool, 

63 ctx_ceiling: int, 

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

65) -> int: 

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

67 

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

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

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

71 

72 *available_bytes* is the budget the caller already scaled 

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

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

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

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

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

78 """ 

79 if available_bytes <= 0: 

80 return _DYNAMIC_CTX_FLOOR 

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

82 

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

84 est = estimate_instance_footprint( 

85 model_path, 

86 ctx=per_slot, 

87 slots=slots, 

88 gpu_layers=gpu_layers, 

89 flash_attn=flash_attn, 

90 kv_cache_type=kv_cache_type, 

91 kv_cache_type_v=kv_cache_type_v, 

92 expert_offload=expert_offload, 

93 ) 

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

95 

96 return _largest_fitting_ctx(upper, _fits) 

97 

98 

99def fit_split_ctx( 

100 model_path: Path, 

101 *, 

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

103 slots: int, 

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

105 per_device_free_bytes: Sequence[int], 

106 gpu_layers: int, 

107 flash_attn: bool, 

108 kv_cache_type: KvCacheType, 

109 kv_cache_type_v: KvCacheType, 

110 ctx_ceiling: int, 

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

112) -> int: 

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

114 

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

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

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

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

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

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

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

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

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

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

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

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

127 

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

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

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

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

132 """ 

133 headrooms = [ 

134 int(free * usable_vram_fraction()) - _MAIN_GPU_SKEW_RESERVE_BYTES 

135 for free in per_device_free_bytes 

136 ] 

137 if min(headrooms) <= 0: 

138 return _DYNAMIC_CTX_FLOOR 

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

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

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

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

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

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

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

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

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

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

149 

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

151 est = estimate_instance_footprint( 

152 model_path, 

153 ctx=per_slot, 

154 slots=slots, 

155 gpu_layers=gpu_layers, 

156 flash_attn=flash_attn, 

157 kv_cache_type=kv_cache_type, 

158 kv_cache_type_v=kv_cache_type_v, 

159 tensor_split=ratio, 

160 expert_offload=expert_offload, 

161 ) 

162 shares = est.per_device_vram 

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

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

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

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

167 

168 return _largest_fitting_ctx(upper, _peak_fits)