| """Batch isolation, video ordering and memory recovery without weights or paid requests.""" |
|
|
| import sys |
| from pathlib import Path |
| from types import SimpleNamespace |
|
|
| import pytest |
| import torch |
|
|
| sys.path.insert(0, str(Path(__file__).parent)) |
| from flora_acceleration import collate_inputs, decode_rows, generate_many, require_fast_kernels |
|
|
|
|
| class Tokenizer: |
| pad_token_id = 0 |
| eos_token_id = 2 |
|
|
| def decode(self, ids, **kwargs): |
| return " ".join(map(str, ids)) |
|
|
|
|
| def request(ids, videos): |
| return { |
| "input_ids": torch.tensor([ids]), |
| "attention_mask": torch.ones(1, len(ids), dtype=torch.long), |
| "mm_token_type_ids": torch.ones(1, len(ids), dtype=torch.long), |
| "pixel_values_videos": torch.tensor(videos).reshape(-1, 1), |
| "video_grid_thw": torch.tensor([[1, 1, 1]] * len(videos)), |
| } |
|
|
|
|
| def test_batch_preserves_video_order_and_masks_padding(): |
| result = collate_inputs([request([5, 6], [10, 11]), request([7, 8, 9], [20, 21, 22])], 0) |
| assert result["input_ids"].tolist() == [[0, 5, 6], [7, 8, 9]] |
| assert result["attention_mask"].tolist() == [[0, 1, 1], [1, 1, 1]] |
| assert result["mm_token_type_ids"].tolist() == [[0, 1, 1], [1, 1, 1]] |
| assert result["pixel_values_videos"].flatten().tolist() == [10, 11, 20, 21, 22] |
| assert result["video_grid_thw"].shape == (5, 3) |
| with pytest.raises(ValueError, match="schema"): |
| collate_inputs([{**request([4], [10]), "position_ids": torch.tensor([[0]])}], 0) |
|
|
|
|
| def test_each_row_has_its_own_eos_and_no_padding_in_caption(): |
| assert decode_rows(torch.tensor([[4, 5, 10, 2, 0], [6, 7, 11, 12, 2]]), 2, Tokenizer()) == ["10", "11 12"] |
| with pytest.raises(ValueError, match="incomplete"): |
| decode_rows(torch.tensor([[4, 5, 10, 0, 0]]), 2, Tokenizer()) |
|
|
|
|
| @pytest.mark.parametrize("fixed", [False, True]) |
| def test_bucketing_returns_original_order_and_splits_oom(monkeypatch, fixed): |
| sizes = [] |
|
|
| class Model: |
| def parameters(self): |
| return iter([torch.zeros(1)]) |
|
|
| def generate(self, input_ids, **kwargs): |
| sizes.append(len(input_ids)) |
| if len(input_ids) > 1: |
| raise torch.OutOfMemoryError("synthetic memory pressure") |
| return torch.cat([input_ids, input_ids[:, -1:], torch.tensor([[2]])], dim=1) |
|
|
| monkeypatch.setattr(torch.cuda, "empty_cache", lambda: None) |
| inputs = [request([5, 6, 7, 8, 9], [3]), request([8, 9, 10, 11], [2]), request([10, 11, 12, 13, 14], [1])] |
| options = ( |
| {"batch_size": 16, "fixed_batch_size": 16, "minimum_width": 128, "pad_batch_power2": True} |
| if fixed |
| else {} |
| ) |
| assert generate_many(Model(), SimpleNamespace(tokenizer=Tokenizer()), inputs, 4, **options) == [ |
| "9", |
| "11", |
| "14", |
| ] |
| assert sizes == [16 if fixed else 3, 1, 2, 1, 1] |
|
|
|
|
| def test_fast_profile_cannot_silently_fall_back_to_torch(monkeypatch): |
| monkeypatch.setattr(torch.cuda, "is_available", lambda: False) |
| with pytest.raises(RuntimeError, match="Accelerated Flora needs"): |
| require_fast_kernels() |
|
|
|
|
| def test_tiny_hybrid_gdn_model_keeps_requests_independent_in_left_padded_batch(): |
| """Real CPU forward/generation through GDN + full attention, with no downloaded assets.""" |
| from transformers import Qwen3_5ForCausalLM, Qwen3_5TextConfig |
|
|
| torch.manual_seed(42) |
| config = Qwen3_5TextConfig( |
| vocab_size=64, |
| hidden_size=32, |
| intermediate_size=64, |
| num_hidden_layers=4, |
| num_attention_heads=4, |
| num_key_value_heads=2, |
| head_dim=8, |
| linear_num_key_heads=2, |
| linear_num_value_heads=2, |
| linear_key_head_dim=8, |
| linear_value_head_dim=8, |
| layer_types=["linear_attention"] * 3 + ["full_attention"], |
| max_position_embeddings=128, |
| pad_token_id=0, |
| eos_token_id=2, |
| rope_parameters={"rope_type": "default", "rope_theta": 10000, "mrope_section": [1, 1, 2]}, |
| ) |
| model = Qwen3_5ForCausalLM(config).eval() |
| model.generation_config.forced_eos_token_id = 2 |
| inputs = [ |
| {"input_ids": torch.tensor([ids]), "attention_mask": torch.ones(1, len(ids), dtype=torch.long)} |
| for ids in ([5, 6, 7, 8], [9, 10, 11]) |
| ] |
| processor = SimpleNamespace(tokenizer=Tokenizer()) |
| single = [generate_many(model, processor, [item], 4, batch_size=1)[0] for item in inputs] |
| assert generate_many(model, processor, inputs, 4, batch_size=4) == single |
|
|
|
|
| @pytest.mark.parametrize( |
| "batch_size,prefix_size,shape_pad,minimum_width,fixed_batch_size", |
| [ |
| (2, 3, False, 0, 0), |
| (8, 3, False, 0, 0), |
| (8, 11, False, 0, 0), |
| (3, 3, True, 0, 0), |
| (2, 11, True, 128, 16), |
| (1, 3, True, 1792, 0), |
| ], |
| ) |
| def test_tiny_multimodal_model_keeps_video_features_with_their_own_prompt( |
| batch_size, prefix_size, shape_pad, minimum_width, fixed_batch_size |
| ): |
| from transformers import Qwen3_5Config, Qwen3_5ForConditionalGeneration |
|
|
| torch.manual_seed(42) |
| config = Qwen3_5Config( |
| text_config={ |
| "vocab_size": 64, |
| "hidden_size": 32, |
| "intermediate_size": 64, |
| "num_hidden_layers": 4, |
| "num_attention_heads": 4, |
| "num_key_value_heads": 2, |
| "head_dim": 8, |
| "linear_num_key_heads": 2, |
| "linear_num_value_heads": 2, |
| "linear_key_head_dim": 8, |
| "linear_value_head_dim": 8, |
| "layer_types": ["linear_attention"] * 3 + ["full_attention"], |
| "pad_token_id": 0, |
| "eos_token_id": 2, |
| "rope_parameters": {"rope_type": "default", "rope_theta": 10000, "mrope_section": [1, 1, 2]}, |
| }, |
| vision_config={ |
| "depth": 1, |
| "hidden_size": 16, |
| "intermediate_size": 32, |
| "num_heads": 2, |
| "in_channels": 3, |
| "patch_size": 2, |
| "temporal_patch_size": 2, |
| "spatial_merge_size": 1, |
| "out_hidden_size": 32, |
| "num_position_embeddings": 64, |
| }, |
| video_token_id=61, |
| image_token_id=59, |
| vision_start_token_id=60, |
| vision_end_token_id=62, |
| ) |
| model = Qwen3_5ForConditionalGeneration(config).eval() |
| model.generation_config.forced_eos_token_id = 2 |
| model.generation_config.min_new_tokens = 2 |
| inputs = [] |
| for prefix in (([5, 6], [7] * prefix_size) * ((batch_size + 1) // 2))[:batch_size]: |
| ids = [*prefix, 60, 61, 61, 61, 61, 62, 10] |
| inputs.append( |
| { |
| "input_ids": torch.tensor([ids]), |
| "attention_mask": torch.ones(1, len(ids), dtype=torch.long), |
| "mm_token_type_ids": torch.tensor([[0] * (len(prefix) + 1) + [2] * 4 + [0, 0]]), |
| "pixel_values_videos": torch.randn(4, 24), |
| "video_grid_thw": torch.tensor([[1, 2, 2]]), |
| } |
| ) |
| processor = SimpleNamespace(tokenizer=Tokenizer()) |
| single = [generate_many(model, processor, [item], 4, batch_size=1)[0] for item in inputs] |
| assert ( |
| generate_many( |
| model, |
| processor, |
| inputs, |
| 4, |
| batch_size=max(batch_size, fixed_batch_size), |
| max_padding_ratio=2.0 if prefix_size == 11 else 1.25, |
| padding_multiple=128 if shape_pad else 1, |
| pad_batch_power2=shape_pad, |
| minimum_width=minimum_width, |
| fixed_batch_size=fixed_batch_size, |
| ) |
| == single |
| ) |
|
|
