wehub-resource-sync
511 lines
21 KiB
Python
511 lines
21 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Tests for memory_monitor module (SSD-only mode)."""
|
|
|
|
from unittest.mock import MagicMock
|
|
|
|
import pytest
|
|
|
|
from omlx.memory_monitor import (
|
|
_SDPA_FALLBACK_SCORE_DTYPE_SIZE,
|
|
_SDPA_FULL_SUPPORTED_HEAD_DIMS,
|
|
_SDPA_VECTOR_QUERY_TOKEN_THRESHOLD,
|
|
_SDPA_VECTOR_SUPPORTED_HEAD_DIMS,
|
|
MemoryInfo,
|
|
MemoryMonitor,
|
|
)
|
|
from omlx.utils.hardware import format_bytes
|
|
|
|
|
|
class TestMemoryInfo:
|
|
"""Tests for MemoryInfo dataclass."""
|
|
|
|
def test_create_memory_info(self):
|
|
"""Test creating MemoryInfo."""
|
|
info = MemoryInfo(
|
|
total_bytes=16 * 1024**3,
|
|
used_bytes=8 * 1024**3,
|
|
available_bytes=8 * 1024**3,
|
|
utilization=0.5,
|
|
)
|
|
assert info.total_bytes == 16 * 1024**3
|
|
assert info.used_bytes == 8 * 1024**3
|
|
assert info.available_bytes == 8 * 1024**3
|
|
assert info.utilization == 0.5
|
|
|
|
def test_memory_info_zero_usage(self):
|
|
"""Test MemoryInfo with zero usage."""
|
|
info = MemoryInfo(
|
|
total_bytes=16 * 1024**3,
|
|
used_bytes=0,
|
|
available_bytes=16 * 1024**3,
|
|
utilization=0.0,
|
|
)
|
|
assert info.used_bytes == 0
|
|
assert info.utilization == 0.0
|
|
|
|
|
|
class TestMemoryMonitor:
|
|
"""Test MemoryMonitor class for SSD-only mode."""
|
|
|
|
def test_init_with_required_params(self):
|
|
"""Test initialization with required parameters."""
|
|
max_kv_cache = 2 * 1024**3 # 2GB
|
|
monitor = MemoryMonitor(max_kv_cache_memory=max_kv_cache)
|
|
assert monitor.max_kv_cache_memory == max_kv_cache
|
|
|
|
def test_init_invalid_max_kv_cache_memory_zero(self):
|
|
"""Test initialization with zero max_kv_cache_memory."""
|
|
with pytest.raises(ValueError, match="max_kv_cache_memory"):
|
|
MemoryMonitor(max_kv_cache_memory=0)
|
|
|
|
def test_init_invalid_max_kv_cache_memory_negative(self):
|
|
"""Test initialization with negative max_kv_cache_memory."""
|
|
with pytest.raises(ValueError, match="max_kv_cache_memory"):
|
|
MemoryMonitor(max_kv_cache_memory=-1)
|
|
|
|
def test_eviction_enabled_property_default_true(self):
|
|
"""The default ``eviction_enabled=True`` makes the
|
|
public-facing predicate True so the existing tiered-cache
|
|
path keeps working without changes."""
|
|
monitor = MemoryMonitor(max_kv_cache_memory=1024**3)
|
|
assert monitor.eviction_enabled is True
|
|
|
|
def test_eviction_enabled_property_false_in_ssd_only_mode(self):
|
|
"""Paged-SSD-only mode passes ``eviction_enabled=False``; the
|
|
public predicate must surface that so Scheduler can branch on
|
|
it (avoiding the RuntimeError from estimate_blocks_to_free)."""
|
|
monitor = MemoryMonitor(max_kv_cache_memory=None, eviction_enabled=False)
|
|
assert monitor.eviction_enabled is False
|
|
|
|
def test_get_memory_info(self):
|
|
"""Test get_memory_info returns valid data."""
|
|
monitor = MemoryMonitor(max_kv_cache_memory=1024**3)
|
|
info = monitor.get_memory_info()
|
|
|
|
assert isinstance(info, MemoryInfo)
|
|
assert info.total_bytes == monitor.max_memory
|
|
# In SSD-only mode, used_bytes is always 0
|
|
assert info.used_bytes == 0
|
|
assert info.available_bytes == monitor.max_memory
|
|
assert info.utilization == 0.0
|
|
|
|
def test_get_memory_info_throttling(self):
|
|
"""Test that memory info checks are throttled."""
|
|
monitor = MemoryMonitor(max_kv_cache_memory=1024**3, check_interval=10.0)
|
|
|
|
# First call
|
|
info1 = monitor.get_memory_info()
|
|
# Second call within interval should return cached value
|
|
info2 = monitor.get_memory_info()
|
|
|
|
# Should be the same object (cached)
|
|
assert info1 is info2
|
|
|
|
def test_is_under_pressure_always_false(self):
|
|
"""Test is_under_pressure always returns False in SSD-only mode."""
|
|
monitor = MemoryMonitor(max_kv_cache_memory=10000)
|
|
# In SSD-only mode, always returns False
|
|
assert not monitor.is_under_pressure()
|
|
|
|
def test_bytes_to_free_always_zero(self):
|
|
"""Test bytes_to_free always returns 0 in SSD-only mode."""
|
|
monitor = MemoryMonitor(max_kv_cache_memory=10000)
|
|
# In SSD-only mode, always returns 0
|
|
assert monitor.bytes_to_free() == 0
|
|
|
|
def test_set_model_info(self):
|
|
"""Test setting model information."""
|
|
monitor = MemoryMonitor(max_kv_cache_memory=1024**3)
|
|
|
|
monitor.set_model_info(
|
|
num_layers=32,
|
|
num_kv_heads=8,
|
|
head_dim=128,
|
|
dtype_size=2,
|
|
)
|
|
|
|
# Internal state should be set
|
|
assert monitor._num_layers == 32
|
|
assert monitor._num_kv_heads == 8
|
|
assert monitor._head_dim == 128
|
|
assert monitor._dtype_size == 2
|
|
|
|
def test_estimate_block_memory(self):
|
|
"""Test block memory estimation."""
|
|
monitor = MemoryMonitor(max_kv_cache_memory=1024**3)
|
|
|
|
# Set model info
|
|
monitor.set_model_info(
|
|
num_layers=32,
|
|
num_kv_heads=8,
|
|
head_dim=128,
|
|
dtype_size=2,
|
|
)
|
|
|
|
# Estimate for 64 tokens
|
|
estimate = monitor.estimate_block_memory(64)
|
|
# Expected: 64 * 8 * 128 * 2 * 2 (keys+values) * 32 layers
|
|
expected = 64 * 8 * 128 * 2 * 2 * 32
|
|
assert estimate == expected
|
|
|
|
def test_estimate_block_memory_default_values(self):
|
|
"""Test block memory estimation with default values."""
|
|
monitor = MemoryMonitor(max_kv_cache_memory=1024**3)
|
|
|
|
# Without setting model info, should use defaults
|
|
estimate = monitor.estimate_block_memory(64)
|
|
# Default: 32 layers, 8 kv_heads, 128 head_dim, 2 dtype_size
|
|
expected = 64 * 8 * 128 * 2 * 2 * 32
|
|
assert estimate == expected
|
|
|
|
def test_estimate_block_memory_with_overrides(self):
|
|
"""Test block memory estimation with parameter overrides."""
|
|
monitor = MemoryMonitor(max_kv_cache_memory=1024**3)
|
|
monitor.set_model_info(
|
|
num_layers=32,
|
|
num_kv_heads=8,
|
|
head_dim=128,
|
|
dtype_size=2,
|
|
)
|
|
|
|
# Override some parameters
|
|
estimate = monitor.estimate_block_memory(
|
|
block_size=32,
|
|
num_layers=16, # Override
|
|
dtype_size=4, # Override
|
|
)
|
|
expected = 32 * 8 * 128 * 4 * 2 * 16
|
|
assert estimate == expected
|
|
|
|
def test_estimate_blocks_to_free(self):
|
|
"""Test estimation of blocks to free."""
|
|
monitor = MemoryMonitor(max_kv_cache_memory=1024**3)
|
|
monitor.set_model_info(
|
|
num_layers=32,
|
|
num_kv_heads=8,
|
|
head_dim=128,
|
|
dtype_size=2,
|
|
)
|
|
|
|
block_size = 64
|
|
block_mem = monitor.estimate_block_memory(block_size)
|
|
|
|
# Need to free 10 blocks worth
|
|
bytes_to_free = block_mem * 10
|
|
num_blocks = monitor.estimate_blocks_to_free(bytes_to_free, block_size)
|
|
assert num_blocks == 10
|
|
|
|
def test_estimate_blocks_to_free_rounds_up(self):
|
|
"""Test that blocks to free rounds up."""
|
|
monitor = MemoryMonitor(max_kv_cache_memory=1024**3)
|
|
monitor.set_model_info(
|
|
num_layers=32,
|
|
num_kv_heads=8,
|
|
head_dim=128,
|
|
dtype_size=2,
|
|
)
|
|
|
|
block_size = 64
|
|
block_mem = monitor.estimate_block_memory(block_size)
|
|
|
|
# Need to free slightly more than 9 blocks
|
|
bytes_to_free = block_mem * 9 + 1
|
|
num_blocks = monitor.estimate_blocks_to_free(bytes_to_free, block_size)
|
|
assert num_blocks == 10 # Should round up
|
|
|
|
def test_get_stats(self):
|
|
"""Test get_stats returns dict with expected keys."""
|
|
monitor = MemoryMonitor(max_kv_cache_memory=1024**3)
|
|
stats = monitor.get_stats()
|
|
|
|
assert "total_bytes" in stats
|
|
assert "used_bytes" in stats
|
|
assert "available_bytes" in stats
|
|
assert "utilization" in stats
|
|
assert "max_kv_cache_memory" in stats
|
|
assert "total_formatted" in stats
|
|
assert "used_formatted" in stats
|
|
assert "available_formatted" in stats
|
|
# In SSD-only mode, used_bytes should be 0
|
|
assert stats["used_bytes"] == 0
|
|
|
|
def test_format_bytes(self):
|
|
"""Test format_bytes utility function."""
|
|
assert "1.00 KB" == format_bytes(1024)
|
|
assert "1.00 MB" == format_bytes(1024 * 1024)
|
|
assert "1.00 GB" == format_bytes(1024 * 1024 * 1024)
|
|
assert "512 B" == format_bytes(512)
|
|
|
|
def test_repr(self):
|
|
"""Test string representation."""
|
|
monitor = MemoryMonitor(max_kv_cache_memory=2 * 1024**3)
|
|
repr_str = repr(monitor)
|
|
assert "MemoryMonitor" in repr_str
|
|
assert "max_kv_cache" in repr_str
|
|
assert "used" in repr_str
|
|
|
|
def test_properties(self):
|
|
"""Test property accessors."""
|
|
max_kv_cache = 2 * 1024**3
|
|
monitor = MemoryMonitor(max_kv_cache_memory=max_kv_cache)
|
|
|
|
assert monitor.max_kv_cache_memory == max_kv_cache
|
|
assert monitor.max_memory > 0
|
|
|
|
def test_set_paged_cache_manager(self):
|
|
"""Test setting paged cache manager."""
|
|
monitor = MemoryMonitor(max_kv_cache_memory=1024**3)
|
|
|
|
mock_manager = MagicMock()
|
|
monitor.set_paged_cache_manager(mock_manager, block_size=128)
|
|
|
|
assert monitor._paged_cache_manager is mock_manager
|
|
assert monitor._block_size == 128
|
|
|
|
def test_set_baseline_memory(self):
|
|
"""Test setting baseline memory."""
|
|
monitor = MemoryMonitor(max_kv_cache_memory=1024**3)
|
|
|
|
# This should not raise (uses MLX if available, otherwise sets to 0)
|
|
monitor.set_baseline_memory()
|
|
|
|
def test_set_request_stats(self):
|
|
"""Test setting request stats."""
|
|
monitor = MemoryMonitor(max_kv_cache_memory=1024**3)
|
|
|
|
monitor.set_request_stats(running=5, waiting=10)
|
|
|
|
assert monitor._running_requests == 5
|
|
assert monitor._waiting_requests == 10
|
|
|
|
def test_check_interval_parameter(self):
|
|
"""Test check_interval parameter."""
|
|
monitor = MemoryMonitor(
|
|
max_kv_cache_memory=1024**3,
|
|
check_interval=5.0,
|
|
)
|
|
|
|
assert monitor._check_interval == 5.0
|
|
|
|
|
|
class TestEstimatePrefillPeakBytes:
|
|
"""Tests for estimate_prefill_peak_bytes (KV + SDPA only)."""
|
|
|
|
def _make_monitor(self, head_dim=128, n_attn=32, n_kv=4, n_layers=62):
|
|
m = MemoryMonitor(max_kv_cache_memory=10 * 1024**3)
|
|
m.set_model_info(
|
|
num_layers=n_layers,
|
|
num_kv_heads=n_kv,
|
|
head_dim=head_dim,
|
|
dtype_size=2,
|
|
num_attention_heads=n_attn,
|
|
)
|
|
return m
|
|
|
|
def _expected_output_sdpa(self, n_q, query_tokens, head_dim):
|
|
return n_q * query_tokens * head_dim * 4
|
|
|
|
def _expected_fallback_sdpa(self, n_q, query_tokens, kv_len, head_dim):
|
|
scores = n_q * query_tokens * kv_len * _SDPA_FALLBACK_SCORE_DTYPE_SIZE
|
|
output = n_q * query_tokens * head_dim * 4
|
|
return scores + output
|
|
|
|
def test_returns_zero_when_model_info_missing(self):
|
|
m = MemoryMonitor(max_kv_cache_memory=10 * 1024**3)
|
|
assert m.estimate_prefill_peak_bytes(32768, 2048) == 0
|
|
|
|
def test_returns_zero_when_no_new_tokens(self):
|
|
# Fully-prefix-cached request: nothing to prefill, peak is 0.
|
|
m = self._make_monitor()
|
|
assert m.estimate_prefill_peak_bytes(0, 2048, cached_tokens=32768) == 0
|
|
|
|
def test_fused_full_prefill_head_dim_128(self):
|
|
# head_dim=128 is supported by the fused full prefill kernel.
|
|
m = self._make_monitor(head_dim=128, n_attn=32, n_kv=4, n_layers=62)
|
|
peak = m.estimate_prefill_peak_bytes(32768, 2048)
|
|
# KV: 62 layers * 4 kv_heads * 128 dim * 2 bytes * 2 (k+v) * 32768 ≈ 4.0 GB
|
|
# SDPA fused: n_attn * chunk * head_dim * 4 = 32*2048*128*4 ≈ 32 MB
|
|
# Total ≈ 4 GB
|
|
assert 3 * 1024**3 < peak < 5 * 1024**3
|
|
|
|
def test_prefill_head_dim_256_uses_full_score_fallback(self):
|
|
# head_dim=256 is vector-kernel-supported, but not full-prefill-supported.
|
|
m = self._make_monitor(head_dim=256, n_attn=8, n_kv=4, n_layers=48)
|
|
peak = m.estimate_prefill_peak_bytes(32768, 2048)
|
|
expected_sdpa = self._expected_fallback_sdpa(8, 2048, 32768, 256)
|
|
expected_kv = m.estimate_prompt_kv_bytes(32768)
|
|
assert peak == expected_sdpa + expected_kv
|
|
assert expected_sdpa > 8 * 2048 * 256 * 2
|
|
|
|
def test_sdpa_fallback_scores_track_compute_dtype(self):
|
|
# The unfused score matrix is materialized at the model's compute
|
|
# dtype, not fp32 and not the (possibly fractional TurboQuant) KV width.
|
|
# fp32 model -> 4 bytes/elem; bf16/fp16 -> 2.
|
|
def _scores(monitor, n_q, chunk, kv, hd):
|
|
out = n_q * chunk * hd * 4
|
|
return monitor._estimate_sdpa_activation_bytes(chunk, kv) - out
|
|
|
|
n_q, chunk, kv, hd = 8, 2048, 32768, 256
|
|
m_bf16 = MemoryMonitor(max_kv_cache_memory=10 * 1024**3)
|
|
m_bf16.set_model_info(
|
|
num_layers=48, num_kv_heads=4, head_dim=hd,
|
|
num_attention_heads=n_q, compute_dtype_size=2,
|
|
)
|
|
m_fp32 = MemoryMonitor(max_kv_cache_memory=10 * 1024**3)
|
|
m_fp32.set_model_info(
|
|
num_layers=48, num_kv_heads=4, head_dim=hd,
|
|
num_attention_heads=n_q, compute_dtype_size=4,
|
|
)
|
|
assert _scores(m_bf16, n_q, chunk, kv, hd) == n_q * chunk * kv * 2
|
|
assert _scores(m_fp32, n_q, chunk, kv, hd) == n_q * chunk * kv * 4
|
|
|
|
def test_sdpa_score_dtype_ignores_fractional_kv_width(self):
|
|
# TurboQuant sets a fractional KV dtype_size; the score matrix must
|
|
# still be charged at the compute dtype, not ~0.5 bytes/elem.
|
|
n_q, chunk, kv, hd = 8, 2048, 32768, 256
|
|
m = MemoryMonitor(max_kv_cache_memory=10 * 1024**3)
|
|
m.set_model_info(
|
|
num_layers=48, num_kv_heads=4, head_dim=hd, dtype_size=0.5,
|
|
num_attention_heads=n_q, compute_dtype_size=2,
|
|
)
|
|
out = n_q * chunk * hd * 4
|
|
scores = m._estimate_sdpa_activation_bytes(chunk, kv) - out
|
|
assert scores == n_q * chunk * kv * 2
|
|
|
|
def test_sdpa_fallback_accounts_for_cached_kv_span(self):
|
|
"""Regression for M3: SDPA fallback spans the FULL prompt (cached + new),
|
|
not just new_tokens. A heavily-cached long-context request previously
|
|
slipped through with under-counted peak.
|
|
"""
|
|
m = self._make_monitor(head_dim=256, n_attn=8, n_kv=4, n_layers=48)
|
|
# Same total prompt (100k), different cache split:
|
|
# - All-new: cached=0, new=100k
|
|
# - Heavy cache: cached=99k, new=1k
|
|
all_new = m.estimate_prefill_peak_bytes(100 * 1024, 2048)
|
|
heavy_cache = m.estimate_prefill_peak_bytes(1024, 2048, cached_tokens=99 * 1024)
|
|
expected_heavy_sdpa = self._expected_fallback_sdpa(8, 1024, 100 * 1024, 256)
|
|
expected_heavy = expected_heavy_sdpa + m.estimate_prompt_kv_bytes(1024)
|
|
assert heavy_cache == expected_heavy
|
|
assert (
|
|
heavy_cache > 900 * 1024**2
|
|
), f"heavy-cache peak under-counted: {heavy_cache / 1024**2:.0f} MB"
|
|
# And the all-new case (larger eff_chunk = 2048 but same kv_len)
|
|
# should be larger overall because both KV growth and scores
|
|
# widen with new_tokens.
|
|
assert all_new > heavy_cache
|
|
|
|
def test_scales_linearly_with_token_count(self):
|
|
m = self._make_monitor()
|
|
p8k = m.estimate_prefill_peak_bytes(8 * 1024, 2048)
|
|
p32k = m.estimate_prefill_peak_bytes(32 * 1024, 2048)
|
|
# KV grows linearly with tokens; SDPA fused doesn't depend on
|
|
# total_tokens. KV dominates here, so 32k/8k ≈ 4x.
|
|
assert p32k > p8k
|
|
ratio = p32k / p8k
|
|
assert 3.5 < ratio < 4.5
|
|
|
|
def test_sdpa_fallback_scales_with_context_length(self):
|
|
# Unsupported full-prefill head dims: SDPA peak ∝ query_len * total_tokens.
|
|
# When chunk is fixed (2048), peak grows linearly with total_tokens
|
|
# plus KV grows linearly too. Doubling tokens should ~double peak.
|
|
m = self._make_monitor(head_dim=256, n_attn=8, n_kv=4, n_layers=48)
|
|
p16k = m.estimate_prefill_peak_bytes(16 * 1024, 2048)
|
|
p32k = m.estimate_prefill_peak_bytes(32 * 1024, 2048)
|
|
ratio = p32k / p16k
|
|
assert 1.8 < ratio < 2.2
|
|
|
|
def test_eff_chunk_capped_at_new_tokens(self):
|
|
"""Short prompts (smaller than chunk_size) must not be charged
|
|
the full chunk_size width — the effective chunk is bounded by
|
|
the number of remaining new tokens. Regression for the constant-
|
|
factor over-count on small prompts.
|
|
"""
|
|
m = self._make_monitor(head_dim=256, n_attn=8, n_kv=4, n_layers=48)
|
|
# 100-token prompt; chunk_size=2048. eff_chunk should be 100,
|
|
# not 2048 — so the query width is 100, not the default step size.
|
|
peak = m.estimate_prefill_peak_bytes(100, 2048)
|
|
# KV: 48*4*256*2*2*100 ≈ 19 MB. SDPA is small here. Total < 25 MB.
|
|
assert peak < 25 * 1024**2, (
|
|
f"short-prompt peak suggests chunk wasn't clamped: "
|
|
f"{peak / 1024**2:.0f} MB"
|
|
)
|
|
|
|
def test_no_python_overhead_constant(self):
|
|
# estimator must NOT include cache_pool_overhead or python_overhead
|
|
# magic constants — those are absorbed by enforcer hard_threshold.
|
|
# If a small prompt returns >2 GB on a small model, that's a sign
|
|
# someone added back the magic constants.
|
|
m = self._make_monitor(head_dim=128, n_attn=8, n_kv=2, n_layers=8)
|
|
peak = m.estimate_prefill_peak_bytes(512, 2048)
|
|
# KV: 8*2*128*2*2*512 ≈ 4 MB. SDPA fused: 8*512*128*4 ≈ 2 MB. Total ≈ 6 MB.
|
|
assert peak < 100 * 1024**2, f"unexpected large peak: {peak / 1024**2:.1f} MB"
|
|
|
|
def test_cached_tokens_extends_sdpa_span(self):
|
|
# Unsupported full-prefill head dims span cached+new tokens.
|
|
# A request with a big prefix-cache hit (small new suffix) must still
|
|
# estimate the SDPA transient over the full span, not just new_tokens.
|
|
m = self._make_monitor(head_dim=256, n_attn=16, n_kv=2, n_layers=40)
|
|
# 2k new on top of 30k cached → SDPA span is 32k, query is 2k.
|
|
with_cache = m.estimate_prefill_peak_bytes(2048, 2048, cached_tokens=30 * 1024)
|
|
# Same new_tokens, no cache → SDPA span is only 2k.
|
|
without_cache = m.estimate_prefill_peak_bytes(2048, 2048, cached_tokens=0)
|
|
# The output buffer and KV growth are identical; only the score-matrix
|
|
# K dimension changes.
|
|
sdpa_with = self._expected_fallback_sdpa(16, 2048, 2048 + 30 * 1024, 256)
|
|
sdpa_without = self._expected_fallback_sdpa(16, 2048, 2048, 256)
|
|
assert with_cache - without_cache == sdpa_with - sdpa_without
|
|
assert with_cache > without_cache * 2
|
|
|
|
def test_cached_tokens_default_matches_no_cache(self):
|
|
# Omitting cached_tokens must reproduce the pre-change behavior so the
|
|
# no-cache path (cached=0) is a strict regression guard.
|
|
m = self._make_monitor(head_dim=256, n_attn=8, n_kv=4, n_layers=48)
|
|
assert m.estimate_prefill_peak_bytes(
|
|
32768, 2048
|
|
) == m.estimate_prefill_peak_bytes(32768, 2048, cached_tokens=0)
|
|
|
|
def test_query_len_capped_at_new_tokens(self):
|
|
# When new_tokens < chunk_size the last (only) chunk's query length is
|
|
# new_tokens, not the full step size.
|
|
m = self._make_monitor(head_dim=256, n_attn=8, n_kv=4, n_layers=48)
|
|
# 512 new on top of 10k cached: query=512, span=10k+512.
|
|
peak = m.estimate_prefill_peak_bytes(512, 2048, cached_tokens=10 * 1024)
|
|
expected_sdpa = self._expected_fallback_sdpa(8, 512, 512 + 10 * 1024, 256)
|
|
expected_kv = m.estimate_prompt_kv_bytes(512)
|
|
assert peak == expected_sdpa + expected_kv
|
|
|
|
def test_sdpa_dispatch_constants_match_mlx_use_fallback(self):
|
|
assert _SDPA_VECTOR_QUERY_TOKEN_THRESHOLD == 8
|
|
assert frozenset({64, 80, 128}) == _SDPA_FULL_SUPPORTED_HEAD_DIMS
|
|
assert frozenset({64, 96, 128, 256}) == _SDPA_VECTOR_SUPPORTED_HEAD_DIMS
|
|
|
|
def test_vector_path_head_dim_256_is_output_only_for_short_query(self):
|
|
m = self._make_monitor(head_dim=256, n_attn=8, n_kv=4, n_layers=48)
|
|
assert m.estimate_chunk_transient_bytes(4, 10_000) == (
|
|
self._expected_output_sdpa(8, 4, 256)
|
|
)
|
|
|
|
def test_vector_path_head_dim_80_falls_back(self):
|
|
m = self._make_monitor(head_dim=80, n_attn=8, n_kv=4, n_layers=48)
|
|
assert m.estimate_chunk_transient_bytes(4, 10_000) == (
|
|
self._expected_fallback_sdpa(8, 4, 10_000, 80)
|
|
)
|
|
|
|
def test_full_prefill_head_dim_80_is_output_only(self):
|
|
m = self._make_monitor(head_dim=80, n_attn=8, n_kv=4, n_layers=48)
|
|
assert m.estimate_chunk_transient_bytes(512, 10_000) == (
|
|
self._expected_output_sdpa(8, 512, 80)
|
|
)
|
|
|
|
def test_full_prefill_head_dim_96_falls_back(self):
|
|
m = self._make_monitor(head_dim=96, n_attn=8, n_kv=4, n_layers=48)
|
|
assert m.estimate_chunk_transient_bytes(512, 10_000) == (
|
|
self._expected_fallback_sdpa(8, 512, 10_000, 96)
|
|
)
|
|
|
|
def test_vector_path_gqa_limit_falls_back(self):
|
|
m = self._make_monitor(head_dim=256, n_attn=64, n_kv=1, n_layers=48)
|
|
assert m.estimate_chunk_transient_bytes(1, 10_000) == (
|
|
self._expected_fallback_sdpa(64, 1, 10_000, 256)
|
|
)
|