[executorch][muse] Export the batching contract for CUDA with --use-offgraph-kv-cache

.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

.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? <img><|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

.github/workflows/cuda.yml

@@ -1063,8 +1063,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

backends/cuda/tests/test_batch_sampler.py

@@ -0,0 +1,89 @@
+# 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)