[Bugfix][MammothModa2] Keep AR layer 0 in BF16 in the FP8 KV preset (#8083)

Signed-off-by: MrlixiangWE <mrdanaer@gmail.com>
Co-authored-by: zijianc2 <157244773+zijianc2@users.noreply.github.com>
Co-authored-by: hsliuustc0106 <222337142+hsliuustc0106@users.noreply.github.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

recipes/MammothModa2/MammothModa2.md

@@ -59,7 +59,9 @@ the measured results are summarized below.
 ### 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
+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.
@@ -78,9 +80,9 @@ python examples/offline_inference/text_to_image/text_to_image.py \
   --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.
+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 |
 | --- | ---: | ---: | ---: |
@@ -97,9 +99,30 @@ 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.
+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
 

tests/engine/test_stage_engine_args.py

@@ -20,6 +20,7 @@ from vllm.config import ParallelConfig as VllmParallelConfig
 from vllm.config import ProfilerConfig as VllmProfilerConfig
 from vllm.config import SchedulerConfig as VllmSchedulerConfig
 from vllm.engine.arg_utils import EngineArgs
+from vllm.model_executor.models.utils import extract_layer_index
 
 from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
 from vllm_omni.config.omni_config import (
@@ -404,7 +405,9 @@ def test_mammoth_fp8_kv_deploy_projects_only_ar_stage(monkeypatch):
     )
 
     assert deploy.stages[0].engine_extras["kv_cache_dtype"] == "fp8_e4m3"
+    assert deploy.stages[0].engine_extras["kv_cache_dtype_skip_layers"] == ["0"]
     assert "kv_cache_dtype" not in deploy.stages[1].engine_extras
+    assert "kv_cache_dtype_skip_layers" not in deploy.stages[1].engine_extras
 
     legacy_args = [build_legacy_engine_args_dict(stage, "test-model") for stage in legacy_stages]
     typed_args = [
@@ -420,14 +423,22 @@ def test_mammoth_fp8_kv_deploy_projects_only_ar_stage(monkeypatch):
 
     assert ar_stage.stage_pipeline_config.execution_type == StageExecutionType.LLM_AR
     assert legacy_args[0]["kv_cache_dtype"] == "fp8_e4m3"
+    assert legacy_args[0]["kv_cache_dtype_skip_layers"] == ["0"]
     assert ar_stage.cache_config.cache_dtype == "fp8_e4m3"
+    assert ar_stage.cache_config.kv_cache_dtype_skip_layers == ["0"]
     assert "cache_dtype" in ar_stage.cache_config._omni_explicit_fields
     assert typed_args[0]["kv_cache_dtype"] == "fp8_e4m3"
+    assert typed_args[0]["kv_cache_dtype_skip_layers"] == ["0"]
+    # vLLM matches the list against the index it parses from each attention prefix.
+    assert str(extract_layer_index("ar.language_model.layers.0.self_attn.attn")) == "0"
 
     assert "kv_cache_dtype" not in legacy_args[1]
+    assert "kv_cache_dtype_skip_layers" not in legacy_args[1]
     assert dit_stage.cache_config.cache_dtype == "auto"
+    assert dit_stage.cache_config.kv_cache_dtype_skip_layers == []
     assert "cache_dtype" not in dit_stage.cache_config._omni_explicit_fields
     assert "kv_cache_dtype" not in typed_args[1]
+    assert "kv_cache_dtype_skip_layers" not in typed_args[1]
 
 
 def test_typed_llm_projection_rejects_explicit_fields_owned_by_another_boundary():

vllm_omni/deploy/mammoth_moda2_fp8_kv.yaml

@@ -5,3 +5,6 @@ base_config: mammoth_moda2.yaml
 stages:
   - stage_id: 0
     kv_cache_dtype: fp8_e4m3
+    # Layer 0 stays BF16: with every layer in FP8, most text-to-image samples
+    # lose the prompt (A800 results in recipes/MammothModa2/MammothModa2.md).
+    kv_cache_dtype_skip_layers: ["0"]