diff --git a/.ci/scripts/check_muse_glimmer_batching_report.py b/.ci/scripts/check_muse_glimmer_batching_report.py new file mode 100644 index 00000000000..69475039b83 --- /dev/null +++ b/.ci/scripts/check_muse_glimmer_batching_report.py @@ -0,0 +1,211 @@ +# 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. + +"""Checks the reports Muse Glimmer's run_solo_batching writes, for the +solo-text-batching CI mode (test_model_e2e.sh). + +Three things must hold: + +1. Every generation stands on its own: it finishes normally, produces text, + and contains its expected answer; and the single-prompt runs with and + without the captured decode graph generate the same tokens. Tokens are not + compared across batches: a prompt prefilled in a wider forward runs GEMMs of + another shape, whose bf16 rounding can steer greedy decoding elsewhere. +2. The model's weights are loaded once: loading costs about the weights, not + one copy per method, and does not grow with the sessions reserved for. +3. GPU memory behaves like an off-graph KV cache: the pool holds what the + sessions used, grows geometrically rather than reserving every session's + context up front, and accounts for what generation added. + + python .ci/scripts/check_muse_glimmer_batching_report.py \\ + --single eager.json --single graph.json --batch batch.json \\ + --expect Paris Tokyo Rome Paris Berlin +""" + +import argparse +import json +import sys +from typing import Dict, List, Optional, Sequence + +MIB = 1024 * 1024 +OK_REASONS = ("stop_token", "token_limit") + + +class ReportError(AssertionError): + pass + + +def _check(condition: bool, message: str) -> None: + if not condition: + raise ReportError(message) + + +def check_generations(report: Dict, expect: Sequence[str], label: str) -> None: + generations = report["generations"] + _check( + len(generations) == len(expect), + f"{label}: {len(generations)} generations, expected {len(expect)}", + ) + for index, (generation, answer) in enumerate(zip(generations, expect)): + name = f"{label}[{index}]" + _check( + generation["finish_reason"] in OK_REASONS, + f"{name}: finished {generation['finish_reason']}: {generation['error']}", + ) + _check(len(generation["tokens"]) > 0, f"{name}: generated nothing") + _check( + answer.lower() in generation["text"].lower(), + f"{name}: expected '{answer}' in {generation['text']!r}", + ) + + +def check_same_tokens(first: Dict, second: Dict, message: str) -> None: + _check(first["tokens"] == second["tokens"], message) + + +def load_cost(report: Dict) -> int: + gpu = report["gpu"] + return gpu["used_after_load_bytes"] - gpu["used_before_load_bytes"] + + +def check_weights_once( + reports: Sequence[Dict], max_weight_ratio: float, max_load_spread_mib: float +) -> None: + for index, report in enumerate(reports): + weights = report["weights_bytes"] + cost = load_cost(report) + _check(weights > 0, f"report {index}: no weights size") + _check( + cost <= max_weight_ratio * weights, + f"report {index}: loading took {cost / MIB:.0f} MiB for " + f"{weights / MIB:.0f} MiB of weights (limit {max_weight_ratio}x): " + "are they loaded once per method?", + ) + costs = [load_cost(report) for report in reports] + _check( + max(costs) - min(costs) <= max_load_spread_mib * MIB, + f"loading cost varies by {(max(costs) - min(costs)) / MIB:.0f} MiB " + f"across runs reserving different session counts (limit " + f"{max_load_spread_mib} MiB)", + ) + + +def check_kv(report: Dict, max_generate_slack_mib: float, label: str) -> Dict: + kv = report["kv"] + config = report["config"] + rows = kv["rows"] + used = kv["cells_in_use"] + _check(rows >= used, f"{label}: {rows} rows hold {used} cells in use") + _check( + kv["allocated_bytes"] == rows * kv["bytes_per_cell"], + f"{label}: {kv['allocated_bytes']} bytes allocated for {rows} rows of " + f"{kv['bytes_per_cell']}", + ) + # Geometric growth: never more than twice what the widest moment needed. + _check( + rows <= max(kv["initial_capacity"], 2 * (used + config["step_width"])), + f"{label}: {rows} rows for {used} cells in use: the pool is not " + "growing on demand", + ) + reserved = config["max_sessions"] * config["max_session_tokens"] + _check( + rows < reserved, + f"{label}: {rows} rows is every session's full context ({reserved})", + ) + gpu = report["gpu"] + added = gpu["used_after_generate_bytes"] - gpu["used_after_load_bytes"] + _check( + added <= kv["allocated_bytes"] + max_generate_slack_mib * MIB, + f"{label}: generation added {added / MIB:.0f} MiB, the KV pool is " + f"{kv['allocated_bytes'] / MIB:.0f} MiB (slack {max_generate_slack_mib} " + "MiB)", + ) + return { + "rows": rows, + "cells_in_use": used, + "kv_mib": kv["allocated_bytes"] / MIB, + "in_graph_mib": reserved * kv["bytes_per_cell"] / MIB, + "generate_added_mib": added / MIB, + } + + +def check_reports( + singles: Sequence[Dict], + batch: Dict, + expect: Sequence[str], + single_expect: Optional[str] = None, + max_weight_ratio: float = 1.25, + max_load_spread_mib: float = 256, + max_generate_slack_mib: float = 3072, +) -> List[str]: + """Raises ReportError on the first violation; returns a summary.""" + single_expect = single_expect or expect[0] + for index, report in enumerate(singles): + check_generations(report, [single_expect], f"single{index}") + for first, second in zip(singles, singles[1:]): + check_same_tokens( + first["generations"][0], + second["generations"][0], + "single-prompt runs disagree: the captured decode graph generated " + "different tokens from eager decode", + ) + check_generations(batch, expect, "batch") + check_weights_once(list(singles) + [batch], max_weight_ratio, max_load_spread_mib) + summary = [] + for label, report in [(f"single{i}", r) for i, r in enumerate(singles)] + [ + ("batch", batch) + ]: + stats = check_kv(report, max_generate_slack_mib, label) + summary.append( + f"{label}: KV {stats['rows']} rows for {stats['cells_in_use']} cells, " + f"{stats['kv_mib']:.0f} MiB (in-graph reservation " + f"{stats['in_graph_mib']:.0f} MiB); load " + f"{load_cost(report) / MIB:.0f} MiB for " + f"{report['weights_bytes'] / MIB:.0f} MiB of weights; generation " + f"added {stats['generate_added_mib']:.0f} MiB" + ) + engine = batch["engine"] + summary.append( + f"batch: {engine['steps']} forwards for {engine['decode_tokens_total']} " + f"decode and {engine['prefill_tokens_total']} prefill tokens" + ) + return summary + + +def main(argv: Optional[Sequence[str]] = None) -> int: + parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + parser.add_argument("--single", action="append", default=[], required=True) + parser.add_argument("--batch", required=True) + parser.add_argument("--expect", nargs="+", required=True) + parser.add_argument("--max-weight-ratio", type=float, default=1.25) + parser.add_argument("--max-load-spread-mib", type=float, default=256) + parser.add_argument("--max-generate-slack-mib", type=float, default=3072) + args = parser.parse_args(argv) + + def load(path): + with open(path) as f: + return json.load(f) + + try: + summary = check_reports( + [load(path) for path in args.single], + load(args.batch), + args.expect, + max_weight_ratio=args.max_weight_ratio, + max_load_spread_mib=args.max_load_spread_mib, + max_generate_slack_mib=args.max_generate_slack_mib, + ) + except ReportError as error: + print(f"FAIL: {error}") + return 1 + for line in summary: + print(line) + print("Success: batching reports check out") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/.ci/scripts/export_model_artifact.sh b/.ci/scripts/export_model_artifact.sh index 206b853242e..04d20f8ba6c 100755 --- a/.ci/scripts/export_model_artifact.sh +++ b/.ci/scripts/export_model_artifact.sh @@ -44,7 +44,7 @@ Arguments: - vr-streaming: Voxtral Realtime streaming mode - vr-offline: Voxtral Realtime offline mode - solo-text: Muse Glimmer solo text mode - - solo-text-offgraph: Muse Glimmer solo text, runtime-owned KV cache + - solo-text-batching: Muse Glimmer batched text (off-graph KV, export_solo_batching.py) - dflash-image: Muse Glimmer DFlash vision mode Examples: @@ -96,7 +96,7 @@ if [ -n "$MODE" ]; then exit 1 fi ;; - solo-text|solo-text-offgraph|dflash-image) + solo-text|solo-text-batching|dflash-image) if [ "$HF_MODEL" != "meta-models/Muse-Glimmer-30B-GGUF" ]; then echo "Error: Mode '$MODE' can only be used with Muse Glimmer model" echo "Provided model: $HF_MODEL" @@ -105,7 +105,7 @@ if [ -n "$MODE" ]; then ;; *) echo "Error: Unsupported mode '$MODE'" - echo "Supported modes: vr-streaming, vr-offline, solo-text, solo-text-offgraph, dflash-image" + echo "Supported modes: vr-streaming, vr-offline, solo-text, solo-text-batching, dflash-image" exit 1 ;; esac @@ -289,8 +289,8 @@ if [ "$MODEL_NAME" = "muse_glimmer" ]; then echo "Error: Muse Glimmer requires quantization 'kquant-17gb' or 'kquant-dynamic'" exit 1 fi - if [ "$MODE" != "solo-text" ] && [ "$MODE" != "solo-text-offgraph" ] && [ "$MODE" != "dflash-image" ]; then - echo "Error: Muse Glimmer requires mode 'solo-text', 'solo-text-offgraph' or 'dflash-image'" + if [ "$MODE" != "solo-text" ] && [ "$MODE" != "solo-text-batching" ] && [ "$MODE" != "dflash-image" ]; then + echo "Error: Muse Glimmer requires mode 'solo-text', 'solo-text-batching' or 'dflash-image'" exit 1 fi fi @@ -566,14 +566,12 @@ if [ "$MODEL_NAME" = "muse_glimmer" ]; then --backend cuda \ --output-dir "${OUTPUT_DIR}" ;; - solo-text-offgraph) + solo-text-batching) EXPORT_START_SECONDS=$SECONDS TMPDIR="$INDUCTOR_TMPDIR" \ TORCHINDUCTOR_CACHE_DIR="$INDUCTOR_CACHE" \ - python -m executorch.examples.models.muse_glimmer.export.export_solo \ + python -m executorch.examples.models.muse_glimmer.export.export_solo_batching \ --gguf "$TARGET_GGUF_PATH" \ - --backend cuda \ - --use-offgraph-kv-cache \ --output-dir "${OUTPUT_DIR}" ;; dflash-image) @@ -592,7 +590,7 @@ if [ "$MODEL_NAME" = "muse_glimmer" ]; then --output-dir "${OUTPUT_DIR}" ;; *) - echo "Error: Muse Glimmer requires mode 'solo-text', 'solo-text-offgraph' or 'dflash-image'" + echo "Error: Muse Glimmer requires mode 'solo-text', 'solo-text-batching' or 'dflash-image'" exit 1 ;; esac diff --git a/.ci/scripts/test_model_e2e.sh b/.ci/scripts/test_model_e2e.sh index 97689231877..1df336b7437 100755 --- a/.ci/scripts/test_model_e2e.sh +++ b/.ci/scripts/test_model_e2e.sh @@ -46,7 +46,7 @@ Arguments: - vr-streaming: Voxtral Realtime streaming mode - vr-offline: Voxtral Realtime offline mode - solo-text: Muse Glimmer solo text mode - - solo-text-offgraph: Muse Glimmer solo text, runtime-owned KV cache + - solo-text-batching: Muse Glimmer batched text (run_solo_batching) - dflash-image: Muse Glimmer DFlash vision mode Environment: @@ -112,7 +112,7 @@ if [ -n "$MODE" ]; then exit 1 fi ;; - solo-text|solo-text-offgraph|dflash-image) + solo-text|solo-text-batching|dflash-image) if [ "$HF_MODEL" != "meta-models/Muse-Glimmer-30B-GGUF" ]; then echo "Error: Mode '$MODE' can only be used with Muse Glimmer model" echo "Provided model: $HF_MODEL" @@ -121,7 +121,7 @@ if [ -n "$MODE" ]; then ;; *) echo "Error: Unsupported mode '$MODE'" - echo "Supported modes: vr-streaming, vr-offline, solo-text, solo-text-offgraph, dflash-image" + echo "Supported modes: vr-streaming, vr-offline, solo-text, solo-text-batching, dflash-image" exit 1 ;; esac @@ -286,18 +286,23 @@ case "$HF_MODEL" in AUDIO_FILE="" IMAGE_PATH="" case "$MODE" in - solo-text|solo-text-offgraph) + solo-text) RUNNER_TARGET="solo_runner" EXPECTED_OUTPUT="Paris" IMAGE_URL="" ;; + solo-text-batching) + RUNNER_TARGET="run_solo_batching" + EXPECTED_OUTPUT="Paris" + IMAGE_URL="" + ;; dflash-image) RUNNER_TARGET="dflash_runner" EXPECTED_OUTPUT="dog" IMAGE_URL="https://github.com/pytorch/hub/raw/master/images/dog.jpg" ;; *) - echo "Error: Muse Glimmer requires mode 'solo-text', 'solo-text-offgraph' or 'dflash-image'" + echo "Error: Muse Glimmer requires mode 'solo-text', 'solo-text-batching' or 'dflash-image'" exit 1 ;; esac @@ -447,18 +452,17 @@ EOF ;; muse_glimmer) PROMPT_FILE="${MODEL_DIR}/muse_glimmer_prompt.txt" - if [ "$MODE" = "solo-text" ] || [ "$MODE" = "solo-text-offgraph" ]; then + if [ "$MODE" = "solo-text" ] || [ "$MODE" = "solo-text-batching" ]; then printf '%s' '<|start|>user<|message|>What is the capital of France?<|eot|><|start|>assistant' > "$PROMPT_FILE" else printf '%s' '<|start|>user<|message|>What animal is in this image? <|eot|><|start|>assistant' > "$PROMPT_FILE" fi - RUNNER_ARGS="$RUNNER_ARGS --tokenizer_path ${MODEL_DIR}/$TOKENIZER_FILE --prompt_file \"$PROMPT_FILE\" --max_new_tokens 512" - if [ "$MODE" = "solo-text-offgraph" ]; then - # Exercise the same artifact on both decode paths: eager, then the - # captured CUDA graph, which is recaptured whenever the cache grows. - RUNNER_VARIANTS="--cuda_graph=false --cuda_graph" + RUNNER_ARGS="$RUNNER_ARGS --tokenizer_path ${MODEL_DIR}/$TOKENIZER_FILE --prompt_file \"$PROMPT_FILE\"" + if [ "$MODE" = "solo-text-batching" ]; then + # The passes, and the reports they are checked by, are run below. + RUNNER_ARGS="$RUNNER_ARGS --max_new_tokens 128 --max_sessions 5 --max_session_tokens 2048" else - RUNNER_ARGS="$RUNNER_ARGS --cuda_graph" + RUNNER_ARGS="$RUNNER_ARGS --max_new_tokens 512 --cuda_graph" fi if [ "$MODE" = "dflash-image" ]; then RUNNER_ARGS="$RUNNER_ARGS --image_path ${MODEL_DIR}/test_image.jpg" @@ -538,7 +542,36 @@ run_and_validate() { esac } -if [ -z "${RUNNER_VARIANTS:-}" ]; then +if [ "$MODEL_NAME" = "muse_glimmer" ] && [ "$MODE" = "solo-text-batching" ]; then + REPORTS="${MODEL_DIR}/batching_reports" + mkdir -p "$REPORTS" + # One prompt, on both decode paths: eager, then the captured CUDA graph, + # which is recaptured whenever the KV pool grows. + run_and_validate "--cuda_graph=false --report_json=${REPORTS}/single_eager.json" + run_and_validate "--cuda_graph=true --report_json=${REPORTS}/single_graph.json" + # Five prompts in one batch, the France prompt first: three more questions, + # the France prompt again, and one longer than a forward, so it prefills in + # slices beside the others' decodes. + BATCH_PROMPTS="${MODEL_DIR}/muse_glimmer_batch_prompts.txt" + NOTES="" + for _ in $(seq 1 60); do + NOTES="${NOTES}Note: the museum opens at nine and closes at five. " + done + printf '%s\n' \ + '<|start|>user<|message|>What is the capital of Japan?<|eot|><|start|>assistant' \ + '<|start|>user<|message|>What is the capital of Italy?<|eot|><|start|>assistant' \ + '<|start|>user<|message|>What is the capital of France?<|eot|><|start|>assistant' \ + "<|start|>user<|message|>${NOTES}Question: what is the capital of Germany?<|eot|><|start|>assistant" \ + > "$BATCH_PROMPTS" + run_and_validate "--cuda_graph=true --prompts_file=${BATCH_PROMPTS} --report_json=${REPORTS}/batch.json" + echo "::group::Check batching reports" + python .ci/scripts/check_muse_glimmer_batching_report.py \ + --single "${REPORTS}/single_eager.json" \ + --single "${REPORTS}/single_graph.json" \ + --batch "${REPORTS}/batch.json" \ + --expect Paris Tokyo Rome Paris Berlin + echo "::endgroup::" +elif [ -z "${RUNNER_VARIANTS:-}" ]; then run_and_validate "" else for VARIANT in $RUNNER_VARIANTS; do diff --git a/.github/workflows/cuda.yml b/.github/workflows/cuda.yml index 39054f27eac..52b35397761 100644 --- a/.github/workflows/cuda.yml +++ b/.github/workflows/cuda.yml @@ -2,7 +2,7 @@ # This workflow tests whether ExecuTorch can be successfully built with CUDA support # across different CUDA versions (12.6, 13.0, 13.4) using the command: # ./install_executorch.sh -# +# # Intentionally skipped CUDA version 13.2 check due to ci image unsupported. # # Note: ExecuTorch automatically detects the system CUDA version using nvcc and @@ -778,7 +778,7 @@ jobs: - variant: kquant-17gb mode: solo-text - variant: kquant-17gb - mode: solo-text-offgraph + mode: solo-text-batching - variant: kquant-17gb mode: dflash-image - variant: kquant-dynamic diff --git a/examples/models/muse-glimmer/export/export_solo.py b/examples/models/muse-glimmer/export/export_solo.py index 0eebc8ecea1..bb020afe6b1 100644 --- a/examples/models/muse-glimmer/export/export_solo.py +++ b/examples/models/muse-glimmer/export/export_solo.py @@ -78,7 +78,6 @@ def export_and_lower( backend: str = "cuda", sample: bool = True, use_turboquant: bool = False, - use_offgraph_kv_cache: bool = False, activation_dtype: torch.dtype = torch.bfloat16, max_prefill_chunk: int = 512, vision_model: nn.Module | None = None, @@ -93,20 +92,12 @@ def export_and_lower( output_dir, sample=sample, use_turboquant=use_turboquant, - use_offgraph_kv_cache=use_offgraph_kv_cache, vision_model=vision_model, pos_embed_table=pos_embed_table, max_vision_patches=max_vision_patches, vision_fp32_mm=vision_fp32_mm, ) elif backend == "mlx": - # The off-graph cache is lowered by the CUDA backend only, and dropping - # the flag silently would hand back an in-graph model that looks like - # what was asked for. - if use_offgraph_kv_cache: - raise ValueError( - "--use-offgraph-kv-cache is not supported by the mlx backend" - ) _export_mlx( model, config, @@ -150,6 +141,7 @@ def _solo_constant_methods( "enable_dynamic_shape": True, "has_vision_encoder": bool(has_vision), } + # A model whose KV cache lives off-graph has no mutable buffers to name. if mutable_buffer_metadata is not None: constant_methods["get_mutable_buffer_metadata"] = mutable_buffer_metadata if has_vision: @@ -164,7 +156,6 @@ def _export_cuda( output_dir: str, sample: bool = True, use_turboquant: bool = False, - use_offgraph_kv_cache: bool = False, vision_model: nn.Module | None = None, pos_embed_table: torch.Tensor | None = None, max_vision_patches: int = 16384, @@ -173,9 +164,6 @@ def _export_cuda( import torch._inductor.config as inductor_config from executorch.backends.cuda.cuda_backend import CudaBackend from executorch.backends.cuda.cuda_partitioner import CudaPartitioner - from executorch.backends.cuda.passes.lower_offgraph_kv import ( - OFFGRAPH_KV_STEP_WIDTH_COMPILE_SPEC, - ) from executorch.exir import ( EdgeCompileConfig, ExecutorchBackendConfig, @@ -205,31 +193,17 @@ def _export_cuda( from executorch.examples.models.muse_glimmer.source_transformations.cuda import ( add_on_device_sampler, cuda_source_transformations, - enable_offgraph_kv_cache, - offgraph_kv_cache_geometry, vision_cuda_source_transformations, ) # Always applied: bounds global-attention SDPA to the valid context via a # runtime kv_len (O(context) decode). With use_turboquant=True it also swaps # the global KV caches for TurboQuant TQ4. - if use_offgraph_kv_cache and use_turboquant: - raise ValueError("off-graph KV cache and TurboQuant are mutually exclusive") + cuda_source_transformations(model, use_turboquant=use_turboquant) - # A prefill chunk is one write step, so the ring must hold the union of the - # step's per-query windows. enable_offgraph_kv_cache sizes it from this. + # Max prefill chunk must fit in the ring buffer (2 * sliding_window) max_prefill = min(config.max_seq_len - 1, model._sliding_window * 2) - offgraph_manifest = None - offgraph_geometry = {} - if use_offgraph_kv_cache: - offgraph_manifest = enable_offgraph_kv_cache(model, max_prefill) - # Read the geometry while the model is alive; it is freed before the - # constant methods are assembled. - offgraph_geometry = offgraph_kv_cache_geometry(model) - else: - cuda_source_transformations(model, use_turboquant=use_turboquant) - has_vision = vision_model is not None programs: dict[str, "torch.export.ExportedProgram"] = {} @@ -295,30 +269,20 @@ def _export_cuda( vision_model, pos_embed_table, max_vision_patches ) - mutable_buffer_metadata = ( - None if offgraph_manifest else common.mutable_buffer_metadata(model) - ) - del model, vision_model + mutable_buffer_metadata = common.mutable_buffer_metadata(model) + del model + if has_vision: + del vision_model gc.collect() torch.cuda.empty_cache() def _partitioner_for(name: str) -> "CudaPartitioner": - compile_specs = [ - CudaBackend.generate_method_name_compile_spec(name), - CompileSpec("low_memory_mode", b"ON"), - ] - if offgraph_manifest is not None: - compile_specs.extend( - ( - CompileSpec("offgraph_kv_manifest", offgraph_manifest.encode()), - # Where the delegate reads this step's token count, as - # "input_index:dim". Lowering compiles the cache op away, so - # the runtime cannot recover the width from the graph; every - # method here takes input_pos second, shaped [T]. - CompileSpec(OFFGRAPH_KV_STEP_WIDTH_COMPILE_SPEC, b"1:0"), - ) - ) - return CudaPartitioner(compile_specs) + return CudaPartitioner( + [ + CudaBackend.generate_method_name_compile_spec(name), + CompileSpec("low_memory_mode", b"ON"), + ] + ) constant_methods = _solo_constant_methods( config=config, @@ -330,13 +294,6 @@ def _partitioner_for(name: str) -> "CudaPartitioner": ) constant_methods["get_min_prefill_chunk"] = _CUDA_MIN_PREFILL_CHUNK constant_methods["use_sampling"] = sample - if offgraph_manifest is not None: - # The manifest is a compile spec (the lowering pass needs geometry at - # partition time). The runtime instead reads the neutral cache-geometry - # constant methods, the same ones the MLX off-graph path publishes. - # Sizing already travels as get_max_seq_len (context) and - # get_max_prefill_chunk (largest step); only the geometry is new. - constant_methods.update(offgraph_geometry) print( f"Lowering {len(programs)} methods to ExecuTorch (CUDA): " @@ -615,11 +572,6 @@ def main() -> None: action="store_true", help="Use TurboQuant TQ4 KV cache on global (NoPE) layers (CUDA).", ) - parser.add_argument( - "--use-offgraph-kv-cache", - action="store_true", - help="Allocate CUDA KV cache at runtime instead of storing it in the PTE/PTD.", - ) parser.add_argument( "--activation-dtype", default=None, @@ -673,10 +625,6 @@ def main() -> None: parser.error("--activation-dtype is only supported with --backend mlx.") if args.backend != "mlx" and args.max_prefill_chunk != 512: parser.error("--max-prefill-chunk is only supported with --backend mlx.") - if args.use_offgraph_kv_cache and args.backend != "cuda": - parser.error("--use-offgraph-kv-cache requires --backend cuda.") - if args.use_offgraph_kv_cache and args.turboquant: - parser.error("--use-offgraph-kv-cache cannot be combined with --turboquant.") if args.gguf: from executorch.examples.models.muse_glimmer.loaders.checkpoint_loader import ( load_gguf_model, @@ -733,7 +681,6 @@ def main() -> None: backend=args.backend, sample=not args.logits, use_turboquant=args.turboquant, - use_offgraph_kv_cache=args.use_offgraph_kv_cache, activation_dtype=activation_dtype, max_prefill_chunk=args.max_prefill_chunk, vision_model=vision_model, diff --git a/examples/models/muse-glimmer/runtime/engine/muse_glimmer_engine.cpp b/examples/models/muse-glimmer/runtime/engine/muse_glimmer_engine.cpp index 6a3532d247f..47fdf9a9e72 100644 --- a/examples/models/muse-glimmer/runtime/engine/muse_glimmer_engine.cpp +++ b/examples/models/muse-glimmer/runtime/engine/muse_glimmer_engine.cpp @@ -28,9 +28,6 @@ #ifdef EXECUTORCH_BUILD_CUDA #include -#include -#include -#include #include #else #include @@ -279,23 +276,11 @@ Result read_sampled_token( #endif } -// The off-graph KV cache is CUDA-only, and so are its headers; elsewhere the -// module is built without one and the guard is always null. -#ifdef EXECUTORCH_BUILD_CUDA -using OffGraphInstallGuard = ::executorch::extension::llm::cache::InstallGuard; -#else -struct OffGraphInstallGuard; -#endif - Result> build_muse_glimmer_module( const MuseGlimmerConfig& config, MuseGlimmerArtifactMode artifact_mode, bool has_vision, - bool multi_session, - const OffGraphInstallGuard* offgraph_guard) { -#ifndef EXECUTORCH_BUILD_CUDA - (void)offgraph_guard; -#endif + bool multi_session) { std::vector data_files; if (!config.data_path.empty()) { data_files.push_back(config.data_path); @@ -330,20 +315,6 @@ Result> build_muse_glimmer_module( #endif const executorch::runtime::LoadBackendOptionsMap* load_options = nullptr; -#ifdef EXECUTORCH_BUILD_CUDA - // Name the installed off-graph cache so CudaBackend::init() can resolve it - // from the registry. Must outlive the load_method calls below, which read it - // during init. Absent for an in-graph model, and its absence is exactly how - // the delegate tells the two apart. - executorch::runtime::BackendOptions<1> cuda_cache_opts; - executorch::runtime::LoadBackendOptionsMap cuda_options_map; - if (offgraph_guard != nullptr) { - ET_CHECK_OK_OR_RETURN_ERROR(offgraph_guard->set_option(cuda_cache_opts)); - ET_CHECK_OK_OR_RETURN_ERROR(cuda_options_map.set_options( - ::executorch::backends::cuda::kCudaBackendId, cuda_cache_opts.view())); - load_options = &cuda_options_map; - } -#endif #ifdef EXECUTORCH_BUILD_MLX // Per-model MLX runtime specs, delivered to MLXBackend::init(). Must outlive // the load_method calls below (they read it during init). @@ -485,65 +456,6 @@ Error register_mutable_fqns( return Error::Ok; } -struct OffGraphKVPlan { - cache::CacheGeometry geometry; - cache::CacheConfig config; -}; - -Result read_offgraph_kv_plan( - Module& module, - const std::unordered_map& metadata, - int64_t initial_capacity) { - // Geometry comes from the neutral constant methods the MLX off-graph path - // also reads. The two sizes are already in the engine's metadata: - // kMaxContextLen is the cache ceiling (get_llm_metadata derives it from - // get_max_seq_len when the model does not publish it), and - // kMaxPrefillChunk is the largest single step, which is what a ring layer - // sizes its slots from. - ET_ASSIGN_OR_RETURN( - geometry, ::executorch::extension::llm::read_cache_geometry(module)); - - const auto capacity_it = metadata.find(kMaxContextLen); - const auto chunk_it = metadata.find(kMaxPrefillChunk); - ET_CHECK_OR_RETURN_ERROR( - capacity_it != metadata.end() && chunk_it != metadata.end(), - InvalidProgram, - "off-graph KV cache needs %s and %s metadata", - kMaxContextLen, - kMaxPrefillChunk); - - // CacheConfig counts in int; a value that does not survive the narrowing - // would otherwise be validated after truncation. - constexpr int64_t kMaxCacheInt = std::numeric_limits::max(); - for (const auto& [name, value] : - {std::pair{kMaxContextLen, capacity_it->second}, - std::pair{kMaxPrefillChunk, chunk_it->second}, - std::pair{"offgraph_initial_capacity", initial_capacity}}) { - ET_CHECK_OR_RETURN_ERROR( - value > 0 && value <= kMaxCacheInt, - InvalidProgram, - "off-graph KV cache: %s is %" PRId64 ", outside [1, %" PRId64 "]", - name, - value, - kMaxCacheInt); - } - - OffGraphKVPlan plan; - plan.geometry = std::move(geometry); - plan.config.capacity = static_cast(capacity_it->second); - plan.config.kv_dtype = - static_cast(executorch::aten::ScalarType::BFloat16); - plan.config.initial_capacity = static_cast(initial_capacity); - plan.config.max_write = static_cast(chunk_it->second); - ET_CHECK_OR_RETURN_ERROR( - cache::valid(plan.geometry, plan.config) && - plan.config.initial_capacity > 0 && - plan.config.initial_capacity <= plan.config.capacity, - InvalidProgram, - "invalid off-graph KV cache configuration"); - return plan; -} - TensorPtr build_decode_pos_table( const std::unordered_map& metadata) { auto ctx_it = metadata.find(kMaxContextLen); @@ -577,9 +489,6 @@ class MuseGlimmerSession : public LLMSession, int64_t min_prefill_chunk, TensorPtr decode_pos_table_dev, MuseGlimmerMutableStateContextOwner* mutable_state, -#ifdef EXECUTORCH_BUILD_CUDA - ::executorch::extension::llm::cache::SequenceControl* offgraph_control, -#endif int session_token) : module_(module), exec_mutex_(exec_mutex), @@ -593,9 +502,6 @@ class MuseGlimmerSession : public LLMSession, decode_pos_table_dev_(std::move(decode_pos_table_dev)), #endif mutable_state_(mutable_state), -#ifdef EXECUTORCH_BUILD_CUDA - offgraph_control_(offgraph_control), -#endif session_token_(session_token) { if (auto it = metadata_.find(kUseSampling); it != metadata_.end()) { use_sampling_ = it->second != 0; @@ -846,11 +752,6 @@ class MuseGlimmerSession : public LLMSession, Error reset() override { pos_ = 0; -#ifdef EXECUTORCH_BUILD_CUDA - if (offgraph_control_ != nullptr) { - offgraph_control_->clear(); - } -#endif pending_.reset(); prev_decode_token_.reset(); if (preserve_staged_image_on_next_reset_) { @@ -1213,10 +1114,6 @@ class MuseGlimmerSession : public LLMSession, TensorPtr decode_pos_table_dev_; #endif MuseGlimmerMutableStateContextOwner* mutable_state_ = nullptr; -#ifdef EXECUTORCH_BUILD_CUDA - ::executorch::extension::llm::cache::SequenceControl* offgraph_control_ = - nullptr; -#endif int session_token_ = kMuseGlimmerNoMutableSession; #ifdef EXECUTORCH_BUILD_CUDA float temp_val_ = 1e-6f; @@ -1502,43 +1399,8 @@ Result> MuseGlimmerEngine::create( } std::unique_ptr mutable_state; - const OffGraphInstallGuard* offgraph_install = nullptr; #ifdef EXECUTORCH_BUILD_CUDA - std::shared_ptr<::executorch::extension::llm::cache::Cache> offgraph_cache; - std::unique_ptr<::executorch::extension::llm::cache::InstallGuard> - offgraph_guard; - ::executorch::extension::llm::cache::SequenceControl* offgraph_control = - nullptr; - if (method_names.count(kNumCaches) != 0) { - ET_CHECK_OR_RETURN_ERROR( - artifact_mode == MuseGlimmerArtifactMode::Autoregressive, - NotSupported, - "off-graph KV cache currently supports autoregressive artifacts only"); - ET_CHECK_OR_RETURN_ERROR( - config.max_sessions == 1, - NotSupported, - "off-graph KV cache currently supports one session"); - auto plan = read_offgraph_kv_plan( - *meta_module, metadata, config.offgraph_initial_capacity); - ET_CHECK_OK_OR_RETURN_ERROR(plan.error()); - // The factory hands back a neutral cache; the guard publishes it under a - // registry key that build_muse_glimmer_module passes to the delegate. - auto built = cache::CacheFactory::global().build( - ::executorch::backends::cuda::kCudaBackendId, - cache::kind::kSingle, - plan.get().geometry, - plan.get().config); - ET_CHECK_OK_OR_RETURN_ERROR(built.error()); - offgraph_cache = built.get(); - offgraph_guard = std::make_unique(offgraph_cache); - offgraph_install = offgraph_guard.get(); - offgraph_control = offgraph_cache->as(); - ET_CHECK_OR_RETURN_ERROR( - offgraph_control != nullptr, - Internal, - "off-graph KV cache is missing a face the engine needs"); - ET_LOG(Info, "MuseGlimmerEngine: dynamic off-graph KV cache enabled"); - } else if (config.enable_cuda_graph) { + if (config.enable_cuda_graph) { ET_LOG( Info, "MuseGlimmerEngine: CUDA graph requested; per-session rebinding " @@ -1571,18 +1433,18 @@ Result> MuseGlimmerEngine::create( // skip_mutable_buffer_init, so the skip flag can never diverge from the // owner. const bool multi_session = mutable_state != nullptr; - auto build_module = [&]() { - return build_muse_glimmer_module( - config, artifact_mode, has_vision, multi_session, offgraph_install); - }; - // The in-graph path still scopes its load, because mutable-buffer rebinding - // has no registry key to travel on. - auto module_res = multi_session ? mutable_state->with_load_scope(build_module) - : build_module(); + auto module_res = multi_session + ? mutable_state->with_load_scope([&]() { + return build_muse_glimmer_module( + config, artifact_mode, has_vision, multi_session); + }) + : build_muse_glimmer_module( + config, artifact_mode, has_vision, multi_session); if (module_res.error() != Error::Ok) { return module_res.error(); } std::unique_ptr shared_module = std::move(module_res.get()); + bool rebind_available = false; rebind_available = mutable_state != nullptr && mutable_state->available(); if (rebind_available && mutable_state->validate_coverage() != Error::Ok) { @@ -1631,12 +1493,6 @@ Result> MuseGlimmerEngine::create( /*vision_runtime=*/nullptr, rebind_available, std::move(mutable_state))); -#ifdef EXECUTORCH_BUILD_CUDA - // Handed over after construction: the cache has to be installed before the - // module loads, which happens above. - engine->offgraph_guard_ = std::move(offgraph_guard); - engine->offgraph_control_ = offgraph_control; -#endif if (has_vision) { MuseGlimmerVisionRuntimeConfig vision_config; vision_config.module = engine->shared_module_.get(); @@ -1767,9 +1623,6 @@ Result> MuseGlimmerEngine::create_session() { min_prefill_chunk_, decode_pos_table_dev_, rebind_available_ ? mutable_state_.get() : nullptr, -#ifdef EXECUTORCH_BUILD_CUDA - offgraph_control_, -#endif token)); } diff --git a/examples/models/muse-glimmer/runtime/engine/muse_glimmer_engine.h b/examples/models/muse-glimmer/runtime/engine/muse_glimmer_engine.h index 49fea8d3cc2..06709f5c0bb 100644 --- a/examples/models/muse-glimmer/runtime/engine/muse_glimmer_engine.h +++ b/examples/models/muse-glimmer/runtime/engine/muse_glimmer_engine.h @@ -30,8 +30,6 @@ #ifdef EXECUTORCH_BUILD_CUDA #include -#include -#include #elif defined(EXECUTORCH_BUILD_MLX) #include #include @@ -82,7 +80,6 @@ struct MuseGlimmerConfig { int32_t max_sessions = 1; int64_t eos_id = 200001; bool enable_cuda_graph = false; - int64_t offgraph_initial_capacity = 512; MuseGlimmerArtifactMode artifact_mode = MuseGlimmerArtifactMode::Auto; int32_t dflash_block_length = 0; int32_t dflash_n_draft = 0; @@ -210,16 +207,6 @@ class ET_EXPERIMENTAL MuseGlimmerEngine : public LLMEngine { std::unique_ptr vision_runtime_; bool rebind_available_ = false; std::unique_ptr mutable_state_; -#ifdef EXECUTORCH_BUILD_CUDA - // The installed off-graph KV cache, or nothing for an in-graph model. The - // guard owns it and keeps it discoverable under its registry key for as long - // as methods may load; the CUDA delegate steps it. The engine only resets - // it, through the neutral SequenceControl face. - std::unique_ptr<::executorch::extension::llm::cache::InstallGuard> - offgraph_guard_; - ::executorch::extension::llm::cache::SequenceControl* offgraph_control_ = - nullptr; -#endif std::atomic live_sessions_{0}; }; diff --git a/examples/models/muse-glimmer/runtime/runners/muse_glimmer_worker.cpp b/examples/models/muse-glimmer/runtime/runners/muse_glimmer_worker.cpp index 769c34c217a..d6bf2001e2c 100644 --- a/examples/models/muse-glimmer/runtime/runners/muse_glimmer_worker.cpp +++ b/examples/models/muse-glimmer/runtime/runners/muse_glimmer_worker.cpp @@ -71,10 +71,6 @@ DEFINE_bool( false, "Capture CUDA graphs for supported methods. Disables CUDA multi-session " "state rebinding."); -DEFINE_int64( - offgraph_initial_capacity, - 512, - "Initial token capacity for growable flat off-graph KV caches."); namespace { namespace llm = ::executorch::extension::llm; @@ -323,7 +319,6 @@ int main(int argc, char** argv) { config.dflash_n_draft = FLAGS_dflash_n_draft; config.dflash_draft_argmax = FLAGS_dflash_draft_argmax; config.enable_cuda_graph = FLAGS_cuda_graph; - config.offgraph_initial_capacity = FLAGS_offgraph_initial_capacity; auto engine_result = llm::MuseGlimmerEngine::create(config); if (engine_result.error() != Error::Ok) { diff --git a/examples/models/muse-glimmer/runtime/runners/solo.cpp b/examples/models/muse-glimmer/runtime/runners/solo.cpp index c57625d97dd..9592b1df871 100644 --- a/examples/models/muse-glimmer/runtime/runners/solo.cpp +++ b/examples/models/muse-glimmer/runtime/runners/solo.cpp @@ -119,10 +119,6 @@ DEFINE_bool( cuda_graph, true, "Enable CUDA graph capture for the decode method. CUDA only; default on."); -DEFINE_int64( - offgraph_initial_capacity, - 512, - "Initial token capacity for growable flat off-graph KV caches."); DEFINE_string( nll_tokens_file, "", @@ -429,7 +425,6 @@ static int run_engine_generation(llm::Stats& stats) { config.max_sessions = 1; config.eos_id = FLAGS_eos_id; config.enable_cuda_graph = FLAGS_cuda_graph; - config.offgraph_initial_capacity = FLAGS_offgraph_initial_capacity; config.artifact_mode = llm::MuseGlimmerArtifactMode::Auto; auto engine_result = llm::MuseGlimmerEngine::create(config);