Buckets:

EmergentTwilight
/

MMCoT

Files

xet

EmergentTwilight/MMCoT / experiments /tests /test_cached_rollout_single_cache.py

EmergentTwilight

19 days ago

download

raw

15.8 kB

import json
import os
import unittest
from pathlib import Path

import torch

from uniar import UniARConfig, UniARForConditionalGeneration
from zebra_sft.cached_rollout import CachedInterleavedRollout, CachedVisualContext


def _tiny_config() -> UniARConfig:
    return UniARConfig(
        text_config={
            "vocab_size": 128,
            "hidden_size": 64,
            "intermediate_size": 128,
            "num_hidden_layers": 2,
            "num_attention_heads": 4,
            "num_key_value_heads": 2,
            "head_dim": 16,
            "max_position_embeddings": 256,
            "rope_theta": 1_000_000,
            "rope_scaling": {
                "mrope_section": [2, 3, 3],
                "rope_type": "default",
            },
            "rms_norm_eps": 1e-6,
            "use_cache": False,
            "tie_word_embeddings": False,
        },
        vision_config={
            "depth": 2,
            "hidden_size": 32,
            "intermediate_size": 64,
            "num_heads": 4,
            "in_channels": 3,
            "patch_size": 16,
            "spatial_merge_size": 2,
            "temporal_patch_size": 2,
            "out_hidden_size": 64,
            "num_position_embeddings": 64,
            "deepstack_visual_indexes": [0],
            "use_bsq": True,
            "bsq_dim": 8,
            "bsq_hidden_dim": 64,
            "bsq_skip_final_layernorm": True,
        },
        vistok_pred=True,
        vistok_pred_layernorm=True,
        visual_transformer_decoder=True,
        visual_transformer_decoder_depth=2,
        image_token_id=100,
        video_token_id=101,
        vision_start_token_id=98,
        vision_end_token_id=99,
    )


def _rollout() -> tuple[CachedInterleavedRollout, CachedVisualContext]:
    torch.manual_seed(13)
    model = UniARForConditionalGeneration(_tiny_config()).bfloat16().eval()
    project_visual = model.visual_transformer_pred

    def cpu_visual_transformer_pred(hidden_states, *, attention_mask=None, **kwargs):
        # Production uses FlashAttention2's 2-D key-padding mask. The CPU SDPA
        # integration reached by this tiny fixture needs the equivalent 4-D
        # additive causal mask because we call decoder layers directly.
        if attention_mask is not None and attention_mask.ndim == 2:
            query = hidden_states.shape[1] if hidden_states.ndim == 3 else 1
            keys = attention_mask.shape[1]
            cached = keys - query
            allowed = torch.arange(keys)[None, :] <= (
                cached + torch.arange(query)[:, None]
            )
            allowed &= attention_mask[0].bool()[None, :]
            converted = torch.zeros(
                (1, 1, query, keys), dtype=hidden_states.dtype
            )
            converted.masked_fill_(~allowed[None, None], torch.finfo(converted.dtype).min)
            attention_mask = converted
        return project_visual(
            hidden_states, attention_mask=attention_mask, **kwargs
        )

    model.visual_transformer_pred = cpu_visual_transformer_pred
    rollout = CachedInterleavedRollout.__new__(CachedInterleavedRollout)
    rollout.ar = model
    rollout.context_limit = 256
    rollout.marker_id = 126
    rollout.im_end_id = 125
    rollout.vision_start_id = model.config.vision_start_token_id
    rollout.vision_end_id = model.config.vision_end_token_id
    rollout.stop_ids = set()
    hidden = model.config.text_config.hidden_size
    dtype = model.get_input_embeddings().weight.dtype
    context = CachedVisualContext(
        problem_main=torch.empty((0, hidden), dtype=dtype),
        problem_deep=[torch.empty((0, hidden), dtype=dtype)],
        problem_grids=torch.empty((0, 3), dtype=torch.long),
    )
    return rollout, context


class CachedRolloutSingleLanguageCacheTest(unittest.TestCase):
    def test_text_visual_text_outputs_match_phase_local_prefill_exactly(self):
        rollout, legacy_context = _rollout()
        cached_context = CachedVisualContext(
            problem_main=legacy_context.problem_main.clone(),
            problem_deep=[value.clone() for value in legacy_context.problem_deep],
            problem_grids=legacy_context.problem_grids.clone(),
        )
        prefix = torch.tensor([[1, 2, 3, 4]], dtype=torch.long)

        legacy_sequence, legacy_attention, _, legacy_first = rollout._generate_tokens(
            prefix.clone(),
            legacy_context,
            max_new_tokens=3,
            stop_on_marker=False,
        )
        state = rollout._initialize_language_cache(prefix.clone(), cached_context)
        cached_sequence, cached_attention, _, cached_first = (
            rollout._generate_tokens_cached(
                prefix.clone(),
                state,
                max_new_tokens=3,
                stop_on_marker=False,
            )
        )
        self.assertEqual(cached_first, legacy_first)
        self.assertTrue(torch.equal(cached_sequence, legacy_sequence))
        self.assertTrue(torch.equal(cached_attention, legacy_attention))

        vision_start = torch.tensor([[rollout.vision_start_id]], dtype=torch.long)
        legacy_sequence = torch.cat([legacy_sequence, vision_start], dim=1)
        legacy_attention = torch.cat(
            [legacy_attention, torch.ones_like(vision_start)], dim=1
        )
        cached_sequence = torch.cat([cached_sequence, vision_start], dim=1)
        rollout._consume_text_token(state, vision_start)

        legacy_codes = rollout._generate_visual(
            legacy_sequence,
            legacy_context,
            1,
            2,
            sequence_attention=legacy_attention,
        )
        cached_codes = rollout._generate_visual_cached(
            cached_sequence, state, cached_context, 1, 2
        )
        self.assertTrue(torch.equal(cached_codes, legacy_codes))

        grid = torch.tensor([[1, 2, 4]], dtype=torch.long)
        rollout._append_visual_features(legacy_context, grid, codes=legacy_codes)
        rollout._append_visual_features(cached_context, grid, codes=cached_codes)
        placeholders = torch.full(
            (1, 2), rollout.ar.config.image_token_id, dtype=torch.long
        )
        vision_end = torch.tensor([[rollout.vision_end_id]], dtype=torch.long)
        legacy_sequence = torch.cat(
            [legacy_sequence, placeholders, vision_end], dim=1
        )
        legacy_attention = torch.cat(
            [legacy_attention, torch.ones_like(placeholders), torch.ones_like(vision_end)],
            dim=1,
        )
        cached_sequence = torch.cat(
            [cached_sequence, placeholders, vision_end], dim=1
        )

        legacy_sequence, legacy_attention, _, legacy_second = rollout._generate_tokens(
            legacy_sequence,
            legacy_context,
            sequence_attention=legacy_attention,
            max_new_tokens=3,
            stop_on_marker=False,
        )
        cached_sequence, cached_attention, _, cached_second = (
            rollout._generate_tokens_cached(
                cached_sequence,
                state,
                max_new_tokens=3,
                stop_on_marker=False,
            )
        )
        self.assertEqual(cached_second, legacy_second)
        self.assertTrue(torch.equal(cached_sequence, legacy_sequence))
        self.assertTrue(torch.equal(cached_attention, legacy_attention))
        self.assertEqual(rollout._cache_length(state), cached_sequence.shape[1])

    def test_intervention_visual_blocks_match_rebuilt_prefix_text(self):
        rollout, _ = _rollout()
        base_prefix = torch.tensor([[1, 2, 3]], dtype=torch.long)
        vision_start = torch.tensor(
            [[rollout.vision_start_id]], dtype=torch.long
        )
        prefix = torch.cat([base_prefix, vision_start], dim=1)
        grid = torch.tensor([[1, 2, 4]], dtype=torch.long)
        levels = 1 + len(
            rollout.ar.model.visual.config.deepstack_visual_indexes
        )
        merge = rollout.ar.model.visual.spatial_merge_unit
        bits = rollout.ar.model.visual.config.bsq_dim
        codes = torch.randint(0, 2, (1, 2, levels, merge, bits))

        for mode in (
            "gray",
            "zero_visual_features",
            "image_pad_only",
            "delete_visual",
        ):
            with self.subTest(mode=mode):
                hidden = rollout.ar.config.text_config.hidden_size
                dtype = rollout.ar.get_input_embeddings().weight.dtype

                def context():
                    return CachedVisualContext(
                        problem_main=torch.empty((0, hidden), dtype=dtype),
                        problem_deep=[torch.empty((0, hidden), dtype=dtype)],
                        problem_grids=torch.empty((0, 3), dtype=torch.long),
                    )

                legacy_context = context()
                cached_context = context()
                state = rollout._initialize_language_cache(
                    base_prefix, cached_context
                )
                rollout._consume_text_token(state, vision_start)
                rollout._consume_serialized_visual_cached(
                    prefix,
                    state,
                    cached_context,
                    1,
                    2,
                    mode=mode,
                    codes=codes if mode == "gray" else None,
                )

                if mode == "gray":
                    rollout._append_visual_features(
                        legacy_context, grid, codes=codes
                    )
                    rollout._append_visual_features(
                        cached_context, grid, codes=codes
                    )
                elif mode == "zero_visual_features":
                    rollout._append_visual_features(legacy_context, grid, zeros=2)
                    rollout._append_visual_features(cached_context, grid, zeros=2)
                elif mode == "image_pad_only":
                    legacy_context.generated_grids.append(grid)
                    cached_context.generated_grids.append(grid)

                placeholders = (
                    torch.empty((1, 0), dtype=torch.long)
                    if mode == "delete_visual"
                    else torch.full(
                        (1, 2), rollout.ar.config.image_token_id, dtype=torch.long
                    )
                )
                vision_end = torch.tensor(
                    [[rollout.vision_end_id]], dtype=torch.long
                )
                sequence = torch.cat([prefix, placeholders, vision_end], dim=1)
                inactive = (
                    [(prefix.shape[1], prefix.shape[1] + 2)]
                    if mode == "image_pad_only"
                    else None
                )
                legacy, legacy_attention, _, legacy_tokens = (
                    rollout._generate_tokens(
                        sequence.clone(),
                        legacy_context,
                        inactive_visual_spans=inactive,
                        max_new_tokens=2,
                        stop_on_marker=False,
                    )
                )
                cached, cached_attention, _, cached_tokens = (
                    rollout._generate_tokens_cached(
                        sequence.clone(),
                        state,
                        max_new_tokens=2,
                        stop_on_marker=False,
                    )
                )
                self.assertEqual(cached_tokens, legacy_tokens)
                self.assertTrue(torch.equal(cached, legacy))
                self.assertTrue(torch.equal(cached_attention, legacy_attention))

    def test_retroactive_thought_mask_is_not_claimed_cache_safe(self):
        self.assertFalse(
            CachedInterleavedRollout._can_reuse_language_cache("mask_prior_thought")
        )
        for ablation in (
            "control",
            "noise",
            "gray",
            "image_pad_only",
            "zero_visual_features",
            "delete_visual",
        ):
            self.assertTrue(
                CachedInterleavedRollout._can_reuse_language_cache(ablation)
            )

    @unittest.skipUnless(
        os.environ.get("SPATIAL_MMCOT_REAL_CACHE_PARITY") == "1"
        and torch.cuda.is_available(),
        "set SPATIAL_MMCOT_REAL_CACHE_PARITY=1 on a CUDA compute node",
    )
    def test_real_flash_attention_production_defaults_are_bit_exact(self):
        from peft import PeftModel
        from transformers import AutoProcessor
        from uniar import configure_processor

        project = Path(
            os.environ.get(
                "SPATIAL_MMCOT_PROJECT_ROOT",
                Path(__file__).resolve().parents[2],
            )
        )
        model_path = project / "models/UniAR-SFT"
        adapter = Path(os.environ["SPATIAL_MMCOT_PARITY_ADAPTER"])
        cache_root = Path(os.environ["SPATIAL_MMCOT_PARITY_CACHE"])
        with (cache_root / "manifest.jsonl").open(encoding="utf-8") as handle:
            rows = [json.loads(next(handle)) for _ in range(2)]
        row = rows[1]
        payload = torch.load(
            cache_root / row["cache_shard"],
            map_location="cpu",
            weights_only=False,
        )
        item = payload["items"][row["cache_index"]]

        processor = configure_processor(
            AutoProcessor.from_pretrained(model_path, use_fast=True)
        )
        base = UniARForConditionalGeneration.from_pretrained(
            model_path,
            dtype=torch.bfloat16,
            attn_implementation="flash_attention_2",
            low_cpu_mem_usage=True,
        ).cuda()
        model = PeftModel.from_pretrained(base, adapter).eval()
        rollout = CachedInterleavedRollout(model, processor)
        prefix, original_context, _ = rollout._prepare(item)

        def clone_context() -> CachedVisualContext:
            return CachedVisualContext(
                problem_main=original_context.problem_main.clone(),
                problem_deep=[
                    value.clone() for value in original_context.problem_deep
                ],
                problem_grids=original_context.problem_grids.clone(),
            )

        legacy_context = clone_context()
        cached_context = clone_context()
        legacy, legacy_attention, _, legacy_tokens = rollout._generate_tokens(
            prefix.clone(),
            legacy_context,
            max_new_tokens=8,
            stop_on_marker=False,
        )
        state = rollout._initialize_language_cache(prefix.clone(), cached_context)
        cached, cached_attention, _, cached_tokens = rollout._generate_tokens_cached(
            prefix.clone(),
            state,
            max_new_tokens=8,
            stop_on_marker=False,
        )
        self.assertEqual(cached_tokens, legacy_tokens)
        self.assertTrue(torch.equal(cached, legacy))

        vision_start = torch.tensor(
            [[rollout.vision_start_id]], dtype=torch.long, device=rollout.ar.device
        )
        legacy = torch.cat([legacy, vision_start], dim=1)
        legacy_attention = torch.cat(
            [legacy_attention, torch.ones_like(vision_start)], dim=1
        )
        cached = torch.cat([cached, vision_start], dim=1)
        rollout._consume_text_token(state, vision_start)
        legacy_codes = rollout._generate_visual(
            legacy,
            legacy_context,
            1,
            2,
            sequence_attention=legacy_attention,
            reuse_visual_decoder_cache=False,
        )
        production_codes = rollout._generate_visual(
            legacy,
            clone_context(),
            1,
            2,
            sequence_attention=legacy_attention,
        )
        self.assertTrue(torch.equal(production_codes, legacy_codes))


if __name__ == "__main__":
    unittest.main()

Xet Storage Details

Size:: 15.8 kB
Xet hash:: cd0831aa25b3dde9322ee0f315e7656fc0d68f738c514a02d0035df4abc76d9e

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.

	import json
	import os
	import unittest
	from pathlib import Path

	import torch

	from uniar import UniARConfig, UniARForConditionalGeneration
	from zebra_sft.cached_rollout import CachedInterleavedRollout, CachedVisualContext


	def _tiny_config() -> UniARConfig:
	return UniARConfig(
	text_config={
	"vocab_size": 128,
	"hidden_size": 64,
	"intermediate_size": 128,
	"num_hidden_layers": 2,
	"num_attention_heads": 4,
	"num_key_value_heads": 2,
	"head_dim": 16,
	"max_position_embeddings": 256,
	"rope_theta": 1_000_000,
	"rope_scaling": {
	"mrope_section": [2, 3, 3],
	"rope_type": "default",
	},
	"rms_norm_eps": 1e-6,
	"use_cache": False,
	"tie_word_embeddings": False,
	},
	vision_config={
	"depth": 2,
	"hidden_size": 32,
	"intermediate_size": 64,
	"num_heads": 4,
	"in_channels": 3,
	"patch_size": 16,
	"spatial_merge_size": 2,
	"temporal_patch_size": 2,
	"out_hidden_size": 64,
	"num_position_embeddings": 64,
	"deepstack_visual_indexes": [0],
	"use_bsq": True,
	"bsq_dim": 8,
	"bsq_hidden_dim": 64,
	"bsq_skip_final_layernorm": True,
	},
	vistok_pred=True,
	vistok_pred_layernorm=True,
	visual_transformer_decoder=True,
	visual_transformer_decoder_depth=2,
	image_token_id=100,
	video_token_id=101,
	vision_start_token_id=98,
	vision_end_token_id=99,
	)


	def _rollout() -> tuple[CachedInterleavedRollout, CachedVisualContext]:
	torch.manual_seed(13)
	model = UniARForConditionalGeneration(_tiny_config()).bfloat16().eval()
	project_visual = model.visual_transformer_pred

	def cpu_visual_transformer_pred(hidden_states, , attention_mask=None, *kwargs):
	# Production uses FlashAttention2's 2-D key-padding mask. The CPU SDPA
	# integration reached by this tiny fixture needs the equivalent 4-D
	# additive causal mask because we call decoder layers directly.
	if attention_mask is not None and attention_mask.ndim == 2:
	query = hidden_states.shape[1] if hidden_states.ndim == 3 else 1
	keys = attention_mask.shape[1]
	cached = keys - query
	allowed = torch.arange(keys)[None, :] <= (
	cached + torch.arange(query)[:, None]
	)
	allowed &= attention_mask[0].bool()[None, :]
	converted = torch.zeros(
	(1, 1, query, keys), dtype=hidden_states.dtype
	)
	converted.masked_fill_(~allowed[None, None], torch.finfo(converted.dtype).min)
	attention_mask = converted
	return project_visual(
	hidden_states, attention_mask=attention_mask, **kwargs
	)

	model.visual_transformer_pred = cpu_visual_transformer_pred
	rollout = CachedInterleavedRollout.__new__(CachedInterleavedRollout)
	rollout.ar = model
	rollout.context_limit = 256
	rollout.marker_id = 126
	rollout.im_end_id = 125
	rollout.vision_start_id = model.config.vision_start_token_id
	rollout.vision_end_id = model.config.vision_end_token_id
	rollout.stop_ids = set()
	hidden = model.config.text_config.hidden_size
	dtype = model.get_input_embeddings().weight.dtype
	context = CachedVisualContext(
	problem_main=torch.empty((0, hidden), dtype=dtype),
	problem_deep=[torch.empty((0, hidden), dtype=dtype)],
	problem_grids=torch.empty((0, 3), dtype=torch.long),
	)
	return rollout, context


	class CachedRolloutSingleLanguageCacheTest(unittest.TestCase):
	def test_text_visual_text_outputs_match_phase_local_prefill_exactly(self):
	rollout, legacy_context = _rollout()
	cached_context = CachedVisualContext(
	problem_main=legacy_context.problem_main.clone(),
	problem_deep=[value.clone() for value in legacy_context.problem_deep],
	problem_grids=legacy_context.problem_grids.clone(),
	)
	prefix = torch.tensor([[1, 2, 3, 4]], dtype=torch.long)

	legacy_sequence, legacy_attention, _, legacy_first = rollout._generate_tokens(
	prefix.clone(),
	legacy_context,
	max_new_tokens=3,
	stop_on_marker=False,
	)
	state = rollout._initialize_language_cache(prefix.clone(), cached_context)
	cached_sequence, cached_attention, _, cached_first = (
	rollout._generate_tokens_cached(
	prefix.clone(),
	state,
	max_new_tokens=3,
	stop_on_marker=False,
	)
	)
	self.assertEqual(cached_first, legacy_first)
	self.assertTrue(torch.equal(cached_sequence, legacy_sequence))
	self.assertTrue(torch.equal(cached_attention, legacy_attention))

	vision_start = torch.tensor([[rollout.vision_start_id]], dtype=torch.long)
	legacy_sequence = torch.cat([legacy_sequence, vision_start], dim=1)
	legacy_attention = torch.cat(
	[legacy_attention, torch.ones_like(vision_start)], dim=1
	)
	cached_sequence = torch.cat([cached_sequence, vision_start], dim=1)
	rollout._consume_text_token(state, vision_start)

	legacy_codes = rollout._generate_visual(
	legacy_sequence,
	legacy_context,
	1,
	2,
	sequence_attention=legacy_attention,
	)
	cached_codes = rollout._generate_visual_cached(
	cached_sequence, state, cached_context, 1, 2
	)
	self.assertTrue(torch.equal(cached_codes, legacy_codes))

	grid = torch.tensor([[1, 2, 4]], dtype=torch.long)
	rollout._append_visual_features(legacy_context, grid, codes=legacy_codes)
	rollout._append_visual_features(cached_context, grid, codes=cached_codes)
	placeholders = torch.full(
	(1, 2), rollout.ar.config.image_token_id, dtype=torch.long
	)
	vision_end = torch.tensor([[rollout.vision_end_id]], dtype=torch.long)
	legacy_sequence = torch.cat(
	[legacy_sequence, placeholders, vision_end], dim=1
	)
	legacy_attention = torch.cat(
	[legacy_attention, torch.ones_like(placeholders), torch.ones_like(vision_end)],
	dim=1,
	)
	cached_sequence = torch.cat(
	[cached_sequence, placeholders, vision_end], dim=1
	)

	legacy_sequence, legacy_attention, _, legacy_second = rollout._generate_tokens(
	legacy_sequence,
	legacy_context,
	sequence_attention=legacy_attention,
	max_new_tokens=3,
	stop_on_marker=False,
	)
	cached_sequence, cached_attention, _, cached_second = (
	rollout._generate_tokens_cached(
	cached_sequence,
	state,
	max_new_tokens=3,
	stop_on_marker=False,
	)
	)
	self.assertEqual(cached_second, legacy_second)
	self.assertTrue(torch.equal(cached_sequence, legacy_sequence))
	self.assertTrue(torch.equal(cached_attention, legacy_attention))
	self.assertEqual(rollout._cache_length(state), cached_sequence.shape[1])

	def test_intervention_visual_blocks_match_rebuilt_prefix_text(self):
	rollout, _ = _rollout()
	base_prefix = torch.tensor([[1, 2, 3]], dtype=torch.long)
	vision_start = torch.tensor(
	[[rollout.vision_start_id]], dtype=torch.long
	)
	prefix = torch.cat([base_prefix, vision_start], dim=1)
	grid = torch.tensor([[1, 2, 4]], dtype=torch.long)
	levels = 1 + len(
	rollout.ar.model.visual.config.deepstack_visual_indexes
	)
	merge = rollout.ar.model.visual.spatial_merge_unit
	bits = rollout.ar.model.visual.config.bsq_dim
	codes = torch.randint(0, 2, (1, 2, levels, merge, bits))

	for mode in (
	"gray",
	"zero_visual_features",
	"image_pad_only",
	"delete_visual",
	):
	with self.subTest(mode=mode):
	hidden = rollout.ar.config.text_config.hidden_size
	dtype = rollout.ar.get_input_embeddings().weight.dtype

	def context():
	return CachedVisualContext(
	problem_main=torch.empty((0, hidden), dtype=dtype),
	problem_deep=[torch.empty((0, hidden), dtype=dtype)],
	problem_grids=torch.empty((0, 3), dtype=torch.long),
	)

	legacy_context = context()
	cached_context = context()
	state = rollout._initialize_language_cache(
	base_prefix, cached_context
	)
	rollout._consume_text_token(state, vision_start)
	rollout._consume_serialized_visual_cached(
	prefix,
	state,
	cached_context,
	1,
	2,
	mode=mode,
	codes=codes if mode == "gray" else None,
	)

	if mode == "gray":
	rollout._append_visual_features(
	legacy_context, grid, codes=codes
	)
	rollout._append_visual_features(
	cached_context, grid, codes=codes
	)
	elif mode == "zero_visual_features":
	rollout._append_visual_features(legacy_context, grid, zeros=2)
	rollout._append_visual_features(cached_context, grid, zeros=2)
	elif mode == "image_pad_only":
	legacy_context.generated_grids.append(grid)
	cached_context.generated_grids.append(grid)

	placeholders = (
	torch.empty((1, 0), dtype=torch.long)
	if mode == "delete_visual"
	else torch.full(
	(1, 2), rollout.ar.config.image_token_id, dtype=torch.long
	)
	)
	vision_end = torch.tensor(
	[[rollout.vision_end_id]], dtype=torch.long
	)
	sequence = torch.cat([prefix, placeholders, vision_end], dim=1)
	inactive = (
	[(prefix.shape[1], prefix.shape[1] + 2)]
	if mode == "image_pad_only"
	else None
	)
	legacy, legacy_attention, _, legacy_tokens = (
	rollout._generate_tokens(
	sequence.clone(),
	legacy_context,
	inactive_visual_spans=inactive,
	max_new_tokens=2,
	stop_on_marker=False,
	)
	)
	cached, cached_attention, _, cached_tokens = (
	rollout._generate_tokens_cached(
	sequence.clone(),
	state,
	max_new_tokens=2,
	stop_on_marker=False,
	)
	)
	self.assertEqual(cached_tokens, legacy_tokens)
	self.assertTrue(torch.equal(cached, legacy))
	self.assertTrue(torch.equal(cached_attention, legacy_attention))

	def test_retroactive_thought_mask_is_not_claimed_cache_safe(self):
	self.assertFalse(
	CachedInterleavedRollout._can_reuse_language_cache("mask_prior_thought")
	)
	for ablation in (
	"control",
	"noise",
	"gray",
	"image_pad_only",
	"zero_visual_features",
	"delete_visual",
	):
	self.assertTrue(
	CachedInterleavedRollout._can_reuse_language_cache(ablation)
	)

	@unittest.skipUnless(
	os.environ.get("SPATIAL_MMCOT_REAL_CACHE_PARITY") == "1"
	and torch.cuda.is_available(),
	"set SPATIAL_MMCOT_REAL_CACHE_PARITY=1 on a CUDA compute node",
	)
	def test_real_flash_attention_production_defaults_are_bit_exact(self):
	from peft import PeftModel
	from transformers import AutoProcessor
	from uniar import configure_processor

	project = Path(
	os.environ.get(
	"SPATIAL_MMCOT_PROJECT_ROOT",
	Path(__file__).resolve().parents[2],
	)
	)
	model_path = project / "models/UniAR-SFT"
	adapter = Path(os.environ["SPATIAL_MMCOT_PARITY_ADAPTER"])
	cache_root = Path(os.environ["SPATIAL_MMCOT_PARITY_CACHE"])
	with (cache_root / "manifest.jsonl").open(encoding="utf-8") as handle:
	rows = [json.loads(next(handle)) for _ in range(2)]
	row = rows[1]
	payload = torch.load(
	cache_root / row["cache_shard"],
	map_location="cpu",
	weights_only=False,
	)
	item = payload["items"][row["cache_index"]]

	processor = configure_processor(
	AutoProcessor.from_pretrained(model_path, use_fast=True)
	)
	base = UniARForConditionalGeneration.from_pretrained(
	model_path,
	dtype=torch.bfloat16,
	attn_implementation="flash_attention_2",
	low_cpu_mem_usage=True,
	).cuda()
	model = PeftModel.from_pretrained(base, adapter).eval()
	rollout = CachedInterleavedRollout(model, processor)
	prefix, original_context, _ = rollout._prepare(item)

	def clone_context() -> CachedVisualContext:
	return CachedVisualContext(
	problem_main=original_context.problem_main.clone(),
	problem_deep=[
	value.clone() for value in original_context.problem_deep
	],
	problem_grids=original_context.problem_grids.clone(),
	)

	legacy_context = clone_context()
	cached_context = clone_context()
	legacy, legacy_attention, _, legacy_tokens = rollout._generate_tokens(
	prefix.clone(),
	legacy_context,
	max_new_tokens=8,
	stop_on_marker=False,
	)
	state = rollout._initialize_language_cache(prefix.clone(), cached_context)
	cached, cached_attention, _, cached_tokens = rollout._generate_tokens_cached(
	prefix.clone(),
	state,
	max_new_tokens=8,
	stop_on_marker=False,
	)
	self.assertEqual(cached_tokens, legacy_tokens)
	self.assertTrue(torch.equal(cached, legacy))

	vision_start = torch.tensor(
	[[rollout.vision_start_id]], dtype=torch.long, device=rollout.ar.device
	)
	legacy = torch.cat([legacy, vision_start], dim=1)
	legacy_attention = torch.cat(
	[legacy_attention, torch.ones_like(vision_start)], dim=1
	)
	cached = torch.cat([cached, vision_start], dim=1)
	rollout._consume_text_token(state, vision_start)
	legacy_codes = rollout._generate_visual(
	legacy,
	legacy_context,
	1,
	2,
	sequence_attention=legacy_attention,
	reuse_visual_decoder_cache=False,
	)
	production_codes = rollout._generate_visual(
	legacy,
	clone_context(),
	1,
	2,
	sequence_attention=legacy_attention,
	)
	self.assertTrue(torch.equal(production_codes, legacy_codes))


	if __name__ == "__main__":
	unittest.main()