Test that retention_ratio prevents skipping even if error is low.
(self)
| 139 | ) |
| 140 | |
| 141 | def test_mag_cache_retention(self): |
| 142 | """Test that retention_ratio prevents skipping even if error is low.""" |
| 143 | model = DummyTransformer() |
| 144 | # Ratios that imply 0 error, so it *would* skip if retention allowed it |
| 145 | ratios = np.array([1.0, 1.0]) |
| 146 | |
| 147 | config = MagCacheConfig( |
| 148 | threshold=100.0, |
| 149 | num_inference_steps=2, |
| 150 | retention_ratio=1.0, # Force retention for ALL steps |
| 151 | mag_ratios=ratios, |
| 152 | ) |
| 153 | |
| 154 | apply_mag_cache(model, config) |
| 155 | self._set_context(model, "test_context") |
| 156 | |
| 157 | # Step 0 |
| 158 | model(torch.tensor([[[10.0]]])) |
| 159 | |
| 160 | # Step 1: Should COMPUTE (44.0) not SKIP (41.0) because of retention |
| 161 | input_t1 = torch.tensor([[[11.0]]]) |
| 162 | output_t1 = model(input_t1) |
| 163 | |
| 164 | self.assertTrue( |
| 165 | torch.allclose(output_t1, torch.tensor([[[44.0]]])), |
| 166 | f"Expected Compute (44.0) due to retention, got {output_t1.item()}", |
| 167 | ) |
| 168 | |
| 169 | def test_mag_cache_tuple_outputs(self): |
| 170 | """Test compatibility with models returning (hidden, encoder_hidden) like Flux.""" |
nothing calls this directly
no test coverage detected