[Performance][MammothModa2] Enable and validate FP8 KV cache for AR stage (#7436)

Signed-off-by: JH <13716265422@139.com>

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

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"),

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