diff --git a/recipes/MammothModa2/MammothModa2.md b/recipes/MammothModa2/MammothModa2.md index 6dc5f12e9..02f8e30bf 100644 --- a/recipes/MammothModa2/MammothModa2.md +++ b/recipes/MammothModa2/MammothModa2.md @@ -49,6 +49,51 @@ cache — see the note under *1x L40S 48GB*. ## GPU +### Optional FP8 AR KV cache + +For CUDA deployments, `mammoth_moda2_fp8_kv.yaml` is an opt-in preset that +stores the Stage 0 AR KV cache as FP8 E4M3. 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 validated on one NVIDIA H800 80GB with CUDA and +FlashAttention 3. 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. + +1024x1024 fixed-seed smoke test completed successfully with no obvious visual +failure. FP8 is lossy, so numerical or image-quality equivalence with BF16 is +not implied. + ### 1x L40S 48GB > **48 GB config adjustment:** the committed diff --git a/tests/engine/test_stage_engine_args.py b/tests/engine/test_stage_engine_args.py index f20538284..7ac5b6b69 100644 --- a/tests/engine/test_stage_engine_args.py +++ b/tests/engine/test_stage_engine_args.py @@ -366,6 +366,44 @@ def test_typed_llm_projection_does_not_emit_inherited_upstream_defaults(): assert inherited_defaults.isdisjoint(engine_args) +def test_mammoth_fp8_kv_deploy_projects_only_ar_stage(monkeypatch): + monkeypatch.setattr(stage_init_utils, "resolve_worker_cls", lambda _engine_args: None) + + pipeline = OMNI_PIPELINES["mammoth_moda2"] + deploy = load_deploy_config(_DEPLOY_DIR / "mammoth_moda2_fp8_kv.yaml") + legacy_stages, omni_config = _legacy_and_typed_stages( + pipeline, + deploy, + model="test-model", + ) + + assert deploy.stages[0].engine_extras["kv_cache_dtype"] == "fp8_e4m3" + assert "kv_cache_dtype" not in deploy.stages[1].engine_extras + + legacy_args = [build_legacy_engine_args_dict(stage, "test-model") for stage in legacy_stages] + typed_args = [ + build_engine_args_dict_from_omni_stage_config( + omni_config.stage_by_id(stage_id), + "test-model", + ) + for stage_id in (0, 1) + ] + + ar_stage = omni_config.stage_by_id(0) + dit_stage = omni_config.stage_by_id(1) + + assert ar_stage.stage_pipeline_config.execution_type == StageExecutionType.LLM_AR + assert legacy_args[0]["kv_cache_dtype"] == "fp8_e4m3" + assert ar_stage.cache_config.cache_dtype == "fp8_e4m3" + assert "cache_dtype" in ar_stage.cache_config._omni_explicit_fields + assert typed_args[0]["kv_cache_dtype"] == "fp8_e4m3" + + assert "kv_cache_dtype" not in legacy_args[1] + assert dit_stage.cache_config.cache_dtype == "auto" + assert "cache_dtype" not in dit_stage.cache_config._omni_explicit_fields + assert "kv_cache_dtype" not in typed_args[1] + + def test_typed_llm_projection_rejects_explicit_fields_owned_by_another_boundary(): stage_config = VllmOmniARStageConfig( stage_pipeline_config=StagePipelineConfig(stage_id=0, model_stage="test"), diff --git a/vllm_omni/deploy/mammoth_moda2_fp8_kv.yaml b/vllm_omni/deploy/mammoth_moda2_fp8_kv.yaml new file mode 100644 index 000000000..8d99226b5 --- /dev/null +++ b/vllm_omni/deploy/mammoth_moda2_fp8_kv.yaml @@ -0,0 +1,7 @@ +# MammothModa2 AR→DiT text-to-image with opt-in FP8 KV cache for the AR stage. + +base_config: mammoth_moda2.yaml + +stages: + - stage_id: 0 + kv_cache_dtype: fp8_e4m3