Skip to content

Commit 52b176f

Browse files
digantdesaidigantdesai
authored
Add Muse Glimmer 30B ExecuTorch example (pytorch#21711)
Export and runtime support for Muse Glimmer 30B on the CUDA and MLX backends, covering text, vision, DFlash speculative decoding, and OpenAI-compatible serving. Co-authored by @Gasoonjia, @metascroy, @digantdesai, @mergennachin
1 parent fb5eedc commit 52b176f

67 files changed

Lines changed: 21435 additions & 5 deletions

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎.gitignore‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,16 @@ build-direct/
2424
dist/
2525
arm-scratch/
2626
executorch.egg-info
27+
28+
# MLX delegate code generated in-tree from backends/mlx/serialization/schema.fbs
29+
# by generate.py, which the MLX backend runs at build time.
30+
backends/mlx/_generated_inspector.py
31+
backends/mlx/runtime/MLXLoader.cpp
32+
backends/mlx/runtime/MLXLoader.h
33+
backends/mlx/runtime/schema_generated.h
34+
backends/mlx/serialization/_generated/
35+
backends/mlx/serialization/_generated_serializers.py
36+
backends/mlx/serialization/mlx_graph_schema.py
2737
pip-out/
2838
build-profiling/
2939
**/ddr_*_temp

‎.lintrunner.toml‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -390,6 +390,9 @@ exclude_patterns = [
390390
'extension/llm/tokenizers',
391391
'extension/llm/tokenizers/**',
392392
'examples/llm_server',
393+
# muse_glimmer is a symlink to this directory; the linter resolves it and
394+
# then fails on a directory.
395+
'examples/models/muse-glimmer',
393396
'backends/cadence/utils/FACTO',
394397
'examples/cuda',
395398
'examples/qualcomm',

‎backends/cuda/runtime/cuda_mutable_state.cpp‎

Lines changed: 11 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -642,10 +642,17 @@ Error mutable_state_rebind_for_execute(CudaDelegateHandle* handle) {
642642
return Error::Ok;
643643
}
644644
auto cit = m.contexts.find(hit->second);
645-
if (cit != m.contexts.end() &&
646-
(!cit->second.sessions.empty() ||
647-
cit->second.rebound_handles.find(handle) !=
648-
cit->second.rebound_handles.end())) {
645+
if (cit != m.contexts.end()) {
646+
const auto desc = cit->second.desc.find(handle);
647+
const bool known_stateless =
648+
desc != cit->second.desc.end() && desc->second.empty();
649+
const bool requires_active_session =
650+
(!known_stateless && !cit->second.sessions.empty()) ||
651+
cit->second.rebound_handles.find(handle) !=
652+
cit->second.rebound_handles.end();
653+
if (!requires_active_session) {
654+
return Error::Ok;
655+
}
649656
ET_LOG(
650657
Error, "mutable_state: active session is required for this handle");
651658
return Error::InvalidState;

‎backends/cuda/runtime/test/test_cuda_mutable_state.cpp‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -371,11 +371,17 @@ TEST(CudaMutableStateTest, CapturesClonesAndRebindsDeviceBuffer) {
371371
container.fqns = {"model.state"};
372372
container.extracted["model.state"] =
373373
reinterpret_cast<aoti::AtenTensorHandle>(source_tensor.get());
374+
FakeContainer stateless_container;
374375
cu::CudaDelegateHandle handle = fake_container_handle(&container);
376+
cu::CudaDelegateHandle stateless_handle =
377+
fake_container_handle(&stateless_container);
375378

376379
cu::MutableStateContextOwner c;
377380
c.register_fqns({"model.state"});
378-
c.with_load_scope([&] { cu::mutable_state_note_handle(&handle); });
381+
c.with_load_scope([&] {
382+
cu::mutable_state_note_handle(&handle);
383+
cu::mutable_state_note_handle(&stateless_handle);
384+
});
379385

380386
ASSERT_TRUE(c.available());
381387
EXPECT_EQ(c.bytes_per_session(), 4 * sizeof(float));
@@ -387,6 +393,7 @@ TEST(CudaMutableStateTest, CapturesClonesAndRebindsDeviceBuffer) {
387393

388394
auto token = c.create_session();
389395
ASSERT_TRUE(token.ok());
396+
EXPECT_EQ(cu::mutable_state_rebind_for_execute(&stateless_handle), Error::Ok);
390397
EXPECT_EQ(cu::mutable_state_rebind_for_execute(&handle), Error::InvalidState);
391398

392399
ASSERT_EQ(

‎examples/models/muse-glimmer/.gitignore‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
__pycache__/
2+
*.pyc

‎examples/models/muse-glimmer/CMakeLists.txt‎

Lines changed: 188 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,188 @@
1+
# Copyright (c) Meta Platforms, Inc. and affiliates.
2+
# All rights reserved.
3+
#
4+
# This source code is licensed under the BSD-style license found in the
5+
# LICENSE file in the root directory of this source tree.
6+
7+
cmake_minimum_required(VERSION 3.24)
8+
project(muse_glimmer)
9+
10+
set(CMAKE_CXX_STANDARD 17)
11+
set(CMAKE_CXX_STANDARD_REQUIRED ON)
12+
13+
if(DEFINED ENV{CONDA_PREFIX} AND EXISTS "$ENV{CONDA_PREFIX}/lib")
14+
list(PREPEND CMAKE_BUILD_RPATH "$ENV{CONDA_PREFIX}/lib")
15+
endif()
16+
17+
set(EXECUTORCH_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/../../..)
18+
19+
include(${EXECUTORCH_ROOT}/tools/cmake/Utils.cmake)
20+
21+
set(_common_include_directories ${EXECUTORCH_ROOT}/..)
22+
set(_json_include
23+
${EXECUTORCH_ROOT}/extension/llm/tokenizers/third-party/json/single_include
24+
)
25+
26+
# gflags
27+
set(gflags_DIR ${EXECUTORCH_ROOT}/cmake-out/third-party/gflags)
28+
find_package(gflags REQUIRED)
29+
30+
# executorch
31+
list(APPEND CMAKE_FIND_ROOT_PATH ${CMAKE_CURRENT_BINARY_DIR}/../../../..)
32+
find_package(executorch CONFIG REQUIRED FIND_ROOT_PATH_BOTH)
33+
executorch_target_link_options_shared_lib(executorch)
34+
35+
set(link_libraries executorch gflags)
36+
37+
# CPU ops (for host-side helpers not delegated to CUDA)
38+
list(APPEND link_libraries optimized_native_cpu_ops_lib cpublas eigen_blas)
39+
executorch_target_link_options_shared_lib(optimized_native_cpu_ops_lib)
40+
41+
# Extensions
42+
list(
43+
APPEND
44+
link_libraries
45+
extension_llm_runner
46+
extension_module
47+
extension_data_loader
48+
extension_tensor
49+
extension_flat_tensor
50+
)
51+
52+
# Backend: CUDA or MLX (exactly one required)
53+
if(EXECUTORCH_BUILD_CUDA)
54+
find_package(CUDAToolkit REQUIRED)
55+
list(APPEND link_libraries aoti_cuda_backend)
56+
executorch_target_link_options_shared_lib(aoti_cuda_backend)
57+
add_compile_definitions(EXECUTORCH_BUILD_CUDA)
58+
elseif(TARGET mlxdelegate)
59+
list(APPEND link_libraries mlxdelegate mlx)
60+
executorch_target_link_options_shared_lib(mlxdelegate)
61+
add_compile_definitions(EXECUTORCH_BUILD_MLX)
62+
else()
63+
message(FATAL_ERROR "Set EXECUTORCH_BUILD_CUDA=ON or EXECUTORCH_BUILD_MLX=ON")
64+
endif()
65+
66+
# Use a supplied or local stb checkout before fetching it. The resize header is
67+
# under deprecated/ in the pinned revision.
68+
if(EXECUTORCH_BUILD_CUDA OR TARGET mlxdelegate)
69+
if(NOT MUSE_GLIMMER_STB_INCLUDE_DIR
70+
AND EXISTS "${EXECUTORCH_ROOT}/third-party/stb/stb_image.h"
71+
)
72+
set(MUSE_GLIMMER_STB_INCLUDE_DIR "${EXECUTORCH_ROOT}/third-party/stb")
73+
endif()
74+
if(MUSE_GLIMMER_STB_INCLUDE_DIR
75+
AND EXISTS "${MUSE_GLIMMER_STB_INCLUDE_DIR}/stb_image.h"
76+
)
77+
message(
78+
STATUS "Using local stb headers from: ${MUSE_GLIMMER_STB_INCLUDE_DIR}"
79+
)
80+
list(APPEND _common_include_directories ${MUSE_GLIMMER_STB_INCLUDE_DIR}
81+
${MUSE_GLIMMER_STB_INCLUDE_DIR}/deprecated
82+
)
83+
else()
84+
include(FetchContent)
85+
FetchContent_Declare(
86+
stb
87+
GIT_REPOSITORY https://github.com/nothings/stb.git
88+
GIT_TAG f0569113c93ad095470c54bf34a17b36646bbbb5
89+
)
90+
FetchContent_MakeAvailable(stb)
91+
list(APPEND _common_include_directories ${stb_SOURCE_DIR}
92+
${stb_SOURCE_DIR}/deprecated
93+
)
94+
endif()
95+
endif()
96+
97+
# Tokenizer (tiktoken via pytorch/tokenizers)
98+
list(APPEND link_libraries tokenizers::tokenizers)
99+
100+
add_executable(
101+
solo_runner runtime/runners/solo.cpp runtime/engine/muse_glimmer_engine.cpp
102+
runtime/engine/dflash_session.cpp
103+
)
104+
if(EXECUTORCH_BUILD_CUDA OR TARGET mlxdelegate)
105+
target_sources(
106+
solo_runner PRIVATE runtime/engine/muse_glimmer_vision_runtime.cpp
107+
)
108+
endif()
109+
target_include_directories(
110+
solo_runner PUBLIC ${_common_include_directories} ${_json_include}
111+
)
112+
target_link_libraries(solo_runner PUBLIC ${link_libraries})
113+
114+
# Export symbols required by the dynamically loaded AOTI CUDA blob.
115+
if(EXECUTORCH_BUILD_CUDA
116+
AND NOT APPLE
117+
AND NOT MSVC
118+
)
119+
target_link_options(solo_runner PRIVATE "LINKER:--export-dynamic")
120+
endif()
121+
122+
add_executable(
123+
dflash_runner
124+
runtime/runners/dflash.cpp runtime/engine/muse_glimmer_engine.cpp
125+
runtime/engine/dflash_session.cpp
126+
)
127+
if(EXECUTORCH_BUILD_CUDA OR TARGET mlxdelegate)
128+
target_sources(
129+
dflash_runner PRIVATE runtime/engine/muse_glimmer_vision_runtime.cpp
130+
)
131+
endif()
132+
target_include_directories(
133+
dflash_runner PUBLIC ${_common_include_directories} ${_json_include}
134+
)
135+
target_link_libraries(dflash_runner PUBLIC ${link_libraries})
136+
137+
add_executable(sampling_test tests/sampling_test.cpp)
138+
target_include_directories(sampling_test PUBLIC ${_common_include_directories})
139+
target_link_libraries(sampling_test PUBLIC extension_tensor executorch)
140+
141+
add_executable(dflash_timing_test tests/dflash_timing_test.cpp)
142+
target_include_directories(
143+
dflash_timing_test PUBLIC ${_common_include_directories}
144+
)
145+
146+
# Worker used by the OpenAI-compatible serving harness.
147+
add_executable(
148+
muse_glimmer_worker
149+
runtime/runners/muse_glimmer_worker.cpp
150+
runtime/engine/muse_glimmer_engine.cpp runtime/engine/dflash_session.cpp
151+
)
152+
if(EXECUTORCH_BUILD_CUDA OR TARGET mlxdelegate)
153+
target_sources(
154+
muse_glimmer_worker PRIVATE runtime/engine/muse_glimmer_vision_runtime.cpp
155+
)
156+
endif()
157+
target_include_directories(
158+
muse_glimmer_worker PUBLIC ${_common_include_directories} ${_json_include}
159+
)
160+
target_link_libraries(muse_glimmer_worker PUBLIC ${link_libraries})
161+
162+
if(EXECUTORCH_BUILD_CUDA
163+
AND NOT APPLE
164+
AND NOT MSVC
165+
)
166+
target_link_options(dflash_runner PRIVATE "LINKER:--export-dynamic")
167+
target_link_options(muse_glimmer_worker PRIVATE "LINKER:--export-dynamic")
168+
endif()
169+
170+
if(NOT CMAKE_BUILD_TYPE STREQUAL "Debug")
171+
target_link_options_gc_sections(solo_runner)
172+
target_link_options_gc_sections(dflash_runner)
173+
target_link_options_gc_sections(muse_glimmer_worker)
174+
if(NOT APPLE AND NOT MSVC)
175+
# CUDA AOTI blobs resolve symbols from the host executable.
176+
if(NOT EXECUTORCH_BUILD_CUDA)
177+
target_link_options(solo_runner PRIVATE "LINKER:-s")
178+
target_link_options(dflash_runner PRIVATE "LINKER:-s")
179+
target_link_options(muse_glimmer_worker PRIVATE "LINKER:-s")
180+
endif()
181+
endif()
182+
endif()
183+
184+
if(TARGET mlxdelegate)
185+
executorch_target_copy_mlx_metallib(solo_runner)
186+
executorch_target_copy_mlx_metallib(dflash_runner)
187+
executorch_target_copy_mlx_metallib(muse_glimmer_worker)
188+
endif()

‎examples/models/muse-glimmer/CMakePresets.json‎

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,83 @@
1+
{
2+
"version": 6,
3+
"configurePresets": [
4+
{
5+
"name": "muse-glimmer-base",
6+
"hidden": true,
7+
"binaryDir": "${sourceDir}/../../../cmake-out/examples/models/muse-glimmer",
8+
"cacheVariables": {
9+
"CMAKE_BUILD_TYPE": "Release",
10+
"CMAKE_FIND_ROOT_PATH": "${sourceDir}/../../../cmake-out",
11+
"CMAKE_PREFIX_PATH": "${sourceDir}/../../../cmake-out"
12+
}
13+
},
14+
{
15+
"name": "muse-glimmer-cuda",
16+
"displayName": "Muse Glimmer runner (CUDA)",
17+
"inherits": ["muse-glimmer-base"],
18+
"cacheVariables": {
19+
"EXECUTORCH_BUILD_CUDA": "ON"
20+
},
21+
"condition": {
22+
"type": "inList",
23+
"string": "${hostSystemName}",
24+
"list": ["Linux", "Windows"]
25+
}
26+
},
27+
{
28+
"name": "muse-glimmer-mlx",
29+
"displayName": "Muse Glimmer runner (MLX)",
30+
"inherits": ["muse-glimmer-base"],
31+
"cacheVariables": {},
32+
"condition": {
33+
"type": "equals",
34+
"lhs": "${hostSystemName}",
35+
"rhs": "Darwin"
36+
}
37+
}
38+
],
39+
"buildPresets": [
40+
{
41+
"name": "muse-glimmer-cuda",
42+
"displayName": "Build Muse Glimmer runner (CUDA)",
43+
"configurePreset": "muse-glimmer-cuda",
44+
"targets": ["solo_runner", "dflash_runner", "muse_glimmer_worker"]
45+
},
46+
{
47+
"name": "muse-glimmer-mlx",
48+
"displayName": "Build Muse Glimmer runner (MLX)",
49+
"configurePreset": "muse-glimmer-mlx",
50+
"targets": ["solo_runner", "dflash_runner", "muse_glimmer_worker"]
51+
}
52+
],
53+
"workflowPresets": [
54+
{
55+
"name": "muse-glimmer-cuda",
56+
"displayName": "Configure and build Muse Glimmer runner (CUDA)",
57+
"steps": [
58+
{
59+
"type": "configure",
60+
"name": "muse-glimmer-cuda"
61+
},
62+
{
63+
"type": "build",
64+
"name": "muse-glimmer-cuda"
65+
}
66+
]
67+
},
68+
{
69+
"name": "muse-glimmer-mlx",
70+
"displayName": "Configure and build Muse Glimmer runner (MLX)",
71+
"steps": [
72+
{
73+
"type": "configure",
74+
"name": "muse-glimmer-mlx"
75+
},
76+
{
77+
"type": "build",
78+
"name": "muse-glimmer-mlx"
79+
}
80+
]
81+
}
82+
]
83+
}

0 commit comments

Comments
 (0)