test(tinytorch): regression coverage for systems/perf module gaps

tinytorch/tests/14_profiling/test_profiler_core.py

@@ -88,6 +88,38 @@ def parameters(self):
         )
 
 
+class TestCountFlopsDispatch:
+    """Test count_flops dispatch logic for non-Linear/Conv2d models."""
+
+    def test_count_flops_routes_unnamed_layers_holder_to_sequential(self):
+        """
+        WHAT: A mock model with a `.layers` attribute but a class name other
+        than 'Sequential' is still routed to the sequential-flops handling.
+
+        WHY: count_flops dispatches on `model_name == 'Sequential' or
+        hasattr(model, 'layers')`, so any object exposing `.layers` should
+        be treated as a container of sub-layers, not fall through to the
+        generic "1 FLOP per element" branch.
+        """
+        class LayerStack:
+            def __init__(self, layers):
+                self.layers = layers
+
+        model = Linear(10, 5)
+        mock = LayerStack([model])
+
+        profiler = Profiler()
+        input_shape = (1, 10)
+
+        dispatched = profiler.count_flops(mock, input_shape)
+        direct = profiler._count_sequential_flops(mock, input_shape)
+
+        assert dispatched == direct, (
+            "count_flops should route objects with a .layers attribute to "
+            "_count_sequential_flops even when the class isn't named 'Sequential'"
+        )
+
+
 class TestLatencyMeasurement:
     """Test timing and latency measurement."""
 

tinytorch/tests/15_quantization/test_quantizer_core.py

@@ -34,6 +34,7 @@
     QuantizedLinear,
     quantize_int8,
     dequantize_int8,
+    quantize_model,
 )
 
 
@@ -143,6 +144,36 @@ def test_quantize_negative_values(self):
                 )
 
 
+class TestQuantizeModelValidation:
+    """Test quantize_model() error paths for unsupported model shapes."""
+
+    def test_quantize_model_rejects_bare_linear(self):
+        """
+        WHAT: Verify quantize_model raises ValueError for a bare Linear
+        layer (not wrapped in a Sequential).
+
+        WHY: quantize_model modifies a model in-place by replacing entries
+        in its .layers list. A bare Linear layer has no container to
+        modify, so it must be rejected instead of silently doing nothing.
+        """
+        linear = Linear(4, 4)
+        with pytest.raises(ValueError):
+            quantize_model(linear)
+
+    def test_quantize_model_rejects_unsupported_type(self):
+        """
+        WHAT: Verify quantize_model raises ValueError for an object with
+        no .layers attribute and that is not a Linear layer.
+
+        WHY: quantize_model only knows how to handle Sequential-style
+        containers (.layers) or a bare Linear layer. Anything else is an
+        unsupported model type and must fail clearly.
+        """
+        unsupported = object()
+        with pytest.raises(ValueError):
+            quantize_model(unsupported)
+
+
 class TestQuantizedLinear:
     """Test the QuantizedLinear layer implementation."""
 

tinytorch/tests/17_acceleration/test_acceleration_core.py

@@ -100,6 +100,61 @@ def test_tiled_matmul_correctness(self):
         )
 
 
+class TestAccelerationShapeValidation:
+    """Test shape validation error paths for matmul acceleration functions."""
+
+    def test_vectorized_matmul_rejects_1d_a(self):
+        """
+        WHAT: Verify vectorized_matmul raises ValueError when a is 1D.
+
+        WHY: Matrix multiplication requires at least 2D operands. A 1D
+        first argument must be rejected with a clear error, not silently
+        misinterpreted.
+        """
+        a = Tensor([1, 2, 3])
+        b = Tensor([[1], [2], [3]])
+        with pytest.raises(ValueError):
+            vectorized_matmul(a, b)
+
+    def test_vectorized_matmul_rejects_inner_dim_mismatch(self):
+        """
+        WHAT: Verify vectorized_matmul raises ValueError on inner dimension
+        mismatch.
+
+        WHY: a.shape[-1] must equal b.shape[-2] for matmul to be valid.
+        A mismatch must be caught, not passed through to a confusing
+        NumPy broadcasting error.
+        """
+        a = Tensor(np.zeros((2, 3)))
+        b = Tensor(np.zeros((4, 2)))
+        with pytest.raises(ValueError):
+            vectorized_matmul(a, b)
+
+    def test_tiled_matmul_rejects_1d_a(self):
+        """
+        WHAT: Verify tiled_matmul raises ValueError when a is 1D.
+
+        WHY: Same shape contract as vectorized_matmul; tiling must not
+        bypass input validation.
+        """
+        a = Tensor([1, 2, 3])
+        b = Tensor([[1], [2], [3]])
+        with pytest.raises(ValueError):
+            tiled_matmul(a, b)
+
+    def test_tiled_matmul_rejects_inner_dim_mismatch(self):
+        """
+        WHAT: Verify tiled_matmul raises ValueError on inner dimension
+        mismatch.
+
+        WHY: Same shape contract as vectorized_matmul.
+        """
+        a = Tensor(np.zeros((2, 3)))
+        b = Tensor(np.zeros((4, 2)))
+        with pytest.raises(ValueError):
+            tiled_matmul(a, b)
+
+
 class TestMemoryOptimization:
     """Test memory-related optimizations."""
 

tinytorch/tests/18_memoization/test_kv_cache_core.py

@@ -28,7 +28,8 @@
 sys.path.insert(0, str(Path(__file__).parent.parent.parent))
 
 from tinytorch.core.tensor import Tensor
-from tinytorch.perf.memoization import KVCache
+from tinytorch.core.transformers import GPT
+from tinytorch.perf.memoization import KVCache, enable_kv_cache, disable_kv_cache
 
 
 class TestKVCacheBasics:
@@ -177,5 +178,45 @@ def test_kv_cache_raises_on_invalid_layer(self):
             cache.update(layer_idx=5, key=K, value=V)  # Invalid layer
 
 
+class TestKVCacheEnableDisable:
+    """Test enable_kv_cache / disable_kv_cache lifecycle edge cases."""
+
+    def test_disable_kv_cache_without_enable_does_not_crash(self):
+        """
+        WHAT: Verify disable_kv_cache on a model that never had
+        enable_kv_cache called on it does not raise, and prints a warning
+        instead.
+
+        WHY: disable_kv_cache should be safe to call defensively (e.g. in
+        cleanup code) even if caching was never turned on.
+        """
+        model = GPT(vocab_size=20, embed_dim=16, num_layers=2, num_heads=2, max_seq_len=32)
+
+        disable_kv_cache(model)
+
+        assert not getattr(model, '_cache_enabled', False)
+
+    def test_disable_kv_cache_twice_is_safe(self):
+        """
+        WHAT: Verify calling disable_kv_cache a second time (after caching
+        was already disabled) also hits the early-return path cleanly
+        instead of raising or corrupting model state.
+
+        WHY: Callers may disable caching defensively more than once (e.g.
+        in overlapping cleanup paths). The second call must be a no-op.
+        """
+        model = GPT(vocab_size=20, embed_dim=16, num_layers=2, num_heads=2, max_seq_len=32)
+
+        enable_kv_cache(model)
+        assert model._cache_enabled
+
+        disable_kv_cache(model)
+        assert not model._cache_enabled
+
+        # Second disable call should not raise and should leave state intact.
+        disable_kv_cache(model)
+        assert not model._cache_enabled
+
+
 if __name__ == "__main__":
     pytest.main([__file__, "-v"])