diff --git a/CLAUDE.md b/CLAUDE.md index 2f14ce4aa..f9e753aab 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -14,10 +14,10 @@ black . && ruff check . # format + lint (CI enforced) ```bash # OpenAI-compatible serving -python -m atom.entrypoints.openai_server --model --kv_cache_dtype fp8 -tp 8 +python -m atom.entrypoints.openai_server --model --kv-cache-dtype fp8 -tp 8 # Offline inference -python -m atom.examples.simple_inference --model --kv_cache_dtype fp8 +python -m atom.examples.simple_inference --model --kv-cache-dtype fp8 ``` - Accuracy validation: see `/ci-pr-guide` for `lm_eval` setup and CI thresholds diff --git a/README.md b/README.md index e0e81a864..1becf7b1d 100644 --- a/README.md +++ b/README.md @@ -126,7 +126,7 @@ hf auth login The default optimization level is 3 (piecewise torch.compile with CUDA graphs). ```bash -python -m atom.examples.simple_inference --model meta-llama/Meta-Llama-3-8B --kv_cache_dtype fp8 +python -m atom.examples.simple_inference --model meta-llama/Meta-Llama-3-8B --kv-cache-dtype fp8 ``` > **Note:** First-time execution may take approximately 10 minutes for model compilation. @@ -137,13 +137,13 @@ Start an OpenAI-compatible server: ```bash # Single GPU -python -m atom.entrypoints.openai_server --model Qwen/Qwen3-0.6B --kv_cache_dtype fp8 +python -m atom.entrypoints.openai_server --model Qwen/Qwen3-0.6B --kv-cache-dtype fp8 # Multi-GPU with tensor parallelism -python -m atom.entrypoints.openai_server --model deepseek-ai/DeepSeek-R1 --kv_cache_dtype fp8 -tp 8 +python -m atom.entrypoints.openai_server --model deepseek-ai/DeepSeek-R1 --kv-cache-dtype fp8 -tp 8 # With MTP speculative decoding -python -m atom.entrypoints.openai_server --model deepseek-ai/DeepSeek-R1 --kv_cache_dtype fp8 -tp 8 \ +python -m atom.entrypoints.openai_server --model deepseek-ai/DeepSeek-R1 --kv-cache-dtype fp8 -tp 8 \ --method mtp --num-speculative-tokens 3 ``` @@ -191,7 +191,7 @@ Launch the server with `--torch-profiler-dir` and `--mark-trace`: ```bash python -m atom.entrypoints.openai_server \ - --model deepseek-ai/DeepSeek-R1 --kv_cache_dtype fp8 -tp 8 \ + --model deepseek-ai/DeepSeek-R1 --kv-cache-dtype fp8 -tp 8 \ --torch-profiler-dir ./trace --mark-trace ``` diff --git a/atom/model_engine/arg_utils.py b/atom/model_engine/arg_utils.py index 2ec68f4cb..34eeda836 100644 --- a/atom/model_engine/arg_utils.py +++ b/atom/model_engine/arg_utils.py @@ -115,7 +115,9 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: help="Engine internal port", ) parser.add_argument( + "--kv-cache-dtype", "--kv_cache_dtype", + dest="kv_cache_dtype", choices=["bf16", "fp8"], type=str, default="bf16", diff --git a/tests/test_arg_utils_spec.py b/tests/test_arg_utils_spec.py index 9f859dec3..473ec2442 100644 --- a/tests/test_arg_utils_spec.py +++ b/tests/test_arg_utils_spec.py @@ -31,9 +31,29 @@ {"NONE": 0, "PIECEWISE": 1, "FULL": 2, "FULL_AND_PIECEWISE": 3}, ) +import argparse # noqa: E402 + from atom.model_engine.arg_utils import EngineArgs # noqa: E402 +class TestKVCacheDtypeCliAlias: + """--kv-cache-dtype and --kv_cache_dtype must both set kv_cache_dtype.""" + + def _parse(self, argv): + parser = argparse.ArgumentParser() + EngineArgs.add_cli_args(parser) + return parser.parse_args(argv) + + def test_dashed_form(self): + assert self._parse(["--kv-cache-dtype", "fp8"]).kv_cache_dtype == "fp8" + + def test_underscore_form(self): + assert self._parse(["--kv_cache_dtype", "fp8"]).kv_cache_dtype == "fp8" + + def test_default(self): + assert self._parse([]).kv_cache_dtype == "bf16" + + class TestEngineArgsSpeculativeValidation: """Regression tests for speculative-config construction in _get_engine_kwargs."""