diff --git a/backends/cuda/batching/cuda_executor.cpp b/backends/cuda/batching/cuda_executor.cpp index 6f357e5443b..b33cb492fbf 100644 --- a/backends/cuda/batching/cuda_executor.cpp +++ b/backends/cuda/batching/cuda_executor.cpp @@ -363,21 +363,46 @@ Result> CudaExecutor::create( NotSupported, "CudaExecutor: the program must select its logits rows"); ET_ASSIGN_OR_RETURN(geometry, metadata::read_cache_geometry(*module)); + // The program fixes its cache layout: a private history per sequence when it + // publishes its sequence count, else one shared pool of cells. ET_ASSIGN_OR_RETURN( - max_cells, - metadata::detail::read_required_positive_int(*module, kMaxCellsMethod)); - - ET_CHECK_OR_RETURN_ERROR( - max_cells > 1 && max_cells <= std::numeric_limits::max(), - InvalidProgram, - "CudaExecutor: cell count must fit an int and leave a padding row"); - if (max_session_tokens == 0) { - max_session_tokens = static_cast( - std::min(max_context_length, (max_cells - 1) / max_sessions)); + declared_seqs, + metadata::detail::read_int_method(*module, kMaxSeqsMethod)); + const bool sequence_layout = declared_seqs.has_value(); + std::int64_t max_cells = 0; + if (sequence_layout) { ET_CHECK_OR_RETURN_ERROR( - max_session_tokens > 0, + *declared_seqs > 0 && + *declared_seqs <= std::numeric_limits::max() && + max_context_length <= std::numeric_limits::max(), + InvalidProgram, + "CudaExecutor: invalid sequence count or context length"); + if (max_session_tokens == 0) { + max_session_tokens = static_cast(max_context_length); + } + ET_CHECK_OR_RETURN_ERROR( + max_session_tokens <= max_context_length, InvalidArgument, - "CudaExecutor: no cell budget remains for each resident session"); + "CudaExecutor: %d session tokens exceed the sequence capacity %" PRId64, + max_session_tokens, + max_context_length); + } else { + ET_ASSIGN_OR_RETURN( + cells, + metadata::detail::read_required_positive_int(*module, kMaxCellsMethod)); + max_cells = cells; + ET_CHECK_OR_RETURN_ERROR( + max_cells > 1 && max_cells <= std::numeric_limits::max(), + InvalidProgram, + "CudaExecutor: cell count must fit an int and leave a padding row"); + if (max_session_tokens == 0) { + max_session_tokens = static_cast( + std::min(max_context_length, (max_cells - 1) / max_sessions)); + ET_CHECK_OR_RETURN_ERROR( + max_session_tokens > 0, + InvalidArgument, + "CudaExecutor: no cell budget remains for each resident session"); + } } // Every forward_{N} and forward_others the program exports, narrowest first. @@ -444,11 +469,11 @@ Result> CudaExecutor::create( "CudaExecutor: more than one dynamic method"); const int max_step_tokens = methods.back().max_tokens; ET_CHECK_OR_RETURN_ERROR( - max_step_tokens < max_cells, + sequence_layout ? max_step_tokens <= max_context_length + : max_step_tokens < max_cells, InvalidProgram, - "CudaExecutor: a %d-token step does not fit %" PRId64 " cells", - max_step_tokens, - max_cells); + "CudaExecutor: a %d-token step does not fit the cache", + max_step_tokens); // A dynamic method whose kernels switch at a small width is exported from // more selected rows than one; its selector is padded up to that. int min_selected_rows = 1; @@ -483,7 +508,9 @@ Result> CudaExecutor::create( // them all, or a step could find no free cell mid-generation. One of the // program's cells is the padding scratch row, never handed out. ET_CHECK_OR_RETURN_ERROR( - static_cast(max_sessions) * max_session_tokens < max_cells, + sequence_layout || + static_cast(max_sessions) * max_session_tokens < + max_cells, InvalidArgument, "CudaExecutor: %d sessions of %d tokens exceed the program's %" PRId64 " cells, one of which is reserved", @@ -493,10 +520,16 @@ Result> CudaExecutor::create( // The pool is the program's: its size and widest step fix the shapes the // program declared for the step buffers. + // A sequence's columns hold the program's context; every sequence has its + // own. llm_cache::CacheConfig cfg{}; - cfg.capacity = static_cast(max_cells); + cfg.capacity = + static_cast(sequence_layout ? max_context_length : max_cells); cfg.kv_dtype = kv_dtype; cfg.max_write = max_step_tokens; + if (sequence_layout) { + cfg.max_seqs = static_cast(*declared_seqs); + } if (initial_capacity >= 0) { cfg.initial_capacity = std::min(initial_capacity, cfg.capacity); } @@ -505,7 +538,11 @@ Result> CudaExecutor::create( InvalidProgram, "CudaExecutor: the program's layout is unusable"); auto built = llm_cache::CacheFactory::global().build( - backend_id, llm_cache::kind::kBatchedCell, geometry, cfg); + backend_id, + sequence_layout ? llm_cache::kind::kBatchedSequence + : llm_cache::kind::kBatchedCell, + geometry, + cfg); ET_CHECK_OK_OR_RETURN_ERROR(built.error()); std::shared_ptr cache = built.get(); auto* const ctl = cache->as(); diff --git a/backends/cuda/batching/cuda_executor.h b/backends/cuda/batching/cuda_executor.h index ccd96afb9eb..4e0cf7bc7da 100644 --- a/backends/cuda/batching/cuda_executor.h +++ b/backends/cuda/batching/cuda_executor.h @@ -48,6 +48,9 @@ inline constexpr char kStaticMethodPrefix[] = "forward_"; inline constexpr char kDynamicMethod[] = "forward_others"; // The cell layout's pool size, which fixes the shape of its step buffers. inline constexpr char kMaxCellsMethod[] = "get_offgraph_kv_max_cells"; +// Published instead of the cell count by a program whose KV cache keeps a +// private history per sequence. +inline constexpr char kMaxSeqsMethod[] = "get_offgraph_kv_max_seqs"; // Optional: the fewest selected rows forward_others was exported for. Its // selector is padded up to it. Defaults to 1. inline constexpr char kMinSelectedRowsMethod[] = "get_min_prefill_chunk"; diff --git a/examples/models/muse-glimmer/export/export_dflash.py b/examples/models/muse-glimmer/export/export_dflash.py index 6c2eb2da01d..015beb5b535 100644 --- a/examples/models/muse-glimmer/export/export_dflash.py +++ b/examples/models/muse-glimmer/export/export_dflash.py @@ -14,6 +14,7 @@ import argparse import gc +import json from typing import TYPE_CHECKING import torch @@ -177,6 +178,8 @@ def export_dflash( draft_max_cells: int | None = None, max_draft_owners: int = 4, forward_sizes=None, + kv_layout: str = "cell", + max_kv_seqs: int = 4, ) -> None: """Export DFlash target + draft to one CUDA or MLX .pte. @@ -282,6 +285,8 @@ def export_dflash( vision_model=vision_model, pos_embed_table=pos_embed_table, max_vision_patches=max_vision_patches, + kv_layout=kv_layout, + max_kv_seqs=max_kv_seqs, ) return @@ -309,9 +314,19 @@ def prepare_cuda_packed_draft( max_step_tokens: int, max_cells: int, max_owners: int, + kv_layout: str = "cell", + max_kv_seqs: int = 4, ): - """Trace the packed draft and publish its independent cell geometry.""" - from executorch.examples.models.muse_glimmer.export.export_solo import cell_manifest + """Trace the packed draft and publish its independent cache geometry. + + On the sequence layout the draft keeps one private history per resident + session, and each owner's block attends all of it and the whole block. + """ + from executorch.examples.models.muse_glimmer.export.export_solo import ( + batched_sequence_manifest, + cell_manifest, + DRAFT_MAX_KV_SEQS_METHOD, + ) from executorch.examples.models.muse_glimmer.model.dflash_model import ( DFlashPackedExportModel, ) @@ -334,7 +349,12 @@ def prepare_cuda_packed_draft( if type(max_step_tokens) is not int or max_step_tokens < 2: raise ValueError("packed DFlash max_step_tokens must be an integer >= 2") max_write = max_step_tokens + max_owners * block - if type(max_cells) is not int or max_cells <= max_write: + if kv_layout not in ("cell", "sequence"): + raise ValueError(f"kv_layout must be 'cell' or 'sequence', got {kv_layout!r}") + sequence_layout = kv_layout == "sequence" + if sequence_layout and (type(max_kv_seqs) is not int or max_kv_seqs <= 0): + raise ValueError("the sequence layout needs max_kv_seqs > 0") + if not sequence_layout and (type(max_cells) is not int or max_cells <= max_write): raise ValueError("draft max_cells must exceed context plus all owner blocks") if config.dim != target_config.dim: raise ValueError("draft and target hidden dimensions must agree") @@ -349,8 +369,13 @@ def prepare_cuda_packed_draft( convert_quantized_tensors_for_cuda(draft_model) draft_model.to(torch.bfloat16) - manifest = cell_manifest( - enable_dflash_offgraph_kv_cache(draft_model, max_write), max_cells + single_manifest = enable_dflash_offgraph_kv_cache(draft_model, max_write) + manifest = ( + batched_sequence_manifest( + single_manifest, max_kv_seqs, causal=False, context_window=True + ) + if sequence_layout + else cell_manifest(single_manifest, max_cells) ) combined = DFlashPackedExportModel(target_model, draft_model) owners = ( @@ -418,7 +443,6 @@ def prepare_cuda_packed_draft( "get_dflash_target_layers": torch.tensor( config.target_layers, dtype=torch.int32 ), - "get_draft_offgraph_kv_max_cells": max_cells, "get_draft_max_step_tokens": max_write, "get_draft_max_owners": max_owners, "get_draft_dflash_block_size": block, @@ -426,6 +450,14 @@ def prepare_cuda_packed_draft( "get_dflash_draft_cache_key": "muse_glimmer_draft", } ) + if sequence_layout: + metadata[DRAFT_MAX_KV_SEQS_METHOD] = max_kv_seqs + # Fixes each sequence column's rows, which the runtime must match. + metadata["get_draft_offgraph_kv_seq_capacity"] = json.loads(single_manifest)[ + "maximum_capacity" + ] + else: + metadata["get_draft_offgraph_kv_max_cells"] = max_cells return programs, manifest, metadata @@ -1130,6 +1162,11 @@ def main() -> None: parser.add_argument("--max-cells", type=int, default=65536) parser.add_argument("--draft-max-cells", type=int, default=None) parser.add_argument("--max-draft-owners", type=int, default=4) + from executorch.examples.models.muse_glimmer.export.export_solo import ( + add_kv_layout_arguments, + ) + + add_kv_layout_arguments(parser) parser.add_argument("--forward-sizes", default="1,2,4,8,16,32,64") args = parser.parse_args() @@ -1165,6 +1202,8 @@ def main() -> None: draft_max_cells=args.draft_max_cells, max_draft_owners=args.max_draft_owners, forward_sizes=[int(n) for n in args.forward_sizes.split(",") if n], + kv_layout=args.kv_layout, + max_kv_seqs=args.max_kv_seqs, ) diff --git a/examples/models/muse-glimmer/export/export_solo.py b/examples/models/muse-glimmer/export/export_solo.py index 7c480b122cf..6dc5fbb6f8d 100644 --- a/examples/models/muse-glimmer/export/export_solo.py +++ b/examples/models/muse-glimmer/export/export_solo.py @@ -124,6 +124,8 @@ def export_and_lower( forward_sizes=None, max_step_tokens: int = 4096, max_cells: int = 65536, + kv_layout: str = "cell", + max_kv_seqs: int = 4, ) -> None: if use_offgraph_kv_cache and use_turboquant: raise ValueError("off-graph KV cache and TurboQuant are mutually exclusive") @@ -145,6 +147,8 @@ def export_and_lower( max_cells=max_cells, forward_sizes=forward_sizes or DEFAULT_FORWARD_SIZES, device_sampling=sample, + kv_layout=kv_layout, + max_kv_seqs=max_kv_seqs, vision_model=vision_model, pos_embed_table=pos_embed_table, max_vision_patches=max_vision_patches, @@ -243,6 +247,13 @@ def _solo_constant_methods( # Read by CudaExecutor: the cell pool's size, which fixes its step buffers. MAX_CELLS_METHOD = "get_offgraph_kv_max_cells" +# Read by CudaExecutor in place of the cell count when each sequence keeps a +# private history: the sequence columns of every layer's storage. +MAX_KV_SEQS_METHOD = "get_offgraph_kv_max_seqs" +# Read by MuseGlimmerCudaDFlashExecutor for the draft's own sequence columns, +# mirroring kMaxSeqs in runtime/cuda_dflash_executor.h. +DRAFT_MAX_KV_SEQS_METHOD = "get_draft_offgraph_kv_max_seqs" +KV_LAYOUTS = ("cell", "sequence") # Static methods by default: powers of two up to the widest GEMM bucket. DEFAULT_FORWARD_SIZES = (1, 2, 4, 8, 16, 32, 64) # The device samplers' widest call: one row per resident session, and the cell @@ -305,6 +316,44 @@ def embeddings_step_forward( return logits if taps is None else (logits, taps) +def add_kv_layout_arguments(parser) -> None: + parser.add_argument( + "--kv-layout", + choices=KV_LAYOUTS, + default="cell", + help="CUDA with --use-offgraph-kv-cache: cell shares one pool of " + "--max-cells among every sequence; sequence keeps --max-kv-seqs private " + "histories of up to --max-seq-len each, as MLX batching does.", + ) + parser.add_argument( + "--max-kv-seqs", + type=int, + default=4, + help="CUDA with --kv-layout sequence: sequences resident at once.", + ) + + +def batched_sequence_manifest( + sequence_manifest: str, + max_seqs: int, + causal: bool = True, + context_window: bool = False, +) -> str: + """The batched-sequence lowering manifest: one column per sequence.""" + from executorch.backends.cuda.passes.lower_offgraph_kv import ( + batched_sequence_max_segments, + ) + + manifest = json.loads(sequence_manifest) + manifest["layout"] = "batched-sequence" + manifest["max_seqs"] = max_seqs + manifest["max_segments"] = batched_sequence_max_segments(max_seqs) + manifest["causal"] = causal + if context_window: + manifest["context_window"] = True + return json.dumps(manifest, sort_keys=True, separators=(",", ":")) + + def cell_manifest(sequence_manifest: str, max_cells: int) -> str: """The cell-layout lowering manifest for the model's off-graph layers.""" manifest = json.loads(sequence_manifest) @@ -372,9 +421,15 @@ def export_cuda_batching( # noqa: C901 vision_model: nn.Module | None = None, pos_embed_table: torch.Tensor | None = None, max_vision_patches: int = 16384, + kv_layout: str = "cell", + max_kv_seqs: int = 4, ): """Exports the forward methods and writes the artifact. + ``kv_layout`` picks the KV cache the program addresses: ``cell`` shares + one pool of ``max_cells`` cells among every sequence, ``sequence`` keeps + ``max_kv_seqs`` private histories of up to ``config.max_seq_len`` each. + With ``vision_model`` the token forwards are unchanged, and image prompts run through three more methods: ``vision_encoder``, ``embed_text`` and a dynamic ``forward_from_embeddings`` over the same cache and weights. @@ -406,10 +461,14 @@ def export_cuda_batching( # noqa: C901 ) from torch.export import Dim, export + if kv_layout not in KV_LAYOUTS: + raise ValueError(f"kv_layout must be one of {KV_LAYOUTS}, got {kv_layout!r}") + sequence_layout = kv_layout == "sequence" + if sequence_layout and max_kv_seqs <= 0: + raise ValueError("the sequence layout needs max_kv_seqs > 0") has_vision = vision_model is not None and pos_embed_table is not None if (vision_model is None) != (pos_embed_table is None): raise ValueError("vision_model and pos_embed_table must be given together") - if dflash_draft is not None: from executorch.examples.models.muse_glimmer.export.export_dflash import ( prepare_cuda_packed_draft, @@ -422,6 +481,8 @@ def export_cuda_batching( # noqa: C901 max_step_tokens=max_step_tokens, max_cells=max_cells if draft_max_cells is None else draft_max_cells, max_owners=max_draft_owners, + kv_layout=kv_layout, + max_kv_seqs=max_kv_seqs, ) forward_sizes = tuple(sorted({int(n) for n in forward_sizes})) @@ -435,9 +496,13 @@ def export_cuda_batching( # noqa: C901 f"the widest forward size must be at least {_CUDA_MIN_PREFILL_CHUNK - 1}, " f"got {widest_static}" ) - if not widest_static < max_step_tokens < max_cells: + if not ( + widest_static < max_step_tokens + and (sequence_layout or max_step_tokens < max_cells) + ): + bound = config.max_seq_len if sequence_layout else max_cells raise ValueError( - f"max_step_tokens must be in ({widest_static}, max_cells), " + f"max_step_tokens must be in ({widest_static}, {bound}), " f"got {max_step_tokens}" ) if max_step_tokens >= config.max_seq_len: @@ -451,14 +516,20 @@ def export_cuda_batching( # noqa: C901 if hasattr(inductor_config, "cpp_cache_precompile_headers"): inductor_config.cpp_cache_precompile_headers = False - manifest = cell_manifest( - enable_offgraph_kv_cache(model, max_step_tokens), max_cells + single_manifest = enable_offgraph_kv_cache(model, max_step_tokens) + manifest = ( + batched_sequence_manifest(single_manifest, max_kv_seqs) + if sequence_layout + else cell_manifest(single_manifest, max_cells) ) if dflash_draft is not None: # Deferred checkpoints still need scalar/RoPE buffers, after KV removal. materialize_runtime_buffers(model, dtype=torch.bfloat16) geometry = offgraph_kv_cache_geometry(model) + # A sequence prefill samples a few rows; tracing all rows creates a + # max_step_tokens * vocab FP32 output during AOTI's JIT validation. + example_rows = _CUDA_MIN_PREFILL_CHUNK if sequence_layout else max_step_tokens width = Dim( "width", min=max(widest_static, _CUDA_MIN_PREFILL_CHUNK), max=max_step_tokens ) @@ -488,7 +559,7 @@ def export_cuda_batching( # noqa: C901 ( torch.zeros((1, max_step_tokens), dtype=torch.long), torch.arange(max_step_tokens, dtype=torch.long), - torch.zeros(max_step_tokens, dtype=torch.long), + torch.zeros(example_rows, dtype=torch.long), ), dynamic_shapes=({1: width}, {0: width}, {0: rows}), strict=True, @@ -512,7 +583,7 @@ def export_cuda_batching( # noqa: C901 ( torch.zeros((1, max_step_tokens, config.dim), dtype=torch.bfloat16), torch.arange(max_step_tokens, dtype=torch.long), - torch.zeros(max_step_tokens, dtype=torch.long), + torch.zeros(example_rows, dtype=torch.long), ), dynamic_shapes=( {1: embeddings_width}, @@ -601,7 +672,10 @@ def partitioner(name: str) -> CudaPartitioner: constant_methods.update(write_logits_to_keep_mode("selected")) constant_methods["get_min_prefill_chunk"] = _CUDA_MIN_PREFILL_CHUNK constant_methods["use_sampling"] = False - constant_methods[MAX_CELLS_METHOD] = max_cells + if sequence_layout: + constant_methods[MAX_KV_SEQS_METHOD] = max_kv_seqs + else: + constant_methods[MAX_CELLS_METHOD] = max_cells if "draft_forward" in programs: constant_methods.update(draft_metadata) @@ -1139,6 +1213,7 @@ def main() -> None: "the tokens every resident session together may hold. Memory grows with " "use, up to this.", ) + add_kv_layout_arguments(parser) parser.add_argument( "--activation-dtype", default=None, @@ -1277,6 +1352,8 @@ def main() -> None: forward_sizes=[int(n) for n in args.forward_sizes.split(",") if n], max_step_tokens=args.max_step_tokens, max_cells=args.max_cells, + kv_layout=args.kv_layout, + max_kv_seqs=args.max_kv_seqs, ) diff --git a/examples/models/muse-glimmer/runtime/cuda_dflash_executor.cpp b/examples/models/muse-glimmer/runtime/cuda_dflash_executor.cpp index b79b28be429..d5cc5776d99 100644 --- a/examples/models/muse-glimmer/runtime/cuda_dflash_executor.cpp +++ b/examples/models/muse-glimmer/runtime/cuda_dflash_executor.cpp @@ -569,7 +569,13 @@ Result MuseGlimmerCudaDFlashExecutor::create( declared_block, positive_int(module, "get_draft_dflash_block_size")); ET_ASSIGN_OR_RETURN(max_owners, positive_int(module, df::kMaxOwners)); ET_ASSIGN_OR_RETURN(max_write, positive_int(module, df::kMaxWrite)); - ET_ASSIGN_OR_RETURN(max_cells, positive_int(module, df::kMaxCells)); + ET_ASSIGN_OR_RETURN( + draft_seqs, detail::read_int_method(module, df::kMaxSeqs)); + const bool sequence_layout = draft_seqs.has_value(); + // A sequence layout sizes each draft history; a cell layout one shared pool. + ET_ASSIGN_OR_RETURN( + max_cells, + positive_int(module, sequence_layout ? df::kSeqCapacity : df::kMaxCells)); ET_ASSIGN_OR_RETURN(tap_layers, positive_int(module, df::kTargetLayers)); ET_ASSIGN_OR_RETURN( layer_ids, @@ -657,7 +663,10 @@ Result MuseGlimmerCudaDFlashExecutor::create( "DFlash draft packed tensor ABI mismatch"); ET_ASSIGN_OR_RETURN(geometry, draft_geometry(module, block)); const int temporary = max_owners * block; - const int draft_budget = (max_cells - 1 - temporary) / config.max_sessions; + // A history column holds its context and, while drafting, its own block. + const int draft_budget = sequence_layout + ? max_cells - block + : (max_cells - 1 - temporary) / config.max_sessions; const int max_context = std::min(target->max_session_tokens(), draft_budget); ET_CHECK_OR_RETURN_ERROR( max_context > block && draft_tokens.sizes()[0] == max_owners && @@ -671,11 +680,19 @@ Result MuseGlimmerCudaDFlashExecutor::create( cache_config.capacity = max_cells; cache_config.kv_dtype = static_cast(aten::ScalarType::BFloat16); cache_config.max_write = max_write; + if (sequence_layout) { + ET_CHECK_OR_RETURN_ERROR( + *draft_seqs > 0 && *draft_seqs <= std::numeric_limits::max(), + InvalidProgram, + "Invalid DFlash draft sequence count"); + cache_config.max_seqs = static_cast(*draft_seqs); + } ET_ASSIGN_OR_RETURN( draft_cache, cache::CacheFactory::global().build( ::executorch::backends::cuda::kCudaBackendId, - cache::kind::kBatchedCell, + sequence_layout ? cache::kind::kBatchedSequence + : cache::kind::kBatchedCell, geometry, cache_config)); const auto* control = draft_cache->as(); diff --git a/examples/models/muse-glimmer/runtime/cuda_dflash_executor.h b/examples/models/muse-glimmer/runtime/cuda_dflash_executor.h index 5614d43e806..57627bc15c5 100644 --- a/examples/models/muse-glimmer/runtime/cuda_dflash_executor.h +++ b/examples/models/muse-glimmer/runtime/cuda_dflash_executor.h @@ -37,6 +37,9 @@ inline constexpr char kTargetLayers[] = "get_n_target_layers"; inline constexpr char kMaxOwners[] = "get_draft_max_owners"; inline constexpr char kMaxWrite[] = "get_draft_max_step_tokens"; inline constexpr char kMaxCells[] = "get_draft_offgraph_kv_max_cells"; +// Published instead of kMaxCells when the draft keeps a history per sequence. +inline constexpr char kMaxSeqs[] = "get_draft_offgraph_kv_max_seqs"; +inline constexpr char kSeqCapacity[] = "get_draft_offgraph_kv_seq_capacity"; // Static draft_forward_{N}: N owners, each with at most one block of context. inline constexpr char kStaticForwardPrefix[] = "draft_forward_"; // On-device verification of [owners, width] candidates; see diff --git a/examples/models/muse-glimmer/tests/test_cuda_sequence_export_inputs.py b/examples/models/muse-glimmer/tests/test_cuda_sequence_export_inputs.py new file mode 100644 index 00000000000..aa0fa11f7ef --- /dev/null +++ b/examples/models/muse-glimmer/tests/test_cuda_sequence_export_inputs.py @@ -0,0 +1,114 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Small export exemplars must preserve the full selected-logit API.""" + +import copy +import tempfile +import unittest +from dataclasses import replace +from unittest.mock import patch + +import torch +from executorch.examples.models.muse_glimmer.export import common +from executorch.examples.models.muse_glimmer.export.export_solo import ( + export_cuda_batching, +) +from executorch.examples.models.muse_glimmer.model.model import MuseGlimmerModel +from executorch.examples.models.muse_glimmer.tests.test_pipeline import TINY_CONFIG +from executorch.extension.llm.cache.reference_cache import ( + CacheConfig, + LayerPolicy, + SequenceReferenceCache, +) +from executorch.extension.llm.cache.update_and_attend import REGISTRY + + +class CapturedPrograms(Exception): + pass + + +class SequenceExportInputsTest(unittest.TestCase): + def test_small_logits_example_preserves_token_and_embedding_capacity(self): + torch.manual_seed(42) + config = replace( + TINY_CONFIG, + dim=32, + n_layers=2, + head_dim=8, + max_seq_len=128, + global_attn_cfg="[16,0]", + ) + target = MuseGlimmerModel(config).to(torch.bfloat16).eval() + target.activation_dtype = torch.bfloat16 + reference = copy.deepcopy(target) + captured = {} + + def capture(programs, **kwargs): + captured.update(programs) + raise CapturedPrograms + + stub = torch.export.export(torch.nn.Identity(), (torch.ones(2, 2),)) + with ( + tempfile.TemporaryDirectory() as out, + patch("executorch.exir.to_edge_transform_and_lower", side_effect=capture), + patch.object(common, "export_vision_encoder", return_value=stub), + self.assertRaises(CapturedPrograms), + ): + export_cuda_batching( + target, + config, + out, + max_step_tokens=32, + max_cells=128, + forward_sizes=(1, 2, 4, 8), + kv_layout="sequence", + max_kv_seqs=4, + vision_model=torch.nn.Identity(), + pos_embed_table=torch.ones(4, 32), + max_vision_patches=4, + ) + + cache_config = CacheConfig( + n_layers=config.n_layers, + n_kv_heads=config.n_kv_heads, + head_dim=config.head_dim, + capacity=config.max_seq_len, + layers=tuple( + ( + LayerPolicy.ring(config.layer_window_size(i)) + if config.layer_window_size(i) + else LayerPolicy.flat() + ) + for i in range(config.n_layers) + ), + ) + tokens = torch.randint(config.vocab_size, (1, 32)) + positions = torch.arange(32) + with torch.no_grad(): + expected = reference(tokens, positions)[0] + embeddings = reference.embed_text(tokens) + for name in ("forward_others", "forward_from_embeddings"): + program = captured[name] + self.assertLessEqual(program.example_inputs[0][2].numel(), 5) + call = program.module() + for count in (5, 9, 32): + with self.subTest(method=name, selected_rows=count): + key = f"export-inputs-{name}-{count}" + REGISTRY.install(key, SequenceReferenceCache(cache_config)) + try: + with torch.no_grad(), REGISTRY.active(key): + actual = call( + tokens if name == "forward_others" else embeddings, + positions, + torch.arange(count), + ) + finally: + REGISTRY.uninstall(key) + self.assertEqual(tuple(actual.shape), (count, config.vocab_size)) + torch.testing.assert_close( + actual, expected[:count], atol=5e-2, rtol=0 + )