From dcde746ef0eb1eaf8bc4d41ce0586d8a87410fd2 Mon Sep 17 00:00:00 2001 From: gasoonjia Date: Mon, 12 Oct 2026 22:21:28 -0700 Subject: [PATCH 1/5] Update [ghstack-poisoned] --- .ci/scripts/export_model_artifact.sh | 21 +- .ci/scripts/test_model_e2e.sh | 19 +- .github/workflows/cuda.yml | 2 - .../models/muse-glimmer/export/export_solo.py | 391 +++++++++++++++--- .../tests/test_checkpoint_context.py | 38 +- .../tests/test_cuda_batching_pipeline.py | 303 ++++++++++++++ 6 files changed, 689 insertions(+), 85 deletions(-) create mode 100644 examples/models/muse-glimmer/tests/test_cuda_batching_pipeline.py diff --git a/.ci/scripts/export_model_artifact.sh b/.ci/scripts/export_model_artifact.sh index 206b853242e..223dd82951c 100755 --- a/.ci/scripts/export_model_artifact.sh +++ b/.ci/scripts/export_model_artifact.sh @@ -44,7 +44,6 @@ 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 - dflash-image: Muse Glimmer DFlash vision mode Examples: @@ -96,7 +95,7 @@ if [ -n "$MODE" ]; then exit 1 fi ;; - solo-text|solo-text-offgraph|dflash-image) + solo-text|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 +104,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, dflash-image" exit 1 ;; esac @@ -289,8 +288,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" != "dflash-image" ]; then + echo "Error: Muse Glimmer requires mode 'solo-text' or 'dflash-image'" exit 1 fi fi @@ -566,16 +565,6 @@ if [ "$MODEL_NAME" = "muse_glimmer" ]; then --backend cuda \ --output-dir "${OUTPUT_DIR}" ;; - solo-text-offgraph) - EXPORT_START_SECONDS=$SECONDS - TMPDIR="$INDUCTOR_TMPDIR" \ - TORCHINDUCTOR_CACHE_DIR="$INDUCTOR_CACHE" \ - python -m executorch.examples.models.muse_glimmer.export.export_solo \ - --gguf "$TARGET_GGUF_PATH" \ - --backend cuda \ - --use-offgraph-kv-cache \ - --output-dir "${OUTPUT_DIR}" - ;; dflash-image) DRAFT_GGUF_FILE="dflash-Muse-Glimmer-30B-Q4_K_M.gguf" MMPROJ_GGUF_FILE="mmproj-Muse-Glimmer-30B-Q4_K_M.gguf" @@ -592,7 +581,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' or 'dflash-image'" exit 1 ;; esac diff --git a/.ci/scripts/test_model_e2e.sh b/.ci/scripts/test_model_e2e.sh index 97689231877..427329ae961 100755 --- a/.ci/scripts/test_model_e2e.sh +++ b/.ci/scripts/test_model_e2e.sh @@ -46,7 +46,6 @@ 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 - dflash-image: Muse Glimmer DFlash vision mode Environment: @@ -112,7 +111,7 @@ if [ -n "$MODE" ]; then exit 1 fi ;; - solo-text|solo-text-offgraph|dflash-image) + solo-text|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 +120,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, dflash-image" exit 1 ;; esac @@ -286,7 +285,7 @@ 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="" @@ -297,7 +296,7 @@ case "$HF_MODEL" in 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' or 'dflash-image'" exit 1 ;; esac @@ -447,19 +446,13 @@ EOF ;; muse_glimmer) PROMPT_FILE="${MODEL_DIR}/muse_glimmer_prompt.txt" - if [ "$MODE" = "solo-text" ] || [ "$MODE" = "solo-text-offgraph" ]; then + if [ "$MODE" = "solo-text" ]; 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" - else - RUNNER_ARGS="$RUNNER_ARGS --cuda_graph" - fi + RUNNER_ARGS="$RUNNER_ARGS --cuda_graph" if [ "$MODE" = "dflash-image" ]; then RUNNER_ARGS="$RUNNER_ARGS --image_path ${MODEL_DIR}/test_image.jpg" fi diff --git a/.github/workflows/cuda.yml b/.github/workflows/cuda.yml index 28f31a8373d..25d35c8fdfb 100644 --- a/.github/workflows/cuda.yml +++ b/.github/workflows/cuda.yml @@ -1061,8 +1061,6 @@ jobs: include: - variant: kquant-17gb mode: solo-text - - variant: kquant-17gb - mode: solo-text-offgraph - 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 4bd82291cb0..e97bbad6106 100644 --- a/examples/models/muse-glimmer/export/export_solo.py +++ b/examples/models/muse-glimmer/export/export_solo.py @@ -8,10 +8,28 @@ Supports GGUF, consolidated BF16, prequantized, and MLX-affine inputs. Vision exports add ``vision_encoder`` and write ``pos_embed.bin``. + +``--use-offgraph-kv-cache`` keeps the KV cache out of the program and exports +the batching runtime's contract. On MLX that is the selected-logits decoder. On +CUDA it is the text-only artifact CudaExecutor (backends/cuda/batching) drives; +without the flag CUDA keeps its in-graph prefill and decode methods: + +- ``forward_{N}`` for each N in ``--forward-sizes`` (default 1, 2, 4, ..., 64): + T = K = N, static, each captured into a CUDA graph at runtime, its quantized + linears traced through the smallest QuantizedGemmFamily bucket holding N. +- ``forward_others``: T in [max N, ``--max-step-tokens``], dynamic. +- ``sample`` / ``sample_argmax`` (extension/llm/batching/sampler.py), so logits + stay on the device; ``--logits`` returns them to the host instead. + +Every sequence of a batch shares one pool of ``--max-cells`` per-token cells. +The sizes stop at 64, the widest quantized GEMM bucket: above it the quantized +linears take dequant + matmul, which ``forward_others`` runs unpadded. """ import argparse import gc +import json +import types import warnings import torch @@ -101,6 +119,9 @@ def export_and_lower( pos_embed_table: torch.Tensor | None = None, max_vision_patches: int = 16384, vision_fp32_mm: str = "none", + forward_sizes=None, + max_step_tokens: int = 4096, + max_cells: int = 65536, ) -> None: if use_offgraph_kv_cache and use_turboquant: raise ValueError("off-graph KV cache and TurboQuant are mutually exclusive") @@ -111,14 +132,25 @@ def export_and_lower( ) if not 1 <= max_prefill_chunk <= config.max_seq_len: raise ValueError("max_prefill_chunk must be within [1, max_seq_len]") - if backend == "cuda": + if backend == "cuda" and use_offgraph_kv_cache: + if vision_model is not None: + raise ValueError("the CUDA batching export is text-only") + export_cuda_batching( + model, + config, + output_dir, + max_step_tokens=max_step_tokens, + max_cells=max_cells, + forward_sizes=forward_sizes or DEFAULT_FORWARD_SIZES, + device_sampling=sample, + ) + elif backend == "cuda": _export_cuda( model, config, 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, @@ -202,13 +234,281 @@ def _solo_constant_methods( return constant_methods +# CUDA batching export (--use-offgraph-kv-cache), over a cell-layout KV cache. + +# Read by CudaExecutor: the cell pool's size, which fixes its step buffers. +MAX_CELLS_METHOD = "get_offgraph_kv_max_cells" +# 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 +# cache holds at most 64 sequences. +MAX_SAMPLE_ROWS = 64 +SAMPLE_METHOD = "sample" +ARGMAX_METHOD = "sample_argmax" + + +def step_width_spec(method: str) -> bytes: + """``delegate_input:dim`` of input_pos[T], the step's token count. + + Every forward method's delegate keeps the exported (tokens, input_pos, + logits_to_keep) input order. The CUDA backend's CheckOffGraphKVStepWidthPass + fails the export if this is wrong. + """ + return b"1:0" + + +def step_forward( + model: MuseGlimmerModel, + tokens: torch.Tensor, + input_pos: torch.Tensor, + logits_to_keep: torch.Tensor, +) -> torch.Tensor: + """Logits for the selected tokens of one packed step: [K, vocab] float32. + + The whole step runs through the decoder so every token's K/V reaches the + cache; only the rows the batch samples from go through the LM head. + """ + x = model._run_blocks(model.embed_text(tokens), input_pos) + x = model.output_norm(x[:, logits_to_keep, :]) + return model._soft_cap(model.lm_head(x))[0] + + +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) + manifest["layout"] = "cell" + manifest["max_cells"] = max_cells + return json.dumps(manifest, sort_keys=True, separators=(",", ":")) + + +def sampler_programs(vocab_size: int, max_rows: int = MAX_SAMPLE_ROWS): + """The device samplers, over up to ``max_rows`` float32 logits rows.""" + from executorch.extension.llm.batching.sampler import ( + BatchArgmax, + BatchSampler, + NUM_PARAMS, + ) + from torch.export import Dim, export + + rows = Dim("sample_rows", min=1, max=max_rows) + logits = torch.zeros(max_rows, vocab_size) + params = torch.zeros(max_rows, NUM_PARAMS) + return { + SAMPLE_METHOD: export( + BatchSampler(), + (logits, params), + dynamic_shapes=({0: rows}, {0: rows}), + strict=True, + ), + ARGMAX_METHOD: export( + BatchArgmax(), (logits,), dynamic_shapes=({0: rows},), strict=True + ), + } + + +def device_residency(forward_methods, device_sampling: bool): + """Per-method ``PropagateDeviceConfig`` for ``to_executorch``. + + With device sampling the forwards keep their logits on the device, planned + in device memory as export_solo's outputs are, and the samplers take them + there without a host round trip. + """ + from executorch.exir.passes.propagate_device_config import PropagateDeviceConfig + + device = { + name: PropagateDeviceConfig(skip_d2h_for_method_outputs=device_sampling) + for name in forward_methods + } + if device_sampling: + for name in (SAMPLE_METHOD, ARGMAX_METHOD): + device[name] = PropagateDeviceConfig(skip_h2d_for_method_inputs=True) + return device + + +def export_cuda_batching( + model: MuseGlimmerModel, + config: MuseGlimmerConfig, + output_dir: str, + max_step_tokens: int, + max_cells: int, + forward_sizes=DEFAULT_FORWARD_SIZES, + device_sampling: bool = True, +): + """Exports the forward methods and writes the artifact. + + Returns the ExecutorchProgramManager so callers can inspect it. + """ + import executorch.backends.cuda.quantize_op_dispatch # noqa: F401 + 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_COMPILE_SPEC, + OFFGRAPH_KV_STEP_WIDTH_COMPILE_SPEC, + ) + from executorch.examples.models.muse_glimmer.source_transformations.cuda import ( + enable_offgraph_kv_cache, + offgraph_kv_cache_geometry, + ) + from executorch.exir import ( + EdgeCompileConfig, + ExecutorchBackendConfig, + to_edge_transform_and_lower, + ) + from executorch.exir.backend.compile_spec_schema import CompileSpec + from executorch.exir.passes import MemoryPlanningPass + from executorch.extension.llm.export.model_metadata import ( + write_logits_to_keep_mode, + write_max_context_len, + ) + from torch.export import Dim, export + + forward_sizes = tuple(sorted({int(n) for n in forward_sizes})) + widest_static = forward_sizes[-1] if forward_sizes else 0 + if not forward_sizes or forward_sizes[0] != 1: + raise ValueError(f"forward_sizes must include 1, got {forward_sizes}") + # forward_others is dynamic from the widest static size, and its linears + # trace from five tokens. + if widest_static < _CUDA_MIN_PREFILL_CHUNK - 1: + raise ValueError( + 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: + raise ValueError( + f"max_step_tokens must be in ({widest_static}, max_cells), " + f"got {max_step_tokens}" + ) + if max_step_tokens >= config.max_seq_len: + raise ValueError("max_step_tokens must be below the model's context") + + inductor_config.coordinate_descent_tuning = False + inductor_config.aot_inductor.compile_wrapper_opt_level = "O0" + # The PCH path shells out to `openssl sha512` and is flaky; it is only a + # compile-time optimization. + inductor_config.aot_inductor.precompile_headers = False + 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 + ) + geometry = offgraph_kv_cache_geometry(model) + + width = Dim( + "width", min=max(widest_static, _CUDA_MIN_PREFILL_CHUNK), max=max_step_tokens + ) + rows = Dim("rows", min=_CUDA_MIN_PREFILL_CHUNK, max=max_step_tokens) + programs = {} + with common.BoundMethodForward( + model, types.MethodType(step_forward, model) + ), torch.no_grad(): + for n in forward_sizes: + print(f"Exporting forward_{n} (T={n})...") + programs[f"forward_{n}"] = export( + model, + ( + torch.zeros((1, n), dtype=torch.long), + torch.arange(n, dtype=torch.long), + torch.zeros(n, dtype=torch.long), + ), + strict=True, + ) + print( + f"Exporting forward_others (T in [{max(widest_static, _CUDA_MIN_PREFILL_CHUNK)}, " + f"{max_step_tokens}])..." + ) + programs["forward_others"] = export( + model, + ( + 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), + ), + dynamic_shapes=({1: width}, {0: width}, {0: rows}), + strict=True, + ) + del model + gc.collect() + torch.cuda.empty_cache() + forward_methods = list(programs) + if device_sampling: + print(f"Exporting {SAMPLE_METHOD} and {ARGMAX_METHOD}...") + programs.update(sampler_programs(config.vocab_size)) + + def partitioner(name: str) -> CudaPartitioner: + if name not in forward_methods: + return CudaPartitioner( + [ + CudaBackend.generate_method_name_compile_spec(name), + CompileSpec("autotune_at_compile_time", b"OFF"), + ] + ) + return CudaPartitioner( + [ + CudaBackend.generate_method_name_compile_spec(name), + CompileSpec("low_memory_mode", b"ON"), + CompileSpec(OFFGRAPH_KV_COMPILE_SPEC, manifest.encode()), + CompileSpec(OFFGRAPH_KV_STEP_WIDTH_COMPILE_SPEC, step_width_spec(name)), + CompileSpec("autotune_at_compile_time", b"OFF"), + ] + ) + + constant_methods = _solo_constant_methods( + config=config, + max_prefill=max_step_tokens, + activation_dtype=torch.bfloat16, + mutable_buffer_metadata=None, + has_vision=False, + max_vision_patches=0, + ) + constant_methods.update(geometry) + constant_methods.update(write_max_context_len(config.max_seq_len)) + 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 + + print(f"Lowering {', '.join(programs)} to ExecuTorch (CUDA)...") + edge = to_edge_transform_and_lower( + programs, + partitioner={name: [partitioner(name)] for name in programs}, + compile_config=EdgeCompileConfig( + _check_ir_validity=False, + _skip_dim_order=True, + ), + constant_methods=constant_methods, + ) + del programs + gc.collect() + torch.cuda.empty_cache() + + et_program = edge.to_executorch( + config=ExecutorchBackendConfig( + extract_delegate_segments=True, + do_quant_fusion_and_const_prop=True, + memory_planning_pass=MemoryPlanningPass(alloc_graph_input=False), + propagate_device_config=device_residency(forward_methods, device_sampling), + emit_mutable_buffer_names=True, + ), + ) + del edge + gc.collect() + + common.save_pte(et_program, output_dir, None) + if torch.cuda.is_available(): + peak_mb = torch.cuda.max_memory_allocated() / (1024**2) + print(f"EXPORT_GPU_PEAK_MEMORY_MB: {peak_mb:.1f}") + print("Done.") + return et_program + + def _export_cuda( model: MuseGlimmerModel, config: MuseGlimmerConfig, 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, @@ -217,9 +517,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, @@ -249,8 +546,6 @@ 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, ) @@ -258,19 +553,8 @@ def _export_cuda( # runtime kv_len (O(context) decode). With use_turboquant=True it also swaps # the global KV caches for TurboQuant TQ4. - # 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 = 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) + cuda_source_transformations(model, use_turboquant=use_turboquant) has_vision = vision_model is not None programs: dict[str, "torch.export.ExportedProgram"] = {} @@ -337,30 +621,18 @@ def _export_cuda( vision_model, pos_embed_table, max_vision_patches ) - mutable_buffer_metadata = ( - None if offgraph_manifest else common.mutable_buffer_metadata(model) - ) + mutable_buffer_metadata = common.mutable_buffer_metadata(model) del model, 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, @@ -372,13 +644,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): " @@ -709,8 +974,31 @@ def main() -> None: parser.add_argument( "--use-offgraph-kv-cache", action="store_true", - help="CUDA and MLX: allocate KV cache at runtime instead of storing it in " - "the PTE/PTD. MLX also exports the batching selected-logits decoder contract.", + help="Allocate the KV cache at runtime and export the batching runtime's " + "contract: MLX's selected-logits decoder, or CUDA's text-only " + "forward_{N} + forward_others over a cell-layout cache. Without it CUDA " + "exports its in-graph prefill and decode methods.", + ) + parser.add_argument( + "--forward-sizes", + default=",".join(map(str, DEFAULT_FORWARD_SIZES)), + help="CUDA with --use-offgraph-kv-cache: comma-separated token counts " + "of the static forward_{N} methods.", + ) + parser.add_argument( + "--max-step-tokens", + type=int, + default=4096, + help="CUDA with --use-offgraph-kv-cache: forward_others' widest step, " + "the most tokens one forward carries.", + ) + parser.add_argument( + "--max-cells", + type=int, + default=65536, + help="CUDA with --use-offgraph-kv-cache: cells in the shared KV pool, " + "the tokens every resident session together may hold. Memory grows with " + "use, up to this.", ) parser.add_argument( "--activation-dtype", @@ -756,6 +1044,10 @@ def main() -> None: args = parser.parse_args() _validate_offgraph_prefill(parser, args) + if args.backend == "cuda" and args.use_offgraph_kv_cache and args.mmproj: + parser.error( + "--use-offgraph-kv-cache on CUDA exports a text-only model; omit --mmproj." + ) if args.backend == "cuda" and not torch.cuda.is_available(): parser.error("CUDA is required for the cuda backend.") @@ -847,6 +1139,9 @@ def main() -> None: pos_embed_table=pos_embed_table, max_vision_patches=args.max_vision_patches, vision_fp32_mm=args.vision_fp32_mm, + 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, ) diff --git a/examples/models/muse-glimmer/tests/test_checkpoint_context.py b/examples/models/muse-glimmer/tests/test_checkpoint_context.py index 022036fce17..355ed4700a7 100644 --- a/examples/models/muse-glimmer/tests/test_checkpoint_context.py +++ b/examples/models/muse-glimmer/tests/test_checkpoint_context.py @@ -304,20 +304,46 @@ def test_legacy_explicit_context(self): class CUDAExportContextTest(unittest.TestCase): - def test_offgraph_cli_dispatch(self): + def run_cli(self, *flags): config = MuseGlimmerConfig(max_seq_len=32) - argv = "export_solo --gguf unused --backend cuda --max-seq-len 32 --use-offgraph-kv-cache".split() + argv = [ + "export_solo", + "--gguf", + "unused", + "--backend", + "cuda", + "--max-seq-len", + "32", + *flags, + ] with ( patch.object(sys, "argv", argv), patch.object(export_solo.torch.cuda, "is_available", return_value=True), patch.object( loaders, "load_gguf_model", return_value=(sentinel.model, config) ) as loader, - patch.object(export_solo, "_export_cuda") as export, + patch.object(export_solo, "_export_cuda") as solo, + patch.object(export_solo, "export_cuda_batching") as batching, ): export_solo.main() self.assertEqual(loader.call_args.kwargs["max_seq_len"], 32) self.assertNotIn("defer_runtime_buffers", loader.call_args.kwargs) - export.assert_called_once() - self.assertEqual(export.call_args.args[:2], (sentinel.model, config)) - self.assertTrue(export.call_args.kwargs["use_offgraph_kv_cache"]) + return config, solo, batching + + def test_offgraph_cli_exports_the_batching_methods(self): + config, solo, batching = self.run_cli( + "--use-offgraph-kv-cache", "--forward-sizes", "1,2,4", "--max-cells", "512" + ) + solo.assert_not_called() + batching.assert_called_once() + self.assertEqual(batching.call_args.args[:2], (sentinel.model, config)) + self.assertEqual(batching.call_args.kwargs["forward_sizes"], [1, 2, 4]) + self.assertEqual(batching.call_args.kwargs["max_cells"], 512) + self.assertTrue(batching.call_args.kwargs["device_sampling"]) + + def test_in_graph_cli_keeps_prefill_and_decode(self): + config, solo, batching = self.run_cli() + batching.assert_not_called() + solo.assert_called_once() + self.assertEqual(solo.call_args.args[:2], (sentinel.model, config)) + self.assertNotIn("use_offgraph_kv_cache", solo.call_args.kwargs) diff --git a/examples/models/muse-glimmer/tests/test_cuda_batching_pipeline.py b/examples/models/muse-glimmer/tests/test_cuda_batching_pipeline.py new file mode 100644 index 00000000000..dae6fac83e6 --- /dev/null +++ b/examples/models/muse-glimmer/tests/test_cuda_batching_pipeline.py @@ -0,0 +1,303 @@ +# 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. + +"""Tests for the batched CUDA export (export_solo.py --kv-layout cell) on a tiny model. + +The eager checks run the off-graph model against the neutral cache's +references; the export check lowers the forward methods and inspects the +artifact. Requires CUDA for the export. + + python -m pytest examples/models/muse-glimmer/tests/test_cuda_batching_pipeline.py -v +""" + +import copy +import json +import os +import re +import tempfile +import unittest + +import executorch.backends.cuda.quantize_op_dispatch as _quantize_op_dispatch # noqa: F401 +import torch +from executorch.backends.cuda.triton.kernels.int4_quantized_gemm import ( + SUPPORTED_BUCKETS as INT4_BUCKETS, +) +from executorch.examples.models.muse_glimmer.export.export_solo import ( + cell_manifest, + export_cuda_batching, + load_prequantized_model, + MAX_CELLS_METHOD, + step_forward, +) +from executorch.examples.models.muse_glimmer.source_transformations.cuda import ( + enable_offgraph_kv_cache, +) +from executorch.examples.models.muse_glimmer.tests.test_pipeline import ( + build_random_tiny_model, + save_checkpoint, + TINY_CONFIG, +) +from executorch.extension.llm.cache.reference_cache import ( + CacheConfig, + CellReferenceCache, + LayerPolicy, + SequenceReferenceCache, +) +from executorch.extension.llm.cache.update_and_attend import REGISTRY + + +def _cache_config(model, capacity: int) -> CacheConfig: + attn = model.layers[0].self_attn + return CacheConfig( + n_layers=TINY_CONFIG.n_layers, + n_kv_heads=attn.n_kv_heads, + head_dim=attn.head_dim, + capacity=capacity, + layers=tuple( + ( + LayerPolicy.ring(layer.self_attn.window_size) + if layer.self_attn.is_sliding + else LayerPolicy.flat() + ) + for layer in model.layers + ), + ) + + +class BatchingStepForwardTest(unittest.TestCase): + """The exported step computes what the model does, per packed sequence.""" + + def setUp(self) -> None: + self.model = build_random_tiny_model() + self.off_graph = copy.deepcopy(self.model) + enable_offgraph_kv_cache(self.off_graph, 16) + + def _install(self, cache) -> str: + key = f"muse-batching-{id(self)}-{id(cache)}" + REGISTRY.install(key, cache) + self.addCleanup(REGISTRY.uninstall, key) + return key + + def test_selected_rows_match_the_full_forward(self) -> None: + key = self._install( + SequenceReferenceCache( + _cache_config(self.off_graph, TINY_CONFIG.max_seq_len) + ) + ) + generator = torch.Generator().manual_seed(0) + start = 0 + for length, keep in ((12, [3, 11]), (1, [0]), (6, [5])): + tokens = torch.randint( + 0, TINY_CONFIG.vocab_size, (1, length), generator=generator + ) + input_pos = torch.arange(start, start + length) + with torch.no_grad(): + expected = self.model(tokens, input_pos)[0, keep] + with REGISTRY.active(key): + actual = step_forward( + self.off_graph, tokens, input_pos, torch.tensor(keep) + ) + self.assertEqual(actual.dtype, torch.float32) + self.assertEqual(tuple(actual.shape), (len(keep), TINY_CONFIG.vocab_size)) + self.assertLess((expected - actual).abs().max().item(), 5e-2) + start += length + + def test_packed_sequences_match_each_alone(self) -> None: + # Two prompts prefilled in one step, then decoded together, over the + # neutral cell cache: each sequence's logits must equal running it on + # its own -- per-token RoPE, windows and caches all kept apart. + generator = torch.Generator().manual_seed(1) + a = torch.randint(0, TINY_CONFIG.vocab_size, (1, 20), generator=generator) + b = torch.randint(0, TINY_CONFIG.vocab_size, (1, 7), generator=generator) + a_next = torch.randint(0, TINY_CONFIG.vocab_size, (1, 1), generator=generator) + b_next = torch.randint(0, TINY_CONFIG.vocab_size, (1, 1), generator=generator) + + def alone(prompt, next_token): + key = self._install( + SequenceReferenceCache( + _cache_config(self.off_graph, TINY_CONFIG.max_seq_len) + ) + ) + length = prompt.shape[1] + with torch.no_grad(), REGISTRY.active(key): + first = step_forward( + self.off_graph, + prompt, + torch.arange(length), + torch.tensor([length - 1]), + ) + second = step_forward( + self.off_graph, + next_token, + torch.tensor([length]), + torch.tensor([0]), + ) + return first[0], second[0] + + a_first, a_second = alone(a, a_next) + b_first, b_second = alone(b, b_next) + + cells = CellReferenceCache(_cache_config(self.off_graph, 64)) + key = self._install(cells) + with torch.no_grad(), REGISTRY.active(key): + cells.declare_step([0] * 20 + [1] * 7) + prefill = step_forward( + self.off_graph, + torch.cat([a, b], dim=1), + torch.cat([torch.arange(20), torch.arange(7)]), + torch.tensor([19, 26]), + ) + cells.declare_step([1, 0]) + decode = step_forward( + self.off_graph, + torch.cat([b_next, a_next], dim=1), + torch.tensor([7, 20]), + torch.tensor([0, 1]), + ) + for actual, expected in ( + (prefill[0], a_first), + (prefill[1], b_first), + (decode[1], a_second), + (decode[0], b_second), + ): + self.assertLess((expected - actual).abs().max().item(), 5e-2) + self.assertEqual(int(expected.argmax()), int(actual.argmax())) + + def test_cell_manifest_keeps_every_layer(self) -> None: + model = build_random_tiny_model() + manifest = json.loads(cell_manifest(enable_offgraph_kv_cache(model, 8), 128)) + self.assertEqual(manifest["layout"], "cell") + self.assertEqual(manifest["max_cells"], 128) + self.assertEqual(manifest["max_write"], 8) + self.assertEqual(len(manifest["layers"]), TINY_CONFIG.n_layers) + + +class BatchingExportTest(unittest.TestCase): + MAX_STEP = 16 + FORWARD_SIZES = (1, 2, 4, 8) + MAX_CELLS = 128 + + def setUp(self) -> None: + if not torch.cuda.is_available(): + self.skipTest("CUDA required") + + def test_exports_forward_methods_over_a_cell_cache(self) -> None: + from executorch.backends.cuda.passes.lower_offgraph_kv import ( + OFFGRAPH_KV_COMPILE_SPEC, + OFFGRAPH_KV_STEP_WIDTH_COMPILE_SPEC, + ) + from executorch.exir.lowered_backend_module import get_lowered_submodules + + with ( + tempfile.TemporaryDirectory() as ckpt_dir, + tempfile.TemporaryDirectory() as out_dir, + ): + save_checkpoint(ckpt_dir) + model, config = load_prequantized_model( + ckpt_dir, max_seq_len=TINY_CONFIG.max_seq_len + ) + program = export_cuda_batching( + model, + config, + out_dir, + max_step_tokens=self.MAX_STEP, + forward_sizes=self.FORWARD_SIZES, + max_cells=self.MAX_CELLS, + ) + self.assertTrue(os.path.exists(os.path.join(out_dir, "model.pte"))) + self.assertTrue(any(name.endswith(".ptd") for name in os.listdir(out_dir))) + + methods = set(program.methods) + forwards = {f"forward_{n}" for n in self.FORWARD_SIZES} | {"forward_others"} + self.assertEqual( + forwards, + methods & (forwards | {m for m in methods if m.startswith("forward_")}), + ) + # Tokens are sampled on the device by default. + self.assertLessEqual({"sample", "sample_argmax"}, methods) + for name in ("sample", "sample_argmax"): + lowered = get_lowered_submodules( + program.exported_program(name).graph_module + ) + self.assertEqual(len(lowered), 1, name) + specs = {spec.key for spec in lowered[0][1].compile_specs} + self.assertNotIn(OFFGRAPH_KV_COMPILE_SPEC, specs, name) + self.assertFalse( + {"embed_text", "forward_from_embeddings", "decode_from_embedding"} + & methods + ) + for name in sorted(forwards): + graph = program.exported_program(name).graph_module + lowered = get_lowered_submodules(graph) + self.assertEqual(len(lowered), 1, name) + # Every method traces its INT4 linears through one + # QuantizedGemmFamily bucket: the smallest that holds its + # width (forward_others: its widest step, MAX_STEP). + inner = { + str(node.target) + for node in lowered[0][1].original_module.graph.nodes + } + bucket_ops = { + m.group(0) + for t in inner + for m in [re.search(r"int4_quantized_gemm_m\d+", t)] + if m + } + width = ( + self.MAX_STEP + if name == "forward_others" + else int(name.split("_")[1]) + ) + bucket = min(b for b in INT4_BUCKETS if b >= width) + self.assertEqual( + bucket_ops, + {f"int4_quantized_gemm_m{bucket}"}, + name, + ) + specs = {spec.key: spec.value for spec in lowered[0][1].compile_specs} + manifest = json.loads(specs[OFFGRAPH_KV_COMPILE_SPEC]) + self.assertEqual(manifest["layout"], "cell") + self.assertEqual(manifest["max_cells"], self.MAX_CELLS) + self.assertEqual(manifest["max_write"], self.MAX_STEP) + # The step width must name the delegate input carrying T, which + # the backend checks is the cache ops' position. The delegate's + # inputs need not follow the method's order: prefill passes + # logits_to_keep ahead of input_pos. + index, dim = map( + int, specs[OFFGRAPH_KV_STEP_WIDTH_COMPILE_SPEC].decode().split(":") + ) + call = next( + n + for n in graph.graph.nodes + if "executorch_call_delegate" in str(n.target) + ) + tokens = next( + n + for n in graph.graph.nodes + if n.op == "placeholder" and n.name == "tokens" + ) + self.assertEqual( + str(call.args[1 + index].meta["val"].shape[dim]), + str(tokens.meta["val"].shape[1]), + name, + ) + + from executorch.runtime import Runtime, Verification + + loaded = Runtime.get().load_program( + os.path.join(out_dir, "model.pte"), verification=Verification.Minimal + ) + + def constant(name): + return loaded.load_method(name).execute([])[0] + + self.assertEqual(constant(MAX_CELLS_METHOD), self.MAX_CELLS) + self.assertEqual(constant("get_max_context_len"), TINY_CONFIG.max_seq_len) + self.assertEqual(constant("get_n_caches"), TINY_CONFIG.n_layers) + self.assertEqual(constant("get_max_prefill_chunk"), self.MAX_STEP) + self.assertEqual(constant("get_min_prefill_chunk"), 5) + # LogitsToKeepMode::Selected. + self.assertEqual(constant("get_logits_to_keep_mode"), 2) From 1c1395f6df359a0360d645204045776d5430eb85 Mon Sep 17 00:00:00 2001 From: gasoonjia Date: Sun, 18 Oct 2026 23:20:33 -0700 Subject: [PATCH 2/5] Update [ghstack-poisoned] --- .../muse-glimmer/tests/test_checkpoint_context.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/examples/models/muse-glimmer/tests/test_checkpoint_context.py b/examples/models/muse-glimmer/tests/test_checkpoint_context.py index 355ed4700a7..8c25754e9b7 100644 --- a/examples/models/muse-glimmer/tests/test_checkpoint_context.py +++ b/examples/models/muse-glimmer/tests/test_checkpoint_context.py @@ -129,9 +129,10 @@ def test_invalid_gguf_context_fails_before_weights(self): if value_type is None else [("muse-glimmer.context_length", value, value_type)] ) - with self.subTest(value=value), patch.object( - loaders, "_atomic_sd_from_gguf" - ) as weights: + with ( + self.subTest(value=value), + patch.object(loaders, "_atomic_sd_from_gguf") as weights, + ): with self.assertRaisesRegex(ValueError, "Native context"): loaders.load_gguf_model(self.write_gguf(fields), max_seq_len=None) weights.assert_not_called() @@ -233,9 +234,10 @@ def test_offgraph_context_and_prefill_routing(self): def test_offgraph_default_context_is_silent(self): for limit in (None, 131072): - with self.subTest(limit=limit), patch.object( - export_solo.warnings, "warn" - ) as warn: + with ( + self.subTest(limit=limit), + patch.object(export_solo.warnings, "warn") as warn, + ): self.run_cli("--gguf", max_seq_len=limit) warn.assert_not_called() self.assertIsNone( From 79a253782c238f68a57183deda3af58383daaa2f Mon Sep 17 00:00:00 2001 From: gasoonjia Date: Mon, 19 Oct 2026 01:14:38 -0700 Subject: [PATCH 3/5] Update [ghstack-poisoned] --- pytest.ini | 1 + 1 file changed, 1 insertion(+) diff --git a/pytest.ini b/pytest.ini index 5ef01fde9a0..db0dc878130 100644 --- a/pytest.ini +++ b/pytest.ini @@ -39,6 +39,7 @@ addopts = # collecting the directory, not when a path is passed on the command line. --ignore=examples/models/muse-glimmer/tests/test_mlx_pipeline.py --ignore=examples/models/muse-glimmer/tests/test_cuda_pipeline.py + --ignore=examples/models/muse-glimmer/tests/test_cuda_batching_pipeline.py # Ignore tests with missing custom_ops_generated_lib dependencies --ignore=exir/backend/test/demos/test_delegate_aten_mode.py From b620e9f57b90e41a3d85db61715c0473e7c5a6f1 Mon Sep 17 00:00:00 2001 From: gasoonjia Date: Mon, 19 Oct 2026 12:50:40 -0700 Subject: [PATCH 4/5] Update [ghstack-poisoned] --- backends/cuda/tests/test_batch_sampler.py | 87 +++++++++++++++++++++++ extension/llm/batching/sampler.py | 6 +- 2 files changed, 92 insertions(+), 1 deletion(-) create mode 100644 backends/cuda/tests/test_batch_sampler.py diff --git a/backends/cuda/tests/test_batch_sampler.py b/backends/cuda/tests/test_batch_sampler.py new file mode 100644 index 00000000000..637a439dcc6 --- /dev/null +++ b/backends/cuda/tests/test_batch_sampler.py @@ -0,0 +1,87 @@ +# 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. + +import math +import tempfile +import unittest +from pathlib import Path + +import torch +from executorch.backends.cuda.cuda_backend import CudaBackend +from executorch.backends.cuda.cuda_partitioner import CudaPartitioner +from executorch.exir import ( + EdgeCompileConfig, + ExecutorchBackendConfig, + to_edge_transform_and_lower, +) +from executorch.exir.passes import MemoryPlanningPass +from executorch.exir.passes.propagate_device_pass import PropagateDeviceConfig +from executorch.extension.llm.batching.sampler import BatchSampler +from executorch.runtime import Runtime + + +def export_sampler(directory: Path, vocab: int): + rows = torch.export.Dim("rows", min=1, max=64) + program = torch.export.export( + BatchSampler(), + (torch.zeros(64, vocab), torch.zeros(64, 4)), + dynamic_shapes=({0: rows}, {0: rows}), + strict=True, + ) + specs = [CudaBackend.generate_method_name_compile_spec("sample")] + edge = to_edge_transform_and_lower( + {"sample": program}, + partitioner={"sample": [CudaPartitioner(specs)]}, + compile_config=EdgeCompileConfig(_check_ir_validity=False, _skip_dim_order=True), + ) + result = edge.to_executorch( + config=ExecutorchBackendConfig( + extract_delegate_segments=True, + do_quant_fusion_and_const_prop=True, + memory_planning_pass=MemoryPlanningPass(alloc_graph_input=False), + propagate_device_config={ + "sample": PropagateDeviceConfig(skip_h2d_for_method_inputs=True) + }, + ) + ) + directory.mkdir(parents=True, exist_ok=True) + with (directory / "model.pte").open("wb") as output: + result.write_to_file(output) + return directory / "model.pte" + + +@unittest.skipUnless(torch.cuda.is_available(), "CUDA required") +class CudaBatchSamplerTest(unittest.TestCase): + def test_dynamic_rows_with_masked_reduction_lanes(self): + # A vocabulary spanning several reduction tiles exposes padding lanes. + vocab = 202048 + with tempfile.TemporaryDirectory() as temporary: + path = export_sampler(Path(temporary), vocab) + method = Runtime.get().load_program(path).load_method("sample") + for count in (1, 3, 4, 5, 64): + with self.subTest(rows=count): + logits = torch.full((count, vocab), -20.0) + logits[:, :8] = torch.arange(8).float() / 2 + settings = [(0.0, 1.0, 0.0, 0.0), (0.7, 0.8, 4.0, 0.37)] + params = torch.tensor([settings[i % 2] for i in range(count)]) + if count == 1: + params[0] = torch.tensor(settings[1]) + expected = [] + for temperature, _, top_k, coin in params.tolist(): + if temperature == 0: + expected.append(7) + continue + order = list(range(7, 7 - int(top_k), -1)) + mass = [math.exp((i / 2 - 3.5) / temperature) for i in order] + threshold = coin * sum(mass) + cumulative = 0.0 + for token, probability in zip(order, mass): + cumulative += probability + if cumulative > threshold: + expected.append(token) + break + actual = method.execute([logits.cuda(), params.cuda()])[0] + self.assertEqual(actual.cpu().tolist(), expected) diff --git a/extension/llm/batching/sampler.py b/extension/llm/batching/sampler.py index 711c2f3527e..63aace75984 100644 --- a/extension/llm/batching/sampler.py +++ b/extension/llm/batching/sampler.py @@ -85,7 +85,11 @@ def sample_rows(logits: torch.Tensor, params: torch.Tensor) -> torch.Tensor: first_pass = torch.where(passes, ranks, torch.full_like(ranks, vocab_size)).amin( dim=-1, keepdim=True ) - sorted_choice = torch.where(first_pass < vocab_size, first_pass, last_kept) + # The valid distribution always keeps a token. Explicit bounds also + # protect Inductor's masked reduction lanes from an unused gather index. + sorted_choice = torch.where(first_pass < vocab_size, first_pass, last_kept).clamp( + min=0, max=vocab_size - 1 + ) truncated = sorted_indices.gather(-1, sorted_choice).squeeze(-1) # The full distribution samples in token order, as the host does. From d2cfc785f704661646a44af4f4497cc8a29b50c0 Mon Sep 17 00:00:00 2001 From: gasoonjia Date: Mon, 19 Oct 2026 13:07:54 -0700 Subject: [PATCH 5/5] Update [ghstack-poisoned] --- backends/cuda/tests/test_batch_sampler.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/backends/cuda/tests/test_batch_sampler.py b/backends/cuda/tests/test_batch_sampler.py index 637a439dcc6..52bcb87d76a 100644 --- a/backends/cuda/tests/test_batch_sampler.py +++ b/backends/cuda/tests/test_batch_sampler.py @@ -35,7 +35,9 @@ def export_sampler(directory: Path, vocab: int): edge = to_edge_transform_and_lower( {"sample": program}, partitioner={"sample": [CudaPartitioner(specs)]}, - compile_config=EdgeCompileConfig(_check_ir_validity=False, _skip_dim_order=True), + compile_config=EdgeCompileConfig( + _check_ir_validity=False, _skip_dim_order=True + ), ) result = edge.to_executorch( config=ExecutorchBackendConfig(