vllm-omni/recipes/MammothModa2/MammothModa2.md

# MammothModa2

> MammothModa2-Preview and MammothModa2-Dev unified understanding and generation

## Summary

- Vendor: ByteDance Research
- Models: `bytedance-research/MammothModa2-Preview`, `bytedance-research/MammothModa2-Dev`
- Tasks: Preview and Dev text-to-image (AR → DiT); Dev text/image understanding
- Mode: Offline inference
- Maintainer: Community

## When to use this recipe

Use this recipe to run MammothModa2-Preview through shared task-oriented
offline examples. Text-to-image uses the shared image example (`text_to_image.py`) instead of a model-specific script.
The generic example formats the AR prompt, drives the AR → DiT stage pipeline,
and forwards MammothModa2-specific generation parameters through the
pipeline-declared `extra_body` contract.

MammothModa2's DiT stage runs in the shared diffusion runtime. The default
`mammoth_moda2.yaml` uses request-level batching with up to eight compatible
requests (`max_num_seqs: 8`) and one image per request
(`num_outputs_per_prompt: 1`). TeaCache and Cache-DiT are optional alternative
cache backends in request mode. Enable TeaCache with `--cache-backend tea_cache`
as shown below, or enable Cache-DiT using the commented settings in that YAML.
Compilation, DiT quantization, parallelism, and offload are not enabled by
these presets.

To use step execution and continuous batching, select
[`mammoth_moda2_step.yaml`](../../vllm_omni/deploy/mammoth_moda2_step.yaml)
with `--deploy-config vllm_omni/deploy/mammoth_moda2_step.yaml` in the
text-to-image commands below. This preset inherits the same stage placement
and capacity, enables step execution for Stage 1, and disables request-batch
admission waiting and diffusion cache acceleration. Compatible requests can
join between denoising steps and finish independently, including requests
with different inference-step counts. Step mode cannot be combined with
TeaCache, Cache-DiT, or other diffusion cache backends.

Image size, seed, guidance, and denoising steps use the standard diffusion
request fields. `cfg_range` remains a MammothModa2-specific `extra_body`
parameter. For compatibility, the runtime also accepts the former
`text_guidance_scale` and `num_inference_steps` keys in `extra_body`; when
present and non-null, those keys take precedence over the standard fields.

## References

- Upstream model:
  [`bytedance-research/MammothModa2-Preview`](https://huggingface.co/bytedance-research/MammothModa2-Preview)
- Dev model:
  [`bytedance-research/MammothModa2-Dev`](https://huggingface.co/bytedance-research/MammothModa2-Dev)
- Related offline example:
  [`examples/offline_inference/text_to_image/text_to_image.py`](../../examples/offline_inference/text_to_image/text_to_image.py)
- Related T2T/I2T example:
  [`examples/offline_inference/x_to_text/x_to_text.py`](../../examples/offline_inference/x_to_text/x_to_text.py)
- Declared parameters:
  [`vllm_omni/model_extras/mammothmodal2_preview.py`](../../vllm_omni/model_extras/mammothmodal2_preview.py)
- Deploy config:
  [`vllm_omni/deploy/mammoth_moda2.yaml`](../../vllm_omni/deploy/mammoth_moda2.yaml)

## Hardware Support

The default deploy config places both the AR and DiT stages on one GPU
(`devices: "0"`). Its committed `gpu_memory_utilization` split is 0.5 for
stage 0 and 0.3 for stage 1. The A800 validation section below also shows a
two-GPU placement with one stage per GPU for attributable timing and memory;
the measured results are summarized below.

## GPU

### Optional FP8 AR KV cache

For CUDA deployments, `mammoth_moda2_fp8_kv.yaml` is an opt-in preset that
keeps the Stage 0 AR KV cache of decoder layer 0 in BF16 and stores the other
27 layers as FP8 E4M3 (`kv_cache_dtype_skip_layers: ["0"]`; write the layer
indices as quoted strings). Stage 1 remains on
`kv_cache_dtype=auto`; its DiT execution is unaffected. This setting quantizes
only the autoregressive KV cache. It is neither FP8 weight/activation
quantization nor vLLM-Omni diffusion KV-cache quantization.

Use the preset in place of the default deploy config:

```bash
python examples/offline_inference/text_to_image/text_to_image.py \
  --model ./MammothModa2-Preview \
  --deploy-config vllm_omni/deploy/mammoth_moda2_fp8_kv.yaml \
  --prompt "A stylish woman riding a motorcycle in NYC, movie poster style" \
  --height 1024 \
  --width 1024 \
  --seed 42 \
  --extra-body '{"text_guidance_scale": 4.0, "cfg_range": [0.0, 1.0], "num_inference_steps": 50}' \
  --output mammoth_t2i.png
```

The preset was first validated on one NVIDIA H800 80GB with CUDA and
FlashAttention 3, with every layer in FP8. The native and FP8 runs used the
same model, code revision, and downstream configuration.

| Metric | Native BF16 (`kv_cache_dtype=auto`) | FP8 E4M3 (`fp8_e4m3`) | Change |
| --- | ---: | ---: | ---: |
| KV cache memory | 15.59 GiB | 15.55 GiB | -0.04 GiB |
| GPU KV cache size | 145,904 tokens | 291,232 tokens | +99.6% |
| Maximum concurrency at 8,192 tokens | 17.81x | 35.55x | +99.6% |
| Steady-state AR median | 71.191 s | 79.816 s | +12.1% |
| Steady-state end-to-end median | 83.840 s | 92.473 s | +10.3% |
| Steady-state DiT median | 12.574 s | 12.561 s | effectively unchanged |

The FP8 run used `kv_cache_dtype=fp8_e4m3` only for Stage 0; Stage 1 used
`kv_cache_dtype=auto`. The nearly unchanged reserved cache memory holds almost
twice as many tokens because FP8 reduces the bytes per cached token. This is a
capacity/concurrency tradeoff: the measured AR and end-to-end latencies were
higher than the native-BF16 baseline.

On H800, a 1024x1024 fixed-seed smoke test with every layer in FP8 completed
successfully with no obvious visual failure. FP8 is lossy, so numerical or
image-quality equivalence with BF16 is not implied.

On one A800 80GB (vLLM 0.30.0), FP8 KV layers run on FlashInfer and BF16 layers
on FlashAttention 2. Text-to-image at 1024x1024, 50 steps, guidance 4.0, seeds 42
and 1-5, with a studio tabby cat prompt and a peephole-view Samoyed prompt;
an image counts when it shows the prompted subject and scene.

| Stage 0 KV cache | GPU KV cache size | Cat images that follow the prompt | Samoyed images that follow the prompt |
| --- | ---: | ---: | ---: |
| BF16 on all 28 layers (`auto`) | 147,408 tokens | 6/6 | 6/6 |
| FP8 E4M3 on all 28 layers | 294,816 tokens | 1/6 | 0/6 |
| Layer 0 BF16, other 27 layers FP8 E4M3 (this preset) | 284,640 tokens | 6/6 | 6/6 |

The KV cache sizes above still count the 28 attention layers of the replaced
Qwen-VL language model, which #8095 removes: with it, BF16 goes from 147,408 to
294,816 tokens and all-FP8 from 294,816 to 589,632. The prompt-following
columns do not depend on it.

With all 28 layers in FP8, most cat images become a framed print on a wall.
Keeping layer 0 in BF16 restores them; keeping only layer 27, which has the
largest key magnitude, does not. The H800 numbers above were measured with all 28
layers in FP8; this preset has been run on A800 only.

### 1x L40S 48GB

> **48 GB config adjustment:** the committed
> `vllm_omni/deploy/mammoth_moda2.yaml` uses
> `gpu_memory_utilization` 0.5 / 0.3 (sized for ~80 GB). To fit on a 48 GB L40S,
> set the stage-0 (AR) value to `0.8` and the stage-1 (DiT) value to `0.16`
> before running. (On an ~80 GB GPU, leave the defaults unchanged.)

### 1x NVIDIA A800 80GB

#### Environment

- OS: Linux
- Python: Match the repository requirements for your checkout
- Driver / runtime: NVIDIA CUDA environment with one A800 80 GB
- vLLM version: Match the repository requirements for your checkout
- vLLM-Omni version or commit: Use the commit you are deploying from

#### Offline Commands

Download the model:

```bash
hf download bytedance-research/MammothModa2-Preview --local-dir ./MammothModa2-Preview
```

Run text-to-image with the shared offline example from the repository root. The
deploy config sets `trust_remote_code`, so no extra flag is needed:

```bash
python examples/offline_inference/text_to_image/text_to_image.py \
  --model ./MammothModa2-Preview \
  --deploy-config vllm_omni/deploy/mammoth_moda2.yaml \
  --prompt "A stylish woman riding a motorcycle in NYC, movie poster style" \
  --height 1024 \
  --width 1024 \
  --seed 42 \
  --guidance-scale 4.0 \
  --num-inference-steps 50 \
  --extra-body '{"cfg_range": [0.0, 1.0]}' \
  --output mammoth_t2i.png
```

The standard diffusion request fields are `height`, `width`, `seed`,
`guidance_scale`, and `num_inference_steps`; use their corresponding CLI flags
shown above. `--height` and `--width` must be multiples of 16.

`cfg_range` is the only recommended MammothModa2 field in `--extra-body`; it
sets the relative step range `[start, end]` over which CFG is applied (default
`[0.0, 1.0]`). For compatibility, `text_guidance_scale` and
`num_inference_steps` remain accepted `extra_body` aliases and, when non-null,
take precedence over the standard request fields. Model extras are filtered
against the declared `extra_body_params`.

TeaCache can be enabled for the DiT stage with the same user-facing sampling
parameters:

```bash
python examples/offline_inference/text_to_image/text_to_image.py \
  --model ./MammothModa2-Preview \
  --deploy-config vllm_omni/deploy/mammoth_moda2.yaml \
  --prompt "A stylish woman riding a motorcycle in NYC, movie poster style" \
  --height 1024 \
  --width 1024 \
  --guidance-scale 4.0 \
  --num-inference-steps 50 \
  --cache-backend tea_cache \
  --extra-body '{"cfg_range": [0.0, 1.0]}' \
  --output mammoth_t2i_teacache.png
```

The bundled TeaCache coefficients were fitted from MammothModa2 full-compute
traces. MammothModa2 uses the model-specific default `rel_l1_thresh=0.075`,
selected for the evaluated 1024x1024, 50-step configuration.

The model-specific keys are declared in
[`vllm_omni/model_extras/mammothmodal2_preview.py`](../../vllm_omni/model_extras/mammothmodal2_preview.py)),
so unknown MammothModa2 extras may be dropped.

Run text-to-text through the shared understanding example. It recognizes the
MammothModa2 checkpoint and automatically selects `mammoth_moda2_ar.yaml`:

```bash
python examples/offline_inference/x_to_text/x_to_text.py \
  --model ./MammothModa2-Preview \
  --prompt "Explain multimodal generation in three sentences."
```

Add an image for image-to-text or image summarization. The shared example
uses MammothModa2's chat and vision-token template:

```bash
python examples/offline_inference/x_to_text/x_to_text.py \
  --model ./MammothModa2-Preview \
  --image /path/to/input.jpg \
  --prompt "Please summarize the content of this image."
```

#### Verification

The example writes the generated image to the `--output` path. Confirm the file
exists and is a valid image:

```bash
ls -lh mammoth_t2i.png
python -c "from PIL import Image; print(Image.open('mammoth_t2i.png').size)"
```

### 2x NVIDIA A800 80GB validation

Use one A800 per stage so AR and DiT memory and timing are attributable. The
per-stage override changes placement only; both stages remain single-rank.

```bash
VLLM_LOGGING_LEVEL=DEBUG vllm serve ./MammothModa2-Preview --omni \
  --deploy-config vllm_omni/deploy/mammoth_moda2.yaml \
  --stage-overrides '{"0":{"devices":"0"},"1":{"devices":"1"}}' \
  --port 8099 \
  --log-stats
```

Startup logs should identify stage 1 as `StageDiffusionClient` and resolve it
to `MammothModa2DiTPipeline`. `DiffusionEngine` step timing is a DEBUG-level,
per-request message, so it appears only after sending a text-to-image request
with `VLLM_LOGGING_LEVEL=DEBUG`; it is not a startup marker. Seeing the legacy
generation model runner for stage 1 is a failed migration.

#### Migration benchmark

The request-mode migration was checked on 2x NVIDIA A800 80GB PCIe with AR on
GPU 0 and DiT on GPU 1. Each revision ran one warmup followed by 10 serial
measured requests in the same initialized process. Both used BF16 eager mode,
1024x1024 output, 50 denoising steps, guidance scale 4.0, seed 42, and no
diffusion cache. The baseline was the pre-migration revision `caed3061`; the
candidate was `19de562a`. Lower latency is better.

| Metric | Baseline p50 | Baseline p95 | Candidate p50 | Candidate p95 |
| --- | ---: | ---: | ---: | ---: |
| End-to-end latency | 105.36 s | 106.01 s | 104.94 s | 105.79 s |
| AR stage latency | 86.97 s | 87.61 s | 86.26 s | 87.11 s |
| DiT stage latency | 18.27 s | 18.41 s | 18.60 s | 18.62 s |

Peak sampled device memory was 39,209 MiB on the AR GPU for both revisions.
The DiT GPU used 11,089 MiB for the baseline and 10,967 MiB for the candidate.
The candidate's shared runtime reported 372.02 ms p50 per denoising step and a
5.89 ms p50 AR-to-diffusion adapter time. All measured requests completed and
both revisions produced valid, prompt-aligned 1024x1024 RGB images. The small
latency differences are regression evidence, not a statistically significant
speedup claim.

### 1x AMD MI300X, MammothModa2 Preview (pre-migration baseline)

#### Environment

- OS: Linux 6.8.0-134-generic, x86_64
- Container: official ROCm image built from `docker/Dockerfile.rocm`
- Python: 3.12.13
- PyTorch: 2.11.0+gitd0c8b1f
- Driver / runtime: AMD 6.19.14.31400000 / ROCm 7.2.53211
- GPU: one AMD Instinct MI300X, `gfx942:sramecc+:xnack-`, 191.69 GiB visible HBM
- vLLM version: 0.27.0+rocm723
- vLLM Omni version or commit: `73e1368c7bb940efe1a025859c9d6c8eeeb2e3f0`
- Installed vLLM Omni package metadata: `0.27.0rc2.dev44+g55abdade9.rocm`

#### Offline Commands

The checked run used the committed stage split, with `gpu_memory_utilization` set to 0.5 for AR and 0.3 for DiT:

```bash
python3 examples/offline_inference/text_to_image/text_to_image.py \
    --model bytedance-research/MammothModa2-Preview \
    --deploy-config vllm_omni/deploy/mammoth_moda2.yaml \
    --prompt "A stylish woman riding a motorcycle in NYC, movie poster style" \
    --height 1024 \
    --width 1024 \
    --seed 42 \
    --extra-body '{"text_guidance_scale": 4.0, "cfg_range": [0.0, 1.0], "num_inference_steps": 50}' \
    --enable-diffusion-pipeline-profiler \
    --log-stats \
    --output mammoth_t2i.png
```

#### Verification

The first request took 85.224 seconds. The AR stage generated 4,161 visual tokens in 72.996 seconds, and the DiT stage took 12.163 seconds. AR weight loading used 21.4 GiB and took 8.250 seconds. DiT weight loading used 5.49 GiB and took 1.824 seconds. The largest one second whole device memory sample was 106.57 GiB, including the AR KV cache reserved by the 0.5 memory setting.

### 1x A800 80GB, MammothModa2 Preview (serving)

#### Environment

- OS: Linux, x86_64
- CPU: 18 cores
- GPU: one NVIDIA A800-SXM4-80GB
- Deploy config: `vllm_omni/deploy/mammoth_moda2.yaml` (default split: AR `gpu_memory_utilization` 0.5, DiT 0.3)
- vLLM-Omni version or commit: use the commit you are deploying from

#### Serving Commands

```bash
vllm-omni serve bytedance-research/MammothModa2-Preview --omni \
    --port 8091 \
    --trust-remote-code

python benchmarks/diffusion/diffusion_benchmark_serving.py \
    --model bytedance-research/MammothModa2-Preview \
    --endpoint /v1/images/generations \
    --host 127.0.0.1 --port 8091 \
    --height 1024 --width 1024 --num-inference-steps 50 \
    --extra-body '{"text_guidance_scale": 9.0, "cfg_range": [0.0, 1.0]}' \
    --num-prompts 8 --seed 142 --warmup-requests 0 \
    --output-file mm2_1024_s50_c1.json
```

#### Verification

On 1024x1024 / 50 steps, single-concurrency end-to-end mean latency is ~96 s
(P99 ~101 s) per image and concurrency-4 mean is ~151 s (~0.024 img/s);
steady-state combined GPU memory is ~50.4 GiB (AR ~39.2 GiB, DiT ~11.1 GiB).
The AR stage dominates (~77 s of a ~96 s request) because it decodes a fixed
4,161-token visual grid per image, so latency is insensitive to prompt length
and scales only partially with DiT step count. See the
[serving performance dashboard](../../benchmarks/diffusion/performance_dashboard/mammoth_moda2_serving_performance.md)
for the full sweep, peak-memory, and component-attribution data.

The output was a valid 1024 by 1024 RGB PNG.

### Startup benchmark

The [startup benchmark](../../benchmarks/mammoth_moda2/README.md) provides
commands, measurement boundaries and a single-A800 eager baseline. On that
16-CPU-quota machine, parallel stage initialization and 16 CPU threads reduced
`Omni()` initialization from about 60 s to 37 s. See the benchmark for the
configuration and limitations; CUDA-graph mode needs separate memory and
quality validation.

## MammothModa2-Dev unified inference

MammothModa2-Dev uses a Qwen3-VL AR backbone, while MammothModa2-Preview uses
Qwen2.5-VL. vLLM-Omni selects the matching implementation from the nested
`llm_config.model_type`; no checkpoint edits or `trust_remote_code` flag are
required.

Text-to-text and image-to-text use the AR-only deploy. Text-to-image loads the
Qwen3 generation experts (`gen_mlp`), extra visual vocabulary and image head,
then sends the generated visual tokens and hidden states to the DiT stage.

Download the checkpoint:

```bash
hf download bytedance-research/MammothModa2-Dev --local-dir ./MammothModa2-Dev
```

Run text-to-text through the shared understanding example. It recognizes the
Dev checkpoint as MammothModa2 and automatically selects
`mammoth_moda2_ar.yaml`:

```bash
python examples/offline_inference/x_to_text/x_to_text.py \
  --model ./MammothModa2-Dev \
  --prompt "Explain multimodal generation in three sentences."
```

Add an image for image-to-text or image summarization:

```bash
python examples/offline_inference/x_to_text/x_to_text.py \
  --model ./MammothModa2-Dev \
  --image ./image.png \
  --prompt "Please summarize the content of this image."
```

The Dev checkpoint is approximately 47.55 GiB on disk. In the verified AR-only
run, loaded model weights used approximately 16.97 GiB of GPU memory before KV
and encoder caches. Allow additional GPU memory for those caches and the input
image.