[executorch][muse-glimmer] Export DFlash sampling methods (pytorch#23145)

This PR was created by the merge bot to help merge the original PR into
the main branch.
ghstack PR number: pytorch#23086 by
@Gasoonjia
^ Please use this as the source of truth for the PR details, comments,
and reviews
ghstack PR base:
https://github.com/pytorch/executorch/tree/gh/Gasoonjia/30/base
ghstack PR head:
https://github.com/pytorch/executorch/tree/gh/Gasoonjia/30/head
Merge bot PR base:
https://github.com/pytorch/executorch/tree/gh/Gasoonjia/29/orig
Merge bot PR head:
https://github.com/pytorch/executorch/tree/gh/Gasoonjia/30/orig
Differential Revision:
[D119463646](https://our.internmc.facebook.com/intern/diff/D119463646/)
@diff-train-skip-merge

---------

Co-authored-by: gasoonjia <gasoonjia@icloud.com>

Author: pytorchbot

examples/models/muse-glimmer/README.md

@@ -119,7 +119,7 @@ advanced inputs.
 |---|---|---|
 | Target | CUDA | `embed_text`, `forward_from_embeddings`, `decode_from_embedding` |
 | Target | MLX | `embed_text`, `forward_from_embeddings` |
-| DFlash | CUDA | `target_forward_from_embeddings`, `target_prefill_from_embeddings`, `embed_text`, `draft_forward`, `draft_prefill` |
+| DFlash | CUDA | `target_forward_from_embeddings`, `target_prefill_from_embeddings`, `embed_text`, `draft_forward`, `draft_prefill`, `dflash_sample_tokens`, `dflash_verify_speculative` |
 | DFlash | MLX | `target_forward_from_embeddings`, `embed_text`, `draft_forward` |
 
 Vision adds `vision_encoder` to each method set.

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

@@ -47,6 +47,51 @@ def _max_draft_prefill_len(draft_config, max_target_prefill: int) -> int:
     return max_target_prefill
 
 
+def _export_cuda_sampler_methods(
+    max_draft_tokens: int, vocab_size: int
+) -> dict[str, torch.export.ExportedProgram]:
+    from executorch.examples.models.muse_glimmer.model.dflash_token_sampler import (
+        DFlashSpeculativeVerifier,
+        DFlashTokenSampler,
+    )
+    from torch.export import Dim, export
+
+    proposals = Dim("dflash_proposals", min=1, max=max_draft_tokens)
+    temperature = torch.tensor([1.0])
+    top_k = torch.tensor([0], dtype=torch.long)
+    top_p = torch.tensor([1.0])
+    return {
+        "dflash_sample_tokens": export(
+            DFlashTokenSampler(),
+            (torch.zeros(max_draft_tokens, vocab_size), temperature, top_k, top_p),
+            dynamic_shapes=({0: proposals}, None, None, None),
+            strict=True,
+        ),
+        "dflash_verify_speculative": export(
+            DFlashSpeculativeVerifier(),
+            (
+                torch.zeros(max_draft_tokens + 1, vocab_size),
+                torch.zeros(max_draft_tokens, vocab_size),
+                torch.zeros(max_draft_tokens + 1, dtype=torch.long),
+                temperature,
+                top_k,
+                top_p,
+                torch.tensor([False]),
+            ),
+            dynamic_shapes=(
+                {0: proposals + 1},
+                {0: proposals},
+                {0: proposals + 1},
+                None,
+                None,
+                None,
+                None,
+            ),
+            strict=True,
+        ),
+    }
+
+
 def export_dflash(
     output_dir: str,
     target_gguf: str | None = None,
@@ -606,6 +651,14 @@ def _export_dflash_cuda(
             strict=True,
         )
 
+    print("=" * 60)
+    print("Exporting DFlash sampling methods...")
+    print("=" * 60)
+    sampler_methods = _export_cuda_sampler_methods(
+        max_draft_tokens=exported_block_size - 1,
+        vocab_size=target_config.vocab_size,
+    )
+
     mutable_buffer_metadata = common.mutable_buffer_metadata(combined)
     del combined, target_model, draft_model
     if vision_model is not None:
@@ -636,6 +689,7 @@ def cuda_partitioner(method: str) -> CudaPartitioner:
         "embed_text": embed_text_ep,
         "draft_forward": draft_ep,
         "draft_prefill": draft_prefill_ep,
+        **sampler_methods,
     }
     if vision_ep is not None:
         methods["vision_encoder"] = vision_ep
@@ -677,6 +731,7 @@ def cuda_partitioner(method: str) -> CudaPartitioner:
         embed_text_ep,
         draft_ep,
         draft_prefill_ep,
+        sampler_methods,
         vision_ep,
     )
     gc.collect()

examples/models/muse-glimmer/tests/test_export_dflash.py

@@ -8,6 +8,8 @@
 
 from unittest import TestCase
 
+import torch
+
 from executorch.examples.models.muse_glimmer.export import export_dflash
 
 
@@ -49,3 +51,24 @@ def test_mlx_vision_options_are_valid(self) -> None:
     def test_callable_rejects_unknown_backend(self) -> None:
         with self.assertRaisesRegex(ValueError, "Unsupported DFlash backend: cpu"):
             export_dflash.validate_dflash_export_options("cpu")
+
+    def test_cuda_sampler_methods_accept_any_proposal_count(self) -> None:
+        methods = export_dflash._export_cuda_sampler_methods(
+            max_draft_tokens=3, vocab_size=5
+        )
+        scalars = (torch.tensor([0.0]), torch.tensor([0]), torch.tensor([1.0]))
+        for proposals in (1, 3):
+            with self.subTest(proposals=proposals):
+                tokens, probabilities = methods["dflash_sample_tokens"].module()(
+                    torch.randn(proposals, 5), *scalars
+                )
+                result = methods["dflash_verify_speculative"].module()(
+                    torch.randn(proposals + 1, 5),
+                    probabilities,
+                    torch.cat([torch.tensor([4]), tokens]),
+                    *scalars,
+                    torch.tensor([False]),
+                )
+                self.assertEqual(tokens.shape, torch.Size([proposals]))
+                self.assertEqual(result.shape, torch.Size([proposals + 3]))
+                self.assertEqual(result.dtype, torch.int64)