"""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 )