[DO NOT MERGE YET] DGX Spark integrations#1705
Draft
rltakashige wants to merge 46 commits intomainfrom
Draft
Conversation
rltakashige
force-pushed
the
leo/dgx-spark-integrations
branch
from
5930f76 to
d6ea684
Compare
rltakashige
force-pushed
the
leo/dgx-spark-integrations
branch
from
d6ea684 to
0c8615f
Compare
rltakashige
force-pushed
the
leo/dgx-spark-integrations
branch
from
07e49ce to
404b976
Compare
rltakashige
force-pushed
the
leo/dgx-spark-integrations
branch
6 times, most recently
from
8c8a96d to
32cc06a
Compare
rltakashige
force-pushed
the
leo/dgx-spark-integrations
branch
from
32cc06a to
f75d36c
Compare
rltakashige
force-pushed
the
leo/dgx-spark-integrations
branch
2 times, most recently
from
95f3e4e to
d5989fd
Compare
rltakashige
force-pushed
the
leo/dgx-spark-integrations
branch
2 times, most recently
from
20685ae to
7492111
Compare
Evanev7
reviewed
Mar 17, 2026
Evanev7
reviewed
Mar 17, 2026
Evanev7
reviewed
Mar 17, 2026
| def format_vllm_prompt( | ||
| engine: LLMEngine, params: TextGenerationTaskParams | ||
| ) -> tuple[list[int], str, int]: | ||
| tokenizer = TokenizerWrapper(engine.get_tokenizer()) |
Member
There was a problem hiding this comment.
we should have our own wrapper so we can just use the tokenizers Tokenizer
Collaborator
Author
There was a problem hiding this comment.
Probably a TODO for a separate PR. Added it as a comment.
Evanev7
reviewed
Mar 17, 2026
Evanev7
reviewed
Mar 17, 2026
Evanev7
reviewed
Mar 17, 2026
Evanev7
reviewed
Mar 17, 2026
Evanev7
reviewed
Mar 17, 2026
Comment on lines
+560
to
+561
| attention_backend="TRITON_ATTN", | ||
| enforce_eager=True, |
Evanev7
reviewed
Mar 17, 2026
Evanev7
reviewed
Mar 17, 2026
Evanev7
reviewed
Mar 17, 2026
Evanev7
reviewed
Mar 17, 2026
Evanev7
reviewed
Mar 17, 2026
Evanev7
reviewed
Mar 17, 2026
| @dataclass(eq=False) | ||
| class BatchGenerator(InferenceGenerator): | ||
| model: Model | ||
| model: Model | None |
Member
There was a problem hiding this comment.
the batch generator should not own the model if possible- what blockers are there to this?
Collaborator
Author
There was a problem hiding this comment.
Made an attempt at this by moving generator creation before the warmup phase.
Evanev7
reviewed
Mar 17, 2026
Evanev7
reviewed
Mar 17, 2026
Evanev7
reviewed
Mar 17, 2026
Evanev7
reviewed
Mar 17, 2026
| case LoadModel() if ( | ||
| ( | ||
| isinstance(self.current_status, RunnerConnected) | ||
| isinstance(self.generator, MlxBuilder) |
Member
There was a problem hiding this comment.
can just be isinstance(Builder)
Collaborator
Author
There was a problem hiding this comment.
I do want the generator to explicitly be an mlx builder if there is an mlx group.
Evanev7
reviewed
Mar 17, 2026
Evanev7
reviewed
Mar 17, 2026
… and vllm batch engine and don't store model on the generators.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
TODO BEFORE MERGE:
nix run exoto just work. Or at leastnix run exo-cudaPart 1 of getting EXO to run well on the DGX Spark and other Linux CUDA machines outside of private betas.
nix run .#exo-cudaor
just sync-cuda && uv run exoChanges
Cache a nix build of the dependencies
Add setup scripts
Make Linux look nice on the dashboard
Caveats
The connectivity between a Mac and Spark is still super dodgy.