From 110fcf1297ef7d7def43843e760b7f305cf67b83 Mon Sep 17 00:00:00 2001 From: Scott Roy Date: Wed, 12 Aug 2026 16:55:43 -0700 Subject: [PATCH] up --- .github/workflows/mlx.yml | 2 + .../test_benchmark_dflash_long_context.py | 174 ------------------ .../muse-glimmer/tests/test_cuda_pipeline.py | 22 ++- .../tests/test_ondevice_serving.py | 11 +- .../models/muse-glimmer/tests/test_serve.py | 19 +- pytest.ini | 5 + 6 files changed, 45 insertions(+), 188 deletions(-) delete mode 100644 examples/models/muse-glimmer/tests/test_benchmark_dflash_long_context.py diff --git a/.github/workflows/mlx.yml b/.github/workflows/mlx.yml index 44bc8b03bf3..4822101bdca 100644 --- a/.github/workflows/mlx.yml +++ b/.github/workflows/mlx.yml @@ -14,6 +14,7 @@ on: - extension/llm/export/** - extension/audio/** - examples/models/gemma4_31b/** + - examples/models/muse-glimmer/** - examples/models/parakeet/** - examples/models/voxtral_realtime/** - examples/models/qwen3_5_moe/** @@ -90,6 +91,7 @@ jobs: backends/mlx/test/test_slot_recycling.py \ backends/mlx/test/test_sample.py \ examples/models/gemma4_31b/tests/test_mlx_pipeline.py \ + examples/models/muse-glimmer/tests/test_mlx_pipeline.py \ -v echo "::endgroup::" diff --git a/examples/models/muse-glimmer/tests/test_benchmark_dflash_long_context.py b/examples/models/muse-glimmer/tests/test_benchmark_dflash_long_context.py deleted file mode 100644 index e7bedeed98d..00000000000 --- a/examples/models/muse-glimmer/tests/test_benchmark_dflash_long_context.py +++ /dev/null @@ -1,174 +0,0 @@ -# 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. - -from __future__ import annotations - -import os -import stat -import tempfile -import unittest -from pathlib import Path - -from executorch.examples.models.muse_glimmer.dev.benchmark_dflash_long_context import ( - BOS_TOKEN_ID, - build_runner_command, - DEFAULT_OUTPUT_TOKENS, - generate_prompt_files, - make_prompt_tokens, - MAX_SEQ_LEN, - parse_runner_metrics, - PROMPT_LENGTHS, - RANDOM_VOCAB_LIMIT, - run_case, - token_file_length, - validate_context_boundary, - validate_generated_tokens_match, -) - - -class BenchmarkDFlashLongContextTest(unittest.TestCase): - def test_prompts_have_exact_lengths_and_common_prefixes(self) -> None: - with tempfile.TemporaryDirectory() as temp_dir: - files = generate_prompt_files( - Path(temp_dir), PROMPT_LENGTHS, seed=17, output_tokens=512 - ) - - contents = {length: path.read_bytes() for length, path in files.items()} - for length, path in files.items(): - self.assertEqual(token_file_length(path), length) - self.assertEqual( - contents[length][:8], BOS_TOKEN_ID.to_bytes(8, "little") - ) - self.assertEqual( - contents[length], contents[PROMPT_LENGTHS[-1]][: length * 8] - ) - - def test_random_payload_excludes_reserved_tokens(self) -> None: - tokens = make_prompt_tokens(4096, seed=9) - self.assertEqual(tokens[0], BOS_TOKEN_ID) - self.assertTrue(all(0 <= token < RANDOM_VOCAB_LIMIT for token in tokens[1:])) - - def test_longest_case_fits_with_transient_verifier(self) -> None: - validate_context_boundary(PROMPT_LENGTHS[-1], DEFAULT_OUTPUT_TOKENS) - with self.assertRaisesRegex(ValueError, "exceeds max sequence length"): - validate_context_boundary(MAX_SEQ_LEN - 1, DEFAULT_OUTPUT_TOKENS) - - def test_runner_commands_share_exact_input_contract(self) -> None: - common = { - "model": Path("/tmp/model.pte"), - "data": Path("/tmp/model.ptd"), - "tokenizer": Path("/tmp/tokenizer"), - "prompt": Path("/tmp/prompt.i64"), - "generated": Path("/tmp/generated.i64"), - "output_tokens": 512, - } - solo = build_runner_command( - implementation="solo", runner=Path("/tmp/solo"), **common - ) - dflash = build_runner_command( - implementation="dflash", runner=Path("/tmp/dflash"), **common - ) - - for command in (solo, dflash): - self.assertIn("--cuda_graph", command) - self.assertNotIn("--cuda_graph_target", command) - self.assertNotIn("--cuda_graph_draft", command) - self.assertIn("--tokens_have_bos", command) - self.assertIn("--temperature=0", command) - self.assertIn("--max_new_tokens=512", command) - self.assertIn("--ignore_eos", command) - self.assertIn("--prompt_tokens_file=/tmp/prompt.i64", command) - - def test_run_case_creates_generated_directory_and_overwrites_old_file( - self, - ) -> None: - with tempfile.TemporaryDirectory() as temp_dir: - root = Path(temp_dir) - runner = root / "fake_runner.py" - runner.write_text( - "#!/usr/bin/env python3\n" - "from pathlib import Path\n" - "import sys\n" - "output = next(a.split('=', 1)[1] for a in sys.argv " - "if a.startswith('--generated_tokens_file='))\n" - "Path(output).write_bytes((7).to_bytes(8, 'little', signed=True))\n" - "print('Prefill: 1 tokens in 1.0 ms (1000.0 tok/s)')\n" - "print('Decode: 1 tokens in 1.0 ms (1000.0 tok/s)')\n", - encoding="utf-8", - ) - runner.chmod(runner.stat().st_mode | stat.S_IXUSR) - generated = root / "nested" / "generated.i64" - metrics = run_case( - [str(runner), f"--generated_tokens_file={generated}"], - root / "logs" / "runner.log", - generated, - expected_output_tokens=1, - environment=os.environ.copy(), - ) - self.assertEqual(generated.stat().st_size, 8) - self.assertEqual(metrics["actual_generated_tokens"], 1) - - def test_run_case_rejects_missing_metrics(self) -> None: - with tempfile.TemporaryDirectory() as temp_dir: - root = Path(temp_dir) - runner = root / "fake_runner.py" - runner.write_text( - "#!/usr/bin/env python3\n" - "from pathlib import Path\n" - "import sys\n" - "output = next(a.split('=', 1)[1] for a in sys.argv " - "if a.startswith('--generated_tokens_file='))\n" - "Path(output).write_bytes((7).to_bytes(8, 'little', signed=True))\n", - encoding="utf-8", - ) - runner.chmod(runner.stat().st_mode | stat.S_IXUSR) - generated = root / "generated.i64" - with self.assertRaisesRegex(RuntimeError, "missing required metrics"): - run_case( - [str(runner), f"--generated_tokens_file={generated}"], - root / "runner.log", - generated, - expected_output_tokens=1, - environment=os.environ.copy(), - ) - - def test_generated_tokens_must_match(self) -> None: - with tempfile.TemporaryDirectory() as temp_dir: - root = Path(temp_dir) - solo = root / "solo.i64" - dflash = root / "dflash.i64" - solo.write_bytes(b"same") - dflash.write_bytes(b"same") - validate_generated_tokens_match(solo, dflash) - - dflash.write_bytes(b"different") - with self.assertRaisesRegex(RuntimeError, "different token sequences"): - validate_generated_tokens_match(solo, dflash) - - def test_parse_runner_metrics(self) -> None: - metrics = parse_runner_metrics( - """ -Prefill: 8192 tokens in 100.0 ms (81920.0 tok/s) -Decode: 512 tokens in 2000.0 ms (256.0 tok/s) -DFlashDecodeTiming {"cycles":128,"speculative_cycles":100,"target_only_cycles":28,"draft_execute_ms":312.0,"target_execute_ms":800.0,"draft_attempts_by_row":[128,128,128],"draft_accepts_by_row":[120,100,80]} -PyTorchObserver {"decode_token_per_sec": 256.0, "inference_start_ms": 1000, "first_token_ms": 1100, "inference_end_ms": 3000} -""" - ) - self.assertEqual(metrics["prompt_tokens"], 8192) - self.assertEqual(metrics["generated_tokens"], 512) - self.assertAlmostEqual(metrics["acceptance_percent"], 78.125) - self.assertEqual(metrics["cycles"], 128) - self.assertEqual(metrics["draft_tokens"], 384) - self.assertEqual(metrics["accepted_draft_tokens"], 300) - self.assertEqual(metrics["draft_execute_ms"], 312.0) - self.assertEqual(metrics["target_execute_ms"], 800.0) - self.assertEqual(metrics["ttft_ms"], 100) - self.assertEqual(metrics["inference_end_to_end_ms"], 2000) - self.assertEqual(metrics["observer"]["decode_token_per_sec"], 256.0) - - -if __name__ == "__main__": - unittest.main() diff --git a/examples/models/muse-glimmer/tests/test_cuda_pipeline.py b/examples/models/muse-glimmer/tests/test_cuda_pipeline.py index 993a6fc51fe..ec3a9264f95 100644 --- a/examples/models/muse-glimmer/tests/test_cuda_pipeline.py +++ b/examples/models/muse-glimmer/tests/test_cuda_pipeline.py @@ -21,9 +21,22 @@ import unittest from dataclasses import replace -import executorch.backends.cuda.quantize_op_dispatch as _quantize_op_dispatch # noqa: F401 import torch -from executorch.backends.cuda.coalesced_int4_tensor import CudaCoalescedInt4Tensor + +try: + import executorch.backends.cuda.quantize_op_dispatch as _quantize_op_dispatch # noqa: F401 + from executorch.backends.cuda.coalesced_int4_tensor import CudaCoalescedInt4Tensor + from executorch.examples.models.muse_glimmer.source_transformations.cuda import ( + add_dflash_hidden_tapping, + add_on_device_sampler, + cuda_source_transformations, + ) +except ImportError as e: + # The CUDA backend pulls in triton, which is absent on non-CUDA builds. + # Every test here is CUDA-specific, so skip the module rather than fail + # collection for anyone running the suite without CUDA. + raise unittest.SkipTest(f"CUDA backend unavailable: {e}") + from executorch.examples.models.muse_glimmer.export.common import ( mutable_buffer_metadata, ) @@ -38,11 +51,6 @@ ) from executorch.examples.models.muse_glimmer.loaders.checkpoint_loader import _finalize from executorch.examples.models.muse_glimmer.model.model import FlatKVCache -from executorch.examples.models.muse_glimmer.source_transformations.cuda import ( - add_dflash_hidden_tapping, - add_on_device_sampler, - cuda_source_transformations, -) from executorch.examples.models.muse_glimmer.tests.test_pipeline import ( build_random_tiny_model, DEFAULT_RECIPE, diff --git a/examples/models/muse-glimmer/tests/test_ondevice_serving.py b/examples/models/muse-glimmer/tests/test_ondevice_serving.py index 9e27eed6444..2c192ba17f0 100644 --- a/examples/models/muse-glimmer/tests/test_ondevice_serving.py +++ b/examples/models/muse-glimmer/tests/test_ondevice_serving.py @@ -10,8 +10,15 @@ import pytest -from executorch.examples.llm_server.python.chat_template import ChatTemplate -from executorch.examples.llm_server.python.protocol import ChatMessage +# The serving stack is built on pydantic, which ships with the llm_server +# extras (examples/llm_server/python/requirements.txt) rather than core +# ExecuTorch, so skip instead of failing collection when it is absent. +pytest.importorskip("pydantic", reason="requires llm_server serving dependencies") + +from executorch.examples.llm_server.python.chat_template import ( # noqa: E402 + ChatTemplate, +) +from executorch.examples.llm_server.python.protocol import ChatMessage # noqa: E402 _SERVER = os.environ.get("MUSE_GLIMMER_SERVER_URL") _HF_DIR = os.environ.get("MUSE_GLIMMER_HF_DIR") diff --git a/examples/models/muse-glimmer/tests/test_serve.py b/examples/models/muse-glimmer/tests/test_serve.py index b7bb62913a1..69ad2a5ad22 100644 --- a/examples/models/muse-glimmer/tests/test_serve.py +++ b/examples/models/muse-glimmer/tests/test_serve.py @@ -13,11 +13,20 @@ import pytest -from executorch.examples.llm_server.python import chat_template -from executorch.examples.llm_server.python.errors import APIError -from executorch.examples.llm_server.python.protocol import ChatCompletionRequest -from executorch.examples.llm_server.python.session_runtime import PromptInput -from executorch.examples.models.muse_glimmer.serving import serve +# The serving stack is built on pydantic, which ships with the llm_server +# extras (examples/llm_server/python/requirements.txt) rather than core +# ExecuTorch, so skip instead of failing collection when it is absent. +pytest.importorskip("pydantic", reason="requires llm_server serving dependencies") + +from executorch.examples.llm_server.python import chat_template # noqa: E402 +from executorch.examples.llm_server.python.errors import APIError # noqa: E402 +from executorch.examples.llm_server.python.protocol import ( # noqa: E402 + ChatCompletionRequest, +) +from executorch.examples.llm_server.python.session_runtime import ( # noqa: E402 + PromptInput, +) +from executorch.examples.models.muse_glimmer.serving import serve # noqa: E402 _HERE = pathlib.Path(serve.__file__).resolve().parent _LLM_SERVER_PYTHON = pathlib.Path(chat_template.__file__).resolve().parent diff --git a/pytest.ini b/pytest.ini index 949d918a963..8f6e0dd3bc3 100644 --- a/pytest.ini +++ b/pytest.ini @@ -98,6 +98,11 @@ testpaths = examples/models/llama3_2_vision/text_decoder/test examples/models/llava/test + # Backend-agnostic Muse Glimmer tests. The MLX pipeline test skips itself + # when the backend is absent and is also run by mlx.yml; the CUDA one is + # ignored above. + examples/models/muse-glimmer/tests + # exir exir/