diff --git a/backends/cuda/batching/cuda_executor.cpp b/backends/cuda/batching/cuda_executor.cpp index 9fd8d4a7a13..db736f8b8b5 100644 --- a/backends/cuda/batching/cuda_executor.cpp +++ b/backends/cuda/batching/cuda_executor.cpp @@ -363,21 +363,40 @@ 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( + *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); + } + } else { + ET_ASSIGN_OR_RETURN( + cells, + metadata::detail::read_required_positive_int(*module, kMaxCellsMethod)); + max_cells = cells; ET_CHECK_OR_RETURN_ERROR( - max_session_tokens > 0, - InvalidArgument, - "CudaExecutor: no cell budget remains for each resident session"); + 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 +463,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 +502,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 +514,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 +532,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(); @@ -700,9 +731,8 @@ bool CudaExecutor::execute_raw( ? slice.width : std::max(selected.size(), min_selected_rows_); selected.resize(rows, selected.empty() ? 0 : selected.back()); - if (!ctl_->declare_step( - std::vector( - step.seq_ids.begin() + off, step.seq_ids.begin() + off + n))) { + if (!ctl_->declare_step(std::vector( + step.seq_ids.begin() + off, step.seq_ids.begin() + off + n))) { return false; } std::vector token_values( @@ -1083,9 +1113,8 @@ bool CudaExecutor::execute_embeddings( const std::size_t rows = std::max(selected.size(), min_selected_rows_); selected.resize(rows, selected.empty() ? 0 : selected.back()); - if (!ctl_->declare_step( - std::vector( - step.seq_ids.begin() + off, step.seq_ids.begin() + off + n))) { + if (!ctl_->declare_step(std::vector( + step.seq_ids.begin() + off, step.seq_ids.begin() + off + n))) { return false; } // Rows past n are padding the cache writes to its scratch cell. 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 7e7b6fc2ac3..39fe125cbfe 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 @@ -173,6 +174,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. @@ -278,6 +281,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 @@ -305,9 +310,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, + MAX_KV_SEQS_METHOD, + ) from executorch.examples.models.muse_glimmer.model.dflash_model import ( DFlashPackedExportModel, ) @@ -330,7 +345,8 @@ 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: + sequence_layout = kv_layout == "sequence" + 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") @@ -345,8 +361,11 @@ 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) + if sequence_layout + else cell_manifest(single_manifest, max_cells) ) combined = DFlashPackedExportModel(target_model, draft_model) owners = ( @@ -414,7 +433,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, @@ -422,6 +440,14 @@ def prepare_cuda_packed_draft( "get_dflash_draft_cache_key": "muse_glimmer_draft", } ) + if sequence_layout: + metadata[MAX_KV_SEQS_METHOD.replace("get_", "get_draft_", 1)] = 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 @@ -1126,6 +1152,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() @@ -1161,6 +1192,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 8391dbcfac0..03084da73d1 100644 --- a/examples/models/muse-glimmer/export/export_solo.py +++ b/examples/models/muse-glimmer/export/export_solo.py @@ -122,6 +122,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") @@ -141,6 +143,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, @@ -239,6 +243,10 @@ 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" +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 @@ -301,6 +309,39 @@ 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 +) -> 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 + 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) @@ -368,9 +409,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. @@ -402,6 +449,11 @@ 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") if dflash_draft is not None: from executorch.examples.models.muse_glimmer.export.export_dflash import ( prepare_cuda_packed_draft, @@ -414,6 +466,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})) @@ -427,7 +481,10 @@ 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) + ): raise ValueError( f"max_step_tokens must be in ({widest_static}, max_cells), " f"got {max_step_tokens}" @@ -443,8 +500,11 @@ 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. @@ -502,13 +562,15 @@ def export_cuda_batching( # noqa: C901 programs[EMBEDDINGS_METHOD] = export( model, ( - torch.zeros( - (1, max_step_tokens, config.dim), dtype=torch.bfloat16 - ), + 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), ), - dynamic_shapes=({1: embeddings_width}, {0: embeddings_width}, {0: rows}), + dynamic_shapes=( + {1: embeddings_width}, + {0: embeddings_width}, + {0: rows}, + ), strict=True, ) with common.BoundMethodForward(model, model.embed_text), torch.no_grad(): @@ -516,9 +578,7 @@ def export_cuda_batching( # noqa: C901 programs["embed_text"] = export( model, (torch.zeros((1, max_step_tokens), dtype=torch.long),), - dynamic_shapes=( - {1: Dim("embed_width", min=1, max=max_step_tokens)}, - ), + dynamic_shapes=({1: Dim("embed_width", min=1, max=max_step_tokens)},), strict=True, ) programs["vision_encoder"] = common.export_vision_encoder( @@ -597,7 +657,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) @@ -1134,6 +1197,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, @@ -1272,6 +1336,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 f053a7784f8..1295fdadac1 100644 --- a/examples/models/muse-glimmer/runtime/cuda_dflash_executor.cpp +++ b/examples/models/muse-glimmer/runtime/cuda_dflash_executor.cpp @@ -555,7 +555,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, @@ -643,7 +649,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 && @@ -657,11 +666,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 96b6bdbad62..4a382576996 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