Skip to content

test(sglang): add FLUX.2 diffusion serving test on SageMaker#6367

Open
Jyothirmaikottu wants to merge 5 commits into
mainfrom
sglang-flux2-diffusion-tests
Open

test(sglang): add FLUX.2 diffusion serving test on SageMaker#6367
Jyothirmaikottu wants to merge 5 commits into
mainfrom
sglang-flux2-diffusion-tests

Conversation

@Jyothirmaikottu

@Jyothirmaikottu Jyothirmaikottu commented Jul 9, 2026

Copy link
Copy Markdown
Contributor

Adds a real SGLang serving test for FLUX.2 diffusion on SageMaker. The SGLang DLC image ships the sglang.multimodal_gen diffusion stack but the SageMaker entrypoint hardcoded the LLM engine, so diffusion pipelines could not be served.

Changes

  • scripts/docker/sglang/sagemaker_entrypoint.sh: new SM_SGLANG_ENGINE gate — default llm (existing behavior unchanged); diffusion launches sglang.multimodal_gen.runtime.launch_server. The control var is excluded from the SM_SGLANG_*--flag mapping.
  • test/sglang/sagemaker/test_sm_endpoint.py: deploy FLUX.2-klein-4B (the ungated, servable FLUX.2 pipeline) with the diffusion engine and assert /v1/images/generations returns image data.

Testing: FLUX.2-klein-4B verified serving end-to-end on an L40S via sglang.multimodal_gen — pipeline recognized (Qwen3 text encoder → Flux2Transformer2DModelAutoencoderKLFlux2), /health ok, /v1/images/generations returned a generated image. Backward-compatible — existing LLM SageMaker images are unaffected (engine defaults to llm).

Add FLUX.2 model inference testing for the SGLang DLC image, which
already ships the sglang.multimodal_gen diffusion stack (diffusers, etc.)
but previously had no diffusion test coverage or SageMaker serving path.

- model-tests: validate FLUX.2-small-decoder (AutoencoderKLFlux2) via a
  VAE decode test, since the distilled decoder is a pipeline component
  rather than a servable text-to-image model.
- sagemaker_entrypoint.sh: add an SM_SGLANG_ENGINE gate (default "llm",
  unchanged) so SM_SGLANG_ENGINE=diffusion launches
  sglang.multimodal_gen.runtime.launch_server instead of the LLM engine.
- sagemaker test: deploy FLUX.2-klein-4B via the diffusion engine and
  assert /v1/images/generations returns image data.
Jyothirmaikottu and others added 2 commits July 9, 2026 10:32
Revert the SageMaker diffusion endpoint test and the SM_SGLANG_ENGINE
entrypoint gate. Those served FLUX.2-klein-4B (the full pipeline), not
FLUX.2-small-decoder, which is what this change is about. The distilled
VAE decoder is a pipeline component and is validated by the model-tests
VAE decode test; no entrypoint change is needed for it.
@Jyothirmaikottu Jyothirmaikottu changed the title test(sglang): add FLUX.2 diffusion inference tests test(sglang): add FLUX.2-small-decoder VAE decode test Jul 9, 2026
Replace the diffusers-only VAE decode test with a real SGLang serving
test. The decode test loaded AutoencoderKLFlux2 via diffusers and never
exercised SGLang, so it did not validate SGLang's diffusion serving path.

- sagemaker_entrypoint.sh: add an SM_SGLANG_ENGINE gate (default "llm",
  unchanged) so SM_SGLANG_ENGINE=diffusion launches
  sglang.multimodal_gen.runtime.launch_server instead of the LLM engine.
- sagemaker test: deploy FLUX.2-klein-4B (the ungated, servable FLUX.2
  pipeline) via the diffusion engine and assert /v1/images/generations
  returns image data.
- drop the sglang_flux_vae_decode_test.sh script and its model-tests entry.
@Jyothirmaikottu Jyothirmaikottu changed the title test(sglang): add FLUX.2-small-decoder VAE decode test test(sglang): add FLUX.2 diffusion serving test on SageMaker Jul 9, 2026
Pre-existing go/stdlib CVE compiled into mooncake's libetcd_wrapper.so,
same class as the other mooncake go/stdlib entries already allowlisted.
Not patchable without an upstream mooncake rebuild on Go 1.26.5+. Added
to both the sglang (ec2) and sglang_server (sagemaker) allowlists.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant