[executorch][muse] Export packed DFlash drafts for CUDA batching

.github/workflows/cuda.yml

@@ -618,6 +618,7 @@ jobs:
         pip install gguf
         python -m pytest -v -o "addopts=" \
           examples/models/muse-glimmer/tests/test_cuda_batching_pipeline.py \
+          examples/models/muse-glimmer/tests/test_cuda_dflash_packed.py \
           examples/models/muse-glimmer/tests/test_checkpoint_context.py
 
   unittest-cuda-quantized-gemm:

examples/models/muse-glimmer/export/export_dflash.py

@@ -24,8 +24,6 @@
 
 
 def validate_dflash_export_options(backend: str, use_offgraph_kv_cache=False) -> None:
-    if use_offgraph_kv_cache and backend != "mlx":
-        raise ValueError("Off-graph DFlash requires MLX")
     if backend not in {"cuda", "mlx"}:
         raise ValueError(f"Unsupported DFlash backend: {backend}")
 
@@ -129,6 +127,11 @@ def export_dflash(
     mmproj: str | None = None,
     max_vision_patches: int = 16384,
     use_offgraph_kv_cache: bool = False,
+    max_step_tokens: int = 4096,
+    max_cells: int = 65536,
+    draft_max_cells: int | None = None,
+    max_draft_owners: int = 4,
+    forward_sizes=None,
 ) -> None:
     """Export DFlash target + draft to one CUDA or MLX .pte.
 
@@ -140,6 +143,8 @@ def export_dflash(
     mmproj loader and vision export contract to add ``vision_encoder``.
     """
     validate_dflash_export_options(backend, use_offgraph_kv_cache)
+    if use_offgraph_kv_cache and backend == "cuda" and mmproj is not None:
+        raise ValueError("the CUDA batching DFlash export is text-only")
 
     import executorch.extension.llm.export.gguf  # noqa: F401
     import executorch.extension.llm.export.int4  # noqa: F401
@@ -197,6 +202,27 @@ def export_dflash(
             activation_dtype=activation_dtype,
         )
 
+    if use_offgraph_kv_cache and backend == "cuda":
+        from executorch.examples.models.muse_glimmer.export.export_solo import (
+            DEFAULT_FORWARD_SIZES,
+            export_cuda_batching,
+        )
+
+        export_cuda_batching(
+            target_model,
+            target_config,
+            output_dir,
+            max_step_tokens=max_step_tokens,
+            max_cells=max_cells,
+            forward_sizes=(
+                DEFAULT_FORWARD_SIZES if forward_sizes is None else forward_sizes
+            ),
+            dflash_draft=draft_model,
+            max_draft_owners=max_draft_owners,
+            draft_max_cells=draft_max_cells,
+        )
+        return
+
     vision_model = None
     pos_embed_table = None
     if mmproj is not None:
@@ -229,6 +255,111 @@ def export_dflash(
     )
 
 
+def prepare_cuda_packed_draft(
+    target_model,
+    target_config,
+    draft_model,
+    *,
+    max_step_tokens: int,
+    max_cells: int,
+    max_owners: int,
+):
+    """Trace the packed draft and publish its independent cell geometry."""
+    from executorch.examples.models.muse_glimmer.export.export_solo import cell_manifest
+    from executorch.examples.models.muse_glimmer.model.dflash_model import (
+        DFlashPackedExportModel,
+    )
+    from executorch.examples.models.muse_glimmer.source_transformations.cuda import (
+        add_dflash_hidden_tapping,
+        enable_dflash_offgraph_kv_cache,
+        offgraph_kv_cache_geometry,
+    )
+    from torch.export import Dim, export
+
+    config = draft_model.config
+    block = config.block_size
+    if (
+        type(block) is not int
+        or block < 2
+        or type(max_owners) is not int
+        or not 1 <= max_owners <= 4
+    ):
+        raise ValueError("packed DFlash requires a trained block >= 2 and 1..4 owners")
+    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:
+        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")
+    if target_model.activation_dtype != torch.bfloat16:
+        raise ValueError("CUDA packed DFlash requires bfloat16 activations")
+    if not 0 <= config.mask_token_id < target_config.vocab_size:
+        raise ValueError("draft mask_token_id must be in the target vocabulary")
+    add_dflash_hidden_tapping(target_model, list(config.target_layers))
+    from executorch.examples.models.gemma4_31b.cuda_packers import (
+        convert_quantized_tensors_for_cuda,
+    )
+
+    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
+    )
+    combined = DFlashPackedExportModel(target_model, draft_model)
+    owners = (
+        Dim.STATIC if max_owners == 1 else Dim("draft_owners", min=1, max=max_owners)
+    )
+    context = Dim("draft_new_context", min=0, max=max_step_tokens)
+    kv_width = (
+        context + block
+        if max_owners == 1
+        else Dim("draft_kv_width", min=block, max=max_write)
+    )
+    sample_owners = min(2, max_owners)
+    sample_queries = sample_owners * block
+    with torch.no_grad():
+        program = export(
+            combined,
+            (
+                torch.zeros((sample_owners, block), dtype=torch.long),
+                torch.zeros(
+                    (1, 2, config.n_target_layers * config.dim), dtype=torch.bfloat16
+                ),
+                torch.arange(2 + sample_queries, dtype=torch.long),
+                torch.arange(2 + sample_queries, dtype=torch.long),
+            ),
+            dynamic_shapes=(
+                {0: owners} if max_owners > 1 else None,
+                {1: context},
+                {0: kv_width},
+                {0: kv_width},
+            ),
+            strict=True,
+        )
+    geometry = offgraph_kv_cache_geometry(draft_model)
+    metadata = {
+        key.replace("get_", "get_draft_", 1): value for key, value in geometry.items()
+    }
+    metadata.update(
+        {
+            "get_block_size": block,
+            "get_mask_token_id": config.mask_token_id,
+            "get_n_target_layers": config.n_target_layers,
+            "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,
+            "get_dflash_target_cache_key": "muse_glimmer_target",
+            "get_dflash_draft_cache_key": "muse_glimmer_draft",
+        }
+    )
+    return {"draft_forward": program}, manifest, metadata
+
+
 def _dflash_constant_methods(
     *,
     target_config,
@@ -926,6 +1057,11 @@ def main() -> None:
         "float16 for MLX and bfloat16 for CUDA.",
     )
     parser.add_argument("--use-offgraph-kv-cache", action="store_true")
+    parser.add_argument("--max-step-tokens", type=int, default=4096)
+    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)
+    parser.add_argument("--forward-sizes", default="1,2,4,8,16,32,64")
     args = parser.parse_args()
 
     try:
@@ -955,6 +1091,11 @@ def main() -> None:
         mmproj=args.mmproj,
         max_vision_patches=args.max_vision_patches,
         use_offgraph_kv_cache=args.use_offgraph_kv_cache,
+        max_step_tokens=args.max_step_tokens,
+        max_cells=args.max_cells,
+        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],
     )
 
 

examples/models/muse-glimmer/export/export_solo.py

@@ -268,9 +268,21 @@ def step_forward(
     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.embed_text(tokens)
+    taps = None
+    if hasattr(model, "_dflash_target_layer_ids"):
+        sliding_mask, global_mask = model._build_masks(input_pos)
+        tapped = []
+        for i, layer in enumerate(model.layers):
+            if i in model._dflash_target_layer_ids:
+                tapped.append(x)
+            x = layer(x, input_pos, sliding_mask, global_mask)
+        taps = torch.cat(tapped, dim=-1)
+    else:
+        x = model._run_blocks(x, input_pos)
     x = model.output_norm(x[:, logits_to_keep, :])
-    return model._soft_cap(model.lm_head(x))[0]
+    logits = model._soft_cap(model.lm_head(x))[0]
+    return logits if taps is None else (logits, taps)
 
 
 def cell_manifest(sequence_manifest: str, max_cells: int) -> str:
@@ -324,14 +336,19 @@ def device_residency(forward_methods, device_sampling: bool):
     return device
 
 
-def export_cuda_batching(
+# Keep target and draft assembly under one shared-weight collector.
+def export_cuda_batching(  # noqa: C901
     model: MuseGlimmerModel,
     config: MuseGlimmerConfig,
     output_dir: str,
     max_step_tokens: int,
     max_cells: int,
     forward_sizes=DEFAULT_FORWARD_SIZES,
     device_sampling: bool = True,
+    *,
+    dflash_draft: nn.Module | None = None,
+    max_draft_owners: int = 4,
+    draft_max_cells: int | None = None,
 ):
     """Exports the forward methods and writes the artifact.
 
@@ -362,6 +379,20 @@ def export_cuda_batching(
     )
     from torch.export import Dim, export
 
+    if dflash_draft is not None:
+        from executorch.examples.models.muse_glimmer.export.export_dflash import (
+            prepare_cuda_packed_draft,
+        )
+
+        draft_programs, draft_manifest, draft_metadata = prepare_cuda_packed_draft(
+            model,
+            config,
+            dflash_draft,
+            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,
+        )
+
     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:
@@ -392,16 +423,20 @@ def export_cuda_batching(
     manifest = cell_manifest(
         enable_offgraph_kv_cache(model, max_step_tokens), max_cells
     )
+    if dflash_draft is not None:
+        # Deferred checkpoints still need scalar/RoPE buffers, after KV removal.
+        materialize_runtime_buffers(model, dtype=torch.bfloat16)
     geometry = offgraph_kv_cache_geometry(model)
 
     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():
+    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(
@@ -431,23 +466,42 @@ def export_cuda_batching(
     gc.collect()
     torch.cuda.empty_cache()
     forward_methods = list(programs)
+    if dflash_draft is not None:
+        programs.update(draft_programs)
+        del draft_programs, dflash_draft
+    cache_methods = forward_methods + (
+        ["draft_forward"] if "draft_forward" in programs else []
+    )
     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:
+        if name not in cache_methods:
             return CudaPartitioner(
                 [CudaBackend.generate_method_name_compile_spec(name)]
             )
-        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)),
-            ]
-        )
+        is_draft = name == "draft_forward"
+        specs = [
+            CudaBackend.generate_method_name_compile_spec(name),
+            CompileSpec("low_memory_mode", b"ON"),
+            CompileSpec(
+                OFFGRAPH_KV_COMPILE_SPEC,
+                (draft_manifest if is_draft else manifest).encode(),
+            ),
+            CompileSpec(
+                OFFGRAPH_KV_STEP_WIDTH_COMPILE_SPEC,
+                b"2:0" if is_draft else step_width_spec(name),
+            ),
+        ]
+        if "draft_forward" in cache_methods:
+            specs.append(
+                CompileSpec(
+                    "llm_cache_registry_key",
+                    b"muse_glimmer_draft" if is_draft else b"muse_glimmer_target",
+                )
+            )
+        return CudaPartitioner(specs)
 
     constant_methods = _solo_constant_methods(
         config=config,
@@ -463,6 +517,8 @@ def partitioner(name: str) -> CudaPartitioner:
     constant_methods["get_min_prefill_chunk"] = _CUDA_MIN_PREFILL_CHUNK
     constant_methods["use_sampling"] = False
     constant_methods[MAX_CELLS_METHOD] = max_cells
+    if "draft_forward" in programs:
+        constant_methods.update(draft_metadata)
 
     print(f"Lowering {', '.join(programs)} to ExecuTorch (CUDA)...")
     edge = to_edge_transform_and_lower(
@@ -483,7 +539,7 @@ def partitioner(name: str) -> CudaPartitioner:
             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),
+            propagate_device_config=device_residency(cache_methods, device_sampling),
             emit_mutable_buffer_names=True,
         ),
     )