diff --git a/.claude/commands/atom-vllm-benchmark-guide.md b/.claude/commands/atom-vllm-benchmark-guide.md new file mode 100644 index 0000000000..a28562f4b8 --- /dev/null +++ b/.claude/commands/atom-vllm-benchmark-guide.md @@ -0,0 +1,283 @@ +# ATOM vLLM Benchmark Guide + +Use this guide when the user asks to benchmark ATOM vLLM Plugin performance, compare with the upstream vLLM (optional), or validate a performance claim. + +## What this guide should do + +1. Confirm benchmark goal and variables (throughput, TTFT, TPOT, E2EL, etc.). +2. Use plugin-on as default setup; optionally benchmark upstream vLLM for performance comparison. +3. Run repeatable measurements and summarize median results. +4. Report reproducible commands, key metrics, and known risks. +5. Prefer model-specific launch settings from `recipes/atom_vllm/` when available; otherwise use the standard workflow in this guide. + +## Inputs To Collect First + +Before running any benchmark, confirm: + +- Model path (HF id or local path) +- Serving backend (`vllm`) +- Hardware shape (GPU type/count) +- `tensor_parallel_size`, KV cache dtype, and scheduler flags +- Request mix (ISL/OSL/concurrency/request rate) +- Comparison target: + - default candidate (plugin on) + - optional baseline/control (plugin off) + - optional runtime baseline (upstream vLLM nightly) + +If any of these are missing, ask before execution. + +## Environment Checklist + +For each benchmark round: + +- Kill old server processes and verify VRAM is released. +- Clear cache when startup behavior is inconsistent: + - `rm -rf /root/.cache/atom/*` + - `rm -rf /root/.cache/vllm/*` +- Set stable logging: + - `export AITER_LOG_LEVEL=WARNING` +- Confirm server is actually ready: + - `curl -sf http://localhost:8000/v1/models` + - `rocm-smi --showmemuse` and verify allocated memory is greater than 0. + +## Benchmark Workflow (Plugin On by Default) + +Run plugin-on by default. + +### Recipe-first execution policy + +When running ATOM vLLM benchmark: + +1. First check whether the target model has a matching recipe under `recipes/atom_vllm/`. +2. If a matching recipe exists, use its model-specific environment variables and `vllm serve` launch args as the default benchmark setup. +3. If no matching recipe exists, follow the standard launch template and workflow in this guide. + +Keep benchmark methodology unchanged (smoke check, concurrency isolation, and fixed comparison variables) regardless of whether the launch config comes from a recipe or from the default template. + +- **Default (candidate):** plugin enabled +- **Optional control (baseline):** plugin disabled (only when A/B comparison is requested) +- **Default ATOM vLLM Plugin image:** `rocm/atom-dev:vllm-latest` + +If both conditions are run, only one variable is allowed to change between them. + +### 1) Launch server + +Use the same launch command template. Keep plugin-on as default, and only change plugin toggle env vars when running the optional plugin-off control. + +```bash +vllm serve \ + --host localhost \ + --port 8000 \ + --tensor-parallel-size \ + --attention-backend ROCM_AITER_FA \ + --trust-remote-code \ + --kv-cache-dtype fp8 \ + --gpu_memory_utilization 0.9 \ + --async-scheduling \ + --compilation-config '{"cudagraph_mode": "FULL_AND_PIECEWISE"}' \ + --no-enable-prefix-caching +``` + +Common plugin toggles: + +- Default plugin-on: do not set `ATOM_DISABLE_VLLM_PLUGIN` or `ATOM_DISABLE_VLLM_PLUGIN_ATTENTION` (or set them to `0`) +- Optional plugin-off control: set `ATOM_DISABLE_VLLM_PLUGIN=1` +- Optional attention-off control: set `ATOM_DISABLE_VLLM_PLUGIN_ATTENTION=1` +- Always set `VLLM_ROCM_USE_AITER=1` in both ATOM plugin and upstream vLLM runs. + +### 2) Run smoke validation + +```bash +curl -sf http://localhost:8000/v1/models +vllm bench serve \ + --backend vllm \ + --base-url http://localhost:8000 \ + --endpoint /v1/completions \ + --model \ + --dataset-name random \ + --random-input-len 32 \ + --random-output-len 32 \ + --temperature 0.0 \ + --max-concurrency \ + --num-prompts \ + --num-warmups \ + --request-rate 1 \ + --disable-tqdm +``` + +Use `smoke_conc_x10 = smoke_conc * 10` and `smoke_conc_x2 = smoke_conc * 2`. + +If smoke test fails, do not continue with performance comparison. + +### 3) Run benchmark workload + +```bash +vllm bench serve \ + --backend vllm \ + --base-url http://localhost:8000 \ + --endpoint /v1/completions \ + --model= \ + --dataset-name=random \ + --random-input-len= \ + --random-output-len= \ + --random-range-ratio=0.8 \ + --temperature=0.0 \ + --num-prompts= \ + --num-warmups= \ + --max-concurrency= \ + --trust-remote-code \ + --request-rate=inf \ + --ignore-eos \ + --disable-tqdm \ + --save-result \ + --percentile-metrics="ttft,tpot,itl,e2el" +``` + +### 4) Concurrency sweep isolation (required) + +When sweeping multiple concurrency points, each point must run in a fresh container lifecycle: + +1. Start a fresh container. +2. Start vLLM server in that container. +3. Run smoke + benchmark for one `--max-concurrency=`. +4. Stop vLLM server. +5. Exit and stop/remove the container. +6. Start a new container for the next concurrency point. + +Do not reuse the same running server/container across different concurrency points. + +Example pattern: + +```bash +for conc in 4 8 16 32 64 128; do + docker run -dt --name "atom_bench_c${conc}" + docker exec "atom_bench_c${conc}" bash -lc 'vllm serve ...' + docker exec "atom_bench_c${conc}" bash -lc "vllm bench serve ... --max-concurrency=${conc} ..." + docker exec "atom_bench_c${conc}" bash -lc 'pkill -f "vllm serve" || true' + docker stop "atom_bench_c${conc}" && docker rm "atom_bench_c${conc}" +done +``` + +### 5) Minimum benchmark matrix + +| Scenario | ISL | OSL | Concurrency | +|---|---:|---:|---:| +| Short in / short out | 1024 | 1024 | 4, 8, 16, 32, 64, 128 | +| Short in / long out | 1024 | 8192 | 4, 8, 16, 32, 64, 128 | +| Long in / short out | 8192 | 1024 | 4, 8, 16, 32, 64, 128 | + +### 6) Repetition and statistics + +- Run each scenario at least 1 times. +- Use median values for comparison. +- If variance is large, increase repeat count. + +## Optional Upstream vLLM Benchmark + +Use this optional path when you need ATOM vs upstream runtime comparison. + +Reference images: + +- ATOM vLLM plugin default image: `rocm/atom-dev:vllm-lastest` +- Upstream vLLM comparison image: `vllm/vllm-openai-rocm:nightly` + +### 1) Pull latest upstream ROCm image + +```bash +docker pull vllm/vllm-openai-rocm:nightly +``` + +### 2) Run the same workflow and config + +For upstream runs, keep everything identical to ATOM runs: + +- Same model path +- Same GPU set and `tensor_parallel_size` +- Same KV cache dtype and scheduler flags +- Same request mix (ISL/OSL/concurrency/request rate) +- Same concurrency sweep isolation policy (restart container per concurrency point) + +Only runtime image is allowed to change. + +### 3) Upstream execution note + +- Use the same smoke and benchmark commands from this guide. +- Do not set ATOM plugin toggle env vars for upstream runs. +- Always set `VLLM_ROCM_USE_AITER=1` for upstream runs. +- Keep result naming explicit (for example, include `upstream-nightly` in filenames). + +## Regression Profiling Workflow + +When regression is observed: + +```bash +# 1) Start (or restart) server with profiler enabled +vllm serve \ + --host localhost \ + --port 8000 \ + --tensor-parallel-size \ + --attention-backend ROCM_AITER_FA \ + --trust-remote-code \ + --kv-cache-dtype fp8 \ + --profiler-config '' + +# 2) Run profiled benchmark workload +vllm bench serve \ + --backend vllm \ + --base-url http://localhost:8000 \ + --endpoint /v1/completions \ + --model \ + --dataset-name random \ + --random-input-len \ + --random-output-len 20 \ + --temperature 0.0 \ + --max-concurrency \ + --num-prompts \ + --num-warmups \ + --request-rate inf \ + --profile \ + --save-result +``` + +- `--profile` requires server-side `--profiler-config`. +- Keep profiler config identical between plugin-off and plugin-on runs. +- For reproducibility, keep benchmark client semantics fixed to completion mode: + `--backend vllm --endpoint /v1/completions` (this path uses `max_tokens`). + +## Reporting Template + +Use this format in final response: + +```markdown +### ATOM vLLM Benchmark Result +- Goal: +- Setup: + - Model: + - Hardware: + - Fixed args: + - Isolation: +- Validation: +- Benchmark: + - Default (plugin on): + - Optional control (plugin off): + - Optional runtime baseline (upstream nightly): + - Delta: +- Confidence: + - Runs: + - Variance: +- Risk: + - +- Repro: + - + - + - +``` + +## Guardrails + +- Never compare runs with different model, TP, KV dtype, or request mix. +- For concurrency sweep, stop server and restart container between concurrency points. +- For ATOM vs upstream comparison, only runtime image may differ. +- Keep `VLLM_ROCM_USE_AITER=1` fixed for both ATOM and upstream runs. +- If GPU memory is not allocated, treat benchmark results as invalid. +- If plugin-on and plugin-off outputs differ semantically, run accuracy checks before claiming performance gains. diff --git a/.claude/commands/ci-pr-guide.md b/.claude/commands/ci-pr-guide.md index cc8226e91c..2916b94a3c 100644 --- a/.claude/commands/ci-pr-guide.md +++ b/.claude/commands/ci-pr-guide.md @@ -37,8 +37,8 @@ PR/Push → pre-checks (black + ruff) | DeepSeek-R1-0528 MTP | `--kv_cache_dtype fp8 -tp 8 --method mtp` | 0.94 | Yes | 8gpu | | DeepSeek-R1-0528 MXFP4 | `--kv_cache_dtype fp8 -tp 8` | 0.93 | Yes | 8gpu | | DeepSeek-R1-0528 MXFP4 MTP | `--kv_cache_dtype fp8 -tp 8 --method mtp` | 0.93 | Yes | 8gpu | -| gpt-oss-120b | `--kv_cache_dtype fp8 --gpu-memory-utilization 0.3` | 0.38 | Yes | mi355-1 | -| gpt-oss-120b (2 GPU) | `-tp 2 --enable-dp-attention --enable-expert-parallel` | 0.38 | No | mi355-4 | +| gpt-oss-120b | `--kv_cache_dtype fp8 --gpu-memory-utilization 0.5` | 0.88 | Yes | mi355-1 | +| gpt-oss-120b (2 GPU) | `-tp 2 --enable-dp-attention --enable-expert-parallel` | 0.88 | No | mi355-4 | | Qwen3-235B FP8 | `-tp 8 --enable-expert-parallel` + `ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1` | 0.87 | Yes | 8gpu | | Qwen3-235B MXFP4 | `-tp 8 --enable-expert-parallel` + `ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1` | 0.87 | No | 8gpu | | Qwen3-Next-80B-A3B | `-tp 8` | 0.65 | Yes | 8gpu | diff --git a/.claude/commands/inferencex-sync.md b/.claude/commands/inferencex-sync.md new file mode 100644 index 0000000000..5a82eb1956 --- /dev/null +++ b/.claude/commands/inferencex-sync.md @@ -0,0 +1,84 @@ +--- +name: inferencex-sync +description: Compare the performance of the latest successful run from ATOM upstream benchmark(https://github.com/ROCm/ATOM/actions/workflows/atom-benchmark.yaml) and that of InferenceX(https://inferencex.semianalysis.com/api/v1/benchmarks?model=$MODEL) and report tput_per_gpu (Total token throughput / number of GPUs), report performance regression of each model on atom framework, create a PR if the performance of ATOM upstream benchamrk is better than that of InferenceX by updating docker image and atom serve arguments (python3 -m atom.entrypoints.openai_server) +memory: project +model: opus +--- + +# Compare ATOM upstream benchmark and InferneceX, and create a PR to InferenceX Instructions + +Use this skill to compare the througput of ATOM upstream benchamrk and InferenceX and create a PR to InferenceX with the udpated docker image and atom serve arguments. + +## Workflow + +```text +- [ ] 1) Create a throughput table from InferenceX benchmark +- [ ] 2) Create a throughput table from ATOM upstream benchmark +- [ ] 3) Chek performance regression of ATOM upstream benchmark from InferenceX benchmark +- [ ] 4) Ask users what model needs a new PR +``` + +## 1) Create a throughput table from InferenceX benchmark + +Check the InferenceX benchmark that is from these URLs and create a throughput performance table from ATOM upstream benchmark. +- https://inferencex.semianalysis.com/api/v1/benchmarks?model=Kimi-K2.5 +- https://inferencex.semianalysis.com/api/v1/benchmarks?model=DeepSeek-V4-Pro +- https://inferencex.semianalysis.com/api/v1/benchmarks?model=DeepSeek-R1-0528 +- https://inferencex.semianalysis.com/api/v1/benchmarks?model=Qwen-3.5-397B-A17B +- https://inferencex.semianalysis.com/api/v1/benchmarks?model=GLM-5 +- https://inferencex.semianalysis.com/api/v1/benchmarks?model=MiniMax-M2.5 +- https://inferencex.semianalysis.com/api/v1/benchmarks?model=gpt-oss-120b + +Rules: +- 1. Use curl to fetch data from URLs +- 2. Use only "hardware":"mi355x","framework":"atom","disagg":false,"is_multinode":false,"spec_method":"none" +- 3. If muliple data of "hardware":"mi355x","framework":"atom","precision":"fp4","disagg":false,"is_multinode":false exist for a same "model", use only the latest "date" +- 4. Create a throughput performance table of each "tput_per_gpu" of each "model","isl", "osl","conc","num_decode_gpus"(number of GPUs),"image","precision" + +## 2) Create a throughput table from ATOM upstream benchmark + +Check the ATOM upstream benchmark that is from the latest successful gh action runs at https://github.com/ROCm/ATOM/actions/workflows/atom-benchmark.yaml, which as all jobs passed. Append the throughput to the InferenceX performance table in Step 1) + +Rules: +- 1. Use main branch. Find the latest job of https://github.com/ROCm/ATOM/actions/workflows/atom-benchmark.yaml where all jobs passed +- 2. "tput_per_gpu" (tput/GPU (tok/s)) is "Total token throughput" divided by TP(number of GPUs). TP(number of GPUs) is the value followed by -tp in https://github.com/ROCm/ATOM/blob/main/.github/benchmark/models.json +- 3. "conc" is "Max Concurrency". Compare same "conc" or "Max Concurrency" between InferenceX benchmark and Atom benchmark +- 4. Model name follows this mapping + + ```python + # inferenceX_model -> (atom_model, precision) + MODEL_MAP = { + "dsv4": [("DeepSeek-V4-Pro", "fp4")], + "dsr1": [("DeepSeek-R1-0528", "fp8"), + ("DeepSeek-R1-0528-MXFP4", "fp4")], + "kimik2.5": [("Kimi-K2.5-MXFP4", "fp4")], + "qwen3.5": [("Qwen3.5-397B-A17B-FP8", "fp8"), + ("Qwen3.5-397B-A17B-MXFP4", "fp4")], + "glm5": [("GLM-5-FP8", "fp8"), + ("GLM-5.1-MXFP4", "fp4")], + "minimaxm2.5": [("MiniMax-M2.7", "fp8"), + ("MiniMax-M2.7-MXFP4", "fp4")], + "gptoss120b": [("gpt-oss-120b", "fp4")], + } + ``` + +If the data is not available, stop and ask for any workaround + +## 3) Chek performance regression of ATOM upstream benchmark from InferenceX benchmark + +Update the performance table to show regression in % of ATOM upstream benchmark from InferenceX benchmark + +## 4) Chek performance regression of ATOM upstream benchmark from InferenceX benchmark + +Based on Step 3) results, ask user what models need a PR to https://github.com/SemiAnalysisAI/InferenceX/pulls and create a PR to each model. + +Rules: +- 1. Find the docker image from the step 2) ATOM upstream benchmark gh action runs. for eample, "Docker: + rocm/atom-dev:nightly_202605141645" from https://github.com/ROCm/ATOM/actions/runs/25875375446/job/76041129934#step:6:26 +- 2. Create a PR and check https://github.com/ROCm/ATOM/blob/main/.github/benchmark/models.json options + +PR example: +- 1. https://github.com/SemiAnalysisAI/InferenceX/pull/1311/changes +- 2. https://github.com/SemiAnalysisAI/InferenceX/blob/main/benchmarks/single_node/dsv4_fp4_mi355x_atom.sh +- 3. https://github.com/SemiAnalysisAI/InferenceX/blob/main/.github/configs/amd-master.yaml#L1646 +- 4. https://github.com/SemiAnalysisAI/InferenceX/blob/main/perf-changelog.yaml#L1824-L1833 diff --git a/.claude/skills/atom-patterns.md b/.claude/skills/atom-patterns/SKILL.md similarity index 100% rename from .claude/skills/atom-patterns.md rename to .claude/skills/atom-patterns/SKILL.md diff --git a/.claude/skills/capture-trace/SKILL.md b/.claude/skills/capture-trace/SKILL.md new file mode 100644 index 0000000000..adff103055 --- /dev/null +++ b/.claude/skills/capture-trace/SKILL.md @@ -0,0 +1,208 @@ +--- +name: capture-trace +description: Capture a PyTorch profiler / kineto trace from a running ATOM server for a short benchmark window. Use when the user asks for "a trace", "profiler trace", "GPU trace", or "抓 trace" for performance investigation — what kernels ran, what's on the critical path, what's slow. Do NOT use for crashes (use debug-agent-locate-kernel) or numerical bugs (use dump-bisect-debug). +version: 1.1.0 +scope: ATOM on AMD ROCm (PyTorch kineto profiler, per-rank `*.pt.trace.json.gz`) +last_updated: 2026-05-17 +--- + +## When to use + +- User asks for a trace, profiler dump, or kineto dump +- Performance analysis: "what's eating the time", "is decode fused", "did kernel X get called", "is this kernel on the critical path" +- Verifying that a code path was actually exercised at runtime (search by kernel name) + +Do NOT use this skill for: + +- Crashes / `Memory access fault` → [[debug-agent-locate-kernel]] +- Wrong outputs / accuracy regression → [[dump-bisect-debug]] +- Suspicion of a hang (no progress) → `scripts/wait_infer_drain.sh` first + +## Critical pre-flight + +1. **Stop the existing server cleanly** — the profiler argument has to be on the launch command line. `start_atom_server.sh` auto-kills the prior atom workers, so just relaunching with the new args is enough. +2. **Pick a SHORT workload** — a trace from a long run is unreadable and OOMs the profiler exporter. Default to `CONC * 1` requests (one prompt per concurrent slot). Never use the production `PROMPT_MULTIPLIER=10` default for a profiling run. +3. **`ATOM_PROFILER_MORE` belongs on the server, not the benchmark client.** The profiler runs inside the model-runner worker processes; an env on the bench client does nothing. +4. **Trace dir must be empty** for a clean per-rank layout. `start_atom_server.sh` does NOT clear it — `rm -rf $TRACE_DIR` before relaunch if you're iterating. + +## Required tools + +```bash +ls /app/ATOM/scripts/start_atom_server.sh # launcher +ls /app/ATOM/scripts/run_benchmark.sh # bench driver (passes --profile when PROFILE=1) +ls /app/ATOM/scripts/wait_server_ready.sh # ready-poll +python3 -c "import torch.profiler" # kineto present +``` + +## Parameters + +Pull these out of the user's request; everything except `MODEL` has a sensible default. + +| Param | Meaning | Typical | +|---|---|---| +| `MODEL` | Model path under `/data/` | `/data/DeepSeek-V4-Pro` | +| `TP` | Tensor-parallel size | `8` (4 for Kimi, 1 for gpt-oss-120b) | +| `ISL` / `OSL` | Random input / output length | `1024 / 1024` | +| `CONC` | Concurrency the bench keeps in flight | `64` or `128` | +| `PROMPT_MULTIPLIER` | Total prompts = `CONC * this` | **`1` for trace runs** (override the script default of 10) | +| `ATOM_PROFILER_MORE` | `1` = shapes + stack + memory (large traces, OOM risk); `0` = kernel-name only | **`0`** unless asked | +| `TRACE_DIR` | Where the kineto `.pt.trace.json.gz` lands | `/app/logs_claude/traces/` | +| `EXTRA_ARGS` | Forwarded to the openai server (MTP, kv-cache, etc.) | See [[atom-patterns]] | + +## Workflow + +### Step 1: Launch the server with the profiler bound + +```bash +TRACE_DIR=/app/logs_claude/traces/ +mkdir -p "$TRACE_DIR" + +# ATOM_PROFILER_MORE on the server env — not the client. +ATOM_PROFILER_MORE=0 \ + bash /app/ATOM/scripts/start_atom_server.sh \ + "$MODEL" "$TP" 8000 \ + --torch-profiler-dir "$TRACE_DIR" \ + $EXTRA_ARGS +``` + +`start_atom_server.sh` blocks until either `Server is ready!` or `Server process died`. Check the tail line; if it died, no point profiling. + +### Step 2: Drive a SHORT bench with `PROFILE=1` + +```bash +bash /app/ATOM/scripts/run_benchmark.sh \ + "$MODEL" 8000 "$ISL" "$OSL" "$CONC" \ + 1 \ # PROMPT_MULTIPLIER — keep this at 1 for traces + 1 \ # PROFILE=1 flag + $BENCH_EXTRA_ARGS +``` + +Position 6 is `PROMPT_MULTIPLIER`; position 7 is `PROFILE`. The bench sends a `start` HTTP call before the run and a `stop` call after, which is what trips the kineto exporter on the server. + +### Step 3: Wait for the exporter to finish (asynchronous on the server) + +Kineto exports lazily on the worker side — `stop_profiler` returns immediately to the bench, but the per-rank `.json` write + gzip can take 10-60 seconds. Poll the output dir: + +```bash +for i in $(seq 1 60); do + GZ=$(find "$TRACE_DIR" -name "*.pt.trace.json.gz" | wc -l) + JSON=$(find "$TRACE_DIR" -name "*.pt.trace.json" -not -name "*.gz" | wc -l) + echo "[t=${i}0s] gz=$GZ json=$JSON" + # Done = expected gz count AND no orphan .json (the .json is deleted after gzip) + [ "$GZ" -ge "$TP" ] && [ "$JSON" -eq 0 ] && break + sleep 10 +done +``` + +The completion signal is **`.gz` exists AND the same-name `.json` is gone**. File size of the `.gz` alone is unreliable (per `feedback_trace_gz_truncated.md`) — the exporter writes the raw `.json`, then gzip + unlink, so an orphan `.json` means it crashed mid-export. + +### Step 4: Verify the layout + +```bash +find "$TRACE_DIR" -type f | xargs ls -la +``` + +Expected: + +- `/rank_0/`, `rank_1/`, …, `rank_/` — one dir per rank +- Each dir has exactly one `*.pt.trace.json.gz` +- `ATOM_PROFILER_MORE=0`: ~50-80 MB per rank +- `ATOM_PROFILER_MORE=1`: ~200-300 MB per rank + +### Step 5: Inspect + +For a quick "did kernel X run" check: + +```bash +zcat "$TRACE_DIR"/rank_0/*.gz | python3 -c " +import json, sys +events = json.load(sys.stdin)['traceEvents'] +names = {e.get('name','') for e in events} +for kw in ['', ...]: + hits = sorted(n for n in names if kw in n) + print(f'{kw}: {len(hits)} matches') + for h in hits[:5]: print(f' {h}') +" +``` + +For counts and aggregate time per kernel: + +```bash +zcat "$TRACE_DIR"/rank_0/*.gz | python3 -c " +import json, sys +events = json.load(sys.stdin)['traceEvents'] +def stat(kw): + m = [e for e in events if kw in e.get('name','') and 'dur' in e] + if m: print(f'{kw}: count={len(m)} total_us={sum(e[\"dur\"] for e in m)}') +stat('aiter::topk_softplus') +stat('aiter::moe_forward') +# ... +" +``` + +For a UI view, drop the `.gz` (decompressed `.json`) into or `chrome://tracing`. + +## `record_function` tag format + +ATOM annotates the critical path with `torch.profiler.record_function`. The tags follow a stable format used by `parse_trace.py`: + +- Prefill: `prefill[bs= tok= ctx=]` +- Decode: `decode[bs= tok= d= spec=]` +- Draft (eagle mid-step): `draft[i/k bs=]` + +Searching by these tags is far more reliable than searching by kernel name (which varies across PyTorch/Triton/AITER versions). + +## `ATOM_PROFILER_MORE` cost + +`ATOM_PROFILER_MORE=1` enables `record_shapes + with_stack + profile_memory`. This multiplies trace size ~3-4x and, more importantly, the **C++ kineto aggregation at `stop_profiler` time scales with the recorded event count × TP**. On 8-rank V4-Pro, a 60-second profile with `PROFILER_MORE=1` has been observed to OOM the worker processes during export. Rules of thumb: + +- Default `=0` (kernel names + durations only — enough for 95% of investigations) +- `=1` only when you specifically need shapes or Python stacks AND you've kept the window short (≤ 5 seconds of bench traffic, ≤ `CONC * 1` prompts) + +## Common model configs + +| Model | Path | TP | Server `EXTRA_ARGS` | Bench notes | +|---|---|---|---|---| +| DeepSeek-V4-Pro | `/data/DeepSeek-V4-Pro` | 8 | `--level 0` (CG not yet enabled) | tokenizer has no chat template — do NOT pass `--use-chat-template` | +| DeepSeek-V4-Pro MTP1 | `/data/DeepSeek-V4-Pro` | 8 | `--level 0 --method mtp --num-speculative-tokens 1` | same — no chat template | +| DeepSeek-R1-0528 | `/data/DeepSeek-R1-0528` | 8 | | | +| DeepSeek-R1-0528 MTP3 | `/data/DeepSeek-R1-0528` | 8 | `--method mtp --num-speculative-tokens 3` | `--use-chat-template` REQUIRED on bench | +| DeepSeek FP4 MTP3 | `/data/DeepSeek-R1-0528-MXFP4-MTP-MoEFP4` | 8 | `--method mtp --num-speculative-tokens 3` | `--use-chat-template` REQUIRED | +| GLM-5 FP8 | `/data/GLM-5-FP8` | 8 | | | +| gpt-oss-120b | `/data/openai/gpt-oss-120b` | 1 | drop `-tp` | | +| Kimi-K2.5-MXFP4 | `/data/Kimi-K2.5-MXFP4` | 4 | `--trust-remote-code` + `HSA_NO_SCRATCH_RECLAIM=1` env | | + +## Anti-patterns + +- **Setting `ATOM_PROFILER_MORE=1` on the bench client.** The profiler runs in the model-runner workers; client env is ignored. Set it on the server launch. +- **Using `PROMPT_MULTIPLIER=10` (the bench default) for a trace run.** That sends 10× `CONC` requests and produces a multi-GB trace that takes forever to export and can OOM the workers. +- **Trusting the `.gz` size to decide if the export finished.** Use `.gz exists AND no orphan .json`. The exporter writes `.json` first, then gzips + unlinks; a leftover `.json` means it crashed mid-export. +- **Forgetting `--use-chat-template` on MTP DeepSeek-R1 benchmarks.** Tokenizer mismatch silently degrades accuracy — the trace will look fine but the workload is wrong. +- **Adding `--mark-trace` or `ENABLE_TORCH_PROFILER=1`.** Neither is needed — `--torch-profiler-dir` on the server + `PROFILE=1` on the bench is the complete handshake. The extras either no-op or interfere. +- **Profiling under `--level 3` for V4-Pro right now.** V4-Pro Inductor + autotune hits a `cluster_dims` bug on AMD — start with `--level 0` until CG support lands. + +## Real example — May 17 2026 + +Captured an MTP-1 vs no-MTP V4-Pro trace at `1024/1024/c=64`, 64 prompts each, to compare whether `_hash_topk` uses the fused `aiter::topk_softplus` kernel or the native `softplus + sqrt` chain: + +```bash +mkdir -p /app/logs_claude/traces/v4_mtp1 +bash /app/ATOM/scripts/start_atom_server.sh /data/DeepSeek-V4-Pro 8 8000 \ + --level 0 --method mtp --num-speculative-tokens 1 \ + --torch-profiler-dir /app/logs_claude/traces/v4_mtp1 +bash /app/ATOM/scripts/run_benchmark.sh /data/DeepSeek-V4-Pro 8000 1024 1024 64 1 1 +# (wait loop on .gz / .json — see Step 3) +zcat /app/logs_claude/traces/v4_mtp1/rank_0/*.gz | python3 -c "..." +``` + +Finding: 12,938 `aiter::topk_softplus` calls (the `elif scoring_func == 'sqrtsoftplus'` +branch in `moe.py:2637` IS taken) plus 12,658 native `softplus_kernel` calls +(from `_hash_topk` in `deepseek_v4.py:1982` — the first 3 hash layers don't go +through the fused path). Trace files: `/app/logs_claude/traces/v4_mtp1/` and +`/app/logs_claude/traces/v4_nomtp/`. + +## Cross-references + +- [[debug-agent-locate-kernel]] — when the server crashes or hangs, this skill is the wrong tool +- [[dump-bisect-debug]] — when the trace shows correct kernels but outputs are wrong +- [[atom-patterns]] — V4 attention buffer/stream conventions referenced from trace kernel names diff --git a/.claude/skills/debug-agent-locate-kernel/SKILL.md b/.claude/skills/debug-agent-locate-kernel/SKILL.md new file mode 100644 index 0000000000..968d0dc490 --- /dev/null +++ b/.claude/skills/debug-agent-locate-kernel/SKILL.md @@ -0,0 +1,248 @@ +--- +name: debug-agent-locate-kernel +description: Use rocm-debug-agent to identify which GPU kernel is faulting/hanging when ATOM server hits a HIP memory access fault, MEMORY_VIOLATION, or silent infinite loop. The agent dumps wave registers, faulting PC, and (with --save-code-objects) the disassembled code object so you can name the exact kernel and trace it back to the Python call site. Use when: server crashes with "Memory access fault by GPU node-N", server hangs with GPU at 100% but no token output, or you need to identify a kernel asserting `s_trap`. Do NOT use for: numerical bugs (use dump-bisect-debug), compile errors, OOM. +version: 1.0.0 +scope: ATOM on AMD ROCm (debug-agent at /opt/rocm/lib/librocm-debug-agent.so.2) +last_updated: 2026-05-14 +--- + +## When to use + +Symptoms that point to this skill: + +- `Memory access fault by GPU node-N (Agent handle: 0x...) on address 0x...` in `atom_server.log` +- Server alive (`curl /v1/models` returns) but `rocm-smi --showuse` shows 100% GPU and no `Engine Core: output send` for >30s — silent kernel hang / livelock +- Workers stuck at `torch.cuda.synchronize()` per `py-spy dump --pid ` — prior kernel never completes +- `HIP_LAUNCH_BLOCKING=1` makes the bug disappear → you have an async race; agent will tell you which kernel +- Reproduces only at certain batch shapes (e.g. MTP-3 + long prefill) + +Do NOT use this skill for: precision bugs (use [[dump-bisect-debug]]), build/compile errors (use `/build-fix`), OOM. + +## Required tools (verify before starting) + +```bash +ls /opt/rocm/lib/librocm-debug-agent.so.2 # the agent itself +ls /opt/rocm/llvm/bin/llvm-objdump # for disassembling code objects +which py-spy && py-spy --version # for stack traces of stuck workers +``` + +If any missing: install `rocm-debug-agent`, `llvm`, `pip install py-spy`. Stop here if not available. + +## Critical pre-flight + +1. **`ulimit -c 0`** — disables gpucore dumps. ROCm fault dumps gpucore files of 30-50 GB each per rank; on 8-GPU TP this fills disk in seconds. The launcher script sets this for you. +2. **`--enforce-eager` / `--level 0` — optional fallbacks, not required.** Try the default launch first; the debug agent runs fine under hipgraph in most cases. Only reach for these flags when symptoms point at graph mode: + - **`--enforce-eager`** disables CUDAGraph capture. Try this when the agent reports faults that don't reproduce in eager mode, or when capture/replay itself crashes under the agent's no-caching-allocator behavior. + - **`--level 0`** disables Inductor. Try this on AMD when you hit the `cluster_dims` autotune bug or other Inductor-side crashes during warmup. + - They are independent — apply only the one(s) the symptom points at. The launcher script does NOT inject either; pass via `EXTRA_ARGS` when you want them. +3. **Clean GPU state** — kill any prior `spawn_main`/`openai_server` processes. Stale `KFD process` entries (`rocm-smi --showpids` showing UNKNOWN PIDs holding VRAM) cause the next launch to OOM at NCCL barrier. `scripts/start_atom_server.sh` does the standard cleanup. +4. **Model-specific env** — pass on the command line or export before calling. Examples: + - V4-Pro requires `ATOM_USE_TRITON_MOE=1` + - Kimi-K2.5-MXFP4 requires `--trust-remote-code` + `HSA_NO_SCRATCH_RECLAIM=1` + +## Launcher scripts (in repo) + +| Script | Purpose | +|--------|---------| +| `scripts/start_atom_server.sh [MODEL] [TP] [PORT] [EXTRA_ARGS...]` | Standard launcher: clears GPU, clears compile cache, backgrounds server, redirects to `atom_server.log`. | +| `scripts/stop_atom_server.sh` | SIGTERM atom.entrypoints, force-kill spawn workers, wait for VRAM release. | +| `scripts/run_debug_agent.sh [MODEL] [TP] [PORT] [EXTRA_ARGS...]` | Wraps `start_atom_server.sh` with `HSA_TOOLS_LIB=librocm-debug-agent.so.2 + --save-code-objects`. Server output goes to `atom_server.log`; code objects land in `/app/logs_claude/debug_run/`. | +| `scripts/run_debug_agent.sh --simple [MODEL] [TP] [EXTRA_ARGS...]` | Same wrapper but invokes `start_simple_inference.sh` (offline, no port). Default log: `/app/logs_claude/simple_inference_debug_agent.log` (override via `LOG_FILE=`). Use for offline batch repros (e.g. V4 MTP-3 prefill hang). | +| `scripts/wait_server_ready.sh [PORT] [MAX_MIN] [POLL] [LOG_FILE]` | Poll `/v1/models` until ready or startup error detected. Allow MAX_MIN ≥ 5 under the agent (3-5× slower than normal). | + +## Workflow + +### Step 1: Reproduce under the agent + +```bash +bash scripts/stop_atom_server.sh # ensure clean +ATOM_USE_TRITON_MOE=1 \ + bash scripts/run_debug_agent.sh \ + /data/DeepSeek-V4-Pro 8 8000 \ + --method mtp --num-speculative-tokens 3 & +# If launch fails / faults look graph-mode-specific, retry with +# `--enforce-eager` (and `--level 0` on AMD for the Inductor cluster_dims bug) +# appended to EXTRA_ARGS. + +bash scripts/wait_server_ready.sh 8000 5 30 # 2-4 min under agent +cd /app/logs_claude && python .py # smallest hang trigger +``` + +Server load is **3-5× slower** under the debug agent. Expect ready at 2-4 min, repro at 30-90s after first big batch. + +### Step 2: Find the fault wave dump + +```bash +grep -E "stopped, reason|Memory access fault|MEMORY_VIOLATION|Disassembly" \ + /app/logs_claude/atom_server.log | head -20 +``` + +Each fault produces a block like: + +``` +wave_27876: pc=0x7f20f5e534c4 (kernel_code_entry=0x7f20f5e52900 ) (stopped, reason: ) + +scalar registers: ... +vector registers: ... ← v0..v? show per-lane values; v6 often holds index values being processed +trap registers: ... +general registers: pc=... + +Disassembly for function : + code object: memory://#offset=&size= + loaded at: [-] + => : +``` + +The `` is the demangled kernel name. **That's the suspect kernel.** Common cases: + +| Kernel name fragment | What it actually is | +|----------------------|---------------------| +| `at::native::index_copy_kernel_impl>` | `Tensor.index_copy_(dim, idx, src)` for dtype with N-byte size (4=int32/float32, 8=int64/float64) | +| `at::native::scatter_kernel` | `Tensor.scatter_(dim, idx, src)` | +| `at::native::index_kernel_impl` | Advanced indexing READ `tensor[idx]` | +| `_swa_write_kernel` / `_update_compressor_states_kernel` | ATOM Triton kernel — name in `state_writes.py` | + +### Step 3: Read the trap reason + +| reason | what it means | +|--------|---------------| +| `ASSERT_TRAP` | Kernel hit `s_trap 2` — almost always a `CUDA_KERNEL_ASSERT(...)` failed device-side. For PyTorch `index_copy_`/`scatter_` this is the bound check `0 <= idx < self.size(dim)`. Recompile PyTorch with `TORCH_USE_HIP_DSA=1` for the assert text — usually unavailable, infer from kernel name. | +| `MEMORY_VIOLATION` | Real OOB load/store. The `pc` instruction is the access; back-trace the address from `s_*`/`v_*` registers. | +| `INVALID_OPCODE` | Corrupted code object — usually an allocator stomp on the kernel binary (very rare). | + +### Step 4: Disassemble the code object + +The trap dump points to `code object: memory://#offset=&size=`. The agent saved it under `/app/logs_claude/debug_run/`. Find it: + +```bash +ls /app/logs_claude/debug_run/ | grep "" | grep "size_" +# Returns e.g.: 7_memory___2188702_offset_0x546c3060_size_4026672 +``` + +Disassemble: + +```bash +/opt/rocm/llvm/bin/llvm-objdump --disassemble-all \ + /app/logs_claude/debug_run/ > /app/logs_claude/fault.s +grep -nE "|s_trap|s_endpgm" /app/logs_claude/fault.s | head -20 +``` + +The PC's surrounding instructions tell you what the kernel was doing. For `s_trap 2` followed by `s_endpgm` you've confirmed an assert (PyTorch `CUDA_KERNEL_ASSERT` lowering). For random other instructions it's a true memory violation — read the address from registers (e.g. `v[0:1]` typically holds the destination address being stored). + +### Step 5: Verify it's actually that kernel (PC can lie) + +Wave debugger PC reports can be **off** when the wave is mid-flight or when the trap fires from a sibling wave. Especially common with Triton — a swa_write trap might be a downstream kernel's fault attributed back. Cross-check: + +- Does the trap reproduce **only when this code path runs**? Disable the call (comment out in Python), retest. +- Does **`HIP_LAUNCH_BLOCKING=1`** make it disappear? Then it's an async race, not a static OOB; the PC kernel is the **victim**, not necessarily the root cause. Bisect for the racer (next step). +- Does inserting `torch.cuda.synchronize()` **right before** this kernel call eliminate the trap? Then root cause is upstream of this point on the same stream. + +### Step 6: Bisect the racer (when PC is racer-victim) + +1. **Comment out one suspect call at a time.** The one whose absence fixes it is the racer (or one of the racing parties). +2. **If neither alone but both together fail**: the race is between them sharing storage / launch slot. Add `torch.cuda.synchronize()` between them as a workaround, but THIS IS NOT A SHIPPABLE FIX — see Step 7. +3. **`py-spy dump --pid `** on stuck ranks: shows the Python frame waiting on the GPU. If it's at your inserted `synchronize()`, the racer is upstream of that line. + +### Step 7: Real fix vs workaround + +Per [[atom-patterns]] / DeepSeek V4 guidance, do not ship `cuda.synchronize()` workarounds without root-causing the race — they mask one workload and surface a worse hang on a larger one. Common real fixes: + +| Symptom | Real fix | +|---------|----------| +| Race involves freshly-allocated transient tensors (e.g. from `torch.where`, `arange`, `.reshape`, `.to(int64)`) | Pre-allocate them in `_alloc_v4_metadata_buffers` (ATOM) or as module-level scratch. Eliminates allocator churn entirely. | +| Multiple `index_copy_` / `scatter_` in sequence | Replace with a single Triton kernel that writes all destinations once. | +| Per-fwd kernel reads stale forward_vars from prior fwd | Switch H2D path off `prep_stream` to default stream (matches ATOM `prepare_mtp_decode` v2 pattern). | +| Cross-rank inconsistency causes one rank to OOB | Ensure all ranks see identical batch shapes before launching kernel; check `cu_seqlens_q` / `state_slot_mapping` parity. | + +## Recovery checklist (after agent run) + +1. `bash scripts/stop_atom_server.sh` — agent leaves zombie KFD entries; if you skip, next launch OOMs at NCCL barrier. +2. `pkill -9 -f spawn_main` — sometimes `stop_atom_server.sh` misses workers stuck in fault state. +3. Wait 10s, then `rocm-smi --showmemuse` — all GPUs must show 0% before relaunching. If not, `rocm-smi --showpids` to find lingering UNKNOWN PIDs (killed but KFD hasn't cleaned yet — wait or escalate). +4. `rm /app/logs_claude/debug_run/memory_*` — code objects are 4 MB each, accumulate fast across runs. +5. Drop `--save-code-objects` from production launches — disk pressure (~500 MB per run). + +## Anti-patterns + +- **Don't assume `--enforce-eager --level 0` is mandatory.** Default launch is fine for most agent runs; reach for these flags only when symptoms point at hipgraph or Inductor (see pre-flight item 2). Adding them blindly hides graph-mode-only bugs. +- **Don't grep `atom_server.log` for "error" or "Traceback"** — agent's wave dump has neither; grep `"stopped, reason"` instead. +- **Don't trust PC literally** — see Step 5. Especially Triton kernels are notorious for cross-wave PC misattribution. Bisect-confirm. +- **Don't leave `--save-code-objects` on for routine runs** — each run dumps ~500 MB. Only enable for the bisect run. +- **Don't add `torch.cuda.synchronize()` "fixes" and ship** — they mask the race for one workload and surface a worse hang (livelock) on a larger one. Find the allocator/stream root cause. + +## Real example (May 14 2026 — V4-Pro MTP-3 prefill hang) + +**Symptom**: `Memory access fault by GPU node-N` on all 8 ranks within 1s of "Scheduled prefill batch: 19 reqs, 9573 tokens". `HIP_LAUNCH_BLOCKING=1` and `ATOM_DUMP_SWA_WRITE=1` (which inserts `.cpu()` sync) both eliminated it. + +**Workflow**: +1. Step 1-2: ran `run_debug_agent.sh`, grep `"stopped, reason"` → got `ASSERT_TRAP` in `at::native::index_copy_kernel_impl>`. +2. Step 4: disassembled code object → confirmed `s_trap 2` followed by `s_endpgm` → PyTorch `CUDA_KERNEL_ASSERT` failed (bound check on index_copy). +3. Step 5: grep'd ATOM source for `index_copy_` of int32 dtype → only csa/hca calls in `deepseek_v4_attn.py:1505-1506`. +4. Step 6: bisected — neither alone trapped, both together did. Added `torch.cuda.synchronize()` between → small workload passed, GSM8K (1319 reqs) eventually livelocked at 1281, py-spy showed all ranks stuck in the synchronize call → confirmed root cause is allocator churn from transient tensors (`csa_win_pos`/`hca_win_pos`/`swa_paged_flat` built via `torch.where`/`arange`/`.reshape`/`.to(int64)`). +5. Real fix (in progress): replace 3 sequential `index_copy_` + transient tensor construction with a single Triton kernel reading only persistent forward_vars buffers. + +Wave dumps and code objects from this run: `/app/logs_claude/atom_server.log` (line 419+), `/app/logs_claude/debug_run/`, `/app/logs_claude/fault.s`. + +## Sample wave dump (what to expect in atom_server.log) + +Trimmed real example from the May 14 2026 V4-Pro MTP-3 run. Key fields are the +mangled function name in `kernel_code_entry=...`, the `stopped, reason: ...` +tag, the `code object: memory://#offset=&size=` line that +points at the saved file, and the `=> : ` arrow showing the +faulting PC. Vector registers (only v0/v1/v6 shown — full dump prints v0..v15 +and beyond) often reveal address / index values that pin down the operand. + +``` +[atom 15:31:09] Scheduled prefill batch: 19 reqs, 9573 tokens, req_ids: (1, 2, ..., 19) +... (some [aiter] type-hints chatter, then the agent's wave dump arrives) ... +-------------------------------------------------------- +wave_27876: pc=0x7f20f5e534c4 (kernel_code_entry=0x7f20f5e52900 >(at::TensorIterator&, long, long, long)::{lambda(int)#1}>(long, at::native::index_copy_kernel_impl >(at::TensorIterator&, long, long, long)::{lambda(int)#1})>) (stopped, reason: ASSERT_TRAP) + +scalar registers: + s0: ffffffff s1: ffffffff s2: 00000000 s3: f8000000 + ... + s32: 0ec00000 s33: 00000002 ... + +system registers: + mode: 000003f0 trapsts: 80000000 status: 80010041 + +trap registers: + ttmp4: 00006ce4 ttmp5: 00000000 ... + +vector registers: + v0: [0] 95f02814 [1] 95f02818 [2] 95f0281c [3] 95f02820 ... [58] 95f028fc [59] 00000a40 ... + v1: [0] 00007f20 [1] 00007f20 ... ← v0:v1 = per-lane dst address + v6: [0] 000080a6 [1] 000080a7 ... [58] 000080e0 [59] 000027ec ← per-lane src VALUES being stored + +general registers: + m0: 000103c0 + pc: 00007f20f5e534c4 exec: f800000000000000 + vcc: ffffffffffffffff + +Disassembly for function void at::native::index_elementwise_kernel<128, 4, at::native::index_copy_kernel_impl >(...)>: + code object: memory://2188702#offset=0x546c3060&size=4026672 + loaded at: [0x7f20f5e00000-0x7f20f615ff09] + => 0x7f20f5e534c4 <+3012>: s_endpgm + 0x7f20f5e534c8 <+3016>: v_cndmask_b32_e32 v0, s0, v0, vcc +``` + +How to read this: + +- **Kernel** = `at::native::index_copy_kernel_impl>` → PyTorch + `Tensor.index_copy_(dim, idx, src)` for 4-byte dtype (int32 / float32). +- **Reason** = `ASSERT_TRAP` → some lane's `index_value` failed + `0 <= idx < self.size(dim)`. Look at `v6` per-lane values to see what was + being processed (if v6 holds the stored value here; the relevant register + varies by kernel). +- **PC** lands on `s_endpgm` because the assert lowering is + `s_trap 2; s_endpgm` — the actual condition test was earlier (look ~10 + instructions back in the disassembly for `s_cbranch_*` + `s_trap`). +- **Code object** at `memory://2188702#offset=0x546c3060&size=4026672` → + saved as `7_memory___2188702_offset_0x546c3060_size_4026672` in + `/app/logs_claude/debug_run/`. Use `llvm-objdump --disassemble-all` on it. + +## Cross-references + +- [[dump-bisect-debug]] — for numerical bugs (wrong output, not crashes) +- [[capture-trace]] — for performance investigation (which kernels eat time) +- [[atom-patterns]] — V4 attention buffer / stream conventions diff --git a/.claude/skills/dump-bisect-debug.md b/.claude/skills/dump-bisect-debug/SKILL.md similarity index 100% rename from .claude/skills/dump-bisect-debug.md rename to .claude/skills/dump-bisect-debug/SKILL.md diff --git a/.claude/skills/run-atom-workload/SKILL.md b/.claude/skills/run-atom-workload/SKILL.md new file mode 100644 index 0000000000..68d329fdb4 --- /dev/null +++ b/.claude/skills/run-atom-workload/SKILL.md @@ -0,0 +1,239 @@ +--- +name: run-atom-workload +description: Run any ATOM workload — accuracy eval (GSM8K via lm_eval), performance benchmark, concurrency sweep, offline simple_inference, or fault repro under rocm-debug-agent. Use when the user asks to "test accuracy", "测精度", "跑 GSM8K", "跑 benchmark", "test performance", "run sweep", "repro the fault", "测一下 MTP1 精度", "跑 simple_inference" — anything that drives an ATOM workload. Encodes the canonical flow (stop → start → workload-in-shell-bg → wait_infer_drain → stop) and the model-family env vars. Same pattern works for both server-based workloads (lm_eval / benchmark client) and offline simple_inference. Do NOT use for profiling traces (use capture-trace). +version: 1.2.0 +scope: ATOM on AMD ROCm; `scripts/` orchestration scripts under repo root +last_updated: 2026-05-18 +--- + +## Path convention + +All script paths in this skill are **project-relative** (`scripts/foo.sh`), +not absolute. The skill lives inside the ATOM repo at `.claude/skills/`, so +it travels with the repo wherever it's cloned. CWD when invoking commands +should be the repo root (Claude Code's default). + +If the user is somewhere else, prefix with `bash $(git rev-parse --show-toplevel)/scripts/foo.sh` +or `cd` to the repo first. + +## Why this skill exists + +Every blocking ATOM workload follows the same 4-step shape: **stop → start → workload-in-bg → wait_infer_drain → stop**. The scripts in `scripts/` are orchestration-grade — chain them via **separate Bash tool calls**, not wrappers, not `&&`. + +Past failure modes this skill prevents (collected from many sessions): + +1. **Writing wrapper scripts in `/app/logs_claude/`** — `start_atom_server.sh` etc. ARE the orchestration layer. Wrapping them is pure noise. The dozens of `run_*.sh` / `start_*_safe.sh` files in `/app/logs_claude/` are session debris — **do not mimic them**. +2. **Chaining all steps with `&&` into one long command** — the user has explicitly forbidden this. Each step gets its own Bash tool call so logs are separate, errors abort cleanly, and the user can interrupt at any boundary. +3. **Double-backgrounding the server** — `start_atom_server.sh` already forks python in background internally and polls ready. Adding `&` after `start_atom_server.sh` is wrong; it blocks until ready or fail. +4. **Skipping `wait_infer_drain.sh`** — without it, GPU faults take the whole timeout to surface, and hangs go undetected. `wait_infer_drain.sh` exits in ~10s on fault and ~1min on hang, with tail-log attached. +5. **Using `curl /health` for liveness** — under heavy load it can false-negative. The flow uses `/v1/models` (start script) and `pgrep` + Engine Core marker (drain script). +6. **Forgetting model-family env vars** — V4-Pro silently regresses on accuracy without `AITER_BF16_FP8_MOE_BOUND=0 ATOM_MOE_GU_ITLV=1`. Pinned in the table below. +7. **Skipping drain for offline simple_inference** — `wait_infer_drain.sh` supports offline mode (process-exit detection + fault scan). Without it you lose early fault visibility. +8. **Passing the wrong LOG_FILE to drain** — historically callers had to know which log carries the "Engine Core: output send" marker. v1.2 of drain auto-discovers the server log via `/proc//fd/1`, so the user LOG_FILE is now only a **secondary** signal (fault scan + mtime progress for clients with tqdm output). Pass any log you want extra coverage on, or pass nothing — drain still works. + +## Backgrounding mechanism — shell `&`, NOT claude task + +In step 3 the workload must run concurrent with step 4's drain monitor. Use **shell-level `&`** (append to the bash command), NOT the Bash tool's `run_in_background: true`: + +- Shell `&`: bash starts workload, returns immediately, workload runs as orphan; drain finds it via `pgrep` — no claude task tracking dependency +- `run_in_background: true`: workload becomes a claude task accessed via TaskOutput — adds complexity, doesn't help drain since drain uses pgrep anyway + +Pattern (literal): +``` +# Step 3 — single Bash tool call, command ends with `&` +bash scripts/run_gsm8k_eval.sh /data/MODEL 30000 3 & + +# Step 4 — single Bash tool call, blocks +bash scripts/wait_infer_drain.sh 30000 30 10 +``` + +The Bash invocation for step 3 returns the instant `&` is processed (`bash -c 'cmd &'` exits as soon as cmd is backgrounded). Step 4 then runs as the next Bash call and blocks on drain. + +## Canonical 4-step flow + +Run each step as a **separate Bash tool call**. Never chain with `&&`. + +### Step 1 — clean GPU (always) + +```bash +bash scripts/stop_atom_server.sh +``` + +Idempotent. SIGTERM → SIGKILL → force-kill GPU PIDs, waits ≤60s for VRAM=0. Always run first even if you believe no server is up — clears orphaned multiprocessing children that hold GPU memory. + +### Step 2 — start workload host (blocks until ready / completion) + +**Server-based workloads** (GSM8K / benchmark / sweep / fault repro): + +```bash + bash scripts/start_atom_server.sh +``` + +- Self-contained: forks python in background internally, polls `/v1/models` + GPU VRAM, returns when ready or 1 on fail +- **Do NOT** add `&` — the script self-backgrounds python and waits for ready +- **Do NOT** chain `wait_server_ready.sh` after — already inlined inside start script +- Log: hard-coded `/app/logs_claude/atom_server.log` (`LOG_FILE` env is NOT respected by this script). Drain auto-discovers it via `/proc//fd/1` regardless of path + +**Offline workload** (simple_inference): step 2 is fused with step 3, see below. + +Model-family env vars (set as `VAR=val VAR=val bash ...` prefix): + +| Model | Required env vars | Required CLI args | +|---|---|---| +| DeepSeek-V4-Pro | `AITER_BF16_FP8_MOE_BOUND=0 ATOM_MOE_GU_ITLV=1 AITER_LOG_LEVEL=WARNING` | `--kv_cache_dtype fp8 --level 0` | +| DeepSeek-R1-0528 (default) | `AITER_LOG_LEVEL=WARNING` | `--kv_cache_dtype fp8` | +| Kimi-K2.5-MXFP4 | `HSA_NO_SCRATCH_RECLAIM=1 AITER_LOG_LEVEL=WARNING` | `--kv_cache_dtype fp8 --trust-remote-code` (tp=4) | + +MTP add-on (any supporting model): append `--method mtp --num-speculative-tokens N` to EXTRA_ARGS. V4-Pro: keep `--level 0`. + +### Step 3 — launch workload in shell background (`&`) + +The workload script must end with shell `&` so the Bash tool returns immediately and step 4 can start monitoring in parallel. + +**Server-based workloads** (PORT is needed): + +| Workload | Command (note trailing `&`) | Optional client log for drain | +|---|---|---| +| GSM8K accuracy | `bash scripts/run_gsm8k_eval.sh MODEL PORT NUM_FEWSHOT &` | `/app/logs_claude/gsm8k_eval.log` (lm_eval is silent during requests; drain's auto-discovered server log carries the engine markers — passing this log only helps fault grep coverage) | +| Single benchmark | `bash scripts/run_benchmark.sh MODEL PORT ISL OSL CONC [PROMPT_MULT] [PROFILE] &` | `/app/logs_claude/benchmark.log` (has tqdm progress, useful mtime signal) | +| Concurrency sweep | `bash scripts/run_benchmark_sweep.sh MODEL PORT ISL OSL "CONC1 CONC2 ..." &` | `/app/logs_claude/benchmark.log` (overwritten per step) | + +**Offline simple_inference** (no PORT; step 2 is skipped since this script IS the workload host): + +```bash + bash scripts/start_simple_inference.sh MODEL TP & +``` + +Optional client log for drain: `/app/logs_claude/simple_inference.log` (drain auto-discovers via /proc anyway; this only helps fault grep redundancy). + +Common workload knobs: +- GSM8K shots: 3 for fast/CI parity, 5 for thorough. Set `LIMIT=50` env for first-50-sample sanity. +- Benchmark `PROMPT_MULTIPLIER` default 10. Profiling: use 2 (CONC × 2 requests). +- MTP benchmark MUST add `--use-chat-template` via EXTRA_ARGS (tokenizer mismatch otherwise). +- Benchmark throughput metric: report **Total Token throughput (tok/s)**, NOT Output throughput. Total = input+output, which users care about. +- Never add `--mark-trace` or `ENABLE_TORCH_PROFILER=1` (handled by capture-trace skill). + +### Step 4 — wait_infer_drain (blocks, with early fault/hang detection) + +```bash +bash scripts/wait_infer_drain.sh PORT MAX_MIN POLL_SEC [LOG_FILE] [STUCK_POLLS] +``` + +Defaults: PORT=8000, MAX_MIN=30, POLL_SEC=10, LOG_FILE=empty (server log auto-discovered via `/proc//fd/1`), STUCK_POLLS=6. + +`LOG_FILE` is **optional** in v1.2+. The drain script discovers the server log itself from the running `atom.entrypoints` process. Pass an additional client/workload log only if you want: +- Extra fault grep coverage (drain scans both) +- Mtime-based progress detection for client tools that write tqdm to a file (benchmark, simple_inference) + +PORT is unused in offline mode but kept positional for API symmetry. + +How drain decides (auto-detects server vs offline by `SERVER_PATTERN` pgrep): +- **Server mode**: client gone (lm_eval / curl / benchmark process exited) + no new "Engine Core: output send" since last poll → exit 0 +- **Offline mode**: simple_inference process exited cleanly (no fault grep) → exit 0 +- **Either mode**: fault grep on auto-discovered server log + optional caller LOG_FILE → exit 2 in ≤10s +- **Server only**: no progress (engine output count flat + caller LOG_FILE mtime flat) AND client still running for STUCK_POLLS × POLL_SEC ≈ 1min → exit 1 (hang) +- **Either mode**: MAX_MIN elapsed without resolution → exit 4 + +If exit ≠ 0: read the printed tail, then run step 5 regardless. + +Typical wait windows: +- GSM8K (1319 samples): MAX_MIN=30 plenty for V4-Pro +- Single benchmark: MAX_MIN=30 +- Sweep (8 conc points): MAX_MIN=60 +- Simple_inference (default ~10 prompts): MAX_MIN=15 plenty +- Fault repro: MAX_MIN=10 (fault should land within first request) + +### Step 5 — teardown (always) + +```bash +bash scripts/stop_atom_server.sh +``` + +Same script as step 1. ALWAYS run, even on fault or for offline workloads — releases GPU for next attempt and kills any lingering multiprocessing children. + +## Reading results + +After step 4 returns 0: + +```bash +# GSM8K +grep -E "flexible-extract|strict-match" /app/logs_claude/gsm8k_eval.log | head -2 +# Benchmark +grep -E "Total Token throughput|Mean TPOT|Mean TTFT" /app/logs_claude/benchmark.log +# Simple_inference +grep -E "^Generated|^Output|tokens/s" /app/logs_claude/simple_inference.log +``` + +GSM8K format: `|gsm8k|3|flexible-extract|3|exact_match|↑|0.XXXX|±|0.00XX|` — flexible-extract is the headline number, ±value is the noise band. Anything within 1σ of baseline = no regression. + +## V4-Pro accuracy baselines (current, 3-shot, no `--limit`) + +- Non-MTP: ~0.9545 ± 0.0057 +- MTP-1: ~0.9492 ± 0.006 +- MTP-3: see `feedback_v4_cg_mtp_status_deadlock.md` — currently hangs on high-concurrency GSM8K + +## Hard rules (do not violate) + +1. **One Bash tool call per script.** No `&&` chains. User has explicitly forbidden chaining. +2. **No wrapper scripts in `/app/logs_claude/`.** Call `scripts/*` directly. +3. **Shell `&`, not `run_in_background: true`** for step 3 (drain finds workload via pgrep, no task tracking needed). +4. **Drain auto-discovers server log via `/proc//fd/1`** — you no longer need to know or pass the canonical server log path. Pass a client log only as supplementary signal. +5. **Always step 5 (`stop_atom_server.sh`)**, even after a fault, even for offline workloads. +6. **Never use `curl /health`** to verify ready — only `/v1/models` (already inlined in start script). +7. **No extra `&`, `wait_server_ready.sh`, `LOG_FILE=` env on `start_atom_server.sh`** — start script self-backgrounds, self-polls, hard-codes log path. + +## Reference: each script in one line + +| Script | What it does | Step | Blocks? | +|---|---|---|---| +| `stop_atom_server.sh` | Kill all atom + multiproc children, wait for VRAM=0 | 1, 5 | Yes ≤60s | +| `start_atom_server.sh MODEL TP PORT [ARGS...]` | Clean GPU, fork python in bg, poll ready | 2 (server) | Yes, until ready or fail | +| `start_simple_inference.sh MODEL TP [ARGS...]` | Offline inference (no server, runs prompts) — wrap with `&` for drain | 3 (offline) | Blocks unless `&` | +| `run_gsm8k_eval.sh MODEL PORT FEWSHOT` | lm_eval local-completions GSM8K — wrap with `&` for drain | 3 (server) | Blocks unless `&` | +| `run_benchmark.sh MODEL PORT ISL OSL CONC [PMULT] [PROF]` | Single perf point — wrap with `&` for drain | 3 (server) | Blocks unless `&` | +| `run_benchmark_sweep.sh MODEL PORT ISL OSL "CONCs"` | Loop run_benchmark — wrap with `&` for drain | 3 (server) | Blocks unless `&` | +| `wait_infer_drain.sh PORT MAX_MIN POLL [LOG] [STUCK]` | Monitor workload for drain / hang / fault (auto-discovers server log) | 4 | Yes, until exit code | +| `wait_server_ready.sh PORT MAX_MIN POLL LOG` | Standalone ready poller (rarely needed; start script self-polls) | (debug) | Yes | +| `run_debug_agent.sh [--simple] MODEL TP [PORT] [ARGS...]` | Server (or simple_inference) under rocm-debug-agent — fault repro | 2 (replaces start) | Yes, until ready or fault | + +## Worked example: V4-Pro MTP1 GSM8K accuracy + +``` +# Step 1 +bash scripts/stop_atom_server.sh + +# Step 2 +AITER_BF16_FP8_MOE_BOUND=0 ATOM_MOE_GU_ITLV=1 AITER_LOG_LEVEL=WARNING \ + bash scripts/start_atom_server.sh /data/DeepSeek-V4-Pro 8 30000 \ + --kv_cache_dtype fp8 --method mtp --num-speculative-tokens 1 --level 0 + +# Step 3 — note trailing `&` +bash scripts/run_gsm8k_eval.sh /data/DeepSeek-V4-Pro 30000 3 & + +# Step 4 — drain auto-discovers server log; no LOG_FILE needed +bash scripts/wait_infer_drain.sh 30000 30 10 + +# Step 5 +bash scripts/stop_atom_server.sh + +# Read result +grep -E "flexible-extract|strict-match" /app/logs_claude/gsm8k_eval.log | head -2 +``` + +## Worked example: V4-Pro offline simple_inference + +``` +# Step 1 +bash scripts/stop_atom_server.sh + +# Step 2+3 fused (simple_inference IS the workload host) — note trailing `&` +AITER_BF16_FP8_MOE_BOUND=0 ATOM_MOE_GU_ITLV=1 AITER_LOG_LEVEL=WARNING \ + bash scripts/start_simple_inference.sh /data/DeepSeek-V4-Pro 8 \ + --kv_cache_dtype fp8 --level 0 & + +# Step 4 — drain auto-discovers via /proc; PORT unused +bash scripts/wait_infer_drain.sh 0 15 10 + +# Step 5 +bash scripts/stop_atom_server.sh +``` diff --git a/.github/benchmark/models.json b/.github/benchmark/models.json index b542d3cf53..f9f2e28c84 100644 --- a/.github/benchmark/models.json +++ b/.github/benchmark/models.json @@ -17,7 +17,17 @@ "bench_args": "", "suffix": "", "runner": "atom-mi355-8gpu.predownload", - "env_vars": "ATOM_USE_TRITON_MOE=1" + "env_vars": "AITER_BF16_FP8_MOE_BOUND=0\nATOM_MOE_GU_ITLV=1" + }, + { + "display": "DeepSeek-V4-Pro MTP1", + "path": "deepseek-ai/DeepSeek-V4-Pro", + "prefix": "deepseek-v4-pro", + "args": "--kv_cache_dtype fp8 -tp 8 --method mtp --num-speculative-tokens 1", + "bench_args": "--use-chat-template", + "suffix": "-mtp1", + "runner": "atom-mi355-8gpu.predownload", + "env_vars": "AITER_BF16_FP8_MOE_BOUND=0\nATOM_MOE_GU_ITLV=1" }, { "display": "DeepSeek-R1-0528 MTP3", @@ -73,7 +83,7 @@ "display": "gpt-oss-120b", "path": "openai/gpt-oss-120b", "prefix": "gpt-oss-120b", - "args": "--kv_cache_dtype fp8", + "args": "--kv_cache_dtype fp8 --gpu-memory-utilization 0.5", "bench_args": "", "suffix": "", "runner": "atom-mi355-8gpu.predownload", @@ -90,9 +100,9 @@ "env_vars": "" }, { - "display": "MiniMax-M2.5", - "path": "MiniMaxAI/MiniMax-M2.5", - "prefix": "MiniMax-M2.5", + "display": "MiniMax-M2.7", + "path": "MiniMaxAI/MiniMax-M2.7", + "prefix": "MiniMax-M2.7", "args": "--kv_cache_dtype fp8 -tp 2 --trust-remote-code", "bench_args": "", "suffix": "", @@ -100,9 +110,9 @@ "env_vars": "" }, { - "display": "MiniMax-M2.5-MXFP4", - "path": "amd/MiniMax-M2.5-MXFP4", - "prefix": "MiniMax-M2.5-MXFP4", + "display": "MiniMax-M2.7-MXFP4", + "path": "amd/MiniMax-M2.7-MXFP4", + "prefix": "MiniMax-M2.7-MXFP4", "args": "--kv_cache_dtype fp8 -tp 1 --trust-remote-code", "bench_args": "", "suffix": "", diff --git a/.github/benchmark/models_accuracy.json b/.github/benchmark/models_accuracy.json index 99efa88f28..fc22c02418 100644 --- a/.github/benchmark/models_accuracy.json +++ b/.github/benchmark/models_accuracy.json @@ -27,7 +27,7 @@ "model_name": "DeepSeek-V4-Pro", "model_path": "deepseek-ai/DeepSeek-V4-Pro", "extraArgs": "--kv_cache_dtype fp8 -tp 8", - "env_vars": "ATOM_USE_TRITON_MOE=1", + "env_vars": "AITER_BF16_FP8_MOE_BOUND=0\nATOM_MOE_GU_ITLV=1", "runner": "atom-mi355-8gpu.predownload", "test_level": "pr", "accuracy_threshold": 0.92, @@ -35,6 +35,18 @@ "accuracy_baseline_model": "deepseek-ai/DeepSeek-V4-Pro", "_baseline_note": "Local 4-run average GSM8K-100 3-shot flexible-extract = 0.96 (runs: 0.96/0.98/0.96/0.94, stderr ~0.024). ATOM_USE_TRITON_MOE=1 is required — without it accuracy drops to ~0.6. Threshold set 4pp below local baseline to absorb full-eval (1319 samples) noise; refresh after first CI measurement." }, + { + "model_name": "DeepSeek-V4-Pro MTP", + "model_path": "deepseek-ai/DeepSeek-V4-Pro", + "extraArgs": "--kv_cache_dtype fp8 -tp 8 --method mtp --num-speculative-tokens 1", + "env_vars": "AITER_BF16_FP8_MOE_BOUND=0\nATOM_MOE_GU_ITLV=1", + "runner": "atom-mi355-8gpu.predownload", + "test_level": "pr", + "accuracy_threshold": 0.92, + "accuracy_baseline": 0.96, + "accuracy_baseline_model": "deepseek-ai/DeepSeek-V4-Pro", + "_baseline_note": "Same base model as DeepSeek-V4-Pro FP8 (MTP-1: 1 speculative token)" + }, { "model_name": "DeepSeek-R1-0528 MTP", "model_path": "deepseek-ai/DeepSeek-R1-0528", @@ -50,12 +62,13 @@ { "model_name": "gpt-oss-120b", "model_path": "openai/gpt-oss-120b", - "extraArgs": "--kv_cache_dtype fp8 --gpu-memory-utilization 0.3", + "extraArgs": "--kv_cache_dtype fp8 --gpu-memory-utilization 0.5", + "client_command": "lm_eval --model local-chat-completions --apply_chat_template --model_args model=${MODEL_PATH},base_url=http://localhost:8000/v1/chat/completions,num_concurrent=65,max_retries=3,max_gen_toks=2048,tokenized_requests=False,trust_remote_code=True --tasks gsm8k --num_fewshot 3 --output_path ${OUTPUT_PATH}", "env_vars": "", "runner": "linux-atom-mi35x-1", "test_level": "pr", - "accuracy_threshold": 0.39, - "accuracy_baseline": 0.41, + "accuracy_threshold": 0.88, + "accuracy_baseline": 0.90, "accuracy_baseline_model": "openai/gpt-oss-120b", "_baseline_note": "No public GSM8K baseline available" }, @@ -110,8 +123,8 @@ { "model_name": "Qwen3-Next-80B-A3B-Thinking", "model_path": "Qwen/Qwen3-Next-80B-A3B-Thinking", - "extraArgs": "-tp 8", - "env_vars": "", + "extraArgs": "-tp 8 --kv_cache_dtype fp8", + "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1", "runner": "atom-mi355-8gpu.predownload", "test_level": "pr", "accuracy_threshold": 0.65, @@ -122,12 +135,13 @@ { "model_name": "gpt-oss-120b (2 GPUs)", "model_path": "openai/gpt-oss-120b", - "extraArgs": "--kv_cache_dtype fp8 -tp 2 --enable-dp-attention --enable-expert-parallel --enable-tbo all --gpu-memory-utilization 0.3", - "env_vars": "", + "extraArgs": "--kv_cache_dtype fp8 -tp 2 --enable-dp-attention --enable-expert-parallel --enable-tbo all --gpu-memory-utilization 0.5", + "client_command": "lm_eval --model local-chat-completions --apply_chat_template --model_args model=${MODEL_PATH},base_url=http://localhost:8000/v1/chat/completions,num_concurrent=65,max_retries=3,max_gen_toks=2048,tokenized_requests=False,trust_remote_code=True --tasks gsm8k --num_fewshot 3 --output_path ${OUTPUT_PATH}", + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4", "runner": "linux-atom-mi35x-4", "test_level": "main", - "accuracy_threshold": 0.39, - "accuracy_baseline": 0.41, + "accuracy_threshold": 0.88, + "accuracy_baseline": 0.90, "accuracy_baseline_model": "openai/gpt-oss-120b", "_baseline_note": "No public GSM8K baseline available" }, @@ -155,6 +169,18 @@ "accuracy_baseline_model": "moonshotai/Kimi-K2.5", "_baseline_note": "HF: amd/Kimi-K2.5-MXFP4 card shows Kimi-K2.5 baseline=0.9409" }, + { + "model_name": "Kimi-K2.5-MXFP4 Eagle3", + "model_path": "amd/Kimi-K2.5-MXFP4", + "extraArgs": "--kv_cache_dtype fp8 -tp 8 --trust-remote-code --method eagle3 --draft-model lightseekorg/kimi-k2.5-eagle3 --num-speculative-tokens 3", + "env_vars": "HSA_NO_SCRATCH_RECLAIM=1", + "runner": "atom-mi355-8gpu.predownload", + "test_level": "nightly", + "accuracy_threshold": 0.91, + "accuracy_baseline": 0.9257, + "accuracy_baseline_model": "amd/Kimi-K2.5-MXFP4 + lightseekorg/kimi-k2.5-eagle3", + "_baseline_note": "Eagle3 spec decode on Kimi-K2.5-MXFP4. Local case_verify_v9_gluon GSM8K 5-shot flexible-extract=0.9257 (vLLM=0.9280, within ±0.71% se). Threshold 0.91 leaves ~1.5pp headroom for noise. -tp 8 (vs base entry's tp=4) because Eagle3 draft KV needs the full 8-rank sharding." + }, { "model_name": "GLM-5-FP8", "model_path": "zai-org/GLM-5-FP8", @@ -179,6 +205,18 @@ "accuracy_baseline_model": "zai-org/GLM-5.1", "_baseline_note": "CI uses 3-shot, not comparable to HF 5-shot baseline" }, + { + "model_name": "GLM-5.1-MXFP4 MTP", + "model_path": "amd/GLM-5.1-MXFP4", + "extraArgs": "--kv_cache_dtype fp8 -tp 4 --default-chat-template-kwargs '{\"enable_thinking\":false}' --method mtp --num-speculative-tokens 3", + "env_vars": "", + "runner": "atom-mi355-8gpu.predownload", + "test_level": "pr", + "accuracy_threshold": 0.87, + "accuracy_baseline": 0.9545, + "accuracy_baseline_model": "zai-org/GLM-5.1", + "_baseline_note": "CI uses 3-shot, not comparable to HF 5-shot baseline" + }, { "model_name": "GLM-5.1-MXFP4", "model_path": "amd/GLM-5.1-MXFP4", @@ -191,11 +229,23 @@ "accuracy_baseline_model": "zai-org/GLM-5.1", "_baseline_note": "CI uses 3-shot, not comparable to HF 5-shot baseline" }, + { + "model_name": "Qwen3.5-397B-A17B", + "model_path": "Qwen/Qwen3.5-397B-A17B", + "extraArgs": "--kv_cache_dtype fp8 -tp 8", + "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1", + "runner": "atom-mi355-8gpu.predownload", + "test_level": "nightly", + "accuracy_threshold": 0.85, + "accuracy_baseline": 0.9538, + "accuracy_baseline_model": "Qwen3.5-397B-A17B-FP8", + "_baseline_note": "BF16 weight, TP8. Placeholder threshold/baseline copied from FP8 entry; refresh after first CI measurement." + }, { "model_name": "Qwen3.5-397B-A17B-FP8", "model_path": "Qwen/Qwen3.5-397B-A17B-FP8", "extraArgs": "--kv_cache_dtype fp8 -tp 4", - "env_vars": "", + "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1", "runner": "atom-mi355-8gpu.predownload", "test_level": "nightly", "accuracy_threshold": 0.85, @@ -207,7 +257,7 @@ "model_name": "Qwen3.5-397B-A17B-FP8 MTP", "model_path": "Qwen/Qwen3.5-397B-A17B-FP8", "extraArgs": "--kv_cache_dtype fp8 -tp 4 --method mtp --num-speculative-tokens 3", - "env_vars": "", + "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1", "runner": "atom-mi355-8gpu.predownload", "test_level": "main", "accuracy_threshold": 0.85, @@ -219,7 +269,7 @@ "model_name": "Qwen3.5-397B-A17B-MXFP4", "model_path": "amd/Qwen3.5-397B-A17B-MXFP4", "extraArgs": "--kv_cache_dtype fp8 -tp 4", - "env_vars": "", + "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1", "runner": "atom-mi355-8gpu.predownload", "test_level": "pr", "accuracy_threshold": 0.835, @@ -231,7 +281,7 @@ "model_name": "Qwen3.5-397B-A17B-MXFP4 MTP", "model_path": "amd/Qwen3.5-397B-A17B-MXFP4", "extraArgs": "--kv_cache_dtype fp8 -tp 4 --method mtp --num-speculative-tokens 3", - "env_vars": "", + "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1", "runner": "atom-mi355-8gpu.predownload", "test_level": "main", "accuracy_threshold": 0.835, @@ -240,27 +290,51 @@ "_baseline_note": "CI baseline=0.8605 (FP8 tp=4, 3-shot completions API, thinking mode active). HF card reports 0.9538 but uses chat API with reasoning_parser" }, { - "model_name": "MiniMax-M2.5", - "model_path": "MiniMaxAI/MiniMax-M2.5", + "model_name": "MiniMax-M2.7", + "model_path": "MiniMaxAI/MiniMax-M2.7", "extraArgs": "--kv_cache_dtype fp8 -tp 2 --trust-remote-code", "env_vars": "", "runner": "atom-mi355-8gpu.predownload", "test_level": "nightly", - "accuracy_threshold": 0.92, - "accuracy_baseline": 0.9401, - "accuracy_baseline_model": "MiniMaxAI/MiniMax-M2.5", - "_baseline_note": "HF: amd/MiniMax-M2.5-MXFP4 card shows baseline=0.9401" + "accuracy_threshold": 0.8872, + "accuracy_baseline": 0.9022, + "accuracy_baseline_model": "MiniMaxAI/MiniMax-M2.7", + "_baseline_note": "ATOM CI measured: 0.9022 (gsm8k 3-shot flexible-extract). Threshold = baseline - 0.015." }, { - "model_name": "MiniMax-M2.5-MXFP4", - "model_path": "amd/MiniMax-M2.5-MXFP4", + "model_name": "MiniMax-M2.7-MXFP4", + "model_path": "amd/MiniMax-M2.7-MXFP4", "extraArgs": "--kv_cache_dtype fp8 --trust-remote-code", "env_vars": "", "runner": "atom-mi355-8gpu.predownload", "test_level": "nightly", - "accuracy_threshold": 0.91, - "accuracy_baseline": 0.9401, - "accuracy_baseline_model": "MiniMaxAI/MiniMax-M2.5", - "_baseline_note": "HF: amd/MiniMax-M2.5-MXFP4 card shows MXFP4=0.9256, baseline=0.9401" + "accuracy_threshold": 0.8872, + "accuracy_baseline": 0.9022, + "accuracy_baseline_model": "MiniMaxAI/MiniMax-M2.7", + "_baseline_note": "ATOM CI measured BF16=0.9022 (gsm8k 3-shot flexible-extract). HF amd/MiniMax-M2.7-MXFP4: MXFP4=91.89, baseline=91.81 (percentage)." + }, + { + "model_name": "MiMo-V2-Flash", + "model_path": "XiaomiMiMo/MiMo-V2-Flash", + "extraArgs": "--kv_cache_dtype fp8 -tp 4 --trust-remote-code --gpu-memory-utilization 0.8", + "env_vars": "", + "runner": "atom-mi355-8gpu.predownload", + "test_level": "nightly", + "accuracy_threshold": 0.81, + "accuracy_baseline": 0.8294, + "accuracy_baseline_model": "XiaomiMiMo/MiMo-V2-Flash", + "_baseline_note": "CI measured GSM8K 5-shot flexible-extract = 0.8294 (±0.0104) under MTP1 greedy decoding (see MiMo-V2-Flash MTP entry); base-only number is identical because MTP only proposes speculative drafts that the target verifies. tp pinned to 4 to match the MTP entry's setup (no separate base measurement)." + }, + { + "model_name": "MiMo-V2-Flash MTP", + "model_path": "XiaomiMiMo/MiMo-V2-Flash", + "extraArgs": "--kv_cache_dtype fp8 -tp 4 --trust-remote-code --gpu-memory-utilization 0.8 --method mtp --num-speculative-tokens 1", + "env_vars": "", + "runner": "atom-mi355-8gpu.predownload", + "test_level": "nightly", + "accuracy_threshold": 0.81, + "accuracy_baseline": 0.8294, + "accuracy_baseline_model": "XiaomiMiMo/MiMo-V2-Flash", + "_baseline_note": "CI measured GSM8K 5-shot flexible-extract = 0.8294 (±0.0104) on FP8 KV + MTP1. tp MUST be 4 and num-speculative-tokens MUST be 1: ATOM only constructs MTP layer 0 (matches vLLM _MIMO_V2_FLASH_NUM_MTP_LAYERS=1); driving more spec tokens would route through layers 1/2 whose KV is never populated and accept rate craters." } -] \ No newline at end of file +] diff --git a/.github/benchmark/oot_benchmark_models.json b/.github/benchmark/oot_benchmark_models.json index d6feac4c6e..63140f7289 100644 --- a/.github/benchmark/oot_benchmark_models.json +++ b/.github/benchmark/oot_benchmark_models.json @@ -1,225 +1,410 @@ -[ - { - "display": "DeepSeek-R1 FP8 TP8", - "source_path": "deepseek-ai/DeepSeek-R1-0528", - "path": "deepseek-ai/DeepSeek-R1-0528", - "prefix": "deepseek-r1-fp8", - "nightly_group": "B", - "extra_args": "--trust-remote-code --tensor-parallel-size 8 --max-num-batched-tokens 16384 --max-model-len 16384", - "bench_args": "", - "runner": "atom-mi355-8gpu-oot-benchmark", - "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4" - }, - { - "display": "DeepSeek-R1 MXFP4 TP8", - "source_path": "amd/DeepSeek-R1-0528-MXFP4-MTP-MoEFP4", - "path": "amd/DeepSeek-R1-0528-MXFP4-MTP-MoEFP4", - "prefix": "deepseek-r1-mxfp4", - "nightly_group": "B", - "extra_args": "--trust-remote-code --tensor-parallel-size 8 --max-num-batched-tokens 16384 --max-model-len 16384", - "bench_args": "", - "runner": "atom-mi355-8gpu-oot-benchmark", - "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4" - }, - { - "display": "DeepSeek-V3.2 FP8 TP8", - "dashboard_model": "DeepSeek-V3.2-FP8", - "source_path": "deepseek-ai/DeepSeek-V3.2", - "path": "deepseek-ai/DeepSeek-V3.2", - "prefix": "deepseek-v3-2-fp8", - "extra_args": "--trust-remote-code --tensor-parallel-size 8 --block-size 1 --max-num-batched-tokens 16384 --max-model-len 16384", - "bench_args": "", - "runner": "atom-mi355-8gpu-oot-benchmark", - "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4" - }, - { - "display": "MiniMax-M2.5 TP2", - "dashboard_model": "MiniMax-M2.5", - "source_path": "MiniMaxAI/MiniMax-M2.5", - "path": "MiniMaxAI/MiniMax-M2.5", - "prefix": "MiniMax-M2.5-tp2", - "extra_args": "--trust-remote-code --kv-cache-dtype fp8 --tensor-parallel-size 2", - "bench_args": "", - "runner": "atom-mi355-8gpu-oot-benchmark", - "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1" - }, - { - "display": "gpt-oss-120b TP1", - "source_path": "openai/gpt-oss-120b", - "path": "openai/gpt-oss-120b", - "prefix": "gpt-oss-120b", - "nightly_group": "C", - "extra_args": "--trust-remote-code --tensor-parallel-size 1 --max-num-batched-tokens 16384 --max-model-len 16384", - "bench_args": "", - "runner": "atom-mi355-8gpu-oot-benchmark", - "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1" - }, - { - "display": "GLM-5.1-FP8 TP8", - "dashboard_model": "GLM-5.1-FP8", - "source_path": "zai-org/GLM-5.1-FP8", - "path": "zai-org/GLM-5.1-FP8", - "prefix": "glm-5-1-fp8-tp8", - "nightly_group": "C", - "extra_args": "--trust-remote-code --tensor-parallel-size 8 --default-chat-template-kwargs '{\"enable_thinking\":false}' --max-num-batched-tokens 16384 --max-model-len 16384", - "bench_args": "", - "runner": "atom-mi355-8gpu-oot-benchmark", - "env_vars": "" - }, - { - "display": "Kimi-K2-Thinking-MXFP4 TP4", - "dashboard_model": "Kimi-K2-Thinking-MXFP4-tp4", - "source_path": "amd/Kimi-K2-Thinking-MXFP4-AttnFP8", - "path": "amd/Kimi-K2-Thinking-MXFP4-AttnFP8", - "prefix": "kimi-k2-thinking-mxfp4-tp4", - "nightly_group": "B", - "extra_args": "--trust-remote-code --tensor-parallel-size 4 --max-num-batched-tokens 16384 --max-model-len 16384", - "bench_args": "", - "runner": "atom-mi355-8gpu-oot-benchmark", - "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4" - }, - { - "display": "Kimi-K2-Thinking-MXFP4 TP8", - "dashboard_model": "Kimi-K2-Thinking-MXFP4-tp8", - "source_path": "amd/Kimi-K2-Thinking-MXFP4-AttnFP8", - "path": "amd/Kimi-K2-Thinking-MXFP4-AttnFP8", - "prefix": "kimi-k2-thinking-mxfp4-tp8", - "nightly_group": "B", - "extra_args": "--trust-remote-code --tensor-parallel-size 8 --max-num-batched-tokens 16384 --max-model-len 16384", - "bench_args": "", - "runner": "atom-mi355-8gpu-oot-benchmark", - "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4" - }, - { - "display": "Kimi-K2.5-MXFP4 TP2", - "dashboard_model": "Kimi-K2.5-MXFP4-tp2", - "source_path": "amd/Kimi-K2.5-MXFP4-AttnFP8", - "path": "amd/Kimi-K2.5-MXFP4-AttnFP8", - "prefix": "kimi-k25-mxfp4-tp2", - "nightly_group": "B", - "extra_args": "--trust-remote-code --tensor-parallel-size 2", - "bench_args": "", - "runner": "atom-mi355-8gpu-oot-benchmark", - "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4" - }, - { - "display": "Kimi-K2.5-MXFP4 TP4", - "dashboard_model": "Kimi-K2.5-MXFP4-tp4", - "source_path": "amd/Kimi-K2.5-MXFP4-AttnFP8", - "path": "amd/Kimi-K2.5-MXFP4-AttnFP8", - "prefix": "kimi-k25-mxfp4-tp4", - "nightly_group": "B", - "extra_args": "--trust-remote-code --tensor-parallel-size 4", - "bench_args": "", - "runner": "atom-mi355-8gpu-oot-benchmark", - "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4" - }, - { - "display": "Kimi-K2.5-MXFP4 TP8", - "dashboard_model": "Kimi-K2.5-MXFP4", - "source_path": "amd/Kimi-K2.5-MXFP4-AttnFP8", - "path": "amd/Kimi-K2.5-MXFP4-AttnFP8", - "prefix": "kimi-k25-mxfp4-tp8", - "nightly_group": "B", - "extra_args": "--trust-remote-code --tensor-parallel-size 8", - "bench_args": "", - "runner": "atom-mi355-8gpu-oot-benchmark", - "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4" - }, - { - "display": "Qwen3.5-397B-A17B-FP8 TP8", - "dashboard_model": "Qwen3.5-397B-A17B-FP8", - "source_path": "Qwen/Qwen3.5-397B-A17B-FP8", - "path": "Qwen/Qwen3.5-397B-A17B-FP8", - "prefix": "qwen3-5-397b-a17b-fp8", - "nightly_group": "A", - "extra_args": "--trust-remote-code --tensor-parallel-size 8 --attention-backend ROCM_AITER_FA --gpu-memory-utilization 0.8 --max-num-batched-tokens 16384 --max-model-len 16384", - "bench_args": "", - "runner": "atom-mi355-8gpu-oot-benchmark", - "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0" - }, - { - "display": "Qwen3.5-397B-A17B-MXFP4 TP4", - "dashboard_model": "Qwen3.5-397B-A17B-MXFP4-tp4", - "source_path": "amd/Qwen3.5-397B-A17B-MXFP4", - "path": "amd/Qwen3.5-397B-A17B-MXFP4", - "prefix": "qwen3-5-397b-a17b-mxfp4", - "nightly_group": "A", - "extra_args": "--trust-remote-code --tensor-parallel-size 4 --attention-backend ROCM_AITER_FA --gpu-memory-utilization 0.8 --max-num-batched-tokens 16384 --max-model-len 16384", - "bench_args": "", - "runner": "atom-mi355-8gpu-oot-benchmark", - "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0" - }, - { - "display": "Qwen3.5-397B-A17B-FP8 TP4", - "dashboard_model": "Qwen3.5-397B-A17B-FP8-tp4", - "source_path": "Qwen/Qwen3.5-397B-A17B-FP8", - "path": "Qwen/Qwen3.5-397B-A17B-FP8", - "prefix": "qwen3-5-397b-a17b-fp8-tp4", - "nightly_group": "A", - "extra_args": "--trust-remote-code --tensor-parallel-size 4 --attention-backend ROCM_AITER_FA --gpu-memory-utilization 0.8 --max-num-batched-tokens 16384 --max-model-len 16384", - "bench_args": "", - "runner": "atom-mi355-8gpu-oot-benchmark", - "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0" - }, - { - "display": "Qwen3.5-397B-A17B TP8", - "dashboard_model": "Qwen3.5-397B-A17B", - "source_path": "Qwen/Qwen3.5-397B-A17B", - "path": "Qwen/Qwen3.5-397B-A17B", - "prefix": "qwen3-5-397b-a17b", - "nightly_group": "A", - "extra_args": "--trust-remote-code --tensor-parallel-size 8 --attention-backend ROCM_AITER_FA --gpu-memory-utilization 0.8 --max-num-batched-tokens 16384 --max-model-len 16384", - "bench_args": "", - "runner": "atom-mi355-8gpu-oot-benchmark", - "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0" - }, - { - "display": "Qwen3.5-397B-A17B TP4", - "dashboard_model": "Qwen3.5-397B-A17B-tp4", - "source_path": "Qwen/Qwen3.5-397B-A17B", - "path": "Qwen/Qwen3.5-397B-A17B", - "prefix": "qwen3-5-397b-a17b-tp4", - "nightly_group": "A", - "extra_args": "--trust-remote-code --tensor-parallel-size 4 --attention-backend ROCM_AITER_FA --gpu-memory-utilization 0.8 --max-num-batched-tokens 16384 --max-model-len 16384", - "bench_args": "", - "runner": "atom-mi355-8gpu-oot-benchmark", - "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0" - }, - { - "display": "Qwen3-Next-80B-A3B-Instruct-FP8 TP1", - "dashboard_model": "Qwen3-Next-80B-A3B-Instruct-FP8-tp1", - "source_path": "Qwen/Qwen3-Next-80B-A3B-Instruct-FP8", - "path": "Qwen/Qwen3-Next-80B-A3B-Instruct-FP8", - "prefix": "qwen3-next-80b-a3b-instruct-fp8-tp1", - "nightly_group": "A", - "extra_args": "--trust-remote-code --tensor-parallel-size 1 --max-num-batched-tokens 32768 --max-model-len 16384", - "bench_args": "", - "runner": "atom-mi355-8gpu-oot-benchmark", - "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0" - }, - { - "display": "Qwen3-Next-80B-A3B-Instruct-FP8 TP2", - "dashboard_model": "Qwen3-Next-80B-A3B-Instruct-FP8-tp2", - "source_path": "Qwen/Qwen3-Next-80B-A3B-Instruct-FP8", - "path": "Qwen/Qwen3-Next-80B-A3B-Instruct-FP8", - "prefix": "qwen3-next-80b-a3b-instruct-fp8-tp2", - "nightly_group": "A", - "extra_args": "--trust-remote-code --tensor-parallel-size 2 --max-num-batched-tokens 32768 --max-model-len 16384", - "bench_args": "", - "runner": "atom-mi355-8gpu-oot-benchmark", - "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0\nATOM_USE_FLYDSL_GDR=1" - }, - { - "display": "Qwen3-Next-80B-A3B-Instruct-FP8 TP4", - "dashboard_model": "Qwen3-Next-80B-A3B-Instruct-FP8-tp4", - "source_path": "Qwen/Qwen3-Next-80B-A3B-Instruct-FP8", - "path": "Qwen/Qwen3-Next-80B-A3B-Instruct-FP8", - "prefix": "qwen3-next-80b-a3b-instruct-fp8-tp4", - "nightly_group": "A", - "extra_args": "--trust-remote-code --tensor-parallel-size 4 --max-num-batched-tokens 32768 --max-model-len 16384", - "bench_args": "", - "runner": "atom-mi355-8gpu-oot-benchmark", - "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0\nATOM_USE_FLYDSL_GDR=1" - } -] +{ + "families": [ + { + "choice_label": "DeepSeek-R1 FP8", + "source_path": "deepseek-ai/DeepSeek-R1-0528", + "path": "deepseek-ai/DeepSeek-R1-0528", + "bench_args": "", + "runner": "atom-mi355-8gpu-oot-benchmark", + "nightly_group": "B", + "variants": [ + { + "tp_size": 8, + "display": "DeepSeek-R1 FP8 TP8 (MET)", + "prefix": "deepseek-r1-fp8-met", + "dashboard_model": "DeepSeek-R1-FP8-tp8", + "extra_args": "--trust-remote-code --tensor-parallel-size 8 --max-num-batched-tokens 16384 --max-model-len 16384", + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4" + } + ] + }, + { + "choice_label": "DeepSeek-R1 MXFP4", + "source_path": "amd/DeepSeek-R1-0528-MXFP4-MTP-MoEFP4", + "path": "amd/DeepSeek-R1-0528-MXFP4-MTP-MoEFP4", + "bench_args": "", + "runner": "atom-mi355-8gpu-oot-benchmark", + "nightly_group": "B", + "variants": [ + { + "tp_size": 8, + "display": "DeepSeek-R1 MXFP4 TP8 (MET)", + "prefix": "deepseek-r1-mxfp4-met", + "dashboard_model": "DeepSeek-R1-MXFP4-tp8", + "extra_args": "--trust-remote-code --tensor-parallel-size 8 --max-num-batched-tokens 16384 --max-model-len 16384", + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4" + } + ] + }, + { + "choice_label": "gpt-oss-120b", + "source_path": "openai/gpt-oss-120b", + "path": "openai/gpt-oss-120b", + "runner": "atom-mi355-8gpu-oot-benchmark", + "nightly_group": "C", + "variants": [ + { + "tp_size": 1, + "display": "gpt-oss-120b TP1 (MET)", + "prefix": "gpt-oss-120b-met", + "dashboard_model": "gpt-oss-120b-tp1", + "bench_args": "", + "extra_args": "--trust-remote-code --tensor-parallel-size 1 --gpu-memory-utilization 0.5 --max-num-batched-tokens 16384 --max-model-len 16384", + "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nGPTOSS_USE_GENERIC_SWIGLU_MXFP4_LAYOUT=1" + }, + { + "tp_size": 1, + "display": "gpt-oss-120b TP1 (AW)", + "dashboard_model": "gpt-oss-120b-aw-tp1", + "prefix": "gpt-oss-120b-aw-tp1", + "bench_args": "", + "extra_args": "--trust-remote-code --tensor-parallel-size 1 --max-num-batched-tokens 16384 --max-model-len 16384", + "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nGPTOSS_USE_GENERIC_SWIGLU_MXFP4_LAYOUT=1" + }, + { + "tp_size": 2, + "display": "gpt-oss-120b TP2 (AW)", + "dashboard_model": "gpt-oss-120b-aw-tp2", + "prefix": "gpt-oss-120b-aw-tp2", + "bench_args": "", + "extra_args": "--trust-remote-code --tensor-parallel-size 2 --gpu-memory-utilization 0.5 --max-num-batched-tokens 16384 --max-model-len 16384", + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nGPTOSS_USE_GENERIC_SWIGLU_MXFP4_LAYOUT=1" + }, + { + "tp_size": 8, + "display": "gpt-oss-120b TP8 (AW)", + "dashboard_model": "gpt-oss-120b-aw-tp8", + "prefix": "gpt-oss-120b-aw-tp8", + "bench_args": "", + "extra_args": "--trust-remote-code --tensor-parallel-size 8 --gpu-memory-utilization 0.5 --max-num-batched-tokens 16384 --max-model-len 16384", + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nGPTOSS_USE_GENERIC_SWIGLU_MXFP4_LAYOUT=1" + } + ] + }, + { + "choice_label": "GLM-5.1-FP8", + "source_path": "zai-org/GLM-5.1-FP8", + "path": "zai-org/GLM-5.1-FP8", + "bench_args": "", + "runner": "atom-mi355-8gpu-oot-benchmark", + "nightly_group": "C", + "variants": [ + { + "tp_size": 8, + "display": "GLM-5.1-FP8 TP8 (OOB)", + "dashboard_model": "GLM-5.1-FP8", + "prefix": "glm-5-1-fp8-tp8-oob", + "excluded_input_output_pairs": [ + "1024x8192" + ], + "extra_args": "--trust-remote-code --tensor-parallel-size 8 --default-chat-template-kwargs '{\"enable_thinking\":false}' --max-num-batched-tokens 16384 --max-model-len 16384", + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4" + } + ] + }, + { + "choice_label": "Kimi-K2-Thinking-MXFP4", + "source_path": "amd/Kimi-K2-Thinking-MXFP4-AttnFP8", + "path": "amd/Kimi-K2-Thinking-MXFP4-AttnFP8", + "bench_args": "", + "runner": "atom-mi355-8gpu-oot-benchmark", + "nightly_group": "B", + "variants": [ + { + "tp_size": 4, + "display": "Kimi-K2-Thinking-MXFP4 TP4 (MET)", + "dashboard_model": "Kimi-K2-Thinking-MXFP4-tp4", + "prefix": "kimi-k2-thinking-mxfp4-tp4-met", + "extra_args": "--trust-remote-code --tensor-parallel-size 4 --max-num-batched-tokens 16384 --max-model-len 16384", + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4" + }, + { + "tp_size": 8, + "display": "Kimi-K2-Thinking-MXFP4 TP8 (MET)", + "dashboard_model": "Kimi-K2-Thinking-MXFP4-tp8", + "prefix": "kimi-k2-thinking-mxfp4-tp8-met", + "extra_args": "--trust-remote-code --tensor-parallel-size 8 --max-num-batched-tokens 16384 --max-model-len 16384", + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4" + } + ] + }, + { + "choice_label": "Kimi-K2.5-MXFP4", + "source_path": "amd/Kimi-K2.5-MXFP4-AttnFP8", + "path": "amd/Kimi-K2.5-MXFP4-AttnFP8", + "runner": "atom-mi355-8gpu-oot-benchmark", + "nightly_group": "B", + "variants": [ + { + "tp_size": 4, + "display": "Kimi-K2.5-MXFP4 TP4 (MET)", + "dashboard_model": "Kimi-K2.5-MXFP4-tp4", + "prefix": "kimi-k25-mxfp4-tp4-met", + "bench_args": "", + "extra_args": "--trust-remote-code --tensor-parallel-size 4 --max-num-batched-tokens 16384 --max-model-len 16384", + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4" + }, + { + "tp_size": 4, + "display": "Kimi-K2.5-MXFP4 TP4 (AW)", + "dashboard_model": "Kimi-K2.5-MXFP4-aw-tp4", + "prefix": "kimi-k2-5-mxfp4-aw-tp4", + "bench_args": "", + "extra_args": "--trust-remote-code --tensor-parallel-size 4 --max-num-batched-tokens 16384 --max-model-len 16384", + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4" + }, + { + "tp_size": 8, + "display": "Kimi-K2.5-MXFP4 TP8 (AW)", + "dashboard_model": "Kimi-K2.5-MXFP4-aw-tp8", + "prefix": "kimi-k2-5-mxfp4-aw-tp8", + "bench_args": "", + "extra_args": "--trust-remote-code --tensor-parallel-size 8 --max-num-batched-tokens 16384 --max-model-len 16384", + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4" + } + ] + }, + { + "choice_label": "Qwen3.5-397B-A17B-FP8", + "source_path": "Qwen/Qwen3.5-397B-A17B-FP8", + "path": "Qwen/Qwen3.5-397B-A17B-FP8", + "bench_args": "", + "runner": "atom-mi355-8gpu-oot-benchmark", + "nightly_group": "A", + "variants": [ + { + "tp_size": 4, + "display": "Qwen3.5-397B-A17B-FP8 TP4 (OOB)", + "dashboard_model": "Qwen3.5-397B-A17B-FP8-tp4", + "prefix": "qwen3-5-397b-a17b-fp8-tp4-oob", + "excluded_input_output_pairs": [ + "1024x8192" + ], + "extra_args": "--trust-remote-code --tensor-parallel-size 4 --attention-backend ROCM_AITER_FA --gpu-memory-utilization 0.8 --max-num-batched-tokens 16384 --max-model-len 16384", + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0\nATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE=0" + }, + { + "tp_size": 8, + "display": "Qwen3.5-397B-A17B-FP8 TP8 (OOB)", + "dashboard_model": "Qwen3.5-397B-A17B-FP8", + "prefix": "qwen3-5-397b-a17b-fp8-oob", + "excluded_input_output_pairs": [ + "1024x8192" + ], + "extra_args": "--trust-remote-code --tensor-parallel-size 8 --attention-backend ROCM_AITER_FA --gpu-memory-utilization 0.8 --max-num-batched-tokens 16384 --max-model-len 16384", + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0\nATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE=0" + } + ] + }, + { + "choice_label": "Qwen3.5-397B-A17B", + "source_path": "Qwen/Qwen3.5-397B-A17B", + "path": "Qwen/Qwen3.5-397B-A17B", + "bench_args": "", + "runner": "atom-mi355-8gpu-oot-benchmark", + "nightly_group": "A", + "variants": [ + { + "tp_size": 8, + "display": "Qwen3.5-397B-A17B TP8 (OOB)", + "dashboard_model": "Qwen3.5-397B-A17B", + "prefix": "qwen3-5-397b-a17b-oob", + "excluded_input_output_pairs": [ + "1024x8192" + ], + "extra_args": "--trust-remote-code --tensor-parallel-size 8 --attention-backend ROCM_AITER_FA --gpu-memory-utilization 0.8 --max-num-batched-tokens 16384 --max-model-len 16384", + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0\nATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE=0" + } + ] + }, + { + "choice_label": "Qwen3-Next-80B-A3B-Instruct-FP8", + "source_path": "Qwen/Qwen3-Next-80B-A3B-Instruct-FP8", + "path": "Qwen/Qwen3-Next-80B-A3B-Instruct-FP8", + "runner": "atom-mi355-8gpu-oot-benchmark", + "nightly_group": "A", + "variants": [ + { + "tp_size": 1, + "display": "Qwen3-Next-80B-A3B-Instruct-FP8 TP1 (MET)", + "dashboard_model": "Qwen3-Next-80B-A3B-Instruct-FP8-tp1", + "prefix": "qwen3-next-80b-a3b-instruct-fp8-tp1-met", + "bench_args": "", + "extra_args": "--trust-remote-code --tensor-parallel-size 1 --max-num-batched-tokens 32768 --max-model-len 16384", + "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0\nATOM_USE_FLYDSL_GDR=0\nATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE=0" + }, + { + "tp_size": 4, + "display": "Qwen3-Next-80B-A3B-Instruct-FP8 TP4 (MET)", + "dashboard_model": "Qwen3-Next-80B-A3B-Instruct-FP8-tp4", + "prefix": "qwen3-next-80b-a3b-instruct-fp8-tp4-met", + "bench_args": "", + "extra_args": "--trust-remote-code --tensor-parallel-size 4 --max-num-batched-tokens 32768 --max-model-len 16384", + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0\nATOM_USE_FLYDSL_GDR=0\nATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE=0" + }, + { + "tp_size": 1, + "display": "Qwen3-Next-80B-A3B-Instruct-FP8 TP1 (AW)", + "dashboard_model": "Qwen3-Next-80B-A3B-Instruct-FP8-aw-tp1", + "prefix": "qwen3-next-80b-a3b-instruct-fp8-aw-tp1", + "bench_args": "", + "extra_args": "--trust-remote-code --tensor-parallel-size 1 --max-num-batched-tokens 32768 --max-model-len 16384", + "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0\nATOM_USE_FLYDSL_GDR=0\nATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE=0" + }, + { + "tp_size": 2, + "display": "Qwen3-Next-80B-A3B-Instruct-FP8 TP2 (AW)", + "dashboard_model": "Qwen3-Next-80B-A3B-Instruct-FP8-aw-tp2", + "prefix": "qwen3-next-80b-a3b-instruct-fp8-aw-tp2", + "bench_args": "", + "extra_args": "--trust-remote-code --tensor-parallel-size 2 --max-num-batched-tokens 32768 --max-model-len 16384", + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0\nATOM_USE_FLYDSL_GDR=0\nATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE=0" + }, + { + "tp_size": 4, + "display": "Qwen3-Next-80B-A3B-Instruct-FP8 TP4 (AW)", + "dashboard_model": "Qwen3-Next-80B-A3B-Instruct-FP8-aw-tp4", + "prefix": "qwen3-next-80b-a3b-instruct-fp8-aw-tp4", + "bench_args": "", + "extra_args": "--trust-remote-code --tensor-parallel-size 4 --max-num-batched-tokens 32768 --max-model-len 16384", + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0\nATOM_USE_FLYDSL_GDR=0\nATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE=0" + }, + { + "tp_size": 1, + "display": "Qwen3-Next-80B-A3B-Instruct-FP8-MTP TP1 (AW)", + "dashboard_model": "Qwen3-Next-80B-A3B-Instruct-FP8-mtp-tp1", + "prefix": "qwen3-next-80b-a3b-instruct-fp8-mtp-tp1-aw", + "bench_args": "", + "extra_args": "--trust-remote-code --tensor-parallel-size 1 --max-num-batched-tokens 32768 --max-model-len 16384 --speculative-config '{\"num_speculative_tokens\":1, \"method\": \"mtp\"}'", + "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0\nATOM_USE_FLYDSL_GDR=0\nATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE=0" + }, + { + "tp_size": 4, + "display": "Qwen3-Next-80B-A3B-Instruct-FP8-MTP TP4 (AW)", + "dashboard_model": "Qwen3-Next-80B-A3B-Instruct-FP8-mtp-tp4", + "prefix": "qwen3-next-80b-a3b-instruct-fp8-mtp-tp4-aw", + "bench_args": "", + "extra_args": "--trust-remote-code --tensor-parallel-size 4 --max-num-batched-tokens 32768 --max-model-len 16384 --speculative-config '{\"num_speculative_tokens\":1, \"method\": \"mtp\"}'", + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0\nATOM_USE_FLYDSL_GDR=0\nATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE=0" + } + ] + }, + { + "choice_label": "DeepSeek-V3.2 FP8", + "source_path": "deepseek-ai/DeepSeek-V3.2", + "path": "deepseek-ai/DeepSeek-V3.2", + "runner": "atom-mi355-8gpu-oot-benchmark", + "nightly_group": "B", + "variants": [ + { + "tp_size": 4, + "display": "DeepSeek-V3.2 FP8 TP4 (AW)", + "dashboard_model": "DeepSeek-V3.2-FP8-aw-tp4", + "prefix": "deepseek-v3-2-fp8-aw-tp4", + "bench_args": "", + "extra_args": "--trust-remote-code --tensor-parallel-size 4 --kv-cache-dtype auto --max-num-batched-tokens 16384 --max-model-len 16384", + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4" + }, + { + "tp_size": 8, + "display": "DeepSeek-V3.2 FP8 TP8 (AW)", + "dashboard_model": "DeepSeek-V3.2-FP8-aw-tp8", + "prefix": "deepseek-v3-2-fp8-aw-tp8", + "bench_args": "", + "extra_args": "--trust-remote-code --tensor-parallel-size 8 --max-num-batched-tokens 16384 --max-model-len 16384", + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4" + } + ] + }, + { + "choice_label": "GLM-4.7-FP8", + "source_path": "zai-org/GLM-4.7-FP8", + "path": "zai-org/GLM-4.7-FP8", + "runner": "atom-mi355-8gpu-oot-benchmark", + "nightly_group": "B", + "variants": [ + { + "tp_size": 4, + "display": "GLM-4.7-FP8 TP4 (AW)", + "dashboard_model": "GLM-4.7-FP8-aw-tp4", + "prefix": "glm-4-7-fp8-aw-tp4", + "bench_args": "", + "extra_args": "--trust-remote-code --tensor-parallel-size 4 --max-num-batched-tokens 16384 --max-model-len 16384", + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4" + }, + { + "tp_size": 8, + "display": "GLM-4.7-FP8 TP8 (AW)", + "dashboard_model": "GLM-4.7-FP8-aw-tp8", + "prefix": "glm-4-7-fp8-aw-tp8", + "bench_args": "", + "extra_args": "--trust-remote-code --tensor-parallel-size 8 --max-num-batched-tokens 16384 --max-model-len 16384", + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4" + } + ] + }, + { + "choice_label": "GLM-4.7-FP8 MTP", + "source_path": "zai-org/GLM-4.7-FP8", + "path": "zai-org/GLM-4.7-FP8", + "runner": "atom-mi355-8gpu-oot-benchmark", + "nightly_group": "B", + "variants": [ + { + "tp_size": 4, + "display": "GLM-4.7-FP8 MTP TP4 (AW)", + "dashboard_model": "GLM-4.7-FP8-mtp-aw-tp4", + "prefix": "glm-4-7-fp8-mtp-aw-tp4", + "bench_args": "", + "extra_args": "--trust-remote-code --tensor-parallel-size 4 --speculative-config.method mtp --speculative-config.num_speculative_tokens 1", + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4" + }, + { + "tp_size": 8, + "display": "GLM-4.7-FP8 MTP TP8 (AW)", + "dashboard_model": "GLM-4.7-FP8-mtp-aw-tp8", + "prefix": "glm-4-7-fp8-mtp-aw-tp8", + "bench_args": "", + "extra_args": "--trust-remote-code --tensor-parallel-size 8 --speculative-config.method mtp --speculative-config.num_speculative_tokens 1", + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4" + } + ] + }, + { + "choice_label": "MiniMax-M2.5", + "source_path": "MiniMaxAI/MiniMax-M2.5", + "path": "MiniMaxAI/MiniMax-M2.5", + "runner": "atom-mi355-8gpu-oot-benchmark", + "nightly_group": "B", + "variants": [ + { + "tp_size": 2, + "display": "MiniMax-M2.5 TP2 (AW)", + "dashboard_model": "MiniMax-M2.5-aw-tp2", + "prefix": "minimax-m2-5-aw-tp2", + "bench_args": "", + "extra_args": "--trust-remote-code --tensor-parallel-size 2 --kv-cache-dtype fp8 --max-num-batched-tokens 16384 --max-model-len 16384", + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1" + }, + { + "tp_size": 4, + "display": "MiniMax-M2.5 TP4 (AW)", + "dashboard_model": "MiniMax-M2.5-aw-tp4", + "prefix": "minimax-m2-5-aw-tp4", + "bench_args": "", + "extra_args": "--trust-remote-code --tensor-parallel-size 4 --kv-cache-dtype fp8 --max-num-batched-tokens 16384 --max-model-len 16384", + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1" + }, + { + "tp_size": 8, + "display": "MiniMax-M2.5 TP8 (AW)", + "dashboard_model": "MiniMax-M2.5-aw-tp8", + "prefix": "minimax-m2-5-aw-tp8", + "bench_args": "", + "extra_args": "--trust-remote-code --tensor-parallel-size 8 --kv-cache-dtype fp8 --max-num-batched-tokens 16384 --max-model-len 16384", + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1" + } + ] + } + ] +} diff --git a/.github/benchmark/oot_models_accuracy.json b/.github/benchmark/oot_models_accuracy.json index 4a259549a4..2e7042077b 100644 --- a/.github/benchmark/oot_models_accuracy.json +++ b/.github/benchmark/oot_models_accuracy.json @@ -3,7 +3,7 @@ "model_name": "Qwen3-235B-A22B-Instruct-2507-FP8 TP8+EP8", "model_path": "Qwen/Qwen3-235B-A22B-Instruct-2507-FP8", "extraArgs": "--tensor-parallel-size 8 --enable-expert-parallel", - "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1", + "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE=0", "runner": "linux-atom-mi35x-8", "test_level": "nightly", "accuracy_threshold": 0.87, @@ -14,7 +14,7 @@ "model_name": "Qwen3-Next-80B-A3B-Instruct-FP8 TP4", "model_path": "Qwen/Qwen3-Next-80B-A3B-Instruct-FP8", "extraArgs": "--tensor-parallel-size 4 --attention-backend ROCM_AITER_FA", - "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0", + "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0\nATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE=0", "runner": "linux-atom-mi35x-4", "test_level": "nightly", "accuracy_threshold": 0.76, @@ -25,7 +25,7 @@ "model_name": "Qwen3.5-397B-A17B-FP8 TP8", "model_path": "Qwen/Qwen3.5-397B-A17B-FP8", "extraArgs": "--tensor-parallel-size 8 --attention-backend ROCM_AITER_FA", - "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0", + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0\nATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE=0", "runner": "linux-atom-mi35x-8", "test_level": "nightly", "accuracy_threshold": 0.83, @@ -36,7 +36,7 @@ "model_name": "Qwen3.5-397B-A17B TP8", "model_path": "Qwen/Qwen3.5-397B-A17B", "extraArgs": "--tensor-parallel-size 8 --attention-backend ROCM_AITER_FA", - "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0", + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0\nATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE=0", "runner": "linux-atom-mi35x-8", "test_level": "nightly", "accuracy_threshold": 0.83, @@ -47,7 +47,7 @@ "model_name": "Qwen3.5-397B-A17B-MXFP4 TP4", "model_path": "amd/Qwen3.5-397B-A17B-MXFP4", "extraArgs": "--tensor-parallel-size 4 --attention-backend ROCM_AITER_FA", - "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0", + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0\nATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE=0", "runner": "linux-atom-mi35x-4", "test_level": "nightly", "accuracy_threshold": 0.82, @@ -55,6 +55,18 @@ "accuracy_baseline_model": "Qwen/Qwen3-235B-A22B-Instruct-2507", "_baseline_note": "Using Qwen3-235B baseline as proxy; needs CI measurement for Qwen3.5 specific baseline" }, + { + "model_name": "Qwen3-Next-80B-A3B-Instruct-FP8-MTP TP4", + "model_path": "Qwen/Qwen3-Next-80B-A3B-Instruct-FP8", + "extraArgs": "--tensor-parallel-size 4 --speculative-config '{\"num_speculative_tokens\":1, \"method\": \"mtp\"}'", + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0\nATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE=0", + "runner": "linux-atom-mi35x-4", + "test_level": "nightly", + "accuracy_threshold": 0.8, + "accuracy_baseline": 0.81, + "accuracy_baseline_model": "Qwen/Qwen3-Next-80B-A3B-Instruct-FP8", + "_baseline_note": "Qwen3-Next-80B-A3B-Instruct-FP8 baseline with TP4 (no MTP) as proxy; needs CI measurement for MTP-specific baseline" + }, { "model_name": "Llama-3.1-8B-Instruct TP1", "model_path": "meta-llama/Llama-3.1-8B-Instruct", @@ -103,16 +115,16 @@ "_baseline_note": "Reference value from recipes/atom_vllm/Kimi-K2.5.md" }, { - "model_name": "MiniMax-M2.5 TP2", - "model_path": "MiniMaxAI/MiniMax-M2.5", + "model_name": "MiniMax-M2.7 TP2", + "model_path": "MiniMaxAI/MiniMax-M2.7", "extraArgs": "--kv_cache_dtype fp8 -tp 2 --trust-remote-code", "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1", "runner": "linux-atom-mi35x-8", "test_level": "nightly", - "accuracy_threshold": 0.92, - "accuracy_baseline": 0.9401, - "accuracy_baseline_model": "MiniMaxAI/MiniMax-M2.5", - "_baseline_note": "HF: amd/MiniMax-M2.5-MXFP4 card shows baseline=0.9401" + "accuracy_threshold": 0.8872, + "accuracy_baseline": 0.9022, + "accuracy_baseline_model": "MiniMaxAI/MiniMax-M2.7", + "_baseline_note": "ATOM CI measured: 0.9022 (gsm8k 3-shot flexible-extract). Threshold = baseline - 0.015." }, { "model_name": "DeepSeek-R1-FP8 TP8", @@ -139,7 +151,7 @@ { "model_name": "DeepSeek-V3.2-FP8 TP8", "model_path": "deepseek-ai/DeepSeek-V3.2", - "extraArgs": "--tensor-parallel-size 8 --block-size 1", + "extraArgs": "--tensor-parallel-size 8", "env_vars": "", "runner": "linux-atom-mi35x-8", "test_level": "nightly", @@ -152,24 +164,37 @@ "model_name": "gpt-oss-120b TP1", "model_path": "openai/gpt-oss-120b", "extraArgs": "--tensor-parallel-size 1", + "client_command": "lm_eval --model local-chat-completions --apply_chat_template --model_args model=${MODEL_PATH},base_url=http://127.0.0.1:${VLLM_PORT}/v1/chat/completions,num_concurrent=65,max_retries=1,max_gen_toks=2048,tokenized_requests=False,trust_remote_code=True --tasks gsm8k --num_fewshot ${LM_EVAL_NUM_FEWSHOT} --output_path ${OUTPUT_PATH}", "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1", "runner": "linux-atom-mi35x-1", "test_level": "nightly", - "accuracy_threshold": 0.38, - "accuracy_baseline": 0.38, + "accuracy_threshold": 0.88, + "accuracy_baseline": 0.9, "accuracy_baseline_model": "openai/gpt-oss-120b" }, { "model_name": "gpt-oss-120b TP2", "model_path": "openai/gpt-oss-120b", "extraArgs": "--tensor-parallel-size 2", - "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1", + "client_command": "lm_eval --model local-chat-completions --apply_chat_template --model_args model=${MODEL_PATH},base_url=http://127.0.0.1:${VLLM_PORT}/v1/chat/completions,num_concurrent=65,max_retries=1,max_gen_toks=2048,tokenized_requests=False,trust_remote_code=True --tasks gsm8k --num_fewshot ${LM_EVAL_NUM_FEWSHOT} --output_path ${OUTPUT_PATH}", + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1", "runner": "linux-atom-mi35x-4", "test_level": "nightly", - "accuracy_threshold": 0.38, - "accuracy_baseline": 0.38, + "accuracy_threshold": 0.88, + "accuracy_baseline": 0.9, "accuracy_baseline_model": "openai/gpt-oss-120b" }, + { + "model_name": "GLM-4.7-FP8 TP8", + "model_path": "zai-org/GLM-4.7-FP8", + "extraArgs": "--tensor-parallel-size 8 --default-chat-template-kwargs '{\"enable_thinking\":false}'", + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1", + "runner": "linux-atom-mi35x-8", + "test_level": "nightly", + "accuracy_threshold": 0.92, + "accuracy_baseline": 0.9386, + "accuracy_baseline_model": "zai-org/GLM-4.7-FP8" + }, { "model_name": "GLM-5.1-FP8 TP8", "model_path": "zai-org/GLM-5.1-FP8", diff --git a/.github/benchmark/sglang_benchmark_models.json b/.github/benchmark/sglang_benchmark_models.json index 8ddee78d0d..d293b64110 100644 --- a/.github/benchmark/sglang_benchmark_models.json +++ b/.github/benchmark/sglang_benchmark_models.json @@ -6,7 +6,8 @@ "prefix": "deepseek-r1-fp8-tp8", "extra_args": "--trust-remote-code --tensor-parallel-size 8", "bench_args": "", - "runner": "atom-mi355-8gpu-oot-benchmark", + "runner": "atom-mi355-8gpu-aac-runner", + "nightly_group": "A", "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nSGLANG_AITER_FP8_PREFILL_ATTN=0\nSGLANG_USE_AITER=1\nATOM_ENABLE_DS_QKNORM_QUANT_FUSION=1" }, { @@ -17,7 +18,8 @@ "prefix": "deepseek-r1-fp8-tp4", "extra_args": "--trust-remote-code --tensor-parallel-size 4", "bench_args": "", - "runner": "atom-mi355-8gpu-oot-benchmark", + "runner": "atom-mi355-8gpu-aac-runner", + "nightly_group": "A", "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nSGLANG_AITER_FP8_PREFILL_ATTN=0\nSGLANG_USE_AITER=1\nATOM_ENABLE_DS_QKNORM_QUANT_FUSION=1" }, { @@ -27,7 +29,8 @@ "prefix": "deepseek-r1-fp4-tp8", "extra_args": "--trust-remote-code --tensor-parallel-size 8", "bench_args": "", - "runner": "atom-mi355-8gpu-oot-benchmark", + "runner": "atom-mi355-8gpu-aac-runner", + "nightly_group": "A", "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nSGLANG_AITER_FP8_PREFILL_ATTN=0\nSGLANG_USE_AITER=1\nATOM_ENABLE_DS_QKNORM_QUANT_FUSION=1" }, { @@ -38,18 +41,44 @@ "prefix": "deepseek-r1-fp4-tp4", "extra_args": "--trust-remote-code --tensor-parallel-size 4", "bench_args": "", - "runner": "atom-mi355-8gpu-oot-benchmark", + "runner": "atom-mi355-8gpu-aac-runner", + "nightly_group": "A", "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nSGLANG_AITER_FP8_PREFILL_ATTN=0\nSGLANG_USE_AITER=1\nATOM_ENABLE_DS_QKNORM_QUANT_FUSION=1" }, { - "display": "DeepSeek-R1-0528-MXFP4 FP4 TP8 EP8", - "dashboard_model": "DeepSeek-R1-0528-MXFP4-tp8-ep8", - "source_path": "amd/DeepSeek-R1-0528-MXFP4-MTP-MoEFP4", - "path": "amd/DeepSeek-R1-0528-MXFP4-MTP-MoEFP4", - "prefix": "deepseek-r1-fp4-tp8-ep8", - "extra_args": "--trust-remote-code --tensor-parallel-size 8 --expert-parallel-size 8", - "bench_args": "", - "runner": "atom-mi355-8gpu-oot-benchmark", - "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nSGLANG_AITER_FP8_PREFILL_ATTN=0\nSGLANG_USE_AITER=1\nATOM_ENABLE_DS_QKNORM_QUANT_FUSION=1" + "display": "DeepSeek-R1-0528-MXFP4 FP4 TP8 EP8", + "dashboard_model": "DeepSeek-R1-0528-MXFP4-tp8-ep8", + "source_path": "amd/DeepSeek-R1-0528-MXFP4-MTP-MoEFP4", + "path": "amd/DeepSeek-R1-0528-MXFP4-MTP-MoEFP4", + "prefix": "deepseek-r1-fp4-tp8-ep8", + "extra_args": "--trust-remote-code --tensor-parallel-size 8 --expert-parallel-size 8", + "bench_args": "", + "runner": "atom-mi355-8gpu-aac-runner", + "nightly_group": "A", + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nSGLANG_AITER_FP8_PREFILL_ATTN=0\nSGLANG_USE_AITER=1\nATOM_ENABLE_DS_QKNORM_QUANT_FUSION=1" + }, + { + "display": "Qwen3.5-397B-A17B-FP8 TP4", + "dashboard_model": "Qwen3.5-397B-A17B-FP8-tp4", + "source_path": "Qwen/Qwen3.5-397B-A17B-FP8", + "path": "Qwen/Qwen3.5-397B-A17B-FP8", + "prefix": "qwen3-5-397b-a17b-fp8-tp4", + "extra_args": "--tensor-parallel-size 4 --mem-fraction-static 0.9 --reasoning-parser qwen3 --disable-radix-cache", + "bench_args": "", + "runner": "atom-mi355-8gpu-aac-runner", + "nightly_group": "B", + "env_vars": "SGLANG_DEFAULT_SERVER_ARGS=\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=0" + }, + { + "display": "Qwen3.5-397B-A17B-FP8 TP8", + "dashboard_model": "Qwen3.5-397B-A17B-FP8", + "source_path": "Qwen/Qwen3.5-397B-A17B-FP8", + "path": "Qwen/Qwen3.5-397B-A17B-FP8", + "prefix": "qwen3-5-397b-a17b-fp8-tp8", + "extra_args": "--tensor-parallel-size 8 --mem-fraction-static 0.9 --reasoning-parser qwen3 --disable-radix-cache", + "bench_args": "", + "runner": "atom-mi355-8gpu-aac-runner", + "nightly_group": "B", + "env_vars": "SGLANG_DEFAULT_SERVER_ARGS=\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=0" } ] diff --git a/.github/benchmark/sglang_models_accuracy.json b/.github/benchmark/sglang_models_accuracy.json index 2580c133d3..e2e8548b7c 100644 --- a/.github/benchmark/sglang_models_accuracy.json +++ b/.github/benchmark/sglang_models_accuracy.json @@ -3,12 +3,96 @@ "model_name": "DeepSeek-R1-FP8 TP4", "model_path": "deepseek-ai/DeepSeek-R1-0528", "extraArgs": "--tensor-parallel-size 4", - "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nSGLANG_AITER_FP8_PREFILL_ATTN=0\nSGLANG_USE_AITER=1\nATOM_ENABLE_DS_QKNORM_QUANT_FUSION=1", + "env_vars": "SGLANG_AITER_FP8_PREFILL_ATTN=0\nSGLANG_USE_AITER=1\nATOM_ENABLE_DS_QKNORM_QUANT_FUSION=1", "runner": "linux-atom-mi35x-4", "test_level": "nightly", - "accuracy_threshold": 0.92, + "accuracy_threshold": 0.91, "accuracy_baseline": null, "accuracy_baseline_model": "deepseek-ai/DeepSeek-R1-0528", "_baseline_note": "Threshold aligned with the SGLANG accuracy validation workflow target for gsm8k." + }, + { + "model_name": "Qwen3.5-35B-A3B-FP8 TP2", + "model_path": "Qwen/Qwen3.5-35B-A3B-FP8", + "extraArgs": "--tensor-parallel-size 2 --mem-fraction-static 0.9 --reasoning-parser qwen3 --disable-radix-cache", + "env_vars": "SGLANG_DEFAULT_SERVER_ARGS=\nSGLANG_EXTERNAL_MODEL_PACKAGE=atom.plugin.sglang.models\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=0", + "runner": "linux-atom-mi35x-4", + "test_level": "nightly", + "accuracy_threshold": 0.76, + "accuracy_baseline": null, + "accuracy_baseline_model": "Qwen/Qwen3.5-35B-A3B-FP8", + "_baseline_note": "Threshold aligned with the SGLANG accuracy validation workflow target for gsm8k." + }, + { + "model_name": "Qwen3.5-35B-A3B TP2", + "model_path": "Qwen/Qwen3.5-35B-A3B", + "extraArgs": "--tensor-parallel-size 2 --mem-fraction-static 0.9 --reasoning-parser qwen3 --disable-radix-cache", + "env_vars": "SGLANG_DEFAULT_SERVER_ARGS=\nSGLANG_EXTERNAL_MODEL_PACKAGE=atom.plugin.sglang.models\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=0", + "runner": "linux-atom-mi35x-4", + "test_level": "nightly", + "accuracy_threshold": 0.83, + "accuracy_baseline": null, + "accuracy_baseline_model": "Qwen/Qwen3.5-35B-A3B", + "_baseline_note": "Threshold aligned with the SGLANG accuracy validation workflow target for gsm8k." + }, + { + "model_name": "Qwen3.5-397B-A17B-FP8 TP4", + "model_path": "Qwen/Qwen3.5-397B-A17B-FP8", + "extraArgs": "--tensor-parallel-size 4 --mem-fraction-static 0.9 --reasoning-parser qwen3 --disable-radix-cache", + "env_vars": "SGLANG_DEFAULT_SERVER_ARGS=\nSGLANG_EXTERNAL_MODEL_PACKAGE=atom.plugin.sglang.models\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=0", + "runner": "linux-atom-mi35x-4", + "test_level": "nightly", + "accuracy_threshold": 0.83, + "accuracy_baseline": null, + "accuracy_baseline_model": "Qwen/Qwen3.5-397B-A17B-FP8", + "_baseline_note": "Threshold aligned with the SGLANG accuracy validation workflow target for gsm8k." + }, + { + "model_name": "Qwen3.5-397B-A17B-FP8 TP8", + "model_path": "Qwen/Qwen3.5-397B-A17B-FP8", + "extraArgs": "--tensor-parallel-size 8 --mem-fraction-static 0.9 --reasoning-parser qwen3 --disable-radix-cache", + "env_vars": "SGLANG_DEFAULT_SERVER_ARGS=\nSGLANG_EXTERNAL_MODEL_PACKAGE=atom.plugin.sglang.models\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=0", + "runner": "linux-atom-mi35x-8", + "test_level": "nightly", + "accuracy_threshold": 0.83, + "accuracy_baseline": null, + "accuracy_baseline_model": "Qwen/Qwen3.5-397B-A17B-FP8", + "_baseline_note": "Threshold aligned with the SGLANG accuracy validation workflow target for gsm8k." + }, + { + "model_name": "DeepSeek-R1-FP8 TP8", + "model_path": "deepseek-ai/DeepSeek-R1-0528", + "extraArgs": "--tensor-parallel-size 8", + "env_vars": "SGLANG_AITER_FP8_PREFILL_ATTN=0\nSGLANG_USE_AITER=1\nATOM_ENABLE_DS_QKNORM_QUANT_FUSION=1", + "runner": "linux-atom-mi35x-8", + "test_level": "nightly", + "accuracy_threshold": 0.93, + "accuracy_baseline": null, + "accuracy_baseline_model": "deepseek-ai/DeepSeek-R1-0528", + "_baseline_note": "Threshold aligned with the SGLANG accuracy validation workflow target for gsm8k." + }, + { + "model_name": "DeepSeek-R1-FP4 TP4", + "model_path": "amd/DeepSeek-R1-0528-MXFP4-MTP-MoEFP4", + "extraArgs": "--tensor-parallel-size 4", + "env_vars": "SGLANG_AITER_FP8_PREFILL_ATTN=0\nSGLANG_USE_AITER=1\nATOM_ENABLE_DS_QKNORM_QUANT_FUSION=1", + "runner": "linux-atom-mi35x-4", + "test_level": "nightly", + "accuracy_threshold": 0.91, + "accuracy_baseline": null, + "accuracy_baseline_model": "amd/DeepSeek-R1-0528-MXFP4-MTP-MoEFP4", + "_baseline_note": "Threshold aligned with the SGLANG accuracy validation workflow target for gsm8k." + }, + { + "model_name": "DeepSeek-R1-FP4 TP8", + "model_path": "amd/DeepSeek-R1-0528-MXFP4-MTP-MoEFP4", + "extraArgs": "--tensor-parallel-size 8", + "env_vars": "SGLANG_AITER_FP8_PREFILL_ATTN=0\nSGLANG_USE_AITER=1\nATOM_ENABLE_DS_QKNORM_QUANT_FUSION=1", + "runner": "linux-atom-mi35x-8", + "test_level": "nightly", + "accuracy_threshold": 0.93, + "accuracy_baseline": null, + "accuracy_baseline_model": "amd/DeepSeek-R1-0528-MXFP4-MTP-MoEFP4", + "_baseline_note": "Threshold aligned with the SGLANG accuracy validation workflow target for gsm8k." } ] diff --git a/.github/runner-config.yml b/.github/runner-config.yml index e0287eccf3..ed11880566 100644 --- a/.github/runner-config.yml +++ b/.github/runner-config.yml @@ -6,6 +6,10 @@ runners: gpu_arch: MI355 gpu_count: 8 + atom-mi355-8gpu-aac-runner: + gpu_arch: MI355 + gpu_count: 8 + linux-atom-mi355-1: gpu_arch: MI355 gpu_count: 1 diff --git a/.github/scripts/atom_oot_test.sh b/.github/scripts/atom_oot_test.sh index fe2984b871..9546582342 100644 --- a/.github/scripts/atom_oot_test.sh +++ b/.github/scripts/atom_oot_test.sh @@ -49,6 +49,7 @@ KEEP_SERVER_ALIVE_ON_EXIT=${KEEP_SERVER_ALIVE_ON_EXIT:-0} EXPLICIT_MODEL_NAME=${OOT_MODEL_NAME:-} EXPLICIT_MODEL_PATH=${OOT_MODEL_PATH:-} EXPLICIT_EXTRA_ARGS=${OOT_EXTRA_ARGS:-} +EXPLICIT_CLIENT_COMMAND=${OOT_CLIENT_COMMAND:-} OOT_DOCKER_IMAGE=${OOT_DOCKER_IMAGE:-} LM_EVAL_NUM_FEWSHOT=${LM_EVAL_NUM_FEWSHOT:-3} LAST_VLLM_LOG_LINE=0 @@ -64,7 +65,7 @@ if [[ -n "${EXPLICIT_MODEL_NAME}" || -n "${EXPLICIT_MODEL_PATH}" || -n "${EXPLIC echo "OOT_MODEL_NAME and OOT_MODEL_PATH must both be set when using explicit model overrides." exit 2 fi - ACTIVE_MODELS=("${EXPLICIT_MODEL_NAME}|${EXPLICIT_MODEL_PATH}|${EXPLICIT_EXTRA_ARGS}") + ACTIVE_MODELS=("${EXPLICIT_MODEL_NAME}|${EXPLICIT_MODEL_PATH}|${EXPLICIT_EXTRA_ARGS}|${EXPLICIT_CLIENT_COMMAND}") else echo "${MODE} mode requires OOT_MODEL_NAME and OOT_MODEL_PATH env vars from the workflow." exit 2 @@ -197,7 +198,6 @@ PY --trust-remote-code \ --kv-cache-dtype fp8 \ "${extra_arg_array[@]}" \ - --gpu-memory-utilization 0.9 \ --no-enable-prefix-caching \ > "${VLLM_LOG_FILE}" 2>&1 & echo $! > "${VLLM_PID_FILE}" @@ -210,6 +210,7 @@ accuracy_one_model() { local model_name="$1" local model_path="$2" local extra_args="$3" + local client_command="${4:-}" local flat_result_file="" local resolved_model_path @@ -231,11 +232,71 @@ accuracy_one_model() { echo "Model name: ${model_name}" echo "Few-shot count: ${LM_EVAL_NUM_FEWSHOT}" - lm_eval --model local-completions \ - --model_args model="${resolved_model_path}",base_url="http://127.0.0.1:${VLLM_PORT}/v1/completions",num_concurrent=65,max_retries=1,tokenized_requests=False,trust_remote_code=True \ - --tasks gsm8k \ - --num_fewshot "${LM_EVAL_NUM_FEWSHOT}" \ - --output_path "${output_path}" 2>&1 | tee -a "${ACCURACY_LOG_FILE}" + if [[ "${client_command}" == "null" ]]; then + client_command="" + fi + + if [[ -n "${client_command}" ]]; then + local -a client_command_args=() + while IFS= read -r -d '' token; do + client_command_args+=("${token}") + done < <( + CLIENT_COMMAND="${client_command}" \ + MODEL_PATH_VALUE="${resolved_model_path}" \ + OUTPUT_PATH_VALUE="${output_path}" \ + LM_EVAL_NUM_FEWSHOT_VALUE="${LM_EVAL_NUM_FEWSHOT}" \ + VLLM_PORT_VALUE="${VLLM_PORT}" \ + python3 - <<'PY' +import os +import shlex +import sys + +client_command = os.environ["CLIENT_COMMAND"] +replacements = { + "${MODEL_PATH}": os.environ["MODEL_PATH_VALUE"], + "$MODEL_PATH": os.environ["MODEL_PATH_VALUE"], + "${OUTPUT_PATH}": os.environ["OUTPUT_PATH_VALUE"], + "$OUTPUT_PATH": os.environ["OUTPUT_PATH_VALUE"], + "${LM_EVAL_NUM_FEWSHOT}": os.environ["LM_EVAL_NUM_FEWSHOT_VALUE"], + "$LM_EVAL_NUM_FEWSHOT": os.environ["LM_EVAL_NUM_FEWSHOT_VALUE"], + "${VLLM_PORT}": os.environ["VLLM_PORT_VALUE"], + "$VLLM_PORT": os.environ["VLLM_PORT_VALUE"], +} +for src, dst in replacements.items(): + client_command = client_command.replace(src, dst) + +for token in shlex.split(client_command): + sys.stdout.write(token) + sys.stdout.write("\0") +PY + ) + + if [[ ${#client_command_args[@]} -eq 0 ]]; then + echo "ERROR: client_command is set but empty after parsing." + return 2 + fi + + for arg in "${client_command_args[@]}"; do + if [[ "${arg}" =~ \$\{[A-Z0-9_]+\} ]] || [[ "${arg}" =~ \$[A-Z_][A-Z0-9_]* ]]; then + echo "ERROR: client_command contains unresolved placeholder after expansion: ${arg}" + return 2 + fi + done + + echo "Using custom lm-eval command from client_command: ${client_command}" + "${client_command_args[@]}" 2>&1 | tee -a "${ACCURACY_LOG_FILE}" + else + echo "Using default lm-eval command." + local lm_args=( + --model_args + model="${resolved_model_path}",base_url="http://127.0.0.1:${VLLM_PORT}/v1/completions",num_concurrent=65,max_retries=1,tokenized_requests=False,trust_remote_code=True + ) + lm_eval --model local-completions \ + "${lm_args[@]}" \ + --tasks gsm8k \ + --num_fewshot "${LM_EVAL_NUM_FEWSHOT}" \ + --output_path "${output_path}" 2>&1 | tee -a "${ACCURACY_LOG_FILE}" + fi # lm-eval output layout differs across versions: output_path may be a file # or a directory containing one/more JSON files. Follow native CI style: @@ -328,7 +389,7 @@ run_for_models() { local matched=0 for entry in "${ACTIVE_MODELS[@]}"; do - IFS='|' read -r model_name model_path extra_args <<< "${entry}" + IFS='|' read -r model_name model_path extra_args client_command <<< "${entry}" if [[ -n "${SELECTED_MODEL}" && "${SELECTED_MODEL}" != "${model_name}" ]]; then continue @@ -342,7 +403,7 @@ run_for_models() { # accuracy mode: launch + evaluate each selected model, then stop server. launch_one_model "${model_name}" "${model_path}" "${extra_args}" - accuracy_one_model "${model_name}" "${model_path}" "${extra_args}" + accuracy_one_model "${model_name}" "${model_path}" "${extra_args}" "${client_command}" stop_server done diff --git a/.github/scripts/atom_sglang_test.sh b/.github/scripts/atom_sglang_test.sh index 1940e4418f..1f1cbeb338 100644 --- a/.github/scripts/atom_sglang_test.sh +++ b/.github/scripts/atom_sglang_test.sh @@ -12,6 +12,7 @@ set -euo pipefail # Optional environment variables: # SGLANG_EXTRA_ARGS # SGLANG_ENV_VARS +# SGLANG_DEFAULT_SERVER_ARGS # SGLANG_PORT # SGLANG_HOST # MAX_WAIT_RETRIES @@ -146,11 +147,6 @@ launch_server() { local resolved_model_path resolved_model_path=$(resolve_model_path "${MODEL_PATH}") - local -a extra_arg_array=() - if [[ -n "${MODEL_EXTRA_ARGS}" ]]; then - read -r -a extra_arg_array <<< "${MODEL_EXTRA_ARGS}" - fi - prepare_runtime_paths export AITER_QUICK_REDUCE_QUANTIZATION="${AITER_QUICK_REDUCE_QUANTIZATION:-INT4}" @@ -168,6 +164,19 @@ launch_server() { done <<< "$(printf '%b' "${MODEL_ENV_VARS}")" fi + local default_server_args + default_server_args=${SGLANG_DEFAULT_SERVER_ARGS---trust-remote-code --kv-cache-dtype fp8_e4m3 --mem-fraction-static 0.8 --page-size 1 --disable-radix-cache} + + local -a default_arg_array=() + if [[ -n "${default_server_args}" ]]; then + read -r -a default_arg_array <<< "${default_server_args}" + fi + + local -a extra_arg_array=() + if [[ -n "${MODEL_EXTRA_ARGS}" ]]; then + read -r -a extra_arg_array <<< "${MODEL_EXTRA_ARGS}" + fi + rm -rf /root/.cache rm -f "${SGLANG_PID_FILE}" "${SGLANG_LOG_FILE}" || true @@ -182,11 +191,7 @@ launch_server() { --model-path "${resolved_model_path}" \ --host "${SGLANG_HOST}" \ --port "${SGLANG_PORT}" \ - --trust-remote-code \ - --kv-cache-dtype fp8_e4m3 \ - --mem-fraction-static 0.8 \ - --page-size 1 \ - --disable-radix-cache \ + "${default_arg_array[@]}" \ "${extra_arg_array[@]}" \ > "${SGLANG_LOG_FILE}" 2>&1 & diff --git a/.github/scripts/atom_test.sh b/.github/scripts/atom_test.sh index 78f6c22d7c..c36eee2971 100755 --- a/.github/scripts/atom_test.sh +++ b/.github/scripts/atom_test.sh @@ -121,12 +121,63 @@ if [ "$TYPE" == "accuracy" ]; then RUN_TAG=$(date +%Y%m%d%H%M%S) OUTPUT_PATH=accuracy_test_results/${RUN_TAG} FLAT_RESULT_FILE=accuracy_test_results/${RUN_TAG}.json - lm_eval --model local-completions \ - --model_args model="$MODEL_PATH",base_url=http://localhost:8000/v1/completions,num_concurrent=65,max_retries=3,tokenized_requests=False,trust_remote_code=True \ - --tasks gsm8k \ - --num_fewshot 3 \ - --output_path "${OUTPUT_PATH}" \ - 2>&1 | tee "$ATOM_CLIENT_LOG" + CLIENT_COMMAND="${CLIENT_COMMAND:-}" + if [[ "${CLIENT_COMMAND}" == "null" ]]; then + CLIENT_COMMAND="" + fi + + if [[ -n "${CLIENT_COMMAND}" ]]; then + CLIENT_COMMAND_ARGS=() + while IFS= read -r -d '' token; do + CLIENT_COMMAND_ARGS+=("${token}") + done < <( + CLIENT_COMMAND="${CLIENT_COMMAND}" \ + MODEL_PATH_VALUE="${MODEL_PATH}" \ + OUTPUT_PATH_VALUE="${OUTPUT_PATH}" \ + python3 - <<'PY' +import os +import shlex +import sys + +client_command = os.environ["CLIENT_COMMAND"] +replacements = { + "${MODEL_PATH}": os.environ["MODEL_PATH_VALUE"], + "$MODEL_PATH": os.environ["MODEL_PATH_VALUE"], + "${OUTPUT_PATH}": os.environ["OUTPUT_PATH_VALUE"], + "$OUTPUT_PATH": os.environ["OUTPUT_PATH_VALUE"], +} +for src, dst in replacements.items(): + client_command = client_command.replace(src, dst) + +for token in shlex.split(client_command): + sys.stdout.write(token) + sys.stdout.write("\0") +PY + ) + + if [[ ${#CLIENT_COMMAND_ARGS[@]} -eq 0 ]]; then + echo "ERROR: CLIENT_COMMAND is set but empty after parsing." + exit 2 + fi + + for arg in "${CLIENT_COMMAND_ARGS[@]}"; do + if [[ "${arg}" =~ \$\{[A-Z0-9_]+\} ]] || [[ "${arg}" =~ \$[A-Z_][A-Z0-9_]* ]]; then + echo "ERROR: CLIENT_COMMAND contains unresolved placeholder after expansion: ${arg}" + exit 2 + fi + done + + echo "Using custom lm-eval command from client_command: ${CLIENT_COMMAND}" + "${CLIENT_COMMAND_ARGS[@]}" 2>&1 | tee "$ATOM_CLIENT_LOG" + else + echo "Using default lm-eval command." + lm_eval --model local-completions \ + --model_args "model=${MODEL_PATH},base_url=http://localhost:8000/v1/completions,num_concurrent=65,max_retries=3,tokenized_requests=False,trust_remote_code=True" \ + --tasks gsm8k \ + --num_fewshot 3 \ + --output_path "${OUTPUT_PATH}" \ + 2>&1 | tee "$ATOM_CLIENT_LOG" + fi RESULT_FILENAME=$( python3 - <&2 + return 1 + fi + + if [ -z "${target_head_sha}" ]; then + echo "Could not determine the target head SHA for the ${CHECKS_WORKFLOW_NAME} workflow run." >&2 + return 1 + fi + + local -a gh_args=( + run list + --repo "${REPO}" + --workflow "${CHECKS_WORKFLOW_NAME}" + --limit 20 + --json databaseId,headSha,headBranch,event,createdAt,status + ) + + if [ -n "${target_branch}" ]; then + gh_args+=(--branch "${target_branch}") + fi + + # Nightly and reusable workflows reuse the Pre Checkin result from the + # original push or pull_request run on the same SHA. + if [ -n "${GITHUB_EVENT_NAME:-}" ] \ + && [ "${GITHUB_EVENT_NAME}" != "schedule" ] \ + && [ "${GITHUB_EVENT_NAME}" != "workflow_call" ]; then + gh_args+=(--event "${GITHUB_EVENT_NAME}") + fi + + gh "${gh_args[@]}" \ + --jq "(map(select(.headSha == \"${target_head_sha}\")) | first | .databaseId) // empty" +} -ARTIFACT_NAME="checks-signal-${TARGET_SHA}" -MAX_RETRIES="${MAX_RETRIES:-10}" -RETRY_DELAY_SECONDS="${RETRY_DELAY_SECONDS:-30}" +# Echoes "\t" for the given run. +fetch_run_state() { + local run_id="$1" + gh api "repos/${REPO}/actions/runs/${run_id}" \ + --jq '[.status, (.conclusion // "")] | @tsv' +} + +print_failed_jobs() { + local run_id="$1" + gh api -X GET "repos/${REPO}/actions/runs/${run_id}/jobs" --paginate \ + --jq '.jobs[] | select(.conclusion != null and .conclusion != "success" and .conclusion != "skipped") | "FAILED: \(.name) (\(.conclusion))"' \ + || true +} for i in $(seq 1 "${MAX_RETRIES}"); do - ATTEMPT_DIR="$(mktemp -d)" - echo "Attempt ${i}: downloading artifact ${ARTIFACT_NAME}..." - - if gh run download \ - --repo "${GITHUB_REPOSITORY}" \ - --name "${ARTIFACT_NAME}" \ - --dir "${ATTEMPT_DIR}"; then - SIGNAL_FILE="${ATTEMPT_DIR}/checks_signal.txt" - if [ -f "${SIGNAL_FILE}" ]; then - echo "Artifact ${ARTIFACT_NAME} downloaded successfully." - SIGNAL="$(head -n 1 "${SIGNAL_FILE}")" - if [ "${SIGNAL}" = "success" ]; then + echo "Attempt ${i}: Locating ${CHECKS_WORKFLOW_NAME} workflow run..." + + RUN_ID="$(find_checks_run_id || true)" + if [ -z "${RUN_ID}" ]; then + echo "Attempt ${i}: Matching ${CHECKS_WORKFLOW_NAME} run not found yet." + else + STATE="$(fetch_run_state "${RUN_ID}" || true)" + STATUS="${STATE%%$'\t'*}" + CONCLUSION="${STATE#*$'\t'}" + # If STATE is unexpected treat that as "no conclusion yet". + [ "${CONCLUSION}" = "${STATE}" ] && CONCLUSION="" + + echo "Attempt ${i}: run ${RUN_ID} status=${STATUS:-unknown} conclusion=${CONCLUSION:-}" + + if [ "${STATUS}" = "completed" ]; then + if [ "${CONCLUSION}" = "success" ]; then echo "Pre Checkin passed, continuing workflow." - rm -rf "${ATTEMPT_DIR}" exit 0 fi - if [ "${SIGNAL}" = "failure" ]; then - echo "Pre Checkin failed, skipping workflow. Details:" - tail -n +2 "${SIGNAL_FILE}" || true - rm -rf "${ATTEMPT_DIR}" - exit 78 - fi - - echo "Unknown signal '${SIGNAL}' in ${SIGNAL_FILE}." - rm -rf "${ATTEMPT_DIR}" - exit 1 + echo "Pre Checkin did not pass (conclusion: ${CONCLUSION:-unknown}), skipping workflow." + print_failed_jobs "${RUN_ID}" + exit 78 # 78 = neutral/skip fi fi - rm -rf "${ATTEMPT_DIR}" - echo "Artifact not found yet, retrying in ${RETRY_DELAY_SECONDS}s..." - sleep "${RETRY_DELAY_SECONDS}" + echo "Pre Checkin not ready yet, retrying in ${RETRY_INTERVAL_SECONDS}s..." + sleep "${RETRY_INTERVAL_SECONDS}" done -echo "Failed to download ${ARTIFACT_NAME} after ${MAX_RETRIES} attempts." +echo "Failed to read Pre Checkin status after ${MAX_RETRIES} attempts. Exiting workflow." exit 1 diff --git a/.github/scripts/collect_gpu_info.sh b/.github/scripts/collect_gpu_info.sh index d4a5723fdd..a98cd1c595 100755 --- a/.github/scripts/collect_gpu_info.sh +++ b/.github/scripts/collect_gpu_info.sh @@ -6,13 +6,17 @@ # 1. `amd-smi static --asic` MARKET_NAME # 2. `rocm-smi --showproductname` Card Series # 3. `rocminfo` Marketing Name -# 4. pattern match (mi355 / mi325 / mi300 / mi250) +# 4. pattern match (mi355 / mi35x / mi325 / mi300 / mi250) # # Step 4 is needed because on freshly-released ASICs (currently MI355X) every # in-container SMI tool can still report "Radeon Graphics" until the # marketing-name table is patched. The CI runner name is operator-asserted # and reliable, so we use it as the final tie-breaker for the dashboard label. # +# Note: `mi35x` is the family-style label used by accuracy-validation runners +# (`linux-atom-mi35x-1` / `-4` / `-8`). The MI35 series in ATOM CI currently +# only contains MI355X, so we map both `mi355` and `mi35x` to MI355X. +# # Usage: # collect_gpu_info.sh # local host # collect_gpu_info.sh # docker exec @@ -58,7 +62,7 @@ if { [ -z "${GPU_NAME:-}" ] || echo "$GPU_NAME" | grep -qi "Radeon Graphics"; } && [ -n "${RUNNER_HINT:-}" ]; then hint_lc=$(echo "$RUNNER_HINT" | tr '[:upper:]' '[:lower:]') case "$hint_lc" in - *mi355*) GPU_NAME="AMD Instinct MI355X" ;; + *mi355*|*mi35x*) GPU_NAME="AMD Instinct MI355X" ;; *mi325*) GPU_NAME="AMD Instinct MI325X" ;; *mi300x*|*mi300*) GPU_NAME="AMD Instinct MI300X" ;; *mi250x*|*mi250*) GPU_NAME="AMD Instinct MI250X" ;; diff --git a/.github/workflows/amd-ci-job-monitor.yml b/.github/workflows/amd-ci-job-monitor.yml index b223c8716b..92e0d7aaca 100644 --- a/.github/workflows/amd-ci-job-monitor.yml +++ b/.github/workflows/amd-ci-job-monitor.yml @@ -219,3 +219,11 @@ jobs: --summary | tee runner-fleet-report.md cat runner-fleet-report.md >> "$GITHUB_STEP_SUMMARY" + + - name: Upload runner fleet report + uses: actions/upload-artifact@v4 + with: + name: runner-fleet-report + path: runner-fleet-report.md + if-no-files-found: error + retention-days: 7 diff --git a/.github/workflows/atom-benchmark.yaml b/.github/workflows/atom-benchmark.yaml index 0466903ac5..ebb67b5415 100644 --- a/.github/workflows/atom-benchmark.yaml +++ b/.github/workflows/atom-benchmark.yaml @@ -22,6 +22,10 @@ on: description: "Benchmark GLM-5-FP8" type: boolean default: true + glm-5.1-fp4: + description: "Benchmark GLM-5.1-MXFP4" + type: boolean + default: true deepseek-r1-0528-mxfp4: description: "Benchmark DeepSeek-R1-0528 MXFP4 (+ MXFP4-MTP3)" type: boolean @@ -34,8 +38,8 @@ on: description: "Benchmark Kimi-K2.5-MXFP4" type: boolean default: true - MiniMax-M2.5: - description: "Benchmark MiniMax-M2.5" + MiniMax-M2.7: + description: "Benchmark MiniMax-M2.7 (+ MXFP4)" type: boolean default: true qwen35-397b-fp8: @@ -78,6 +82,10 @@ on: Example (multiple sets): 1024,1024,128,0.8;2048,1024,256,0.7" type: string default: "1024,1024,128,0.8" + atom_commit: + description: "ATOM commit SHA to checkout; leave empty to use the workflow ref" + type: string + default: "" jobs: parse-param-lists: @@ -87,6 +95,8 @@ jobs: matrix_json: ${{ steps.parse-param-lists.outputs.matrix_json }} steps: - uses: actions/checkout@v6 + with: + ref: ${{ inputs.atom_commit || github.ref }} - name: Parse parameter lists id: parse-param-lists @@ -118,6 +128,8 @@ jobs: models_json: ${{ steps.load.outputs.models_json }} steps: - uses: actions/checkout@v6 + with: + ref: ${{ inputs.atom_commit || github.ref }} - id: load env: EVENT_NAME: ${{ github.event_name }} @@ -182,6 +194,8 @@ jobs: - name: Checkout ATOM repo uses: actions/checkout@v6 + with: + ref: ${{ inputs.atom_commit || github.ref }} - name: Start CI container run: | @@ -190,13 +204,11 @@ jobs: MODEL_MOUNT="" [ -d "/models" ] && MODEL_MOUNT="-v /models:/models" - # Build env var flags from model config - ENV_FLAGS="" - if [ -n "${{ matrix.model.env_vars }}" ]; then - for ev in ${{ matrix.model.env_vars }}; do - ENV_FLAGS="$ENV_FLAGS -e $ev" - done - fi + # Write env_vars via env block (avoids expression injection — newlines + # in `matrix.model.env_vars` would otherwise break the host bash). + # Empty file when MODEL_ENV_VARS unset is fine — docker --env-file + # accepts it. + printenv MODEL_ENV_VARS 2>/dev/null | grep -v '^$' > /tmp/atom_benchmark_env.txt || true docker run -dt --device=/dev/kfd $DEVICE_FLAG \ -v "${GITHUB_WORKSPACE:-$PWD}":/workspace $MODEL_MOUNT \ @@ -210,11 +222,12 @@ jobs: -e CONC=${{ env.CONC }} -e RANDOM_RANGE_RATIO=${{ env.RANDOM_RANGE_RATIO }} \ -e ENABLE_TORCH_PROFILER=${{ inputs.enable_profiler && '1' || '0' }} \ -e ENABLE_RTL_PROFILER=${{ inputs.enable_rtl && '1' || '0' }} \ - $ENV_FLAGS \ + --env-file /tmp/atom_benchmark_env.txt \ --name atom-benchmark \ ${{ inputs.image || 'rocm/atom-dev:latest' }} env: GITHUB_WORKSPACE: ${{ github.workspace }} + MODEL_ENV_VARS: ${{ matrix.model.env_vars }} - name: Collect GPU info (inside container) id: gpu-info @@ -346,6 +359,8 @@ jobs: steps: - name: Checkout ATOM repo uses: actions/checkout@v6 + with: + ref: ${{ inputs.atom_commit || github.ref }} - name: Download all benchmark results uses: actions/download-artifact@v8 @@ -494,6 +509,8 @@ jobs: has_matrix: ${{ steps.gen.outputs.has_matrix }} steps: - uses: actions/checkout@v6 + with: + ref: ${{ inputs.atom_commit || github.ref }} - uses: actions/download-artifact@v8 with: @@ -549,6 +566,8 @@ jobs: - name: Checkout ATOM repo uses: actions/checkout@v6 + with: + ref: ${{ inputs.atom_commit || github.ref }} - name: Docker Login run: echo "${{ secrets.DOCKER_PASSWORD }}" | docker login -u ${{ secrets.DOCKER_USERNAME }} --password-stdin @@ -559,13 +578,9 @@ jobs: MODEL_MOUNT="" [ -d "/models" ] && MODEL_MOUNT="-v /models:/models" - # Build env var flags from model config - ENV_FLAGS="" - if [ -n "${{ matrix.cell.env_vars }}" ]; then - for ev in ${{ matrix.cell.env_vars }}; do - ENV_FLAGS="$ENV_FLAGS -e $ev" - done - fi + # Write env_vars via env block (avoids expression injection — newlines + # in `matrix.cell.env_vars` would otherwise break the host bash). + printenv CELL_ENV_VARS 2>/dev/null | grep -v '^$' > /tmp/atom_regression_env.txt || true docker run -dt --device=/dev/kfd $DEVICE_FLAG \ -v "${GITHUB_WORKSPACE:-$PWD}":/workspace $MODEL_MOUNT \ @@ -575,11 +590,12 @@ jobs: --security-opt seccomp=unconfined \ --ulimit memlock=-1 --ulimit stack=67108864 --pull always \ -e ATOM_DISABLE_MMAP=true \ - $ENV_FLAGS \ + --env-file /tmp/atom_regression_env.txt \ --name atom-regression \ ${{ inputs.image || 'rocm/atom-dev:latest' }} env: GITHUB_WORKSPACE: ${{ github.workspace }} + CELL_ENV_VARS: ${{ matrix.cell.env_vars }} - name: Download model run: | @@ -595,8 +611,12 @@ jobs: if [ -d "/models" ]; then MODEL_DIR="/models/${{ matrix.cell.model_path }}" else MODEL_DIR="${{ matrix.cell.model_path }}"; fi + # cell.env_vars are already injected into the container env via + # --env-file in "Start CI container" — `docker exec` inherits them, + # so no need to re-prefix them on the command line (the previous + # inline form silently dropped the first var when env_vars contained + # a literal newline). docker exec atom-regression bash -lc "set -euo pipefail - ${{ matrix.cell.env_vars }} \ ENABLE_TORCH_PROFILER=1 \ .github/scripts/atom_test.sh launch $MODEL_DIR ${{ matrix.cell.server_args }}" @@ -740,6 +760,8 @@ jobs: steps: - name: Checkout ATOM repo uses: actions/checkout@v6 + with: + ref: ${{ inputs.atom_commit || github.ref }} - name: Download regression traces if: needs.collect-regression-traces.result == 'success' diff --git a/.github/workflows/atom-sglang-accuracy-validation.yaml b/.github/workflows/atom-sglang-accuracy-validation.yaml index 7e6b9baa12..d51a3cff1d 100644 --- a/.github/workflows/atom-sglang-accuracy-validation.yaml +++ b/.github/workflows/atom-sglang-accuracy-validation.yaml @@ -14,6 +14,41 @@ on: required: false type: boolean default: false + run_qwen35_35b_a3b_fp8_tp2: + description: "Qwen3.5-35B-A3B-FP8 TP2" + required: false + type: boolean + default: false + run_qwen35_35b_a3b_tp2: + description: "Qwen3.5-35B-A3B TP2" + required: false + type: boolean + default: false + run_qwen35_397b_a17b_fp8_tp4: + description: "Qwen3.5-397B-A17B-FP8 TP4" + required: false + type: boolean + default: false + run_qwen35_397b_a17b_fp8_tp8: + description: "Qwen3.5-397B-A17B-FP8 TP8" + required: false + type: boolean + default: false + run_dsr1_fp8_tp8: + description: "DeepSeek-R1-FP8 TP8" + required: false + type: boolean + default: false + run_dsr1_fp4_tp4: + description: "DeepSeek-R1-FP4 TP4" + required: false + type: boolean + default: false + run_dsr1_fp4_tp8: + description: "DeepSeek-R1-FP4 TP8" + required: false + type: boolean + default: false upload_accuracy_to_dashboard: description: "Optional: upload SGLANG accuracy results to dashboard after this manual run" required: false @@ -55,6 +90,13 @@ jobs: id: meta env: RUN_DSR1_FP8_TP4: ${{ inputs.run_dsr1_fp8_tp4 }} + RUN_QWEN35_35B_A3B_FP8_TP2: ${{ inputs.run_qwen35_35b_a3b_fp8_tp2 }} + RUN_QWEN35_35B_A3B_TP2: ${{ inputs.run_qwen35_35b_a3b_tp2 }} + RUN_QWEN35_397B_A17B_FP8_TP4: ${{ inputs.run_qwen35_397b_a17b_fp8_tp4 }} + RUN_QWEN35_397B_A17B_FP8_TP8: ${{ inputs.run_qwen35_397b_a17b_fp8_tp8 }} + RUN_DSR1_FP8_TP8: ${{ inputs.run_dsr1_fp8_tp8 }} + RUN_DSR1_FP4_TP4: ${{ inputs.run_dsr1_fp4_tp4 }} + RUN_DSR1_FP4_TP8: ${{ inputs.run_dsr1_fp4_tp8 }} run: | set -euo pipefail @@ -72,10 +114,73 @@ jobs: "model_name": "DeepSeek-R1-FP8 TP4", "model_path": "deepseek-ai/DeepSeek-R1-0528", "extra_args": "--tensor-parallel-size 4", - "accuracy_test_threshold": 0.92, - "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nSGLANG_AITER_FP8_PREFILL_ATTN=0\nSGLANG_USE_AITER=1\nATOM_ENABLE_DS_QKNORM_QUANT_FUSION=1", + "accuracy_test_threshold": 0.91, + "env_vars": "SGLANG_AITER_FP8_PREFILL_ATTN=0\nSGLANG_USE_AITER=1\nATOM_ENABLE_DS_QKNORM_QUANT_FUSION=1", "runner": "linux-atom-mi35x-4", }, + { + "toggle_env": "RUN_QWEN35_35B_A3B_FP8_TP2", + "model_name": "Qwen3.5-35B-A3B-FP8 TP2", + "model_path": "Qwen/Qwen3.5-35B-A3B-FP8", + "extra_args": "--tensor-parallel-size 2 --mem-fraction-static 0.9 --reasoning-parser qwen3 --disable-radix-cache", + "accuracy_test_threshold": 0.76, + "env_vars": "SGLANG_DEFAULT_SERVER_ARGS=\nSGLANG_EXTERNAL_MODEL_PACKAGE=atom.plugin.sglang.models\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=0", + "runner": "linux-atom-mi35x-4", + }, + { + "toggle_env": "RUN_QWEN35_35B_A3B_TP2", + "model_name": "Qwen3.5-35B-A3B TP2", + "model_path": "Qwen/Qwen3.5-35B-A3B", + "extra_args": "--tensor-parallel-size 2 --mem-fraction-static 0.9 --reasoning-parser qwen3 --disable-radix-cache", + "accuracy_test_threshold": 0.83, + "env_vars": "SGLANG_DEFAULT_SERVER_ARGS=\nSGLANG_EXTERNAL_MODEL_PACKAGE=atom.plugin.sglang.models\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=0", + "runner": "linux-atom-mi35x-4", + }, + { + "toggle_env": "RUN_QWEN35_397B_A17B_FP8_TP4", + "model_name": "Qwen3.5-397B-A17B-FP8 TP4", + "model_path": "Qwen/Qwen3.5-397B-A17B-FP8", + "extra_args": "--tensor-parallel-size 4 --mem-fraction-static 0.9 --reasoning-parser qwen3 --disable-radix-cache", + "accuracy_test_threshold": 0.83, + "env_vars": "SGLANG_DEFAULT_SERVER_ARGS=\nSGLANG_EXTERNAL_MODEL_PACKAGE=atom.plugin.sglang.models\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=0", + "runner": "linux-atom-mi35x-4", + }, + { + "toggle_env": "RUN_QWEN35_397B_A17B_FP8_TP8", + "model_name": "Qwen3.5-397B-A17B-FP8 TP8", + "model_path": "Qwen/Qwen3.5-397B-A17B-FP8", + "extra_args": "--tensor-parallel-size 8 --mem-fraction-static 0.9 --reasoning-parser qwen3 --disable-radix-cache", + "accuracy_test_threshold": 0.83, + "env_vars": "SGLANG_DEFAULT_SERVER_ARGS=\nSGLANG_EXTERNAL_MODEL_PACKAGE=atom.plugin.sglang.models\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=0", + "runner": "linux-atom-mi35x-8", + }, + { + "toggle_env": "RUN_DSR1_FP8_TP8", + "model_name": "DeepSeek-R1-FP8 TP8", + "model_path": "deepseek-ai/DeepSeek-R1-0528", + "extra_args": "--tensor-parallel-size 8", + "accuracy_test_threshold": 0.93, + "env_vars": "SGLANG_AITER_FP8_PREFILL_ATTN=0\nSGLANG_USE_AITER=1\nATOM_ENABLE_DS_QKNORM_QUANT_FUSION=1", + "runner": "linux-atom-mi35x-8", + }, + { + "toggle_env": "RUN_DSR1_FP4_TP4", + "model_name": "DeepSeek-R1-FP4 TP4", + "model_path": "amd/DeepSeek-R1-0528-MXFP4-MTP-MoEFP4", + "extra_args": "--tensor-parallel-size 4", + "accuracy_test_threshold": 0.91, + "env_vars": "SGLANG_AITER_FP8_PREFILL_ATTN=0\nSGLANG_USE_AITER=1\nATOM_ENABLE_DS_QKNORM_QUANT_FUSION=1", + "runner": "linux-atom-mi35x-4", + }, + { + "toggle_env": "RUN_DSR1_FP4_TP8", + "model_name": "DeepSeek-R1-FP4 TP8", + "model_path": "amd/DeepSeek-R1-0528-MXFP4-MTP-MoEFP4", + "extra_args": "--tensor-parallel-size 8", + "accuracy_test_threshold": 0.93, + "env_vars": "SGLANG_AITER_FP8_PREFILL_ATTN=0\nSGLANG_USE_AITER=1\nATOM_ENABLE_DS_QKNORM_QUANT_FUSION=1", + "runner": "linux-atom-mi35x-8", + }, ] selected = [] diff --git a/.github/workflows/atom-sglang-benchmark.yaml b/.github/workflows/atom-sglang-benchmark.yaml index 66720b9241..84680999b9 100644 --- a/.github/workflows/atom-sglang-benchmark.yaml +++ b/.github/workflows/atom-sglang-benchmark.yaml @@ -5,6 +5,9 @@ concurrency: cancel-in-progress: ${{ github.ref != 'refs/heads/main' }} on: + schedule: + # Weekday nightly at 23:00 Beijing time (15:00 UTC) + - cron: '0 15 * * 1-5' workflow_dispatch: inputs: deepseek-r1-fp8-tp8: @@ -27,6 +30,14 @@ on: description: "DeepSeek-R1-0528-MXFP4 FP4 TP8 EP8" type: boolean default: false + qwen3-5-397b-a17b-fp8-tp4: + description: "Qwen3.5-397B-A17B-FP8 TP4" + type: boolean + default: false + qwen3-5-397b-a17b-fp8-tp8: + description: "Qwen3.5-397B-A17B-FP8 TP8" + type: boolean + default: false sglang_image: description: "Optional SGLang benchmark image override. Leave empty to use sglang-latest on main or rebuild from the selected non-main branch." type: string @@ -83,6 +94,10 @@ jobs: REBUILD_SGLANG_IMAGE=false SGLANG_IMAGE_SOURCE="prebuilt" USER_PROVIDED_SGLANG_IMAGE=false + RESOLVED_PREBUILT_ALIAS="" + RESOLVED_PREBUILT_DIGEST="" + PREBUILT_RESOLUTION_TYPE="" + RESOLVED_NIGHTLY_TAG="" if [[ -n "${INPUT_SGLANG_IMAGE}" ]]; then PREBUILT_SGLANG_IMAGE="${INPUT_SGLANG_IMAGE}" SGLANG_IMAGE_SOURCE="user-provided" @@ -116,6 +131,39 @@ jobs: SELECTED_SGLANG_REF="${VLLM_META[0]}" SELECTED_SGLANG_VERSION="${VLLM_META[1]}" + if [[ "${SGLANG_IMAGE_SOURCE}" == "prebuilt" ]]; then + if RESOLUTION_JSON="$( + python3 .github/scripts/resolve_atom_image.py \ + --repository rocm/atom-dev \ + --reference-tag sglang-latest \ + --preferred-version "${SELECTED_SGLANG_VERSION}" + )"; then + mapfile -t RESOLVED_IMAGE_META < <( + RESOLUTION_JSON="${RESOLUTION_JSON}" python3 - <<'PY' + import json + import os + + resolution = json.loads(os.environ["RESOLUTION_JSON"]) + print(resolution["resolved_image"]) + print(resolution["reference_image"]) + print(resolution["reference_digest"]) + print(resolution["match_type"]) + print(resolution.get("matched_tag", "")) + PY + ) + PREBUILT_SGLANG_IMAGE="${RESOLVED_IMAGE_META[0]}" + RESOLVED_PREBUILT_ALIAS="${RESOLVED_IMAGE_META[1]}" + RESOLVED_PREBUILT_DIGEST="${RESOLVED_IMAGE_META[2]}" + PREBUILT_RESOLUTION_TYPE="${RESOLVED_IMAGE_META[3]}" + RESOLVED_NIGHTLY_TAG="${RESOLVED_IMAGE_META[4]}" + else + echo "::warning::Failed to resolve rocm/atom-dev:sglang-latest to an immutable digest; falling back to the floating tag." + PREBUILT_SGLANG_IMAGE="rocm/atom-dev:sglang-latest" + RESOLVED_PREBUILT_ALIAS="rocm/atom-dev:sglang-latest" + PREBUILT_RESOLUTION_TYPE="resolution-failed" + fi + fi + { echo "atom_repository=${ATOM_REPOSITORY}" echo "atom_ref=${ATOM_REF}" @@ -153,6 +201,32 @@ jobs: "${PREBUILT_SGLANG_IMAGE}" \ "${SELECTED_SGLANG_VERSION}" \ "${SELECTED_SGLANG_REF}" >> "$GITHUB_STEP_SUMMARY" + if [[ -n "${RESOLVED_PREBUILT_ALIAS}" ]]; then + printf -- '- Floating alias: `%s`\n' \ + "${RESOLVED_PREBUILT_ALIAS}" >> "$GITHUB_STEP_SUMMARY" + fi + if [[ -n "${RESOLVED_PREBUILT_DIGEST}" ]]; then + printf -- '- Resolved digest: `%s`\n' \ + "${RESOLVED_PREBUILT_DIGEST}" >> "$GITHUB_STEP_SUMMARY" + fi + case "${PREBUILT_RESOLUTION_TYPE}" in + matched-nightly-tag) + printf -- '- Resolution: reverse-matched same digest to nightly tag `%s`\n' \ + "${RESOLVED_NIGHTLY_TAG}" >> "$GITHUB_STEP_SUMMARY" + ;; + digest-pinned-latest) + printf -- '- Resolution: no nightly tag with the same digest was found, so this run pins `sglang-latest` to its current digest instead.\n' \ + >> "$GITHUB_STEP_SUMMARY" + ;; + reverse-lookup-failed-digest-pinned-latest) + printf -- '- Resolution: reverse lookup could not complete, so this run pins `sglang-latest` to its current digest instead of falling back to the floating tag.\n' \ + >> "$GITHUB_STEP_SUMMARY" + ;; + resolution-failed) + printf -- '- Resolution: digest pinning itself failed, so this run falls back to floating `sglang-latest`.\n' \ + >> "$GITHUB_STEP_SUMMARY" + ;; + esac fi printf -- '- Upload to dashboard: `%s`\n' \ @@ -208,6 +282,7 @@ jobs: outputs: models_json: ${{ steps.load.outputs.models_json }} has_enabled_models: ${{ steps.load.outputs.has_enabled_models }} + selected_group: ${{ steps.load.outputs.selected_group }} steps: - uses: actions/checkout@v6 - id: load @@ -217,24 +292,80 @@ jobs: ENABLE_DEEPSEEK_R1_FP4_TP8: ${{ inputs.deepseek-r1-fp4-tp8 }} ENABLE_DEEPSEEK_R1_FP4_TP4: ${{ inputs.deepseek-r1-fp4-tp4 }} ENABLE_DEEPSEEK_R1_FP4_TP8_EP8: ${{ inputs.deepseek-r1-fp4-tp8-ep8 }} + ENABLE_QWEN3_5_397B_A17B_FP8_TP4: ${{ inputs.qwen3-5-397b-a17b-fp8-tp4 }} + ENABLE_QWEN3_5_397B_A17B_FP8_TP8: ${{ inputs.qwen3-5-397b-a17b-fp8-tp8 }} run: | - MODELS_JSON="$(jq -c ' - map(select( - (.prefix == "deepseek-r1-fp8-tp8" and env.ENABLE_DEEPSEEK_R1_FP8_TP8 == "true") - or (.prefix == "deepseek-r1-fp8-tp4" and env.ENABLE_DEEPSEEK_R1_FP8_TP4 == "true") - or (.prefix == "deepseek-r1-fp4-tp8" and env.ENABLE_DEEPSEEK_R1_FP4_TP8 == "true") - or (.prefix == "deepseek-r1-fp4-tp4" and env.ENABLE_DEEPSEEK_R1_FP4_TP4 == "true") - or (.prefix == "deepseek-r1-fp4-tp8-ep8" and env.ENABLE_DEEPSEEK_R1_FP4_TP8_EP8 == "true") - )) - ' .github/benchmark/sglang_benchmark_models.json)" - echo "models_json=${MODELS_JSON}" >> "$GITHUB_OUTPUT" - if [ "${MODELS_JSON}" = "[]" ]; then - echo "has_enabled_models=false" >> "$GITHUB_OUTPUT" - echo "No models selected for SGLang benchmark." - else - echo "has_enabled_models=true" >> "$GITHUB_OUTPUT" - echo "Selected models: ${MODELS_JSON}" + set -euo pipefail + + export BEIJING_WEEKDAY="$(TZ=Asia/Shanghai date +%u)" + python3 - <<'PY' >> "$GITHUB_OUTPUT" + import json + import os + import sys + from pathlib import Path + + models = json.loads( + Path(".github/benchmark/sglang_benchmark_models.json").read_text( + encoding="utf-8" + ) + ) + + event = os.environ["GITHUB_EVENT_NAME"] + selected_group = "" + + if event == "schedule": + weekday = int(os.environ["BEIJING_WEEKDAY"]) + if weekday in (1, 3): + selected_group = "A-DEEPSEEK" + selected = [m for m in models if m.get("nightly_group", "A") == "A"] + elif weekday in (2, 4): + selected_group = "B-QWEN35" + selected = [m for m in models if m.get("nightly_group") == "B"] + elif weekday == 5: + selected_group = "C-ALL" + selected = list(models) + else: + selected_group = "SKIP-WEEKEND" + selected = [] + else: + enabled_by_prefix = { + "deepseek-r1-fp8-tp8": os.environ.get("ENABLE_DEEPSEEK_R1_FP8_TP8", ""), + "deepseek-r1-fp8-tp4": os.environ.get("ENABLE_DEEPSEEK_R1_FP8_TP4", ""), + "deepseek-r1-fp4-tp8": os.environ.get("ENABLE_DEEPSEEK_R1_FP4_TP8", ""), + "deepseek-r1-fp4-tp4": os.environ.get("ENABLE_DEEPSEEK_R1_FP4_TP4", ""), + "deepseek-r1-fp4-tp8-ep8": os.environ.get("ENABLE_DEEPSEEK_R1_FP4_TP8_EP8", ""), + "qwen3-5-397b-a17b-fp8-tp4": os.environ.get("ENABLE_QWEN3_5_397B_A17B_FP8_TP4", ""), + "qwen3-5-397b-a17b-fp8-tp8": os.environ.get("ENABLE_QWEN3_5_397B_A17B_FP8_TP8", ""), + } + selected = [ + model + for model in models + if enabled_by_prefix.get(str(model.get("prefix")), "").lower() == "true" + ] + + if selected_group: + print(f"Scheduled SGLang benchmark group: {selected_group}", file=sys.stderr) + if not selected: + print("No models selected for SGLang benchmark.", file=sys.stderr) + else: + print("Selected SGLang benchmark models:", file=sys.stderr) + for model in selected: + print(f" - {model['display']} ({model['prefix']})", file=sys.stderr) + + print(f"models_json={json.dumps(selected, separators=(',', ':'))}") + print(f"selected_group={selected_group}") + print(f"has_enabled_models={'true' if selected else 'false'}") + PY + + - name: Print selected models + env: + SELECTED_GROUP: ${{ steps.load.outputs.selected_group }} + MODELS_JSON: ${{ steps.load.outputs.models_json }} + run: | + if [[ -n "${SELECTED_GROUP}" ]]; then + echo "Scheduled SGLang benchmark group: ${SELECTED_GROUP}" fi + printf 'Selected models: %s\n' "${MODELS_JSON}" build-benchmark-matrix: name: Build SGLang benchmark matrix @@ -469,12 +600,19 @@ jobs: - name: Check if model is enabled id: check run: | + if [ "${GITHUB_EVENT_NAME}" = "schedule" ]; then + echo "enabled=true" >> "$GITHUB_OUTPUT" + exit 0 + fi + case "${{ matrix.model.prefix }}" in deepseek-r1-fp8-tp8) echo "enabled=${{ inputs.deepseek-r1-fp8-tp8 }}" >> "$GITHUB_OUTPUT" ;; deepseek-r1-fp8-tp4) echo "enabled=${{ inputs.deepseek-r1-fp8-tp4 }}" >> "$GITHUB_OUTPUT" ;; deepseek-r1-fp4-tp8) echo "enabled=${{ inputs.deepseek-r1-fp4-tp8 }}" >> "$GITHUB_OUTPUT" ;; deepseek-r1-fp4-tp4) echo "enabled=${{ inputs.deepseek-r1-fp4-tp4 }}" >> "$GITHUB_OUTPUT" ;; deepseek-r1-fp4-tp8-ep8) echo "enabled=${{ inputs.deepseek-r1-fp4-tp8-ep8 }}" >> "$GITHUB_OUTPUT" ;; + qwen3-5-397b-a17b-fp8-tp4) echo "enabled=${{ inputs.qwen3-5-397b-a17b-fp8-tp4 }}" >> "$GITHUB_OUTPUT" ;; + qwen3-5-397b-a17b-fp8-tp8) echo "enabled=${{ inputs.qwen3-5-397b-a17b-fp8-tp8 }}" >> "$GITHUB_OUTPUT" ;; *) echo "enabled=true" >> "$GITHUB_OUTPUT" ;; esac diff --git a/.github/workflows/atom-sglang-test.yaml b/.github/workflows/atom-sglang-test.yaml index c0ac04638a..da23aa4423 100644 --- a/.github/workflows/atom-sglang-test.yaml +++ b/.github/workflows/atom-sglang-test.yaml @@ -35,7 +35,7 @@ jobs: - name: Checkout ATOM repo uses: actions/checkout@v6 - - name: Download and check pre-checkin signal + - name: Wait for Pre Checkin workflow run: bash ./.github/scripts/check_signal.sh env: GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} @@ -50,55 +50,140 @@ jobs: aiter_artifact_id: ${{ steps.download.outputs.aiter_artifact_id }} aiter_wheel_name: ${{ steps.download.outputs.aiter_wheel_name }} steps: - - name: Find and download latest aiter wheel + - name: Prefer latest main aiter wheel manifest and fallback to artifact id: download run: | set -euo pipefail - echo "=== Finding latest aiter-whl-main artifact from ROCm/aiter ===" + echo "=== Trying latest main aiter wheel manifest from S3 first ===" + S3_MAIN_MANIFEST_URL="https://rocm.frameworks-nightlies.amd.com/whl-staging/gfx942-gfx950/main/latest.json" API_URL="https://api.github.com" AUTH_HEADER="Authorization: token ${{ secrets.GITHUB_TOKEN }}" AITER_TEST_WORKFLOW_ID=179476100 - RUNS=$(curl -s -H "$AUTH_HEADER" \ - "$API_URL/repos/ROCm/aiter/actions/workflows/$AITER_TEST_WORKFLOW_ID/runs?per_page=100&branch=main&event=push") - ARTIFACT_ID="" ARTIFACT_NAME="" - for RUN_ID in $(echo "$RUNS" | jq -r '.workflow_runs[].id'); do - ARTIFACT_JSON=$(curl -s -H "$AUTH_HEADER" \ - "$API_URL/repos/ROCm/aiter/actions/runs/$RUN_ID/artifacts" \ - | jq '[.artifacts[] | select(.name | startswith("aiter-whl-main")) | select(.expired == false)] | first') - - if [ "$ARTIFACT_JSON" != "null" ] && [ -n "$ARTIFACT_JSON" ]; then - ARTIFACT_ID=$(echo "$ARTIFACT_JSON" | jq -r '.id') - ARTIFACT_NAME=$(echo "$ARTIFACT_JSON" | jq -r '.name') - echo "Found artifact in run $RUN_ID: $ARTIFACT_NAME (ID: $ARTIFACT_ID)" - break + ARTIFACT_RUN_ID="" + ARTIFACT_RUN_SHA="" + ARTIFACT_RUN_CREATED_AT="" + + resolve_download_url() { + python3 -c 'import sys + from urllib.parse import quote, unquote, urlsplit, urlunsplit + parts = urlsplit(sys.argv[1]) + encoded_path = "/".join(quote(unquote(segment), safe="") for segment in parts.path.split("/")) + print(urlunsplit((parts.scheme, parts.netloc, encoded_path, parts.query, parts.fragment)))' "$1" + } + + find_latest_artifact() { + local runs_json artifact_json run_id + + if [ -n "$ARTIFACT_ID" ] && [ "$ARTIFACT_ID" != "null" ]; then + return 0 fi - done - if [ -z "$ARTIFACT_ID" ] || [ "$ARTIFACT_ID" = "null" ]; then - echo "ERROR: No aiter-whl-main artifact found in recent Aiter Test runs" - exit 1 - fi + echo "=== Finding latest aiter-whl-* artifact from ROCm/aiter ===" + runs_json=$(curl -fsSL -H "$AUTH_HEADER" \ + "$API_URL/repos/ROCm/aiter/actions/workflows/$AITER_TEST_WORKFLOW_ID/runs?per_page=100&branch=main&event=push") + + for run_id in $(echo "$runs_json" | jq -r '.workflow_runs[].id'); do + artifact_json=$(curl -fsSL -H "$AUTH_HEADER" \ + "$API_URL/repos/ROCm/aiter/actions/runs/$run_id/artifacts" \ + | jq '[.artifacts[] | select(.name | startswith("aiter-whl-")) | select(.expired == false)] | sort_by(.created_at) | last') + + if [ "$artifact_json" != "null" ] && [ -n "$artifact_json" ]; then + ARTIFACT_ID=$(echo "$artifact_json" | jq -r '.id') + ARTIFACT_NAME=$(echo "$artifact_json" | jq -r '.name') + ARTIFACT_RUN_ID="$run_id" + ARTIFACT_RUN_SHA=$(echo "$runs_json" | jq -r --arg run_id "$run_id" '.workflow_runs[] | select((.id | tostring) == $run_id) | .head_sha') + ARTIFACT_RUN_CREATED_AT=$(echo "$runs_json" | jq -r --arg run_id "$run_id" '.workflow_runs[] | select((.id | tostring) == $run_id) | .created_at') + echo "Found artifact in run $ARTIFACT_RUN_ID: $ARTIFACT_NAME (ID: $ARTIFACT_ID, SHA: $ARTIFACT_RUN_SHA)" + return 0 + fi + done + + return 1 + } + + download_from_s3_manifest() { + local manifest_file manifest_fetch_url manifest_branch manifest_timestamp manifest_commit wheel_name wheel_url resolved_wheel_url + + mkdir -p aiter-whl + rm -f aiter-whl/amd_aiter*.whl + + manifest_file=$(mktemp) + trap 'rm -f "$manifest_file"' RETURN + manifest_fetch_url="${S3_MAIN_MANIFEST_URL}?ts=$(date +%s)" + curl -fsSL -H "Cache-Control: no-cache" "$manifest_fetch_url" -o "$manifest_file" || return 1 + + manifest_branch=$(jq -r '.branch // empty' "$manifest_file") + manifest_timestamp=$(jq -r '.timestamp // empty' "$manifest_file") + manifest_commit=$(jq -r '.commit // empty' "$manifest_file") + wheel_name=$(jq -r '.wheel_name // empty' "$manifest_file") + wheel_url=$(jq -r '.wheel_url // empty' "$manifest_file") + + if [ "$manifest_branch" != "main" ] || [ -z "$manifest_timestamp" ] || [ -z "$manifest_commit" ] || [ -z "$wheel_name" ] || [ -z "$wheel_url" ]; then + echo "Invalid latest main wheel manifest" + return 1 + fi - echo "=== Downloading artifact ===" - mkdir -p aiter-whl - curl -s -L -H "$AUTH_HEADER" \ - "$API_URL/repos/ROCm/aiter/actions/artifacts/$ARTIFACT_ID/zip" \ - -o aiter-whl.zip - unzip -o aiter-whl.zip -d aiter-whl - rm -f aiter-whl.zip + if find_latest_artifact; then + if [ -n "$ARTIFACT_RUN_SHA" ] && [ "$manifest_commit" != "$ARTIFACT_RUN_SHA" ]; then + if [ -n "$ARTIFACT_RUN_CREATED_AT" ] && [[ "$manifest_timestamp" < "$ARTIFACT_RUN_CREATED_AT" ]]; then + echo "Manifest commit $manifest_commit is older than latest artifact run $ARTIFACT_RUN_ID ($ARTIFACT_RUN_SHA); treating manifest as stale" + return 1 + fi + echo "Manifest commit $manifest_commit differs from latest artifact run $ARTIFACT_RUN_ID ($ARTIFACT_RUN_SHA), but manifest timestamp is not older" + fi + else + echo "No GitHub fallback artifact found while checking manifest freshness" + fi + + resolved_wheel_url=$(resolve_download_url "$wheel_url") + + echo "Selected latest main wheel manifest: $S3_MAIN_MANIFEST_URL" + echo "Manifest timestamp: $manifest_timestamp" + echo "Manifest commit: $manifest_commit" + echo "Manifest wheel: $wheel_name" + echo "Downloading manifest-selected wheel: $resolved_wheel_url" + curl -fsSL "$resolved_wheel_url" -o "aiter-whl/$wheel_name" || return 1 + echo "Downloaded wheel from manifest: aiter-whl/$wheel_name" + + rm -f "$manifest_file" + trap - RETURN + } + + download_from_artifact() { + echo "=== Falling back to latest aiter-whl-* artifact from ROCm/aiter ===" + find_latest_artifact || { + echo "ERROR: No aiter-whl-* artifact found in recent Aiter Test runs" + return 1 + } + + mkdir -p aiter-whl + rm -f aiter-whl/amd_aiter*.whl + curl -fsSL -H "$AUTH_HEADER" \ + "$API_URL/repos/ROCm/aiter/actions/artifacts/$ARTIFACT_ID/zip" \ + -o aiter-whl.zip + unzip -o aiter-whl.zip -d aiter-whl + rm -f aiter-whl.zip + } + + if download_from_s3_manifest; then + echo "Using wheel from S3 main manifest" + else + echo "Main wheel manifest download failed, falling back to GitHub artifact" + download_from_artifact + fi AITER_WHL=$(ls -t aiter-whl/amd_aiter*.whl 2>/dev/null | head -1) if [ -z "$AITER_WHL" ]; then - echo "ERROR: No amd_aiter wheel found in artifact" - ls -la aiter-whl/ + echo "ERROR: No amd_aiter wheel available after S3/artifact attempts" + ls -la aiter-whl/ || true exit 1 fi - echo "Downloaded wheel: $AITER_WHL" + echo "Selected wheel: $AITER_WHL" echo "aiter_artifact_id=${ARTIFACT_ID}" >> "$GITHUB_OUTPUT" echo "aiter_wheel_name=$(basename "$AITER_WHL")" >> "$GITHUB_OUTPUT" @@ -107,7 +192,7 @@ jobs: with: name: aiter-whl path: aiter-whl/amd_aiter*.whl - retention-days: 1 + retention-days: 7 atom-sglang-accuracy: needs: [check-signal, download_aiter_wheel] @@ -121,11 +206,28 @@ jobs: model_path: "deepseek-ai/DeepSeek-R1-0528" extra_args: "--tensor-parallel-size 4" env_vars: | - AITER_QUICK_REDUCE_QUANTIZATION=INT4 SGLANG_AITER_FP8_PREFILL_ATTN=0 SGLANG_USE_AITER=1 ATOM_ENABLE_DS_QKNORM_QUANT_FUSION=1 - accuracy_test_threshold: 0.92 + accuracy_test_threshold: 0.91 + runner: linux-atom-mi35x-4 + - model_name: "DeepSeek-R1-FP4 TP4" + model_path: "amd/DeepSeek-R1-0528-MXFP4-MTP-MoEFP4" + extra_args: "--tensor-parallel-size 4" + env_vars: | + SGLANG_AITER_FP8_PREFILL_ATTN=0 + SGLANG_USE_AITER=1 + ATOM_ENABLE_DS_QKNORM_QUANT_FUSION=1 + accuracy_test_threshold: 0.91 + runner: linux-atom-mi35x-4 + - model_name: "Qwen3.5-35B-A3B-FP8 TP2" + model_path: "Qwen/Qwen3.5-35B-A3B-FP8" + extra_args: "--tensor-parallel-size 2 --mem-fraction-static 0.9 --reasoning-parser qwen3 --disable-radix-cache" + env_vars: | + SGLANG_DEFAULT_SERVER_ARGS= + SGLANG_EXTERNAL_MODEL_PACKAGE=atom.plugin.sglang.models + ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=0 + accuracy_test_threshold: 0.76 runner: linux-atom-mi35x-4 runs-on: ${{ matrix.runner }} timeout-minutes: 180 @@ -434,4 +536,3 @@ jobs: rm -f Dockerfile.mod || true docker rmi "atom_sglang_base:ci" || true docker rmi "atom_sglang:ci" || true - diff --git a/.github/workflows/atom-test.yaml b/.github/workflows/atom-test.yaml index 769d34cc4b..83cdb49fc5 100644 --- a/.github/workflows/atom-test.yaml +++ b/.github/workflows/atom-test.yaml @@ -47,7 +47,7 @@ jobs: if: ${{ github.event_name != 'workflow_dispatch' }} uses: actions/checkout@v6 - - name: Download and check pre-checkin signal + - name: Wait for Pre Checkin workflow if: ${{ github.event_name != 'workflow_dispatch' }} run: bash ./.github/scripts/check_signal.sh env: @@ -199,7 +199,7 @@ jobs: with: name: aiter-whl path: aiter-whl/amd_aiter*.whl - retention-days: 1 + retention-days: 7 load-test-models: name: Load test model configs @@ -545,6 +545,7 @@ jobs: timeout-minutes: 30 env: MODEL_EXTRA_ARGS: ${{ matrix.extraArgs }} + CLIENT_COMMAND: ${{ matrix.client_command || '' }} run: | set -euo pipefail echo "" @@ -561,6 +562,7 @@ jobs: echo "" echo "========== Running accuracy test ==========" docker exec \ + -e CLIENT_COMMAND="${CLIENT_COMMAND}" \ -e GPU_NAME="${{ steps.gpu-info.outputs.gpu_name }}" \ -e GPU_VRAM_GB="${{ steps.gpu-info.outputs.gpu_vram_gb }}" \ -e ROCM_VERSION="${{ steps.gpu-info.outputs.rocm_version }}" \ diff --git a/.github/workflows/atom-vllm-accuracy-validation.yaml b/.github/workflows/atom-vllm-accuracy-validation.yaml index 97d9a6c8f3..6d82397d9c 100644 --- a/.github/workflows/atom-vllm-accuracy-validation.yaml +++ b/.github/workflows/atom-vllm-accuracy-validation.yaml @@ -24,6 +24,11 @@ on: required: false type: boolean default: false + run_qwen3_next_80b_mtp_tp4: + description: "Qwen3-Next-80B-A3B-Instruct-FP8-MTP TP4" + required: false + type: boolean + default: false run_qwen35_397b_fp8_tp8: description: "Qwen3.5-397B-A17B-FP8 TP8" required: false @@ -84,6 +89,11 @@ on: required: false type: boolean default: false + run_glm4_7_fp8_mtp_tp8: + description: "GLM-4.7-FP8 MTP TP8" + required: false + type: boolean + default: false run_glm5_1_fp8_tp8: description: "GLM-5.1-FP8 TP8" required: false @@ -137,6 +147,7 @@ jobs: RUN_QWEN3_MOE_TP8: ${{ inputs.run_qwen3_moe_tp8 }} RUN_QWEN3_NEXT_80B_TP1: ${{ inputs.run_qwen3_next_80b_tp1 }} RUN_QWEN3_NEXT_80B_TP4: ${{ inputs.run_qwen3_next_80b_tp4 }} + RUN_QWEN3_NEXT_80B_MTP_TP4: ${{ inputs.run_qwen3_next_80b_mtp_tp4 }} RUN_QWEN35_397B_FP8_TP8: ${{ inputs.run_qwen35_397b_fp8_tp8 }} RUN_QWEN35_397B_TP8: ${{ inputs.run_qwen35_397b_tp8 }} RUN_QWEN35_397B_FP4_TP4: ${{ inputs.run_qwen35_397b_fp4_tp4 }} @@ -149,6 +160,7 @@ jobs: RUN_DSV32_FP8_TP8: ${{ inputs.run_dsv32_fp8_tp8 }} RUN_GPT_OSS_120B_TP1: ${{ inputs.run_gpt_oss_120b_tp1 }} RUN_GPT_OSS_120B_TP2: ${{ inputs.run_gpt_oss_120b_tp2 }} + RUN_GLM4_7_FP8_MTP_TP8: ${{ inputs.run_glm4_7_fp8_mtp_tp8 }} RUN_GLM5_1_FP8_TP8: ${{ inputs.run_glm5_1_fp8_tp8 }} INPUT_OOT_IMAGE: ${{ inputs.oot_image || '' }} run: | @@ -169,7 +181,7 @@ jobs: "model_path": "Qwen/Qwen3-235B-A22B-Instruct-2507-FP8", "extra_args": "--tensor-parallel-size 8 --enable-expert-parallel", "accuracy_test_threshold": 0.87, - "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1", + "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE=0", "runner": "linux-atom-mi35x-8", }, { @@ -178,7 +190,7 @@ jobs: "model_path": "Qwen/Qwen3-Next-80B-A3B-Instruct-FP8", "extra_args": "--tensor-parallel-size 1", "accuracy_test_threshold": 0.83, - "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0\nATOM_USE_FLYDSL_GDR=1", + "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0\nATOM_USE_FLYDSL_GDR=1\nATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE=0", "runner": "linux-atom-mi35x-1", }, { @@ -187,7 +199,16 @@ jobs: "model_path": "Qwen/Qwen3-Next-80B-A3B-Instruct-FP8", "extra_args": "--tensor-parallel-size 4", "accuracy_test_threshold": 0.83, - "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0\nATOM_USE_FLYDSL_GDR=1", + "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0\nATOM_USE_FLYDSL_GDR=1\nATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE=0", + "runner": "linux-atom-mi35x-4", + }, + { + "toggle_env": "RUN_QWEN3_NEXT_80B_MTP_TP4", + "model_name": "Qwen3-Next-80B-A3B-Instruct-FP8-MTP TP4", + "model_path": "Qwen/Qwen3-Next-80B-A3B-Instruct-FP8", + "extra_args": "--tensor-parallel-size 4 --speculative-config '{\"num_speculative_tokens\":1, \"method\": \"mtp\"}'", + "accuracy_test_threshold": 0.80, + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0\nATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE=0", "runner": "linux-atom-mi35x-4", }, { @@ -196,7 +217,7 @@ jobs: "model_path": "Qwen/Qwen3.5-397B-A17B-FP8", "extra_args": "--tensor-parallel-size 8", "accuracy_test_threshold": 0.83, - "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0", + "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0\nATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE=0", "runner": "linux-atom-mi35x-8", }, { @@ -205,7 +226,7 @@ jobs: "model_path": "Qwen/Qwen3.5-397B-A17B", "extra_args": "--tensor-parallel-size 8", "accuracy_test_threshold": 0.83, - "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0", + "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0\nATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE=0", "runner": "linux-atom-mi35x-8", }, { @@ -214,7 +235,7 @@ jobs: "model_path": "amd/Qwen3.5-397B-A17B-MXFP4", "extra_args": "--tensor-parallel-size 4", "accuracy_test_threshold": 0.83, - "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0", + "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0\nATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE=0", "runner": "linux-atom-mi35x-4", }, { @@ -286,7 +307,8 @@ jobs: "model_name": "gpt-oss-120b TP1", "model_path": "openai/gpt-oss-120b", "extra_args": "--tensor-parallel-size 1", - "accuracy_test_threshold": 0.38, + "client_command": "lm_eval --model local-chat-completions --apply_chat_template --model_args model=${MODEL_PATH},base_url=http://127.0.0.1:${VLLM_PORT}/v1/chat/completions,num_concurrent=65,max_retries=1,max_gen_toks=2048,tokenized_requests=False,trust_remote_code=True --tasks gsm8k --num_fewshot ${LM_EVAL_NUM_FEWSHOT} --output_path ${OUTPUT_PATH}", + "accuracy_test_threshold": 0.88, "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1", "runner": "linux-atom-mi35x-1", }, @@ -295,10 +317,20 @@ jobs: "model_name": "gpt-oss-120b TP2", "model_path": "openai/gpt-oss-120b", "extra_args": "--tensor-parallel-size 2", - "accuracy_test_threshold": 0.38, - "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1", + "client_command": "lm_eval --model local-chat-completions --apply_chat_template --model_args model=${MODEL_PATH},base_url=http://127.0.0.1:${VLLM_PORT}/v1/chat/completions,num_concurrent=65,max_retries=1,max_gen_toks=2048,tokenized_requests=False,trust_remote_code=True --tasks gsm8k --num_fewshot ${LM_EVAL_NUM_FEWSHOT} --output_path ${OUTPUT_PATH}", + "accuracy_test_threshold": 0.88, + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1", "runner": "linux-atom-mi35x-4", }, + { + "toggle_env": "RUN_GLM4_7_FP8_MTP_TP8", + "model_name": "GLM-4.7-FP8 MTP TP8", + "model_path": "zai-org/GLM-4.7-FP8", + "extra_args": "--tensor-parallel-size 8 --speculative-config.method mtp --speculative-config.num_speculative_tokens 1", + "accuracy_test_threshold": 0.92, + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1", + "runner": "linux-atom-mi35x-8", + }, { "toggle_env": "RUN_GLM5_1_FP8_TP8", "model_name": "GLM-5.1-FP8 TP8", @@ -600,6 +632,7 @@ jobs: OOT_MODEL_NAME: ${{ matrix.model_name }} OOT_MODEL_PATH: ${{ matrix.model_path }} OOT_EXTRA_ARGS: ${{ matrix.extra_args }} + OOT_CLIENT_COMMAND: ${{ matrix.client_command || '' }} OOT_ENV_VARS: ${{ matrix.env_vars }} LM_EVAL_NUM_FEWSHOT: ${{ matrix.lm_eval_num_fewshot }} MAX_WAIT_RETRIES: "120" @@ -609,6 +642,7 @@ jobs: -e OOT_MODEL_NAME="${OOT_MODEL_NAME}" \ -e OOT_MODEL_PATH="${OOT_RESOLVED_MODEL_PATH:-$OOT_MODEL_PATH}" \ -e OOT_EXTRA_ARGS="${OOT_EXTRA_ARGS}" \ + -e OOT_CLIENT_COMMAND="${OOT_CLIENT_COMMAND}" \ -e OOT_ENV_VARS="${OOT_ENV_VARS}" \ -e OOT_DOCKER_IMAGE="${OOT_IMAGE_TAG}" \ -e GPU_NAME="${{ steps.gpu-info.outputs.gpu_name }}" \ diff --git a/.github/workflows/atom-vllm-benchmark.yaml b/.github/workflows/atom-vllm-benchmark.yaml index 455cdb0276..1261f2f362 100644 --- a/.github/workflows/atom-vllm-benchmark.yaml +++ b/.github/workflows/atom-vllm-benchmark.yaml @@ -7,82 +7,331 @@ concurrency: on: schedule: - # Nightly at 23:00 Beijing time (15:00 UTC) - - cron: '0 15 * * *' + # Mon–Sat 21:00 Beijing (13:00 UTC): Mon/Wed/Fri AWS P0; Tue/Thu AWS P1+Meta P0; Sat AWS P2+OOB (see load-models). + - cron: '0 13 * * 1-6' + # Sat/Sun 10:00 Beijing (02:00 UTC): Sat AWS P1+Meta P0; Sun Meta P1 + - cron: '0 2 * * 0,6' workflow_dispatch: inputs: - deepseek-r1-fp8: - description: "DeepSeek-R1 FP8 TP8" - type: boolean - default: false - deepseek-r1-mxfp4: - description: "DeepSeek-R1 MXFP4 TP8" - type: boolean - default: false - deepseek-v3-2-fp8: - description: "DeepSeek-V3.2 FP8 TP8" - type: boolean - default: false - gpt-oss-120b: - description: "gpt-oss-120b TP1" - type: boolean - default: false - glm-5-1-fp8-tp8: - description: "GLM-5.1-FP8 TP8" - type: boolean - default: false - kimi-k2-thinking-mxfp4-tp4: - description: "Kimi-K2-Thinking-MXFP4 TP4" - type: boolean - default: false - kimi-k2-thinking-mxfp4-tp8: - description: "Kimi-K2-Thinking-MXFP4 TP8" - type: boolean - default: false - kimi-k2-5-mxfp4-tp2: - description: "Kimi-K2.5-MXFP4 TP2" - type: boolean - default: false - kimi-k2-5-mxfp4-tp4: - description: "Kimi-K2.5-MXFP4 TP4" - type: boolean - default: false - kimi-k2-5-mxfp4-tp8: - description: "Kimi-K2.5-MXFP4 TP8" - type: boolean - default: false - qwen3-5-397b-a17b-tp4: - description: "Qwen3.5-397B-A17B TP4" - type: boolean - default: false - qwen3-5-397b-a17b: - description: "Qwen3.5-397B-A17B TP8" - type: boolean - default: false - qwen3-5-397b-a17b-fp8-tp4: - description: "Qwen3.5-397B-A17B-FP8 TP4" - type: boolean - default: false - qwen3-5-397b-a17b-fp8: - description: "Qwen3.5-397B-A17B-FP8 TP8" - type: boolean - default: false - qwen3-5-397b-a17b-mxfp4: - description: "Qwen3.5-397B-A17B-MXFP4 TP4" - type: boolean - default: false - qwen3-next-80b-a3b-instruct-fp8-tp1: - description: "Qwen3-Next-80B-A3B-Instruct-FP8 TP1" - type: boolean - default: false - qwen3-next-80b-a3b-instruct-fp8-tp2: - description: "Qwen3-Next-80B-A3B-Instruct-FP8 TP2" - type: boolean - default: false - qwen3-next-80b-a3b-instruct-fp8-tp4: - description: "Qwen3-Next-80B-A3B-Instruct-FP8 TP4" - type: boolean - default: false + benchmark_client: + description: "Benchmark client command. `InferenceMax bench` uses the repo benchmark script. `vLLM bench` uses the vLLM CLI benchmark command." + type: choice + default: InferenceMax bench + options: + - InferenceMax bench + - vLLM bench + model_slot_1: + description: "Manual selection slot 1: choose a model + TP size candidate." + type: choice + default: none + options: + - none + - DeepSeek-R1 FP8 TP8 (MET) + - DeepSeek-R1 MXFP4 TP8 (MET) + - gpt-oss-120b TP1 (MET) + - GLM-5.1-FP8 TP8 (OOB) + - Kimi-K2-Thinking-MXFP4 TP4 (MET) + - Kimi-K2-Thinking-MXFP4 TP8 (MET) + - Kimi-K2.5-MXFP4 TP4 (MET) + - Qwen3.5-397B-A17B-FP8 TP4 (OOB) + - Qwen3.5-397B-A17B-FP8 TP8 (OOB) + - Qwen3.5-397B-A17B TP8 (OOB) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP1 (MET) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP4 (MET) + - Qwen3-Next-80B-A3B-Instruct-FP8-MTP TP1 (AW) + - Qwen3-Next-80B-A3B-Instruct-FP8-MTP TP4 (AW) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP1 (AW) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP2 (AW) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP4 (AW) + - DeepSeek-V3.2 FP8 TP4 (AW) + - DeepSeek-V3.2 FP8 TP8 (AW) + - GLM-4.7-FP8 TP4 (AW) + - GLM-4.7-FP8 TP8 (AW) + - GLM-4.7-FP8 MTP TP4 (AW) + - GLM-4.7-FP8 MTP TP8 (AW) + - gpt-oss-120b TP1 (AW) + - gpt-oss-120b TP2 (AW) + - gpt-oss-120b TP8 (AW) + - Kimi-K2.5-MXFP4 TP4 (AW) + - Kimi-K2.5-MXFP4 TP8 (AW) + - MiniMax-M2.5 TP2 (AW) + - MiniMax-M2.5 TP4 (AW) + - MiniMax-M2.5 TP8 (AW) + model_slot_2: + description: "Manual selection slot 2: choose a model + TP size candidate." + type: choice + default: none + options: + - none + - DeepSeek-R1 FP8 TP8 (MET) + - DeepSeek-R1 MXFP4 TP8 (MET) + - gpt-oss-120b TP1 (MET) + - GLM-5.1-FP8 TP8 (OOB) + - Kimi-K2-Thinking-MXFP4 TP4 (MET) + - Kimi-K2-Thinking-MXFP4 TP8 (MET) + - Kimi-K2.5-MXFP4 TP4 (MET) + - Qwen3.5-397B-A17B-FP8 TP4 (OOB) + - Qwen3.5-397B-A17B-FP8 TP8 (OOB) + - Qwen3.5-397B-A17B TP8 (OOB) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP1 (MET) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP4 (MET) + - Qwen3-Next-80B-A3B-Instruct-FP8-MTP TP1 (AW) + - Qwen3-Next-80B-A3B-Instruct-FP8-MTP TP4 (AW) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP1 (AW) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP2 (AW) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP4 (AW) + - DeepSeek-V3.2 FP8 TP4 (AW) + - DeepSeek-V3.2 FP8 TP8 (AW) + - GLM-4.7-FP8 TP4 (AW) + - GLM-4.7-FP8 TP8 (AW) + - GLM-4.7-FP8 MTP TP4 (AW) + - GLM-4.7-FP8 MTP TP8 (AW) + - gpt-oss-120b TP1 (AW) + - gpt-oss-120b TP2 (AW) + - gpt-oss-120b TP8 (AW) + - Kimi-K2.5-MXFP4 TP4 (AW) + - Kimi-K2.5-MXFP4 TP8 (AW) + - MiniMax-M2.5 TP2 (AW) + - MiniMax-M2.5 TP4 (AW) + - MiniMax-M2.5 TP8 (AW) + model_slot_3: + description: "Manual selection slot 3: choose a model + TP size candidate." + type: choice + default: none + options: + - none + - DeepSeek-R1 FP8 TP8 (MET) + - DeepSeek-R1 MXFP4 TP8 (MET) + - gpt-oss-120b TP1 (MET) + - GLM-5.1-FP8 TP8 (OOB) + - Kimi-K2-Thinking-MXFP4 TP4 (MET) + - Kimi-K2-Thinking-MXFP4 TP8 (MET) + - Kimi-K2.5-MXFP4 TP4 (MET) + - Qwen3.5-397B-A17B-FP8 TP4 (OOB) + - Qwen3.5-397B-A17B-FP8 TP8 (OOB) + - Qwen3.5-397B-A17B TP8 (OOB) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP1 (MET) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP4 (MET) + - Qwen3-Next-80B-A3B-Instruct-FP8-MTP TP1 (AW) + - Qwen3-Next-80B-A3B-Instruct-FP8-MTP TP4 (AW) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP1 (AW) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP2 (AW) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP4 (AW) + - DeepSeek-V3.2 FP8 TP4 (AW) + - DeepSeek-V3.2 FP8 TP8 (AW) + - GLM-4.7-FP8 TP4 (AW) + - GLM-4.7-FP8 TP8 (AW) + - GLM-4.7-FP8 MTP TP4 (AW) + - GLM-4.7-FP8 MTP TP8 (AW) + - gpt-oss-120b TP1 (AW) + - gpt-oss-120b TP2 (AW) + - gpt-oss-120b TP8 (AW) + - Kimi-K2.5-MXFP4 TP4 (AW) + - Kimi-K2.5-MXFP4 TP8 (AW) + - MiniMax-M2.5 TP2 (AW) + - MiniMax-M2.5 TP4 (AW) + - MiniMax-M2.5 TP8 (AW) + model_slot_4: + description: "Manual selection slot 4: choose a model + TP size candidate." + type: choice + default: none + options: + - none + - DeepSeek-R1 FP8 TP8 (MET) + - DeepSeek-R1 MXFP4 TP8 (MET) + - gpt-oss-120b TP1 (MET) + - GLM-5.1-FP8 TP8 (OOB) + - Kimi-K2-Thinking-MXFP4 TP4 (MET) + - Kimi-K2-Thinking-MXFP4 TP8 (MET) + - Kimi-K2.5-MXFP4 TP4 (MET) + - Qwen3.5-397B-A17B-FP8 TP4 (OOB) + - Qwen3.5-397B-A17B-FP8 TP8 (OOB) + - Qwen3.5-397B-A17B TP8 (OOB) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP1 (MET) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP4 (MET) + - Qwen3-Next-80B-A3B-Instruct-FP8-MTP TP1 (AW) + - Qwen3-Next-80B-A3B-Instruct-FP8-MTP TP4 (AW) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP1 (AW) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP2 (AW) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP4 (AW) + - DeepSeek-V3.2 FP8 TP4 (AW) + - DeepSeek-V3.2 FP8 TP8 (AW) + - GLM-4.7-FP8 TP4 (AW) + - GLM-4.7-FP8 TP8 (AW) + - GLM-4.7-FP8 MTP TP4 (AW) + - GLM-4.7-FP8 MTP TP8 (AW) + - gpt-oss-120b TP1 (AW) + - gpt-oss-120b TP2 (AW) + - gpt-oss-120b TP8 (AW) + - Kimi-K2.5-MXFP4 TP4 (AW) + - Kimi-K2.5-MXFP4 TP8 (AW) + - MiniMax-M2.5 TP2 (AW) + - MiniMax-M2.5 TP4 (AW) + - MiniMax-M2.5 TP8 (AW) + model_slot_5: + description: "Manual selection slot 5: choose a model + TP size candidate." + type: choice + default: none + options: + - none + - DeepSeek-R1 FP8 TP8 (MET) + - DeepSeek-R1 MXFP4 TP8 (MET) + - gpt-oss-120b TP1 (MET) + - GLM-5.1-FP8 TP8 (OOB) + - Kimi-K2-Thinking-MXFP4 TP4 (MET) + - Kimi-K2-Thinking-MXFP4 TP8 (MET) + - Kimi-K2.5-MXFP4 TP4 (MET) + - Qwen3.5-397B-A17B-FP8 TP4 (OOB) + - Qwen3.5-397B-A17B-FP8 TP8 (OOB) + - Qwen3.5-397B-A17B TP8 (OOB) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP1 (MET) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP4 (MET) + - Qwen3-Next-80B-A3B-Instruct-FP8-MTP TP1 (AW) + - Qwen3-Next-80B-A3B-Instruct-FP8-MTP TP4 (AW) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP1 (AW) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP2 (AW) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP4 (AW) + - DeepSeek-V3.2 FP8 TP4 (AW) + - DeepSeek-V3.2 FP8 TP8 (AW) + - GLM-4.7-FP8 TP4 (AW) + - GLM-4.7-FP8 TP8 (AW) + - GLM-4.7-FP8 MTP TP4 (AW) + - GLM-4.7-FP8 MTP TP8 (AW) + - gpt-oss-120b TP1 (AW) + - gpt-oss-120b TP2 (AW) + - gpt-oss-120b TP8 (AW) + - Kimi-K2.5-MXFP4 TP4 (AW) + - Kimi-K2.5-MXFP4 TP8 (AW) + - MiniMax-M2.5 TP2 (AW) + - MiniMax-M2.5 TP4 (AW) + - MiniMax-M2.5 TP8 (AW) + model_slot_6: + description: "Manual selection slot 6: choose a model + TP size candidate." + type: choice + default: none + options: + - none + - DeepSeek-R1 FP8 TP8 (MET) + - DeepSeek-R1 MXFP4 TP8 (MET) + - gpt-oss-120b TP1 (MET) + - GLM-5.1-FP8 TP8 (OOB) + - Kimi-K2-Thinking-MXFP4 TP4 (MET) + - Kimi-K2-Thinking-MXFP4 TP8 (MET) + - Kimi-K2.5-MXFP4 TP4 (MET) + - Qwen3.5-397B-A17B-FP8 TP4 (OOB) + - Qwen3.5-397B-A17B-FP8 TP8 (OOB) + - Qwen3.5-397B-A17B TP8 (OOB) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP1 (MET) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP4 (MET) + - Qwen3-Next-80B-A3B-Instruct-FP8-MTP TP1 (AW) + - Qwen3-Next-80B-A3B-Instruct-FP8-MTP TP4 (AW) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP1 (AW) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP2 (AW) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP4 (AW) + - DeepSeek-V3.2 FP8 TP4 (AW) + - DeepSeek-V3.2 FP8 TP8 (AW) + - GLM-4.7-FP8 TP4 (AW) + - GLM-4.7-FP8 TP8 (AW) + - GLM-4.7-FP8 MTP TP4 (AW) + - GLM-4.7-FP8 MTP TP8 (AW) + - gpt-oss-120b TP1 (AW) + - gpt-oss-120b TP2 (AW) + - gpt-oss-120b TP8 (AW) + - Kimi-K2.5-MXFP4 TP4 (AW) + - Kimi-K2.5-MXFP4 TP8 (AW) + - MiniMax-M2.5 TP2 (AW) + - MiniMax-M2.5 TP4 (AW) + - MiniMax-M2.5 TP8 (AW) + model_slot_7: + description: "Manual selection slot 7: choose a model + TP size candidate." + type: choice + default: none + options: + - none + - DeepSeek-R1 FP8 TP8 (MET) + - DeepSeek-R1 MXFP4 TP8 (MET) + - gpt-oss-120b TP1 (MET) + - GLM-5.1-FP8 TP8 (OOB) + - Kimi-K2-Thinking-MXFP4 TP4 (MET) + - Kimi-K2-Thinking-MXFP4 TP8 (MET) + - Kimi-K2.5-MXFP4 TP4 (MET) + - Qwen3.5-397B-A17B-FP8 TP4 (OOB) + - Qwen3.5-397B-A17B-FP8 TP8 (OOB) + - Qwen3.5-397B-A17B TP8 (OOB) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP1 (MET) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP4 (MET) + - Qwen3-Next-80B-A3B-Instruct-FP8-MTP TP1 (AW) + - Qwen3-Next-80B-A3B-Instruct-FP8-MTP TP4 (AW) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP1 (AW) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP2 (AW) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP4 (AW) + - DeepSeek-V3.2 FP8 TP4 (AW) + - DeepSeek-V3.2 FP8 TP8 (AW) + - GLM-4.7-FP8 TP4 (AW) + - GLM-4.7-FP8 TP8 (AW) + - GLM-4.7-FP8 MTP TP4 (AW) + - GLM-4.7-FP8 MTP TP8 (AW) + - gpt-oss-120b TP1 (AW) + - gpt-oss-120b TP2 (AW) + - gpt-oss-120b TP8 (AW) + - Kimi-K2.5-MXFP4 TP4 (AW) + - Kimi-K2.5-MXFP4 TP8 (AW) + - MiniMax-M2.5 TP2 (AW) + - MiniMax-M2.5 TP4 (AW) + - MiniMax-M2.5 TP8 (AW) + model_slot_8: + description: "Manual selection slot 8: choose a model + TP size candidate." + type: choice + default: none + options: + - none + - DeepSeek-R1 FP8 TP8 (MET) + - DeepSeek-R1 MXFP4 TP8 (MET) + - gpt-oss-120b TP1 (MET) + - GLM-5.1-FP8 TP8 (OOB) + - Kimi-K2-Thinking-MXFP4 TP4 (MET) + - Kimi-K2-Thinking-MXFP4 TP8 (MET) + - Kimi-K2.5-MXFP4 TP4 (MET) + - Qwen3.5-397B-A17B-FP8 TP4 (OOB) + - Qwen3.5-397B-A17B-FP8 TP8 (OOB) + - Qwen3.5-397B-A17B TP8 (OOB) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP1 (MET) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP4 (MET) + - Qwen3-Next-80B-A3B-Instruct-FP8-MTP TP1 (AW) + - Qwen3-Next-80B-A3B-Instruct-FP8-MTP TP4 (AW) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP1 (AW) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP2 (AW) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP4 (AW) + - DeepSeek-V3.2 FP8 TP4 (AW) + - DeepSeek-V3.2 FP8 TP8 (AW) + - GLM-4.7-FP8 TP4 (AW) + - GLM-4.7-FP8 TP8 (AW) + - GLM-4.7-FP8 MTP TP4 (AW) + - GLM-4.7-FP8 MTP TP8 (AW) + - gpt-oss-120b TP1 (AW) + - gpt-oss-120b TP2 (AW) + - gpt-oss-120b TP8 (AW) + - Kimi-K2.5-MXFP4 TP4 (AW) + - Kimi-K2.5-MXFP4 TP8 (AW) + - MiniMax-M2.5 TP2 (AW) + - MiniMax-M2.5 TP4 (AW) + - MiniMax-M2.5 TP8 (AW) + isl_osl_pairs: + description: | + Benchmark ISL/OSL pairs. + Input as semicolon-separated `input_length,output_length` pairs. + Example: 1024,1024;1024,8192;8192,1024 + Additional copy-ready example: 1000,100;5000,500;10000,1000 + type: string + default: "1024,1024;1024,8192;8192,1024" + concurrency_values: + description: "Benchmark concurrency values. Use comma-separated values from 4,8,16,32,64,128,256,512. The workflow will benchmark exactly the values you provide." + type: string + default: "4,8,16,32,64,128,256,512" + random_range_ratios: + description: "Benchmark random range ratios. Use comma-separated values like 0.8,1.0." + type: string + default: "0.8" oot_image: description: "Optional OOT benchmark image override. Leave empty to resolve main to the same-digest published nightly tag behind vllm-latest, or rebuild from the selected non-main branch." type: string @@ -91,15 +340,10 @@ on: description: "Upload benchmark results to the dashboard" type: boolean default: true - param_lists: - description: | - "Benchmark parameter lists. - Input as a single or multiple sets (comma-separated, semicolon between sets), - format: input_length,output_length,concurrency,random_range_ratio. - Example (single set): 1024,1024,64,0.8 - Example (multiple sets): 1024,1024,64,0.8;2048,1024,32,0.7" - type: string - default: "1024,1024,4,0.8;1024,1024,8,0.8;1024,1024,16,0.8;1024,1024,32,0.8;1024,1024,64,0.8;1024,8192,4,0.8;1024,8192,8,0.8;1024,8192,16,0.8;1024,8192,32,0.8;1024,8192,64,0.8;8192,1024,4,0.8;8192,1024,8,0.8;8192,1024,16,0.8;8192,1024,32,0.8;8192,1024,64,0.8" + upload_to_custom_dashboard: + description: "Upload benchmark results to the LLM-Booster dashboard" + type: boolean + default: true jobs: wait-for-older-manual-runs: @@ -231,6 +475,8 @@ jobs: oot_base_image: ${{ steps.resolve.outputs.oot_base_image }} oot_image_source: ${{ steps.resolve.outputs.oot_image_source }} publish_to_dashboard: ${{ steps.resolve.outputs.publish_to_dashboard }} + upload_to_custom_dashboard: ${{ steps.resolve.outputs.upload_to_custom_dashboard }} + benchmark_client: ${{ steps.resolve.outputs.benchmark_client }} steps: - name: Checkout selected branch uses: actions/checkout@v6 @@ -240,6 +486,8 @@ jobs: env: INPUT_OOT_IMAGE: ${{ inputs.oot_image || '' }} INPUT_PUBLISH_TO_DASHBOARD: ${{ inputs.publish_to_dashboard }} + INPUT_UPLOAD_TO_CUSTOM_DASHBOARD: ${{ inputs.upload_to_custom_dashboard }} + INPUT_BENCHMARK_CLIENT: ${{ inputs.benchmark_client || 'InferenceMax bench' }} run: | set -euo pipefail @@ -261,10 +509,14 @@ jobs: if [[ "${GITHUB_EVENT_NAME}" == "schedule" ]]; then PUBLISH_TO_DASHBOARD=true + UPLOAD_TO_CUSTOM_DASHBOARD=true else PUBLISH_TO_DASHBOARD="${INPUT_PUBLISH_TO_DASHBOARD:-true}" + UPLOAD_TO_CUSTOM_DASHBOARD="${INPUT_UPLOAD_TO_CUSTOM_DASHBOARD:-true}" fi + BENCHMARK_CLIENT="${INPUT_BENCHMARK_CLIENT:-InferenceMax bench}" + mapfile -t VLLM_META < <(python3 - <<'PY' import re from pathlib import Path @@ -336,6 +588,8 @@ jobs: echo "oot_base_image=${OOT_BASE_IMAGE}" echo "oot_image_source=${OOT_IMAGE_SOURCE}" echo "publish_to_dashboard=${PUBLISH_TO_DASHBOARD}" + echo "upload_to_custom_dashboard=${UPLOAD_TO_CUSTOM_DASHBOARD}" + echo "benchmark_client=${BENCHMARK_CLIENT}" } >> "$GITHUB_OUTPUT" printf '### OOT benchmark source\n- Repository: `%s`\n- Ref: `%s`\n- Image mode: `%s`\n' \ @@ -390,53 +644,140 @@ jobs: esac fi - printf -- '- Upload to dashboard: `%s`\n' \ - "${PUBLISH_TO_DASHBOARD}" >> "$GITHUB_STEP_SUMMARY" + printf -- '- Benchmark client: `%s`\n- Upload to dashboard: `%s`\n- Upload to custom dashboard: `%s`\n' \ + "${BENCHMARK_CLIENT}" \ + "${PUBLISH_TO_DASHBOARD}" \ + "${UPLOAD_TO_CUSTOM_DASHBOARD}" >> "$GITHUB_STEP_SUMMARY" if [[ "${PUBLISH_TO_DASHBOARD}" != "true" ]]; then printf '\nDashboard upload is disabled by workflow input, so results are kept in artifacts and run summaries only.\n' \ >> "$GITHUB_STEP_SUMMARY" fi + if [[ "${UPLOAD_TO_CUSTOM_DASHBOARD}" != "true" ]]; then + printf 'Custom dashboard payload upload is disabled by workflow input.\n' \ + >> "$GITHUB_STEP_SUMMARY" + fi + parse-param-lists: - name: Parse parameter lists + name: Parse benchmark parameter inputs needs: [wait-for-older-manual-runs] runs-on: ubuntu-latest outputs: - matrix_json: ${{ steps.parse-param-lists.outputs.matrix_json }} - param_lists: ${{ steps.parse-param-lists.outputs.param_lists }} + matrix_json: ${{ steps.parse-benchmark-params.outputs.matrix_json }} + isl_osl_pairs: ${{ steps.parse-benchmark-params.outputs.isl_osl_pairs }} + concurrency_values: ${{ steps.parse-benchmark-params.outputs.concurrency_values }} + random_range_ratios: ${{ steps.parse-benchmark-params.outputs.random_range_ratios }} steps: - - name: Parse parameter lists - id: parse-param-lists + - name: Parse benchmark parameter inputs + id: parse-benchmark-params + env: + ISL_OSL_PAIRS: ${{ inputs.isl_osl_pairs || '1024,1024;1024,8192;8192,1024' }} + CONCURRENCY_VALUES: ${{ inputs.concurrency_values || '4,8,16,32,64,128,256,512' }} + RANDOM_RANGE_RATIOS: ${{ inputs.random_range_ratios || '0.8' }} run: | - PARAM_LISTS="${{ inputs.param_lists || '1024,1024,4,0.8;1024,1024,8,0.8;1024,1024,16,0.8;1024,1024,32,0.8;1024,1024,64,0.8;1024,8192,4,0.8;1024,8192,8,0.8;1024,8192,16,0.8;1024,8192,32,0.8;1024,8192,64,0.8;8192,1024,4,0.8;8192,1024,8,0.8;8192,1024,16,0.8;8192,1024,32,0.8;8192,1024,64,0.8' }}" - echo "Using param_lists: ${PARAM_LISTS}" - printf 'param_lists=%s\n' "${PARAM_LISTS}" >> "$GITHUB_OUTPUT" - IFS=';' read -ra SETS <<< "${PARAM_LISTS}" - MATRIX_JSON="[" - SEP="" - for SET in "${SETS[@]}"; do - IFS=',' read -ra PARAMS <<< "$SET" - INPUT_LEN="${PARAMS[0]}" - OUTPUT_LEN="${PARAMS[1]}" - CONCURRENCY="${PARAMS[2]}" - RANDOM_RANGE_RATIO="${PARAMS[3]}" - case "${CONCURRENCY}" in - 4|8|16|32|64|128|256) ;; - *) - echo "Unsupported concurrency: ${CONCURRENCY}. Allowed values: 4,8,16,32,64,128,256" - exit 1 - ;; - esac - MATRIX_JSON="${MATRIX_JSON}${SEP}{\"input_length\":${INPUT_LEN},\"output_length\":${OUTPUT_LEN},\"concurrency\":${CONCURRENCY},\"random_range_ratio\":${RANDOM_RANGE_RATIO}}" - SEP="," - done - MATRIX_JSON="${MATRIX_JSON}]" - echo "matrix_json=${MATRIX_JSON}" >> "$GITHUB_OUTPUT" + set -euo pipefail + echo "Using isl_osl_pairs: ${ISL_OSL_PAIRS}" + echo "Using concurrency_values: ${CONCURRENCY_VALUES}" + echo "Using random_range_ratios: ${RANDOM_RANGE_RATIOS:-}" + printf 'isl_osl_pairs=%s\n' "${ISL_OSL_PAIRS}" >> "$GITHUB_OUTPUT" + printf 'concurrency_values=%s\n' "${CONCURRENCY_VALUES}" >> "$GITHUB_OUTPUT" + printf 'random_range_ratios=%s\n' "${RANDOM_RANGE_RATIOS}" >> "$GITHUB_OUTPUT" + python3 - <<'PY' >> "$GITHUB_OUTPUT" + import json + import os + + allowed_concurrencies = {4, 8, 16, 32, 64, 128, 256, 512} + + def dedupe_preserve_order(values): + deduped = [] + seen = set() + for value in values: + marker = ( + json.dumps(value, sort_keys=True) if isinstance(value, dict) else value + ) + if marker in seen: + continue + seen.add(marker) + deduped.append(value) + return deduped + + pairs = [] + for raw_pair in os.environ["ISL_OSL_PAIRS"].split(";"): + pair = raw_pair.strip() + if not pair: + continue + parts = [part.strip() for part in pair.split(",")] + if len(parts) != 2 or not all(part.isdigit() for part in parts): + raise SystemExit( + f"Invalid ISL/OSL pair {pair!r}. Use semicolon-separated pairs like 1024,1024;2048,1024." + ) + pairs.append( + { + "input_length": int(parts[0]), + "output_length": int(parts[1]), + } + ) + if not pairs: + raise SystemExit("At least one ISL/OSL pair is required.") + pairs = dedupe_preserve_order(pairs) + + concurrencies = [] + for raw_value in os.environ["CONCURRENCY_VALUES"].split(","): + value = raw_value.strip() + if not value: + continue + if not value.isdigit(): + raise SystemExit( + f"Invalid concurrency {value!r}. Allowed values: 4,8,16,32,64,128,256,512" + ) + parsed = int(value) + if parsed not in allowed_concurrencies: + raise SystemExit( + f"Unsupported concurrency: {parsed}. Allowed values: 4,8,16,32,64,128,256,512" + ) + concurrencies.append(parsed) + if not concurrencies: + raise SystemExit("At least one concurrency value is required.") + concurrencies = dedupe_preserve_order(concurrencies) + + ratio_text = os.environ["RANDOM_RANGE_RATIOS"].strip() + if not ratio_text: + ratios = ["0.8"] + else: + ratios = [] + for raw_ratio in ratio_text.split(","): + ratio = raw_ratio.strip() + if not ratio: + continue + try: + float(ratio) + except ValueError as exc: + raise SystemExit( + f"Invalid random range ratio {ratio!r}. Use values like 0.8 or 1.0." + ) from exc + ratios.append(ratio) + ratios = dedupe_preserve_order(ratios) or ["0.8"] + + matrix = [] + for pair in pairs: + for concurrency in concurrencies: + for ratio in ratios: + matrix.append( + { + "input_length": pair["input_length"], + "output_length": pair["output_length"], + "concurrency": concurrency, + "random_range_ratio": ratio, + } + ) - - name: Print matrix JSON + print(f"matrix_json={json.dumps(matrix, separators=(',', ':'))}") + PY + + - name: Print benchmark parameter matrix JSON run: | - echo "Matrix JSON: ${{ steps.parse-param-lists.outputs.matrix_json }}" + echo "Matrix JSON: ${{ steps.parse-benchmark-params.outputs.matrix_json }}" load-models: name: Load OOT model configs @@ -453,113 +794,249 @@ jobs: - uses: actions/checkout@v6 - id: load env: - ENABLE_DEEPSEEK_R1_FP8: ${{ inputs.deepseek-r1-fp8 }} - ENABLE_DEEPSEEK_R1_MXFP4: ${{ inputs.deepseek-r1-mxfp4 }} - ENABLE_DEEPSEEK_V3_2_FP8: ${{ inputs.deepseek-v3-2-fp8 }} - ENABLE_GPT_OSS_120B: ${{ inputs.gpt-oss-120b }} - ENABLE_GLM_5_1_FP8_TP8: ${{ inputs.glm-5-1-fp8-tp8 }} - ENABLE_KIMI_K2_TP4: ${{ inputs.kimi-k2-thinking-mxfp4-tp4 }} - ENABLE_KIMI_K2_TP8: ${{ inputs.kimi-k2-thinking-mxfp4-tp8 }} - ENABLE_KIMI_K25_TP2: ${{ inputs.kimi-k2-5-mxfp4-tp2 }} - ENABLE_KIMI_K25_TP4: ${{ inputs.kimi-k2-5-mxfp4-tp4 }} - ENABLE_KIMI_K25_TP8: ${{ inputs.kimi-k2-5-mxfp4-tp8 }} - ENABLE_QWEN3_5_397B_A17B_TP4: ${{ inputs.qwen3-5-397b-a17b-tp4 }} - ENABLE_QWEN3_5_397B_A17B: ${{ inputs.qwen3-5-397b-a17b }} - ENABLE_QWEN3_5_397B_A17B_FP8_TP4: ${{ inputs.qwen3-5-397b-a17b-fp8-tp4 }} - ENABLE_QWEN3_5_397B_A17B_FP8: ${{ inputs.qwen3-5-397b-a17b-fp8 }} - ENABLE_QWEN3_5_397B_A17B_MXFP4: ${{ inputs.qwen3-5-397b-a17b-mxfp4 }} - ENABLE_QWEN3_NEXT_80B_A3B_INSTRUCT_FP8_TP1: ${{ inputs.qwen3-next-80b-a3b-instruct-fp8-tp1 }} - ENABLE_QWEN3_NEXT_80B_A3B_INSTRUCT_FP8_TP2: ${{ inputs.qwen3-next-80b-a3b-instruct-fp8-tp2 }} - ENABLE_QWEN3_NEXT_80B_A3B_INSTRUCT_FP8_TP4: ${{ inputs.qwen3-next-80b-a3b-instruct-fp8-tp4 }} + MODEL_SLOT_1: ${{ inputs.model_slot_1 || 'none' }} + MODEL_SLOT_2: ${{ inputs.model_slot_2 || 'none' }} + MODEL_SLOT_3: ${{ inputs.model_slot_3 || 'none' }} + MODEL_SLOT_4: ${{ inputs.model_slot_4 || 'none' }} + MODEL_SLOT_5: ${{ inputs.model_slot_5 || 'none' }} + MODEL_SLOT_6: ${{ inputs.model_slot_6 || 'none' }} + MODEL_SLOT_7: ${{ inputs.model_slot_7 || 'none' }} + MODEL_SLOT_8: ${{ inputs.model_slot_8 || 'none' }} GITHUB_TOKEN: ${{ github.token }} run: | + set -euo pipefail + + if [[ "${GITHUB_EVENT_NAME}" == "schedule" ]]; then + SCHEDULE_CREATED_AT="$( + python3 - <<'PY' + import json + import os + from urllib.error import HTTPError, URLError + from urllib.request import Request, urlopen + + api_url = os.environ.get("GITHUB_API_URL", "https://api.github.com") + repository = os.environ["GITHUB_REPOSITORY"] + run_id = os.environ["GITHUB_RUN_ID"] + token = os.environ.get("GITHUB_TOKEN", "") + request = Request( + f"{api_url}/repos/{repository}/actions/runs/{run_id}", + headers={ + "Accept": "application/vnd.github+json", + "Authorization": f"Bearer {token}", + "User-Agent": "atom-vllm-benchmark-nightly-rotation", + }, + ) + try: + with urlopen(request, timeout=30) as response: + payload = json.load(response) + except HTTPError as exc: + detail = exc.read().decode("utf-8", errors="replace").strip() + raise RuntimeError( + f"Failed to fetch workflow run metadata: HTTP {exc.code} {detail}" + ) from exc + except URLError as exc: + raise RuntimeError( + f"Failed to reach GitHub API for workflow run metadata: {exc}" + ) from exc + + created_at = payload.get("created_at") + if not created_at: + raise RuntimeError("Workflow run metadata did not include created_at") + print(created_at) + PY + )" + export SCHEDULE_CREATED_AT + else + export SCHEDULE_CREATED_AT="" + fi + python3 - <<'PY' >> "$GITHUB_OUTPUT" import json import os import sys - from datetime import date, datetime, timedelta, timezone + from datetime import datetime, timedelta, timezone from pathlib import Path - from urllib.error import HTTPError, URLError - from urllib.request import Request, urlopen - models = json.loads( - Path(".github/benchmark/oot_benchmark_models.json").read_text(encoding="utf-8") + none_choice = "none" + + catalog = json.loads( + Path(".github/benchmark/oot_benchmark_models.json").read_text( + encoding="utf-8" + ) ) + families = catalog.get("families") + if not isinstance(families, list): + raise SystemExit("Benchmark catalog must define a 'families' list.") + + def flatten_variant(family, variant): + flattened = { + key: value + for key, value in family.items() + if key not in {"choice_label", "variants"} + } + flattened.update(variant) + return flattened + + selected_group = "" event = os.environ["GITHUB_EVENT_NAME"] + selected = [] - manual_toggles = { - "deepseek-r1-fp8": os.environ.get("ENABLE_DEEPSEEK_R1_FP8", "").lower() == "true", - "deepseek-r1-mxfp4": os.environ.get("ENABLE_DEEPSEEK_R1_MXFP4", "").lower() == "true", - "deepseek-v3-2-fp8": os.environ.get("ENABLE_DEEPSEEK_V3_2_FP8", "").lower() == "true", - "gpt-oss-120b": os.environ.get("ENABLE_GPT_OSS_120B", "").lower() == "true", - "glm-5-1-fp8-tp8": os.environ.get("ENABLE_GLM_5_1_FP8_TP8", "").lower() == "true", - "kimi-k2-thinking-mxfp4-tp4": os.environ.get("ENABLE_KIMI_K2_TP4", "").lower() == "true", - "kimi-k2-thinking-mxfp4-tp8": os.environ.get("ENABLE_KIMI_K2_TP8", "").lower() == "true", - "kimi-k25-mxfp4-tp2": os.environ.get("ENABLE_KIMI_K25_TP2", "").lower() == "true", - "kimi-k25-mxfp4-tp4": os.environ.get("ENABLE_KIMI_K25_TP4", "").lower() == "true", - "kimi-k25-mxfp4-tp8": os.environ.get("ENABLE_KIMI_K25_TP8", "").lower() == "true", - "qwen3-5-397b-a17b-tp4": os.environ.get("ENABLE_QWEN3_5_397B_A17B_TP4", "").lower() == "true", - "qwen3-5-397b-a17b": os.environ.get("ENABLE_QWEN3_5_397B_A17B", "").lower() == "true", - "qwen3-5-397b-a17b-fp8-tp4": os.environ.get("ENABLE_QWEN3_5_397B_A17B_FP8_TP4", "").lower() == "true", - "qwen3-5-397b-a17b-fp8": os.environ.get("ENABLE_QWEN3_5_397B_A17B_FP8", "").lower() == "true", - "qwen3-5-397b-a17b-mxfp4": os.environ.get("ENABLE_QWEN3_5_397B_A17B_MXFP4", "").lower() == "true", - "qwen3-next-80b-a3b-instruct-fp8-tp1": os.environ.get("ENABLE_QWEN3_NEXT_80B_A3B_INSTRUCT_FP8_TP1", "").lower() == "true", - "qwen3-next-80b-a3b-instruct-fp8-tp2": os.environ.get("ENABLE_QWEN3_NEXT_80B_A3B_INSTRUCT_FP8_TP2", "").lower() == "true", - "qwen3-next-80b-a3b-instruct-fp8-tp4": os.environ.get("ENABLE_QWEN3_NEXT_80B_A3B_INSTRUCT_FP8_TP4", "").lower() == "true", - } + all_variants = [] + for family in families: + for variant in family["variants"]: + all_variants.append(flatten_variant(family, variant)) - def fetch_run_created_at() -> datetime: - api_url = os.environ.get("GITHUB_API_URL", "https://api.github.com") - repository = os.environ["GITHUB_REPOSITORY"] - run_id = os.environ["GITHUB_RUN_ID"] - token = os.environ.get("GITHUB_TOKEN", "") - request = Request( - f"{api_url}/repos/{repository}/actions/runs/{run_id}", - headers={ - "Accept": "application/vnd.github+json", - "Authorization": f"Bearer {token}", - "User-Agent": "atom-vllm-benchmark-nightly-rotation", - }, - ) - try: - with urlopen(request, timeout=30) as response: - payload = json.load(response) - except HTTPError as exc: - detail = exc.read().decode("utf-8", errors="replace").strip() - raise RuntimeError( - f"Failed to fetch workflow run metadata: HTTP {exc.code} {detail}" - ) from exc - except URLError as exc: - raise RuntimeError( - f"Failed to reach GitHub API for workflow run metadata: {exc}" - ) from exc + def select_by_display(variants, displays): + return [ + m for m in variants if str(m.get("display", "")) in displays + ] - created_at = payload.get("created_at") - if not created_at: - raise RuntimeError("Workflow run metadata did not include created_at") - return datetime.fromisoformat(created_at.replace("Z", "+00:00")) + AWS_P0 = { + "MiniMax-M2.5 TP2 (AW)", + "MiniMax-M2.5 TP4 (AW)", + "DeepSeek-V3.2 FP8 TP8 (AW)", + "gpt-oss-120b TP2 (AW)", + "gpt-oss-120b TP8 (AW)", + } + AWS_P1 = { + "GLM-4.7-FP8 TP4 (AW)", + "Kimi-K2.5-MXFP4 TP4 (AW)", + "Kimi-K2.5-MXFP4 TP8 (AW)", + "GLM-4.7-FP8 MTP TP4 (AW)", + "GLM-4.7-FP8 MTP TP8 (AW)", + "Qwen3-Next-80B-A3B-Instruct-FP8-MTP TP1 (AW)", + "Qwen3-Next-80B-A3B-Instruct-FP8-MTP TP4 (AW)", + } + AWS_P2 = { + "DeepSeek-V3.2 FP8 TP4 (AW)", + "GLM-4.7-FP8 TP8 (AW)", + "Qwen3-Next-80B-A3B-Instruct-FP8 TP1 (AW)", + "Qwen3-Next-80B-A3B-Instruct-FP8 TP2 (AW)", + "Qwen3-Next-80B-A3B-Instruct-FP8 TP4 (AW)", + "gpt-oss-120b TP1 (AW)", + } + META_P0 = { + "Kimi-K2.5-MXFP4 TP4 (MET)", + "gpt-oss-120b TP1 (MET)", + } + META_P1 = { + "Kimi-K2-Thinking-MXFP4 TP4 (MET)", + "Kimi-K2-Thinking-MXFP4 TP8 (MET)", + "Qwen3-Next-80B-A3B-Instruct-FP8 TP1 (MET)", + "Qwen3-Next-80B-A3B-Instruct-FP8 TP4 (MET)", + "Qwen3.5-397B-A17B-FP8 TP4 (OOB)", + "DeepSeek-R1 FP8 TP8 (MET)", + "DeepSeek-R1 MXFP4 TP8 (MET)", + } + OOB_DISPLAYS = { + str(m.get("display", "")) + for m in all_variants + if str(m.get("display", "")).endswith("(OOB)") + } - selected_group = "" if event == "schedule": + created_at = os.environ.get("SCHEDULE_CREATED_AT", "") + if not created_at: + raise SystemExit( + "Scheduled selection requires a run creation timestamp." + ) beijing_tz = timezone(timedelta(hours=8)) - beijing_date = fetch_run_created_at().astimezone(beijing_tz).date() - # Rotation anchor: 2026-05-06 runs group A, then B, then C. - anchor_date = date(2026, 5, 6) - rotation = ("A", "B", "C") - selected_group = rotation[(beijing_date - anchor_date).days % len(rotation)] - selected = [ - model for model in models if model.get("nightly_group") == selected_group - ] - if not selected: + run_date = datetime.fromisoformat(created_at.replace("Z", "+00:00")) + beijing_dt = run_date.astimezone(beijing_tz) + beijing_weekday = beijing_dt.weekday() + beijing_hour = beijing_dt.hour + + # Beijing weekday: Mon=0 .. Sun=6. Hour buckets use 15:00 to split + # ~10:00 daytime vs ~21:00 evening scheduled runs. + # Mon/Wed/Fri evening: AWS P0. Tue/Thu evening: AWS P1 + Meta P0. + # Fri evening: AWS P0. Sat daytime (~10:00): AWS P1 + Meta P0. + # Sat evening (~21:00): AWS P2 + OOB. Sun daytime (~10:00): Meta P1. + is_daytime_slot = beijing_hour < 15 + + if beijing_weekday == 5: # Saturday + if is_daytime_slot: + selected_group = "AWS-P1+META-P0" + target_displays = AWS_P1 | META_P0 + else: + selected_group = "AWS-P2+OOB" + target_displays = AWS_P2 | OOB_DISPLAYS + elif beijing_weekday == 6: # Sunday + if is_daytime_slot: + selected_group = "META-P1" + target_displays = META_P1 + else: + selected_group = "SKIP" + target_displays = set() + elif beijing_weekday in (0, 2, 4): # Mon / Wed / Fri + if is_daytime_slot: + selected_group = "SKIP" + target_displays = set() + else: + selected_group = "AWS-P0" + target_displays = AWS_P0 + elif beijing_weekday in (1, 3): # Tue / Thu + if is_daytime_slot: + selected_group = "SKIP" + target_displays = set() + else: + selected_group = "AWS-P1+META-P0" + target_displays = AWS_P1 | META_P0 + else: + selected_group = "SKIP" + target_displays = set() + + if target_displays is None: + selected = list(all_variants) + else: + selected = select_by_display(all_variants, target_displays) + missing = target_displays - { + str(m.get("display", "")) for m in selected + } + if missing: + print( + "Warning: scheduled targets missing from catalog: " + + ", ".join(sorted(missing)), + file=sys.stderr, + ) + else: + variant_map = {} + for family in families: + for variant in family["variants"]: + variant_map[str(variant["display"])] = flatten_variant( + family, variant + ) + + seen_prefixes = set() + for slot_idx in range(1, 9): + label = os.environ.get(f"MODEL_SLOT_{slot_idx}", "").strip() + if not label or label == none_choice: + continue + selected_variant = variant_map.get(label) + if selected_variant is None: + available = ", ".join(sorted(variant_map)) + raise SystemExit( + f"Unknown benchmark model variant choice {label!r}. Available choices: {available}" + ) + prefix = str(selected_variant["prefix"]) + if prefix in seen_prefixes: + continue + selected.append(selected_variant) + seen_prefixes.add(prefix) + + if selected_group: + print( + f"Scheduled nightly benchmark group: {selected_group}", + file=sys.stderr, + ) + if not selected: + print("No benchmark model variants were selected.", file=sys.stderr) + else: + print("Selected benchmark model variants:", file=sys.stderr) + for model in selected: print( - f"No models configured for nightly benchmark group {selected_group}.", + f" - {model['display']} ({model['prefix']})", file=sys.stderr, ) - sys.exit(1) - else: - selected = [ - model for model in models if manual_toggles.get(model["prefix"], False) - ] print(f"models_json={json.dumps(selected, separators=(',', ':'))}") print(f"selected_group={selected_group}") @@ -567,11 +1044,14 @@ jobs: PY - name: Print selected models + env: + SELECTED_GROUP: ${{ steps.load.outputs.selected_group }} + MODELS_JSON: ${{ steps.load.outputs.models_json }} run: | - if [ -n "${{ steps.load.outputs.selected_group }}" ]; then - echo "Scheduled nightly benchmark group: ${{ steps.load.outputs.selected_group }}" + if [[ -n "${SELECTED_GROUP}" ]]; then + echo "Scheduled nightly benchmark group: ${SELECTED_GROUP}" fi - echo "Selected models: ${{ steps.load.outputs.models_json }}" + printf 'Selected models: %s\n' "${MODELS_JSON}" build-benchmark-matrix: name: Build OOT benchmark matrix @@ -600,19 +1080,37 @@ jobs: include = [] for model in models: + prefix = str(model.get("prefix", "")) + is_aw_model = "-aw-" in prefix + aw_pairs = ( + (1000, 100), + (5000, 500), + (10000, 1000), + ) supported_pairs = set(model.get("supported_input_output_pairs", [])) excluded_pairs = set(model.get("excluded_input_output_pairs", [])) model_params = [] seen = set() - # Run all benchmark cases at the two highest supported concurrency levels too. extra_concurrency = (128, 256) for param in params: - variants = [param] - for concurrency in extra_concurrency: - variant = dict(param) - variant["concurrency"] = concurrency - variants.append(variant) + base_params = [param] + if is_aw_model: + base_params = [] + for input_length, output_length in aw_pairs: + aws_param = dict(param) + aws_param["input_length"] = input_length + aws_param["output_length"] = output_length + aws_param["random_range_ratio"] = "" + base_params.append(aws_param) + + variants = [] + for base_param in base_params: + variants.append(base_param) + for concurrency in extra_concurrency: + variant = dict(base_param) + variant["concurrency"] = concurrency + variants.append(variant) for variant in variants: key = ( @@ -793,6 +1291,8 @@ jobs: VLLM_COMMIT_USED: ${{ needs.resolve-atom-source.outputs.selected_vllm_commit }} VLLM_VERSION_USED: ${{ needs.resolve-atom-source.outputs.selected_vllm_version }} PUBLISH_TO_DASHBOARD: ${{ needs.resolve-atom-source.outputs.publish_to_dashboard }} + UPLOAD_TO_CUSTOM_DASHBOARD: ${{ needs.resolve-atom-source.outputs.upload_to_custom_dashboard }} + BENCHMARK_CLIENT: ${{ contains(matrix.model.prefix, '-aw-') && 'vLLM bench' || needs.resolve-atom-source.outputs.benchmark_client }} steps: - name: Detect container engine run: | @@ -807,38 +1307,7 @@ jobs: exit 1 fi - - name: Check if model is enabled - id: check - run: | - if [ "${GITHUB_EVENT_NAME}" = "schedule" ]; then - echo "enabled=true" >> "$GITHUB_OUTPUT" - exit 0 - fi - - case "${{ matrix.model.prefix }}" in - deepseek-r1-fp8) echo "enabled=${{ inputs.deepseek-r1-fp8 }}" >> "$GITHUB_OUTPUT" ;; - deepseek-r1-mxfp4) echo "enabled=${{ inputs.deepseek-r1-mxfp4 }}" >> "$GITHUB_OUTPUT" ;; - deepseek-v3-2-fp8) echo "enabled=${{ inputs.deepseek-v3-2-fp8 }}" >> "$GITHUB_OUTPUT" ;; - gpt-oss-120b) echo "enabled=${{ inputs.gpt-oss-120b }}" >> "$GITHUB_OUTPUT" ;; - glm-5-1-fp8-tp8) echo "enabled=${{ inputs.glm-5-1-fp8-tp8 }}" >> "$GITHUB_OUTPUT" ;; - kimi-k2-thinking-mxfp4-tp4) echo "enabled=${{ inputs.kimi-k2-thinking-mxfp4-tp4 }}" >> "$GITHUB_OUTPUT" ;; - kimi-k2-thinking-mxfp4-tp8) echo "enabled=${{ inputs.kimi-k2-thinking-mxfp4-tp8 }}" >> "$GITHUB_OUTPUT" ;; - kimi-k25-mxfp4-tp2) echo "enabled=${{ inputs.kimi-k2-5-mxfp4-tp2 }}" >> "$GITHUB_OUTPUT" ;; - kimi-k25-mxfp4-tp4) echo "enabled=${{ inputs.kimi-k2-5-mxfp4-tp4 }}" >> "$GITHUB_OUTPUT" ;; - kimi-k25-mxfp4-tp8) echo "enabled=${{ inputs.kimi-k2-5-mxfp4-tp8 }}" >> "$GITHUB_OUTPUT" ;; - qwen3-5-397b-a17b-tp4) echo "enabled=${{ inputs.qwen3-5-397b-a17b-tp4 }}" >> "$GITHUB_OUTPUT" ;; - qwen3-5-397b-a17b) echo "enabled=${{ inputs.qwen3-5-397b-a17b }}" >> "$GITHUB_OUTPUT" ;; - qwen3-5-397b-a17b-fp8-tp4) echo "enabled=${{ inputs.qwen3-5-397b-a17b-fp8-tp4 }}" >> "$GITHUB_OUTPUT" ;; - qwen3-5-397b-a17b-fp8) echo "enabled=${{ inputs.qwen3-5-397b-a17b-fp8 }}" >> "$GITHUB_OUTPUT" ;; - qwen3-5-397b-a17b-mxfp4) echo "enabled=${{ inputs.qwen3-5-397b-a17b-mxfp4 }}" >> "$GITHUB_OUTPUT" ;; - qwen3-next-80b-a3b-instruct-fp8-tp1) echo "enabled=${{ inputs.qwen3-next-80b-a3b-instruct-fp8-tp1 }}" >> "$GITHUB_OUTPUT" ;; - qwen3-next-80b-a3b-instruct-fp8-tp2) echo "enabled=${{ inputs.qwen3-next-80b-a3b-instruct-fp8-tp2 }}" >> "$GITHUB_OUTPUT" ;; - qwen3-next-80b-a3b-instruct-fp8-tp4) echo "enabled=${{ inputs.qwen3-next-80b-a3b-instruct-fp8-tp4 }}" >> "$GITHUB_OUTPUT" ;; - *) echo "enabled=true" >> "$GITHUB_OUTPUT" ;; - esac - - name: Clean up containers and workspace - if: steps.check.outputs.enabled == 'true' run: | echo "=== Cleaning up containers on $(hostname) ===" containers=$($CONTAINER_ENGINE ps -q) @@ -849,7 +1318,6 @@ jobs: $CONTAINER_ENGINE run --rm -v "${GITHUB_WORKSPACE:-$PWD}":/workspace -w /workspace --privileged docker.io/rocm/pytorch:latest bash -lc "shopt -s dotglob && ls -la /workspace/ && rm -rf /workspace/*" || true - name: Checkout benchmark ATOM source - if: steps.check.outputs.enabled == 'true' uses: actions/checkout@v6 with: repository: ${{ env.ATOM_SOURCE_REPOSITORY }} @@ -857,14 +1325,12 @@ jobs: fetch-depth: 1 - name: Record benchmark source revision - if: steps.check.outputs.enabled == 'true' run: | SOURCE_SHA="$(git rev-parse HEAD)" echo "ATOM_SOURCE_SHA=${SOURCE_SHA}" >> "$GITHUB_ENV" echo "Benchmarking ${ATOM_SOURCE_REPOSITORY}@${ATOM_SOURCE_REF} (${SOURCE_SHA}) with ${OOT_IMAGE_SOURCE} image ${OOT_IMAGE_TAG}" - name: Container Engine Login - if: steps.check.outputs.enabled == 'true' run: | # Podman requires an explicit registry when registries.conf has no search list; # Docker Hub short names (e.g. rocm/foo) must use docker.io here. @@ -876,11 +1342,9 @@ jobs: echo "${{ secrets.DOCKER_PASSWORD }}" | $CONTAINER_ENGINE login "$REG" -u "${{ secrets.DOCKER_USERNAME }}" --password-stdin - name: Set HF_TOKEN - if: steps.check.outputs.enabled == 'true' run: echo "HF_TOKEN=${HF_TOKEN:-${{ secrets.AMD_HF_TOKEN }}}" >> "$GITHUB_ENV" - name: Pull OOT benchmark image - if: steps.check.outputs.enabled == 'true' run: | set -euo pipefail IMG="${OOT_IMAGE_TAG}" @@ -899,7 +1363,6 @@ jobs: $CONTAINER_ENGINE pull "${IMG}" - name: Prepare model cache mount - if: steps.check.outputs.enabled == 'true' run: | MODEL_CACHE_MOUNT="" MODEL_CACHE_DESC="container-local /models (no host cache mount)" @@ -921,7 +1384,6 @@ jobs: echo "MODEL_CACHE_DESC=${MODEL_CACHE_DESC}" >> "$GITHUB_ENV" - name: Start OOT benchmark container - if: steps.check.outputs.enabled == 'true' run: | if [ -f "/etc/podinfo/gha-render-devices" ]; then DEVICE_FLAG=$(cat /etc/podinfo/gha-render-devices) @@ -958,12 +1420,10 @@ jobs: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Collect GPU info (inside container) - if: steps.check.outputs.enabled == 'true' id: gpu-info run: bash .github/scripts/collect_gpu_info.sh "$CONTAINER_NAME" "$CONTAINER_ENGINE" "${{ matrix.model.runner }}" - name: Download model if needed - if: steps.check.outputs.enabled == 'true' run: | model_dir="${MODEL_PATH}" if [[ "${model_dir}" = /models/* ]]; then @@ -986,7 +1446,6 @@ jobs: fi - name: Resolve OOT model path - if: steps.check.outputs.enabled == 'true' run: | model_dir="${MODEL_PATH}" if [[ "${model_dir}" = /models/* ]]; then @@ -1006,19 +1465,24 @@ jobs: echo "Using model id: ${MODEL_SOURCE_PATH}" fi - - name: Prepare OOT benchmark runner in container - if: steps.check.outputs.enabled == 'true' + - name: Prepare benchmark client in container run: | - $CONTAINER_ENGINE exec "$CONTAINER_NAME" bash -lc " - set -euo pipefail - rm -rf \"${CONTAINER_BENCH_SERVING_DIR%/bench_serving}\" - mkdir -p \"${CONTAINER_BENCH_SERVING_DIR%/bench_serving}\" - git clone --depth 1 \"${BENCH_SERVING_REPO_URL}\" \"${CONTAINER_BENCH_SERVING_DIR}\" - test -f \"${CONTAINER_BENCH_SERVING_DIR}/benchmark_serving.py\" - " + if [[ "${BENCHMARK_CLIENT}" == "InferenceMax bench" ]]; then + $CONTAINER_ENGINE exec "$CONTAINER_NAME" bash -lc " + set -euo pipefail + rm -rf \"${CONTAINER_BENCH_SERVING_DIR%/bench_serving}\" + mkdir -p \"${CONTAINER_BENCH_SERVING_DIR%/bench_serving}\" + git clone --depth 1 \"${BENCH_SERVING_REPO_URL}\" \"${CONTAINER_BENCH_SERVING_DIR}\" + test -f \"${CONTAINER_BENCH_SERVING_DIR}/benchmark_serving.py\" + " + else + $CONTAINER_ENGINE exec "$CONTAINER_NAME" bash -lc " + set -euo pipefail + command -v vllm + " + fi - name: Launch OOT benchmark server - if: steps.check.outputs.enabled == 'true' env: MAX_WAIT_RETRIES: "120" STREAM_VLLM_LOGS: "1" @@ -1038,53 +1502,93 @@ jobs: - name: Run OOT benchmark client id: run_benchmark_client - if: steps.check.outputs.enabled == 'true' timeout-minutes: 240 continue-on-error: true run: | set -euo pipefail TRUST_REMOTE_CODE_ARG="" if [[ "${OOT_EXTRA_ARGS}" == *"--trust-remote-code"* ]]; then - TRUST_REMOTE_CODE_ARG="--trust-remote-code" + if [[ "${BENCHMARK_CLIENT}" == "vLLM bench" ]]; then + TRUST_REMOTE_CODE_ARG="--trust_remote_code" + else + TRUST_REMOTE_CODE_ARG="--trust-remote-code" + fi fi { - echo "=== Benchmark config: ${MODEL_NAME} ISL=${ISL} OSL=${OSL} CONC=${CONC} RANDOM_RANGE_RATIO=${RANDOM_RANGE_RATIO} ===" - $CONTAINER_ENGINE exec \ - -e ISL="${ISL}" \ - -e OSL="${OSL}" \ - -e CONC="${CONC}" \ - -e RANDOM_RANGE_RATIO="${RANDOM_RANGE_RATIO}" \ - -e RESULT_FILENAME="${RESULT_FILENAME}" \ - -e BENCH_EXTRA_ARGS="${BENCH_EXTRA_ARGS}" \ - "$CONTAINER_NAME" bash -lc " - set -euo pipefail - rm -rf \"${CONTAINER_RESULT_DIR}\" - mkdir -p \"${CONTAINER_RESULT_DIR}\" - PYTHONDONTWRITEBYTECODE=1 python \"${CONTAINER_BENCH_SERVING_DIR}/benchmark_serving.py\" \ - --model=\"${OOT_RESOLVED_MODEL_PATH:-$MODEL_PATH}\" \ - --backend=vllm \ - --base-url=http://127.0.0.1:8000 \ - --dataset-name=random \ - --random-input-len=\"${ISL}\" \ - --random-output-len=\"${OSL}\" \ - --random-range-ratio \"${RANDOM_RANGE_RATIO}\" \ - --num-prompts=\"$(( CONC * 10 ))\" \ - --max-concurrency=\"${CONC}\" \ - ${TRUST_REMOTE_CODE_ARG} \ - --num-warmups=\"$(( 2 * CONC ))\" \ - --request-rate=inf \ - --ignore-eos \ - --save-result \ - --percentile-metrics=\"ttft,tpot,itl,e2el\" \ - --result-dir=\"${CONTAINER_RESULT_DIR}\" \ - --result-filename=\"${RESULT_FILENAME}.json\" \ - ${BENCH_EXTRA_ARGS:-} - " + echo "=== Benchmark config: ${MODEL_NAME} CLIENT=${BENCHMARK_CLIENT} ISL=${ISL} OSL=${OSL} CONC=${CONC} RANDOM_RANGE_RATIO=${RANDOM_RANGE_RATIO} ===" + if [[ "${BENCHMARK_CLIENT}" == "InferenceMax bench" ]]; then + $CONTAINER_ENGINE exec \ + -e ISL="${ISL}" \ + -e OSL="${OSL}" \ + -e CONC="${CONC}" \ + -e RANDOM_RANGE_RATIO="${RANDOM_RANGE_RATIO}" \ + -e RESULT_FILENAME="${RESULT_FILENAME}" \ + -e BENCH_EXTRA_ARGS="${BENCH_EXTRA_ARGS}" \ + "$CONTAINER_NAME" bash -lc " + set -euo pipefail + rm -rf \"${CONTAINER_RESULT_DIR}\" + mkdir -p \"${CONTAINER_RESULT_DIR}\" + PYTHONDONTWRITEBYTECODE=1 python \"${CONTAINER_BENCH_SERVING_DIR}/benchmark_serving.py\" \ + --model=\"${OOT_RESOLVED_MODEL_PATH:-$MODEL_PATH}\" \ + --backend=vllm \ + --base-url=http://127.0.0.1:8000 \ + --dataset-name=random \ + --random-input-len=\"${ISL}\" \ + --random-output-len=\"${OSL}\" \ + --random-range-ratio \"${RANDOM_RANGE_RATIO}\" \ + --num-prompts=\"$(( CONC * 10 ))\" \ + --max-concurrency=\"${CONC}\" \ + ${TRUST_REMOTE_CODE_ARG} \ + --num-warmups=\"$(( 2 * CONC ))\" \ + --request-rate=inf \ + --ignore-eos \ + --save-result \ + --percentile-metrics=\"ttft,tpot,itl,e2el\" \ + --result-dir=\"${CONTAINER_RESULT_DIR}\" \ + --result-filename=\"${RESULT_FILENAME}.json\" \ + ${BENCH_EXTRA_ARGS:-} + " + else + $CONTAINER_ENGINE exec \ + -e ISL="${ISL}" \ + -e OSL="${OSL}" \ + -e CONC="${CONC}" \ + -e RANDOM_RANGE_RATIO="${RANDOM_RANGE_RATIO}" \ + -e RESULT_FILENAME="${RESULT_FILENAME}" \ + -e BENCH_EXTRA_ARGS="${BENCH_EXTRA_ARGS}" \ + "$CONTAINER_NAME" bash -lc " + set -euo pipefail + rm -rf \"${CONTAINER_RESULT_DIR}\" + mkdir -p \"${CONTAINER_RESULT_DIR}\" + vllm bench serve \ + --backend vllm \ + --base-url http://127.0.0.1:8000 \ + --endpoint /v1/completions \ + --model \"${OOT_RESOLVED_MODEL_PATH:-$MODEL_PATH}\" \ + --dataset-name random \ + --random-input-len \"${ISL}\" \ + --random-output-len \"${OSL}\" \ + ${RANDOM_RANGE_RATIO:+--random-range-ratio \"${RANDOM_RANGE_RATIO}\"} \ + --temperature 0.0 \ + --num-prompts \"$(( CONC * 10 ))\" \ + --max-concurrency \"${CONC}\" \ + ${TRUST_REMOTE_CODE_ARG} \ + --num-warmups \"$(( 2 * CONC ))\" \ + --request-rate inf \ + --ignore-eos \ + --disable-tqdm \ + --save-result \ + --percentile-metrics ttft,tpot,itl,e2el \ + --result-dir \"${CONTAINER_RESULT_DIR}\" \ + --result-filename \"${RESULT_FILENAME}.json\" \ + ${BENCH_EXTRA_ARGS:-} + " + fi } > "./${RESULT_FILENAME}.client.log" 2>&1 - name: Copy OOT benchmark outputs - if: always() && steps.check.outputs.enabled == 'true' + if: always() run: | set -euo pipefail if $CONTAINER_ENGINE exec "$CONTAINER_NAME" bash -lc "test -f \"${CONTAINER_RESULT_DIR}/${RESULT_FILENAME}.json\"" >/dev/null 2>&1; then @@ -1095,7 +1599,7 @@ jobs: fi - name: Enrich copied benchmark result - if: always() && steps.check.outputs.enabled == 'true' + if: always() run: | set -euo pipefail if [ ! -f "${RESULT_FILENAME}.json" ]; then @@ -1115,6 +1619,8 @@ jobs: OOT_IMAGE_SOURCE="${OOT_IMAGE_SOURCE}" \ OOT_IMAGE_TAG_USED="${OOT_IMAGE_TAG}" \ PUBLISH_TO_DASHBOARD="${PUBLISH_TO_DASHBOARD}" \ + UPLOAD_TO_CUSTOM_DASHBOARD="${UPLOAD_TO_CUSTOM_DASHBOARD}" \ + BENCHMARK_CLIENT="${BENCHMARK_CLIENT}" \ python3 - <<'PY' import json import os @@ -1146,16 +1652,20 @@ jobs: data["vllm_version"] = os.environ.get("VLLM_VERSION_USED", "") data["oot_image_source"] = os.environ.get("OOT_IMAGE_SOURCE", "") data["oot_image_tag"] = os.environ.get("OOT_IMAGE_TAG_USED", "") + data["benchmark_client"] = os.environ.get("BENCHMARK_CLIENT", "") data["dashboard_publish_allowed"] = ( os.environ.get("PUBLISH_TO_DASHBOARD", "false").lower() == "true" ) + data["custom_dashboard_publish_allowed"] = ( + os.environ.get("UPLOAD_TO_CUSTOM_DASHBOARD", "false").lower() == "true" + ) with open(result_path, "w", encoding="utf-8") as f: json.dump(data, f, indent=2) PY - name: Dump OOT benchmark diagnostics - if: always() && steps.check.outputs.enabled == 'true' + if: always() run: | set -euo pipefail RESULT_PATH="${RESULT_FILENAME}.json" @@ -1210,7 +1720,6 @@ jobs: echo "::endgroup::" - name: Inject GPU metadata into benchmark results - if: steps.check.outputs.enabled == 'true' run: | shopt -s nullglob for f in "${RESULT_PREFIX}"-*.json; do @@ -1233,7 +1742,6 @@ jobs: - name: Validate benchmark result completeness id: validate_benchmark_result - if: steps.check.outputs.enabled == 'true' run: | set -euo pipefail RESULT_PATH="${RESULT_FILENAME}.json" @@ -1292,7 +1800,7 @@ jobs: PY - name: Upload benchmark artifacts - if: always() && steps.check.outputs.enabled == 'true' && steps.validate_benchmark_result.outcome == 'success' + if: always() && steps.validate_benchmark_result.outcome == 'success' uses: actions/upload-artifact@v7 with: name: oot-benchmark-${{ env.RESULT_FILENAME }} @@ -1303,7 +1811,7 @@ jobs: ${{ env.RESULT_FILENAME }}.vllm.log - name: Clean up OOT benchmark container - if: always() && steps.check.outputs.enabled == 'true' + if: always() run: | $CONTAINER_ENGINE exec "$CONTAINER_NAME" bash -lc "if [ -f /tmp/vllm_oot.pid ]; then kill \$(cat /tmp/vllm_oot.pid) || true; fi" || true $CONTAINER_ENGINE stop "$CONTAINER_NAME" || true @@ -1345,16 +1853,28 @@ jobs: - name: Note summary behavior run: | + printf '### Summary mode\nResults from `%s@%s` using `%s` OOT image mode are shown below as a workflow summary table.\n- Benchmark client: `%s`\n- Upload to dashboard: `%s`\n- Upload to custom dashboard: `%s`\n' \ + "${{ needs.resolve-atom-source.outputs.atom_repository }}" \ + "${{ needs.resolve-atom-source.outputs.atom_ref }}" \ + "${{ needs.resolve-atom-source.outputs.oot_image_source }}" \ + "${{ needs.resolve-atom-source.outputs.benchmark_client }}" \ + "${{ needs.resolve-atom-source.outputs.publish_to_dashboard }}" \ + "${{ needs.resolve-atom-source.outputs.upload_to_custom_dashboard }}" >> "$GITHUB_STEP_SUMMARY" + if [[ "${{ needs.resolve-atom-source.outputs.publish_to_dashboard }}" == "true" ]]; then - printf '### Summary mode\nResults from `%s@%s` using `%s` OOT image mode are shown below as a workflow summary table. Because dashboard upload is enabled for this run, this final Ubuntu summary job collects benchmark artifacts from all completed GPU cases and publishes the available successful results in one pass, then reconciles the completed run for regression comparison.\n' \ - "${{ needs.resolve-atom-source.outputs.atom_repository }}" \ - "${{ needs.resolve-atom-source.outputs.atom_ref }}" \ - "${{ needs.resolve-atom-source.outputs.oot_image_source }}" >> "$GITHUB_STEP_SUMMARY" + printf '\nDashboard upload is enabled for this run, so this final Ubuntu summary job also reconciles the completed run for regression comparison and publishes dashboard data.\n' \ + >> "$GITHUB_STEP_SUMMARY" + else + printf '\nDashboard upload is disabled for this run.\n' \ + >> "$GITHUB_STEP_SUMMARY" + fi + + if [[ "${{ needs.resolve-atom-source.outputs.upload_to_custom_dashboard }}" == "true" ]]; then + printf 'LLM-Booster dashboard publish is enabled for this run.\n' \ + >> "$GITHUB_STEP_SUMMARY" else - printf '### Summary mode\nResults from `%s@%s` using `%s` OOT image mode are shown below as a workflow summary table only. Dashboard upload is disabled for this run.\n' \ - "${{ needs.resolve-atom-source.outputs.atom_repository }}" \ - "${{ needs.resolve-atom-source.outputs.atom_ref }}" \ - "${{ needs.resolve-atom-source.outputs.oot_image_source }}" >> "$GITHUB_STEP_SUMMARY" + printf 'LLM-Booster dashboard publish is disabled for this run.\n' \ + >> "$GITHUB_STEP_SUMMARY" fi - name: Generate OOT benchmark summary table @@ -1385,6 +1905,265 @@ jobs: name: oot-benchmark-summary path: oot_benchmark_summary.json + - name: Note skipped LLM-Booster dashboard publish + if: needs.resolve-atom-source.outputs.upload_to_custom_dashboard == 'true' && steps.summary-stats.outputs.passed_cases == '0' + run: | + printf '\nNo successful benchmark result JSONs were produced, so LLM-Booster dashboard publish was skipped.\n' \ + >> "$GITHUB_STEP_SUMMARY" + + - name: Check LLM-Booster deploy key + id: llmbooster-key + if: needs.resolve-atom-source.outputs.upload_to_custom_dashboard == 'true' && steps.summary-stats.outputs.passed_cases != '0' + env: + LLM_BOOSTER_DEPLOY_KEY: ${{ secrets.LLM_BOOSTER_DEPLOY_KEY }} + run: | + if [[ -n "${LLM_BOOSTER_DEPLOY_KEY}" ]]; then + echo "available=true" >> "$GITHUB_OUTPUT" + else + echo "available=false" >> "$GITHUB_OUTPUT" + echo "::warning::LLM_BOOSTER_DEPLOY_KEY is not configured; skipping LLM-Booster dashboard publish." + printf '\nLLM-Booster dashboard publish skipped because secret `LLM_BOOSTER_DEPLOY_KEY` is not configured.\n' \ + >> "$GITHUB_STEP_SUMMARY" + fi + + - name: Checkout LLM-Booster dashboard repo + if: needs.resolve-atom-source.outputs.upload_to_custom_dashboard == 'true' && steps.summary-stats.outputs.passed_cases != '0' && steps.llmbooster-key.outputs.available == 'true' + uses: actions/checkout@v6 + with: + repository: ROCm/LLM-Booster + ref: main + path: llm-booster-repo + ssh-key: ${{ secrets.LLM_BOOSTER_DEPLOY_KEY }} + fetch-depth: 1 + + - name: Build LLM-Booster dashboard files + if: needs.resolve-atom-source.outputs.upload_to_custom_dashboard == 'true' && steps.summary-stats.outputs.passed_cases != '0' && steps.llmbooster-key.outputs.available == 'true' + env: + RUN_URL: https://github.com/${{ github.repository }}/actions/runs/${{ github.run_id }} + ATOM_SOURCE_REPOSITORY: ${{ needs.resolve-atom-source.outputs.atom_repository }} + ATOM_SOURCE_REF: ${{ needs.resolve-atom-source.outputs.atom_ref }} + OOT_IMAGE_SOURCE: ${{ needs.resolve-atom-source.outputs.oot_image_source }} + BENCHMARK_CLIENT: ${{ needs.resolve-atom-source.outputs.benchmark_client }} + run: | + python3 - <<'PY' + import json + import os + import re + import sys + import time + from collections import defaultdict + from datetime import datetime, timezone + from pathlib import Path + + root = Path(".") + llm_booster_dashboard = root / "llm-booster-repo" / "dashboard" + llm_booster_data_dir = llm_booster_dashboard / "data" / "vllm-atom" + llm_booster_data_dir.mkdir(parents=True, exist_ok=True) + + sys.path.insert(0, str(llm_booster_dashboard)) + from update_data import index_root_for, rebuild_index # noqa: E402 + + catalog = json.loads( + (root / ".github" / "benchmark" / "oot_benchmark_models.json").read_text( + encoding="utf-8" + ) + ) + + def slugify(text: str) -> str: + return re.sub(r"[^a-z0-9]+", "-", text.lower()).strip("-") or "unknown" + + def precision_from_model(model: str): + upper = model.upper() + if "MXFP4" in upper or "FP4" in upper: + return "fp4" + if "FP8" in upper: + return "fp8" + return None + + def model_name_for_variant(variant: dict) -> str: + if variant.get("dashboard_model"): + return str(variant["dashboard_model"]) + display = str(variant["display"]) + return display.replace(" TP", "-tp").replace(" ", "-") + + prefix_map = {} + for family in catalog["families"]: + base = {k: v for k, v in family.items() if k not in {"choice_label", "variants"}} + for variant in family["variants"]: + item = dict(base) + item.update(variant) + prefix_map[str(variant["prefix"])] = item + + def parse_result_filename(path: Path): + stem = path.stem + prefix, isl, osl, concurrency, ratio = stem.rsplit("-", 4) + return prefix, int(isl), int(osl), int(concurrency), float(ratio) if ratio else 0.0 + + default_benchmark_client = os.environ["BENCHMARK_CLIENT"] + + def resolve_client(payload): + bc = str(payload.get("benchmark_client", default_benchmark_client)) + slug = "vllmbench" if bc == "vLLM bench" else "inferencemax" + label = "vllm bench" if bc == "vLLM bench" else "inferencemax" + return slug, label + + grouped = defaultdict(list) + for path in sorted(root.glob("*.json")): + if path.name in {"oot_benchmark_summary.json", "regression_report.json"}: + continue + try: + payload = json.loads(path.read_text(encoding="utf-8")) + except (OSError, UnicodeDecodeError, json.JSONDecodeError): + continue + if "output_throughput" not in payload: + continue + custom_dashboard_flag = str( + payload.get("custom_dashboard_publish_allowed", True) + ).strip().lower() + if custom_dashboard_flag in {"0", "false", "no"}: + continue + + prefix, isl, osl, concurrency, ratio = parse_result_filename(path) + variant = prefix_map.get(prefix) + if variant is None: + raise SystemExit(f"Unknown benchmark prefix in result file {path.name}: {prefix}") + + model = model_name_for_variant(variant) + precision = precision_from_model(model) + atom_source_sha = str(payload.get("atom_source_sha", "")) or "unknown" + sha7 = atom_source_sha[:7] if atom_source_sha else "unknown" + day = time.strftime("%Y-%m-%d", time.gmtime()) + client_slug, _ = resolve_client(payload) + grouped[(day, sha7, model, isl, osl, client_slug)].append((path, payload, precision, concurrency, ratio)) + + written_files = set() + for (day, sha7, model, isl, osl, client_slug), entries in sorted(grouped.items()): + entries.sort(key=lambda item: (item[3], item[4])) + client_label = "vllm bench" if client_slug == "vllmbench" else "inferencemax" + timestamp_ms = int(time.time() * 1000) + precision = entries[0][2] + first_payload = entries[0][1] + tp = first_payload.get("tensor_parallel_size") + tp_label = f"_tp{int(tp)}" if tp not in (None, "", 0) else "" + model_slug = slugify(model) + run_id = ( + f"atom_vllm_{day.replace('-', '')}_{model_slug}_isl{isl}_osl{osl}_" + f"{sha7}_{os.environ['GITHUB_RUN_ID']}_{client_slug}" + ) + config_label = f"atom-vllm_{model_slug}" + if precision: + config_label += f"_{precision}" + if tp not in (None, "", 0): + config_label += tp_label + if client_slug == "inferencemax": + config_label += "_inferencemax" + + points = [] + for _, payload, precision, concurrency, ratio in entries: + output_tput = payload.get("output_throughput") + total_tput = payload.get("total_token_throughput") + total_gpu = payload.get("tensor_parallel_size") + try: + total_gpu = int(total_gpu) if total_gpu is not None else None + except (TypeError, ValueError): + total_gpu = None + points.append( + { + "isl": isl, + "osl": osl, + "concurrency": concurrency, + "ratio": ratio, + "client_bench": client_label, + "ttft_ms": payload.get("mean_ttft_ms"), + "ttft_p99": payload.get("p99_ttft_ms"), + "tpot_ms": payload.get("mean_tpot_ms"), + "tpot_p99": payload.get("p99_tpot_ms"), + "itl_ms": payload.get("mean_itl_ms"), + "e2el_ms": payload.get("mean_e2el_ms"), + "interactivity": (output_tput / concurrency) if output_tput and concurrency else None, + "tput_per_gpu": (total_tput / total_gpu) if total_tput and total_gpu else None, + "output_tput_per_gpu": (output_tput / total_gpu) if output_tput and total_gpu else None, + "input_tput_per_gpu": None, + "output_tput": output_tput, + "total_tput": total_tput, + "req_tput": payload.get("request_throughput"), + "completed": payload.get("completed"), + "duration": payload.get("duration"), + "num_prompts": payload.get("num_prompts"), + "prefill_tp": None, + "decode_tp": None, + "tensor_parallel_size": total_gpu, + "num_prefill_gpu": None, + "num_decode_gpu": None, + "total_gpu": total_gpu, + "gpu_name": payload.get("gpu_name"), + "image": payload.get("oot_image_tag") or payload.get("docker_image"), + "rocm": payload.get("rocm_version"), + "run_url": os.environ["RUN_URL"], + "model_path": payload.get("model_id"), + "point_date": day, + "precision": precision, + "hardware": "mi355x", + } + ) + + run = { + "run_id": run_id, + "date": day, + "timestamp": timestamp_ms, + "model": model, + "backend": "vLLM-ATOM", + "config_label": config_label, + "precision": precision, + "source": "ATOM Dashboard", + "client_bench": client_label, + "hardware": "mi355x", + "commit": { + "id": first_payload.get("atom_source_sha"), + "message": f"ATOM benchmark run {os.environ['GITHUB_RUN_ID']}", + "timestamp": datetime.now(timezone.utc).isoformat(), + "url": f"https://github.com/{os.environ['ATOM_SOURCE_REPOSITORY']}/commit/{first_payload.get('atom_source_sha', '')}", + "author": {}, + "committer": {}, + }, + "points": points, + "gsm8k": None, + } + + run_file = llm_booster_data_dir / f"{run_id}.json" + run_file.write_text(json.dumps(run, indent=2) + "\n", encoding="utf-8") + written_files.add(run_file) + + rebuild_index(index_root_for(llm_booster_data_dir)) + print(f"Wrote {len(written_files)} LLM-Booster ATOM vLLM run files.") + PY + + - name: Upload LLM-Booster dashboard preview artifact + if: needs.resolve-atom-source.outputs.upload_to_custom_dashboard == 'true' && steps.summary-stats.outputs.passed_cases != '0' && steps.llmbooster-key.outputs.available == 'true' + uses: actions/upload-artifact@v7 + with: + name: oot-benchmark-llm-booster-preview + path: | + llm-booster-repo/dashboard/data/index.json + llm-booster-repo/dashboard/data/vllm-atom/*.json + + - name: Push LLM-Booster dashboard data + if: needs.resolve-atom-source.outputs.upload_to_custom_dashboard == 'true' && steps.summary-stats.outputs.passed_cases != '0' && steps.llmbooster-key.outputs.available == 'true' + working-directory: llm-booster-repo + env: + GIT_AUTHOR_NAME: github-actions[bot] + GIT_AUTHOR_EMAIL: github-actions[bot]@users.noreply.github.com + GIT_COMMITTER_NAME: github-actions[bot] + GIT_COMMITTER_EMAIL: github-actions[bot]@users.noreply.github.com + run: | + git add dashboard/data/index.json dashboard/data/vllm-atom + if git diff --cached --quiet; then + echo "No LLM-Booster dashboard data changes to push." + exit 0 + fi + git commit -m "Update ATOM vLLM benchmark dashboard data" + git push origin HEAD:main + - name: Download baseline from previous successful main run if: needs.resolve-atom-source.outputs.publish_to_dashboard == 'true' && steps.summary-stats.outputs.passed_cases != '0' id: baseline @@ -1394,7 +2173,7 @@ jobs: mapfile -t CANDIDATE_RUN_IDS < <( { gh run list \ - --workflow="ATOM vLLM OOT Benchmark" \ + --workflow="ATOM vLLM Benchmark" \ --branch=main \ --event=workflow_dispatch \ --status=success \ diff --git a/.github/workflows/atom-vllm-test.yaml b/.github/workflows/atom-vllm-test.yaml index 798ec69622..732741c255 100644 --- a/.github/workflows/atom-vllm-test.yaml +++ b/.github/workflows/atom-vllm-test.yaml @@ -35,7 +35,7 @@ jobs: - name: Checkout ATOM repo uses: actions/checkout@v6 - - name: Download and check pre-checkin signal + - name: Wait for Pre Checkin workflow run: bash ./.github/scripts/check_signal.sh env: GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} @@ -192,7 +192,7 @@ jobs: with: name: aiter-whl path: aiter-whl/amd_aiter*.whl - retention-days: 1 + retention-days: 7 atom-vllm-oot: needs: [check-signal, download_aiter_wheel] @@ -212,9 +212,10 @@ jobs: - display_name: "gpt-oss-120b TP1" model_name: "gpt-oss-120b" model_path: "openai/gpt-oss-120b" - extra_args: "--tensor-parallel-size 1" + extra_args: "--tensor-parallel-size 1 --gpu-memory-utilization 0.5" + client_command: "lm_eval --model local-chat-completions --apply_chat_template --model_args model=${MODEL_PATH},base_url=http://127.0.0.1:${VLLM_PORT}/v1/chat/completions,num_concurrent=65,max_retries=3,max_gen_toks=2048,tokenized_requests=False,trust_remote_code=True --tasks gsm8k --num_fewshot ${LM_EVAL_NUM_FEWSHOT} --output_path ${OUTPUT_PATH}" env_vars: "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1" - accuracy_test_threshold: 0.38 + accuracy_test_threshold: 0.88 runner: linux-atom-mi35x-1 - display_name: "Kimi-K2-Thinking-MXFP4 TP4" model_name: "Kimi-K2-Thinking-MXFP4" @@ -468,6 +469,7 @@ jobs: OOT_MODEL_NAME: ${{ matrix.model_name }} OOT_MODEL_PATH: ${{ matrix.model_path }} OOT_EXTRA_ARGS: ${{ matrix.extra_args }} + OOT_CLIENT_COMMAND: ${{ matrix.client_command || '' }} OOT_ENV_VARS: ${{ matrix.env_vars }} LM_EVAL_NUM_FEWSHOT: ${{ matrix.lm_eval_num_fewshot }} MAX_WAIT_RETRIES: "120" @@ -477,6 +479,7 @@ jobs: -e OOT_MODEL_NAME="${OOT_MODEL_NAME}" \ -e OOT_MODEL_PATH="${OOT_RESOLVED_MODEL_PATH:-$OOT_MODEL_PATH}" \ -e OOT_EXTRA_ARGS="${OOT_EXTRA_ARGS}" \ + -e OOT_CLIENT_COMMAND="${OOT_CLIENT_COMMAND}" \ -e OOT_ENV_VARS="${OOT_ENV_VARS}" \ -e LM_EVAL_NUM_FEWSHOT="${LM_EVAL_NUM_FEWSHOT}" \ -e MAX_WAIT_RETRIES="${MAX_WAIT_RETRIES}" \ diff --git a/.github/workflows/docker-release.yaml b/.github/workflows/docker-release.yaml index ebe6041afb..4cca15e75e 100644 --- a/.github/workflows/docker-release.yaml +++ b/.github/workflows/docker-release.yaml @@ -67,6 +67,9 @@ on: description: "Release the OOT vLLM image (skips native image test/push). If both this and SGLang are checked, both will be released." type: boolean default: false + custom_tag: + description: "Custom tag for the native Docker image (e.g. 'my-feature'). Produces rocm/atom-dev:. When set, branch-based tagging and 'latest' update are skipped." + default: "" only_release_sglang: description: "Release the SGLang+ATOM image (skips native image test/push). If both this and OOT are checked, both will be released." type: boolean @@ -110,12 +113,14 @@ jobs: - name: Echo environment variables run: | echo "ATOM_REPO: ${{ inputs.atom_repo || env.ATOM_REPO }}" + echo "ATOM_BRANCH: ${{ github.ref_name }}" echo "ATOM_COMMIT: ${{ inputs.atom_commit || env.ATOM_COMMIT }}" echo "AITER_REPO: ${{ inputs.aiter_repo || env.AITER_REPO }}" echo "AITER_COMMIT: ${{ inputs.aiter_commit || env.AITER_COMMIT }}" echo "RCCL_REPO: ${{ inputs.rccl_repo || env.RCCL_REPO }}" echo "RCCL_BRANCH: ${{ inputs.rccl_branch || env.RCCL_BRANCH }}" echo "OOT_TAG_SUFFIX: ${{ inputs.oot_tag_suffix || '' }}" + echo "CUSTOM_TAG: ${{ inputs.custom_tag || '' }}" echo "SGLANG_REPO: ${{ inputs.sglang_repo || env.SGLANG_REPO }}" echo "SGLANG_REF: ${{ inputs.sglang_ref || env.SGLANG_REF }}" echo "SGLANG_VERSION: ${{ inputs.sglang_version || env.SGLANG_VERSION }}" @@ -123,12 +128,19 @@ jobs: - name: Build Docker image timeout-minutes: 180 run: | + ATOM_COMMIT="${{ inputs.atom_commit || env.ATOM_COMMIT }}" + if [ "${ATOM_COMMIT}" = "HEAD" ]; then + ATOM_CHECKOUT="${{ github.ref_name }}" + else + ATOM_CHECKOUT="${ATOM_COMMIT}" + fi + echo "ATOM checkout ref: ${ATOM_CHECKOUT} (branch: ${{ github.ref_name }})" DOCKER_BUILDKIT=1 docker build --pull --network=host -t atom_release:ci \ --target atom_image \ --build-arg BASE_IMAGE="${{ matrix.base_image }}" \ --build-arg GPU_ARCH="${{ env.GPU_ARCH }}" \ --build-arg ATOM_REPO="${{ inputs.atom_repo || env.ATOM_REPO }}" \ - --build-arg ATOM_COMMIT="${{ inputs.atom_commit || env.ATOM_COMMIT }}" \ + --build-arg ATOM_COMMIT="${ATOM_CHECKOUT}" \ --build-arg AITER_REPO="${{ inputs.aiter_repo || env.AITER_REPO }}" \ --build-arg AITER_COMMIT="${{ inputs.aiter_commit || env.AITER_COMMIT }}" \ --build-arg RCCL_REPO="${{ inputs.rccl_repo || env.RCCL_REPO }}" \ @@ -189,12 +201,32 @@ jobs: - name: Push Docker image if: ${{ success() && inputs.only_release_oot != true && inputs.only_release_sglang != true }} run: | - # Push both the versioned tag and update the 'latest' tag for the Docker image - TAG=nightly_$(date +%Y%m%d%H%M) - docker tag atom_release:ci rocm/atom-dev:latest - docker push rocm/atom-dev:latest - docker tag atom_release:ci rocm/atom-dev:${TAG} - docker push rocm/atom-dev:${TAG} + CUSTOM_TAG="${{ inputs.custom_tag || '' }}" + if [ -n "${CUSTOM_TAG}" ]; then + if [[ ! "${CUSTOM_TAG}" =~ ^[A-Za-z0-9_.-]+$ ]]; then + echo "Invalid custom_tag '${CUSTOM_TAG}'. Allowed characters: letters, numbers, '.', '_' and '-'." + exit 1 + fi + docker tag atom_release:ci rocm/atom-dev:${CUSTOM_TAG} + docker push rocm/atom-dev:${CUSTOM_TAG} + echo "Custom tag: rocm/atom-dev:${CUSTOM_TAG}. Skipping latest update." + else + BRANCH="${{ github.ref_name }}" + DEFAULT_BRANCH="${{ github.event.repository.default_branch }}" + TAG=nightly_$(date +%Y%m%d%H%M) + if [ "${BRANCH}" != "${DEFAULT_BRANCH}" ]; then + SAFE_BRANCH=$(echo "${BRANCH}" | sed 's/[^A-Za-z0-9._-]/-/g') + TAG="${TAG}-${SAFE_BRANCH}" + docker tag atom_release:ci rocm/atom-dev:${TAG} + docker push rocm/atom-dev:${TAG} + echo "Branch '${BRANCH}'; tag: rocm/atom-dev:${TAG}. Skipping latest update." + else + docker tag atom_release:ci rocm/atom-dev:latest + docker push rocm/atom-dev:latest + docker tag atom_release:ci rocm/atom-dev:${TAG} + docker push rocm/atom-dev:${TAG} + fi + fi - name: Build OOT Docker image if: ${{ success() && (inputs.only_release_oot == true || (inputs.only_release_sglang != true && (github.event_name == 'schedule' || inputs.build_oot_image == true))) }} diff --git a/.github/workflows/pre-checks.yaml b/.github/workflows/pre-checks.yaml index 4d4e799e2d..8f6426afb1 100644 --- a/.github/workflows/pre-checks.yaml +++ b/.github/workflows/pre-checks.yaml @@ -62,40 +62,3 @@ jobs: -reporter=github-pr-review \ -filter-mode=diff_context \ -fail-on-error=true - - upload-success-artifact: - name: Upload Success Signal - runs-on: ubuntu-latest - needs: [black, ruff] - if: ${{ needs.black.result == 'success' && needs.ruff.result == 'success' }} - steps: - - name: Create success signal file - run: echo "success" > checks_signal.txt - - - name: Upload success artifact - uses: actions/upload-artifact@v4 - with: - name: checks-signal-${{ github.sha }} - path: checks_signal.txt - - upload-failure-artifact: - name: Upload Failure Signal - runs-on: ubuntu-latest - needs: [black, ruff] - if: ${{ always() && (needs.black.result != 'success' || needs.ruff.result != 'success') }} - steps: - - name: Create failure signal file with failed jobs - run: | - echo "failure" > checks_signal.txt - if [ "${{ needs.black.result }}" != "success" ]; then - echo "FAILED: black (${{ needs.black.result }})" >> checks_signal.txt - fi - if [ "${{ needs.ruff.result }}" != "success" ]; then - echo "FAILED: ruff (${{ needs.ruff.result }})" >> checks_signal.txt - fi - - - name: Upload failure artifact - uses: actions/upload-artifact@v4 - with: - name: checks-signal-${{ github.sha }} - path: checks_signal.txt diff --git a/README.md b/README.md index d6ae9bb8bd..2805acb7ec 100644 --- a/README.md +++ b/README.md @@ -14,6 +14,7 @@ ## 📢 News - **[2026/05]** [Dissecting DeepSeek V4 Compressor](https://rocm.github.io/ATOM/dissecting_dsv4_compressor) — interactive animation visualizing how the CSA/HCA compressor state cache works (overlap mechanism, prefill vs decode, bulk compression vs sequential accumulation). +- **[2026/05]** ATOM now supports **Prefill/Decode (P/D) disaggregation** with [Mooncake](https://github.com/kvcache-ai/Mooncake) RDMA push-mode KV cache transfer. See [PD disaggregation guide](docs/pd_disaggregation_guide.md). - **[2026/03]** ATOM now supports **Prefill/Decode (P/D) disaggregation** — run prefill and decode on separate GPU nodes with RDMA-based KV cache transfer via [MORI-IO](https://github.com/ROCm/mori). See [disaggregation docs](atom/kv_transfer/disaggregation/README.md). ## 🚀 Features diff --git a/atom/benchmarks/benchmark_serving.py b/atom/benchmarks/benchmark_serving.py index 7f584444cd..35b4302c49 100644 --- a/atom/benchmarks/benchmark_serving.py +++ b/atom/benchmarks/benchmark_serving.py @@ -28,6 +28,7 @@ import argparse import asyncio import contextlib +import functools import gc import json import os @@ -37,12 +38,17 @@ from argparse import ArgumentParser as FlexibleArgumentParser from dataclasses import dataclass from datetime import datetime -from typing import Any, AsyncGenerator, Dict, List, Optional, Tuple +from typing import Any, AsyncGenerator, Callable, Dict, List, Optional, Tuple import numpy as np from tqdm.asyncio import tqdm from transformers import PreTrainedTokenizerBase +from atom.entrypoints.openai.chat_encoders import ( + apply_chat_template, + load_custom_message_encoder, +) + from .backend_request_func import ( ASYNC_REQUEST_FUNCS, RequestFuncInput, @@ -92,16 +98,15 @@ def sample_random_requests( range_ratio: float, tokenizer: PreTrainedTokenizerBase, use_chat_template: bool = False, + apply_chat_template_fn: Callable = lambda x: x, ) -> List[Tuple[str, int, int]]: prefix_token_ids = np.random.randint( 0, tokenizer.vocab_size, size=prefix_len ).tolist() if use_chat_template: - chat_template_dummy = tokenizer.apply_chat_template( + chat_template_dummy = apply_chat_template_fn( [{"role": "user", "content": "a"}], - add_generation_prompt=True, - tokenize=False, ) tokenized_chat_template_dummy = tokenizer.encode( chat_template_dummy, add_special_tokens=False @@ -143,10 +148,8 @@ def sample_uniform(seq_len): prompt = tokenizer.decode(prompt_token_ids) if use_chat_template: - prompt = tokenizer.apply_chat_template( + prompt = apply_chat_template_fn( [{"role": "user", "content": prompt}], - add_generation_prompt=True, - tokenize=False, ) prompt_len = len(tokenizer.encode(prompt, add_special_tokens=False)) @@ -705,6 +708,10 @@ def main(args: argparse.Namespace): tokenizer_mode=tokenizer_mode, trust_remote_code=args.trust_remote_code, ) + custom_encoder = load_custom_message_encoder(model_id) + apply_chat_template_fn = functools.partial( + apply_chat_template, tokenizer, custom_encoder + ) if args.dataset_name == "random": input_requests = sample_random_requests( @@ -714,6 +721,7 @@ def main(args: argparse.Namespace): num_prompts=args.num_prompts, range_ratio=args.random_range_ratio, tokenizer=tokenizer, + apply_chat_template_fn=apply_chat_template_fn, use_chat_template=args.use_chat_template, ) diff --git a/atom/config.py b/atom/config.py index 2f42471389..223f6f4557 100644 --- a/atom/config.py +++ b/atom/config.py @@ -723,10 +723,14 @@ class SpeculativeConfig: model: Optional[str] = None num_speculative_tokens: Optional[int] = None draft_model_hf_config: Optional[PretrainedConfig] = None + use_aux_hidden_state: bool = False + eagle3_aux_layer_ids: list[int] = field(default_factory=list) # model_type → mtp_model_type mapping _MTP_TYPE_MAP: ClassVar[dict[str, str]] = { "deepseek_v3": "deepseek_mtp", + "deepseek_v32": "deepseek_mtp", + "deepseek_v4": "deepseek_v4_mtp", "glm_moe_dsa": "deepseek_mtp", "qwen3_next": "qwen3_next_mtp", "qwen3_5": "qwen3_5_mtp", @@ -739,22 +743,51 @@ class SpeculativeConfig: # mtp_model_type → (n_predict_attr, architecture) _MTP_CONFIG: ClassVar[dict[str, tuple[str, str]]] = { "deepseek_mtp": ("num_nextn_predict_layers", "DeepSeekMTPModel"), + "deepseek_v4_mtp": ("num_nextn_predict_layers", "DeepseekV4MTPModel"), "qwen3_next_mtp": ("num_nextn_predict_layers", "Qwen3NextMTPModel"), "qwen3_5_mtp": ("mtp_num_hidden_layers", "Qwen3_5MTPModel"), } def __post_init__(self): if self.draft_model_hf_config is None: - self.draft_model_hf_config = AutoConfig.from_pretrained( + self.draft_model_hf_config = get_hf_config( self.model, trust_remote_code=True ) # For multimodal models, extract text_config if hasattr(self.draft_model_hf_config, "text_config"): self.draft_model_hf_config = self.draft_model_hf_config.text_config - self.hf_config_override(self.draft_model_hf_config) + self.hf_config_override(self.draft_model_hf_config, self.model) + + if self.method == "eagle3": + if getattr(self.draft_model_hf_config, "kv_lora_rank", None): + raise NotImplementedError( + "Eagle3 draft model with MLA attention is not supported" + ) + # Aux hidden state layers: prefer the draft checkpoint's + # eagle_config; if absent or the list is empty, ModelRunner + # falls back to model.get_eagle3_aux_hidden_state_layers(), + # which defaults to 3 layers — early / middle / late + # (see DeepseekV2ForCausalLM.get_eagle3_aux_hidden_state_layers, + # returns `(2, num_layers // 2, num_layers - 3)`, aligned with vLLM). + eagle_cfg = getattr(self.draft_model_hf_config, "eagle_config", None) + if eagle_cfg: + self.use_aux_hidden_state = eagle_cfg.get("use_aux_hidden_state", False) + if self.use_aux_hidden_state and not self.eagle3_aux_layer_ids: + self.eagle3_aux_layer_ids = eagle_cfg.get( + "eagle_aux_hidden_state_layer_ids", [] + ) + else: + self.use_aux_hidden_state = True @staticmethod - def hf_config_override(hf_config: PretrainedConfig) -> None: + def hf_config_override( + hf_config: PretrainedConfig, model_path: Optional[str] = None + ) -> None: + # Eagle3 architecture mapping (architecture-level, not model_type) + arch = (getattr(hf_config, "architectures", None) or [""])[0] + if arch == "LlamaForCausalLMEagle3": + hf_config.architectures = ["Eagle3LlamaModel"] + # Step 1: resolve model_type → mtp model_type mtp_type = SpeculativeConfig._MTP_TYPE_MAP.get(hf_config.model_type) if mtp_type is not None: @@ -777,10 +810,33 @@ def hf_config_override(hf_config: PretrainedConfig) -> None: "num_nextn_predict_layers": n_predict, "architectures": [arch], } - # Qwen3.5 MTP needs expert counts for MoE layer construction - if hf_config.model_type == "qwen3_5_mtp": - updates["n_shared_experts"] = 1 - updates["n_routed_experts"] = getattr(hf_config, "num_experts", 0) + # Naming differs across families: + # DeepSeek / GLM → already have `n_routed_experts` + # Qwen3.5 / Qwen3-Next → only carry `num_experts` + # non-MoE / unknown → leave unset (no MoE = no field) + n_routed = getattr( + hf_config, + "n_routed_experts", + getattr(hf_config, "num_experts", None), + ) + if n_routed is not None: + updates["n_routed_experts"] = n_routed + # n_shared_experts: prefer the field's own value if it exists + # (DeepSeek / GLM ship it natively); else scan the checkpoint + # for `shared_expert` weights and use the parsed count (1 for + # the flat-block layout every released model uses). Leaving + # the field unset for ckpts without shared experts avoids + # fabricating a phantom shared block (e.g. R1 MTP eh_proj + # would otherwise pull a stale BF16 weight_scale). + existing_n_shared = getattr(hf_config, "n_shared_experts", None) + if existing_n_shared is not None: + updates["n_shared_experts"] = existing_n_shared + else: + from atom.models.utils import ckpt_shared_expert_count + + n_shared = ckpt_shared_expert_count(model_path) + if n_shared > 0: + updates["n_shared_experts"] = n_shared hf_config.update(updates) @@ -824,7 +880,7 @@ class Config: kv_cache_block_size: int = 16 num_kvcache_blocks: int = -1 kv_cache_dtype: str = "bf16" - enable_prefix_caching: bool = False + enable_prefix_caching: bool = True port: int = 8006 torch_profiler_dir: str | None = field( default_factory=lambda: envs.ATOM_TORCH_PROFILER_DIR @@ -963,10 +1019,26 @@ def __post_init__(self): # tokens. ATOM's BlockManager + slot_mapping math assume one global # block_size, so we override `kv_cache_block_size` here when V4 is # detected; the V4 attention builder enforces the same value. - if getattr(self.hf_config, "model_type", None) == "deepseek_v4": + # + # NOTE: cannot use `hf_config.model_type` for detection — `_CONFIG_REGISTRY` + # maps "deepseek_v4" → "deepseek_v3" so model_type reads as "deepseek_v3". + # Use the preserved `architectures` field (re-injected by get_hf_config, + # line 567) which keeps the original "DeepseekV4ForCausalLM[NextN]" name. + arches = getattr(self.hf_config, "architectures", None) or [] + if any("DeepseekV4" in str(a) for a in arches): v4_block_size = 128 if self.kv_cache_block_size != v4_block_size: self.kv_cache_block_size = v4_block_size + # TODO: V4's per-request SWA buffer cannot be restored from the classical + # KV pool on prefix cache hit, so disable prefix caching silently. + if self.enable_prefix_caching: + import logging + + logging.getLogger(__name__).warning( + "DeepSeek-V4 does not support prefix caching " + "(SWA buffer is not cacheable); disabling automatically." + ) + self.enable_prefix_caching = False def compute_hash(self) -> str: """ @@ -1011,6 +1083,30 @@ def set_current_atom_config(atom_config: Config): _current_atom_config = atom_config +def _get_current_atom_config_from_vllm_forward_context() -> Optional[Config]: + # In vLLM plugin mode (especially speculative decode), main/draft models + # can coexist in one process. Resolve per-forward config first to avoid + # reading a stale global singleton. + try: + from vllm.forward_context import ( + get_forward_context as get_vllm_forward_context, + is_forward_context_available, + ) + except Exception: + return None + if not is_forward_context_available(): + return None + try: + return get_vllm_forward_context().additional_kwargs.get("atom_config") + except Exception: + return None + + def get_current_atom_config() -> Config: + # Try to get the atom config from forward context first in vLLM plugin mode. + if is_vllm(): + forward_atom_config = _get_current_atom_config_from_vllm_forward_context() + if forward_atom_config is not None: + return forward_atom_config assert _current_atom_config is not None, "Current atom config is not set" return _current_atom_config diff --git a/atom/entrypoints/openai/api_server.py b/atom/entrypoints/openai/api_server.py index 01ef080815..bcad5e7083 100644 --- a/atom/entrypoints/openai/api_server.py +++ b/atom/entrypoints/openai/api_server.py @@ -31,6 +31,7 @@ from fastapi.responses import JSONResponse, StreamingResponse from transformers import AutoTokenizer +from .chat_encoders import apply_chat_template, load_custom_message_encoder from .protocol import ( ChatCompletionRequest, CompletionRequest, @@ -66,6 +67,7 @@ tokenizer: Optional[AutoTokenizer] = None model_name: str = "" default_chat_template_kwargs: Dict[str, Any] = {} +custom_message_encoder: Optional[Any] = None _stream_queues: Dict[str, asyncio.Queue] = {} _seq_id_to_request_id: Dict[int, str] = {} _stream_loops: Dict[str, AbstractEventLoop] = {} @@ -590,14 +592,12 @@ async def chat_completions(request: ChatCompletionRequest): merged_kwargs = dict(default_chat_template_kwargs) if request.chat_template_kwargs: merged_kwargs.update(request.chat_template_kwargs) - merged_kwargs["tokenize"] = False - merged_kwargs["add_generation_prompt"] = True - # Pass tools so the chat template can inject tool declarations - if request.tools: - merged_kwargs["tools"] = request.tools - prompt = tokenizer.apply_chat_template( + prompt = apply_chat_template( + tokenizer, + custom_message_encoder, [msg.to_template_dict() for msg in messages], + tools=request.tools, **merged_kwargs, ) @@ -825,6 +825,7 @@ async def stop_profile(): def main(): """Main entry point for the server.""" global engine, tokenizer, model_name, default_chat_template_kwargs, _request_logger + global custom_message_encoder parser = argparse.ArgumentParser(description="ATOM OpenAI API Server") EngineArgs.add_cli_args(parser) @@ -869,6 +870,7 @@ def main(): logger.info(f"Loading tokenizer from {args.model}...") tokenizer = _load_tokenizer(args.model, args.trust_remote_code) model_name = args.model + custom_message_encoder = load_custom_message_encoder(args.model) logger.info(f"Initializing engine with model {args.model}...") engine_args = EngineArgs.from_cli_args(args) diff --git a/atom/entrypoints/openai/chat_encoders.py b/atom/entrypoints/openai/chat_encoders.py new file mode 100644 index 0000000000..d3fb0462ce --- /dev/null +++ b/atom/entrypoints/openai/chat_encoders.py @@ -0,0 +1,118 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Chat-template dispatch for the OpenAI chat endpoint. + +Some models (e.g. DeepSeek V4) ship a Python encoder under ``/encoding/`` +instead of a Jinja ``chat_template``. This module discovers such encoders at +server startup and provides :func:`apply_chat_template`, a single entry point +that the request handler calls — it transparently routes to the custom encoder +when one was found, or to ``tokenizer.apply_chat_template`` otherwise. +""" + +import glob +import importlib.util +import logging +import os +from typing import Any, Callable, List, Optional + +from huggingface_hub import snapshot_download + +logger = logging.getLogger("atom") + +MessageEncoder = Callable[..., str] + + +def _resolve_model_path(model: str) -> str: + if os.path.isdir(model): + return model + try: + return snapshot_download(model, local_files_only=True, allow_patterns=[]) + except Exception: + return model + + +def _load_encoder_from_dir(model_path: str) -> Optional[MessageEncoder]: + """Look for ``/encoding/encoding_*.py`` and load ``encode_messages``. + + Returns ``None`` when the directory or matching file is absent (model uses + the standard Jinja path). Returns ``None`` and warns on ambiguity (multiple + matches) or load failures. + + """ + enc_dir = os.path.join(model_path, "encoding") + if not os.path.isdir(enc_dir): + return None + + candidates = sorted(glob.glob(os.path.join(enc_dir, "encoding_*.py"))) + if not candidates: + return None + if len(candidates) > 1: + logger.warning( + f"Multiple encoding_*.py found in {enc_dir}, refusing to guess: " + f"{[os.path.basename(p) for p in candidates]}" + ) + return None + + enc_path = candidates[0] + module_name = os.path.splitext(os.path.basename(enc_path))[0] + try: + spec = importlib.util.spec_from_file_location(module_name, enc_path) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + raw = getattr(mod, "encode_messages") + except Exception as e: + logger.warning(f"Failed to load encoder from {enc_path}: {e}") + return None + + # also valid is "chat" (non-thinking short-form). May need to add as an option. + # Revisit when a second model ships an encode_*.py — the default may need to be per-model. + def encode(messages, **kwargs): + kwargs.setdefault("thinking_mode", "thinking") + return raw(messages, **kwargs) + + logger.info(f"Loaded message encoder from {enc_path}") + return encode + + +def load_custom_message_encoder(model_path: str) -> Optional[MessageEncoder]: + """Probe ``model_path`` once at startup for a custom message encoder. + + Returns the encoder, or ``None`` when the model uses the standard Jinja + ``chat_template`` path. Result should be cached by the caller — this does + filesystem IO and a Python import. + """ + return _load_encoder_from_dir(_resolve_model_path(model_path)) + + +def apply_chat_template( + tokenizer: Any, + custom_encoder: Optional[MessageEncoder], + messages: List[dict], + *, + tools: Optional[List[dict]] = None, + **kwargs: Any, +) -> str: + """Render ``messages`` to a prompt string. + + Dispatches to ``custom_encoder`` if one was discovered for this model, + otherwise to ``tokenizer.apply_chat_template``. Jinja-only kwargs + (``tokenize``, ``add_generation_prompt``) are stripped on the custom + path; ``tools`` are forwarded only on the Jinja path (custom encoders + don't currently have a tools API — caller is warned and tools are + dropped). + """ + if custom_encoder is not None: + for k in ("tokenize", "add_generation_prompt"): + kwargs.pop(k, None) + if tools: + logger.warning( + "tools= is not supported with the custom message encoder; ignoring." + ) + return custom_encoder(messages, **kwargs) + + kwargs["tokenize"] = False + kwargs["add_generation_prompt"] = True + if tools: + kwargs["tools"] = tools + return tokenizer.apply_chat_template(messages, **kwargs) diff --git a/atom/entrypoints/openai/reasoning.py b/atom/entrypoints/openai/reasoning.py index 7759a04816..8e5cd6f3d1 100644 --- a/atom/entrypoints/openai/reasoning.py +++ b/atom/entrypoints/openai/reasoning.py @@ -23,13 +23,22 @@ def separate_reasoning(text: str) -> Tuple[Optional[str], str]: Tuple of (reasoning_content, content). reasoning_content is None if no thinking block was found. """ - # Check for closed thinking block + # Check for closed thinking block: ... match = re.match(r"(.*?)\s*(.*)", text, flags=re.DOTALL) if match: reasoning = match.group(1).strip() content = match.group(2).strip() return (reasoning if reasoning else None, content) + # Check for without — models like MiniMax M2.7 don't + # generate (the chat template injects it as part of the prompt), + # so the model output starts with reasoning text directly. + if "" in text: + reasoning, _, content = text.partition("") + reasoning = reasoning.strip() + content = content.strip() + return (reasoning if reasoning else None, content) + # Check for unclosed thinking block (truncated response) match = re.match(r"(.*)", text, flags=re.DOTALL) if match: @@ -89,7 +98,23 @@ def process(self, text: str) -> list: elif self.buf: results.append(("reasoning_content", self.buf)) self.buf = "" + elif "" in self.buf: + # No but found — model started reasoning + # without tag (e.g., MiniMax M2.7 where the chat + # template injects as part of the prompt). + reasoning = self.buf.split("", 1)[0] + after = self.buf.split("", 1)[1].lstrip("\n") + if reasoning: + results.append(("reasoning_content", reasoning)) + self.state = 2 + self.buf = "" + if after: + results.extend(self._process_content(after)) elif len(self.buf) > 7 and "<" not in self.buf: + # No tag found — emit as content. For models that + # don't emit (MiniMax), streaming reasoning separation + # requires buffering the entire response, which is impractical. + # Non-streaming path handles this correctly via separate_reasoning(). results.append(("content", self.buf)) self.buf = "" diff --git a/atom/examples/simple_inference.py b/atom/examples/simple_inference.py index fda0de7c63..73d51c2e69 100644 --- a/atom/examples/simple_inference.py +++ b/atom/examples/simple_inference.py @@ -2,9 +2,12 @@ # Copyright (C) 2024-2025, Advanced Micro Devices, Inc. All rights reserved. import argparse -import os from atom import SamplingParams +from atom.entrypoints.openai.chat_encoders import ( + apply_chat_template, + load_custom_message_encoder, +) from atom.model_engine.arg_utils import EngineArgs from transformers import AutoTokenizer @@ -61,41 +64,11 @@ def main(): temperature=args.temperature, max_tokens=args.max_tokens ) - # Apply chat template. DeepSeek-V4 ships without a HuggingFace - # chat_template; use its custom encoding_dsv4.py if available. - if getattr(tokenizer, "chat_template", None): - prompts = [ - tokenizer.apply_chat_template( - [{"role": "user", "content": prompt}], - tokenize=False, - add_generation_prompt=True, - enable_thinking=True, - ) - for prompt in prompts - ] - else: - try: - import importlib.util - - enc_path = os.path.join(args.model, "encoding", "encoding_dsv4.py") - if os.path.exists(enc_path): - spec = importlib.util.spec_from_file_location("encoding_dsv4", enc_path) - enc_mod = importlib.util.module_from_spec(spec) - spec.loader.exec_module(enc_mod) - prompts = [ - enc_mod.encode_messages( - [{"role": "user", "content": p}], thinking_mode="chat" - ) - for p in prompts - ] - print( - f" (applied V4 encoding_dsv4, prompt tokens: " - f"{[len(tokenizer.encode(p)) for p in prompts]})" - ) - else: - print(" (tokenizer has no chat_template — feeding raw prompts as-is)") - except Exception as e: - print(f" (V4 encoding failed: {e} — feeding raw prompts as-is)") + custom_encoder = load_custom_message_encoder(args.model) + prompts = [ + apply_chat_template(tokenizer, custom_encoder, [{"role": "user", "content": p}]) + for p in prompts + ] print("This is prompts:", prompts) # print("Warming up...") # _ = llm.generate(["warmup"], sampling_params) diff --git a/atom/kv_transfer/disaggregation/aggregator.py b/atom/kv_transfer/disaggregation/aggregator.py index 4c4afe59d4..bbe74f0f0e 100644 --- a/atom/kv_transfer/disaggregation/aggregator.py +++ b/atom/kv_transfer/disaggregation/aggregator.py @@ -28,10 +28,10 @@ class KVOutputAggregator: """Aggregates :class:`KVConnectorOutput` from all TP workers. - Uses a countdown approach: when a request ID first appears, a counter - is initialized to ``world_size``. Each worker that reports it as - finished decrements the counter. When the counter reaches zero, - the request is considered globally complete and is emitted. + Tracks which unique worker indices have reported each request as + finished. A request is globally complete only when all + ``world_size`` workers have reported it — duplicate reports from + the same worker (e.g. from retried notifications) are ignored. Args: world_size: Number of TP workers to aggregate over. @@ -48,8 +48,8 @@ def __init__(self, world_size: int = 8) -> None: if world_size <= 0: raise ValueError(f"world_size must be positive, got {world_size}") self._world_size = world_size - self._remaining_sending: dict[str, int] = {} - self._remaining_recving: dict[str, int] = {} + self._seen_sending: dict[str, set[int]] = {} + self._seen_recving: dict[str, set[int]] = {} @property def world_size(self) -> int: @@ -60,6 +60,7 @@ def aggregate(self, worker_outputs: list[KVConnectorOutput]) -> KVConnectorOutpu Args: worker_outputs: One :class:`KVConnectorOutput` per worker. + The list index is the worker index. Returns: A new :class:`KVConnectorOutput` containing only request IDs @@ -68,40 +69,29 @@ def aggregate(self, worker_outputs: list[KVConnectorOutput]) -> KVConnectorOutpu if not worker_outputs: return KVConnectorOutput() - all_sending_ids: set[str] = set() - all_recving_ids: set[str] = set() - for wo in worker_outputs: - if wo.finished_sending: - all_sending_ids.update(wo.finished_sending) - if wo.finished_recving: - all_recving_ids.update(wo.finished_recving) - - for rid in all_sending_ids: - if rid not in self._remaining_sending: - self._remaining_sending[rid] = self._world_size - - for rid in all_recving_ids: - if rid not in self._remaining_recving: - self._remaining_recving[rid] = self._world_size - - for wo in worker_outputs: + for worker_idx, wo in enumerate(worker_outputs): if wo.finished_sending: for rid in wo.finished_sending: - if rid in self._remaining_sending: - self._remaining_sending[rid] -= 1 - + self._seen_sending.setdefault(rid, set()).add(worker_idx) if wo.finished_recving: for rid in wo.finished_recving: - if rid in self._remaining_recving: - self._remaining_recving[rid] -= 1 - - done_sending = {rid for rid, cnt in self._remaining_sending.items() if cnt <= 0} - done_recving = {rid for rid, cnt in self._remaining_recving.items() if cnt <= 0} + self._seen_recving.setdefault(rid, set()).add(worker_idx) + + done_sending = { + rid + for rid, workers in self._seen_sending.items() + if len(workers) >= self._world_size + } + done_recving = { + rid + for rid, workers in self._seen_recving.items() + if len(workers) >= self._world_size + } for rid in done_sending: - del self._remaining_sending[rid] + del self._seen_sending[rid] for rid in done_recving: - del self._remaining_recving[rid] + del self._seen_recving[rid] return KVConnectorOutput( finished_sending=done_sending, @@ -110,10 +100,10 @@ def aggregate(self, worker_outputs: list[KVConnectorOutput]) -> KVConnectorOutpu def reset(self) -> None: """Clear all internal tracking state.""" - self._remaining_sending.clear() - self._remaining_recving.clear() + self._seen_sending.clear() + self._seen_recving.clear() @property def pending_count(self) -> tuple[int, int]: """Return ``(num_pending_sending, num_pending_recving)``.""" - return len(self._remaining_sending), len(self._remaining_recving) + return len(self._seen_sending), len(self._seen_recving) diff --git a/atom/kv_transfer/disaggregation/base.py b/atom/kv_transfer/disaggregation/base.py index a798c04b3c..c32902237c 100644 --- a/atom/kv_transfer/disaggregation/base.py +++ b/atom/kv_transfer/disaggregation/base.py @@ -53,6 +53,15 @@ def get_finished(self) -> tuple[set, set]: """ ... + def get_finished_recv_blocks(self) -> list[int]: + """Return block IDs from recently completed receives for GPU memory fence. + + RDMA writes to HBM may not be immediately visible to GPU compute + kernels. Connectors using RDMA should override this to return + blocks that need a GPU-side read-write cycle to ensure coherence. + """ + return [] + class KVConnectorSchedulerBase(ABC): """Scheduler-side KV connector interface.""" diff --git a/atom/kv_transfer/disaggregation/factory.py b/atom/kv_transfer/disaggregation/factory.py index e6867c254d..d18d621c40 100644 --- a/atom/kv_transfer/disaggregation/factory.py +++ b/atom/kv_transfer/disaggregation/factory.py @@ -31,10 +31,10 @@ class KVConnectorFactory: # Registration (happens once, typically at import time) KVConnectorFactory.register( "moriio", - worker_module="atom.kv_transfer.disaggregation.kv_transfer_engine", - worker_class="KVConnector", - scheduler_module="atom.kv_transfer.disaggregation.kv_transfer_engine", - scheduler_class="KVConnectorScheduler", + worker_module="atom.kv_transfer.disaggregation.moriio.moriio_connector", + worker_class="MoRIIOConnector", + scheduler_module="atom.kv_transfer.disaggregation.moriio.moriio_connector", + scheduler_class="MoRIIOConnectorScheduler", ) # Instantiation (called from forward_context.py) @@ -121,8 +121,16 @@ def create_connector( KVConnectorFactory.register( "moriio", - worker_module="atom.kv_transfer.disaggregation.kv_transfer_engine", - worker_class="KVConnector", - scheduler_module="atom.kv_transfer.disaggregation.kv_transfer_engine", - scheduler_class="KVConnectorScheduler", + worker_module="atom.kv_transfer.disaggregation.moriio.moriio_connector", + worker_class="MoRIIOConnector", + scheduler_module="atom.kv_transfer.disaggregation.moriio.moriio_connector", + scheduler_class="MoRIIOConnectorScheduler", +) + +KVConnectorFactory.register( + "mooncake", + worker_module="atom.kv_transfer.disaggregation.mooncake.mooncake_connector", + worker_class="MooncakeConnector", + scheduler_module="atom.kv_transfer.disaggregation.mooncake.mooncake_connector", + scheduler_class="MooncakeConnectorScheduler", ) diff --git a/atom/kv_transfer/disaggregation/mooncake/__init__.py b/atom/kv_transfer/disaggregation/mooncake/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/atom/kv_transfer/disaggregation/mooncake/mooncake_connector.py b/atom/kv_transfer/disaggregation/mooncake/mooncake_connector.py new file mode 100644 index 0000000000..1599ddf187 --- /dev/null +++ b/atom/kv_transfer/disaggregation/mooncake/mooncake_connector.py @@ -0,0 +1,1051 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2025, Advanced Micro Devices, Inc. All rights reserved. + +""" +Worker-side and scheduler-side KV cache connectors for disaggregated P/D. + +Uses Mooncake TransferEngine for RDMA-based push (WRITE) transfers of +KV cache data from producer (prefill) to consumer (decode) nodes. +""" + +from __future__ import annotations + +import logging +import os +import threading +import time +from concurrent.futures import ThreadPoolExecutor +from typing import Any + +import msgpack +import msgspec +import torch +import zmq + +from atom.config import Config +from atom.kv_transfer.disaggregation.base import ( + KVConnectorBase, + KVConnectorSchedulerBase, +) +from atom.kv_transfer.disaggregation.types import ( + ConnectorMetadata, + ReqId, + TransferId, +) +from atom.model_engine.sequence import Sequence +from atom.utils import get_open_port, make_zmq_path, zmq_socket_ctx +from atom.utils.network import get_ip +from aiter.dist.parallel_state import get_dp_group, get_tp_group + +logger = logging.getLogger("atom") + +# --------------------------------------------------------------------------- +# Mooncake availability check +# --------------------------------------------------------------------------- + +_MOONCAKE_AVAILABLE = False +try: + from mooncake.engine import TransferEngine + + _MOONCAKE_AVAILABLE = True + logger.info("Mooncake TransferEngine loaded successfully") +except ImportError: + logger.warning( + "Mooncake is not available — KV cache disaggregation via mooncake " + "will not work. Install the mooncake package to enable push-mode " + "RDMA transfers." + ) + + +# --------------------------------------------------------------------------- +# Constants +# --------------------------------------------------------------------------- + +MOONCAKE_PING_INTERVAL_SECONDS = 5 +MOONCAKE_MAX_PING_RETRIES = 100 +MOONCAKE_DEFAULT_PROTOCOL = "rdma" +PREFILL_LOOKUP_TIMEOUT = 60 +PREFILL_LOOKUP_POLL_INTERVAL = 0.01 + +# ZMQ side-channel message types +MSG_WRITE_REQUEST = b"write_request" +MSG_WRITE_DONE = b"write_done" +MSG_GET_META = b"get_meta" + + +# --------------------------------------------------------------------------- +# Metadata struct for bootstrap handshake +# --------------------------------------------------------------------------- + + +class MooncakeAgentMetadata( + msgspec.Struct, + omit_defaults=True, + dict=True, + kw_only=True, +): + """Serializable metadata exchanged during the mooncake bootstrap.""" + + engine_id: str + rpc_port: int + kv_caches_base_addr: list[int] | None = None + num_blocks: int = 0 + block_len: int = 0 + + +def _port_offset(dp_rank: int, tp_rank: int, tp_size: int = 1) -> int: + return dp_rank * tp_size + tp_rank + + +# =================================================================== +# MooncakeConnectorScheduler — scheduler-side connector +# =================================================================== + + +class MooncakeConnectorScheduler(KVConnectorSchedulerBase): + def __init__(self, config: Config) -> None: + kv_transfer_config = config.kv_transfer_config + self.is_producer = ( + kv_transfer_config.get("kv_role", "kv_producer") == "kv_producer" + ) + self.handshake_port = get_open_port() + self.engine_id = "None" + self.tp_size = config.tensor_parallel_size + self.dp_size = config.parallel_config.data_parallel_size + self.dp_rank = config.parallel_config.data_parallel_rank + self.host_ip = get_ip() + + # Pending requests: req_id -> (Sequence, block_table) + self._reqs_need_recv: dict[ReqId, tuple[Any, list[int]]] = {} + self._reqs_need_save: dict[ReqId, tuple[Any, list[int]]] = {} + + # Bidirectional transfer_id <-> request_id mapping + self.request_id_to_transfer_id: dict[ReqId, TransferId] = {} + self.transfer_id_to_request_id: dict[TransferId, ReqId] = {} + + def get_num_new_matched_tokens(self, seq: Sequence) -> tuple[int, bool]: + params = seq.kv_transfer_params or {} + + if params.get("do_remote_prefill") and not getattr( + seq, "kv_async_tagged", False + ): + return len(seq.prompt_token_ids), True + + return 0, False + + def build_connector_meta(self) -> ConnectorMetadata: + meta = ConnectorMetadata() + meta.request_id_to_transfer_id = self.request_id_to_transfer_id + + for req_id, (req, block_ids) in self._reqs_need_recv.items(): + assert req.kv_transfer_params is not None + meta.add_new_req_to_recv( + request_id=req_id, + local_block_ids=block_ids, + kv_transfer_params=req.kv_transfer_params, + ) + + # Producer side: pass completed prefill block_ids to worker + for req_id, (req, block_ids) in self._reqs_need_save.items(): + assert req.kv_transfer_params is not None + meta.add_new_req_to_save( + request_id=req_id, + local_block_ids=block_ids, + kv_transfer_params=req.kv_transfer_params, + ) + + if self._reqs_need_recv or self._reqs_need_save: + logger.info( + "[SCHEDULER] build_connector_meta: %d recv, %d save, " "id_map=%s", + len(self._reqs_need_recv), + len(self._reqs_need_save), + meta.request_id_to_transfer_id, + ) + self._reqs_need_recv.clear() + self._reqs_need_save.clear() + return meta + + def update_state_after_alloc(self, seq: Sequence) -> None: + params = seq.kv_transfer_params or {} + + if not self.is_producer: + transfer_id = params.get("transfer_id") + if transfer_id is not None: + self.transfer_id_to_request_id[transfer_id] = seq.id + self.request_id_to_transfer_id[seq.id] = transfer_id + + # Consumer side: queue for remote KV loading + if params.get("do_remote_prefill"): + assert ( + not self.is_producer + ), "Only the decode (consumer) side handles do_remote_prefill" + self._reqs_need_recv[seq.id] = (seq, seq.block_table) + params["do_remote_prefill"] = False + logger.info( + "[SCHEDULER-CONSUMER] Queued req %s for remote KV recv " + "(%d blocks), transfer_id=%s, remote_host=%s, " + "remote_handshake_port=%s", + seq.id, + len(seq.block_table), + params.get("transfer_id"), + params.get("remote_host"), + params.get("remote_handshake_port"), + ) + + # Producer side: queue block_ids for the write listener to look up + if params.get("do_remote_decode"): + assert self.is_producer, "Only the producer side handles do_remote_decode" + self._reqs_need_save[seq.id] = (seq, seq.block_table) + logger.debug( + "Queued req %s for KV save (%d blocks)", + seq.id, + len(seq.block_table), + ) + + def request_finished(self, seq: Sequence) -> None: + first_token_id = seq.output_tokens[0] if seq.output_tokens else None + seq.kv_transfer_params_output = { + "do_remote_prefill": True, + "do_remote_decode": False, + "remote_block_ids": seq.block_table.copy(), + "remote_engine_id": self.engine_id, + "remote_host": self.host_ip, + "remote_port": self.handshake_port, + "tp_size": self.tp_size, + "dp_rank": self.dp_rank, + "transfer_id": seq.id, + "first_token_id": first_token_id, + } + + if not self.is_producer: + transfer_id = self.request_id_to_transfer_id.pop(seq.id, None) + if transfer_id is not None: + self.transfer_id_to_request_id.pop(transfer_id, None) + + +# =================================================================== +# MooncakeConnector — worker-side connector (runs inside each TP rank) +# =================================================================== + + +class MooncakeConnector(KVConnectorBase): + """Worker-side KV cache connector using Mooncake push-mode RDMA. + + Mooncake uses a push/WRITE model: the prefill (producer) node writes + KV cache data directly into the decode (consumer) node's registered + GPU memory via ``batch_transfer_sync_write``. + """ + + def __init__(self, config: Config) -> None: + self.tp_rank = get_tp_group().rank_in_group + self.dp_rank = get_dp_group().rank_in_group + self.tp_size = get_tp_group().world_size + self.dp_size = get_dp_group().world_size + + kv_transfer_config = config.kv_transfer_config + self.local_ip = get_ip() + self._local_ping_port = get_open_port() + + self.is_producer = ( + kv_transfer_config.get("kv_role", "kv_producer") == "kv_producer" + ) + self.is_consumer = not self.is_producer + + # Networking / service discovery config + self.http_port = kv_transfer_config.get("http_port", 8000) + self.proxy_ping_port = kv_transfer_config.get("proxy_ping_port", 36367) + self.proxy_ip = kv_transfer_config.get("proxy_ip") + self.request_address = f"{self.local_ip}:{self.http_port}" + self.protocol = kv_transfer_config.get("protocol", MOONCAKE_DEFAULT_PROTOCOL) + + # Side channel port (ZMQ) — deterministic from config for proxy relay + self.base_handshake_port = kv_transfer_config.get("handshake_port", 6301) + self._side_channel_port = self.base_handshake_port + _port_offset( + self.dp_rank, self.tp_rank, self.tp_size + ) + + # --- Mooncake TransferEngine initialization --- + if not _MOONCAKE_AVAILABLE: + raise RuntimeError( + "Mooncake is not installed but kv_connector='mooncake' was requested. " + "Install the mooncake package to use push-mode RDMA transfers." + ) + + # Determine which RDMA device this TP rank should use. + # On AMD MI300X, each GPU has a paired RDMA NIC: GPU N → rdmaN. + # Registering GPU memory with a non-local RDMA NIC fails with + # EINVAL. Pass the device name as a filter so Mooncake only + # creates a context for the local NIC. + ib_device = kv_transfer_config.get("ib_device", "") + if not ib_device: + ib_device = os.environ.get("ATOM_MOONCAKE_IB_DEVICE", "") + if not ib_device: + gpu_idx = torch.cuda.current_device() + ib_device = f"rdma{gpu_idx}" + logger.info( + "Auto-selecting RDMA device %s for GPU %d (tp_rank=%d)", + ib_device, + gpu_idx, + self.tp_rank, + ) + + self.transfer_engine = TransferEngine() + ret = self.transfer_engine.initialize( + self.local_ip, + "P2PHANDSHAKE", + self.protocol, + ib_device, + ) + if ret != 0: + raise RuntimeError( + f"Mooncake TransferEngine.initialize() failed (ret={ret}) " + f"on ip={self.local_ip}, protocol={self.protocol}, " + f"ib_device={ib_device}" + ) + self.rpc_port = self.transfer_engine.get_rpc_port() + self.engine_id = f"{self.local_ip}:{self.rpc_port}" + logger.info( + "Mooncake TransferEngine initialized: ip=%s, protocol=%s, " + "ib_device=%s, rpc_port=%d", + self.local_ip, + self.protocol, + ib_device, + self.rpc_port, + ) + + # --- KV cache state (populated in register_kv_caches) --- + self.kv_caches: dict[str, Any] | None = None + self.kv_caches_base_addr: list[int] = [] + self._per_block_bytes_list: list[int] = [] + self.kv_cache_shape: tuple[int, ...] | None = None + self.block_len: int = config.kv_cache_block_size + self.num_blocks: int = 0 + self._per_block_bytes: int = 0 + + # --- Producer: completed prefill block_ids cache --- + # Populated from ConnectorMetadata.reqs_to_save each step. + # The write listener looks up block_ids here when consumer requests a write. + self._completed_prefills: dict[ReqId, list[int]] = {} + self._completed_prefills_lock = threading.Lock() + self._completed_prefills_cv = threading.Condition(self._completed_prefills_lock) + self._transfer_refcount: dict[ReqId, int] = {} + self._transfer_refcount_lock = threading.Lock() + + # --- Consumer: pending receive tracking --- + self._pending_recv: set[ReqId] = set() + self._pending_recv_blocks: dict[ReqId, list[int]] = {} + self._notification_port = get_open_port() + + # --- Completion tracking --- + self.done_sending: set[str] = set() + self.done_recving: set[str] = set() + self._completion_lock = threading.Lock() + + # --- GPU memory fence: blocks pending coherence enforcement --- + self._blocks_pending_fence: list[int] = [] + self._fence_lock = threading.Lock() + + # --- Transfer ID mapping (worker side) --- + self.request_id_to_transfer_id: dict[ReqId, TransferId] = {} + + # --- Producer: thread pool for RDMA writes --- + if self.is_producer: + self._send_executor = ThreadPoolExecutor( + max_workers=kv_transfer_config.get("num_worker_threads", 16), + thread_name_prefix="mooncake-send-worker", + ) + + # --- ZMQ for metadata exchange --- + self.zmq_context = zmq.Context() + + # --- Producer: persistent socket cache for write-done notifications --- + self._notify_sockets: dict[str, zmq.Socket] = {} + self._notify_sockets_lock = threading.Lock() + + # --- Msgspec encoder/decoder for bootstrap metadata --- + self._encoder = msgspec.msgpack.Encoder() + self._decoder = msgspec.msgpack.Decoder(MooncakeAgentMetadata) + + # --- Service discovery ping (rank 0 only) --- + if self.tp_rank == 0 and self.dp_rank == 0: + self._ping_thread = threading.Thread( + target=self._service_discovery_ping, + args=(self.zmq_context,), + daemon=True, + name="mooncake-ping", + ) + self._ping_thread.start() + + # ----------------------------------------------------------------- + # Service discovery + # ----------------------------------------------------------------- + + def _service_discovery_ping(self, zmq_context: zmq.Context) -> None: + """Periodically register with the proxy (rank 0 only).""" + grpc_endpoint = f"http://{self.request_address}/v1/completions" + role_code = "P" if self.is_producer else "D" + retry_count = 0 + msg_index = 1 + proxy_path = f"tcp://{self.proxy_ip}:{self.proxy_ping_port}" + + with zmq_context.socket(zmq.DEALER) as sock: + sock.connect(proxy_path) + + while True: + try: + registration_data = { + "type": "register", + "role": role_code, + "index": str(msg_index), + "request_address": grpc_endpoint, + "rpc_port": self.rpc_port, + "handshake_port": self.base_handshake_port, + "dp_size": self.dp_size, + "tp_size": self.tp_size, + "transfer_mode": "write", + } + sock.send(msgpack.dumps(registration_data)) + logger.debug( + "Ping #%d sent to %s (role=%s)", + msg_index, + proxy_path, + role_code, + ) + retry_count = 0 + + except ConnectionRefusedError: + logger.info( + "Proxy connection refused: %s -> %s", + self.local_ip, + proxy_path, + ) + retry_count += 1 + + except OSError as e: + logger.info("OS error during ping: %s", e) + retry_count += 1 + + except Exception as e: + logger.info("Unexpected ping error: %s", e) + retry_count += 1 + if retry_count >= MOONCAKE_MAX_PING_RETRIES: + logger.error( + "Ping failed after %d retries, aborting", + MOONCAKE_MAX_PING_RETRIES, + ) + raise RuntimeError( + f"Service discovery ping failed after " + f"{retry_count} retries" + ) from e + + finally: + time.sleep(MOONCAKE_PING_INTERVAL_SECONDS) + msg_index += 1 + + # ----------------------------------------------------------------- + # KVConnectorBase: register_kv_caches + # ----------------------------------------------------------------- + + # RDMA MR size limit — must stay under device max_mr_size (2 GiB on + # AMD ROCm RDMA). Mooncake truncates silently, leaving addresses + # beyond the limit unregistered. + _MAX_RDMA_CHUNK_BYTES = 2 * 1024 * 1024 * 1024 - 64 * 1024 + + def register_kv_caches(self, kv_caches: dict[str, Any]) -> None: + """Register all KV cache tensors with the Mooncake TransferEngine. + + Must be called once after model loading and KV cache allocation. + Starts the side-channel listener threads for write coordination. + + Large tensors (> 2 GiB) are split into chunks for RDMA memory + registration to stay within the device max_mr_size. The logical + per-tensor base addresses (used for transfer offset computation) + remain unchanged. + + For fp8 KV caches, k_scale and v_scale tensors are also registered + so that quantization scale factors are transferred alongside data. + """ + self.kv_caches = kv_caches + + # Logical per-tensor info (for transfer address computation) + tensor_ptrs: list[int] = [] + tensor_sizes: list[int] = [] + + # Chunked regions (for RDMA registration) + reg_ptrs: list[int] = [] + reg_sizes: list[int] = [] + + def _register_tensor(tensor): + self.kv_caches_base_addr.append(tensor.data_ptr()) + sz = tensor.element_size() + if is_mla: + bpb = self.block_len * tensor.stride(0) * sz + else: + bpb = tensor.stride(0) * sz + self._per_block_bytes_list.append(bpb) + + base = tensor.data_ptr() + total = tensor.numel() * sz + tensor_ptrs.append(base) + tensor_sizes.append(total) + + # Chunk for RDMA registration + offset = 0 + while offset < total: + chunk = min(self._MAX_RDMA_CHUNK_BYTES, total - offset) + reg_ptrs.append(base + offset) + reg_sizes.append(chunk) + offset += chunk + + for layer_name, kv_cache in kv_caches.items(): + k_cache = kv_cache.k_cache + v_cache = kv_cache.v_cache + is_mla = v_cache is None + + if self.kv_cache_shape is None: + self.kv_cache_shape = k_cache.shape + + _register_tensor(k_cache) + + if not is_mla: + _register_tensor(v_cache) + + if kv_cache.k_scale is not None: + _register_tensor(kv_cache.k_scale) + if kv_cache.v_scale is not None: + _register_tensor(kv_cache.v_scale) + + logger.info( + "Registering %d tensors as %d RDMA chunks " "(max_chunk=%.2f GiB)", + len(tensor_ptrs), + len(reg_ptrs), + self._MAX_RDMA_CHUNK_BYTES / (1024**3), + ) + for i, (ptr, sz_bytes) in enumerate(zip(reg_ptrs, reg_sizes)): + logger.info( + " reg_chunk[%d] ptr=0x%x size=%d (%.2f GiB)", + i, + ptr, + sz_bytes, + sz_bytes / (1024**3), + ) + + ret = self.transfer_engine.batch_register_memory(reg_ptrs, reg_sizes) + if ret != 0: + logger.error( + "batch_register_memory FAILED (ret=%d). " + "Trying individual registration as fallback...", + ret, + ) + for i, (ptr, sz_bytes) in enumerate(zip(reg_ptrs, reg_sizes)): + r = self.transfer_engine.register_memory(ptr, sz_bytes) + logger.info( + " register_memory[%d] ptr=0x%x size=%d (%.2f GiB) " "=> ret=%d", + i, + ptr, + sz_bytes, + sz_bytes / (1024**3), + r, + ) + else: + logger.info( + "batch_register_memory OK (%d chunks from %d tensors)", + len(reg_ptrs), + len(tensor_ptrs), + ) + + # Block geometry from last layer + last_kv = list(kv_caches.values())[-1] + k = last_kv.k_cache + is_mla = last_kv.v_cache is None + sz = k.element_size() + + if is_mla: + self.num_blocks = k.shape[0] // self.block_len + self._per_block_bytes = self.block_len * k.stride(0) * sz + else: + self.num_blocks = k.shape[0] + self._per_block_bytes = k.stride(0) * sz + + logger.info( + "Mooncake registered %d tensors (%d RDMA chunks, %d blocks, " + "per_block_bytes=%s, block_len=%d, is_mla=%s, " + "k_shape=%s, k_stride=%s, elem_sz=%d, " + "tensor_sizes_MB=%s, chunk_sizes_MB=%s)", + len(self.kv_caches_base_addr), + len(reg_ptrs), + self.num_blocks, + sorted(set(self._per_block_bytes_list)), + self.block_len, + is_mla, + k.shape, + k.stride(), + sz, + [s // (1024 * 1024) for s in sorted(set(tensor_sizes))], + [s // (1024 * 1024) for s in sorted(set(reg_sizes))], + ) + + logger.info( + "Mooncake KV registration complete: role=%s, engine_id=%s, " + "rpc_port=%d, num_regions=%d, num_blocks=%d, " + "base_addrs=[%s]", + "PRODUCER" if self.is_producer else "CONSUMER", + self.engine_id, + self.rpc_port, + len(self.kv_caches_base_addr), + self.num_blocks, + ", ".join(f"0x{a:x}" for a in self.kv_caches_base_addr[:4]), + ) + + # Build metadata for bootstrap exchange + self._local_metadata = MooncakeAgentMetadata( + engine_id=self.engine_id, + rpc_port=self.rpc_port, + kv_caches_base_addr=self.kv_caches_base_addr, + num_blocks=self.num_blocks, + block_len=self.block_len, + ) + + # Start side channel threads + if self.is_producer: + self._write_listener_thread = threading.Thread( + target=self._write_listener, + daemon=True, + name="mooncake-write-listener", + ) + self._write_listener_thread.start() + else: + self._notification_listener_thread = threading.Thread( + target=self._notification_listener, + daemon=True, + name="mooncake-notify-listener", + ) + self._notification_listener_thread.start() + + # ----------------------------------------------------------------- + # KVConnectorBase: start_load_kv + # ----------------------------------------------------------------- + + def start_load_kv(self, metadata: ConnectorMetadata) -> None: + """Initiate KV transfers for pending requests. + + **Producer side**: Cache completed prefill block_ids from + ``metadata.reqs_to_save`` so the write listener can look them up. + + **Consumer side**: For each pending recv request, connect to the + producer's ZMQ side channel and send a write request with our + memory addresses and block allocation. + """ + if metadata is None: + return + + self.request_id_to_transfer_id = metadata.request_id_to_transfer_id + + # Producer: cache block_ids from completed prefills + if self.is_producer: + for req_id, meta in metadata.reqs_to_save.items(): + with self._completed_prefills_cv: + self._completed_prefills[req_id] = meta.local_block_ids + self._completed_prefills_cv.notify_all() + logger.info( + "[PRODUCER] Cached %d prefill blocks for req %s", + len(meta.local_block_ids), + req_id, + ) + return + + # Consumer: send write requests to producer + if not metadata.reqs_to_recv: + return + + logger.info( + "[CONSUMER] start_load_kv: %d reqs_to_recv, id_map=%s", + len(metadata.reqs_to_recv), + metadata.request_id_to_transfer_id, + ) + + for req_id, meta in metadata.reqs_to_recv.items(): + remote_tp_size = meta.tp_size + if remote_tp_size != self.tp_size: + remote_tp_rank = self.tp_rank % remote_tp_size + else: + remote_tp_rank = self.tp_rank + remote_port = meta.remote_handshake_port + _port_offset( + meta.remote_dp_rank, remote_tp_rank, remote_tp_size + ) + remote_addr = make_zmq_path("tcp", meta.remote_host, remote_port) + + unique_bpb = sorted(set(self._per_block_bytes_list)) + logger.info( + "[CONSUMER] Sending write_request for req %s (transfer_id=%s) " + "to %s (handshake_port=%d, dp_rank=%d, " + "local_tp=%d, remote_tp=%d/%d), " + "dst_block_ids=%s, num_regions=%d, " + "bytes/block=%s, num_blocks=%d", + req_id, + meta.transfer_id, + remote_addr, + meta.remote_handshake_port, + meta.remote_dp_rank, + self.tp_rank, + remote_tp_rank, + remote_tp_size, + meta.local_block_ids[:10], + len(self.kv_caches_base_addr), + unique_bpb, + self.num_blocks, + ) + + write_request = msgpack.dumps( + { + "request_id": req_id, + "transfer_id": meta.transfer_id, + "consumer_host": self.local_ip, + "consumer_rpc_port": self.rpc_port, + "consumer_base_addrs": self.kv_caches_base_addr, + "dst_block_ids": meta.local_block_ids, + "notify_host": self.local_ip, + "notify_port": self._notification_port, + "consumer_tp_size": self.tp_size, + } + ) + + with self._notify_sockets_lock: + sock = self._notify_sockets.get(remote_addr) + if sock is None: + sock = self.zmq_context.socket(zmq.DEALER) + sock.setsockopt(zmq.LINGER, 5000) + sock.setsockopt(zmq.SNDHWM, 0) + sock.connect(remote_addr) + self._notify_sockets[remote_addr] = sock + sock.send_multipart([MSG_WRITE_REQUEST, write_request]) + + self._pending_recv.add(req_id) + self._pending_recv_blocks[req_id] = list(meta.local_block_ids) + logger.info( + "[CONSUMER] write_request sent for req %s to %s", + req_id, + remote_addr, + ) + + # ----------------------------------------------------------------- + # KVConnectorBase: get_finished + # ----------------------------------------------------------------- + + def get_finished(self) -> tuple[set, set]: + """Return ``(done_sending, done_recving)`` and clear internal sets.""" + with self._completion_lock: + ds = self.done_sending.copy() + dr = self.done_recving.copy() + self.done_sending.clear() + self.done_recving.clear() + if ds or dr: + logger.info( + "[%s] get_finished: sending=%s, recving=%s", + "PRODUCER" if self.is_producer else "CONSUMER", + ds, + dr, + ) + return ds, dr + + def get_finished_recv_blocks(self) -> list[int]: + """Return block IDs from recently completed RDMA receives.""" + with self._fence_lock: + blocks = self._blocks_pending_fence + self._blocks_pending_fence = [] + return blocks + + # ----------------------------------------------------------------- + # Producer: write listener (ZMQ ROUTER) + # ----------------------------------------------------------------- + + def _write_listener(self) -> None: + """Accept write requests from consumers and dispatch RDMA writes.""" + path = make_zmq_path("tcp", "*", self._side_channel_port) + logger.info("Mooncake write listener bound to %s", path) + + with zmq_socket_ctx(path, zmq.ROUTER, bind=True) as sock: + while True: + parts = sock.recv_multipart() + identity, msg_type = parts[0], parts[1] + + if msg_type == MSG_GET_META: + encoded = self._encoder.encode(self._local_metadata) + sock.send_multipart([identity, b"", encoded]) + logger.debug("Sent metadata to peer") + + elif msg_type == MSG_WRITE_REQUEST: + request_data = msgpack.loads(parts[2]) + logger.info( + "[PRODUCER] Received write_request for req %s " + "(transfer_id=%s, consumer=%s:%s)", + request_data["request_id"], + request_data.get("transfer_id"), + request_data.get("consumer_host"), + request_data.get("consumer_rpc_port"), + ) + self._send_executor.submit(self._execute_transfer, request_data) + + else: + logger.error("Unknown message type: %s", msg_type) + + # ----------------------------------------------------------------- + # Producer: execute RDMA write + # ----------------------------------------------------------------- + + def _execute_transfer(self, request_data: dict) -> None: + """Compute offsets and perform RDMA write for a single request.""" + req_id = request_data["request_id"] + transfer_id = request_data.get("transfer_id", req_id) + consumer_host = request_data["consumer_host"] + consumer_rpc_port = request_data["consumer_rpc_port"] + consumer_base_addrs = request_data["consumer_base_addrs"] + dst_block_ids = request_data["dst_block_ids"] + notify_host = request_data["notify_host"] + notify_port = request_data["notify_port"] + consumer_tp_size = request_data.get("consumer_tp_size", self.tp_size) + consumers_per_rank = max(1, consumer_tp_size // self.tp_size) + + logger.info( + "[PRODUCER] _execute_transfer: req_id=%s, transfer_id=%s, " + "consumer=%s:%s, dst_blocks=%d", + req_id, + transfer_id, + consumer_host, + consumer_rpc_port, + len(dst_block_ids), + ) + + # Look up producer's block_ids by transfer_id (= prefill's seq.id), + # not request_id (= decode's seq.id). + src_block_ids = self._wait_for_prefill_blocks(transfer_id) + if src_block_ids is None: + logger.error( + "[PRODUCER] Timed out waiting for prefill blocks for " + "transfer_id=%s (req_id=%s). Available keys: %s", + transfer_id, + req_id, + list(self._completed_prefills.keys()), + ) + return + + target = f"{consumer_host}:{consumer_rpc_port}" + + # Verify the consumer's segment is discoverable via P2P handshake + if hasattr(self.transfer_engine, "get_first_buffer_address"): + remote_buf = self.transfer_engine.get_first_buffer_address(target) + logger.info( + "[PRODUCER] P2P segment probe for %s: " "get_first_buffer_address=0x%x", + target, + remote_buf, + ) + if remote_buf == 0: + logger.error( + "[PRODUCER] Consumer %s has NO registered buffers visible " + "via P2P handshake! Consumer batch_register_memory may " + "have failed, or segment descriptors haven't propagated.", + target, + ) + + src_addrs: list[int] = [] + dst_addrs: list[int] = [] + sizes: list[int] = [] + + num_regions = len(self.kv_caches_base_addr) + if num_regions != len(consumer_base_addrs): + logger.error( + "[PRODUCER] REGION COUNT MISMATCH: producer has %d regions, " + "consumer has %d regions", + num_regions, + len(consumer_base_addrs), + ) + if len(src_block_ids) != len(dst_block_ids): + logger.error( + "[PRODUCER] BLOCK COUNT MISMATCH: src has %d blocks, " + "dst has %d blocks", + len(src_block_ids), + len(dst_block_ids), + ) + for region_idx in range(num_regions): + src_base = self.kv_caches_base_addr[region_idx] + dst_base = consumer_base_addrs[region_idx] + bpb = self._per_block_bytes_list[region_idx] + for sb, db in zip(src_block_ids, dst_block_ids): + src_addrs.append(src_base + sb * bpb) + dst_addrs.append(dst_base + db * bpb) + sizes.append(bpb) + + unique_bpb = sorted(set(self._per_block_bytes_list)) + logger.info( + "[PRODUCER] Starting RDMA write: req=%s, target=%s, " + "%d regions × %d blocks, bytes/block=%s, " + "src_block_ids=%s, dst_block_ids=%s, " + "total_transfer_entries=%d, total_bytes=%d", + req_id, + target, + num_regions, + len(src_block_ids), + unique_bpb, + src_block_ids[:10], + dst_block_ids[:10], + len(src_addrs), + sum(sizes), + ) + + if dst_addrs: + logger.info( + "[PRODUCER] dst_addr range: 0x%x -- 0x%x (first region), " + "consumer_base_addrs[0]=0x%x", + min(dst_addrs[: len(dst_block_ids)]), + max(dst_addrs[: len(dst_block_ids)]) + sizes[0], + consumer_base_addrs[0], + ) + + max_entries_per_batch = 4096 + total_entries = len(src_addrs) + max_retries = 3 + + for chunk_start in range(0, total_entries, max_entries_per_batch): + chunk_end = min(chunk_start + max_entries_per_batch, total_entries) + chunk_src = src_addrs[chunk_start:chunk_end] + chunk_dst = dst_addrs[chunk_start:chunk_end] + chunk_sizes = sizes[chunk_start:chunk_end] + + retry_delay = 2.0 + for attempt in range(max_retries): + try: + ret = self.transfer_engine.batch_transfer_sync_write( + target, chunk_src, chunk_dst, chunk_sizes + ) + if ret == 0: + break + logger.error( + "[PRODUCER] RDMA write chunk error %d for req %s → %s " + "(entries %d-%d/%d, attempt %d/%d)", + ret, + req_id, + target, + chunk_start, + chunk_end, + total_entries, + attempt + 1, + max_retries, + ) + except Exception: + logger.exception( + "[PRODUCER] RDMA write chunk FAILED for req %s " + "(entries %d-%d/%d, attempt %d/%d)", + req_id, + chunk_start, + chunk_end, + total_entries, + attempt + 1, + max_retries, + ) + ret = -1 + if ret == 0: + break + if attempt < max_retries - 1: + time.sleep(retry_delay) + retry_delay *= 2 + else: + return + + logger.info( + "[PRODUCER] RDMA write completed for req %s → %s (%d entries)", + req_id, + target, + total_entries, + ) + + # Notify this consumer immediately — its data is written. + self._send_write_done(notify_host, notify_port, req_id) + + # Track how many consumers still need this transfer_id's blocks. + # Only mark done_sending (which triggers block deallocation on the + # scheduler) after ALL consumers sharing this producer rank have + # completed their transfers. + all_done = False + with self._transfer_refcount_lock: + if transfer_id not in self._transfer_refcount: + self._transfer_refcount[transfer_id] = consumers_per_rank + self._transfer_refcount[transfer_id] -= 1 + if self._transfer_refcount[transfer_id] <= 0: + self._transfer_refcount.pop(transfer_id) + all_done = True + + if all_done: + with self._completion_lock: + self.done_sending.add(transfer_id) + with self._completed_prefills_lock: + self._completed_prefills.pop(transfer_id, None) + logger.info( + "[PRODUCER] All %d consumers served for transfer_id=%s, " + "blocks released", + consumers_per_rank, + transfer_id, + ) + + def _wait_for_prefill_blocks(self, req_id: str) -> list[int] | None: + """Wait until prefill block_ids are available for this request.""" + with self._completed_prefills_cv: + ready = self._completed_prefills_cv.wait_for( + lambda: req_id in self._completed_prefills, + timeout=PREFILL_LOOKUP_TIMEOUT, + ) + if ready: + return self._completed_prefills[req_id] + return None + + def _send_write_done(self, host: str, port: int, req_id: str) -> None: + """Send write-done notification to consumer via persistent socket. + + Sends the notification multiple times for reliability — the consumer + uses a set so duplicates are harmless. + """ + path = make_zmq_path("tcp", host, port) + notification = msgpack.dumps({"request_id": req_id}) + with self._notify_sockets_lock: + sock = self._notify_sockets.get(path) + if sock is None: + sock = self.zmq_context.socket(zmq.DEALER) + sock.setsockopt(zmq.LINGER, 5000) + sock.setsockopt(zmq.SNDHWM, 0) + sock.connect(path) + self._notify_sockets[path] = sock + for _ in range(3): + sock.send_multipart([MSG_WRITE_DONE, notification]) + logger.info("[PRODUCER] write-done sent for req %s", req_id) + + # ----------------------------------------------------------------- + # Consumer: notification listener (ZMQ ROUTER) + # ----------------------------------------------------------------- + + def _notification_listener(self) -> None: + """Receive write-done notifications from producers.""" + path = make_zmq_path("tcp", "*", self._notification_port) + logger.info("Mooncake notification listener bound to %s", path) + + with zmq_socket_ctx(path, zmq.ROUTER, bind=True) as sock: + while True: + parts = sock.recv_multipart() + msg_type = parts[1] + + if msg_type == MSG_WRITE_DONE: + data = msgpack.loads(parts[2]) + req_id = data["request_id"] + dst_blocks = self._pending_recv_blocks.pop(req_id, None) + if dst_blocks: + with self._fence_lock: + self._blocks_pending_fence.extend(dst_blocks) + with self._completion_lock: + self.done_recving.add(req_id) + self._pending_recv.discard(req_id) + logger.info( + "[CONSUMER] Write-done received for req %s, " + "done_recving now: %s", + req_id, + self.done_recving, + ) + else: + logger.error("Unknown notification type: %s", msg_type) diff --git a/atom/kv_transfer/disaggregation/moriio/__init__.py b/atom/kv_transfer/disaggregation/moriio/__init__.py new file mode 100644 index 0000000000..8fbac11053 --- /dev/null +++ b/atom/kv_transfer/disaggregation/moriio/__init__.py @@ -0,0 +1,12 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2025, Advanced Micro Devices, Inc. All rights reserved. + +from atom.kv_transfer.disaggregation.moriio.moriio_connector import ( + MoRIIOConnector, + MoRIIOConnectorScheduler, +) + +__all__ = [ + "MoRIIOConnector", + "MoRIIOConnectorScheduler", +] diff --git a/atom/kv_transfer/disaggregation/moriio/moriio_common.py b/atom/kv_transfer/disaggregation/moriio/moriio_common.py new file mode 100644 index 0000000000..40d7db2c8f --- /dev/null +++ b/atom/kv_transfer/disaggregation/moriio/moriio_common.py @@ -0,0 +1,178 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2025, Advanced Micro Devices, Inc. All rights reserved. + +""" +Shared types, constants, enums, and helpers for the MoRIIO KV connector. + +This module has no heavy dependencies (no torch at import time, no RDMA +engine instances) so it can be imported freely by the other moriio +submodules. +""" + +from __future__ import annotations + +import logging +import threading +from enum import Enum +from typing import Optional + +import msgspec + +from atom.kv_transfer.disaggregation.utils import ( # noqa: F401 + MAX_RDMA_CHUNK_BYTES, + chunk_tensor_for_rdma, +) + +logger = logging.getLogger("atom") + +_chunk_tensor_for_rdma = chunk_tensor_for_rdma + +# --------------------------------------------------------------------------- +# MoRIIO availability check +# --------------------------------------------------------------------------- + +_MORIIO_AVAILABLE = False +try: + import mori.io # noqa: F401 + + _MORIIO_AVAILABLE = True + logger.info("MoRIIO RDMA library loaded successfully") +except ImportError: + logger.warning( + "MoRIIO is not available — KV cache disaggregation will not work. " + "Install the mori package to enable RDMA transfers." + ) + + +# --------------------------------------------------------------------------- +# Msgspec metadata structs +# --------------------------------------------------------------------------- + + +class MoRIIOAgentMetadata( + msgspec.Struct, + omit_defaults=True, + dict=True, + kw_only=True, +): + """Serializable metadata exchanged during the RDMA handshake.""" + + engine_id: str + agent_metadata: bytes + kv_caches_base_addr: Optional[list[int]] = None + num_blocks: int = 0 + block_len: int = 0 + attn_backend_name: str = "aiter" + + +# --------------------------------------------------------------------------- +# Enums & role management +# --------------------------------------------------------------------------- + + +class Role(Enum): + """Role of the current engine instance in the P/D architecture.""" + + PRODUCER = "producer" + CONSUMER = "consumer" + NOT_INITIALIZED = "not_initialized" + + +def convert_virtual_to_physical_pages( + virtual_pages: list[int], + virtual_block_size: int = 16, + physical_block_size: int = 1, +) -> list[int]: + """Expand virtual (coarse) block IDs into physical (fine-grained) page IDs. + + In paged-attention the scheduler works with *virtual* blocks of + ``virtual_block_size`` tokens, but the RDMA transfer operates at + ``physical_block_size`` granularity. + + Args: + virtual_pages: List of virtual block IDs. + virtual_block_size: Tokens per virtual block. + physical_block_size: Tokens per physical block. + + Returns: + Expanded list of physical page IDs. + """ + block_ratio = virtual_block_size // physical_block_size + physical_pages: list[int] = [] + for vp in virtual_pages: + start = vp * block_ratio + physical_pages.extend(range(start, start + block_ratio)) + return physical_pages + + +class _RoleManager: + """Thread-safe singleton that tracks the P/D role of this process. + + Use the module-level :func:`get_role` / :func:`set_role` helpers + instead of accessing this class directly. + """ + + _instance: Optional[_RoleManager] = None + _lock = threading.Lock() + + def __init__(self) -> None: + self._role: Role = Role.NOT_INITIALIZED + + @classmethod + def get_instance(cls) -> _RoleManager: + """Return the singleton, creating it on first call.""" + if cls._instance is None: + with cls._lock: + if cls._instance is None: + instance = object.__new__(cls) + instance.__init__() + cls._instance = instance + return cls._instance + + def set_role(self, role: Role) -> None: + with self._lock: + self._role = role + + @property + def role(self) -> Role: + return self._role + + +def set_role(role: Role) -> None: + """Set the global P/D role for this process.""" + _RoleManager.get_instance().set_role(role) + + +def get_role() -> Role: + """Get the global P/D role for this process.""" + return _RoleManager.get_instance().role + + +# --------------------------------------------------------------------------- +# Constants +# --------------------------------------------------------------------------- + + +class MoRIIOConstants: + """Protocol constants for the MoRIIO-based KV connector.""" + + # ZMQ handshake message types + GET_META_MSG = b"get_meta_msg" + POP_DONE_RECV = b"pop_done_recv" + OVER = b"OVER" + COMPLETION_PREFIX = "cmpl" + + # Service discovery + PING_INTERVAL_SECONDS = 5 + MAX_PING_RETRIES = 100 + + # Networking + DEFAULT_HANDSHAKE_PORT = 6301 + DEFAULT_NOTIFY_PORT = "61005" + + # Timeouts + ABORT_REQUEST_TIMEOUT = 3600 + + +def get_port_offset(dp_rank: int, tp_rank: int, tp_size: int = 1) -> int: + return (dp_rank) * tp_size + tp_rank diff --git a/atom/kv_transfer/disaggregation/kv_transfer_engine.py b/atom/kv_transfer/disaggregation/moriio/moriio_connector.py similarity index 67% rename from atom/kv_transfer/disaggregation/kv_transfer_engine.py rename to atom/kv_transfer/disaggregation/moriio/moriio_connector.py index 7b83833590..d287a135d8 100644 --- a/atom/kv_transfer/disaggregation/kv_transfer_engine.py +++ b/atom/kv_transfer/disaggregation/moriio/moriio_connector.py @@ -2,26 +2,10 @@ # Copyright (C) 2024-2025, Advanced Micro Devices, Inc. All rights reserved. """ -KV Cache Connector for Disaggregated Prefill-Decode (P/D) Architecture. +Worker-side and scheduler-side KV cache connectors for disaggregated P/D. -This module implements the KV cache transfer mechanism for disaggregated -inference, where prefill and decode stages run on separate GPU instances. -It uses RDMA-based zero-copy transfers via the MoRIIO library for efficient +Uses RDMA-based zero-copy transfers via the MoRIIO library for efficient KV cache migration between producer (prefill) and consumer (decode) nodes. - -Key Components: - - KVConnector: Worker-side connector managing RDMA transfers and handshakes. - - KVConnectorScheduler: Scheduler-side connector coordinating transfer state. - - MoRIIOWrapper: Abstraction over the MoRIIO RDMA engine. - - ConnectorMetadata: Transfer metadata exchanged between scheduler and workers. - -Transfer Modes: - - Read mode: The decode instance reads KV cache directly from prefill memory. - Prefill completes first, then decode reads the blocks via RDMA. - -Architecture:: - - Proxy (service discovery) <-> Prefill Instance <--RDMA--> Decode Instance """ from __future__ import annotations @@ -32,26 +16,29 @@ import time from collections import defaultdict from concurrent.futures import Future, ThreadPoolExecutor -from enum import Enum -from typing import TYPE_CHECKING, Any, Optional +from typing import Any import msgpack import msgspec import numpy as np import zmq -if TYPE_CHECKING: - import torch - from atom.config import Config from atom.kv_transfer.disaggregation.base import ( KVConnectorBase, KVConnectorSchedulerBase, ) +from atom.kv_transfer.disaggregation.moriio.moriio_common import ( + MoRIIOAgentMetadata, + MoRIIOConstants, + _MORIIO_AVAILABLE, + get_port_offset, +) +from atom.kv_transfer.disaggregation.utils import chunk_tensor_for_rdma +from atom.kv_transfer.disaggregation.moriio.moriio_engine import MoRIIOWrapper from atom.kv_transfer.disaggregation.types import ( ConnectorMetadata, EngineId, - RemoteAllocInfo, ReqId, ReqMeta, TransferId, @@ -60,542 +47,32 @@ from atom.utils import ( get_open_port, make_zmq_path, - make_zmq_socket, zmq_socket_ctx, ) from atom.utils.network import get_ip from aiter.dist.parallel_state import get_dp_group, get_tp_group -logger = logging.getLogger("atom") - -# --------------------------------------------------------------------------- -# MoRIIO availability check -# --------------------------------------------------------------------------- - -_MORIIO_AVAILABLE = False -try: +if _MORIIO_AVAILABLE: from mori.io import ( BackendType, - EngineDesc, IOEngine, IOEngineConfig, - MemoryDesc, - MemoryLocationType, PollCqMode, RdmaBackendConfig, ) - _MORIIO_AVAILABLE = True - logger.info("MoRIIO RDMA library loaded successfully") -except ImportError: - logger.warning( - "MoRIIO is not available — KV cache disaggregation will not work. " - "Install the mori package to enable RDMA transfers." - ) - - -# --------------------------------------------------------------------------- -# Msgspec metadata structs -# --------------------------------------------------------------------------- - - -class MoRIIOAgentMetadata( - msgspec.Struct, - omit_defaults=True, - dict=True, - kw_only=True, -): - """Serializable metadata exchanged during the RDMA handshake.""" - - engine_id: str - agent_metadata: bytes - kv_caches_base_addr: Optional[list[int]] = None - num_blocks: int = 0 - block_len: int = 0 - attn_backend_name: str = "aiter" - - -# --------------------------------------------------------------------------- -# Enums & role management -# --------------------------------------------------------------------------- - - -class Role(Enum): - """Role of the current engine instance in the P/D architecture.""" - - PRODUCER = "producer" - CONSUMER = "consumer" - NOT_INITIALIZED = "not_initialized" - - -def convert_virtual_to_physical_pages( - virtual_pages: list[int], - virtual_block_size: int = 16, - physical_block_size: int = 1, -) -> list[int]: - """Expand virtual (coarse) block IDs into physical (fine-grained) page IDs. - - In paged-attention the scheduler works with *virtual* blocks of - ``virtual_block_size`` tokens, but the RDMA transfer operates at - ``physical_block_size`` granularity. - - Args: - virtual_pages: List of virtual block IDs. - virtual_block_size: Tokens per virtual block. - physical_block_size: Tokens per physical block. - - Returns: - Expanded list of physical page IDs. - """ - block_ratio = virtual_block_size // physical_block_size - physical_pages: list[int] = [] - for vp in virtual_pages: - start = vp * block_ratio - physical_pages.extend(range(start, start + block_ratio)) - return physical_pages - - -class _RoleManager: - """Thread-safe singleton that tracks the P/D role of this process. - - Use the module-level :func:`get_role` / :func:`set_role` helpers - instead of accessing this class directly. - """ - - _instance: Optional[_RoleManager] = None - _lock = threading.Lock() - - def __init__(self) -> None: - self._role: Role = Role.NOT_INITIALIZED - - @classmethod - def get_instance(cls) -> _RoleManager: - """Return the singleton, creating it on first call.""" - if cls._instance is None: - with cls._lock: - if cls._instance is None: - instance = object.__new__(cls) - instance.__init__() - cls._instance = instance - return cls._instance - - def set_role(self, role: Role) -> None: - with self._lock: - self._role = role - - @property - def role(self) -> Role: - return self._role - - -def set_role(role: Role) -> None: - """Set the global P/D role for this process.""" - _RoleManager.get_instance().set_role(role) - - -def get_role() -> Role: - """Get the global P/D role for this process.""" - return _RoleManager.get_instance().role - - -# --------------------------------------------------------------------------- -# Constants -# --------------------------------------------------------------------------- - - -class MoRIIOConstants: - """Protocol constants for the MoRIIO-based KV connector.""" - - # ZMQ handshake message types - GET_META_MSG = b"get_meta_msg" - POP_DONE_RECV = b"pop_done_recv" - OVER = b"OVER" - COMPLETION_PREFIX = "cmpl" - - # Service discovery - PING_INTERVAL_SECONDS = 5 - MAX_PING_RETRIES = 100 - - # Networking - DEFAULT_HANDSHAKE_PORT = 6301 - DEFAULT_NOTIFY_PORT = "61005" - - # Timeouts - ABORT_REQUEST_TIMEOUT = 3600 - - -def get_port_offset(dp_rank: int, tp_rank: int, tp_size: int = 1) -> int: - return (dp_rank) * tp_size + tp_rank - - -MAX_RDMA_CHUNK_BYTES = 2 * 1024 * 1024 * 1024 - 64 * 1024 # just under 2 GiB - - -def _chunk_tensor_for_rdma( - tensor: torch.Tensor, block_size_in_dim0: int = 1 -) -> tuple[list[tuple[int, int]], int]: - """Split a tensor into <2 GiB RDMA-registrable chunks along dim 0. - - Args: - tensor: contiguous torch.Tensor whose dim-0 is the block (or - token) axis. - block_size_in_dim0: elements per logical block in dim 0. - Non-MLA: 1 (dim 0 = num_blocks). - MLA: block_size (dim 0 = num_blocks * block_size). - - Returns: - ``(chunks, blocks_per_chunk)`` where *chunks* is a list of - ``(data_ptr, size_bytes)`` pairs and *blocks_per_chunk* is - the number of logical blocks in each full chunk. - """ - elem_sz = tensor.element_size() - per_block_bytes = block_size_in_dim0 * tensor.stride(0) * elem_sz - total_blocks = tensor.shape[0] // block_size_in_dim0 - bpc = max(1, MAX_RDMA_CHUNK_BYTES // per_block_bytes) - chunks: list[tuple[int, int]] = [] - base = tensor.data_ptr() - for start in range(0, total_blocks, bpc): - end = min(start + bpc, total_blocks) - chunks.append((base + start * per_block_bytes, (end - start) * per_block_bytes)) - return chunks, bpc - - -# =================================================================== -# MoRIIOWrapper — thin abstraction over the RDMA engine -# =================================================================== - - -class MoRIIOWrapper: - """Low-level wrapper around a MoRIIO ``IOEngine``. - - Provides helper methods for memory registration, session management, - and asynchronous RDMA read/write operations. Both producer and - consumer code paths share this wrapper. - - Thread-safety: - ``transfer_status``, ``done_req_ids``, and ``done_write_cache_req_ids`` - are guarded by ``self.lock``. The ZMQ socket cache ``self._sockets`` - is *not* thread-safe — callers must ensure ``send_notify`` is invoked - from a single thread. - - Args: - moriio_engine: MoRIIO IOEngine instance. - tp_rank: Tensor parallel rank. - dp_rank: Data parallel rank. - """ - - def __init__( - self, - moriio_engine: Any = None, - tp_rank: int = 0, - dp_rank: int = 0, - ): - self.tp_rank = tp_rank - self.dp_rank = dp_rank - self.moriio_engine = moriio_engine - - self.remote_memory_metadata: Any = None - self.local_memory_registered: bool = False - # Raw-pointer registration produces multiple MemoryDesc per layer - # (e.g. K + V on MHA) to keep each ibv_reg_mr below the AINIC ~2 GiB - # limit. We retain handles here purely so they aren't GC'd; nothing - # in the active session/transfer path indexes into this list. - self.local_memory_descs: list[Any] = [] - self.transfer_status: list[Any] = [] - self.remote_engine_ip: str | None = None - self.notify_port: int | None = None - - self.lock = threading.Lock() - self.done_req_ids: list[str] = [] - self.done_write_cache_req_ids: list[str] = [] - self.notify_thread: threading.Thread | None = None - - # ZMQ socket cache keyed by endpoint path - self._sockets: dict[str, zmq.Socket] = {} - - def set_moriio_engine(self, moriio_engine: Any) -> None: - """Assign the MoRIIO engine (must not be None).""" - if moriio_engine is None: - raise ValueError("Cannot assign a None MoRIIO engine") - self.moriio_engine = moriio_engine - - def set_backend_type(self, backend_type, backend_config=None): - assert self.moriio_engine is not None, "MoRIIO engine must be set first" - if backend_config is not None: - self.moriio_engine.create_backend(backend_type, backend_config) - else: - self.moriio_engine.create_backend(backend_type) - - def get_agent_metadata(self): - assert self.moriio_engine is not None, "MoRIIO engine must be set first" - engine_metadata = self.moriio_engine.get_engine_desc() - engine_metadata_packed = engine_metadata.pack() - return engine_metadata_packed - - def register_remote_engine(self, remote_packed_engine_metadata): - assert self.moriio_engine is not None, "MoRIIO engine must be set first" - consumer_engine_metadata = EngineDesc.unpack(remote_packed_engine_metadata) - self.moriio_engine.register_remote_engine(consumer_engine_metadata) - logger.info( - "Registered remote engine with key: %s", consumer_engine_metadata.key - ) - return consumer_engine_metadata.key - - def register_local_buffer(self, ptr: int, size: int, device_id: int) -> bytes: - """Register one raw GPU memory region with MoRIIO. - - Using ``register_memory(ptr, size, ...)`` directly (instead of - ``register_torch_tensor``) lets callers split a single tensor into - multiple smaller regions, which is required to stay under the - AINIC ~2 GiB ``ibv_reg_mr`` limit. - """ - assert self.moriio_engine is not None, "MoRIIO engine must be set first" - try: - desc = self.moriio_engine.register_memory( - ptr, size, device_id, MemoryLocationType.GPU - ) - assert desc is not None, "register_memory returned None" - packed = desc.pack() - except Exception as e: - raise ValueError(f"Failed to register local memory: {e}") from e - self.local_memory_descs.append(desc) - self.local_memory_registered = True - return packed - - def register_local_tensor(self, tensor): - """Back-compat helper: register an entire contiguous torch tensor. - - Prefer :meth:`register_local_buffer` from new code so that callers - explicitly choose how to chunk large tensors before registration. - """ - if not tensor.is_contiguous(): - raise RuntimeError("input tensor must be contiguous") - ptr = tensor.data_ptr() - size = tensor.numel() * tensor.element_size() - device_id = tensor.device.index if tensor.device.index is not None else -1 - return self.register_local_buffer(ptr, size, device_id) - - def get_unpack_memory_metadata(self, packed_memory_metadata): - return MemoryDesc.unpack(packed_memory_metadata) - - def build_session(self, local_memory_metadata, remote_memory_metadata): - assert self.moriio_engine is not None, "MoRIIO engine must be set first" - tmp = self.moriio_engine.create_session( - local_memory_metadata, remote_memory_metadata - ) - - return tmp - - def read_remote_data( - self, transfer_size_byte, local_offset=0, remote_offset=0, session=None - ): - assert self.local_memory_registered, "You have not register local memory data!" - assert self.moriio_engine is not None, "MoRIIO engine must be set first" - transfer_status = session.batch_read( - local_offset, - remote_offset, - transfer_size_byte, - self.moriio_engine.allocate_transfer_uid(), - ) - return transfer_status - - def write_remote_data( - self, transfer_size_byte, local_offset=0, remote_offset=0, session=None - ): - assert self.local_memory_registered, "You have not register local memory data!" - assert self.moriio_engine is not None, "MoRIIO engine must be set first" - write_uid = self.moriio_engine.allocate_transfer_uid() - - transfer_status = session.batch_write( - local_offset, remote_offset, transfer_size_byte, write_uid - ) - with self.lock: - self.transfer_status.append(transfer_status) - - def write_remote_data_single( - self, transfer_size_byte, local_offset=0, remote_offset=0, sess_idx=0 - ): - assert self.local_memory_registered, "You have not register local memory data!" - assert self.moriio_engine is not None, "MoRIIO engine must be set first" - transfer_status = self.sessions[sess_idx].write( - local_offset, - remote_offset, - transfer_size_byte, - self.moriio_engine.allocate_transfer_uid(), - ) - with self.lock: - self.transfer_status.append(transfer_status) - - def waiting_for_transfer_complete(self): - if not self.transfer_status: - return - - transfers_to_wait = [] - with self.lock: - transfers_to_wait = self.transfer_status[:] - self.transfer_status.clear() - - for status in transfers_to_wait: - try: - status.Wait() - if not status.Succeeded(): - logger.error( - "Transfer failed: %s, Code: %s", status.Message(), status.Code() - ) - raise ValueError("MoRIIO transfer failed!") - except Exception as e: - logger.error("Transfer %s failed: %s", status, e) - raise - - def async_wait_reqid(self): - assert self.notify_port is not None, "Notify port cannot be None" - - if self.notify_thread is not None: - return - - def _async_wait(): - host = "*" - path = make_zmq_path("tcp", host, self.notify_port) - logger.info("Node starting to listen notify from path = %s", path) - - with _zmq_ctx(zmq.ROUTER, path) as sock: - while True: - try: - identity, msg = sock.recv_multipart() - self._dispatch_message(msg) - except Exception as e: - logger.error("Error processing message: %s", e) - raise ValueError(f"Error processing message: {e}") from e - - self.notify_thread = threading.Thread( - target=_async_wait, daemon=True, name="moriio-notify-listener" - ) - self.notify_thread.start() - - def _dispatch_message(self, msg: bytes) -> None: - """Route an incoming ZMQ message to the appropriate handler. - - Message formats: - - msgpack dict with ``req_id``: remote block allocation (producer only) - - UTF-8 string prefixed with ``cmpl``: transfer completion signal - """ - # Try msgpack structured message first (block allocation from decode) - try: - data = msgpack.loads(msg) - if isinstance(data, dict) and "req_id" in data: - self._handle_block_alloc_message(data) - return - except (msgpack.exceptions.ExtraData, msgpack.exceptions.UnpackException): - pass - - # Fall back to string-encoded completion message - try: - msg_str = msg.decode("utf-8") - except UnicodeDecodeError: - logger.warning("Received non-decodable message of %d bytes", len(msg)) - return - - if msg_str.startswith(MoRIIOConstants.COMPLETION_PREFIX): - self._handle_completion_message(msg_str) - else: - raise ValueError(f"Unrecognized message format: {msg_str!r}") - - def _handle_block_alloc_message(self, data: dict) -> None: - """Process a remote block allocation notification (producer side).""" - assert ( - get_role() == Role.PRODUCER - ), "Only producer should receive block alloc messages" - req_id = data["req_id"] - block_notify_list = data.get("block_notify_list", []) - decode_dp_rank = data.get("decode_rank", 0) - assert ( - len(block_notify_list) > 0 - ), "block_notify_list cannot be empty in remote allocate message" - - with self.lock: - self.done_remote_allocate_req_dict[req_id] = RemoteAllocInfo( - block_ids=block_notify_list, decode_dp_rank=decode_dp_rank - ) - - def _handle_completion_message(self, msg: str) -> None: - """Record a transfer completion notification.""" - with self.lock: - if get_role() == Role.PRODUCER: - self.done_req_ids.append(msg) - else: - self.done_write_cache_req_ids.append(msg) - - def send_notify( - self, - req_ids: str | int | list[str | int], - remote_ip: str, - remote_port: str | int, - ) -> None: - """Notify a remote engine that transfer(s) have completed.""" - if not remote_ip or not remote_port: - logger.warning("Cannot send notification: missing remote_ip or remote_port") - return - - path = make_zmq_path("tcp", remote_ip, int(remote_port)) - - if path not in self._sockets: - ctx = zmq.Context.instance() - self._sockets[path] = make_zmq_socket( - ctx=ctx, path=path, socket_type=zmq.DEALER, bind=False - ) - - id_list = req_ids if isinstance(req_ids, list) else [req_ids] - sock = self._sockets[path] - try: - for rid in id_list: - rid_str = str(rid) if isinstance(rid, int) else rid - if not isinstance(rid_str, str): - logger.warning("Skipping non-string req_id of type %s", type(rid)) - continue - sock.send_multipart( - [MoRIIOConstants.POP_DONE_RECV, rid_str.encode("utf-8")] - ) - except Exception as e: - logger.error("Failed to send notification to %s: %s", path, e) - self._sockets.pop(path, None) - raise - - def pop_finished_req_ids(self) -> set[str]: - """Return and clear the set of completed send-side request IDs.""" - with self.lock: - result = set(self.done_req_ids) - self.done_req_ids.clear() - return result - - def pop_finished_write_req_ids(self) -> set[str]: - """Return and clear the set of completed write-side request IDs.""" - with self.lock: - result = set(self.done_write_cache_req_ids) - self.done_write_cache_req_ids.clear() - return result - - def shutdown(self) -> None: - """Close all cached ZMQ sockets and release resources.""" - logger.debug( - "Shutting down MoRIIOWrapper, closing %d sockets", len(self._sockets) - ) - for path, sock in self._sockets.items(): - try: - sock.close(linger=0) - except Exception as e: - logger.warning("Error closing socket for %s: %s", path, e) - self._sockets.clear() +logger = logging.getLogger("atom") # =================================================================== -# KVConnector — worker-side connector (runs inside each TP rank) +# MoRIIOConnector — worker-side connector (runs inside each TP rank) # =================================================================== -class KVConnector(KVConnectorBase): +class MoRIIOConnector(KVConnectorBase): """Worker-side KV cache connector for disaggregated P/D inference. - Each tensor-parallel worker instantiates one ``KVConnector``. It is + Each tensor-parallel worker instantiates one ``MoRIIOConnector``. It is responsible for: 1. Registering local KV cache tensors for RDMA access. @@ -727,7 +204,7 @@ def register_kv_caches(self, kv_caches: dict[str, Any]) -> None: transfers can occur. Each K (and V, when present) tensor is split into block-aligned - chunks of < 2 GiB via :func:`_chunk_tensor_for_rdma` and each + chunks of < 2 GiB via :func:`chunk_tensor_for_rdma` and each chunk is registered independently with ``ibv_reg_mr``. Per-layer metadata list layout: ``[k_chunk0, k_chunk1, ..., v_chunk0, v_chunk1, ...]``. @@ -754,7 +231,7 @@ def register_kv_caches(self, kv_caches: dict[str, Any]) -> None: # equals the KV cache block_size (typically 16). # Non-MLA: dim 0 = num_blocks directly, so 1. bsd0 = self.kv_cache_block_size if is_mla else 1 - k_chunks, bpc = _chunk_tensor_for_rdma(cache_tensor, bsd0) + k_chunks, bpc = chunk_tensor_for_rdma(cache_tensor, bsd0) if self.blocks_per_chunk is None: self.blocks_per_chunk = bpc @@ -770,7 +247,7 @@ def register_kv_caches(self, kv_caches: dict[str, Any]) -> None: v_device_id = ( v_cache.device.index if v_cache.device.index is not None else -1 ) - v_chunks, _ = _chunk_tensor_for_rdma(v_cache, 1) + v_chunks, _ = chunk_tensor_for_rdma(v_cache, 1) for ptr, size in v_chunks: meta_list.append( self.moriio_wrapper.register_local_buffer( @@ -1454,11 +931,11 @@ def get_finished(self) -> tuple[set[int], set[str]]: # =================================================================== -# KVConnectorScheduler — scheduler-side connector +# MoRIIOConnectorScheduler — scheduler-side connector # =================================================================== -class KVConnectorScheduler(KVConnectorSchedulerBase): +class MoRIIOConnectorScheduler(KVConnectorSchedulerBase): """Scheduler-side KV connector that tracks transfer lifecycle. Runs in the scheduler process (not in TP workers). Responsible for: @@ -1564,6 +1041,7 @@ def request_finished(self, seq: Sequence) -> None: the transfer_id mapping. """ # Attach output metadata for the proxy to relay + first_token_id = seq.output_tokens[0] if seq.output_tokens else None seq.kv_transfer_params_output = { "do_remote_prefill": True, "do_remote_decode": False, @@ -1574,6 +1052,7 @@ def request_finished(self, seq: Sequence) -> None: "tp_size": self.tp_size, "dp_rank": self.dp_rank, "transfer_id": seq.id, + "first_token_id": first_token_id, } # Clean up transfer ID mapping on the consumer side diff --git a/atom/kv_transfer/disaggregation/moriio/moriio_engine.py b/atom/kv_transfer/disaggregation/moriio/moriio_engine.py new file mode 100644 index 0000000000..beda8dc46a --- /dev/null +++ b/atom/kv_transfer/disaggregation/moriio/moriio_engine.py @@ -0,0 +1,360 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2025, Advanced Micro Devices, Inc. All rights reserved. + +""" +Low-level RDMA wrapper around a MoRIIO ``IOEngine``. + +Provides helper methods for memory registration, session management, +and asynchronous RDMA read/write operations. Both producer and +consumer code paths share this wrapper. +""" + +from __future__ import annotations + +import logging +import threading +from typing import Any + +import msgpack +import zmq + +from atom.kv_transfer.disaggregation.moriio.moriio_common import ( + MoRIIOConstants, + Role, + _MORIIO_AVAILABLE, + get_role, +) +from atom.kv_transfer.disaggregation.types import RemoteAllocInfo +from atom.utils import make_zmq_path, make_zmq_socket + +if _MORIIO_AVAILABLE: + from mori.io import EngineDesc, MemoryDesc, MemoryLocationType + +logger = logging.getLogger("atom") + + +class MoRIIOWrapper: + """Low-level wrapper around a MoRIIO ``IOEngine``. + + Provides helper methods for memory registration, session management, + and asynchronous RDMA read/write operations. Both producer and + consumer code paths share this wrapper. + + Thread-safety: + ``transfer_status``, ``done_req_ids``, ``done_write_cache_req_ids``, + and ``done_remote_allocate_req_dict`` are guarded by ``self.lock``. The ZMQ socket cache ``self._sockets`` + is *not* thread-safe — callers must ensure ``send_notify`` is invoked + from a single thread. + + Args: + moriio_engine: MoRIIO IOEngine instance. + tp_rank: Tensor parallel rank. + dp_rank: Data parallel rank. + """ + + def __init__( + self, + moriio_engine: Any = None, + tp_rank: int = 0, + dp_rank: int = 0, + ): + self.tp_rank = tp_rank + self.dp_rank = dp_rank + self.moriio_engine = moriio_engine + + self.remote_memory_metadata: Any = None + self.local_memory_registered: bool = False + # Raw-pointer registration produces multiple MemoryDesc per layer + # (e.g. K + V on MHA) to keep each ibv_reg_mr below the AINIC ~2 GiB + # limit. We retain handles here purely so they aren't GC'd; nothing + # in the active session/transfer path indexes into this list. + self.local_memory_descs: list[Any] = [] + self.transfer_status: list[Any] = [] + self.remote_engine_ip: str | None = None + self.notify_port: int | None = None + + self.lock = threading.Lock() + self.done_req_ids: list[str] = [] + self.done_write_cache_req_ids: list[str] = [] + self.done_remote_allocate_req_dict: dict[str, Any] = {} + self.notify_thread: threading.Thread | None = None + + # ZMQ socket cache keyed by endpoint path + self._sockets: dict[str, zmq.Socket] = {} + + def set_moriio_engine(self, moriio_engine: Any) -> None: + """Assign the MoRIIO engine (must not be None).""" + if moriio_engine is None: + raise ValueError("Cannot assign a None MoRIIO engine") + self.moriio_engine = moriio_engine + + def set_backend_type(self, backend_type, backend_config=None): + assert self.moriio_engine is not None, "MoRIIO engine must be set first" + if backend_config is not None: + self.moriio_engine.create_backend(backend_type, backend_config) + else: + self.moriio_engine.create_backend(backend_type) + + def get_agent_metadata(self): + assert self.moriio_engine is not None, "MoRIIO engine must be set first" + engine_metadata = self.moriio_engine.get_engine_desc() + engine_metadata_packed = engine_metadata.pack() + return engine_metadata_packed + + def register_remote_engine(self, remote_packed_engine_metadata): + assert self.moriio_engine is not None, "MoRIIO engine must be set first" + consumer_engine_metadata = EngineDesc.unpack(remote_packed_engine_metadata) + self.moriio_engine.register_remote_engine(consumer_engine_metadata) + logger.info( + "Registered remote engine with key: %s", consumer_engine_metadata.key + ) + return consumer_engine_metadata.key + + def register_local_buffer(self, ptr: int, size: int, device_id: int) -> bytes: + """Register one raw GPU memory region with MoRIIO. + + Using ``register_memory(ptr, size, ...)`` directly (instead of + ``register_torch_tensor``) lets callers split a single tensor into + multiple smaller regions, which is required to stay under the + AINIC ~2 GiB ``ibv_reg_mr`` limit. + """ + assert self.moriio_engine is not None, "MoRIIO engine must be set first" + try: + desc = self.moriio_engine.register_memory( + ptr, size, device_id, MemoryLocationType.GPU + ) + assert desc is not None, "register_memory returned None" + packed = desc.pack() + except Exception as e: + raise ValueError(f"Failed to register local memory: {e}") from e + self.local_memory_descs.append(desc) + self.local_memory_registered = True + return packed + + def register_local_tensor(self, tensor): + """Back-compat helper: register an entire contiguous torch tensor. + + Prefer :meth:`register_local_buffer` from new code so that callers + explicitly choose how to chunk large tensors before registration. + """ + if not tensor.is_contiguous(): + raise RuntimeError("input tensor must be contiguous") + ptr = tensor.data_ptr() + size = tensor.numel() * tensor.element_size() + device_id = tensor.device.index if tensor.device.index is not None else -1 + return self.register_local_buffer(ptr, size, device_id) + + def get_unpack_memory_metadata(self, packed_memory_metadata): + return MemoryDesc.unpack(packed_memory_metadata) + + def build_session(self, local_memory_metadata, remote_memory_metadata): + assert self.moriio_engine is not None, "MoRIIO engine must be set first" + tmp = self.moriio_engine.create_session( + local_memory_metadata, remote_memory_metadata + ) + + return tmp + + def read_remote_data( + self, transfer_size_byte, local_offset=0, remote_offset=0, session=None + ): + assert self.local_memory_registered, "You have not register local memory data!" + assert self.moriio_engine is not None, "MoRIIO engine must be set first" + transfer_status = session.batch_read( + local_offset, + remote_offset, + transfer_size_byte, + self.moriio_engine.allocate_transfer_uid(), + ) + return transfer_status + + def write_remote_data( + self, transfer_size_byte, local_offset=0, remote_offset=0, session=None + ): + assert self.local_memory_registered, "You have not register local memory data!" + assert self.moriio_engine is not None, "MoRIIO engine must be set first" + write_uid = self.moriio_engine.allocate_transfer_uid() + + transfer_status = session.batch_write( + local_offset, remote_offset, transfer_size_byte, write_uid + ) + with self.lock: + self.transfer_status.append(transfer_status) + + def write_remote_data_single( + self, transfer_size_byte, local_offset=0, remote_offset=0, sess_idx=0 + ): + assert self.local_memory_registered, "You have not register local memory data!" + assert self.moriio_engine is not None, "MoRIIO engine must be set first" + transfer_status = self.sessions[sess_idx].write( + local_offset, + remote_offset, + transfer_size_byte, + self.moriio_engine.allocate_transfer_uid(), + ) + with self.lock: + self.transfer_status.append(transfer_status) + + def waiting_for_transfer_complete(self): + if not self.transfer_status: + return + + transfers_to_wait = [] + with self.lock: + transfers_to_wait = self.transfer_status[:] + self.transfer_status.clear() + + for status in transfers_to_wait: + try: + status.Wait() + if not status.Succeeded(): + logger.error( + "Transfer failed: %s, Code: %s", status.Message(), status.Code() + ) + raise ValueError("MoRIIO transfer failed!") + except Exception as e: + logger.error("Transfer %s failed: %s", status, e) + raise + + def async_wait_reqid(self): + assert self.notify_port is not None, "Notify port cannot be None" + + if self.notify_thread is not None: + return + + def _async_wait(): + from atom.kv_transfer.disaggregation.moriio.moriio_connector import ( + _zmq_ctx, + ) + + host = "*" + path = make_zmq_path("tcp", host, self.notify_port) + logger.info("Node starting to listen notify from path = %s", path) + + with _zmq_ctx(zmq.ROUTER, path) as sock: + while True: + try: + identity, msg = sock.recv_multipart() + self._dispatch_message(msg) + except Exception as e: + logger.error("Error processing message: %s", e) + raise ValueError(f"Error processing message: {e}") from e + + self.notify_thread = threading.Thread( + target=_async_wait, daemon=True, name="moriio-notify-listener" + ) + self.notify_thread.start() + + def _dispatch_message(self, msg: bytes) -> None: + """Route an incoming ZMQ message to the appropriate handler. + + Message formats: + - msgpack dict with ``req_id``: remote block allocation (producer only) + - UTF-8 string prefixed with ``cmpl``: transfer completion signal + """ + # Try msgpack structured message first (block allocation from decode) + try: + data = msgpack.loads(msg) + if isinstance(data, dict) and "req_id" in data: + self._handle_block_alloc_message(data) + return + except (msgpack.exceptions.ExtraData, msgpack.exceptions.UnpackException): + pass + + # Fall back to string-encoded completion message + try: + msg_str = msg.decode("utf-8") + except UnicodeDecodeError: + logger.warning("Received non-decodable message of %d bytes", len(msg)) + return + + if msg_str.startswith(MoRIIOConstants.COMPLETION_PREFIX): + self._handle_completion_message(msg_str) + else: + raise ValueError(f"Unrecognized message format: {msg_str!r}") + + def _handle_block_alloc_message(self, data: dict) -> None: + """Process a remote block allocation notification (producer side).""" + assert ( + get_role() == Role.PRODUCER + ), "Only producer should receive block alloc messages" + req_id = data["req_id"] + block_notify_list = data.get("block_notify_list", []) + decode_dp_rank = data.get("decode_rank", 0) + assert ( + len(block_notify_list) > 0 + ), "block_notify_list cannot be empty in remote allocate message" + + with self.lock: + self.done_remote_allocate_req_dict[req_id] = RemoteAllocInfo( + block_ids=block_notify_list, decode_dp_rank=decode_dp_rank + ) + + def _handle_completion_message(self, msg: str) -> None: + """Record a transfer completion notification.""" + with self.lock: + if get_role() == Role.PRODUCER: + self.done_req_ids.append(msg) + else: + self.done_write_cache_req_ids.append(msg) + + def send_notify( + self, + req_ids: str | int | list[str | int], + remote_ip: str, + remote_port: str | int, + ) -> None: + """Notify a remote engine that transfer(s) have completed.""" + if not remote_ip or not remote_port: + logger.warning("Cannot send notification: missing remote_ip or remote_port") + return + + path = make_zmq_path("tcp", remote_ip, int(remote_port)) + + if path not in self._sockets: + ctx = zmq.Context.instance() + self._sockets[path] = make_zmq_socket( + ctx=ctx, path=path, socket_type=zmq.DEALER, bind=False + ) + + id_list = req_ids if isinstance(req_ids, list) else [req_ids] + sock = self._sockets[path] + try: + for rid in id_list: + rid_str = str(rid) if isinstance(rid, int) else rid + if not isinstance(rid_str, str): + logger.warning("Skipping non-string req_id of type %s", type(rid)) + continue + sock.send_multipart( + [MoRIIOConstants.POP_DONE_RECV, rid_str.encode("utf-8")] + ) + except Exception as e: + logger.error("Failed to send notification to %s: %s", path, e) + self._sockets.pop(path, None) + raise + + def pop_finished_req_ids(self) -> set[str]: + """Return and clear the set of completed send-side request IDs.""" + with self.lock: + result = set(self.done_req_ids) + self.done_req_ids.clear() + return result + + def pop_finished_write_req_ids(self) -> set[str]: + """Return and clear the set of completed write-side request IDs.""" + with self.lock: + result = set(self.done_write_cache_req_ids) + self.done_write_cache_req_ids.clear() + return result + + def shutdown(self) -> None: + """Close all cached ZMQ sockets and release resources.""" + logger.debug( + "Shutting down MoRIIOWrapper, closing %d sockets", len(self._sockets) + ) + for path, sock in self._sockets.items(): + try: + sock.close(linger=0) + except Exception as e: + logger.warning("Error closing socket for %s: %s", path, e) + self._sockets.clear() diff --git a/atom/kv_transfer/disaggregation/proxy.py b/atom/kv_transfer/disaggregation/proxy.py index 81c062074f..f3c6d192d1 100644 --- a/atom/kv_transfer/disaggregation/proxy.py +++ b/atom/kv_transfer/disaggregation/proxy.py @@ -213,13 +213,42 @@ async def stream_decode_response( request_id: str, ): """Yield response chunks from a decode instance, then close the session.""" + chunk_count = 0 try: if response.status != 200: raise RuntimeError( f"Decode request {request_id} failed with status {response.status}" ) - async for chunk_bytes in response.content.iter_chunked(1024): + while True: + try: + chunk_bytes = await asyncio.wait_for( + response.content.readany(), timeout=120.0 + ) + except asyncio.TimeoutError: + logger.warning( + "[PROXY] stream_decode TIMEOUT for %s after %d chunks " + "(120s no data), closing", + request_id, + chunk_count, + ) + break + if not chunk_bytes: + break + chunk_count += 1 yield chunk_bytes + logger.info( + "[PROXY] stream_decode finished for %s, %d chunks forwarded", + request_id, + chunk_count, + ) + except Exception as e: + logger.exception( + "[PROXY] stream_decode ERROR for %s after %d chunks: %s", + request_id, + chunk_count, + e, + ) + raise finally: await session.close() @@ -254,6 +283,7 @@ async def handle_request(): req_data = await request.get_json() request_id = str(uuid.uuid4()) + logger.info("[PROXY] handle_request #%d, id=%s", request_nums, request_id) if not prefill_instances or not decode_instances: return await make_response( @@ -300,7 +330,10 @@ async def handle_request(): ) # --- Decode request --- - req_data["max_tokens"] -= 1 + # NOTE: Do NOT decrement max_tokens. The decode engine's first token + # (T0 re-prediction) is overridden with prefill's T0 via first_token_id, + # but that decode step still counts toward max_tokens. Keeping the + # original max_tokens ensures the total output length matches non-PD. req_data["kv_transfer_params"] = { "do_remote_decode": False, "do_remote_prefill": True, @@ -313,26 +346,30 @@ async def handle_request(): logger.info("Transfer type: %s", TRANSFER_TYPE) + prefill_response = await send_prefill_task + logger.info("Prefill response received for request %s", request_id) + prefill_kv = prefill_response["kv_transfer_params"] + req_data["kv_transfer_params"]["transfer_id"] = prefill_kv["transfer_id"] + if "first_token_id" in prefill_kv: + req_data["kv_transfer_params"]["first_token_id"] = prefill_kv[ + "first_token_id" + ] + + actual_dp_rank = prefill_kv.get("dp_rank", selected_prefill_dp_rank) + if actual_dp_rank is not None: + selected_prefill_dp_rank = actual_dp_rank + if TRANSFER_TYPE == "read": - prefill_response = await send_prefill_task - logger.info("Prefill response received for request %s", request_id) - prefill_kv = prefill_response["kv_transfer_params"] req_data["kv_transfer_params"]["remote_engine_id"] = prefill_kv[ "remote_engine_id" ] req_data["kv_transfer_params"]["remote_block_ids"] = prefill_kv[ "remote_block_ids" ] - req_data["kv_transfer_params"]["transfer_id"] = prefill_kv["transfer_id"] - - # Use the actual dp_rank from the prefill response (authoritative), - # falling back to the proxy's round-robin selection. - actual_dp_rank = prefill_kv.get("dp_rank", selected_prefill_dp_rank) - if actual_dp_rank is not None: - selected_prefill_dp_rank = actual_dp_rank req_data["kv_transfer_params"]["remote_dp_size"] = prefill_ep["dp_size"] req_data["kv_transfer_params"]["remote_tp_size"] = prefill_ep["tp_size"] + req_data["kv_transfer_params"]["tp_size"] = prefill_ep["tp_size"] if selected_prefill_dp_rank is not None: req_data["kv_transfer_params"]["remote_dp_rank"] = selected_prefill_dp_rank @@ -341,12 +378,19 @@ async def handle_request(): decode_ep["request_address"], req_data, request_id ) stream_gen = stream_decode_response(session, decode_response, request_id) + logger.info( + "[PROXY] decode response status=%d for %s", + decode_response.status, + request_id, + ) response = await make_response(stream_gen) response.headers["Content-Type"] = "application/json; charset=utf-8" return response except Exception as e: - logger.exception("Error handling request: %s", e) + logger.exception( + "[PROXY] Error handling request #%d id=%s: %s", request_nums, request_id, e + ) resp = await make_response((f"Internal Server Error: {e!s}", 500)) resp.headers["Content-Type"] = "application/json; charset=utf-8" return resp @@ -355,11 +399,38 @@ async def handle_request(): _DEFAULT_DISCOVERY_PORT = 36367 +@app.route("/start_profile", methods=["POST"]) +@app.route("/stop_profile", methods=["POST"]) +async def proxy_profile(): + """Forward profiler start/stop to all registered prefill and decode instances.""" + path = request.path # "/start_profile" or "/stop_profile" + results = {} + async with aiohttp.ClientSession() as session: + for tag, instances in [ + ("prefill", prefill_instances), + ("decode", decode_instances), + ]: + for i, ep in enumerate(instances): + ip, port = _extract_ip_port(ep["request_address"]) + url = f"http://{ip}:{port}{path}" + try: + async with session.post( + url, timeout=aiohttp.ClientTimeout(total=600) + ) as resp: + results[f"{tag}_{i}"] = { + "status": resp.status, + "body": await resp.json(), + } + except Exception as e: + results[f"{tag}_{i}"] = {"status": "error", "detail": str(e)} + return results + + def main(port: int = 10001): """Launch the P/D proxy service.""" discovery_thread = start_service_discovery("0.0.0.0", _DEFAULT_DISCOVERY_PORT) - app.debug = True + app.debug = False app.config["BODY_TIMEOUT"] = 360000 app.config["RESPONSE_TIMEOUT"] = 360000 diff --git a/atom/kv_transfer/disaggregation/types.py b/atom/kv_transfer/disaggregation/types.py index fdcc6ab1ab..ca24b0f111 100644 --- a/atom/kv_transfer/disaggregation/types.py +++ b/atom/kv_transfer/disaggregation/types.py @@ -129,7 +129,11 @@ def _build_req_meta( remote_handshake_port=kv_transfer_params["remote_handshake_port"], remote_dp_size=kv_transfer_params.get("remote_dp_size", 1), remote_dp_rank=kv_transfer_params.get("remote_dp_rank", 0), - tp_size=kv_transfer_params.get("tp_size", 1), + tp_size=( + kv_transfer_params.get("tp_size") + if "tp_size" in kv_transfer_params + else kv_transfer_params.get("remote_tp_size", 1) + ), transfer_id=kv_transfer_params.get("transfer_id", 0), ) diff --git a/atom/kv_transfer/disaggregation/utils.py b/atom/kv_transfer/disaggregation/utils.py new file mode 100644 index 0000000000..ea705d1ccc --- /dev/null +++ b/atom/kv_transfer/disaggregation/utils.py @@ -0,0 +1,116 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2025, Advanced Micro Devices, Inc. All rights reserved. + +""" +Shared utilities for KV cache disaggregation backends (MoRIIO, Mooncake, etc.). +""" + +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING + +import triton +import triton.language as tl + +if TYPE_CHECKING: + import torch + +logger = logging.getLogger("atom") + +MAX_RDMA_CHUNK_BYTES = 2 * 1024 * 1024 * 1024 - 64 * 1024 # just under 2 GiB + + +# --------------------------------------------------------------------------- +# GPU memory fence for RDMA-written KV cache blocks +# +# Not needed when producer and consumer are in the same network partition +# (same-partition RDMA completes before decode starts). +# Enable for cross-partition deployments where higher RDMA latency +# can cause stale GPU reads (improves accuracy from ~0.85 to ~0.90). +# --------------------------------------------------------------------------- + + +@triton.jit +def _gpu_fence_kernel( + kv_ptr, + block_ids_ptr, + total_blocks: tl.int64, + elems_per_block: tl.int64, + NUM_CL: tl.constexpr, +): + blk_idx = tl.program_id(0) + layer_idx = tl.program_id(1).to(tl.int64) + blk = tl.load(block_ids_ptr + blk_idx).to(tl.int64) + base = (layer_idx * total_blocks + blk) * elems_per_block + cl_offsets = tl.arange(0, NUM_CL).to(tl.int64) * 32 + mask = cl_offsets < elems_per_block + tl.atomic_or( + kv_ptr + base + cl_offsets, tl.zeros([NUM_CL], dtype=tl.int32), mask=mask + ) + + +def gpu_memory_fence( + kv_cache: "torch.Tensor", + block_ids: list[int], + use_mla: bool, +) -> None: + """Force GPU memory coherence for RDMA-written KV cache blocks. + + Single Triton kernel launch -- one atomic_or(0) per cache line (128 bytes) + across all (layer, block) pairs. + + Args: + kv_cache: The full KV cache tensor. + block_ids: Physical block IDs that received RDMA writes. + use_mla: Whether the model uses Multi-head Latent Attention. + """ + import torch + + block_ids_t = torch.tensor(block_ids, dtype=torch.int32, device=kv_cache.device) + kv_flat = kv_cache.view(torch.int32) + if use_mla: + num_layers = kv_cache.shape[0] + total_blocks = kv_cache.shape[1] + else: + num_layers = 2 * kv_cache.shape[1] + total_blocks = kv_cache.shape[2] + elems_per_block = kv_flat.numel() // (num_layers * total_blocks) + num_cl = triton.cdiv(elems_per_block, 32) + NUM_CL = triton.next_power_of_2(num_cl) + _gpu_fence_kernel[(len(block_ids), num_layers)]( + kv_flat, + block_ids_t, + total_blocks, + elems_per_block, + NUM_CL=NUM_CL, + ) + + +def chunk_tensor_for_rdma( + tensor: torch.Tensor, block_size_in_dim0: int = 1 +) -> tuple[list[tuple[int, int]], int]: + """Split a tensor into <2 GiB RDMA-registrable chunks along dim 0. + + Args: + tensor: contiguous torch.Tensor whose dim-0 is the block (or + token) axis. + block_size_in_dim0: elements per logical block in dim 0. + Non-MLA: 1 (dim 0 = num_blocks). + MLA: block_size (dim 0 = num_blocks * block_size). + + Returns: + ``(chunks, blocks_per_chunk)`` where *chunks* is a list of + ``(data_ptr, size_bytes)`` pairs and *blocks_per_chunk* is + the number of logical blocks in each full chunk. + """ + elem_sz = tensor.element_size() + per_block_bytes = block_size_in_dim0 * tensor.stride(0) * elem_sz + total_blocks = tensor.shape[0] // block_size_in_dim0 + bpc = max(1, MAX_RDMA_CHUNK_BYTES // per_block_bytes) + chunks: list[tuple[int, int]] = [] + base = tensor.data_ptr() + for start in range(0, total_blocks, bpc): + end = min(start + bpc, total_blocks) + chunks.append((base + start * per_block_bytes, (end - start) * per_block_bytes)) + return chunks, bpc diff --git a/atom/model_config/kimi_k25.py b/atom/model_config/kimi_k25.py index f9b32dc573..fb85d257b1 100644 --- a/atom/model_config/kimi_k25.py +++ b/atom/model_config/kimi_k25.py @@ -11,6 +11,13 @@ from transformers.configuration_utils import PretrainedConfig +def _first_non_none(*values): + for value in values: + if value is not None: + return value + return None + + class KimiK25VisionConfig(PretrainedConfig): model_type = "kimi_k25_vision" @@ -22,10 +29,10 @@ def __init__( init_pos_emb_width: int = 64, init_pos_emb_time: int = 4, pos_emb_type: str = "divided_fixed", - num_attention_heads: int = 16, - num_hidden_layers: int = 27, - hidden_size: int = 1152, - intermediate_size: int = 4304, + num_attention_heads: int | None = None, + num_hidden_layers: int | None = None, + hidden_size: int | None = None, + intermediate_size: int | None = None, merge_kernel_size: tuple[int, int] = (2, 2), video_attn_type: str = "spatial_temporal", merge_type: str = "sd2_tpool", @@ -34,9 +41,23 @@ def __init__( mm_hidden_size: int | None = None, projector_hidden_act: str = "gelu", projector_ln_eps: float = 1e-5, + # moonshotai/Kimi-K2.5 remote-code field names + vt_num_attention_heads: int | None = None, + vt_num_hidden_layers: int | None = None, + vt_hidden_size: int | None = None, + vt_intermediate_size: int | None = None, + text_hidden_size: int | None = None, **kwargs, ): super().__init__(**kwargs) + num_attention_heads = _first_non_none( + num_attention_heads, vt_num_attention_heads, 16 + ) + num_hidden_layers = _first_non_none(num_hidden_layers, vt_num_hidden_layers, 27) + hidden_size = _first_non_none(hidden_size, vt_hidden_size, 1152) + intermediate_size = _first_non_none( + intermediate_size, vt_intermediate_size, 4304 + ) # Vision Tower self.patch_size = patch_size self.init_pos_emb_height = init_pos_emb_height @@ -47,6 +68,11 @@ def __init__( self.num_hidden_layers = num_hidden_layers self.hidden_size = hidden_size self.intermediate_size = intermediate_size + # Preserve the repo-specific aliases so either schema can be consumed. + self.vt_num_attention_heads = num_attention_heads + self.vt_num_hidden_layers = num_hidden_layers + self.vt_hidden_size = hidden_size + self.vt_intermediate_size = intermediate_size self.merge_kernel_size = merge_kernel_size self.video_attn_type = video_attn_type self.merge_type = merge_type @@ -58,6 +84,25 @@ def __init__( self.mm_hidden_size = hidden_size self.projector_hidden_act = projector_hidden_act self.projector_ln_eps = projector_ln_eps + if text_hidden_size is not None: + self.text_hidden_size = text_hidden_size + + @classmethod + def from_remote_config( + cls, + vision_config: dict | PretrainedConfig | None, + ) -> "KimiK25VisionConfig": + if vision_config is None: + return cls() + if isinstance(vision_config, cls): + return vision_config + if isinstance(vision_config, dict): + config_dict = dict(vision_config) + elif hasattr(vision_config, "to_dict"): + config_dict = dict(vision_config.to_dict()) + else: + config_dict = dict(vars(vision_config)) + return cls(**config_dict) class KimiK25Config(PretrainedConfig): @@ -89,11 +134,7 @@ def __init__( **kwargs, ): # Vision config - if vision_config is None: - vision_config = KimiK25VisionConfig() - elif isinstance(vision_config, dict): - vision_config = KimiK25VisionConfig(**vision_config) - self.vision_config: KimiK25VisionConfig = vision_config + self.vision_config = KimiK25VisionConfig.from_remote_config(vision_config) # Text config if text_config is None: @@ -105,6 +146,8 @@ def __init__( # Set mm_hidden_size to text hidden size if not explicitly set if self.vision_config.mm_hidden_size == self.vision_config.hidden_size: self.vision_config.mm_hidden_size = self.text_config.hidden_size + if getattr(self.vision_config, "text_hidden_size", None) is None: + self.vision_config.text_hidden_size = self.text_config.hidden_size # Other config self.ignore_index = ignore_index diff --git a/atom/model_engine/arg_utils.py b/atom/model_engine/arg_utils.py index 1358c04052..7dae8df47c 100644 --- a/atom/model_engine/arg_utils.py +++ b/atom/model_engine/arg_utils.py @@ -31,7 +31,7 @@ class EngineArgs: tensor_parallel_size: int = 1 data_parallel_size: int = 1 enforce_eager: bool = False - enable_prefix_caching: bool = False + enable_prefix_caching: bool = True port: int = 8006 kv_cache_dtype: str = "bf16" block_size: int = 16 @@ -51,6 +51,7 @@ class EngineArgs: method: Optional[str] = None num_speculative_tokens: int = 1 kv_transfer_config: str = "{}" + draft_model: Optional[str] = None mark_trace: bool = False @staticmethod @@ -86,8 +87,10 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: ) parser.add_argument( "--enable_prefix_caching", - action="store_true", - help="Enable prefix caching.", + action=argparse.BooleanOptionalAction, + default=True, + help="Enable prefix caching (default: enabled). " + "Use --no-enable_prefix_caching to disable.", ) parser.add_argument( "--port", @@ -114,7 +117,7 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: parser.add_argument( "--cudagraph-capture-sizes", type=str, - default="[1,2,4,8,16,32,48,64,128,256]", + default="[1,2,4,8,16,32,48,64,128,256,512]", help="Sizes to capture cudagraph. Example: [1,2,4,8,16]", ) parser.add_argument( @@ -163,7 +166,7 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: "--method", type=str, default=None, - choices=["mtp"], + choices=["mtp", "eagle3"], help="Speculative method", ) parser.add_argument( @@ -172,6 +175,12 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: default=1, help="Number of speculative tokens to generate per iteration (draft model runs this many times autoregressively)", ) + parser.add_argument( + "--draft-model", + type=str, + default=None, + help="Path to external Eagle3 draft model. Required when --method eagle3.", + ) parser.add_argument( "--max-num-batched-tokens", type=int, @@ -243,14 +252,25 @@ def _get_engine_kwargs(self) -> dict: ), ) if self.method and self.num_speculative_tokens > 0: - kwargs["speculative_config"] = SpeculativeConfig( - method=kwargs.pop("method"), - model=self.model, - num_speculative_tokens=kwargs.pop("num_speculative_tokens"), - ) + method = kwargs.pop("method") + num_spec_tokens = kwargs.pop("num_speculative_tokens") + draft_model = kwargs.pop("draft_model") + if method == "eagle3": + kwargs["speculative_config"] = SpeculativeConfig( + method=method, + model=draft_model, + num_speculative_tokens=num_spec_tokens, + ) + else: + kwargs["speculative_config"] = SpeculativeConfig( + method=method, + model=self.model, + num_speculative_tokens=num_spec_tokens, + ) else: kwargs.pop("method") kwargs.pop("num_speculative_tokens") + kwargs.pop("draft_model") kwargs["speculative_config"] = None # --enable-tbo [prefill|all] → enable_tbo + enable_tbo_decode diff --git a/atom/model_engine/block_manager.py b/atom/model_engine/block_manager.py index d62a2c413c..7e081a7253 100644 --- a/atom/model_engine/block_manager.py +++ b/atom/model_engine/block_manager.py @@ -118,6 +118,15 @@ def can_allocate(self, seq: Sequence) -> bool: block_id = self.hash_to_block_id.get(h, -1) if block_id == -1 or self.blocks[block_id].token_ids != token_ids: cache_miss = True + # If the entire prompt would be cached, force the last full block + # to recompute so prefill has at least one token to forward and + # produce logits for the next-token sampler. + if ( + not cache_miss + and i == seq.num_blocks - 1 + and len(token_ids) == self.block_size + ): + cache_miss = True if cache_miss: needed_free += 1 return ( @@ -142,6 +151,16 @@ def allocate(self, seq: Sequence): ) if block_id == -1 or self.blocks[block_id].token_ids != token_ids: cache_miss = True + # If the entire prompt would be cached, force the last full block + # to recompute so prefill has at least one token to forward and + # produce logits for the next-token sampler. Must mirror the same + # condition in can_allocate() so the block budget agrees. + if ( + not cache_miss + and i == seq.num_blocks - 1 + and len(token_ids) == self.block_size + ): + cache_miss = True if cache_miss: block_id = self._pop_free_block() block = self._allocate_block(block_id) diff --git a/atom/model_engine/engine_core.py b/atom/model_engine/engine_core.py index f220231343..ad57c56e66 100644 --- a/atom/model_engine/engine_core.py +++ b/atom/model_engine/engine_core.py @@ -192,6 +192,15 @@ def busy_loop(self): def _process_engine_step(self): result = self.scheduler.schedule() + + # Surface admit-rejected seqs (those `_unschedulable_reason` flags in + # the scheduler) through the same finished-seq path as normal seqs. + # Without this, `llm.generate()` blocks forever waiting for an output + # the rejected seq will never produce. + rejected = self.scheduler.take_rejected() + if rejected: + self.output_queue.put_nowait(rejected) + if result is None: return False scheduled_batch, seqs = result diff --git a/atom/model_engine/model_runner.py b/atom/model_engine/model_runner.py index ce607625df..1adbf9df25 100644 --- a/atom/model_engine/model_runner.py +++ b/atom/model_engine/model_runner.py @@ -85,8 +85,8 @@ def __init__( """Asynchronously copy the sampled_token_ids tensor to the host.""" # Deferred output is disabled when running in P/D disaggregation mode # (kv_transfer_config is set), enabled otherwise. - # TODO: enable deferred output in P/D disaggregation mode - self.is_deferred_out = not bool(runner.config.kv_transfer_config) + # TODO: In P/D disaggregation mode, if have issue, we can disable it + self.is_deferred_out = True self.runner = runner device = runner.device @@ -313,6 +313,15 @@ def prepare_input_ids( self.input_ids.np[:total_tokens_decode] = token_ids return self.input_ids.copy_to_gpu(total_tokens_decode) + # PD consumer first decode: no prior prefill step initialized + # prev_batch, so use scheduled_tokens directly for this step. + if self.prev_batch is None: + token_ids = scheduled_tokens[ + total_tokens_prefill : total_tokens_prefill + total_tokens_decode + ] + self.input_ids.np[:total_tokens_decode] = token_ids + return self.input_ids.copy_to_gpu(total_tokens_decode) + """for decode: input ids are from prev_sampled_token_ids""" deferred_curr_indices, deferred_prev_indices, new_curr_indices, is_all_same = ( self.get_token_locations(batch) @@ -559,6 +568,13 @@ def __init__(self, rank: int, config: Config): self.num_spec_tokens = ( self.config.speculative_config.num_speculative_tokens if use_spec else 0 ) + self.eagle3_mode = ( + self.config.speculative_config is not None + and self.config.speculative_config.method == "eagle3" + ) + + self.use_aux_hidden_state_outputs = False + self._aux_hidden_states = None self.tokenID_processor = tokenIDProcessor( self, self.config.max_num_batched_tokens, @@ -621,6 +637,18 @@ def __init__(self, rank: int, config: Config): torch.set_default_device(None) logger.info("Loading drafter model...") self.drafter.load_model(self.model) + + if self.eagle3_mode and self.config.speculative_config.use_aux_hidden_state: + aux_ids = self.config.speculative_config.eagle3_aux_layer_ids + if not aux_ids and hasattr( + self.model, "get_eagle3_aux_hidden_state_layers" + ): + aux_ids = list(self.model.get_eagle3_aux_hidden_state_layers()) + if aux_ids: + self.model.set_aux_hidden_state_layers(tuple(aux_ids)) + self.use_aux_hidden_state_outputs = True + logger.info(f"Eagle3 aux hidden state layers: {aux_ids}") + torch.set_default_device(self.device) self.async_execute_stream = torch.cuda.Stream(self.device) self.allocate_forward_vars() @@ -707,9 +735,18 @@ def is_qwen_next(self) -> bool: return False def is_deepseek_v4(self) -> bool: - if not hasattr(self.hf_text_config, "model_type"): - return False - return self.hf_text_config.model_type == "deepseek_v4" + # NOTE: `hf_text_config.model_type` reads "deepseek_v3" for V4 because + # `_CONFIG_REGISTRY` maps deepseek_v4 → deepseek_v3 (V4 reuses V3 schema). + # Use `architectures` (preserved by get_hf_config:567) instead. Covers + # both target (DeepseekV4ForCausalLM[NextN]) and draft (whose model_type + # SpeculativeConfig stamps as deepseek_v4_mtp). + arches = getattr(self.hf_text_config, "architectures", None) or [] + if any("DeepseekV4" in str(a) for a in arches): + return True + return getattr(self.hf_text_config, "model_type", None) in ( + "deepseek_v4", + "deepseek_v4_mtp", + ) def is_mimo_v2(self) -> bool: if not hasattr(self.hf_text_config, "model_type"): @@ -1054,8 +1091,15 @@ def allocate_forward_vars(self): "top_ks": CpuGpuBuffer(self.max_bs, **i32_kwargs), "top_ps": CpuGpuBuffer(self.max_bs, **f32_kwargs), # Keep enough space for MTP decode (max_q_len > 1). + # `extra_output_dims` lets a model insert dims between N and dim + # (e.g. DeepSeek-V4 returns the un-reduced mHC residual + # [N, hc_mult, dim] from forward, with hc_head + LM head deferred + # to compute_logits). Default `()` keeps the standard 2D layout. "outputs": torch.empty( - self.max_num_batched_tokens, hidden_size, dtype=hidden_type + self.max_num_batched_tokens, + *getattr(self.model, "extra_output_dims", ()), + hidden_size, + dtype=hidden_type, ), } if hasattr(self, "drafter"): @@ -1075,24 +1119,35 @@ def _get_num_kv_heads(self): return 1 def _get_total_num_layers(self): - """Return total layer count including draft (MTP) layers.""" + """Return total layer count including draft (MTP) layers. + + Drafts that own an independent KV cache via their own builder + (e.g. Eagle3 MHA draft on an MLA target) account for their layers + through that builder, so they are NOT added here. Only MTP-style + drafts that share the target's KV pool contribute. + """ total = self.config.hf_config.num_hidden_layers if self.config.speculative_config and hasattr(self, "drafter"): - draft_hf = self.config.speculative_config.draft_model_hf_config - total += getattr(draft_hf, "num_nextn_predict_layers", 1) + if not hasattr(self, "eagle3_draft_builder"): + draft_hf = self.config.speculative_config.draft_model_hf_config + total += getattr(draft_hf, "num_nextn_predict_layers", 1) return total def _compute_block_bytes(self): """Per-block bytes for the unified KV pool budget. - Delegates to the attention builder, which knows its own tensor - layout (MLA 576-dim packed, GDN-hybrid full-attn-only, MiMo-V2 - per-layer-type, standard MHA split-K/V). Mirror of - `attn_metadata_builder.allocate_kv_cache_tensors()` so the budget - math matches what's actually allocated. Per-request cache bytes - are accounted for separately via `compute_per_req_cache_bytes()`. + Sum across all attention builders attached to this runner: the + target builder always, plus an optional `eagle3_draft_builder` + when a heterogeneous spec-decode draft owns its own KV pool. Each + builder knows its own tensor layout (MLA 576-dim packed, GDN-hybrid + full-attn-only, MiMo-V2 per-layer-type, standard MHA split-K/V, + Eagle3 independent MHA). Per-request cache bytes are accounted + for separately via `compute_per_req_cache_bytes()`. """ - return self.attn_metadata_builder.compute_block_bytes() + block_bytes = self.attn_metadata_builder.compute_block_bytes() + if hasattr(self, "eagle3_draft_builder"): + block_bytes += self.eagle3_draft_builder.compute_block_bytes() + return block_bytes def _estimate_cudagraph_overhead(self): """Estimate GPU memory consumed by CUDA graph capture. @@ -1255,13 +1310,24 @@ def allocate_kv_cache(self, num_kvcache_blocks): num_draft_layers = 0 if self.config.speculative_config and hasattr(self, "drafter"): draft_hf_config = self.config.speculative_config.draft_model_hf_config - # For MTP, use num_nextn_predict_layers instead of num_hidden_layers - num_draft_layers = getattr(draft_hf_config, "num_nextn_predict_layers", 1) - total_num_layers += num_draft_layers - logger.info( - f"Allocating KV cache for {hf_config.num_hidden_layers} target layers + " - f"{num_draft_layers} draft (MTP) layers = {total_num_layers} total layers" - ) + if hasattr(self, "eagle3_draft_builder"): + # Heterogeneous draft (e.g. Eagle3 MHA on MLA target) owns + # its own KV pool via its builder; don't add to target's count. + num_draft_layers = draft_hf_config.num_hidden_layers + logger.info( + f"Allocating KV cache for {hf_config.num_hidden_layers} target layers + " + f"{num_draft_layers} Eagle3 draft layers (separate non-MLA cache)" + ) + else: + # For MTP, use num_nextn_predict_layers instead of num_hidden_layers + num_draft_layers = getattr( + draft_hf_config, "num_nextn_predict_layers", 1 + ) + total_num_layers += num_draft_layers + logger.info( + f"Allocating KV cache for {hf_config.num_hidden_layers} target layers + " + f"{num_draft_layers} draft (MTP) layers = {total_num_layers} total layers" + ) # Primary KV cache allocation (model-agnostic, delegated to the # attention builder). Each builder owns its tensor layout: MLA → @@ -1277,6 +1343,16 @@ def allocate_kv_cache(self, num_kvcache_blocks): for name, value in main_kv.items(): setattr(self, name, value) + # Heterogeneous draft (e.g. Eagle3 MHA alongside an MLA target) owns + # its own KV pool through a sibling builder; same protocol as above, + # tensors land under namespaced keys (eagle3_kv_cache, eagle3_kv_scale). + if hasattr(self, "eagle3_draft_builder"): + draft_kv = self.eagle3_draft_builder.allocate_kv_cache_tensors( + num_kv_heads, num_draft_layers + ) + for name, value in draft_kv.items(): + setattr(self, name, value) + # Per-request cache allocation (model-agnostic, delegated to the # attention metadata builder). For GDN this returns # `{"mamba_k_cache": ..., "mamba_v_cache": ...}`; for stateless @@ -1302,10 +1378,12 @@ def allocate_kv_cache(self, num_kvcache_blocks): kv_cache_tensors = [] layer_id = 0 # Promote to self so the attention builder's build_kv_cache_tensor() - # can access it without recomputing from drafter state. + # can access it without recomputing from drafter state. Heterogeneous + # drafts (Eagle3) own their own layer space via their builder, so + # leave mtp_start_layer_idx at hf_config.num_hidden_layers in that mode. self.mtp_start_layer_idx = ( self.drafter.model.model.mtp_start_layer_idx - if hasattr(self, "drafter") + if hasattr(self, "drafter") and not hasattr(self, "eagle3_draft_builder") else hf_config.num_hidden_layers ) for model_name, model in models_to_bind: @@ -1314,6 +1392,18 @@ def allocate_kv_cache(self, num_kvcache_blocks): ) for module in model.modules(): + # Drafts that own an independent KV pool (Eagle3) bind through + # their sibling builder first; for unrecognized modules it + # returns None and we fall through to the target builder. + if model_name == "draft" and hasattr(self, "eagle3_draft_builder"): + kv_cache_tensor = self.eagle3_draft_builder.build_kv_cache_tensor( + layer_id, module + ) + if kv_cache_tensor is not None: + kv_cache_tensors.append(kv_cache_tensor) + layer_id += 1 + continue + # Per-attention-type binding is owned by the attention # metadata builder; ModelRunner only walks modules and # collects the resulting KVCacheTensor entries. The builder @@ -1625,7 +1715,12 @@ def run_model( label += f" tok={batch.total_tokens_num} ctx={ctx_str}" label += "]" with record_function(label): - hidden_states = self.model(input_ids, positions) + model_output = self.model(input_ids, positions) + if self.use_aux_hidden_state_outputs: + hidden_states, self._aux_hidden_states = model_output + else: + hidden_states = model_output + self._aux_hidden_states = None logits = self.model.compute_logits(hidden_states) else: # decode[bs=128 tok=128 d=128] or decode[bs=128 tok=128 p=2 d=126 spec=3] @@ -1645,6 +1740,12 @@ def run_model( self.graphs[graph_key].replay() num_tokens = context.batch_size * max_q_len hidden_states = self.forward_vars["outputs"][:num_tokens] + if graph_key in self.graph_aux_hidden: + self._aux_hidden_states = [ + aux[:num_tokens] for aux in self.graph_aux_hidden[graph_key] + ] + else: + self._aux_hidden_states = None if self.logits_in_graph: logits = self.graph_logits[graph_key][:num_tokens] else: @@ -1803,6 +1904,7 @@ def async_proc_aggregation(self) -> KVConnectorOutput: if connector is None: return KVConnectorOutput(finished_sending=[], finished_recving=[]) done_sending, done_recving = connector.get_finished() + return KVConnectorOutput( finished_sending=done_sending, finished_recving=done_recving ) @@ -1833,6 +1935,7 @@ def propose_draft_token_ids( num_reject_tokens=num_reject_tokens, next_token_ids=next_token_ids, last_token_indices=last_token_indices, + aux_hidden_states=self._aux_hidden_states, ) return self.tokenID_processor.prepare_draft_ids(batch, draft_token) @@ -1877,9 +1980,12 @@ def capture_cudagraph(self): positions = self.forward_vars["positions"].gpu outputs = self.forward_vars["outputs"] self.forward_vars["kv_indptr"].gpu.zero_() + if self.is_deepseek_v32 and "sparse_kv_indptr" in self.forward_vars: + self.forward_vars["sparse_kv_indptr"].gpu.zero_() self.graphs: dict[tuple[int, int], torch.cuda.CUDAGraph] = dict() self.graph_logits: dict[tuple[int, int], torch.Tensor] = dict() + self.graph_aux_hidden: dict[tuple[int, int], list[torch.Tensor]] = dict() self.graph_pool = None is_tbo = self.config.enable_tbo and isinstance(self.model, UBatchWrapper) # TBO graphs don't capture compute_logits, so disable logits_in_graph. @@ -1927,12 +2033,17 @@ def capture_cudagraph(self): num_tokens=num_tokens, num_tokens_across_dp=num_tokens_across_dp, ubatch_slices=ubatch_slices, + in_hipgraph=True, ) # Warmup - outputs[:num_tokens] = self.model( + model_output = self.model( input_ids[:num_tokens], positions[:num_tokens] ) + if self.use_aux_hidden_state_outputs: + outputs[:num_tokens] = model_output[0] + else: + outputs[:num_tokens] = model_output if self.logits_in_graph: self.model.compute_logits(outputs[:num_tokens]) @@ -1951,13 +2062,20 @@ def capture_cudagraph(self): gc.stream, output_buffer=outputs[:num_tokens], ) + graph_aux = None else: # Standard single-stream capture graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph, self.graph_pool, stream=gc.stream): - outputs[:num_tokens] = self.model( + model_output = self.model( input_ids[:num_tokens], positions[:num_tokens] ) + if self.use_aux_hidden_state_outputs: + outputs[:num_tokens] = model_output[0] + graph_aux = model_output[1] + else: + outputs[:num_tokens] = model_output + graph_aux = None if self.logits_in_graph: graph_logits = self.model.compute_logits( outputs[:num_tokens] @@ -1967,6 +2085,8 @@ def capture_cudagraph(self): self.graphs[(bs, max_q_len)] = graph if self.logits_in_graph and ubatch_slices is None: self.graph_logits[(bs, max_q_len)] = graph_logits + if graph_aux is not None: + self.graph_aux_hidden[(bs, max_q_len)] = graph_aux torch.cuda.synchronize() self.graph_bs.sort(reverse=False) diff --git a/atom/model_engine/scheduler.py b/atom/model_engine/scheduler.py index 9e9acd6ced..a977c95eb4 100644 --- a/atom/model_engine/scheduler.py +++ b/atom/model_engine/scheduler.py @@ -224,9 +224,13 @@ def __init__( is_dummy_run: bool = False, num_spec_step: int = 0, scheduled_spec_decode_tokens: dict[int, np.ndarray] | None = None, + remote_kv_block_ids: list[int] | None = None, + remote_kv_seq_blocks: dict[int, list[int]] | None = None, ): if scheduled_spec_decode_tokens is None: scheduled_spec_decode_tokens = {} + self.remote_kv_block_ids = remote_kv_block_ids or [] + self.remote_kv_seq_blocks = remote_kv_seq_blocks or {} self.req_ids = list(seqs.keys()) # self.scheduled_tokens = [ @@ -368,6 +372,7 @@ class Scheduler: def __init__(self, config: Config): self.max_num_seqs = config.max_num_seqs self.max_num_batched_tokens = config.max_num_batched_tokens + self.max_model_len = config.max_model_len self.bos_token_id = config.bos_token_id self.eos_token_id = config.eos_token_id self.stop_token_ids = config.stop_token_ids @@ -376,6 +381,11 @@ def __init__(self, config: Config): self.running: deque[Sequence] = deque() self.config = config + # Admit-rejected seqs (those `_unschedulable_reason` flags). Drained + # by `take_rejected` each EngineCore step; routed through the same + # output_queue path as forward-finished seqs. + self._rejected: list[Sequence] = [] + # KV transfer bookkeeping self.finished_recving_kv_req_ids: list[int] = [] self.deferred_free_blocks: dict[int, Sequence] = {} @@ -405,7 +415,13 @@ def __init__(self, config: Config): self.kv_connector = get_kvconnector("scheduler", config) def is_finished(self): - return not self.waiting and not self.running + # `_rejected` must be considered too: if a batch of seqs is all + # oversized, schedule() moves them straight from `waiting` to + # `_rejected`, leaving both `waiting` and `running` empty. Without + # this check, busy_loop's `is_finished()` short-circuits to True + # before EngineCore drains `_rejected` via take_rejected(), and + # llm.generate() blocks forever. + return not self.waiting and not self.running and not self._rejected def add(self, seq: Sequence): self._warn_if_unschedulable(seq) @@ -416,52 +432,73 @@ def extend(self, seqs: list[Sequence]): self._warn_if_unschedulable(seq) self.waiting.extend(seqs) - def _warn_if_unschedulable(self, seq: Sequence) -> None: - """Detect requests that exceed static scheduler/KV-pool capacity. - - These requests would otherwise sit in `waiting` forever (head-of-line - blocking the prefill loop, which `break`s on the first oversized seq) - with no log output. Surface a single warning at submit time so the - caller knows the request will never be picked up. - - Three permanent failure modes: + def _unschedulable_reason(self, seq: Sequence) -> Optional[str]: + """Return a human-readable reason if `seq` is permanently unschedulable. + + Only checks static (configuration-time) capacity. Dynamic conditions + that can clear up as other seqs finish (e.g. transiently full + per-req-cache pool) are NOT checked here — they're warned at submit + time (`_warn_if_unschedulable`) but not eligible for permanent drop + at schedule time, since the prefill loop's existing `can_allocate` + check will retry them later. + + Permanent failure modes (each leaves the seq stuck in `waiting` + forever and would head-of-line block the prefill loop, which + `break`s on the first oversized seq): + - prompt longer than `max_model_len` → exceeds per-seq KV cache + geometry; attention backends size `block_tables` as + `max_model_len // block_size` cols and would crash with a + broadcast error at prepare-time. (Checked first since it's the + usual actionable cause.) - prompt longer than `max_num_batched_tokens` → no single prefill forward can ever fit it - prompt's KV blocks (+ per-req cache reservation) exceed the total pool size → never fits even on a fully empty pool - - request needs a per-req cache slot but the model was started with - zero slots (e.g. GDN/V4 with `max_num_seqs=0`) + + Called at submit time (`_warn_if_unschedulable`, which logs the + reason and adds extra dynamic warnings) and at schedule time + (drops the seq before it reaches the attention backend). """ num_tokens = seq.num_tokens + if num_tokens > self.max_model_len: + return ( + f"input tokens={num_tokens} > max_model_len={self.max_model_len}. " + f"Increase --max-model-len or shorten the prompt." + ) if num_tokens > self.max_num_batched_tokens: - logger.warning( - "Request %s will never be scheduled: input tokens=%d > " - "max_num_batched_tokens=%d. Increase --max-num-batched-tokens " - "or shorten the prompt.", - seq.id, - num_tokens, - self.max_num_batched_tokens, + return ( + f"input tokens={num_tokens} > max_num_batched_tokens=" + f"{self.max_num_batched_tokens}. Increase --max-num-batched-tokens " + f"or shorten the prompt." ) - return - bm = self.block_manager per_req_cost = bm.per_req_cache_equiv_blocks if seq.has_per_req_cache else 0 total_blocks = len(bm.blocks) if seq.num_blocks + per_req_cost > total_blocks: - logger.warning( - "Request %s will never be scheduled: needs %d KV blocks " - "(%d for %d input tokens + %d for per-req cache) > " - "total pool blocks=%d. Reduce prompt length, lower " - "--max-num-seqs, or raise --gpu-memory-utilization.", - seq.id, - seq.num_blocks + per_req_cost, - seq.num_blocks, - num_tokens, - per_req_cost, - total_blocks, + return ( + f"needs {seq.num_blocks + per_req_cost} KV blocks " + f"({seq.num_blocks} for {num_tokens} input tokens + " + f"{per_req_cost} for per-req cache) > total pool blocks=" + f"{total_blocks}. Reduce prompt length, lower --max-num-seqs, " + f"or raise --gpu-memory-utilization." ) - return + return None + def _warn_if_unschedulable(self, seq: Sequence) -> None: + """Log a single warning at submit time for permanently-unschedulable + sequences. The seq still enters `waiting`; the prefill scheduler drops + it later (see `schedule`). + + Also surfaces a dynamic configuration-time-only warning when the + model was started with zero per-req-cache slots (max_num_seqs=0) — + this is permanent if it holds at submit time, but is NOT eligible + for schedule-time drop (a future config change could create slots). + """ + reason = self._unschedulable_reason(seq) + if reason is not None: + logger.warning("Request %s will never be scheduled: %s", seq.id, reason) + return + bm = self.block_manager if ( seq.has_per_req_cache and not bm.free_per_req_cache_groups @@ -469,11 +506,23 @@ def _warn_if_unschedulable(self, seq: Sequence) -> None: ): logger.warning( "Request %s will never be scheduled: needs per-req cache slot " - "but no slots were allocated (max_num_seqs=0 for this model " - "type).", + "but no slots were allocated (max_num_seqs=0 for this model type).", seq.id, ) + def take_rejected(self) -> list[Sequence]: + """Pop and return any seqs the prefill scheduler dropped because + `_unschedulable_reason` flagged them (oversized prompt, exhausted + pool, etc.). Caller (EngineCore) pushes them onto the same + output_queue as forward-finished seqs so `llm.generate()` returns + an output for them instead of blocking forever. + """ + if not self._rejected: + return [] + out = self._rejected + self._rejected = [] + return out + def schedule(self) -> tuple[ScheduledBatch, dict[int, Sequence]]: """Select the next batch of sequences for a forward pass. @@ -494,6 +543,22 @@ def schedule(self) -> tuple[ScheduledBatch, dict[int, Sequence]]: while self.waiting and num_seqs_prefill < self.max_num_seqs: seq = self.waiting.popleft() + # Drop seqs the static-capacity check at submit-time flagged as + # permanently unschedulable (oversized prompt, exhausted pool, + # etc.). They've already been warned; mark FINISHED + record the + # rejection reason and route them to `_rejected` so EngineCore + # surfaces them through the same output_queue as forward-finished + # seqs. Without this they'd reach the attention backend (where an + # oversized prompt crashes with a broadcast error) AND + # `llm.generate()` would block forever waiting for an output. + # Re-check here (not just at submit) since pool state may change. + unschedulable = self._unschedulable_reason(seq) + if unschedulable is not None: + seq.status = SequenceStatus.FINISHED + seq.leave_reason = f"unschedulable: {unschedulable}" + self._rejected.append(seq) + continue + # KV Transfer: skip request if still waiting for remote KVs waiting_remote_to_waiting_ready = False if seq.status == SequenceStatus.WAITING_FOR_REMOTE_KVS: @@ -512,6 +577,27 @@ def schedule(self) -> tuple[ScheduledBatch, dict[int, Sequence]]: self.kv_connector.get_num_new_matched_tokens(seq) ) + if waiting_remote_to_waiting_ready: + seq.status = SequenceStatus.RUNNING + seq.is_first_decode = True + first_token_id = (seq.kv_transfer_params or {}).get("first_token_id") + if first_token_id is not None: + seq.append_token(first_token_id) + seq._injected_t0 = first_token_id + logger.info( + "[PD-TRANSITION] seq %s: num_tokens=%d, " + "num_prompt=%d, blocks=%d, first_token=%s, " + "last_5_tids=%s", + seq.id, + seq.num_tokens, + seq.num_prompt_tokens, + len(seq.block_table), + first_token_id, + seq.token_ids[-5:], + ) + self.running.append(seq) + continue + num_new_tokens = seq.num_tokens - seq.num_cached_tokens if ( num_batched_tokens + num_new_tokens > self.max_num_batched_tokens @@ -520,8 +606,7 @@ def schedule(self) -> tuple[ScheduledBatch, dict[int, Sequence]]: self.waiting.appendleft(seq) break - if not waiting_remote_to_waiting_ready: - self.block_manager.allocate(seq) + self.block_manager.allocate(seq) if self.kv_connector is not None: self.kv_connector.update_state_after_alloc(seq) @@ -531,12 +616,6 @@ def schedule(self) -> tuple[ScheduledBatch, dict[int, Sequence]]: seq.status = SequenceStatus.WAITING_FOR_REMOTE_KVS continue - if waiting_remote_to_waiting_ready: - seq.status = SequenceStatus.RUNNING - seq.is_first_decode = True - self.running.append(seq) - continue - num_seqs_prefill += 1 num_new_tokens = seq.num_tokens - seq.num_cached_tokens if self.cache_stats: @@ -558,9 +637,11 @@ def schedule(self) -> tuple[ScheduledBatch, dict[int, Sequence]]: total_tokens_num_prefill = sum(num_scheduled_tokens) if num_seqs_prefill > 0: + cached_per_req = [s.num_cached_tokens for s in scheduled_seqs.values()] logger.info( f"Scheduled prefill batch: {num_seqs_prefill} reqs, " - f"{total_tokens_num_prefill} tokens, " + f"{total_tokens_num_prefill} new tokens " + f"(cached: {cached_per_req}, new: {num_scheduled_tokens}), " f"req_ids: {tuple(scheduled_seqs.keys())}" ) self.prev_prompt = True @@ -585,6 +666,8 @@ def schedule(self) -> tuple[ScheduledBatch, dict[int, Sequence]]: # --- Decode scheduling --- num_seqs_decode = 0 num_new_tokens = self.mtp_k + 1 + remote_kv_blocks: set[int] = set() + remote_kv_seq_blocks: dict[int, list[int]] = {} while self.running and num_seqs_decode < self.max_num_seqs: seq = self.running.popleft() while not self.block_manager.can_append(seq, num_new_tokens): @@ -597,10 +680,31 @@ def schedule(self) -> tuple[ScheduledBatch, dict[int, Sequence]]: if seq.spec_token_ids.size > 0: scheduled_spec_decode_tokens[seq.id] = seq.spec_token_ids num_seqs_decode += 1 - # Skip block append for the first decode step after remote - # prefill — blocks were already allocated during prefill. - if not getattr(seq, "is_first_decode", False): + # For PD first-decode: if T0 was injected, may_append is + # needed for the new position N. Without T0 injection, + # blocks were already allocated during prefill. + is_first = getattr(seq, "is_first_decode", False) + if is_first and seq.block_table: + remote_kv_blocks.update(seq.block_table) + remote_kv_seq_blocks[seq.id] = list(seq.block_table) + has_injected_t0 = ( + is_first + and (seq.kv_transfer_params or {}).get("first_token_id") is not None + ) + if not is_first or has_injected_t0: self.block_manager.may_append(seq, num_new_tokens) + if is_first: + logger.info( + "[PD-FIRST-DECODE] seq %s: num_tokens=%d, " + "blocks=%d, injected_t0=%s, " + "last_block_num=%d, context_will_be=%d", + seq.id, + seq.num_tokens, + len(seq.block_table), + has_injected_t0, + seq.last_block_num_tokens, + seq.num_tokens, + ) scheduled_seqs[seq.id] = seq seq.type = SequenceType.DECODE num_scheduled_tokens.append(num_new_tokens) @@ -626,6 +730,8 @@ def schedule(self) -> tuple[ScheduledBatch, dict[int, Sequence]]: connector_meta_output=connector_meta_output, num_spec_step=self.mtp_k, scheduled_spec_decode_tokens=scheduled_spec_decode_tokens, + remote_kv_block_ids=sorted(remote_kv_blocks) if remote_kv_blocks else [], + remote_kv_seq_blocks=remote_kv_seq_blocks, ) return (decode_batch, scheduled_seqs) @@ -643,6 +749,7 @@ def preempt(self, seq: Sequence): seq.num_rejected = 0 seq.num_bonus_tokens = 0 seq.spec_token_ids = np.array([], dtype=np.int32) + seq.is_first_decode = False self.block_manager.deallocate(seq) self.waiting.appendleft(seq) @@ -676,12 +783,21 @@ def postprocess( continue token_ids = prev_token_ids[idx] num_new_token = len(token_ids) - if self.spec_stats: - self.spec_stats.update(num_new_token) if is_deferred_out or self.use_spec: num_rejected = fwd_output.num_rejected[idx] num_bonus = fwd_output.num_bonus[idx] offset = 0 if (num_new_token + num_rejected) == 1 else self.mtp_k + # Align stats with vLLM: only count steps that actually ran + # speculation (drafts proposed and validated). Skip the + # prefill-only step where no draft tokens were scored against + # the target — vLLM gates this via + # `if scheduled_spec_token_ids and generated_token_ids`. + if ( + self.spec_stats + and num_new_token > 0 + and (num_new_token + num_rejected) > 1 + ): + self.spec_stats.update(num_new_token) seq.num_rejected = num_rejected seq.num_bonus_tokens = num_bonus for i, el in enumerate(token_ids): @@ -697,15 +813,41 @@ def postprocess( seq.append_token(token_id) new_tokens = token_ids + injected_t0 = getattr(seq, "_injected_t0", None) + if injected_t0 is not None: + new_tokens = [injected_t0] + list(new_tokens) + seq._injected_t0 = None + if self.mtp_k > 0: # idx already resolved above via get_idx seq.spec_token_ids = draft_token_ids[idx] - if seq.num_completion_tokens == 1 and seq.first_token_time == 0.0: + if seq.num_completion_tokens <= 3 and seq.kv_transfer_params: + logger.info( + "[PD-DECODE] seq %s: comp_tokens=%d, " + "new_token=%s, num_tokens=%d, blocks=%d", + seq.id, + seq.num_completion_tokens, + token_ids, + seq.num_tokens, + len(seq.block_table), + ) + if seq.num_completion_tokens >= 1 and seq.first_token_time == 0.0: seq.first_token_time = time.time() num_tokens = seq.num_tokens - self.mtp_k - num_rejected leave_reason = None + # MTP edge case: `rejection_sampler` does NOT inspect EOS — it + # only compares draft vs target_argmax for acceptance. So when + # the verified token is EOS the kernel still emits 1+ accepted + # bonus tokens after EOS (often BOS, since the model naturally + # starts a new sentence). Without truncating, those post-EOS + # tokens leak into the detokenized output (e.g. "...6."). + # Empirically confirmed via DIAG: `token_ids=[EOS=1, BOS=0]`, + # `eos_idx=0`, `num_new=2`, `num_rejected=0` for V4-Pro MTP-1. + # Track the earliest stop position so `num_tokens` can drop the + # spurious tail below. + stop_at_idx: Optional[int] = None # Check if sequence ends with any stop sequence for stop_seq in seq.stop_token_sequences: stop_len = len(stop_seq) @@ -715,6 +857,10 @@ def postprocess( offset = num_tokens - i if seq.token_ids[offset - stop_len : offset] == stop_seq: is_stop = True + # `i` counts back from the last sampled token + # (i=0 = last). Truncate to include this stop + # sequence (drop everything after it). + stop_at_idx = num_new_token - 1 - i break if is_stop: leave_reason = "stop_sequence" @@ -723,16 +869,31 @@ def postprocess( # Check the last token in the list for EOS if token_ids and not seq.ignore_eos and self.eos_token_id in token_ids: leave_reason = "eos" + stop_at_idx = token_ids.index(self.eos_token_id) elif not seq.ignore_eos and any( t in self.stop_token_ids for t in token_ids ): - first_stop_token = next( - t for t in token_ids if t in self.stop_token_ids + stop_at_idx = next( + i for i, t in enumerate(token_ids) if t in self.stop_token_ids ) - leave_reason = f"stop_{first_stop_token}" - elif seq.num_completion_tokens >= seq.max_tokens: + leave_reason = f"stop_{token_ids[stop_at_idx]}" + elif (num_tokens - seq.num_prompt_tokens) >= seq.max_tokens: + # Use the local `num_tokens` (= seq.num_tokens - mtp_k - + # num_rejected, set at line 716) instead of the property + # `seq.num_completion_tokens` which still reflects the + # raw mtp_k+1 placeholder bump from `prepare_decode`. The + # property over-counts by `mtp_k + num_rejected`, causing + # max_tokens to trip that many tokens early (visible as + # `output tokens=95` for max_tokens=100, mtp_k=3). Non-MTP + # path: mtp_k = num_rejected = 0 → behavior unchanged. leave_reason = "max_tokens" + # Drop accepted-draft tokens past the stop position (MTP only — + # for non-spec the sampler emits exactly 1 token so this is a + # no-op). + if stop_at_idx is not None and stop_at_idx < num_new_token - 1: + num_tokens -= (num_new_token - 1) - stop_at_idx + # Prepare stream output if stream_output_queue is not None and new_tokens: if self.kv_connector is not None and leave_reason is not None: diff --git a/atom/model_loader/loader.py b/atom/model_loader/loader.py index 987b05e18a..057f007d97 100644 --- a/atom/model_loader/loader.py +++ b/atom/model_loader/loader.py @@ -196,6 +196,8 @@ def load_model_in_plugin_mode( prefix: str = "", weights_mapper: WeightsMapper | None = None, load_fused_expert_weights_fn=None, + spec_decode: bool = False, + hf_config_override: AutoConfig | None = None, ) -> set[str]: # during loading model, the outplace operation may consume more @@ -216,17 +218,24 @@ def _empty_cache(): model_name_or_path = config.plugin_config.model_config.model_path _empty_cache() - config_for_loading = ( - config.hf_config.text_config - if hasattr(config.hf_config, "text_config") - else config.hf_config - ) + if hf_config_override is not None: + config_for_loading = getattr( + hf_config_override, "hf_config", hf_config_override + ) + if hasattr(config_for_loading, "text_config"): + config_for_loading = config_for_loading.text_config + else: + config_for_loading = ( + config.hf_config.text_config + if hasattr(config.hf_config, "text_config") + else config.hf_config + ) loaded_weights_record = load_model( model=model, model_name_or_path=model_name_or_path, hf_config=config_for_loading, load_dummy=config.load_dummy, - spec_decode=False, + spec_decode=spec_decode, prefix=prefix, is_plugin_mode=True, weights_mapper=weights_mapper, @@ -287,6 +296,12 @@ def extract_expert_target_and_id(name: str) -> Tuple[str, int] | None: if weights_mapper is None: weights_mapper = getattr(model, "weights_mapper", None) params_dict = dict(model.named_parameters()) + # Load `mtp.*` ckpt entries only when this is a spec-decode draft model + # AND it actually carries `mtp.*` params. The two-condition gate handles + # both directions: target loads (spec_decode=False) skip mtp regardless, + # and spec drafts that don't use the standard `mtp.*` param naming + # (e.g. Eagle3 with its own arch) also skip cleanly. + need_load_mtp: bool = spec_decode and any("mtp" in n for n in params_dict) # Pre-index expert_mapping by weight_name_part for O(1) lookup. # Original code does O(N) scan of expert_mapping (768 entries) per tensor, @@ -318,6 +333,20 @@ def extract_expert_target_and_id(name: str) -> Tuple[str, int] | None: with concurrent.futures.ThreadPoolExecutor() as executor: futures = [] + use_threadpool = envs.ATOM_LOADER_USE_THREADPOOL + if use_threadpool: + executor = concurrent.futures.ThreadPoolExecutor() + else: + executor = None + futures = [] + + def _submit(fn, *args): + if executor is not None: + futures.append(executor.submit(fn, *args)) + else: + fn(*args) + + try: disable_mmap = envs.ATOM_DISABLE_MMAP for name, weight_tensor in safetensors_weights_iterator( model_name_or_path, disable_mmap=disable_mmap @@ -330,7 +359,7 @@ def extract_expert_target_and_id(name: str) -> Tuple[str, int] | None: name = mapped_name if load_dummy: continue - if "mtp" in name and not spec_decode: + if "mtp" in name and not need_load_mtp: continue if name.endswith("kv_scale") or "inv_freq" in name: continue @@ -389,11 +418,7 @@ def extract_expert_target_and_id(name: str) -> Tuple[str, int] | None: ) continue weight_loader = getattr(param, "weight_loader") - futures.append( - executor.submit( - weight_loader, param, weight_tensor, shard_idx - ) - ) + _submit(weight_loader, param, weight_tensor, shard_idx) loaded_weights_record.add(prefix + param_name) else: # Checkpoint has separate weights, load into fused param @@ -407,12 +432,7 @@ def extract_expert_target_and_id(name: str) -> Tuple[str, int] | None: dropped_ckpt_keys.append((_orig_ckpt_name, param_name)) break weight_loader = getattr(param, "weight_loader") - # weight_loader(param, weight_tensor, shard_id) - futures.append( - executor.submit( - weight_loader, param, weight_tensor, shard_id - ) - ) + _submit(weight_loader, param, weight_tensor, shard_id) loaded_weights_record.add(prefix + param_name) break else: @@ -471,7 +491,7 @@ def extract_expert_target_and_id(name: str) -> Tuple[str, int] | None: ) and name not in params_dict: matched = True break - if "mtp" in name and not spec_decode: + if "mtp" in name and not need_load_mtp: matched = True break try: @@ -482,21 +502,19 @@ def extract_expert_target_and_id(name: str) -> Tuple[str, int] | None: matched = True break weight_loader = getattr(param, "weight_loader") - futures.append( - executor.submit( - weight_loader, - param, - weight_tensor, - name, - shard_id, - expert_id, - ) + _submit( + weight_loader, + param, + weight_tensor, + name, + shard_id, + expert_id, ) loaded_weights_record.add(prefix + name) matched = True break if not matched: - if "mtp" in name and not spec_decode: + if "mtp" in name and not need_load_mtp: continue if merged_target := extract_expert_target_and_id(name): fused_name, expert_id = merged_target @@ -508,15 +526,13 @@ def extract_expert_target_and_id(name: str) -> Tuple[str, int] | None: weight_loader = getattr( param, "weight_loader", default_weight_loader ) - futures.append( - executor.submit( - weight_loader, - param, - weight_tensor, - "", # use merged moe loader - "", - expert_id, - ) + _submit( + weight_loader, + param, + weight_tensor, + "", # use merged moe loader + "", + expert_id, ) loaded_weights_record.add(prefix + name) try: @@ -527,9 +543,7 @@ def extract_expert_target_and_id(name: str) -> Tuple[str, int] | None: weight_loader = getattr( param, "weight_loader", default_weight_loader ) - futures.append( - executor.submit(weight_loader, param, weight_tensor) - ) + _submit(weight_loader, param, weight_tensor) loaded_weights_record.add(prefix + name) else: # Model doesn't have expert mapping, use generic loading @@ -541,12 +555,12 @@ def extract_expert_target_and_id(name: str) -> Tuple[str, int] | None: weight_loader = getattr( param, "weight_loader", default_weight_loader ) - # weight_loader(param, weight_tensor) - futures.append(executor.submit(weight_loader, param, weight_tensor)) + _submit(weight_loader, param, weight_tensor) loaded_weights_record.add(prefix + name) - # Wait for all tasks to complete and raise any exceptions. - for future in concurrent.futures.as_completed(futures): - future.result() + finally: + if executor is not None: + concurrent.futures.wait(futures) + executor.shutdown(wait=True) # Verify every model parameter actually got loaded from the checkpoint. # Without this check, weights_mapping bugs (e.g. a substring rule diff --git a/atom/model_ops/__init__.py b/atom/model_ops/__init__.py index 3883dd6944..4670a5cdb2 100644 --- a/atom/model_ops/__init__.py +++ b/atom/model_ops/__init__.py @@ -1,5 +1,7 @@ from .paged_attention import PagedAttention -from atom.plugin.sglang.attention_backend.radix_attention import RadixAttention +from atom.plugin.sglang.attention_backend.full_attention.radix_attention import ( + RadixAttention, +) # This global class is used to construct the attention op in model, # it can be assigned to different attention ops. diff --git a/atom/model_ops/attention_gdn.py b/atom/model_ops/attention_gdn.py index 26581188c3..d6230e9eba 100644 --- a/atom/model_ops/attention_gdn.py +++ b/atom/model_ops/attention_gdn.py @@ -31,17 +31,19 @@ def fused_gdn_gating_kernel( dt_bias, seq_len, NUM_HEADS: tl.constexpr, + stride_a_batch, + stride_b_batch, beta: tl.constexpr, threshold: tl.constexpr, BLK_HEADS: tl.constexpr, ): i_b, i_s, i_d = tl.program_id(0), tl.program_id(1), tl.program_id(2) head_off = i_d * BLK_HEADS + tl.arange(0, BLK_HEADS) - off = i_b * seq_len * NUM_HEADS + i_s * NUM_HEADS + head_off + out_off = i_b * seq_len * NUM_HEADS + i_s * NUM_HEADS + head_off mask = head_off < NUM_HEADS blk_A_log = tl.load(A_log + head_off, mask=mask) - blk_a = tl.load(a + off, mask=mask) - blk_b = tl.load(b + off, mask=mask) + blk_a = tl.load(a + i_b * stride_a_batch + head_off, mask=mask) + blk_b = tl.load(b + i_b * stride_b_batch + head_off, mask=mask) blk_bias = tl.load(dt_bias + head_off, mask=mask) # If the model is loaded in fp16, without the .float() here, A might be -inf x = blk_a.to(tl.float32) + blk_bias.to(tl.float32) @@ -49,11 +51,13 @@ def fused_gdn_gating_kernel( beta * x <= threshold, (1 / beta) * tl.log(1 + tl.exp(beta * x)), x ) blk_g = -tl.exp(blk_A_log.to(tl.float32)) * softplus_x - tl.store(g + off, blk_g.to(g.dtype.element_ty), mask=mask) + tl.store(g + out_off, blk_g.to(g.dtype.element_ty), mask=mask) # compute beta_output = sigmoid(b) blk_beta_output = tl.sigmoid(blk_b.to(tl.float32)) tl.store( - beta_output + off, blk_beta_output.to(beta_output.dtype.element_ty), mask=mask + beta_output + out_off, + blk_beta_output.to(beta_output.dtype.element_ty), + mask=mask, ) @@ -85,6 +89,8 @@ def fused_gdn_gating( dt_bias, seq_len, num_heads, + a.stride(0), + b.stride(0), beta, threshold, 8, @@ -161,6 +167,7 @@ def forward( fwd_ctx.attn_metadata, "gdn_metadata", None ) if gdn_metadata is None: + core_attn_out.zero_() return core_attn_out gdn_cache = fwd_ctx.kv_cache_data @@ -373,4 +380,8 @@ def forward( else: core_attn_out[:num_actual_tokens] = core_attn_out_non_spec.squeeze(0) + # Zero padding tail for CUDA graph replay safety + if num_actual_tokens < core_attn_out.shape[0]: + core_attn_out[num_actual_tokens:].zero_() + return core_attn_out diff --git a/atom/model_ops/attention_mha.py b/atom/model_ops/attention_mha.py index c3ebd525a8..9e0d5dbb0b 100644 --- a/atom/model_ops/attention_mha.py +++ b/atom/model_ops/attention_mha.py @@ -77,6 +77,10 @@ def __init__( self.rotary_emb = rotary_emb self.q_norm = q_norm self.k_norm = k_norm + # Set by the attention backend's build_kv_cache_tensor when KV cache is + # allocated in flash layout [num_blocks, block_size, num_kv_heads, head_dim] + # for aiter triton unified_attention. AiterBackend keeps this False. + self.use_flash_layout = False # for plugin mode(vllm), the query quant is disabled for now if is_vllm(): @@ -127,7 +131,12 @@ def rope_cache(self, q, k, v, qkv, position, fwd_ctx: ForwardContext): k_scale = kv_cache_data[f"layer_{self.layer_num}"].k_scale v_scale = kv_cache_data[f"layer_{self.layer_num}"].v_scale - use_triton_attn = self.sliding_window != -1 or self.head_dim != 128 + # MTP MHA must go through triton/gluon; aiter ASM non-persistent path may have some unexpected behavior. + use_triton_attn = ( + self.sliding_window != -1 + or self.head_dim != 128 + or self.num_heads == self.num_kv_heads + ) self.use_triton_attn = use_triton_attn if ( @@ -229,7 +238,7 @@ def rope_cache(self, q, k, v, qkv, position, fwd_ctx: ForwardContext): k_scale, v_scale, self.rotary_emb.is_neox_style, - flash_layout=False, + flash_layout=self.use_flash_layout, apply_scale=self.kv_cache_dtype.startswith("fp8"), offs=None, q_out=q, @@ -345,6 +354,13 @@ def _gather_prefix_and_concat_kv( block_size = k_cache.shape[1] block_tables = attn_metadata.block_tables + per_token_quant = ( + self.kv_cache_dtype.startswith("fp8") + and k_scale is not None + and v_scale is not None + and k_scale.numel() > 1 + and v_scale.numel() > 1 + ) cp_mha_gather_cache( key_cache=k_cache_gather, value_cache=v_cache_gather, @@ -359,6 +375,7 @@ def _gather_prefix_and_concat_kv( dequant=self.kv_cache_dtype.startswith("fp8"), kv_cache_layout="SHUFFLE" if use_shuffle else "NHD", total_tokens=total_tokens, + per_token_quant=per_token_quant, ) return q, k_full, v_full, k_cache, v_cache, k_scale, v_scale @@ -372,71 +389,101 @@ def paged_attention_triton( o = torch.empty_like(q) num_seqs = attn_metadata.context_lens.shape[0] - _, num_q_heads_total, head_size = q.shape - num_blocks, num_kv_heads, _, block_size, _ = k_cache.shape - # assume all query have same length - query_group_size = attn_metadata.max_seqlen_q * ( - num_q_heads_total // num_kv_heads - ) - assert num_q_heads_total % num_kv_heads == 0 - max_context_partition_num = get_recommended_splits(num_seqs, num_kv_heads) + if self.use_flash_layout: + sliding_window = ( + (self.sliding_window - 1, 0) if self.sliding_window > 0 else (-1, -1) + ) - context_partition_size = 256 - if self.sliding_window > 0: - max_context_partition_num = 1 - context_partition_size = 128 + # KV cache is already in flash layout (4D), allocated by + # TritonMHAMetadataBuilder.build_kv_cache_tensor. + nkv = k_cache.shape[2] + descale_shape = (num_seqs, nkv) - # Output buffers (same as Triton) - intermediate_shape = ( - num_seqs, - num_kv_heads, - max_context_partition_num, - query_group_size, - ) - exp_sums = torch.empty(intermediate_shape, dtype=torch.float32, device=q.device) - max_logits = torch.empty( - intermediate_shape, dtype=torch.float32, device=q.device - ) - temporary_output = torch.empty( - *intermediate_shape, - head_size, - dtype=q.dtype, - device=q.device, - ) + unified_attention( + q, + k_cache, + v_cache, + o, + cu_seqlens_q=attn_metadata.cu_seqlens_q, + seqused_k=attn_metadata.context_lens, + max_seqlen_q=attn_metadata.max_seqlen_q, + max_seqlen_k=attn_metadata.max_seqlen_k, + softmax_scale=self.scale, + causal=True, + alibi_slopes=None, + window_size=sliding_window, + block_table=attn_metadata.block_tables, + softcap=0, + q_descale=None, + k_descale=self.kv_scale.expand(descale_shape), + v_descale=self.kv_scale.expand(descale_shape), + sinks=self.sinks, + ) + else: + _, num_q_heads_total, head_size = q.shape + num_blocks, num_kv_heads, _, block_size, _ = k_cache.shape + query_group_size = attn_metadata.max_seqlen_q * ( + num_q_heads_total // num_kv_heads + ) + assert num_q_heads_total % num_kv_heads == 0 - if k_scale is not None and k_scale.numel() > 1: - k_scale = k_scale.unsqueeze(-1) - v_scale = v_scale.unsqueeze(-1) + max_context_partition_num = get_recommended_splits(num_seqs, num_kv_heads) - compute_type = ( - torch.bfloat16 - if self.kv_cache_dtype == "bf16" # or per_tensor - else aiter.dtypes.fp8 - ) - torch.ops.aiter.pa_decode_gluon( - o, - q, - k_cache, - v_cache, - attn_metadata.context_lens, - attn_metadata.block_tables, - self.scale, - attn_metadata.max_seqlen_q, - max_context_partition_num, - context_partition_size, - compute_type, - None, # q_scale - None if self.kv_cache_dtype == "bf16" else k_scale, - None if self.kv_cache_dtype == "bf16" else v_scale, - exp_sums=exp_sums, - max_logits=max_logits, - temporary_output=temporary_output, - alibi_slopes=None, - sinks=self.sinks, - sliding_window=self.sliding_window, - ps=True, - ) + context_partition_size = 256 + if self.sliding_window > 0: + max_context_partition_num = 1 + context_partition_size = 128 + + intermediate_shape = ( + num_seqs, + num_kv_heads, + max_context_partition_num, + query_group_size, + ) + exp_sums = torch.empty( + intermediate_shape, dtype=torch.float32, device=q.device + ) + max_logits = torch.empty( + intermediate_shape, dtype=torch.float32, device=q.device + ) + temporary_output = torch.empty( + *intermediate_shape, + head_size, + dtype=q.dtype, + device=q.device, + ) + + if k_scale is not None and k_scale.numel() > 1: + k_scale = k_scale.unsqueeze(-1) + v_scale = v_scale.unsqueeze(-1) + + compute_type = ( + torch.bfloat16 if self.kv_cache_dtype == "bf16" else aiter.dtypes.fp8 + ) + torch.ops.aiter.pa_decode_gluon( + o, + q, + k_cache, + v_cache, + attn_metadata.context_lens, + attn_metadata.block_tables, + self.scale, + attn_metadata.max_seqlen_q, + max_context_partition_num, + context_partition_size, + compute_type, + None, # q_scale + None if self.kv_cache_dtype == "bf16" else k_scale, + None if self.kv_cache_dtype == "bf16" else v_scale, + exp_sums=exp_sums, + max_logits=max_logits, + temporary_output=temporary_output, + alibi_slopes=None, + sinks=self.sinks, + sliding_window=self.sliding_window, + ps=True, + ) return o @@ -501,9 +548,7 @@ def prefill_attention( # variable lenth attention use key value as input attn_metadata = fwd_ctx.attn_metadata sliding_window = ( - (self.sliding_window, 0, 0) - if self.sliding_window is not None - else (-1, -1, 0) + (self.sliding_window, 0, 0) if self.sliding_window > 0 else (-1, -1, 0) ) o = aiter.flash_attn_varlen_func( q, @@ -540,19 +585,32 @@ def prefill_attention_triton( # value: [num_blocks, 1, num_kv_heads, head_size] attn_metadata = fwd_ctx.attn_metadata - block_tables = attn_metadata.block_tables o = torch.empty_like(q) - descale_shape = (attn_metadata.cu_seqlens_q.shape[0] - 1, k.shape[1]) + num_seqs = attn_metadata.cu_seqlens_q.shape[0] - 1 + descale_shape = (num_seqs, k.shape[1]) sliding_window = ( - (self.sliding_window - 1, 0) - if self.sliding_window is not None - else (-1, -1) + (self.sliding_window - 1, 0) if self.sliding_window > 0 else (-1, -1) ) + + # `block_tables` is always populated by TritonMHAMetadataBuilder. + # For pure prefill (no cached tokens) it is the fake table built in + # prepare_prefill that maps seq i to token indices + # [cu_seqlens_k[i], ..., cu_seqlens_k[i+1]-1], paired with raw K/V + # treated as kv_cache with block_size=1. + if attn_metadata.has_cached: + k_for_attn = k_cache + v_for_attn = v_cache + else: + # k: [total_tokens, num_kv_heads, head_size] + # -> [total_tokens, 1, num_kv_heads, head_size] + k_for_attn = k.unsqueeze(1) + v_for_attn = v.unsqueeze(1) + unified_attention( q, - k_cache, - v_cache, + k_for_attn, + v_for_attn, o, cu_seqlens_q=attn_metadata.cu_seqlens_q, seqused_k=attn_metadata.context_lens, @@ -562,7 +620,7 @@ def prefill_attention_triton( causal=True, alibi_slopes=None, window_size=sliding_window, - block_table=block_tables, + block_table=attn_metadata.block_tables, softcap=0, q_descale=None, k_descale=self.kv_scale.expand(descale_shape), @@ -577,9 +635,11 @@ def dispatch_backend(self, fwd_ctx: ForwardContext): ctx = fwd_ctx.context if ctx.is_prefill: + if self.use_flash_layout: + return self.prefill_attention_triton return self.prefill_attention else: - if self.use_triton_attn: + if self.use_triton_attn or self.use_flash_layout: return self.paged_attention_triton else: # Only use pa persistent when block_size == 1024 diff --git a/atom/model_ops/attention_mla.py b/atom/model_ops/attention_mla.py index b2a8aa67d1..c92c231c0d 100644 --- a/atom/model_ops/attention_mla.py +++ b/atom/model_ops/attention_mla.py @@ -484,6 +484,7 @@ def _forward_prefill_mla( attn_metadata.block_tables, attn_metadata.cu_seqlens_k, NUM_TOPK_TOKENS=self.topk_indices_buffer.shape[1], + PAGE_SIZE=get_current_atom_config().kv_cache_block_size, ) paged_cu_seqlens_q = attn_metadata.sparse_cu_seqlens_q paged_kv_indptr = attn_metadata.sparse_kv_indptr @@ -546,64 +547,124 @@ def _forward_decode( device=q.device, ) - kv_buffer = kv_c_and_k_pe_cache.unsqueeze(2) - paged_kv_indptr = attn_metadata.kv_indptr - paged_kv_indices = attn_metadata.kv_indices - if self.topk_indices_buffer is not None: - paged_kv_indptr = attn_metadata.sparse_kv_indptr - paged_kv_indices = triton_convert_req_index_to_global_index( - attn_metadata.cu_seqlens_q, - attn_metadata.kv_indptr, - paged_kv_indptr, - attn_metadata.kv_indices, - self.topk_indices_buffer[:B], - NUM_TOPK_TOKENS=self.topk_indices_buffer.shape[1], + if hasattr(attn_metadata, "triton_block_table"): + from aiter.ops.triton.attention.mla_decode import decode_attention_fwd + + k_buffer = kv_c_and_k_pe_cache.unsqueeze(2) + v_buffer = k_buffer[..., : self.kv_lora_rank] + page_size = k_buffer.shape[1] + + q_for_triton = ( + q.to(torch.bfloat16) + if q.dtype.is_floating_point and q.element_size() == 1 + else q ) - # q_scale = kv_scale = None - # if self.kv_cache_dtype.startswith("fp8"): - # q = q.to(dtypes.fp8) - # q_scale = kv_scale = self.one_scale - - dp_size = get_dp_group().world_size - use_persistent_mode = not (dp_size > 1) - - if not use_persistent_mode: - # DP : disable persistent mode to avoid overflow - work_meta_data = None - work_indptr = None - work_info_set = None - reduce_indptr = None - reduce_final_map = None - reduce_partial_map = None + # Use pre-built dense block_table from prepare_decode() + decode_attention_fwd( + q_for_triton, + k_buffer, + v_buffer, + o, + attn_metadata.triton_lse, + attn_metadata.triton_block_table, + attn_metadata.context_lens, + attn_metadata.triton_attn_logits, + 4, # num_kv_splits + self.scale, + page_size, + k_scale=self._k_scale, + v_scale=self._k_scale, + ) else: - work_meta_data = attn_metadata.work_meta_data - work_indptr = attn_metadata.work_indptr - work_info_set = attn_metadata.work_info_set - reduce_indptr = attn_metadata.reduce_indptr - reduce_final_map = attn_metadata.reduce_final_map - reduce_partial_map = attn_metadata.reduce_partial_map - - mla_decode_fwd( - q, - kv_buffer.view(-1, 1, 1, q.shape[-1]), - o, - attn_metadata.cu_seqlens_q, - paged_kv_indptr, - paged_kv_indices, - attn_metadata.kv_last_page_lens, - attn_metadata.max_seqlen_q, - num_kv_splits=16, - sm_scale=self.scale, - work_meta_data=work_meta_data, - work_indptr=work_indptr, - work_info_set=work_info_set, - reduce_indptr=reduce_indptr, - reduce_final_map=reduce_final_map, - reduce_partial_map=reduce_partial_map, - q_scale=self._q_scale, - kv_scale=self._k_scale, - ) + kv_buffer = kv_c_and_k_pe_cache.unsqueeze(2) + paged_cu_seqlens_q = attn_metadata.cu_seqlens_q + paged_kv_indptr = attn_metadata.kv_indptr + paged_kv_indices = attn_metadata.kv_indices + paged_kv_last_page_lens = attn_metadata.kv_last_page_lens + max_q_len = attn_metadata.max_seqlen_q + if self.topk_indices_buffer is not None: + if attn_metadata.max_seqlen_q > 1: + # MTP verify: per-token layout with max_q_len=1. + # Persistent metadata is per-token (from _set_mla_persistent_worker_buffers_sparse_mtp). + paged_cu_seqlens_q = attn_metadata.sparse_cu_seqlens_q + paged_kv_indptr = attn_metadata.sparse_kv_indptr + paged_kv_last_page_lens = attn_metadata.sparse_kv_last_page_lens + # Gather physical page indices from kv_indices using topk positions. + # block_tables contains large-block IDs (block_ratio > 1) that + # need expansion; kv_indices already has per-token page indices. + paged_kv_indices = triton_gather_kv_indices_sparse( + paged_kv_indptr, + attn_metadata.token_to_seq_idxs, + self.topk_indices_buffer[:B], + attn_metadata.kv_indices, + attn_metadata.kv_indptr, + NUM_TOPK_TOKENS=self.topk_indices_buffer.shape[1], + ) + max_q_len = 1 + else: + paged_kv_indptr = attn_metadata.sparse_kv_indptr + paged_kv_indices = triton_convert_req_index_to_global_index( + attn_metadata.cu_seqlens_q, + attn_metadata.kv_indptr, + paged_kv_indptr, + attn_metadata.kv_indices, + self.topk_indices_buffer[:B], + NUM_TOPK_TOKENS=self.topk_indices_buffer.shape[1], + ) + + dp_size = get_dp_group().world_size + use_persistent_mode = not (dp_size > 1) + + # Sparse layers in MTP verify use separate persistent metadata + # (per-token, max_seqlen_qo=1) while dense layers use normal metadata + # (max_seqlen_qo=2). + is_sparse_mtp = ( + self.topk_indices_buffer is not None and attn_metadata.max_seqlen_q > 1 + ) + + if not use_persistent_mode: + work_meta_data = None + work_indptr = None + work_info_set = None + reduce_indptr = None + reduce_final_map = None + reduce_partial_map = None + elif is_sparse_mtp: + work_meta_data = attn_metadata.sparse_mtp_work_meta_data + work_indptr = attn_metadata.sparse_mtp_work_indptr + work_info_set = attn_metadata.sparse_mtp_work_info_set + reduce_indptr = attn_metadata.sparse_mtp_reduce_indptr + reduce_final_map = attn_metadata.sparse_mtp_reduce_final_map + reduce_partial_map = attn_metadata.sparse_mtp_reduce_partial_map + else: + work_meta_data = attn_metadata.work_meta_data + work_indptr = attn_metadata.work_indptr + work_info_set = attn_metadata.work_info_set + reduce_indptr = attn_metadata.reduce_indptr + reduce_final_map = attn_metadata.reduce_final_map + reduce_partial_map = attn_metadata.reduce_partial_map + + mla_decode_fwd( + q, + kv_buffer.view(-1, 1, 1, q.shape[-1]), + o, + paged_cu_seqlens_q, + paged_kv_indptr, + paged_kv_indices, + paged_kv_last_page_lens, + max_q_len, + num_kv_splits=16, + sm_scale=self.scale, + work_meta_data=work_meta_data, + work_indptr=work_indptr, + work_info_set=work_info_set, + reduce_indptr=reduce_indptr, + reduce_final_map=reduce_final_map, + reduce_partial_map=reduce_partial_map, + q_scale=self._q_scale, + kv_scale=self._k_scale, + ) if self.head_repeat_factor > 1: o = o[:, :: self.head_repeat_factor, :].contiguous() @@ -637,12 +698,6 @@ def forward_impl_server_mode( kv_cache = kv_cache_data[f"layer_{self.layer_num}"].k_cache if context.is_prefill and not use_prefill_mla: - use_prefix_cache = ( - attn_metadata.has_cached - and not is_rocm_aiter_fp4bmm_enabled() - and self.qk_nope_head_dim == self.v_head_dim - ) - prefill_q = self.q_proj(q, x_scale=q_scale).view( -1, self.num_heads, self.qk_head_dim ) @@ -659,7 +714,7 @@ def forward_impl_server_mode( scale=self._k_scale, ) - if use_prefix_cache: + if attn_metadata.has_cached: # k_full/v_full are used for attention compute; gather_kv_b_proj reads # fp8 from cache and dequantizes internally, so output must be model dtype k_full = torch.empty( @@ -675,7 +730,7 @@ def forward_impl_server_mode( ( attn_metadata.total_kv, self.num_heads, - self.qk_nope_head_dim, + self.v_head_dim, ), device=q.device, dtype=self.dtype, @@ -691,7 +746,9 @@ def forward_impl_server_mode( self.kv_b_proj.weight_scale, k_full, v_full, - weight_preshuffle=True, + weight_preshuffle=getattr( + self.kv_b_proj.weight, "is_shuffled", False + ), ) output = flash_attn_varlen_func( q=prefill_q, @@ -878,7 +935,7 @@ def triton_convert_req_index_to_global_index( kv_indices_c = kv_indices.contiguous() token_indices_c = token_indices.contiguous() page_kv_indptr_c = page_kv_indptr.contiguous() - # TODO: not support mtp + # NOTE: MTP (max_seqlen_q > 1) uses triton_convert_req_index_to_global_index_dsa_prefill instead new_kv_indices = torch.empty_like(kv_indices) # Strides in elements @@ -916,7 +973,7 @@ def _convert_req_index_to_global_index_dsa_prefill_kernel( out_kv_indices, # int32 # shapes (compile-time where possible) NUM_TOPK_TOKENS: tl.constexpr, - BLOCK_SIZE: tl.constexpr, + PAGE_SIZE: tl.constexpr, BLOCK_N: tl.constexpr, # tile width along columns # strides (in elements) ti_stride0: tl.int64, # topk_indices stride 0 @@ -941,14 +998,19 @@ def _convert_req_index_to_global_index_dsa_prefill_kernel( ) # int32 pre_seqlens_q = tl.load(cu_seqlens_q + req_id) + seq_token_idx = indice - pre_seqlens_q + block_id = seq_token_idx // PAGE_SIZE + inblock_offset = seq_token_idx % PAGE_SIZE + # Guard block_table access store_mask = (col_id < kv_len) & (col_id < NUM_TOPK_TOKENS) valid_mask = store_mask & (indice >= 0) - out_val = tl.load( - block_table + req_id * bt_stride0 + (indice - pre_seqlens_q) * bt_stride1, + physical_block = tl.load( + block_table + req_id * bt_stride0 + block_id * bt_stride1, mask=valid_mask, other=-1, ) + out_val = tl.where(valid_mask, physical_block * PAGE_SIZE + inblock_offset, -1) # Store results out_ptr_ij = out_kv_indices + kv_start + col_id @@ -967,7 +1029,7 @@ def triton_convert_req_index_to_global_index_dsa_prefill( block_table: torch.Tensor, # int32 [num_req, max_num_blocks_per_req] cu_seqlens_q: torch.Tensor, # int32 [num_tokens + 1] # dsa_kv_indices: torch.Tensor, # int32 [total_kv_seqlen] -->>> output for this kernel - PAGE_SIZE: int = 1, # page_block_size = 1 for now + PAGE_SIZE: int = 1, NUM_TOPK_TOKENS: int = 2048, BLOCK_N: int = 1024, # tile width along columns ): @@ -1010,3 +1072,85 @@ def triton_convert_req_index_to_global_index_dsa_prefill( bt_stride1, ) return new_kv_indices + + +@triton.jit +def _gather_kv_indices_sparse_kernel( + sparse_kv_indptr, + token_to_seq_idxs, + topk_indices, + kv_indices, + kv_indptr, + out_kv_indices, + NUM_TOPK_TOKENS: tl.constexpr, + BLOCK_N: tl.constexpr, + ti_stride0: tl.int64, + ti_stride1: tl.constexpr, +): + token_id = tl.program_id(0) + tile_id = tl.program_id(1) + col_id = tile_id * BLOCK_N + tl.arange(0, BLOCK_N) + + req_id = tl.load(token_to_seq_idxs + token_id) + + out_start = tl.load(sparse_kv_indptr + token_id) + out_end = tl.load(sparse_kv_indptr + token_id + 1) + kv_len = out_end - out_start + + pos = tl.load(topk_indices + token_id * ti_stride0 + col_id * ti_stride1) + + kv_base = tl.load(kv_indptr + req_id) + kv_end = tl.load(kv_indptr + req_id + 1) + req_kv_len = kv_end - kv_base + + store_mask = (col_id < kv_len) & (col_id < NUM_TOPK_TOKENS) + valid_mask = store_mask & (pos >= 0) & (pos < req_kv_len) + + out_val = tl.load( + kv_indices + kv_base + pos, + mask=valid_mask, + other=0, + ) + + tl.store( + out_kv_indices + out_start + col_id, + out_val, + mask=store_mask, + ) + + +def triton_gather_kv_indices_sparse( + sparse_kv_indptr: torch.Tensor, + token_to_seq_idxs: torch.Tensor, + topk_indices: torch.Tensor, + kv_indices: torch.Tensor, + kv_indptr: torch.Tensor, + NUM_TOPK_TOKENS: int = 2048, + BLOCK_N: int = 1024, +): + assert topk_indices.shape[1] == NUM_TOPK_TOKENS + assert NUM_TOPK_TOKENS % BLOCK_N == 0 + + num_tokens = token_to_seq_idxs.shape[0] + tiles_per_row = NUM_TOPK_TOKENS // BLOCK_N + + out = torch.empty( + num_tokens * NUM_TOPK_TOKENS, dtype=torch.int32, device=topk_indices.device + ) + + ti_stride0, ti_stride1 = topk_indices.stride() + grid = (num_tokens, tiles_per_row) + + _gather_kv_indices_sparse_kernel[grid]( + sparse_kv_indptr, + token_to_seq_idxs, + topk_indices, + kv_indices, + kv_indptr, + out, + NUM_TOPK_TOKENS, + BLOCK_N, + ti_stride0, + ti_stride1, + ) + return out diff --git a/atom/model_ops/attentions/aiter_attention.py b/atom/model_ops/attentions/aiter_attention.py index c6defd2fd6..ba2bec9687 100644 --- a/atom/model_ops/attentions/aiter_attention.py +++ b/atom/model_ops/attentions/aiter_attention.py @@ -16,7 +16,9 @@ import atom.model_ops as ops from atom.model_ops.paged_attention import PagedAttention from atom.model_ops.attention_mha import PagedAttentionImpl -from atom.plugin.sglang.attention_backend.radix_attention import RadixAttention +from atom.plugin.sglang.attention_backend.full_attention.radix_attention import ( + RadixAttention, +) from atom.utils.forward_context import AttentionMetaData, Context from .backends import AttentionBackend, CommonAttentionBuilder @@ -71,9 +73,7 @@ def __init__( CommonAttentionBuilder.__init__(self, model_runner) config = model_runner.config hf_config = config.hf_config - self.num_attention_heads = ( - hf_config.num_attention_heads // get_tp_group().world_size - ) + # `self.num_attention_heads` set by CommonAttentionBuilder.__init__. # For speculative decode (MTP), max_qlen = num_speculative_tokens + 1 if ( config.speculative_config is not None diff --git a/atom/model_ops/attentions/aiter_mla.py b/atom/model_ops/attentions/aiter_mla.py index 1f1aa0f86d..03b44b2fb3 100644 --- a/atom/model_ops/attentions/aiter_mla.py +++ b/atom/model_ops/attentions/aiter_mla.py @@ -12,7 +12,6 @@ get_mla_metadata_info_v1, get_mla_metadata_v1, ) -from aiter.dist.parallel_state import get_tp_group from atom.model_engine.scheduler import ScheduledBatch from atom.model_ops.attention_mla import _MLA_MIN_HEADS, MLAAttention from atom.plugin.attention import ( @@ -22,7 +21,6 @@ from atom.plugin.prepare import is_plugin_mode from atom.utils import CpuGpuBuffer from atom.utils.block_convert import ( - block_table_convert_triton, kv_indices_generate_triton, ) from atom.utils.forward_context import AttentionMetaData, Context @@ -61,9 +59,7 @@ def __init__(self, model_runner): CommonAttentionBuilder.__init__(self, model_runner) config = model_runner.config hf_config = config.hf_config - self.num_attention_heads = ( - hf_config.num_attention_heads // get_tp_group().world_size - ) + # `self.num_attention_heads` set by CommonAttentionBuilder.__init__. self.padded_num_attention_heads = max(self.num_attention_heads, _MLA_MIN_HEADS) self.is_sparse = model_runner.is_deepseek_v32 self.index_topk = hf_config.index_topk if self.is_sparse else -1 @@ -143,7 +139,54 @@ def __init__(self, model_runner): mla_metadata["sparse_kv_last_page_lens"].np[:] = 1 mla_metadata["sparse_kv_last_page_lens"].copy_to_gpu() + if self.is_sparse and max_seqlen_qo > 1: + # Allocate a second set of persistent work buffers for sparse MTP + # per-token layout: max_bs*max_seqlen_qo virtual seqs, each q_len=1. + smt_max_bs = self.max_bs * max_seqlen_qo + ( + (smt_wmd_size, smt_wmd_type), + (smt_wi_size, smt_wi_type), + (smt_wis_size, smt_wis_type), + (smt_ri_size, smt_ri_type), + (smt_rfm_size, smt_rfm_type), + (smt_rpm_size, smt_rpm_type), + ) = get_mla_metadata_info_v1( + smt_max_bs, + 1, # max_seqlen_qo=1 for per-token + self.padded_num_attention_heads, + self.dtype_q, + self.dtype_kv, + is_sparse=True, + fast_mode=True, + ) + mla_metadata["sparse_mtp_work_meta_data"] = torch.empty( + smt_wmd_size, dtype=smt_wmd_type, device=self.device + ) + mla_metadata["sparse_mtp_work_indptr"] = torch.empty( + smt_wi_size, dtype=smt_wi_type, device=self.device + ) + mla_metadata["sparse_mtp_work_info_set"] = torch.empty( + smt_wis_size, dtype=smt_wis_type, device=self.device + ) + mla_metadata["sparse_mtp_reduce_indptr"] = torch.empty( + smt_ri_size, dtype=smt_ri_type, device=self.device + ) + mla_metadata["sparse_mtp_reduce_final_map"] = torch.empty( + smt_rfm_size, dtype=smt_rfm_type, device=self.device + ) + mla_metadata["sparse_mtp_reduce_partial_map"] = torch.empty( + smt_rpm_size, dtype=smt_rpm_type, device=self.device + ) + self.model_runner.forward_vars.update(mla_metadata) + + if self.is_sparse: + self._token_to_seq_idxs_gpu = torch.zeros( + self.max_num_batched_tokens, + dtype=torch.int32, + device=self.device, + ) + # Per-ubatch buffers for CUDAGraph TBO if config.enable_tbo: self._allocate_ubatch_buffers( @@ -213,7 +256,6 @@ def _allocate_ubatch_buffers( **i32_kwargs, ) var[f"{p}cu_seqlens_q"] = CpuGpuBuffer(ub_max_bs + 1, **i32_kwargs) - # cu_seqlens_q for decode is constant [0, q, 2q, ..., bs*q] var[f"{p}cu_seqlens_q"].cpu.copy_( torch.arange( 0, @@ -267,6 +309,60 @@ def prep_stream(self): # return self.model_runner.tokenID_processor.async_copy_stream return self.model_runner.async_execute_stream + def _set_mla_persistent_worker_buffers_sparse_mtp( + self, + num_tokens: int, + ): + """Compute persistent metadata for sparse MTP per-token layout. + + B = batch_size * max_seqlen_q tokens are treated as B independent + virtual sequences each with q_len=1. cu_seqlens_q = [0,1,...,B], + kv_indptr = per-token sparse_kv_indptr, kv_last_page_lens = all 1s. + + Uses separate sparse_mtp_* buffers so dense layers can keep + their own persistent metadata (max_seqlen_qo=2) intact. + """ + var = self.model_runner.forward_vars + split_params = { + "kv_granularity": max(self.block_size, 16), + "max_seqlen_qo": 1, + "uni_seqlen_qo": 1, + "fast_mode": 1, + "max_split_per_batch": 16, + } + work_meta_data = var["sparse_mtp_work_meta_data"] + work_info_set = var["sparse_mtp_work_info_set"] + work_indptr = var["sparse_mtp_work_indptr"] + reduce_indptr = var["sparse_mtp_reduce_indptr"] + reduce_final_map = var["sparse_mtp_reduce_final_map"] + reduce_partial_map = var["sparse_mtp_reduce_partial_map"] + get_mla_metadata_v1( + var["sparse_cu_seqlens_q"].gpu[: num_tokens + 1], + var["sparse_kv_indptr"].gpu[: num_tokens + 1], + var["sparse_kv_last_page_lens"].gpu[:num_tokens], + self.padded_num_attention_heads, + 1, # nhead_kv + True, + work_meta_data, + work_info_set, + work_indptr, + reduce_indptr, + reduce_final_map, + reduce_partial_map, + page_size=self.block_size, + dtype_q=self.dtype_q, + dtype_kv=self.dtype_kv, + **split_params, + ) + return { + "sparse_mtp_work_meta_data": work_meta_data, + "sparse_mtp_work_info_set": work_info_set, + "sparse_mtp_work_indptr": work_indptr, + "sparse_mtp_reduce_indptr": reduce_indptr, + "sparse_mtp_reduce_final_map": reduce_final_map, + "sparse_mtp_reduce_partial_map": reduce_partial_map, + } + def set_mla_persistent_worker_buffers( self, bs: int, @@ -288,14 +384,19 @@ def set_mla_persistent_worker_buffers( reduce_indptr = var["reduce_indptr"] reduce_final_map = var["reduce_final_map"] reduce_partial_map = var["reduce_partial_map"] + # Dense layers use kv_indptr (full KV lengths per seq). + # sparse_kv_indptr is per-token in MTP mode and must NOT be + # indexed with [:bs+1] here — that misinterprets the per-token + # cumsum as per-seq, producing wrong KV lengths and OOB metadata. + kv_indptr_for_metadata = ( + var["sparse_kv_indptr"].gpu[: bs + 1] + if self.is_sparse and max_q_len == 1 + else var["kv_indptr"].gpu[: bs + 1] + ) if only_update: decode_update_mla_metadata_v1( var["cu_seqlens_q"].gpu[: bs + 1], - ( - var["sparse_kv_indptr"].gpu[: bs + 1] - if self.is_sparse - else var["kv_indptr"].gpu[: bs + 1] - ), + kv_indptr_for_metadata, var["kv_last_page_lens"].gpu[:bs], self.padded_num_attention_heads, 1, # nhead_kv, @@ -316,11 +417,7 @@ def set_mla_persistent_worker_buffers( else: get_mla_metadata_v1( var["cu_seqlens_q"].gpu[: bs + 1], - ( - var["sparse_kv_indptr"].gpu[: bs + 1] - if self.is_sparse - else var["kv_indptr"].gpu[: bs + 1] - ), + kv_indptr_for_metadata, var["kv_last_page_lens"].gpu[:bs], self.padded_num_attention_heads, 1, # nhead_kv, @@ -356,7 +453,14 @@ def prepare_mtp_decode( var = self.model_runner.forward_vars kv_indptr = var["kv_indptr"].gpu[: bs + 1] if self.is_sparse: - assert False, "TODO: MTP decode is not supported for sparse attention yet" + # Update dense kv_indptr (needed for kv_indices generation and slot_mapping) + kv_indptr += var["cu_seqlens_q"].gpu[: bs + 1] + # Recompute sparse_kv_indptr: per-seq sparse count = min(dense_kv_count, index_topk) + sparse_kv_indptr = var["sparse_kv_indptr"].gpu[: bs + 1] + kv_counts = kv_indptr[1 : bs + 1] - kv_indptr[:bs] + sparse_counts = torch.clamp(kv_counts, max=self.index_topk) + sparse_kv_indptr[0] = 0 + sparse_kv_indptr[1 : bs + 1] = torch.cumsum(sparse_counts, dim=0) else: assert self.block_size == 1 kv_indptr += var["cu_seqlens_q"].gpu[: bs + 1] @@ -368,9 +472,12 @@ def prepare_mtp_decode( self.block_ratio, max_seqlen_k, ) - return self.set_mla_persistent_worker_buffers( + result = self.set_mla_persistent_worker_buffers( bs, max_seqlen_q, only_update, num_reject_tokens ) + if self.is_sparse: + result["sparse_kv_indptr"] = sparse_kv_indptr + return result def compute_block_bytes(self) -> int: """MLA per-block bytes: single 576-dim packed tensor per layer @@ -378,7 +485,7 @@ def compute_block_bytes(self) -> int: V cache or kv_scale). DeepSeek-V3.2 sparse variant adds an indexer cache contribution - (`hf_config.num_hidden_layers` rows, aligned-to-16 fp8). + for every bound layer, including draft/MTP layers. """ runner = self.model_runner config = runner.config @@ -391,7 +498,7 @@ def compute_block_bytes(self) -> int: index_dim = hf_config.index_head_dim + 4 aligned_index_dim = ((index_dim + 15) // 16) * 16 block_bytes += ( - hf_config.num_hidden_layers + total_num_layers * runner.block_size * aligned_index_dim * dtypes.fp8.itemsize @@ -429,7 +536,7 @@ def allocate_kv_cache_tensors( aligned = ((index_dim + 15) // 16) * 16 out["aligned_index_dim"] = aligned out["index_cache"] = torch.zeros( - hf_config.num_hidden_layers, + total_num_layers, runner.num_physical_kvcache_blocks, runner.physical_block_size, aligned, @@ -489,22 +596,14 @@ def prepare_prefill(self, batch: ScheduledBatch): if attn_metadata.block_tables is None: self.prepare_block_tables(batch) attn_metadata.block_tables = var["block_tables"].copy_to_gpu(bs) - if self.block_ratio > 1: - block_table_convert_triton( - var["block_tables"].gpu[:bs], - var["block_tables_converted"].gpu[:bs], - var["context_lens"].gpu[:bs], - self.block_ratio, - ) - attn_metadata.block_tables = var["block_tables_converted"].gpu[:bs] counts = var["cu_seqlens_q"].np[1 : bs + 1] - var["cu_seqlens_q"].np[:bs] if attn_metadata.has_cached: # Full context (cached + new): use cu_seqlens_k for indexer var["cu_seqlen_ks"].np[:sum_scheduled_tokens] = np.repeat( - var["cu_seqlens_k"].np[:-1], counts + var["cu_seqlens_k"].np[:bs], counts ) var["cu_seqlen_ke"].np[:sum_scheduled_tokens] = np.repeat( - var["cu_seqlens_k"].np[1:], counts + var["cu_seqlens_k"].np[1 : bs + 1], counts ) else: var["cu_seqlen_ke"].np[:sum_scheduled_tokens] = ( @@ -624,6 +723,7 @@ def prepare_decode(self, batch: ScheduledBatch, bs: int): var["slot_mapping"].np[:sum_scheduled_tokens] = slot_mapping var["positions"].np[:sum_scheduled_tokens] = positions var["context_lens"].np[:scheduled_bs] = context_lens + var["context_lens"].np[scheduled_bs:bs] = 0 if any(batch.is_first_decode_without_local_prefill): num_blocks_per_seq = [ @@ -666,15 +766,39 @@ def prepare_decode(self, batch: ScheduledBatch, bs: int): if self.is_sparse: index_topk = self.index_topk - sparse_context_lens = np.clip(var["context_lens"].np[:bs], None, index_topk) - var["sparse_kv_indptr"].np[1 : bs + 1] = np.cumsum( - sparse_context_lens, dtype=np.int32 - ) - var["sparse_kv_indptr"].np[scheduled_bs : bs + 1] = var[ - "sparse_kv_indptr" - ].np[scheduled_bs] - vars_used.append(("sparse_kv_indptr", bs + 1)) - metadata_deps.add("sparse_kv_indptr") + if max_seqlen_q > 1: + # MTP verify: per-token sparse metadata + # Each token at offset j in seq s sees (context_lens[s] - max_seqlen_q + j + 1) KV entries + per_token_kv_lens = ( + np.repeat(context_lens[:scheduled_bs], max_seqlen_q) + - max_seqlen_q + + np.tile( + np.arange(1, max_seqlen_q + 1, dtype=np.int32), scheduled_bs + ) + ) + sparse_per_token_lens = np.clip(per_token_kv_lens, 0, index_topk) + var["sparse_kv_indptr"].np[1 : sum_scheduled_tokens + 1] = np.cumsum( + sparse_per_token_lens, dtype=np.int32 + ) + sum_tokens = bs * max_seqlen_q + var["sparse_kv_indptr"].np[ + sum_scheduled_tokens + 1 : sum_tokens + 1 + ] = var["sparse_kv_indptr"].np[sum_scheduled_tokens] + vars_used.append(("sparse_kv_indptr", sum_tokens + 1)) + vars_used.append(("sparse_cu_seqlens_q", sum_tokens + 1)) + metadata_deps.add("sparse_kv_indptr") + else: + sparse_context_lens = np.clip( + var["context_lens"].np[:bs], None, index_topk + ) + var["sparse_kv_indptr"].np[1 : bs + 1] = np.cumsum( + sparse_context_lens, dtype=np.int32 + ) + var["sparse_kv_indptr"].np[scheduled_bs : bs + 1] = var[ + "sparse_kv_indptr" + ].np[scheduled_bs] + vars_used.append(("sparse_kv_indptr", bs + 1)) + metadata_deps.add("sparse_kv_indptr") prep_stream = self.prep_stream vars_for_metadata = [(el, num) for el, num in vars_used if el in metadata_deps] @@ -698,21 +822,22 @@ def prepare_decode(self, batch: ScheduledBatch, bs: int): max_seqlen_k, ) + is_sparse_mtp = self.is_sparse and max_seqlen_q > 1 # metadata copies on main_stream positions = var["positions"].copy_to_gpu(sum_scheduled_tokens) ctx.update({el: var[el].copy_to_gpu(num) for el, num in vars_for_metadata}) - ctx_mla_ps = self.set_mla_persistent_worker_buffers(bs, max_seqlen_q) + + if is_sparse_mtp: + sum_tokens = bs * max_seqlen_q + ctx_mla_ps = self.set_mla_persistent_worker_buffers(bs, max_seqlen_q) + ctx_mla_ps_sparse = self._set_mla_persistent_worker_buffers_sparse_mtp( + sum_tokens + ) + else: + ctx_mla_ps = self.set_mla_persistent_worker_buffers(bs, max_seqlen_q) + ctx_mla_ps_sparse = None ctx.update(ctx_mla_ps) current_stream.wait_stream(prep_stream) - # if self.block_ratio > 1: - # if "block_tables" in ctx: - # block_table_convert_triton( - # var["block_tables"].gpu[:bs], - # var["block_tables_converted"].gpu[:bs], - # var["context_lens"].gpu[:bs], - # self.block_ratio, - # ) - # ctx["block_tables_converted"] = var["block_tables_converted"].gpu[:bs] attn_metadata = AttentionMetaData( dropout_p=dropout_p, max_seqlen_q=max_seqlen_q, @@ -721,6 +846,24 @@ def prepare_decode(self, batch: ScheduledBatch, bs: int): ) attn_metadata.dtype_q = self.dtype_q + if ctx_mla_ps_sparse is not None: + for k, v in ctx_mla_ps_sparse.items(): + setattr(attn_metadata, k, v) + + if is_sparse_mtp: + sum_tokens = bs * max_seqlen_q + attn_metadata.sparse_cu_seqlens_q = var["sparse_cu_seqlens_q"].gpu[ + : sum_tokens + 1 + ] + attn_metadata.sparse_kv_last_page_lens = var[ + "sparse_kv_last_page_lens" + ].gpu[:sum_tokens] + self._token_to_seq_idxs_gpu[:sum_scheduled_tokens] = torch.arange( + scheduled_bs, dtype=torch.int32, device=self.device + ).repeat_interleave(max_seqlen_q) + self._token_to_seq_idxs_gpu[sum_scheduled_tokens:sum_tokens] = 0 + attn_metadata.token_to_seq_idxs = self._token_to_seq_idxs_gpu[:sum_tokens] + # Use bs (graph_bs) >= 2 instead of scheduled_bs >= 2 to avoid accuracy issue: if self.model_runner.config.enable_tbo_decode and bs >= 2: self._prepare_ubatch_decode( @@ -884,9 +1027,19 @@ def build_for_cudagraph_capture(self, bs: int) -> AttentionMetaData: var = self.model_runner.forward_vars sparse_kv_indptr = var["sparse_kv_indptr"].gpu if self.is_sparse else None max_q_len = var["mtp_k"] + 1 if "mtp_k" in var else 1 - ctx_mla_ps = self.set_mla_persistent_worker_buffers(bs, max_q_len) + sum_tokens = bs * max_q_len + is_sparse_mtp = self.is_sparse and max_q_len > 1 + if is_sparse_mtp: + # Two sets: normal for dense layers, sparse_mtp for sparse layers + ctx_mla_ps = self.set_mla_persistent_worker_buffers(bs, max_q_len) + ctx_mla_ps_sparse = self._set_mla_persistent_worker_buffers_sparse_mtp( + sum_tokens + ) + else: + ctx_mla_ps = self.set_mla_persistent_worker_buffers(bs, max_q_len) + ctx_mla_ps_sparse = None attn_matadata = AttentionMetaData( - slot_mapping=var["slot_mapping"].gpu[: bs * max_q_len], + slot_mapping=var["slot_mapping"].gpu[:sum_tokens], context_lens=var["context_lens"].gpu[:bs], block_tables=var["block_tables"].gpu[:bs], max_seqlen_q=max_q_len, @@ -895,15 +1048,27 @@ def build_for_cudagraph_capture(self, bs: int) -> AttentionMetaData: kv_indices=var["kv_indices"].gpu, kv_last_page_lens=var["kv_last_page_lens"].gpu[:bs], sparse_kv_indptr=sparse_kv_indptr, - block_tables_converted=( - var["block_tables_converted"].gpu[:bs] - if "block_tables_converted" in var - else None - ), **ctx_mla_ps, ) attn_matadata.dtype_q = self.dtype_q - positions = var["positions"].copy_to_gpu(bs * max_q_len) + if ctx_mla_ps_sparse is not None: + for k, v in ctx_mla_ps_sparse.items(): + setattr(attn_matadata, k, v) + if is_sparse_mtp: + attn_matadata.sparse_cu_seqlens_q = var["sparse_cu_seqlens_q"].gpu[ + : sum_tokens + 1 + ] + attn_matadata.sparse_kv_indptr = var["sparse_kv_indptr"].gpu[ + : sum_tokens + 1 + ] + attn_matadata.sparse_kv_last_page_lens = var[ + "sparse_kv_last_page_lens" + ].gpu[:sum_tokens] + self._token_to_seq_idxs_gpu[:sum_tokens] = torch.arange( + bs, dtype=torch.int32, device=self.device + ).repeat_interleave(max_q_len) + attn_matadata.token_to_seq_idxs = self._token_to_seq_idxs_gpu[:sum_tokens] + positions = var["positions"].copy_to_gpu(sum_tokens) context = Context( positions=positions, is_prefill=False, batch_size=bs, graph_bs=bs ) diff --git a/atom/model_ops/attentions/backends.py b/atom/model_ops/attentions/backends.py index ca1b77ddc7..9c7963c551 100644 --- a/atom/model_ops/attentions/backends.py +++ b/atom/model_ops/attentions/backends.py @@ -6,6 +6,7 @@ from typing import Any, Dict, Generic, Optional, Type, TypeVar import torch +from aiter.dist.parallel_state import get_tp_group from atom.model_engine.scheduler import ScheduledBatch from atom.model_ops.attention_mla import MLAModules from atom.utils import CpuGpuBuffer @@ -208,6 +209,13 @@ def __init__(self, model_runner): self.max_num_blocks_per_seq = ( config.max_model_len + self.block_size - 1 ) // self.block_size + # Per-rank attention head count. eagle.propose's mid-step path reads + # this to gate the `do_attn_metadata_update` branch. Subclasses that + # need a kernel-minimum-padded count set `self.padded_num_attention_heads` + # separately (it does NOT replace this attribute). + self.num_attention_heads = ( + hf_config.num_attention_heads // get_tp_group().world_size + ) i64_kwargs = {"dtype": torch.int64, "device": self.device} i32_kwargs = {"dtype": torch.int32, "device": self.device} diff --git a/atom/model_ops/attentions/deepseek_v4_attn.py b/atom/model_ops/attentions/deepseek_v4_attn.py index 80a668748e..6fd548369c 100644 --- a/atom/model_ops/attentions/deepseek_v4_attn.py +++ b/atom/model_ops/attentions/deepseek_v4_attn.py @@ -47,8 +47,13 @@ AttentionMetadataBuilder, CommonAttentionBuilder, ) +from atom.model_ops.v4_kernels import write_v4_paged_decode_indices from atom.utils import CpuGpuBuffer -from atom.utils.forward_context import AttentionMetaData, Context +from atom.utils.forward_context import ( + AttentionMetaData, + Context, + get_forward_context, +) # --------------------------------------------------------------------------- # Typed metadata surface for V4. The base AttentionMetaData class is shared @@ -89,8 +94,15 @@ class AttentionMetaData_DSV4(AttentionMetaData): # ----- CPU mirrors (avoid GPU→CPU `.item()` / `.tolist()` syncs) ----- state_slot_mapping_cpu: Optional[Any] = None """[bs] np.int32 — per-seq state cache slot id (host copy).""" - start_pos_per_seq_cpu: Optional[Any] = None - """[bs] np.int64 — absolute start position per seq.""" + n_committed_csa_per_seq_cpu: Optional[Any] = None + """[bs] np.int32 — `ctx_len // 4` (CSA committed K per seq). Built once + in `_attach_v4_per_fwd_meta` from `var["context_lens"].np`; consumed by + `_attach_v4_paged_decode_meta`, `_build_paged_prefill_meta`, and + `_build_v4_indexer_meta` (indptr cumsums). Single source of truth so + those callers don't each re-read context_lens + divide.""" + n_committed_hca_per_seq_cpu: Optional[Any] = None + """[bs] np.int32 — `ctx_len // 128` (HCA committed compress entries per + seq). Same lifecycle as `n_committed_csa_per_seq_cpu`.""" # ----- Per-seq GPU scalars (single-source-of-truth, shared by kernels) ----- state_slot_mapping: Optional[torch.Tensor] = None @@ -110,9 +122,11 @@ class AttentionMetaData_DSV4(AttentionMetaData): kernels skip on `bid < 0`. All other per-token quantities resolved as `per_seq_data[batch_id_per_token[t]]` — no [T] aliases of seq data.""" swa_write_indices: Optional[torch.Tensor] = None - """[padded_T] int64 GPU — src row id into per-fwd KV for swa_write. - `[0:num_write]` = real (last `win` tokens per seq); - `[num_write:padded_T]` = -1 sentinel. `None` for warmup / empty fwd.""" + """[W] int64 GPU — src row id into per-fwd KV for swa_write. + `[0:num_write]` = real (last `win` tokens per seq); trailing entries + `[num_write:W]` = -1 sentinel (only present on decode/MTP CG paths + where `W = padded_bs * (1 + max_spec_steps)`; prefill is eager and + uses `W = num_write` exactly, no padding). `None` for warmup / empty fwd.""" compress_plans: Optional[Dict[int, Any]] = None """dict[ratio:int -> CompressPlan] — packed plan tensors per compress_ratio (4=CSA, 128=HCA).""" @@ -141,24 +155,25 @@ class AttentionMetaData_DSV4(AttentionMetaData): (= `win + n_committed_hca[bid]`). Padded tail = last value.""" swa_pages: int = 0 """Boundary in `unified_kv`: index < swa_pages → SWA region; index >= - swa_pages → compress region. Equal to `max_per_req_cache_slots * win`.""" + swa_pages → compress region. Equal to + `max_per_req_cache_slots * win_with_spec` (per-slot SWA region holds + `win + mtp_k` ring entries; reduces to `win` when MTP is off).""" # ----- Indexer / sparse-layout side metadata ----- indexer_meta: Optional[Dict[str, Any]] = None """dict — `Indexer.forward_batched` per-fwd GPU tensors. Notable keys: cu_committed_gpu [bs+1] int32 per-seq committed cumsum - compress_slot_mapping_gpu [num_compress] int64 FP8 cache write slots seq_base_per_token_gpu [T] int32 prefill subtract base (also aliased as cu_starts_gpu for fp8_mqa_logits) cu_ends_gpu [T] int32 per-token end offset for fp8_mqa_logits (causal cap) - decode_logits_gpu [max_bs, ...] pre-allocated decode buffer - decode_topk_indices_gpu [max_bs, index_topk] int32 decode topk - output buffer (raw seq-local - with -1 sentinels in tail) total_committed int sum of n_committed_csa_per_seq + Note: decode logits / topk-indices scratch are allocated per-fwd inside + `Indexer._score_topk_decode` (write-once, no CPU mirror, CG-stable via + the captured graph's private memory pool). + The indexer's downstream contract: `_score_topk_*` returns RAW seq-local `[T, index_topk] int32` with kernel-native -1 in tail cols (cells past the per-token visibility cap). `csa_translate_pack` consumes this layout @@ -226,71 +241,15 @@ def _segment_indices( total = int(lens.sum()) if total == 0: return ( - np.empty(0, dtype=np.int64), - np.empty(0, dtype=np.int64), + np.empty(0, dtype=np.int32), + np.empty(0, dtype=np.int32), ) - token_seq_ids = np.repeat(seq_ids.astype(np.int64), lens) - cum = np.concatenate(([0], np.cumsum(lens.astype(np.int64))[:-1])) - local_pos = np.arange(total, dtype=np.int64) - np.repeat(cum, lens) + token_seq_ids = np.repeat(seq_ids.astype(np.int32), lens) + cum = np.concatenate(([0], np.cumsum(lens, dtype=np.int32)[:-1])) + local_pos = np.arange(total, dtype=np.int32) - np.repeat(cum, lens) return token_seq_ids, local_pos -def _build_window_topk_batched( - positions: torch.Tensor, # [total_tokens] int (abs token positions) - start_pos_per_token: torch.Tensor, # [total_tokens] int (each token's seq start_pos) - window_size: int, - out: Optional[ - torch.Tensor - ] = None, # [total_tokens, window_size] int32 — CG-stable buffer -) -> torch.Tensor: # [total_tokens, window_size] int32 - """Per-token sliding-window topk indices for the whole batch. - - Three-branch semantics: - - sp == 0 (fresh prefill): matrix entries = abs positions in the window - [pos-win+1, pos] clamped to [0, pos]; mask future via -1. - - 0 < sp < win-1 (prefix mode): all tokens in the seq share a single - matrix [0..sp, -1, ..., -1] (matches original semantics, including - MTP-N where the same start_pos is reused). - - sp >= win-1 (cyclic mode): cyclic ring offsets starting at sp+1 mod win. - - `out`: when supplied, the final int32 result is written into `out` - (sliced to `[total_tokens, window_size]`) so its storage address is - stable across captures — required for CUDAGraph replay correctness. - Intermediates remain transient (not consumed by the captured graph). - """ - device = positions.device - total = positions.size(0) - arange_w = torch.arange(window_size, device=device, dtype=positions.dtype).view( - 1, window_size - ) - pos_col = positions.view(total, 1) - sp_col = start_pos_per_token.view(total, 1) - neg1 = torch.tensor(-1, device=device, dtype=positions.dtype) - - # Case A: sp == 0 (fresh prefill) — abs positions [pos-win+1, pos] clamped. - case_a = (pos_col - window_size + 1).clamp(min=0) + arange_w - case_a = torch.where(case_a > pos_col, neg1, case_a) - - # Case B: 0 < sp < win-1 (prefix mode) — shared per-seq matrix. - case_b = arange_w.expand(total, window_size).clone() - case_b = torch.where(arange_w > sp_col, neg1, case_b) - - # Case C: sp >= win-1 (cyclic mode) — ring offsets. - sp_mod = sp_col % window_size - case_c = (sp_mod + 1 + arange_w) % window_size - - sp_eq_0 = sp_col == 0 - sp_in_prefix = (sp_col > 0) & (sp_col < window_size - 1) - - res = case_c - res = torch.where(sp_in_prefix.expand_as(res), case_b, res) - res = torch.where(sp_eq_0.expand_as(res), case_a, res) - if out is None: - return res.to(torch.int32) - out[:total].copy_(res, non_blocking=True) - return out[:total] - - class DeepseekV4Backend(AttentionBackend): """Backend selector entry for V4 hybrid attention. @@ -354,13 +313,6 @@ def __init__(self, model_runner): # max-density ratio (1 indexer slot per 4 source tokens). self.max_model_len_idx = model_runner.config.max_model_len // 4 - # Compressor state shape: [coff * ratio, coff * head_dim], fp32. - # CSA: ratio=4, overlap=True -> coff=2 -> [8, 2*head_dim] - # HCA: ratio=128, overlap=False -> coff=1 -> [128, head_dim] - self.csa_main_state_shape = (2 * 4, 2 * self.head_dim) - self.csa_idx_state_shape = (2 * 4, 2 * self.index_head_dim) - self.hca_main_state_shape = (128, self.head_dim) - # Classical KV pool geometry. block_size=128 original tokens means # each V4 block holds k1=128/4=32 CSA entries and k2=128/128=1 HCA # entry per layer (paper §3.6.1). @@ -383,7 +335,40 @@ def __init__(self, model_runner): self.max_spec_steps = ( int(model_runner.drafter.mtp_k) if hasattr(model_runner, "drafter") else 0 ) + + # Compressor state shape: [ring_size, coff * head_dim], fp32. + # ring_size = K_pool + (max_spec_steps + 1), where K_pool = coff * ratio. + # The extra (max_spec_steps + 1) slots prevent reject K/V written to the + # ring in round R from being borrow-read by the round-R+1 re-commit of + # the same boundary (slot index = pos % ring_size). With ring_size = + # K_pool the slot of pos P+K_pool aliases the slot of pos P, so a + # rejected K_{P+K_pool} written in R overwrites the K_P that R+1's + # commit pool window [P..P+K_pool-1] still needs. Bumping by one + # verify window's worth of positions decouples the two. + # CSA: ratio=4, overlap=True → K_pool=8; spec ring_size=8 + (mtp_k+1) + # HCA: ratio=128, overlap=False → K_pool=128; spec ring_size=128+(mtp_k+1) + # Non-spec (max_spec_steps=0) → ring_size = K_pool + 1 (effectively + # equivalent to old K_pool layout for correctness; one extra slot is + # the algebraic minimum and trivial in memory). + ring_extra = self.max_spec_steps + 1 + self.csa_main_state_shape = (2 * 4 + ring_extra, 2 * self.head_dim) + self.csa_idx_state_shape = (2 * 4 + ring_extra, 2 * self.index_head_dim) + self.hca_main_state_shape = (128 + ring_extra, self.head_dim) self.max_decode_tokens = self.max_bs * (1 + self.max_spec_steps) + # SWA ring-buffer slots per req. Distinct from `window_size`: + # * `window_size` = SWA attention window = topk count per token + # (each query attends to W consecutive K/V positions). + # * `win_with_spec` = `window_size + max_spec_steps` = ring-buffer + # slot count per req. With MTP-k the per-fwd writes the verified + # token + k draft tokens at positions [p_0..p_k]; if the cache + # were only sized W, draft slots `p_(i+1)..p_k` would alias into + # [p_0-W+1..p_0] and the verified query at `p_0` would read + # future tokens (silent correctness bug). MTP off → max_spec_steps + # == 0 → win_with_spec == window_size, identical bytes layout. + # Used as: SWA `unified_kv` per-slot stride, `swa_kv` ring-buffer dim, + # `swa_write` modulo, and the ring-index modulo `cs` in the V4 + # paged-decode index-write kernel. + self.win_with_spec = self.window_size + self.max_spec_steps # Worst-case HCA per-token committed compress count # (= max_model_len // 128 for V4-Pro = 8192 at 1M context). self.max_committed_hca = model_runner.config.max_model_len // 128 @@ -395,6 +380,10 @@ def __init__(self, model_runner): # `torch.as_tensor(arr)` allocations. self._alloc_v4_metadata_buffers() + @property + def prep_stream(self): + return self.model_runner.async_execute_stream + # ------------------------------------------------------------------ # # AttentionMetadataBuilder hooks (per-request cache abstraction). # # ------------------------------------------------------------------ # @@ -413,8 +402,9 @@ def compute_per_req_cache_bytes(self) -> int: csa_main = self._numel(self.csa_main_state_shape) * 2 * elem_state csa_idx = self._numel(self.csa_idx_state_shape) * 2 * elem_state hca_main = self._numel(self.hca_main_state_shape) * 2 * elem_state - # SWA window per layer. - swa_per_layer = self.window_size * self.head_dim * elem_swa + # SWA window per layer. Cache holds `win_with_spec = win + mtp_k` + # slots so MTP draft tokens don't alias verified-token slots. + swa_per_layer = self.win_with_spec * self.head_dim * elem_swa return ( len(self.csa_layers) * (csa_main + csa_idx) + len(self.hca_layers) * hca_main @@ -422,7 +412,9 @@ def compute_per_req_cache_bytes(self) -> int: ) def slots_per_req(self) -> int: - # No spec decoding lookahead in V4 (pre-PR3-main). + # State cache is one slot per req regardless of MTP. The MTP draft + # lookahead bytes are absorbed into per-slot SWA size via + # `win_with_spec` (above), not into a slots_per_req multiplier. return 1 def compute_block_bytes(self) -> int: @@ -482,10 +474,12 @@ def allocate_per_req_cache(self, num_slots: int) -> dict[str, object]: Per-layer `unified_kv` layout (decode-time paged_decode kernel reads a single base ptr; offsets `[0, swa_pages)` are SWA, `[swa_pages, ..)` - are compress): - Dense layer: [num_slots*win, head_dim] BF16 - CSA layer: [num_slots*win + NB*k1, head_dim] BF16 - HCA layer: [num_slots*win + NB*k2, head_dim] BF16 + are compress). Per-slot SWA region is `win_with_spec = win + mtp_k` + (extra slack so MTP draft tokens don't overwrite the verified token's + ring slot mid-fwd): + Dense layer: [num_slots*win_with_spec, head_dim] BF16 + CSA layer: [num_slots*win_with_spec + NB*k1, head_dim] BF16 + HCA layer: [num_slots*win_with_spec + NB*k2, head_dim] BF16 `build_kv_cache_tensor` slices per-layer views to bind into `attn.swa_kv` (SWA portion, reshape to [num_slots, win, head_dim]) @@ -508,7 +502,7 @@ def allocate_per_req_cache(self, num_slots: int) -> dict[str, object]: num_blocks = self.model_runner.num_physical_kvcache_blocks n_csa = len(self.csa_layers) n_hca = len(self.hca_layers) - swa_pages = num_slots * self.window_size + swa_pages = num_slots * self.win_with_spec head_dim = self.head_dim dtype = self._swa_dtype @@ -578,18 +572,20 @@ def build_kv_cache_tensor(self, layer_id: int, module): runner = self.model_runner num_slots = self.model_runner.max_per_req_cache_slots - swa_pages = num_slots * self.window_size + swa_pages = num_slots * self.win_with_spec if isinstance(module, _V4Attention): # Bind both: # - `attn.unified_kv`: the full per-layer pool (paged_decode reads). - # - `attn.swa_kv`: a [num_slots, win, head_dim] view onto the - # SWA prefix (compatibility with `swa_write` and the existing - # prefill/non-CG path that index_select's by state slot). + # - `attn.swa_kv`: a [num_slots, win_with_spec, head_dim] view + # onto the SWA prefix. Per-slot dim is `win + mtp_k` so MTP + # draft tokens have their own ring slots; `swa_write` modulo + # and `paged_decode` per-row case_c modulo both use this + # dim (= `swa_kv.shape[1]`). unified = runner.v4_unified_kv[module.layer_id] module.unified_kv = unified module.swa_kv = unified[:swa_pages].view( - num_slots, self.window_size, self.head_dim + num_slots, self.win_with_spec, self.head_dim ) return None @@ -618,22 +614,38 @@ def build_kv_cache_tensor(self, layer_id: int, module): is_indexer_inner = "indexer" in parts ratio = module.compress_ratio - # Per-kind shared compress output buffer (CUDAGraph: stable - # data pointer + fixed shape across captures of the same kind). - # Read from forward_vars (allocated in _alloc_v4_metadata_buffers). - compress_out_buffers = self.model_runner.forward_vars if is_indexer_inner: assert ratio == 4, "Indexer-inner Compressor only on CSA layers" pos = self.layer_id_to_csa_pos[layer_id_from_prefix] module.kv_state = runner.v4_csa_idx_kv_state[pos] module.score_state = runner.v4_csa_idx_score_state[pos] # Inner compressor writes target the SAME storage as the - # outer Indexer.kv_cache (csa_idx_kv). Same 3D shape — write - # via slot_mapping is shape-agnostic (flat indexing), but we - # keep [NB, k1_csa, aligned_dim] for symmetry with the read - # binding above. - module.kv_cache = runner.v4_csa_idx_kv[pos] - module.compress_out = compress_out_buffers["v4_csa_idx_compress_out"] + # outer Indexer.kv_cache (csa_idx_kv). Same [NB, k1_csa, + # aligned_dim] FP8 shape — `Compressor.forward` resolves + # slot via block_table+ci internally (no flat slot_mapping + # needed; matches CSA Main's path). + idx_kv = runner.v4_csa_idx_kv[pos] + module.kv_cache = idx_kv + # FP8 quant path: bind a strided fp32 view of the per-block + # scale region. Layout per block: [k1*head_dim FP8 region] + # then [k1 fp32 scale region] then padding (cache_kernels.cu + # :1209-1239). Strides expressed in fp32 elements. + nb, k1, aligned_dim = idx_kv.shape + head_dim = self.index_head_dim + assert ( + k1 * aligned_dim + ) % 4 == 0, f"per-block bytes ({k1 * aligned_dim}) must be 4-aligned" + block_fp32_stride = (k1 * aligned_dim) // 4 + scale_fp32_offset = (k1 * head_dim) // 4 + module.cache_scale = ( + idx_kv.view(torch.float32) + .view(-1) + .as_strided( + size=(nb, k1), + stride=(block_fp32_stride, 1), + storage_offset=scale_fp32_offset, + ) + ) elif ratio == 4: pos = self.layer_id_to_csa_pos[layer_id_from_prefix] module.kv_state = runner.v4_csa_main_kv_state[pos] @@ -647,7 +659,6 @@ def build_kv_cache_tensor(self, layer_id: int, module): module.kv_cache = unified[swa_pages:].view( num_blocks, self.k1_csa, self.head_dim ) - module.compress_out = compress_out_buffers["v4_csa_main_compress_out"] elif ratio == 128: pos = self.layer_id_to_hca_pos[layer_id_from_prefix] module.kv_state = runner.v4_hca_main_kv_state[pos] @@ -657,7 +668,6 @@ def build_kv_cache_tensor(self, layer_id: int, module): module.kv_cache = unified[swa_pages:].view( num_blocks, self.k2_hca, self.head_dim ) - module.compress_out = compress_out_buffers["v4_hca_main_compress_out"] else: raise ValueError( f"Unknown V4 compress_ratio={ratio} on Compressor at " @@ -676,8 +686,6 @@ def build_kv_cache_tensor(self, layer_id: int, module): def _attach_v4_indexer_meta( self, attn_metadata: AttentionMetaData_DSV4, - cu_seqlens_q_np, - start_pos_per_seq, scheduled_bs: int, total_tokens: int, positions_gpu=None, @@ -685,40 +693,26 @@ def _attach_v4_indexer_meta( """Build and attach the CSA Indexer per-fwd GPU metadata. Hoists per-CSA-layer H2D calls (batch_id_per_token / cu_committed / - n_committed / seq_base_per_token / cu_ends / compress_slot_mapping) - into a single per-fwd build. compress_plans[4] (CSA, ratio=4) carries - the compress rows that the indexer-FP8 path needs to derive its - write-side slot_mapping. None for warmup or empty fwd; - `_build_v4_indexer_meta` handles both. + n_committed / seq_base_per_token / cu_ends) into a single per-fwd + build. None for warmup or empty fwd; `_build_v4_indexer_meta` + handles both. """ - import numpy as np - - csa_compress_plan_cpu = None - plans = getattr(attn_metadata, "compress_plans", None) or {} - if 4 in plans: - csa_compress_plan_cpu = plans[4].compress_plan_cpu attn_metadata.indexer_meta = self._build_v4_indexer_meta( attn_metadata=attn_metadata, - cu_seqlens_q_np=cu_seqlens_q_np, - start_pos_per_seq_np=np.asarray(start_pos_per_seq, dtype=np.int64), positions_gpu=positions_gpu, scheduled_bs=scheduled_bs, total_tokens=total_tokens, device=self.device, - csa_compress_plan_cpu=csa_compress_plan_cpu, ) def _build_v4_indexer_meta( self, *, attn_metadata: AttentionMetaData_DSV4, - cu_seqlens_q_np, - start_pos_per_seq_np, positions_gpu, scheduled_bs: int, total_tokens: int, device, - csa_compress_plan_cpu, ): """Build per-fwd GPU index tensors consumed by `Indexer.forward_batched`. @@ -726,40 +720,29 @@ def _build_v4_indexer_meta( inline H2D path) or when CSA / Indexer is not on the model. CSA ratio is fixed at 4; we always build under that assumption. - Reuses two shared GPU tensors set by `_attach_v4_per_fwd_meta` - (which MUST be called before this helper): - - `attn_metadata.batch_id_per_token` [padded_T] int32 + Reads pre-computed `attn_metadata.n_committed_csa_per_seq_cpu` + (set by `_attach_v4_per_fwd_meta`, which MUST run first) for the + per-seq committed count and cumsums it on CPU. + + Reuses two shared GPU tensors also set by `_attach_v4_per_fwd_meta`: + - `attn_metadata.batch_id_per_token` [padded_T] int64 - `attn_metadata.n_committed_csa_per_seq` [bs] int32 - These were previously re-staged here as `v4_indexer_batch_id_per_token` - (int64) and `v4_indexer_n_committed_per_seq` (int64) — one extra H2D - each per fwd. Now we read int32 → cast to int64 on GPU. - - The FP8-cache write side (`indexer_k_quant_and_cache`) needs a flat - `compress_slot_mapping` int64 tensor — one entry per row in - `csa_compress_plan_cpu` mapping the compress entry to a global slot - in the `[n_csa, NB, k1, aligned_dim=144]` cache pool (layer-major, - FP8 + 4-byte fp32 scale per row, 16B-aligned). Computed here host-side - because the plan rows + block_tables_cpu are both already on host. + + The FP8 indexer K-cache write happens inside `fused_compress_attn` + (the unified Indexer-inner Compressor path) via the same block_tables + that CSA Main uses; no separate slot_mapping is built here. """ - from atom.models.deepseek_v4 import _V4_BLOCK_SIZE # Caller contract: scheduled_bs >= 1, total_tokens >= 1 (same # invariants as `_attach_v4_per_fwd_meta` — guaranteed by every # prepare_*/CG-capture path). bs = scheduled_bs - ratio = 4 # CSA - k_per_block = _V4_BLOCK_SIZE // ratio # 32 - cu_seqlens_q_arr = cu_seqlens_q_np[: bs + 1].astype(np.int64) - token_num_per_seq = (cu_seqlens_q_arr[1:] - cu_seqlens_q_arr[:bs]).astype( - np.int64 - ) - start_pos_per_seq = start_pos_per_seq_np[:bs].astype(np.int64) - end_pos_per_seq = start_pos_per_seq + token_num_per_seq - n_committed_per_seq = end_pos_per_seq // ratio # host copy for cumsum + ratio = 4 # CSA — also referenced by `visible_end_gpu` below + n_committed_per_seq = attn_metadata.n_committed_csa_per_seq_cpu[:bs] cu_committed_cpu = np.concatenate( [ - np.zeros(1, dtype=np.int64), - np.cumsum(n_committed_per_seq, dtype=np.int64), + np.zeros(1, dtype=np.int32), + np.cumsum(n_committed_per_seq, dtype=np.int32), ] ) # Empty-batch guard: when no seq has committed K yet @@ -779,14 +762,6 @@ def _build_v4_indexer_meta( cu_committed_cpu[-1] = max(int(cu_committed_cpu[-1]), 1) total_committed = int(cu_committed_cpu[-1]) - # FP8 write-side slot_mapping (independent of `total_committed` — - # written even when there's nothing to read because num_compress can - # be > 0 on the same fwd that crosses a compress boundary for the - # FIRST committed entry of a fresh seq). - compress_slot_mapping_gpu = self._build_indexer_compress_slot_mapping( - csa_compress_plan_cpu, scheduled_bs, k_per_block, ratio - ) - # batch_id_per_token + n_committed_csa: reuse the shared GPU # tensors set in `_attach_v4_per_fwd_meta` (which MUST run before # this helper — see prepare_decode/prefill ordering). int64 @@ -799,9 +774,7 @@ def _build_v4_indexer_meta( # cu_committed_gpu is consumed both as `cu_starts/cu_ends` for the # fp8_mqa_logits per-token range AND as `cu_seq_lens` for the # cp_gather_indexer_k_quant_cache call (per-seq cumulative committed K). - cu_committed_gpu = self._stage( - "v4_indexer_cu_committed", cu_committed_cpu.astype(np.int32) - ) + cu_committed_gpu = self._stage("v4_indexer_cu_committed", cu_committed_cpu) # Layer-invariant GPU derivations (each was previously rebuilt ~30x # per fwd inside the per-CSA-layer body). @@ -822,23 +795,10 @@ def _build_v4_indexer_meta( seq_base_per_token_gpu + visible_end_gpu ) # [total_tokens] int32 — fp8_mqa_logits per-token end offset - # Pre-allocated decode-path buffers (full [max_bs, ...] views). Decode - # helper slices each to [:total_tokens]. Returned full because - # `_build_v4_indexer_meta` is called for both prefill and decode - # batches; prefill's `total_tokens` may exceed `max_bs` so builder - # can't pre-slice. Prefill path doesn't read these. - decode_logits_gpu = self.model_runner.forward_vars[ - "v4_indexer_decode_logits" - ].gpu # [max_bs, max_model_len_idx] fp32 - decode_topk_indices_gpu = self.model_runner.forward_vars[ - "v4_indexer_decode_topk_indices" - ].gpu # [max_bs, index_topk] int32 - return { "total_committed": total_committed, "cu_committed_gpu": cu_committed_gpu, "n_committed_per_seq_gpu": n_committed_per_seq_gpu, # int32, [bs] - "compress_slot_mapping_gpu": compress_slot_mapping_gpu, "batch_id_per_token_gpu": batch_id_per_token_gpu, # int64, [total_tokens] # Prefill-only fields below — decode never consults them. NOT # in pre-allocated buffers (per-fwd derived); CG capture path @@ -847,95 +807,168 @@ def _build_v4_indexer_meta( "seq_base_per_token_gpu": seq_base_per_token_gpu, "cu_starts_gpu": seq_base_per_token_gpu, # alias for fp8_mqa_logits "cu_ends_gpu": cu_ends_gpu, - "decode_logits_gpu": decode_logits_gpu, - "decode_topk_indices_gpu": decode_topk_indices_gpu, } - def _build_indexer_compress_slot_mapping( + def prepare_mtp_decode( self, - csa_compress_plan_cpu, - scheduled_bs: int, - k_per_block: int, - ratio: int, + bs: int, + max_seqlen_q: int, + max_seqlen_k: int, + only_update: bool = False, + num_reject_tokens: torch.Tensor = None, ): - """Compute the per-compress-row flat slot in `v4_csa_idx_kv` pool. - - For each row in `csa_compress_plan_cpu` (= `(ragged_id, batch_id, - position, window_len)`): - ci = position // ratio # compress entry idx in seq - block_in_seq = ci // k_per_block - slot_in_block = ci % k_per_block - physical_block = block_tables_cpu[batch_id, block_in_seq] - slot = physical_block * k_per_block + slot_in_block - - Returns None when the plan is empty (no boundary crossed this fwd) — - the caller skips the `indexer_k_quant_and_cache` write entirely. + """Per-draft-step V4 region metadata rebuild for 1-token-per-seq shape. + + Called by EagleProposer.propose at mid-step iters (i < mtp_k - 1). + Eagle has already bumped attn_metadata.context_lens GPU by +1 and + max_seqlen_k by +1 before calling us. We mirror the +1 on CPU (zero + D2H: structural invariant of eagle's loop) and rebuild + v4_kv_indices_{swa,csa,hca}, batch_id_per_token, n_committed_csa_per_seq, + indexer meta, and compress_plans. + + ``only_update`` / ``num_reject_tokens`` are MLA-specific (no V4 + analog — V4 has no incremental-update kernel and the ctx rollback + is already applied at the top of prepare_decode for the verify + shape). Ignored. """ - if csa_compress_plan_cpu is None or csa_compress_plan_cpu.shape[0] == 0: - return None + # `max_per_req_cache_slots` is set inside `model_runner.get_num_blocks`, + # which runs AFTER `warmup_model`. The full per-fwd meta rebuild below + # eventually reads it via `_attach_v4_paged_decode_meta`, so during + # warmup (attr unset) we no-op — warmup discards draft output anyway, + # and the verify-shape attn_metadata from the main forward stays valid + # for the rest of eagle.propose. + if not getattr(self.model_runner, "max_per_req_cache_slots", 0): + return {} + var = self.model_runner.forward_vars - block_tables_np = var["block_tables"].np[:scheduled_bs] - bid = csa_compress_plan_cpu[:, 1] - pos = csa_compress_plan_cpu[:, 2] - ci = pos // ratio - block_in_seq = ci // k_per_block - slot_in_block = ci % k_per_block - physical_block = block_tables_np[bid, block_in_seq] - # int64 — `indexer_k_quant_and_cache` kernel signature is `int64_t*` - # (matches V3.2's `attn_metadata.slot_mapping` dtype). int32 caused - # GPU memory access faults from 2x stride mis-stepping. - slot_mapping_np = physical_block.astype(np.int64) * k_per_block + slot_in_block - return self._stage("v4_indexer_compress_slot_mapping", slot_mapping_np) + attn_metadata = get_forward_context().attn_metadata + + # 1. CPU mirror of eagle's GPU `context_lens[:bs] += 1` bump. Zero + # D2H: we know the offset by construction (+1 per call). + var["context_lens"].np[:bs] += 1 + context_lens_np = var["context_lens"].np[:bs] + + # 2. 1-token-per-seq shape: positions = ctx-1, cu_seqlens_q = arange. + positions_np = (context_lens_np - 1).astype(np.int32) + cu_seqlens_q_np = np.arange(bs + 1, dtype=np.int32) + var["positions"].np[:bs] = positions_np + var["cu_seqlens_q"].np[: bs + 1] = cu_seqlens_q_np + + # 3. H2D staging. context_lens already on GPU (eagle bumped it); + # CPU is now mirrored. + positions_gpu = var["positions"].copy_to_gpu(bs) + var["cu_seqlens_q"].copy_to_gpu(bs + 1) + + # 4. CPU numpy: extend_lens (=1 per seq). + extend_lens_np = np.ones(bs, dtype=np.int32) + + # 5. Rebuild V4 region metadata via existing helpers (numpy + H2D, + # no D2H — `_build_compress_plans` only triggers `.cpu()` when + # given torch tensors; we pass numpy below). + self._attach_v4_per_fwd_meta( + attn_metadata, + cu_seqlens_q_np, + extend_lens_np, # = np.ones(bs) — MTP draft step is 1 token per seq + state_slot_mapping_cpu=attn_metadata.state_slot_mapping_cpu, + scheduled_bs=bs, + total_tokens=bs, + padded_bs=bs, + max_q_len=1, + ) + self._attach_v4_indexer_meta( + attn_metadata, + scheduled_bs=bs, + total_tokens=bs, + positions_gpu=positions_gpu, + ) + + # 6. Compress plans for state-ring write of each new draft token. + attn_metadata.compress_plans = self._build_compress_plans( + extend_lens_np, + context_lens_np, + for_decode_cg=True, + ) + + # All updates done in-place on attn_metadata; eagle's + # `for k, v in workinfos.items(): __dict__[k] = v` loop is a no-op. + return {} def prepare_decode(self, batch: ScheduledBatch, bs: int): """V4-style decode prep: populates positions, cu_seqlens_q, block_tables, and state_slot_mapping. - For PR3-main multi-seq: - - cu_seqlens_q gives V4 forward per-seq token-range slicing. - - block_tables[i] is the per-seq block list for compressor scatter - and sparse_attn gather. - - state_slot_mapping[i] is the per-seq state-cache slot (swa_kv + - Compressor.kv_state row index). Distinct from `slot_mapping` - which is per-token paged-KV-pool index. - - Also publishes CPU mirrors (`v4_*_cpu` attrs) so the V4 forward path - can read per-seq metadata without GPU→CPU `.tolist()`/.item() syncs - (PR-A Phase 2: required to unlock CUDAGraph). + Uses stream overlap (like AiterMLAMetadataBuilder) to hide H2D + latency behind CPU numpy work: basic H2D copies fire on + ``prep_stream`` while ``_build_compress_plans`` runs on the CPU. """ - import numpy as np - var = self.model_runner.forward_vars scheduled_bs = batch.total_seqs_num_decode - # Naming: `_np` suffix marks numpy arrays, `_gpu` marks GPU tensors. - # Both representations of context_lens coexist (CPU mirror feeds the - # plan builder; GPU copy goes into attn_metadata for kernels). context_lens_np = np.asarray(batch.context_lens, dtype=np.int32) max_seqlen_q = batch.num_spec_step + 1 + # MTP: roll back ctx by `num_rejected` so this fwd's positions overwrite + # last fwd's rejected-draft slots (matches aiter_mla.py:701 / + # aiter_attention.py:542). `batch.context_lens` = `seq.num_tokens` + # which the scheduler advances by `mtp_k - num_rejected` placeholders + # per fwd (scheduler.py:789); without this rollback, MTP-k positions + # would skip ahead by `num_rejected` and the rejected slots would + # never be overwritten with the corrected K/V. `num_rejected` is None + # on dummy runs and on the first fwd before any sampler output. + # Bound n_committed_csa/hca via the rolled-back ctx (n_committed_* = + # ctx // 4 / 128 in `_attach_v4_paged_decode_meta`), so block_tables + # truncation isn't needed here — the per-token kv_len already shrinks. + if not batch.is_dummy_run and max_seqlen_q > 1: + num_rejected = self.model_runner.tokenID_processor.num_rejected + if num_rejected is not None: + context_lens_np = context_lens_np - num_rejected.astype(np.int32) positions_np = np.tile( np.arange(max_seqlen_q, dtype=np.int32), scheduled_bs ) + np.repeat(context_lens_np - max_seqlen_q, max_seqlen_q) sum_scheduled_tokens = batch.total_tokens_num_decode + var["positions"].np[:sum_scheduled_tokens] = positions_np - positions = var["positions"].copy_to_gpu(sum_scheduled_tokens) - # cu_seqlens_q for decode: each seq has max_seqlen_q tokens. cu_seqlens_q_np = np.arange( 0, (scheduled_bs + 1) * max_seqlen_q, max_seqlen_q, dtype=np.int32 ) var["cu_seqlens_q"].np[: scheduled_bs + 1] = cu_seqlens_q_np - cu_seqlens_q_gpu = var["cu_seqlens_q"].copy_to_gpu(scheduled_bs + 1) - # PR-A: V4 Compressor state-cache update reads context_lens to - # discriminate fresh prefill vs decode/prefix-cache. Parent's - # prepare_prefill pushes this; decode path must do the same. var["context_lens"].np[:scheduled_bs] = context_lens_np - context_lens_gpu = var["context_lens"].copy_to_gpu(scheduled_bs) - block_tables_gpu = self._populate_block_tables(batch, scheduled_bs) - state_slot_gpu, state_slot_np = self._populate_state_slot_mapping( - batch, scheduled_bs, return_cpu=True + # Inline block_tables CPU fill (H2D deferred to prep_stream). + self.prepare_block_tables(batch) + + state_slot_np = np.asarray( + batch.per_req_cache_groups[:scheduled_bs], dtype=np.int32 + ) + if len(state_slot_np) < scheduled_bs: + state_slot_np = np.zeros(scheduled_bs, dtype=np.int32) + ss_buf = var["v4_meta_state_slot_groups"] + ss_buf.np[:scheduled_bs] = state_slot_np + + # ---- fire H2D on prep_stream ---- + # NB: this runs inside attn_metadata_builder.build(), BEFORE + # set_forward_context() — can't read main_stream from the context yet. + prep_stream = self.prep_stream + current_stream = torch.cuda.current_stream() + prep_stream.wait_stream(current_stream) + with torch.cuda.stream(prep_stream): + positions = var["positions"].copy_to_gpu(sum_scheduled_tokens) + cu_seqlens_q_gpu = var["cu_seqlens_q"].copy_to_gpu(scheduled_bs + 1) + context_lens_gpu = var["context_lens"].copy_to_gpu(scheduled_bs) + block_tables_gpu = var["block_tables"].copy_to_gpu(scheduled_bs) + state_slot_gpu = ss_buf.copy_to_gpu(scheduled_bs) + + # ---- CPU numpy work, overlapped with prep_stream H2D ---- + extend_lens_np = np.full(scheduled_bs, max_seqlen_q, dtype=np.int32) + compress_plans = self._build_compress_plans( + extend_lens_np, + context_lens_np, + for_decode_cg=True, ) + + # ---- sync, build attn_metadata, per-fwd meta ---- + current_stream.wait_stream(prep_stream) + attn_metadata = AttentionMetaData_DSV4( cu_seqlens_q=cu_seqlens_q_gpu, cu_seqlens_k=None, @@ -948,45 +981,18 @@ def prepare_decode(self, batch: ScheduledBatch, bs: int): num_cached_tokens=None, block_tables=block_tables_gpu, context_lens=context_lens_gpu, + is_pure_decode=True, ) - # Attach V4-specific per-seq metadata via dynamic attribute. - # AttentionMetaData is a regular class (not slotted), so this works - # without modifying its schema. attn_metadata.state_slot_mapping = state_slot_gpu - # PR-A Phase 2: CPU mirrors so the V4 forward path can read per-seq - # metadata without `.tolist()` / `.item()` GPU→CPU syncs. Generic - # naming (no `v4_` prefix) — these are not V4-specific concepts and - # other backends with stateful per-request buffers can adopt them. attn_metadata.state_slot_mapping_cpu = state_slot_np - attn_metadata.start_pos_per_seq_cpu = positions_np[ - cu_seqlens_q_np[:scheduled_bs] - ] - # Compress plans (per ratio) for batched fused_compress + update_states. - # Decode batch: extend_lens = max_seqlen_q for all seqs (uniform). - # `context_lens_np` is post-extend (from batch.context_lens, set by - # scheduler after appending this fwd's tokens) — this is what plan - # generation needs as `seq_lens`. Must run BEFORE - # `_attach_v4_indexer_meta` since the indexer consumes - # plan.compress_plan_cpu to derive its FP8 write-side slot_mapping. - extend_lens_np = np.full(scheduled_bs, max_seqlen_q, dtype=np.int32) - attn_metadata.compress_plans = self._build_compress_plans( - extend_lens_np, context_lens_np, positions.device - ) - # CG capture path: model_runner pads cu_seqlens_q to graph_bs. The - # captured kernels still iterate `bs * max_q_len` tokens at replay, - # so per-fwd metadata MUST be sized to that padded T (otherwise - # padded slots read stale buffer entries → GPU memory access fault). - # `bs` parameter is the padded graph_bs (or actual bs in eager). - # NOTE: must run BEFORE `_attach_v4_indexer_meta` because the latter's - # `_build_v4_indexer_meta` consumes the per-fwd shared tensors - # (batch_id_per_token, n_committed_csa_per_seq) staged here. + attn_metadata.compress_plans = compress_plans + padded_bs = int(bs) self._attach_v4_per_fwd_meta( attn_metadata, - positions, cu_seqlens_q_np, - attn_metadata.start_pos_per_seq_cpu, - attn_metadata.state_slot_mapping_cpu, + extend_lens_np, # = np.full(scheduled_bs, max_seqlen_q) for decode + state_slot_np, scheduled_bs, sum_scheduled_tokens, padded_bs=padded_bs, @@ -994,8 +1000,6 @@ def prepare_decode(self, batch: ScheduledBatch, bs: int): ) self._attach_v4_indexer_meta( attn_metadata, - cu_seqlens_q_np, - attn_metadata.start_pos_per_seq_cpu, scheduled_bs, sum_scheduled_tokens, positions_gpu=positions, @@ -1013,14 +1017,17 @@ def prepare_prefill(self, batch: ScheduledBatch): Also publishes CPU mirrors (`v4_*_cpu`) consumed by the V4 forward path to avoid `.item()` / `.tolist()` syncs (PR-A Phase 2). """ - import numpy as np - base_md, positions = super().prepare_prefill(batch) # Promote to V4 typed metadata so V4-specific attribute assignments # below are well-typed. Safe because AttentionMetaData_DSV4 only adds # fields with defaults; the parent dataclass is non-slotted. base_md.__class__ = AttentionMetaData_DSV4 attn_metadata = cast(AttentionMetaData_DSV4, base_md) + # Prefill is by definition not pure decode (fresh-prefill seqs have + # start_pos == 0). Class-promotion via __class__ doesn't run the + # dataclass __init__, so V4-specific defaults aren't applied — set + # explicitly. + attn_metadata.is_pure_decode = False scheduled_bs = batch.total_seqs_num_prefill if attn_metadata.block_tables is None: attn_metadata.block_tables = self._populate_block_tables( @@ -1037,9 +1044,10 @@ def prepare_prefill(self, batch: ScheduledBatch): positions_np = np.asarray(var["positions"].np[:sum_scheduled_tokens]) cu_seqlens_q_np = np.asarray(var["cu_seqlens_q"].np[: scheduled_bs + 1]) attn_metadata.state_slot_mapping_cpu = state_slot_np - attn_metadata.start_pos_per_seq_cpu = positions_np[ - cu_seqlens_q_np[:scheduled_bs] - ] + # `start_pos_per_seq` = position of FIRST token of each seq in this fwd. + # Only consumed by `_build_paged_prefill_meta` below; not stashed on + # attn_metadata (no other reader, no inter-fwd reuse). + start_pos_per_seq_np = positions_np[cu_seqlens_q_np[:scheduled_bs]] # Compress plans (per ratio) for batched fused_compress + update_states. # Prefill batch: extend_lens read from cu_seqlens_q_np. # Must run BEFORE `_attach_v4_indexer_meta` (the indexer consumes @@ -1047,13 +1055,16 @@ def prepare_prefill(self, batch: ScheduledBatch): extend_lens_np = ( cu_seqlens_q_np[1 : scheduled_bs + 1] - cu_seqlens_q_np[:scheduled_bs] ).astype(np.int32) - # context_lens for prefill = positions[seq_start_in] + extend_lens - # (= absolute seq_len incl. this fwd's tokens). - context_lens_np = (attn_metadata.start_pos_per_seq_cpu + extend_lens_np).astype( - np.int32 + # context_lens already populated on host by `super().prepare_prefill` + # (backends.py: `var["context_lens"].np[:bs] = batch.context_lens`). + # Mathematically equals `start_pos + extend_lens` but reading the + # canonical buffer avoids drift if scheduler/batch semantics ever + # change. + context_lens_np = np.asarray( + var["context_lens"].np[:scheduled_bs], dtype=np.int32 ) attn_metadata.compress_plans = self._build_compress_plans( - extend_lens_np, context_lens_np, positions.device + extend_lens_np, context_lens_np, for_decode_cg=False ) # Prefill goes through eager (no CG): defaults make padded_total_tokens # collapse to total_tokens — no padding logic kicks in. Must still run @@ -1061,17 +1072,14 @@ def prepare_prefill(self, batch: ScheduledBatch): # reuse the shared GPU tensors (batch_id_per_token, n_committed_csa). self._attach_v4_per_fwd_meta( attn_metadata, - positions, cu_seqlens_q_np, - attn_metadata.start_pos_per_seq_cpu, + extend_lens_np, # = cu_seqlens_q[1:] - cu_seqlens_q[:bs] attn_metadata.state_slot_mapping_cpu, scheduled_bs, sum_scheduled_tokens, ) self._attach_v4_indexer_meta( attn_metadata, - cu_seqlens_q_np, - attn_metadata.start_pos_per_seq_cpu, scheduled_bs, sum_scheduled_tokens, positions_gpu=positions, @@ -1084,19 +1092,137 @@ def prepare_prefill(self, batch: ScheduledBatch): attn_metadata, positions_np, cu_seqlens_q_np, - attn_metadata.start_pos_per_seq_cpu, + extend_lens_np, + start_pos_per_seq_np, attn_metadata.state_slot_mapping_cpu, scheduled_bs, sum_scheduled_tokens, ) return attn_metadata, positions + def build_ubatch_prefill_metadata( + self, + attn_metadata: AttentionMetaData, + ub_slice, + padded_bs: int, + ) -> AttentionMetaData_DSV4: + """Split prefill AttentionMetaData for V4 TBO micro-batches.""" + from atom.utils.tbo.ubatch_splitting import split_attn_metadata + + ub_attn = split_attn_metadata(attn_metadata, ub_slice, padded_bs) + ub_attn.__class__ = AttentionMetaData_DSV4 + + src = cast(AttentionMetaData_DSV4, attn_metadata) + rs = ub_slice.request_slice + ts = ub_slice.token_slice + ub_num_reqs = rs.stop - rs.start + ub_num_tokens = ts.stop - ts.start + + if src.state_slot_mapping is not None: + ub_attn.state_slot_mapping = src.state_slot_mapping[rs] + if src.state_slot_mapping_cpu is not None: + ub_attn.state_slot_mapping_cpu = src.state_slot_mapping_cpu[rs] + + var = self.model_runner.forward_vars + positions_np = np.asarray(var["positions"].np[ts.start : ts.stop]) + full_cu = var["cu_seqlens_q"].np + ub_cu = np.asarray(full_cu[rs.start : rs.stop + 1], dtype=np.int32).copy() + ub_cu -= ub_cu[0] + + extend_lens_np = (ub_cu[1:] - ub_cu[:ub_num_reqs]).astype(np.int32) + # Slice the full batch's `var["context_lens"]` (populated by + # `super().prepare_prefill`) for this ubatch — mathematically equals + # `start_pos + extend_lens` but reads from the canonical source. + # `.copy()` so the swap-into-front below doesn't alias. + context_lens_np = np.asarray( + var["context_lens"].np[rs.start : rs.stop], dtype=np.int32 + ).copy() + device = src.state_slot_mapping.device + + from atom.model_ops.v4_kernels import make_compress_plans + + if self._unique_compress_ratios_overlap: + ub_attn.compress_plans = make_compress_plans( + np.ascontiguousarray(extend_lens_np, dtype=np.int32), + np.ascontiguousarray(context_lens_np, dtype=np.int32), + self._unique_compress_ratios_overlap, + device, + plan_buffers=None, + decode_capacity_per_ratio=None, + ) + else: + ub_attn.compress_plans = {} + + # Multiple helpers read context_lens and block_tables from + # forward_vars by position [0:scheduled_bs]. For ubatch 1 the + # relevant rows live at [rs.start:rs.stop], not [0:ub_num_reqs]. + # Temporarily place the ubatch's slices at the front so helpers + # see the right values. + bt = var["block_tables"].np + saved_ctx = var["context_lens"].np[:ub_num_reqs].copy() + saved_bt = bt[:ub_num_reqs].copy() + try: + var["context_lens"].np[:ub_num_reqs] = context_lens_np + bt[:ub_num_reqs] = bt[rs.start : rs.stop].copy() + + self._attach_v4_per_fwd_meta( + ub_attn, + ub_cu, + extend_lens_np, # ubatch's per-seq token counts + ub_attn.state_slot_mapping_cpu, + ub_num_reqs, + ub_num_tokens, + ) + + positions_gpu = var["positions"].gpu[ts.start : ts.stop] + self._attach_v4_indexer_meta( + ub_attn, + ub_num_reqs, + ub_num_tokens, + positions_gpu=positions_gpu, + ) + + # start_pos = position of first token of each seq in this ubatch. + ub_start_pos_per_seq_np = positions_np[ub_cu[:ub_num_reqs]] + self._build_paged_prefill_meta( + ub_attn, + positions_np, + ub_cu, + extend_lens_np, + ub_start_pos_per_seq_np, + ub_attn.state_slot_mapping_cpu, + ub_num_reqs, + ub_num_tokens, + ) + finally: + bt[:ub_num_reqs] = saved_bt + var["context_lens"].np[:ub_num_reqs] = saved_ctx + + # Clone all GPU tensors that are views into shared CpuGpuBuffers. + # Without this, building the next ubatch overwrites this ubatch's + # data via the same underlying buffer. + if ub_attn.batch_id_per_token is not None: + ub_attn.batch_id_per_token = ub_attn.batch_id_per_token.clone() + if ub_attn.n_committed_csa_per_seq is not None: + ub_attn.n_committed_csa_per_seq = ub_attn.n_committed_csa_per_seq.clone() + if ub_attn.swa_write_indices is not None: + ub_attn.swa_write_indices = ub_attn.swa_write_indices.clone() + if ub_attn.indexer_meta is not None: + im = ub_attn.indexer_meta + if im.get("cu_committed_gpu") is not None: + im["cu_committed_gpu"] = im["cu_committed_gpu"].clone() + if im.get("batch_id_per_token_gpu") is not None: + im["batch_id_per_token_gpu"] = im["batch_id_per_token_gpu"].clone() + if im.get("n_committed_per_seq_gpu") is not None: + im["n_committed_per_seq_gpu"] = im["n_committed_per_seq_gpu"].clone() + + return ub_attn + def _attach_v4_per_fwd_meta( self, attn_metadata: AttentionMetaData_DSV4, - positions: torch.Tensor, cu_seqlens_q_np, - start_pos_per_seq_cpu, + token_num_per_seq, state_slot_mapping_cpu, scheduled_bs: int, total_tokens: int, @@ -1111,10 +1237,14 @@ def _attach_v4_per_fwd_meta( them once per fwd saves ~64 redundant constructions for V4-Pro. Sets: - - `attn_metadata.swa_write_indices`: [padded_T] int64 row ids - into per-token KV. Real entries are the last `win` tokens per seq; - tail (entries beyond the per-fwd `num_write`) is sentinel-filled - with -1 so a fixed grid `padded_T` is always safe (CUDAGraph). + - `attn_metadata.swa_write_indices`: [W] int64 row ids into + per-token KV. Real entries are the last `win` tokens per seq. + For decode/MTP (CG-captured): `W = padded_bs * (1 + max_spec_steps)`, + trailing entries `[num_write:W]` get -1 sentinel so the fixed + grid is CUDAGraph-safe. For prefill (eager, no CG): + `W = num_write` exactly — no padding, no wasted grid programs + (long-prefill chunks would otherwise launch up to ~64× more + programs that all bail on `src_id < 0`). - `attn_metadata.batch_id_per_token`: [padded_T] int32 batch id per token (single per-token mapping; consumed by `swa_write`, Phase B/C/E paged-decode kernels, and the indexer). @@ -1132,215 +1262,211 @@ def _attach_v4_per_fwd_meta( """ win = self.window_size + # cu_seqlens_q_arr still needed for write_starts/write_ends below; + # token_num_per_seq is now passed in by caller (== batch.num_scheduled_tokens + # for prepare_decode/prefill, np.ones for MTP draft, etc.) — no longer + # re-derived from cu_seqlens_q here. cu_seqlens_q_arr = np.asarray( - cu_seqlens_q_np[: scheduled_bs + 1], dtype=np.int64 - ) - token_num_per_seq = cu_seqlens_q_arr[1:] - cu_seqlens_q_arr[:scheduled_bs] - start_pos_per_seq_np = np.asarray( - start_pos_per_seq_cpu[:scheduled_bs], dtype=np.int64 + cu_seqlens_q_np[: scheduled_bs + 1], dtype=np.int32 ) - # CG-padding-aware token count: padded_T == captured kernels' grid size. - # Real per-fwd data fills [0:total_tokens], rest is sentinel-padded so - # padded slot kernels skip cleanly. Eager / prefill paths leave both - # kwargs as None → no padding (padded_total_tokens == total_tokens). - if padded_bs is None or max_q_len is None: - padded_total_tokens = total_tokens - else: - padded_total_tokens = max(int(padded_bs) * int(max_q_len), total_tokens) - - # ----- window_topk (per-token) ----- - # Builder-internal intermediate: consumed below by - # `_attach_v4_paged_decode_meta` to derive SWA paged offsets. Not - # exposed on attn_metadata — no kernel in V4Attention.forward reads it. - # Stored in a stable forward_vars buffer (CG requires fixed addresses). - start_pos_per_seq_gpu = self._stage( - "v4_meta_start_pos_per_seq", start_pos_per_seq_np - ) - token_num_per_seq_gpu = self._stage( - "v4_meta_token_num_per_seq", token_num_per_seq - ) - # Decode-only fast path: token_num_per_seq is all-1 so the - # repeat_interleave below is identity. - is_decode_only = scheduled_bs == total_tokens and bool( - (token_num_per_seq == 1).all() - ) - if is_decode_only: - start_pos_per_token = start_pos_per_seq_gpu - else: - start_pos_per_token = torch.repeat_interleave( - start_pos_per_seq_gpu, token_num_per_seq_gpu + # is_pure_decode is set by the caller at AttentionMetaData_DSV4 + # construction time (single source of truth — prepare_decode / + # prepare_prefill / build_for_cudagraph_capture each know their + # own semantics). Consumed for: padded_total_tokens (CG bucket vs + # eager-tight), arange shortcut (write_indices), SWA grid bucketing, + # and by `_attach_v4_paged_decode_meta` for the index-buffer skip. + is_pure_decode = attn_metadata.is_pure_decode + + # padded_total_tokens: CG-captured decode/MTP pads to the fixed + # bucket `padded_bs * (1+max_spec_steps)` so per-token buffers + # (batch_id_per_token, swa_write_indices) have a stable shape across + # captures. Non-pure-decode (prefill / mixed / fresh-prefill) is + # eager and uses `total_tokens` exactly — no wasted padding (a long + # prefill chunk doesn't need to be padded up to a bucket that + # doesn't exist for it). + if is_pure_decode: + assert padded_bs is not None and max_q_len is not None, ( + "is_pure_decode requires padded_bs + max_q_len from caller " + "(CG bucket size — fixed at capture)" ) - window_topk_gpu = _build_window_topk_batched( - positions[:total_tokens].to(torch.long), - start_pos_per_token, - win, - out=self.model_runner.forward_vars["v4_meta_window_topk"].gpu, - ) + padded_total_tokens = int(padded_bs) * int(max_q_len) + else: + padded_total_tokens = total_tokens + + var = self.model_runner.forward_vars - # ----- batch_id_per_token (single per-token mapping) ----- - # Built unconditionally. Consumed by swa_write, Phase B/C/E - # paged-decode kernels, AND the indexer (cast to long inline). - # CG: real data [0:total_tokens]; padded slots get -1 sentinel so - # consumer kernels (csa_translate_pack, swa_write) skip on `bid<0`. - # int64 dtype: PyTorch fancy-index requires int64 INDEX (used by - # the indexer in `_build_v4_indexer_meta`); triton kernels (swa_write, - # csa_translate_pack) read whatever dtype the pointer carries — int64 - # works for them too. Single buffer for all consumers, no dtype mirror. + # ---- CPU numpy work (all on main thread) ---- batch_id_per_token_np = np.full(padded_total_tokens, -1, dtype=np.int64) batch_id_per_token_np[:total_tokens] = np.repeat( np.arange(scheduled_bs, dtype=np.int64), token_num_per_seq ) - attn_metadata.batch_id_per_token = self._stage( - "v4_batch_id_per_token", batch_id_per_token_np - ) - - # ----- n_committed_csa_per_seq (per-seq CSA committed count) ----- - # = ctx_len // 4. Built unconditionally so both csa_translate_pack - # (via batch_id_per_token lookup) and the indexer (uses raw per-seq - # vector) read from this single H2D-staged buffer. Kernel-side - # consumers do their own `min(., index_topk)` clamp via mask. - ctx_per_seq_np = np.asarray( - self.model_runner.forward_vars["context_lens"].np[:scheduled_bs], - dtype=np.int64, - ) - n_committed_csa_per_seq_np = (ctx_per_seq_np // 4).astype(np.int32) - attn_metadata.n_committed_csa_per_seq = self._stage( - "v4_n_committed_csa_per_seq", n_committed_csa_per_seq_np - ) - # ----- SWA write indices (last `win` tokens per seq) ----- - # `state_slot_mapping_cpu` is always sized to scheduled_bs by - # `_populate_state_slot_mapping` (zeros-fill for dummy/warmup batches); - # `token_num_per_seq.sum() == total_tokens >= 1` from the caller - # contract. write_starts/write_ends sum is `min(token_num, win).sum()`, - # also >= 1. No empty-write fallback needed. - write_starts = cu_seqlens_q_arr[:scheduled_bs] + np.maximum( - 0, token_num_per_seq - win - ) - write_ends = cu_seqlens_q_arr[1:] - # Build write_indices_np (compact, num_write entries). - if is_decode_only: - # 1 tok/seq: every row is the last-win token of its seq. - write_indices_np = np.arange(total_tokens, dtype=np.int64) + # context_lens is int32 on the buffer; keep dtype through divide so + # n_committed_{csa,hca} stay int32 (max value ~max_model_len // 4 ≪ 2^31). + ctx_per_seq_np = var["context_lens"].np[:scheduled_bs] + # Single source of truth for n_committed_{csa,hca}_per_seq on CPU. + # Stashed on attn_metadata so paged_decode_meta / paged_prefill_meta / + # v4_indexer_meta can read instead of each re-running `ctx // k`. + n_committed_csa_per_seq_np = ctx_per_seq_np // 4 + n_committed_hca_per_seq_np = ctx_per_seq_np // 128 + attn_metadata.n_committed_csa_per_seq_cpu = n_committed_csa_per_seq_np + attn_metadata.n_committed_hca_per_seq_cpu = n_committed_hca_per_seq_np + + # Build write_indices for the SWA-write kernel. + # is_pure_decode ⇒ source is `self._swa_iota` (static GPU tensor) + + # GPU fill_(-1) tail. No CPU numpy needed; `num_write = total_tokens`, + # `swa_write_grid = padded_total_tokens` (CG bucket). + # else ⇒ per-seq concat tail (filter `last win tokens per seq`). + # Built as a fresh local numpy (NOT into the shared pinned + # `wi_buf.np`, which would race with the next fwd's rewrite during + # in-flight DMA). `swa_write_grid = num_write` (tight, no padding — + # prefill is eager). + if is_pure_decode: + write_indices_np = None + num_write = total_tokens + swa_write_grid = padded_total_tokens else: + write_starts = cu_seqlens_q_arr[:scheduled_bs] + np.maximum( + 0, token_num_per_seq - win + ) + write_ends = cu_seqlens_q_arr[1:] write_indices_np = np.concatenate( [ np.arange(s, e, dtype=np.int64) for s, e in zip(write_starts, write_ends) ] ) - # Sentinel-pad to `padded_total_tokens` so the kernel grid matches the - # captured size at replay (CG bs may exceed real bs → padded slots - # need write_indices=-1 to bail). - num_write = int(write_indices_np.shape[0]) - wi_buf = self.model_runner.forward_vars["v4_meta_swa_write_indices"] - cap = wi_buf.np.shape[0] - assert ( - padded_total_tokens <= cap - ), f"v4_meta_swa_write_indices too small: need {padded_total_tokens}, have {cap}" - if num_write > 0: - wi_buf.np[:num_write] = write_indices_np - if num_write < padded_total_tokens: - wi_buf.np[num_write:padded_total_tokens].fill(-1) - attn_metadata.swa_write_indices = wi_buf.copy_to_gpu(padded_total_tokens) - - # ----- Phase B: paged-decode metadata (layer-invariant) ----- - # Builds per-fwd kv_indices/kv_indptr buffers for SWA / CSA / HCA - # plus valid_count + block_tables_per_token consumed by Phase C - # (V4Attention.forward CSA per-layer translation) and Phase E - # (sparse_attn_v4_paged_decode dispatch). Skipped (None defaults) for - # non-pure-decode batches (prefill, mixed, fresh-prefill). + num_write = int(write_indices_np.shape[0]) + swa_write_grid = num_write + + # ---- Stage all buffers to GPU ---- + # window_topk used to be CPU-built here ([T, win] of ring indices with + # -1 sentinels) and staged via v4_meta_window_topk. Now the ring index + # is computed inline inside `write_v4_paged_decode_indices` kernel + # from `var["positions"].gpu` — saves O(T·win) numpy work + 4 MB + # staging buffer. The `positions` H2D is already done by the caller. + attn_metadata.batch_id_per_token = self._stage( + "v4_batch_id_per_token", batch_id_per_token_np + ) + # Stage n_committed to GPU. For CG-replay safety: aiter + # `top_k_per_row_decode` iterates the CAPTURED grid (= padded_bs * + # next_n rows) and reads `rowEnds[batch_id]` for every row. Its + # per-row length formula is + # `row_len = rowEnds[bid] - next_n + (r % next_n) + 1` + # — for pad rows `bid ∈ [scheduled_bs, padded_bs)` the buffer slot + # carries a stale value from a prior fwd; if that stale value is + # `< next_n - 1` (easy with MTP3 next_n=4 if a prior fwd had a seq + # in early prefill with ctx ≤ 11), row_len becomes negative and the + # kernel's radix loop runs unbounded → GPU hang. The downstream + # `batch_id_per_token = -1` sentinel masks pad rows out of + # `csa_translate_pack`, so the value just needs to be "big enough" + # to keep row_len non-negative. Use `index_topk` (≥ 1024 ≫ next_n). + n_csa_buf = var["v4_n_committed_csa_per_seq"] + n_csa_buf.np[:scheduled_bs] = n_committed_csa_per_seq_np + if is_pure_decode and padded_bs is not None and padded_bs > scheduled_bs: + n_csa_buf.np[scheduled_bs:padded_bs] = self.index_topk + attn_metadata.n_committed_csa_per_seq = n_csa_buf.copy_to_gpu(padded_bs) + else: + attn_metadata.n_committed_csa_per_seq = n_csa_buf.copy_to_gpu(scheduled_bs) + # Race-free write into the shared `v4_meta_swa_write_indices` GPU + # buffer (stable pointer → CG-capture-safe). Bypasses `_stage` + # because `_stage`'s `buf.np[:n] = arr; copy_to_gpu(non_blocking=True)` + # pattern uses the shared pinned alias as the H2D source — the next + # fwd's CPU rewrite of `buf.np` can land mid-DMA and tear the GPU + # contents, producing torn `src_id` values that exceed `kv.shape[0]` + # and MEMORY_VIOLATION the `tl.load(kv_ptr + src_id * ...)` in + # `_swa_write_kernel`. + wi_gpu = var["v4_meta_swa_write_indices"].gpu + assert swa_write_grid <= wi_gpu.shape[0], ( + f"v4_meta_swa_write_indices too small: need {swa_write_grid}, " + f"have {wi_gpu.shape[0]}" + ) + if is_pure_decode: + if num_write > 0: + wi_gpu[:num_write].copy_(self._swa_iota[:num_write]) + if num_write < swa_write_grid: + wi_gpu[num_write:swa_write_grid].fill_(-1) + elif num_write > 0: + wi_gpu[:num_write].copy_(torch.from_numpy(write_indices_np)) + attn_metadata.swa_write_indices = wi_gpu[:swa_write_grid] + self._attach_v4_paged_decode_meta( attn_metadata=attn_metadata, token_num_per_seq=token_num_per_seq, - start_pos_per_seq_np=start_pos_per_seq_np, state_slot_mapping_cpu=state_slot_mapping_cpu, - window_topk_gpu=window_topk_gpu, scheduled_bs=scheduled_bs, total_tokens=total_tokens, padded_total_tokens=padded_total_tokens, ) - def _clear_v4_paged_decode_meta(self, attn_metadata) -> None: - """Set all Phase-B paged-decode fields to default (non-decode batches). - - Note: `batch_id_per_token` and `n_committed_csa_per_seq` are - intentionally NOT cleared here — both are also consumed by - `swa_write` / the indexer on prefill / mixed paths and are built - unconditionally in `_attach_v4_per_fwd_meta`. - """ - attn_metadata.is_pure_decode = False - attn_metadata.swa_pages = 0 - for k in ( - "kv_indices_swa", - "kv_indices_csa", - "kv_indices_hca", - "kv_indptr_swa", - "kv_indptr_csa", - "kv_indptr_hca", - ): - setattr(attn_metadata, k, None) - def _attach_v4_paged_decode_meta( self, attn_metadata, token_num_per_seq, - start_pos_per_seq_np, state_slot_mapping_cpu, - window_topk_gpu, scheduled_bs: int, total_tokens: int, padded_total_tokens: Optional[int] = None, ) -> None: """Phase B: build per-fwd paged-decode index buffers (layer-invariant). + All three per-token regions are RAGGED-PACKED — same layout family as + the prefill path (`_build_paged_prefill_meta`). Per-token slot count: + SWA: actual_swa_count = min(positions[t]+1, win) + CSA: actual_swa_count + min(n_committed_csa, index_topk) + HCA: actual_swa_count + n_committed_hca + Writes into stable forward_vars buffers (attn_metadata fields are the V4-namespaced counterparts on `AttentionMetaData_DSV4`): - - kv_indices_swa : SWA paged offsets, uniform stride win + - kv_indices_swa : per-token SWA paged offsets, ragged-packed - kv_indices_csa : SWA prefix written here; CSA compress section left UNINITIALIZED — V4Attention.forward fills it per-layer via csa_translate_pack (Phase C) - kv_indices_hca : SWA prefix + HCA compress section, both fully written (HCA is layer-invariant) - - kv_indptr_{swa,csa,hca} : 3 indptr cumsums (SWA uniform, CSA / - HCA packed; per doc §6.3). Padded tail repeats + - kv_indptr_{swa,csa,hca} : 3 ragged cumsums. Padded tail repeats last value → kv_len=0 sentinel for CG-padded slots. + - skip_prefix_len_csa : per-token actual_swa_count (offset where + csa_translate_pack starts writing CSA topk + within each token's region). Matches prefill + semantics (where it equals prefix_swa_count[t]). Reuses (built earlier in `_attach_v4_per_fwd_meta`): - batch_id_per_token : single per-token mapping (with -1 sentinel) - n_committed_csa_per_seq : per-seq `ctx_len // 4` - - Skipped (defaults via `_clear_v4_paged_decode_meta`) when not - is_pure_decode (prefill, mixed, fresh-prefill). + - var["positions"] : global token positions (already H2D-copied by + the caller; consumed by the index-write kernel + + CPU-side actual_swa_count cumsum here) + + Skipped when not is_pure_decode (prefill, mixed, fresh-prefill). The + Phase-B fields (kv_indices_*, kv_indptr_*, swa_pages) stay at their + dataclass defaults (None / 0) for non-decode batches; downstream + consumers gate on `is_pure_decode`. """ if scheduled_bs == 0 or total_tokens == 0: - self._clear_v4_paged_decode_meta(attn_metadata) - return + return # fields stay at dataclass defaults - # is_pure_decode = uniform tokens-per-seq AND no fresh prefill - # (per doc §7.4). Catches both regular decode (1 tok/seq) and MTP-1 - # (2 tok/seq); rejects mixed batches and fresh-prefill seqs. - tok_first = int(token_num_per_seq[0]) - is_uniform_tok = bool((token_num_per_seq == tok_first).all()) - no_fresh_prefill = bool((start_pos_per_seq_np > 0).all()) - is_pure_decode = is_uniform_tok and no_fresh_prefill + if not attn_metadata.is_pure_decode: + return # caller already marked non-decode; nothing to build - if not is_pure_decode or len(state_slot_mapping_cpu) < scheduled_bs: - self._clear_v4_paged_decode_meta(attn_metadata) + if len(state_slot_mapping_cpu) < scheduled_bs: + # Warmup / dummy_run carve-out: caller asserted pure_decode but + # state_slot_mapping is incomplete. Flip the flag so downstream + # consumers (V4Attention.forward) take the non-decode codepath + # instead of dereferencing the unbuilt kv_indices_* buffers. + attn_metadata.is_pure_decode = False return var = self.model_runner.forward_vars - win = self.window_size - # swa_pages = num_slots * win, layer-invariant; matches the boundary + win = self.window_size # per-token max SWA prefix slots + cs = self.win_with_spec # SWA region per-slot stride (W + mtp_k) + # swa_pages = num_slots * cs, layer-invariant; matches the boundary # between SWA and compress regions in unified_kv (Phase A). - swa_pages = self.model_runner.max_per_req_cache_slots * win + swa_pages = self.model_runner.max_per_req_cache_slots * cs T = total_tokens # ----- Per-seq scalars (CPU numpy) ----- - # context_lens / block_tables CPU mirrors are populated by the caller - # (prepare_decode / build_for_cudagraph_capture) BEFORE this runs. - ctx_per_seq = np.asarray(var["context_lens"].np[:scheduled_bs], dtype=np.int64) # The single per-token mapping. Built once in `_attach_v4_per_fwd_meta` # (so swa_write / indexer can also consume it). Pull the GPU view from # attn_metadata; recompute the np copy here only for cumsum math below. @@ -1349,21 +1475,31 @@ def _attach_v4_paged_decode_meta( ) # [T] int32 — host copy for indptr cumsums batch_id_per_token_gpu = attn_metadata.batch_id_per_token - n_committed_csa_per_seq = ctx_per_seq // 4 # host copy for cumsum below - n_committed_hca_per_seq = (ctx_per_seq // 128).astype(np.int64) - - # ----- 3 indptr cumsums (CPU numpy, per doc §6.3) ----- - # Per-token kv_len = f(per-seq quantity)[batch_id_per_token]. CSA - # length is clamped to index_topk because csa_translate_pack only - # writes that many rows per seq (kernel mask `(k < n_committed) & - # (k < index_topk)`); the host-side clamp here just keeps the indices - # buffer correctly sized — the staged GPU n_committed_csa stays raw. + # Read pre-computed `ctx // {4,128}` from attn_metadata — populated by + # `_attach_v4_per_fwd_meta` (always runs first). int32. + n_committed_csa_per_seq = attn_metadata.n_committed_csa_per_seq_cpu + n_committed_hca_per_seq = attn_metadata.n_committed_hca_per_seq_cpu + + # ----- 3 indptr cumsums (CPU numpy, ragged) ----- + # Per-token kv_len = actual_swa_count + n_compress. CSA length is + # clamped to index_topk because csa_translate_pack only writes that + # many rows per seq (kernel mask `(k < n_committed) & (k < index_topk)`); + # the host-side clamp here just keeps the indices buffer correctly + # sized — the staged GPU n_committed_csa stays raw. index_topk = self.index_topk n_committed_csa_clamped_per_token = np.minimum( n_committed_csa_per_seq[batch_id_per_token_np], index_topk ) n_committed_hca_per_token = n_committed_hca_per_seq[batch_id_per_token_np] + # actual_swa_count[t] = min(positions[t]+1, win). Matches the kernel's + # inline `n = tl.minimum(pos+1, win)` so SWA-prefix segment sizes line + # up perfectly. `var["positions"]` is the int64 CpuGpuBuffer populated + # + H2D-copied by the caller (prepare_decode / build_for_cudagraph_capture). + actual_swa_count_np = np.minimum(var["positions"].np[:T] + 1, win).astype( + np.int32 + ) + # CG-padding-aware T_for_indptr: indptr buffer must size to the # captured kernel grid (= padded_total_tokens) so padded slots see # `kv_len = indptr[t+1] - indptr[t] = 0` and the inner loop bails. @@ -1373,23 +1509,24 @@ def _attach_v4_paged_decode_meta( if T_pad < T: T_pad = T - # SWA: uniform stride for real [0:T+1]. Padded entries [T+1:T_pad+1] - # repeat the last value (T*win) → kv_len = 0 for padded tokens, kernel - # bails out of the inner loop without reading kv_indices. - swa_indptr_np = np.empty(T_pad + 1, dtype=np.int32) - swa_indptr_np[: T + 1] = np.arange(T + 1, dtype=np.int64) * win + # All three indptr cumsums output int32 directly. Values are bounded + # (T ≤ mnbt=8192, per-tok ≤ win + index_topk ≈ 2200 → max cumsum ~18M, + # well within int32). + # SWA: ragged, per-token len = actual_swa_count[t]. + swa_indptr_np = np.zeros(T_pad + 1, dtype=np.int32) + swa_indptr_np[1 : T + 1] = np.cumsum(actual_swa_count_np, dtype=np.int32) if T_pad > T: - swa_indptr_np[T + 1 :].fill(int(T * win)) - # CSA: packed, per-token len = win + min(n_committed_csa, index_topk) - csa_per_tok = win + n_committed_csa_clamped_per_token.astype(np.int64) + swa_indptr_np[T + 1 :].fill(int(swa_indptr_np[T])) + # CSA: ragged, per-token len = actual_swa_count + min(n_csa, index_topk) + csa_per_tok = actual_swa_count_np + n_committed_csa_clamped_per_token csa_indptr_np = np.zeros(T_pad + 1, dtype=np.int32) - csa_indptr_np[1 : T + 1] = np.cumsum(csa_per_tok).astype(np.int32) + csa_indptr_np[1 : T + 1] = np.cumsum(csa_per_tok, dtype=np.int32) if T_pad > T: csa_indptr_np[T + 1 :].fill(int(csa_indptr_np[T])) - # HCA: packed, per-token len = win + n_committed_hca - hca_per_tok = win + n_committed_hca_per_token + # HCA: ragged, per-token len = actual_swa_count + n_committed_hca + hca_per_tok = actual_swa_count_np + n_committed_hca_per_token hca_indptr_np = np.zeros(T_pad + 1, dtype=np.int32) - hca_indptr_np[1 : T + 1] = np.cumsum(hca_per_tok).astype(np.int32) + hca_indptr_np[1 : T + 1] = np.cumsum(hca_per_tok, dtype=np.int32) if T_pad > T: hca_indptr_np[T + 1 :].fill(int(hca_indptr_np[T])) @@ -1399,87 +1536,78 @@ def _attach_v4_paged_decode_meta( # batch_id_per_token + n_committed_csa_per_seq already staged in # `_attach_v4_per_fwd_meta`. - # ----- HCA compress paged offsets (CPU numpy) ----- - # `paged_offset = swa_pages + block_tables[bid, idx]` for - # idx in [0, n_committed_hca[bid]). HCA is layer-invariant so we fully - # build the packed flat array here and stage once. + # ----- HCA compress paged offsets (CPU numpy, vectorized) ----- block_tables_np_full = var["block_tables"].np[:scheduled_bs] hca_total_indices = int(hca_indptr_np[T]) - hca_indices_np = np.empty(hca_total_indices, dtype=np.int32) - # Layout per token: [w0..w_{win-1}, h0..h_{n_committed_hca[bid]-1}]. - # Window prefix is overwritten from GPU swa_paged below; pre-fill -1 - # so any uninitialized prefix slot is a sentinel the kernel skips. - hca_indices_np.fill(-1) - for t in range(T): - bid = int(batch_id_per_token_np[t]) - n_h = int(n_committed_hca_per_seq[bid]) - if n_h == 0: - continue - base_t = int(hca_indptr_np[t]) + win - hca_indices_np[base_t : base_t + n_h] = ( - swa_pages + block_tables_np_full[bid, :n_h] + hca_indices_np = np.full(hca_total_indices, -1, dtype=np.int32) + # n_committed_hca_per_seq is int32; gather stays int32. + n_h_per_token = n_committed_hca_per_seq[batch_id_per_token_np[:T]] + total_hca_entries = int(n_h_per_token.sum()) + if total_hca_entries > 0: + token_indices = np.repeat(np.arange(T, dtype=np.int32), n_h_per_token) + cu_n_h = np.zeros(T + 1, dtype=np.int32) + np.cumsum(n_h_per_token, out=cu_n_h[1:], dtype=np.int32) + entry_offsets = np.arange(total_hca_entries, dtype=np.int32) - np.repeat( + cu_n_h[:T], n_h_per_token ) + # HCA compress section starts at `actual_swa_count[t]` (was `win` + # under the old uniform layout) — ragged-packed offset matches + # what the kernel writes for the SWA prefix segment. + write_pos = ( + hca_indptr_np[token_indices] + + actual_swa_count_np[token_indices] + + entry_offsets + ) + bid_expanded = batch_id_per_token_np[token_indices] + hca_indices_np[write_pos] = ( + swa_pages + block_tables_np_full[bid_expanded, entry_offsets] + ).astype(np.int32) # Stage to GPU (HCA compress tail; window prefix scattered below). hca_indices_gpu = self._stage("v4_kv_indices_hca", hca_indices_np) - # ----- SWA paged offsets (GPU-side from window_topk_gpu) ----- - # `state_slot_per_token = state_slot_per_seq[batch_id_per_token]` — - # done as a transient GPU gather (no persistent per-token buffer). - # Reuse `attn_metadata.state_slot_mapping` (= forward_vars[ - # "v4_meta_state_slot_groups"] gpu view, set by `_populate_state_slot_mapping`) - # — same data as the legacy `v4_meta_state_slot_i32` re-staging, no - # second H2D needed. - device = self.device - state_slot_per_seq_gpu = attn_metadata.state_slot_mapping - # Use only the real-data prefix of batch_id_per_token here — the - # padded tail carries -1 sentinels, and `swa_paged_2d` is computed - # against `window_topk_gpu` of shape [T, win] only. - state_slot_per_token_gpu = state_slot_per_seq_gpu[ - batch_id_per_token_gpu[:T].long() - ] # [T] int32 transient - # `window_topk_gpu` [T, win] int32 — passed in from caller - # (`_attach_v4_per_fwd_meta` builds it once per fwd). Builder-internal - # intermediate; not exposed on attn_metadata since no kernel reads it. - swa_paged_2d = torch.where( - window_topk_gpu >= 0, - state_slot_per_token_gpu[:, None] * win + window_topk_gpu, - torch.full_like(window_topk_gpu, -1), - ) # [T, win] int32 - swa_paged_flat = swa_paged_2d.reshape(-1) # [T*win] int32 - - # ----- Write SWA paged offsets to all 3 buffer heads (GPU) ----- - # SWA buffer: uniform stride win → simple flat copy. + # ----- Write SWA / CSA / HCA window-prefix paged offsets (1 kernel) ----- + # Kernel computes `n = min(positions[t]+1, win)` and ring-index + # `(positions[t] - n + 1 + i) % cs` inline — no window_topk staging. + # See `write_v4_paged_decode_indices` docstring and plan + # `sequential-noodling-turing.md` for the motivation. Reads only + # persistent forward_vars buffers — no allocator churn (the prior + # `index_copy_` chain raced under MTP-3 long-prefill; this kernel + # also fixes that, see skill `debug-agent-locate-kernel`). swa_indices_gpu = var["v4_kv_indices_swa"].gpu - swa_indices_gpu[: T * win].copy_(swa_paged_flat) - - # CSA / HCA buffers: window prefix at packed positions - # [indptr[t], indptr[t]+win) — scatter via index_copy_. - win_arange = torch.arange(win, device=device, dtype=torch.int64) - csa_win_pos = ( - csa_indptr_gpu[:T].to(torch.int64).unsqueeze(1) + win_arange.unsqueeze(0) - ).reshape(-1) - hca_win_pos = ( - hca_indptr_gpu[:T].to(torch.int64).unsqueeze(1) + win_arange.unsqueeze(0) - ).reshape(-1) csa_indices_gpu = var["v4_kv_indices_csa"].gpu - csa_indices_gpu.index_copy_(0, csa_win_pos, swa_paged_flat) - hca_indices_gpu.index_copy_(0, hca_win_pos, swa_paged_flat) - - # ----- skip_prefix_len_csa: decode = full SWA window per token ----- - # csa_translate_pack consumes this to know where the CSA section - # starts within each per-token region. Decode path fills [0:T] with - # `win`; padded tail [T:padded_T] = 0 (kernel bails on bid<0). - # Pre-allocated buffer gives a stable address for CG capture. - skip_csa_gpu = var["v4_skip_prefix_len_csa"].gpu - skip_csa_gpu[:T].fill_(win) - skip_csa_gpu[T:].zero_() + write_v4_paged_decode_indices( + state_slot_per_seq=attn_metadata.state_slot_mapping, + batch_id_per_token=batch_id_per_token_gpu, + positions=var["positions"].gpu, + swa_indptr=swa_indptr_gpu, + csa_indptr=csa_indptr_gpu, + hca_indptr=hca_indptr_gpu, + swa_indices=swa_indices_gpu, + csa_indices=csa_indices_gpu, + hca_indices=hca_indices_gpu, + T=T, + win=win, + cs=cs, + ) + + # ----- skip_prefix_len_csa: per-token actual SWA-prefix length ----- + # csa_translate_pack consumes this as the offset within each token's + # `kv_indices_csa` region where the CSA topk section starts (after + # the SWA prefix segment). Decode + prefill now share this semantics + # (was `win` in decode, `prefix_swa_count[t]` in prefill). + skip_csa_buf = var["v4_skip_prefix_len_csa"] + skip_csa_buf.np[:T] = actual_swa_count_np + skip_csa_buf.np[T:T_pad].fill(0) + skip_csa_gpu = skip_csa_buf.copy_to_gpu(T_pad) # ----- Stash on attn_metadata for V4Attention.forward consumption ----- # batch_id_per_token + n_committed_csa_per_seq already set in # `_attach_v4_per_fwd_meta` (single source of truth, also consumed by # swa_write / indexer outside the is_pure_decode branch). - attn_metadata.is_pure_decode = True - attn_metadata.kv_indices_swa = swa_indices_gpu[: T * win] + # is_pure_decode was set by the caller at AttentionMetaData_DSV4 + # construction time; we only flip it (True→False) above when the + # warmup carve-out fires (incomplete state_slot_mapping_cpu). + attn_metadata.kv_indices_swa = swa_indices_gpu[: int(swa_indptr_np[T])] attn_metadata.kv_indices_csa = csa_indices_gpu[: int(csa_indptr_np[T])] attn_metadata.kv_indices_hca = hca_indices_gpu # already exact len attn_metadata.kv_indptr_swa = swa_indptr_gpu @@ -1493,6 +1621,7 @@ def _build_paged_prefill_meta( attn_metadata: AttentionMetaData_DSV4, positions_np: np.ndarray, cu_seqlens_q_np: np.ndarray, + token_num_per_seq: np.ndarray, start_pos_per_seq_np: np.ndarray, state_slot_mapping_cpu: np.ndarray, scheduled_bs: int, @@ -1533,11 +1662,14 @@ def _build_paged_prefill_meta( - skip_prefix_len_csa = prefix_swa_count_per_token (per-token) - swa_pages """ - if scheduled_bs == 0 or total_tokens == 0: - return + assert scheduled_bs >= 1 and total_tokens >= 1, ( + "scheduled_bs and total_tokens must be positive for prefill meta " + "build (got scheduled_bs={scheduled_bs}, total_tokens={total_tokens})" + ) device = self.device - win = self.window_size + win = self.window_size # per-token topk count + cs = self.win_with_spec # SWA region per-slot stride (W + mtp_k) index_topk = self.index_topk T = total_tokens # warmup_model runs BEFORE allocate_kv_cache binds the paged pool @@ -1548,60 +1680,65 @@ def _build_paged_prefill_meta( num_slots = getattr(self.model_runner, "max_per_req_cache_slots", 0) if num_slots == 0: return - swa_pages = num_slots * win + swa_pages = num_slots * cs + var = self.model_runner.forward_vars # used for block_tables + plan buffers # ----- Per-token quantities (CPU numpy) ----- + # cu_seqlens_q_arr still needed below (ext_cu_q gather); + # token_num_per_seq is now passed in by caller — no longer re-derived. + # All quantities here fit in int32: positions ≤ max_model_len ≪ 2^31, + # counts (win, index_topk) are small, cumsums bounded by T·max_per_tok + # ≈ 18M. Keeping int32 throughout avoids needless widen/narrow churn. cu_seqlens_q_arr = np.asarray( - cu_seqlens_q_np[: scheduled_bs + 1], dtype=np.int64 + cu_seqlens_q_np[: scheduled_bs + 1], dtype=np.int32 ) - token_num_per_seq = ( - cu_seqlens_q_arr[1:] - cu_seqlens_q_arr[:scheduled_bs] - ).astype(np.int64) + token_num_per_seq = np.asarray(token_num_per_seq, dtype=np.int32) chunk_start_per_seq = np.asarray( - start_pos_per_seq_np[:scheduled_bs], dtype=np.int64 + start_pos_per_seq_np[:scheduled_bs], dtype=np.int32 ) # batch_id_per_token_np mirrors what _attach_v4_per_fwd_meta computed - # but we need a CPU copy for the cumsum / segment math below. + # but we need a CPU copy for the cumsum / segment math below. CPU-only + # — the GPU copy (int64 for PyTorch fancy index) is staged separately. batch_id_per_token_np = np.repeat( - np.arange(scheduled_bs, dtype=np.int64), token_num_per_seq - ) # [T] int64 - positions_arr = np.asarray(positions_np[:T], dtype=np.int64) + np.arange(scheduled_bs, dtype=np.int32), token_num_per_seq + ) # [T] int32 + positions_arr = np.asarray(positions_np[:T], dtype=np.int32) token_pos_in_chunk = ( positions_arr - chunk_start_per_seq[batch_id_per_token_np] - ) # [T] int64 + ) # [T] int32 # SWA window low bound (clamped at 0); used to derive prefix_swa_count # AND the absolute global positions inside the prefix SWA section. - swa_window_low_global = np.maximum(0, positions_arr - win + 1) # [T] int64 + swa_window_low_global = np.maximum(0, positions_arr - win + 1) # [T] int32 - extend_count_np = np.minimum(token_pos_in_chunk + 1, win).astype(np.int64) + extend_count_np = np.minimum(token_pos_in_chunk + 1, win).astype(np.int32) prefix_swa_count_np = np.maximum( 0, chunk_start_per_seq[batch_id_per_token_np] - swa_window_low_global - ).astype(np.int64) + ).astype(np.int32) - # n_committed_{csa,hca}[bid] from per-seq context_lens (already populated - # by parent). HCA: ctx_len // 128, CSA already exposed as int32 GPU tensor; - # rebuild host copy for cumsum below. - var = self.model_runner.forward_vars - ctx_per_seq_np = np.asarray( - var["context_lens"].np[:scheduled_bs], dtype=np.int64 - ) - n_committed_csa_per_seq_np = ctx_per_seq_np // 4 - n_committed_hca_per_seq_np = ctx_per_seq_np // 128 + # Read pre-computed `ctx // {4,128}` from attn_metadata — populated by + # `_attach_v4_per_fwd_meta` (always runs first). int32. + n_committed_csa_per_seq_np = attn_metadata.n_committed_csa_per_seq_cpu + n_committed_hca_per_seq_np = attn_metadata.n_committed_hca_per_seq_cpu n_csa_per_token = np.minimum( n_committed_csa_per_seq_np[batch_id_per_token_np], index_topk - ) # [T] int64 — clamped because csa_translate_pack writes at most + ) # [T] int32 — clamped because csa_translate_pack writes at most # index_topk per token (kernel mask `(k < n) & (k < index_topk)`) - n_hca_per_token = n_committed_hca_per_seq_np[batch_id_per_token_np] # [T] int64 + n_hca_per_token = n_committed_hca_per_seq_np[batch_id_per_token_np] # [T] int32 # ----- indptr cumsums (CPU) ----- - prefix_swa_indptr_np = np.zeros(T + 1, dtype=np.int64) - prefix_swa_indptr_np[1:] = np.cumsum(prefix_swa_count_np) - prefix_csa_indptr_np = np.zeros(T + 1, dtype=np.int64) - prefix_csa_indptr_np[1:] = np.cumsum(prefix_swa_count_np + n_csa_per_token) - prefix_hca_indptr_np = np.zeros(T + 1, dtype=np.int64) - prefix_hca_indptr_np[1:] = np.cumsum(prefix_swa_count_np + n_hca_per_token) - extend_indptr_np = np.zeros(T + 1, dtype=np.int64) - extend_indptr_np[1:] = np.cumsum(extend_count_np) + # All output int32 directly (downstream H2D stages int32 GPU buffers). + prefix_swa_indptr_np = np.zeros(T + 1, dtype=np.int32) + prefix_swa_indptr_np[1:] = np.cumsum(prefix_swa_count_np, dtype=np.int32) + prefix_csa_indptr_np = np.zeros(T + 1, dtype=np.int32) + prefix_csa_indptr_np[1:] = np.cumsum( + prefix_swa_count_np + n_csa_per_token, dtype=np.int32 + ) + prefix_hca_indptr_np = np.zeros(T + 1, dtype=np.int32) + prefix_hca_indptr_np[1:] = np.cumsum( + prefix_swa_count_np + n_hca_per_token, dtype=np.int32 + ) + extend_indptr_np = np.zeros(T + 1, dtype=np.int32) + extend_indptr_np[1:] = np.cumsum(extend_count_np, dtype=np.int32) # ----- Extend kv_indices (in `kv` tensor): per-token rows ----- # extend window for token t: kv rows @@ -1609,10 +1746,10 @@ def _build_paged_prefill_meta( # ... cu_seqlens_q[bid] + token_pos_in_chunk[t]]`. Build via segment # expansion. ext_seg_tok, ext_seg_k = _segment_indices( - np.arange(T, dtype=np.int64), extend_count_np + np.arange(T, dtype=np.int32), extend_count_np ) # Pre-gather per-token vars for the segment positions (vectorised). - ext_cu_q = cu_seqlens_q_arr[:scheduled_bs][batch_id_per_token_np] # [T] int64 + ext_cu_q = cu_seqlens_q_arr[:scheduled_bs][batch_id_per_token_np] # [T] int32 ext_indices_np = ( ext_cu_q[ext_seg_tok] + ( @@ -1627,19 +1764,23 @@ def _build_paged_prefill_meta( # For each token's prefix SWA segment: positions # `[swa_window_low_global[t] + k for k in range(prefix_swa_count[t])]`, # paged into `unified_kv` SWA region: - # paged = state_slot[bid] * win + (global_pos % win) + # paged = state_slot[bid] * cs + (global_pos % cs) + # where cs = win_with_spec (per-slot ring size). MUST match the same + # stride/modulo used by `swa_write` and the decode-path + # `_attach_v4_paged_decode_meta`, otherwise prefill prefix reads would + # land in the wrong slot once MTP makes cs > win. swa_seg_tok, swa_seg_k = _segment_indices( - np.arange(T, dtype=np.int64), prefix_swa_count_np + np.arange(T, dtype=np.int32), prefix_swa_count_np ) state_slot_arr = np.asarray( - state_slot_mapping_cpu[:scheduled_bs], dtype=np.int64 + state_slot_mapping_cpu[:scheduled_bs], dtype=np.int32 ) prefix_swa_global_pos = ( swa_window_low_global[swa_seg_tok] + swa_seg_k - ) # [sum prefix_swa_count] int64 + ) # [sum prefix_swa_count] int32 prefix_swa_paged_np = ( - state_slot_arr[batch_id_per_token_np[swa_seg_tok]] * win - + (prefix_swa_global_pos % win) + state_slot_arr[batch_id_per_token_np[swa_seg_tok]] * cs + + (prefix_swa_global_pos % cs) ).astype(np.int32) # ----- HCA compress paged offsets (layer-invariant, fully built here) ----- @@ -1647,7 +1788,7 @@ def _build_paged_prefill_meta( # mapped to `swa_pages + phys` (HCA block_capacity = 1). block_tables_np = var["block_tables"].np[:scheduled_bs] hca_seg_tok, hca_seg_k = _segment_indices( - np.arange(T, dtype=np.int64), n_hca_per_token + np.arange(T, dtype=np.int32), n_hca_per_token ) hca_phys = block_tables_np[ batch_id_per_token_np[hca_seg_tok], hca_seg_k @@ -1695,50 +1836,64 @@ def _build_paged_prefill_meta( # ----- One-shot H2D for everything ----- attn_metadata.kv_indices_extend = torch.from_numpy(ext_indices_np).to(device) - attn_metadata.kv_indptr_extend = torch.from_numpy( - extend_indptr_np.astype(np.int32) - ).to(device) + attn_metadata.kv_indptr_extend = torch.from_numpy(extend_indptr_np).to(device) attn_metadata.kv_indices_prefix_swa = torch.from_numpy( kv_indices_prefix_swa_np ).to(device) - attn_metadata.kv_indptr_prefix_swa = torch.from_numpy( - prefix_swa_indptr_np.astype(np.int32) - ).to(device) + attn_metadata.kv_indptr_prefix_swa = torch.from_numpy(prefix_swa_indptr_np).to( + device + ) attn_metadata.kv_indices_prefix_csa = torch.from_numpy( kv_indices_prefix_csa_np ).to(device) - attn_metadata.kv_indptr_prefix_csa = torch.from_numpy( - prefix_csa_indptr_np.astype(np.int32) - ).to(device) + attn_metadata.kv_indptr_prefix_csa = torch.from_numpy(prefix_csa_indptr_np).to( + device + ) attn_metadata.kv_indices_prefix_hca = torch.from_numpy( kv_indices_prefix_hca_np ).to(device) - attn_metadata.kv_indptr_prefix_hca = torch.from_numpy( - prefix_hca_indptr_np.astype(np.int32) - ).to(device) - attn_metadata.skip_prefix_len_csa = torch.from_numpy( - prefix_swa_count_np.astype(np.int32) - ).to(device) + attn_metadata.kv_indptr_prefix_hca = torch.from_numpy(prefix_hca_indptr_np).to( + device + ) + attn_metadata.skip_prefix_len_csa = torch.from_numpy(prefix_swa_count_np).to( + device + ) attn_metadata.swa_pages = swa_pages - def _build_compress_plans(self, extend_lens_np, seq_lens_np, device): + def _build_compress_plans( + self, extend_lens_np, context_lens_np, *, for_decode_cg: bool + ): """Build per-ratio CompressPlan dict consumed by batched compressor. - Reuse this from both prepare_decode and prepare_prefill — caller - supplies extend_lens / seq_lens (np int32) and target device. Plan - tensors are written into the pre-allocated `v4_compress_plan_{ratio}` - / `v4_write_plan_{ratio}` CpuGpuBuffers (fixed pointers for - CUDAGraph capture); the kernels skip sentinel-marked tail rows. + Reuse this from prepare_decode / prepare_prefill / prepare_capture — + caller supplies extend_lens / context_lens (np int32). context_lens + is the absolute per-seq length AFTER the new extend tokens (i.e. + prefix + extend); `make_compress_plans` reads it as `context_lens_cpu` + and reconstructs prefix internally. + Plan tensors are written into the pre-allocated + `v4_compress_plan_{ratio}` / `v4_write_plan_{ratio}` CpuGpuBuffers + (fixed pointers for CUDAGraph capture); the kernels skip + sentinel-marked tail rows. + + `for_decode_cg`: True for decode runtime AND decode CG capture — + the returned plan_gpu is sliced to a fixed `_decode_compress_cap` + per ratio so capture/replay shapes match. False for eager prefill — + the plan_gpu is sliced to the actual `n_compress` (smallest grid). """ from atom.model_ops.v4_kernels import make_compress_plans if not self._unique_compress_ratios_overlap: return {} - # Ensure inputs are np int32 (callers may pass torch tensors / lists). - if isinstance(extend_lens_np, torch.Tensor): - extend_lens_np = extend_lens_np.cpu().numpy().astype(np.int32) - if isinstance(seq_lens_np, torch.Tensor): - seq_lens_np = seq_lens_np.cpu().numpy().astype(np.int32) + # Inputs MUST be numpy int32 — torch tensors would force a D2H sync. + # Callers are responsible for staging from forward_vars np mirrors. + assert isinstance(extend_lens_np, np.ndarray), ( + f"extend_lens_np must be np.ndarray, got {type(extend_lens_np).__name__} " + "— passing torch.Tensor here would trigger a hidden D2H sync" + ) + assert isinstance(context_lens_np, np.ndarray), ( + f"context_lens_np must be np.ndarray, got {type(context_lens_np).__name__} " + "— passing torch.Tensor here would trigger a hidden D2H sync" + ) var = self.model_runner.forward_vars plan_buffers = { ratio: { @@ -1748,11 +1903,13 @@ def _build_compress_plans(self, extend_lens_np, seq_lens_np, device): for ratio, _ in self._unique_compress_ratios_overlap } return make_compress_plans( - np.ascontiguousarray(extend_lens_np, dtype=np.int32), - np.ascontiguousarray(seq_lens_np, dtype=np.int32), + extend_lens_np, + context_lens_np, self._unique_compress_ratios_overlap, - device, plan_buffers=plan_buffers, + decode_capacity_per_ratio=( + self._decode_compress_cap if for_decode_cg else None + ), ) def _populate_block_tables( @@ -1785,8 +1942,6 @@ def _populate_state_slot_mapping( copy is consumed by the V4 forward path to avoid `.tolist()` syncs (PR-A Phase 2). """ - import numpy as np - groups_np = np.asarray( batch.per_req_cache_groups[:scheduled_bs], dtype=np.int32 ) @@ -1817,15 +1972,18 @@ def build_for_cudagraph_capture( buffers in `forward_vars`. Replay-time prepare_decode writes into the SAME buffers — captured graph reads stable addresses. - NOTE on dynamic-shape kernels (fused_compress_attn / update_compressor_states / - swa_write): these currently use variable kernel grids (`grid=(num_compress,)`), + NOTE on dynamic-shape kernels (`update_compressor_states` / `swa_write`): + these currently use variable kernel grids (`grid=(num_compress,)`), which CUDAGraph capture rejects. A follow-up PR converts them to fixed grid + sentinel masking. Until then, capture itself can succeed (the helpers run on CPU + small H2D), but model.forward inside torch.cuda.graph will likely fail at the first such kernel launch — the user can detect - this via capture log output. + this via capture log output. (`fused_compress_attn` is already + CG-safe: launches at the decode-tight slice + (`_decode_compress_cap[ratio]`, baked at capture) and + sentinel-skips inactive rows internally for both BF16 Main and FP8 + Indexer paths.) """ - device = self.model_runner.device var = self.model_runner.forward_vars # Honor MTP at capture time: V4-Pro `mtp_k=1` → 2 tokens/req. The # outer `model_runner.capture_cudagraph` populates cu_seqlens_q with @@ -1861,6 +2019,9 @@ def build_for_cudagraph_capture( block_tables_gpu = var["block_tables"].copy_to_gpu(bs) state_slot_gpu = self._stage("v4_meta_state_slot_groups", state_slot_np) + # Synthetic decode batch: start_pos = win > 0 and uniform + # max_q_len tokens per seq, so is_pure_decode is True by + # construction (capture replays the decode codepath). attn_metadata = AttentionMetaData_DSV4( cu_seqlens_q=cu_seqlens_q_gpu, cu_seqlens_k=None, @@ -1873,25 +2034,24 @@ def build_for_cudagraph_capture( num_cached_tokens=None, block_tables=block_tables_gpu, context_lens=context_lens_gpu, + is_pure_decode=True, ) attn_metadata.state_slot_mapping = state_slot_gpu attn_metadata.state_slot_mapping_cpu = state_slot_np - attn_metadata.start_pos_per_seq_cpu = positions_np[cu_seqlens_q_np[:bs]] # Build compress_plans + per-fwd meta + indexer meta via the same # helpers used at runtime — guarantees addresses match. extend_lens_np = np.full(bs, max_q_len, dtype=np.int32) attn_metadata.compress_plans = self._build_compress_plans( - extend_lens_np, context_lens_np, device + extend_lens_np, context_lens_np, for_decode_cg=True ) # Capture: padded_bs == scheduled_bs == bs (synthetic batch is full). # Must run BEFORE `_attach_v4_indexer_meta` so the indexer-side meta # builder can reuse the shared per-fwd GPU tensors. self._attach_v4_per_fwd_meta( attn_metadata, - positions, cu_seqlens_q_np, - attn_metadata.start_pos_per_seq_cpu, + extend_lens_np, # = np.full(bs, max_q_len) — synthetic uniform decode batch attn_metadata.state_slot_mapping_cpu, bs, total_tokens, @@ -1900,8 +2060,6 @@ def build_for_cudagraph_capture( ) self._attach_v4_indexer_meta( attn_metadata, - cu_seqlens_q_np, - attn_metadata.start_pos_per_seq_cpu, bs, total_tokens, positions_gpu=positions, @@ -1925,7 +2083,6 @@ def _alloc_v4_metadata_buffers(self) -> None: Bounds: - per-seq: max_bs - per-token: max_num_batched_tokens - - window_topk: max_num_batched_tokens * window_size - csa compress: max_num_batched_tokens * index_topk - hca compress: max_num_batched_tokens * max_num_blocks_per_seq - csa gather: max_bs * max_num_blocks_per_seq * (block_size // 4) @@ -1955,18 +2112,19 @@ def _alloc_v4_metadata_buffers(self) -> None: # `_populate_state_slot_mapping`); attn_metadata.state_slot_mapping # exposes that GPU view to all downstream consumers (no second # H2D-staged copy). - bufs["v4_meta_start_pos_per_seq"] = CpuGpuBuffer(bs, **i64) - bufs["v4_meta_token_num_per_seq"] = CpuGpuBuffer(bs, **i64) bufs["v4_meta_state_slot_groups"] = CpuGpuBuffer(bs, **i32) - bufs["v4_meta_swa_write_indices"] = CpuGpuBuffer(mnbt, **i64) - # Window-topk batched output buffer (per-token sliding-window indices). - # Sized [mnbt, win] int32 — `_build_window_topk_batched(out=)` writes - # into this; the [total_tokens, win] view feeds `_attach_v4_paged_decode_meta`'s - # SWA paged-offsets derivation. Builder-internal intermediate, but the - # buffer must be CG-capture-stable since it backs a tensor read by the - # captured graph (the SWA paged offsets derived from it land in - # `kv_indices_swa/csa/hca` which the captured kernels do read). - bufs["v4_meta_window_topk"] = CpuGpuBuffer(mnbt, win, **i32) + # swa_write_indices: tight bound = `max_bs * window_size`. Universal + # worst-case across paths — prefill compact write is `sum(min(num_i, + # win)) ≤ bs * win`; decode/MTP CG is `bs * (1+max_spec_steps) ≤ + # bs * win` (since `1+max_spec_steps ≪ win`). Legacy `mnbt` sizing + # was over-padded to the prefill total-token worst case. + bufs["v4_meta_swa_write_indices"] = CpuGpuBuffer(bs * win, **i64) + # Static GPU iota used by `_attach_v4_per_fwd_meta` for the + # is_pure_decode arange shortcut. Pre-allocated once so the per-fwd + # write into `v4_meta_swa_write_indices` is a pure GPU-to-GPU copy + # (no CPU intermediate → no pinned-buffer H2D race) and the source + # pointer stays stable for CUDAGraph capture. + self._swa_iota = torch.arange(bs * win, dtype=torch.int64, device=self.device) # Phase B: paged-decode index buffers (consumed by Phase C/E). # Sized to worst-case decode shape `T = max_bs * (1 + max_spec_steps)` @@ -2008,33 +2166,18 @@ def _alloc_v4_metadata_buffers(self) -> None: # for cu_seq_lens. Also reused as cu_starts/cu_ends for fp8_mqa_logits # (which accepts both int32 and int64). bufs["v4_indexer_cu_committed"] = CpuGpuBuffer(bs + 1, **i32) - # Decode-path logits buffer for `deepgemm_fp8_paged_mqa_logits`. - # Sized [max_bs, max_model_len_idx] fp32 — assumes V4-Pro next_n=1. - # deepgemm writes valid cols [0, n_committed_per_seq[batch]) per row; - # padding cols carry stale data but `top_k_per_row_decode` honors - # `n_committed_per_seq` per-row so unwritten cols are never selected. - bufs["v4_indexer_decode_logits"] = CpuGpuBuffer( - bs, self.max_model_len_idx, dtype=torch.float32, device=self.device - ) - # Decode-path top-k indices buffer (consumed by `top_k_per_row_decode`). - # Sized [max_bs, index_topk] int32 — V4-Pro next_n=1; multi-token decode - # would scale rows by next_n. Replaces the per-fwd torch.topk allocation. - bufs["v4_indexer_decode_topk_indices"] = CpuGpuBuffer( - bs, self.index_topk, dtype=torch.int32, device=self.device - ) + # NOTE: decode-path `logits` ([T, max_model_len_idx] fp32) and + # `topk_indices` ([T, index_topk] int32) are NOT pre-allocated — + # they are write-once GPU scratch with no CPU mirror, allocated + # per-fwd inside `Indexer._score_topk_decode` via `torch.empty`. + # Under CUDAGraph capture they land in the graph's private pool + # and replay reuses the same address; eager keeps the standard + # caching-allocator fast path. # Per-token write offset consumed by `csa_translate_pack` (decode path # fills with `window_size`; prefill path will fill with per-token # prior_swa_count once the new dual-source kernel lands). Allocated # `mnbt` (worst-case prefill) so prefill writes don't overflow. bufs["v4_skip_prefix_len_csa"] = CpuGpuBuffer(mnbt, **i32) - # FP8 cache write-side slot mapping (one entry per compress row). - # Bound = ⌈mnbt/4⌉ + bs (worst-case num_compress for ratio=4 CSA; - # matches v4_compress_plan_4 row count). - # int64 — `indexer_k_quant_and_cache` requires int64 slot_mapping. - idx_compress_bound = mnbt // 4 + bs - bufs["v4_indexer_compress_slot_mapping"] = CpuGpuBuffer( - idx_compress_bound, **i64 - ) # Compress plan buffers (per-ratio) — pre-allocated for CUDAGraph # plan-tensor address stability. `make_compress_plans(..., plan_buffers=)` @@ -2042,12 +2185,26 @@ def _alloc_v4_metadata_buffers(self) -> None: # sizes: num_compress ≤ ⌈mnbt/ratio⌉ + bs (one boundary per seq plus # alignment slack); num_write ≤ bs * STATE_SIZE (per-seq ring window # carries STATE_SIZE rows per fwd at most). - max_compress_per_ratio = {} + # + # The decode CG path uses a much tighter capacity than the prefill + # worst case — the kernel grid is dictated by the slice of this + # buffer that we hand to the kernel, and decode only ever needs + # `max_decode_tokens // ratio + max_bs` rows (vs `mnbt // ratio + bs` + # for prefill, which is ~13× larger at typical config). We still + # allocate the full prefill capacity (eager prefill needs it), but + # both decode capture and replay slice down to `_decode_compress_cap` + # so the captured grid is the decode-tight bound. capture and + # replay MUST use the same value (CG kernel call args are baked). + self._decode_compress_cap: dict[int, int] = {} for ratio, is_overlap in self._unique_compress_ratios_overlap: - state_size = (2 if is_overlap else 1) * ratio + # NOTE: this is the pool-window size (algorithm constant), NOT the + # state ring buffer size. The ring buffer is now K_pool + max_spec_steps + 1 + # to avoid R+1 re-commit borrow-reads (see csa_main_state_shape comment), + # but write_plan still emits ≤ K_pool rows per seq per fwd because + # `write_starts = max(0, context_lens - K_pool)` in make_compress_plans. + K_pool = (2 if is_overlap else 1) * ratio max_compress = mnbt // ratio + bs - max_write = min(mnbt, bs * state_size) - max_compress_per_ratio[ratio] = max_compress + max_write = min(mnbt, bs * K_pool) bufs[f"v4_compress_plan_{ratio}"] = CpuGpuBuffer(max_compress, 4, **i32) bufs[f"v4_write_plan_{ratio}"] = CpuGpuBuffer(max_write, 4, **i32) # Pre-fill with sentinel so capture-time buffer state is valid @@ -2056,24 +2213,12 @@ def _alloc_v4_metadata_buffers(self) -> None: bufs[f"v4_compress_plan_{ratio}"].copy_to_gpu() bufs[f"v4_write_plan_{ratio}"].cpu.fill_(-1) bufs[f"v4_write_plan_{ratio}"].copy_to_gpu() - - # Compressor output buffers (one per kind, shared across same-kind - # layers within a single fwd — Compressor outputs are consumed - # immediately by the layer's sparse_attn before the next layer runs). - # Sized to (max_compress_per_ratio, head_dim) so fused_compress_attn - # can launch with full-capacity grid + sentinel-skip; output rows - # past the actual num_compress carry stale data but are never read - # (consumer slices via `cu_compress_cpu`). - bf16 = {"dtype": torch.bfloat16, "device": self.device} - if 4 in max_compress_per_ratio: - mc = max_compress_per_ratio[4] - bufs["v4_csa_main_compress_out"] = torch.empty((mc, self.head_dim), **bf16) - bufs["v4_csa_idx_compress_out"] = torch.empty( - (mc, self.index_head_dim), **bf16 - ) - if 128 in max_compress_per_ratio: - mc = max_compress_per_ratio[128] - bufs["v4_hca_main_compress_out"] = torch.empty((mc, self.head_dim), **bf16) + # Decode-tight bound. Worst case = total_tokens_decode boundaries + # all firing simultaneously, each in its own ratio-aligned slot. + # `total_tokens_decode = max_decode_tokens` (= max_bs * (1+spec)). + # The `+ max_bs` covers per-seq alignment slack (each seq can hit + # at most one extra boundary when extend_len isn't ratio-aligned). + self._decode_compress_cap[ratio] = self.max_decode_tokens // ratio + bs self.model_runner.forward_vars.update(bufs) @@ -2084,15 +2229,19 @@ def _stage(self, name: str, arr) -> torch.Tensor: """ buf = self.model_runner.forward_vars[name] n = arr.shape[0] if arr.ndim > 0 else 1 - if n == 0: - return buf.gpu[:0] + assert ( + n > 0 + ), f"Cannot stage empty array for {name!r} — ensure the input array has at least one element." cap = buf.np.shape[0] assert n <= cap, ( f"V4 buffer {name!r} too small: need {n}, have {cap}. " f"Increase the corresponding bound in _alloc_v4_metadata_buffers." ) - if arr.dtype != buf.np.dtype: - arr = arr.astype(buf.np.dtype, copy=False) + assert arr.dtype == buf.np.dtype, ( + f"V4 buffer {name!r} dtype mismatch: buffer is {buf.np.dtype}, " + f"but got arr with dtype {arr.dtype}. Cast arr to the correct " + f"dtype before calling _stage." + ) buf.np[:n] = arr return buf.copy_to_gpu(n) diff --git a/atom/model_ops/attentions/triton_mha.py b/atom/model_ops/attentions/triton_mha.py new file mode 100644 index 0000000000..90b2babce5 --- /dev/null +++ b/atom/model_ops/attentions/triton_mha.py @@ -0,0 +1,131 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2025, Advanced Micro Devices, Inc. All rights reserved. + +import logging +from typing import Type + +import torch + +import atom.model_ops as ops +from atom.config import KVCacheTensor +from atom.model_engine.scheduler import ScheduledBatch +from atom.model_ops.attention_mha import PagedAttentionImpl +from atom.model_ops.paged_attention import PagedAttention + +from .aiter_attention import AiterAttentionMetadataBuilder +from .backends import AttentionBackend + +logger = logging.getLogger("atom") + + +class TritonMHABackend(AttentionBackend): + @staticmethod + def get_name() -> str: + return "ROCM_TRITON_MHA" + + @staticmethod + def get_builder_cls() -> Type["TritonMHAMetadataBuilder"]: + return TritonMHAMetadataBuilder + + @staticmethod + def get_impl_cls(): + attn_cls = ops.Attention + if attn_cls == PagedAttention: + return PagedAttentionImpl + raise NotImplementedError( + f"TritonMHABackend does not support attention class {attn_cls!r}" + ) + + +class TritonMHAMetadataBuilder(AiterAttentionMetadataBuilder): + """MHA metadata builder that allocates KV cache in flash layout. + + Flash layout: K/V both [num_blocks, block_size, num_kv_heads, head_dim]. + Consumed directly by aiter triton `unified_attention` for prefill+decode. + """ + + def prepare_prefill(self, batch: ScheduledBatch): + attn_metadata, positions = super().prepare_prefill(batch) + + # When there are no cached tokens, the base builder leaves + # `block_tables=None` because AiterBackend's prefill consumes raw q/k/v + # via flash_attn_varlen_func. The unified_attention path used by + # TritonMHABackend instead requires a block_table even for pure prefill, + # so build a fake one here that treats raw K/V as a kv_cache with + # block_size=1: row i = [cu_seqlens_k[i], ..., cu_seqlens_k[i]+max-1]. + if attn_metadata.block_tables is None: + cu_k = attn_metadata.cu_seqlens_k + num_seqs = cu_k.shape[0] - 1 + offsets = cu_k[:num_seqs] + attn_metadata.block_tables = offsets.unsqueeze(1) + torch.arange( + attn_metadata.max_seqlen_k, dtype=torch.int32, device=cu_k.device + ) + + return attn_metadata, positions + + def build_kv_cache_tensor(self, layer_id: int, module): + if not ( + hasattr(module, "base_attention") + and hasattr(module, "use_mla") + and not module.use_mla + ): + return None + + runner = self.model_runner + config = runner.config + hf_config = config.hf_config + + if runner.is_mimo_v2(): + raise NotImplementedError( + "TritonMHABackend does not support MiMo-V2 (per-layer alloc path)" + ) + + impl = getattr(module, "impl", None) + if impl is not None and ( + getattr(impl, "rotary_emb", None) is not None + and getattr(impl, "q_norm", None) is not None + and getattr(impl, "k_norm", None) is not None + ): + raise NotImplementedError( + "TritonMHABackend is incompatible with the fused qk_norm+rope+shuffle " + "cache path; use AiterBackend for this model." + ) + + if runner.is_qwen_next(): + mtp_start = runner.mtp_start_layer_idx + if layer_id < mtp_start: + attn_idx = layer_id // runner.full_attention_interval + else: + attn_idx = runner.num_full_attn + (layer_id - mtp_start) + else: + attn_idx = layer_id + + k_cache = runner.kv_cache[0, attn_idx].view( + runner.num_physical_kvcache_blocks, + runner.physical_block_size, + runner.num_kv_heads, + hf_config.head_dim, + ) + v_cache = runner.kv_cache[1, attn_idx].view( + runner.num_physical_kvcache_blocks, + runner.physical_block_size, + runner.num_kv_heads, + hf_config.head_dim, + ) + if config.kv_cache_dtype == "fp8": + module.k_scale = runner.kv_scale[0, attn_idx] + module.v_scale = runner.kv_scale[1, attn_idx] + + module.max_model_len = config.max_model_len + module.k_cache = k_cache + module.v_cache = v_cache + if impl is not None: + impl.use_flash_layout = True + + return KVCacheTensor( + layer_num=layer_id, + k_cache=k_cache, + v_cache=v_cache, + k_scale=module.k_scale, + v_scale=module.v_scale, + ) diff --git a/atom/model_ops/attentions/triton_mla.py b/atom/model_ops/attentions/triton_mla.py new file mode 100644 index 0000000000..3a52dfd4f0 --- /dev/null +++ b/atom/model_ops/attentions/triton_mla.py @@ -0,0 +1,101 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2025, Advanced Micro Devices, Inc. All rights reserved. + +import logging +from typing import Type + +import torch +from aiter.ops.triton.attention.mla_decode import csr_to_dense_block_table +from atom.model_engine.scheduler import ScheduledBatch +from atom.model_ops.attention_mla import MLAAttention +from atom.utils.forward_context import AttentionMetaData + +from .aiter_mla import AiterMLAMetadataBuilder +from .backends import AttentionBackend + +logger = logging.getLogger("atom") + + +class TritonMLABackend(AttentionBackend): + @staticmethod + def get_name() -> str: + return "ROCM_TRITON_MLA" + + @staticmethod + def get_builder_cls() -> Type["TritonMLAMetadataBuilder"]: + return TritonMLAMetadataBuilder + + @staticmethod + def get_impl_cls() -> Type["MLAAttention"]: + return MLAAttention + + +class TritonMLAMetadataBuilder(AiterMLAMetadataBuilder): + + def __init__(self, model_runner): + super().__init__(model_runner) + + hf = model_runner.config.hf_config + kv_lora_rank = hf.kv_lora_rank + num_kv_splits = 4 + triton_mla_buffers = { + "triton_block_table": torch.zeros( + self.max_bs, + self.max_num_blocks_per_seq, + dtype=torch.int32, + device=self.device, + ), + "triton_attn_logits": torch.empty( + self.max_bs, + self.padded_num_attention_heads, + num_kv_splits, + kv_lora_rank + 1, + dtype=torch.float32, + device=self.device, + ), + "triton_lse": torch.empty( + self.max_bs, + self.padded_num_attention_heads, + dtype=torch.float32, + device=self.device, + ), + } + self.model_runner.forward_vars.update(triton_mla_buffers) + + def set_mla_persistent_worker_buffers( + self, bs, max_q_len, only_update=False, num_reject_tokens=None + ): + # Triton MLA does not use aiter persistent worker buffers + return {} + + def prepare_decode(self, batch: ScheduledBatch, bs: int): + attn_metadata, positions = super().prepare_decode(batch, bs) + + scheduled_bs = batch.total_seqs_num_decode + max_seqlen_k = attn_metadata.max_seqlen_k + var = self.model_runner.forward_vars + + triton_bt = var["triton_block_table"][:scheduled_bs, :max_seqlen_k] + triton_bt.zero_() + csr_to_dense_block_table( + attn_metadata.kv_indices, + attn_metadata.kv_indptr, + triton_bt, + max_seqlen_k, + scheduled_bs, + ) + attn_metadata.triton_block_table = triton_bt + attn_metadata.triton_attn_logits = var["triton_attn_logits"][:scheduled_bs] + attn_metadata.triton_lse = var["triton_lse"][:scheduled_bs] + + return attn_metadata, positions + + def build_for_cudagraph_capture(self, bs: int) -> AttentionMetaData: + attn_metadata, context = super().build_for_cudagraph_capture(bs) + + var = self.model_runner.forward_vars + attn_metadata.triton_block_table = var["triton_block_table"][:bs] + attn_metadata.triton_attn_logits = var["triton_attn_logits"][:bs] + attn_metadata.triton_lse = var["triton_lse"][:bs] + + return attn_metadata, context diff --git a/atom/model_ops/base_attention.py b/atom/model_ops/base_attention.py index 3bf7760a03..8a7cd45c50 100644 --- a/atom/model_ops/base_attention.py +++ b/atom/model_ops/base_attention.py @@ -100,10 +100,8 @@ def cp_mha_gather_cache_kernel( # per-tensor: one scale per ptr, no offset k_scale = tl.load(k_scale_ptr) v_scale = tl.load(v_scale_ptr) - k_dtype = k_reg.dtype - v_dtype = v_reg.dtype - k_reg = (k_reg.to(tl.float32) * k_scale).to(k_dtype) - v_reg = (v_reg.to(tl.float32) * v_scale).to(v_dtype) + k_reg = k_reg.to(tl.float32) * k_scale + v_reg = v_reg.to(tl.float32) * v_scale tl.store(key_ptr_offset + col_offsets, k_reg) tl.store(value_ptr_offset + col_offsets, v_reg) @@ -167,6 +165,12 @@ def cp_mha_gather_cache( ], "kv_cache_layout only support NHD, SHUFFLE" if dequant: assert k_scales is not None and v_scales is not None + if k_scales.numel() == 1 and v_scales.numel() == 1: + per_token_quant = False + else: + assert ( + k_scales.numel() > 1 and v_scales.numel() > 1 + ), "k_scales and v_scales must both be scalar or per-token" head_dim = key.shape[2] x = 16 // key_cache.element_size() diff --git a/atom/model_ops/fla_ops/fused_sigmoid_gating.py b/atom/model_ops/fla_ops/fused_sigmoid_gating.py index 5f4df00484..236504ddc1 100644 --- a/atom/model_ops/fla_ops/fused_sigmoid_gating.py +++ b/atom/model_ops/fla_ops/fused_sigmoid_gating.py @@ -48,6 +48,8 @@ def fused_sigmoid_gating_delta_rule_update_kernel( V: tl.constexpr, BK: tl.constexpr, BV: tl.constexpr, + stride_a_token, + stride_b_token, stride_init_state_token: tl.constexpr, stride_final_state_token: tl.constexpr, stride_indices_seq: tl.constexpr, @@ -87,13 +89,13 @@ def fused_sigmoid_gating_delta_rule_update_kernel( p_A_log = A_log + i_hv if not IS_KDA: - p_a = a + bos * HV + i_hv + p_a = a + bos * stride_a_token + i_hv p_dt_bias = dt_bias + i_hv else: p_a = a + (bos * HV + i_hv) * K + o_k p_dt_bias = dt_bias + i_hv * K + o_k - p_b = b + bos * HV + i_hv + p_b = b + bos * stride_b_token + i_hv p_o = o + ((i_k * all + bos) * HV + i_hv) * V + o_v mask_k = o_k < K @@ -175,8 +177,8 @@ def fused_sigmoid_gating_delta_rule_update_kernel( p_k += H * K p_o += HV * V p_v += HV * V - p_b += HV - p_a += HV + p_b += stride_b_token + p_a += stride_a_token def fused_sigmoid_gating_delta_rule_update( @@ -246,8 +248,8 @@ def fused_sigmoid_gating_delta_rule_update( grid = (NK, NV, N * HV) fused_sigmoid_gating_delta_rule_update_kernel[grid]( A_log=A_log, - a=a.contiguous(), - b=b.contiguous(), + a=a, + b=b, dt_bias=dt_bias, beta=beta, threshold=threshold, @@ -270,6 +272,8 @@ def fused_sigmoid_gating_delta_rule_update( V=V, BK=BK, BV=BV, + stride_a_token=a.stride(-2), + stride_b_token=b.stride(-2), stride_init_state_token=stride_init_state_token, stride_final_state_token=stride_final_state_token, stride_indices_seq=stride_indices_seq, diff --git a/atom/model_ops/fused_moe_triton.py b/atom/model_ops/fused_moe_triton.py index 96aafcfe9b..8bb1be4c30 100644 --- a/atom/model_ops/fused_moe_triton.py +++ b/atom/model_ops/fused_moe_triton.py @@ -24,7 +24,12 @@ from typing import Any import logging from math import prod +from aiter import ActivationType from aiter.jit.utils.chip_info import get_gfx +from aiter.ops.triton.fusions.fused_routing_from_topk import ( + fused_routing_from_topk as _aiter_fused_routing_from_topk, +) +from aiter.ops.triton.fusions.fused_clamp_act_mul import fused_clamp_act_mul from atom.model_ops.utils import has_triton_kernels logger = logging.getLogger("atom") @@ -32,9 +37,14 @@ if has_triton_kernels(): try: - from triton_kernels.matmul_ogs import matmul_ogs + import triton_kernels.swiglu + from triton_kernels.matmul_ogs import ( + FnSpecs, + FusedActivation, + PrecisionConfig, + matmul_ogs, + ) from triton_kernels.routing import routing - from triton_kernels.matmul_ogs import PrecisionConfig except (AttributeError, ImportError) as e: logger.error( "Failed to import Triton kernels. Please make sure your triton " @@ -73,7 +83,7 @@ def _amd_smem_safe_tile(): # Defaults chosen so BLOCK_M*BLOCK_N stays ≤ 16384 entries (64 KiB FP32 # acc), comfortably fitting MI355X's register file. Override via env if # a future compiler/kernel update relaxes the budget. - block_m = int(os.getenv("ATOM_TRITON_MOE_BLOCK_M", "64")) + block_m = int(os.getenv("ATOM_TRITON_MOE_BLOCK_M", "32")) block_n = int(os.getenv("ATOM_TRITON_MOE_BLOCK_N", "256")) update_opt_flags_constraints({"block_m": block_m, "block_n": block_n}) try: @@ -108,6 +118,56 @@ def _swizzle_mxfp4(quant_tensor, scale): return quant_tensor, InFlexData(), scale +def fused_routing_from_topk_triton(topk_weights, topk_ids, n_expts_tot): + """Build matmul_ogs routing data via the AITER fused-routing kernel. + + Thin bridge over ``aiter.ops.triton.fused_routing_from_topk``: invokes + the single-CTA counting-sort kernel for small NK and packages the + resulting indices into the ``RoutingData`` / ``GatherIndx`` / + ``ScatterIndx`` structures consumed by + ``triton_kernels.matmul_ogs``. For ``NK = n_tokens * n_expts_act`` + above the kernel's single-CTA budget (prefill-shaped inputs), falls + back to the multi-kernel ``routing_from_topk`` reference defined + below — that path does the per-row sort + global stable argsort in + plain torch and is correctness-stable at any NK. + + Equivalence vs reference: the fused kernel skips the per-row sort, + so ``topk_indx`` / ``gate_indx`` differ at intra-expert ordering. + ``hist`` and the per-(token, expert, weight) bucket assignments + match exactly; ``matmul_ogs`` is commutative over per-expert slices + so the MoE output is unchanged (up to FP non-associativity). + """ + if not has_triton_kernels(): + return routing_from_topk(topk_weights, topk_ids, n_expts_tot) + + n_tokens, n_expts_act = topk_weights.shape + n_gates_pad = n_tokens * n_expts_act + + if n_gates_pad > 4096: + # Single-CTA design exceeded; fall back rather than degrading + # silently. Typically only hit during prefill. + return routing_from_topk(topk_weights, topk_ids, n_expts_tot) + + hist, topk_indx, gate_indx, gate_scal = _aiter_fused_routing_from_topk( + topk_weights, topk_ids, n_expts_tot + ) + + # Package as the matmul_ogs routing data structures. + from triton_kernels.routing import ( + RoutingData, + GatherIndx, + ScatterIndx, + compute_expt_data, + ) + + gather_indx = GatherIndx(src_indx=topk_indx, dst_indx=gate_indx) + scatter_indx = ScatterIndx(src_indx=gate_indx, dst_indx=topk_indx) + expt_data = compute_expt_data(hist, n_expts_tot, n_gates_pad) + + routing_data = RoutingData(gate_scal, hist, n_expts_tot, n_expts_act, expt_data) + return routing_data, gather_indx, scatter_indx + + def routing_from_topk(topk_weights, topk_ids, n_expts_tot): """Convert FusedMoE.select_experts output to triton routing data structures. @@ -283,28 +343,46 @@ def triton_kernel_fused_experts( ) with _amd_smem_safe_tile(): - matmul_ogs( - hidden_states, - w1, - w1_bias, - routing_data, - gather_indx=gather_indx, - precision_config=w13_precision_config, - gammas=gammas if apply_router_weight_on_input else None, - y=raw_intermediate, - ) - - # Standard SiLU/SwiGLU activation: silu(gate) * up - # With optional swiglu_limit clamping (V4: limit=10.0) - raw_2d = raw_intermediate.view(M * topk, N) - gate = raw_2d[:, :half_N] - up = raw_2d[:, half_N:] - if swiglu_limit > 0: - gate = gate.clamp(max=swiglu_limit) - up = up.clamp(-swiglu_limit, swiglu_limit) - intermediate_cache[0] = torch.nn.functional.silu(gate) * up + if activation == ActivationType.Swiglu: + # SwiGLU (GPT OSS): fused activation with interleaved [gate, up] layout + act = FusedActivation( + FnSpecs("swiglu", triton_kernels.swiglu.swiglu_fn, ("alpha", "limit")), + (swiglu_alpha, swiglu_limit), + 2, + ) + matmul_ogs( + hidden_states, + w1, + w1_bias, + routing_data, + gather_indx=gather_indx, + precision_config=w13_precision_config, + gammas=gammas if apply_router_weight_on_input else None, + fused_activation=act, + y=intermediate_cache, + ) + else: + # SiLU (DeepSeek): concatenated [gate | up] layout, manual activation + raw_intermediate = matmul_ogs( + hidden_states, + w1, + w1_bias, + routing_data, + gather_indx=gather_indx, + precision_config=w13_precision_config, + gammas=gammas if apply_router_weight_on_input else None, + ) + raw_2d = raw_intermediate.view(M * topk, N) + intermediate_cache = intermediate_cache.view(M * topk, half_N) + fused_clamp_act_mul( + raw_2d, + out=intermediate_cache, + swiglu_limit=swiglu_limit, + activation="silu", + dtype_quant=None, + ) + intermediate_cache = intermediate_cache.view(batch_dim, M * topk, half_N) - with _amd_smem_safe_tile(): matmul_ogs( intermediate_cache.view(M * topk, half_N), w2, diff --git a/atom/model_ops/layernorm.py b/atom/model_ops/layernorm.py index 030bcaae7a..14898b2001 100644 --- a/atom/model_ops/layernorm.py +++ b/atom/model_ops/layernorm.py @@ -21,6 +21,7 @@ from atom.model_ops.utils import atom_parameter from atom.quant_spec import LayerQuantConfig from atom.utils.decorators import mark_trace +from atom.utils import envs from torch import Tensor, nn from torch.overrides import handle_torch_function, has_torch_function_unary @@ -118,51 +119,93 @@ def fused_add_rmsnorm_pad_( return fused_add_rmsnorm_pad(x, weight, epsilon, res, x_pad_to_multiple) -def mxfp4_rms_quant_fuse_fake( +# --------------------------------------------------------------------------- +# Aiter dynamic RMSNorm + quant — single dispatch covering per_1x32 (MXFP4), +# per_1x128 (FP8 block), and per_Token (FP8). All three reach +# aiter.{rmsnorm_quant, add_rmsnorm_quant} (HIP) which both normalizes and +# emits a freshly-computed scale, so callers must have x_scale=None +# (static-scale FP8 stays on its own branch). Per-quant params (out_dtype, +# scale shape, group_size, shuffle_scale) are derived from quant_type_value +# inside the fake helper so torch.compile's schema infer sees a single, +# stable signature. +# +# `mutates_args=[]` keeps torch.compile from functionalizing the out-buffers +# — same pattern as the legacy mxfp4 fuse helper. +# --------------------------------------------------------------------------- + +_QV_PER_1X32 = QuantType.per_1x32.value +_QV_PER_1X128 = QuantType.per_1x128.value +_QV_PER_TOKEN = QuantType.per_Token.value +_AITER_RMS_QUANT_TYPE_VALUES = frozenset({_QV_PER_1X32, _QV_PER_1X128, _QV_PER_TOKEN}) + + +def _aiter_rms_quant_fake( x: torch.Tensor, weight: torch.Tensor, eps: float, - shuffle: bool = False, + quant_type_value: int, + transpose_scale: bool, res1: Optional[torch.Tensor] = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + from aiter.utility.dtypes import fp8 + M, N = x.shape - out = torch.empty((M, N // 2), dtype=torch.float4_e2m1fn_x2, device=x.device) - MXFP4_QUANT_BLOCK_SIZE = 32 - SCALE_N_valid = (N + MXFP4_QUANT_BLOCK_SIZE - 1) // MXFP4_QUANT_BLOCK_SIZE - use_scale_shuffle_padding = shuffle - if use_scale_shuffle_padding: - SCALE_M = ((M + 255) // 256) * 256 - SCALE_N = ((SCALE_N_valid + 7) // 8) * 8 - else: - SCALE_M = M - SCALE_N = SCALE_N_valid - scale = torch.empty( - (SCALE_M, SCALE_N), - dtype=torch.float8_e8m0fnu, - device=x.device, - ) - out_res1 = None - if res1 is not None: - out_res1 = torch.empty_like(res1) + if quant_type_value == _QV_PER_1X32: + # MXFP4: out=(M, N/2) fp4x2; scale=(⌈M/256⌉*256, ⌈⌈N/32⌉/8⌉*8) UE8M0 + # bytes (kernel writes one uint8 per group; passing fp8_e8m0fnu + # directly yields the matching byte layout — no fp32 view). + out = torch.empty((M, N // 2), dtype=torch.float4_e2m1fn_x2, device=x.device) + scale_m = ((M + 255) // 256) * 256 + scale_n = ((((N + 31) // 32) + 7) // 8) * 8 + scale = torch.empty( + (scale_m, scale_n), dtype=torch.float8_e8m0fnu, device=x.device + ) + elif quant_type_value == _QV_PER_1X128: + # FP8 per-block: scale=(M, ⌈N/128⌉) fp32. Preshuffle GEMM expects + # column-major; allocate (num_groups, M) row-major then view as + # (M, num_groups). Matches GemmaRMSNorm._forward_fused_fp8. + out = torch.empty((M, N), dtype=fp8, device=x.device) + num_groups = N // 128 + if transpose_scale: + scale = torch.empty( + (num_groups, M), dtype=torch.float32, device=x.device + ).view(M, num_groups) + else: + scale = torch.empty((M, num_groups), dtype=torch.float32, device=x.device) + else: # _QV_PER_TOKEN + out = torch.empty((M, N), dtype=fp8, device=x.device) + scale = torch.empty((M, 1), dtype=torch.float32, device=x.device) + out_res1 = torch.empty_like(res1) if res1 is not None else None return (out, scale, out_res1) -# It's important to use mutates_args=[] to avoid functionized_v2 op generation -@torch_compile_guard(gen_fake=mxfp4_rms_quant_fuse_fake, mutates_args=[]) -def mxfp4_rms_quant_fuse( +@torch_compile_guard(gen_fake=_aiter_rms_quant_fake, mutates_args=[]) +def _aiter_rms_quant( x: torch.Tensor, weight: torch.Tensor, eps: float, - shuffle: bool = False, + quant_type_value: int, + transpose_scale: bool, res1: Optional[torch.Tensor] = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - from aiter.ops.triton.fused_mxfp4_quant import fused_rms_mxfp4_quant + from aiter import add_rmsnorm_quant, rmsnorm_quant - (x_quant, x_scale), _, _, residual_out = fused_rms_mxfp4_quant( - x, weight, eps, shuffle=shuffle, res1=res1 + out, scale, out_res1 = _aiter_rms_quant_fake( + x, weight, eps, quant_type_value, transpose_scale, res1 ) - - return x_quant, x_scale, residual_out + if quant_type_value == _QV_PER_1X32: + group_size, shuffle = 32, True + elif quant_type_value == _QV_PER_1X128: + group_size, shuffle = 128, transpose_scale + else: # _QV_PER_TOKEN + group_size, shuffle = 0, False + if res1 is None: + rmsnorm_quant(out, x, scale, weight, eps, group_size, shuffle) + else: + add_rmsnorm_quant( + out, x, res1, out_res1, scale, weight, eps, group_size, shuffle + ) + return out, scale, out_res1 class RMSNorm(nn.Module): @@ -194,6 +237,14 @@ def __init__( params_dtype = layer_quant_config.quant_dtype self.quant_type = quant_type self.params_dtype = params_dtype + # transpose_scale (column-major scale) only applies to per_1x128 with + # the preshuffle GEMM consumer; resolve the env once at init time so + # forward sees a hot static bool instead of an env lookup per call. + self._aiter_transpose_scale = ( + fused_quant + and quant_type.value == _QV_PER_1X128 + and envs.ATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE + ) @mark_trace(prefix="rmsnorm", torch_compile=True) def forward( @@ -262,19 +313,24 @@ def forward( res1=residual, ) return (x, x_scale), residual - elif self.use_fused_quant and ( - x_scale is None and self.quant_type.value == QuantType.per_1x32.value + elif ( + self.use_fused_quant + and x_scale is None + and self.quant_type.value in _AITER_RMS_QUANT_TYPE_VALUES ): + # Dynamic-scale fused RMSNorm + quant via aiter HIP kernels. + # Static FP8 (x_scale provided) stays on the branch above. + x, x_scale, residual_out = _aiter_rms_quant( + x, + self.weight, + self.eps, + self.quant_type.value, + self._aiter_transpose_scale, + residual, + ) if residual is None: - x, x_scale, _ = mxfp4_rms_quant_fuse( - x, self.weight, self.eps, shuffle=True - ) return x, x_scale - else: - x, x_scale, residual = mxfp4_rms_quant_fuse( - x, self.weight, self.eps, shuffle=True, res1=residual - ) - return (x, x_scale), residual + return (x, x_scale), residual_out else: if residual is None: # return rmsnorm2d_fwd(x, self.weight, self.eps).view(ori_shape) @@ -347,8 +403,9 @@ def __init__( # Extract group size from quant type if quant_type == QuantType.per_1x128: self.group_size_quant = 128 - # per_1x128 blockscale GEMM requires transposed scale layout - self.transpose_scale = True + # preshuffle GEMM expects column-major x_scale; + # non-preshuffle GEMM expects row-major x_scale + self.transpose_scale = envs.ATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE elif quant_type == QuantType.per_1x32: self.group_size_quant = 32 self.transpose_scale = False @@ -547,25 +604,59 @@ def forward_cuda( x: torch.Tensor, residual: torch.Tensor | None = None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - from atom.model_ops.triton_gemma_rmsnorm import gemma_rmsnorm_triton + # Use the aiter HIP fused_qk_rmsnorm_group_quant kernel in no-quant mode + # (q_out_scale=None) to perform Gemma RMSNorm + optional residual add. + # Same math as the Triton kernel: out = rmsnorm(x [+ residual]) * (1 + w), + # but executed by the aiter kernel for higher achieved bandwidth. + from aiter.ops.fused_qk_rmsnorm_group_quant import fused_qk_rmsnorm_group_quant + + ori_shape = x.shape + x_2d = x.view(-1, ori_shape[-1]) - return gemma_rmsnorm_triton( - x, self.weight.data, self.variance_epsilon, residual + out = torch.empty_like(x_2d) + if residual is not None: + residual_2d = residual.view(-1, ori_shape[-1]) + res_out = torch.empty_like(x_2d) + else: + residual_2d = None + res_out = None + + fused_qk_rmsnorm_group_quant( + q=x_2d, + q_weight=self.weight.data, + q_epsilon=self.variance_epsilon, + q_out_unquantized=out, + q_res_out=res_out, + q_residual=residual_2d, + gemma_norm=True, ) + out = out.view(ori_shape) + if residual is not None: + return out, res_out.view(ori_shape) + return out + def _forward_fused_fp8(self, x, residual=None): from aiter.ops.fused_qk_rmsnorm_group_quant import fused_qk_rmsnorm_group_quant from aiter.utility.dtypes import fp8 + transpose_scale = envs.ATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE group_size = 128 M = x.shape[0] N = x.shape[1] num_groups = N // group_size out_fp8 = torch.empty((M, N), dtype=fp8, device=x.device) - out_scale = torch.empty( - (num_groups, M), dtype=torch.float32, device=x.device - ).view(M, num_groups) + if transpose_scale: + # column-major: allocate (num_groups, M) then view as (M, num_groups) + out_scale = torch.empty( + (num_groups, M), dtype=torch.float32, device=x.device + ).view(M, num_groups) + else: + # row-major: allocate (M, num_groups) directly + out_scale = torch.empty( + (M, num_groups), dtype=torch.float32, device=x.device + ) out_bf16 = ( torch.empty((M, N), dtype=x.dtype, device=x.device) if self.write_bf16 @@ -583,7 +674,7 @@ def _forward_fused_fp8(self, x, residual=None): q_res_out=res_out, q_residual=residual, group_size=group_size, - transpose_scale=True, + transpose_scale=transpose_scale, gemma_norm=True, ) if residual is not None: @@ -625,6 +716,7 @@ def _fused_qk_norm_single_kernel( num_q_heads, num_k_heads, ADD_UNIT_OFFSET: tl.constexpr, + Q_HAS_WEIGHT: tl.constexpr, RBLOCK: tl.constexpr, XBLOCK: tl.constexpr, ): @@ -654,15 +746,22 @@ def _fused_qk_norm_single_kernel( mask = xmask & col_mask - # Weight: load both, select via is_q - qw = tl.load( - q_weight_ptr + cols, mask=col_mask, other=0.0, eviction_policy="evict_last" - ).to(tl.float32) + # Weight: load both (or use ones for Q when Q_HAS_WEIGHT=False), select via is_q. + # Q_HAS_WEIGHT=False is for callers whose Q-side norm has implicit identity + # weight (e.g. V4's per-head Q normalization, equivalent to the prior + # `_rmsnorm_nw` helper) — saves a load + register row. + if Q_HAS_WEIGHT: + qw = tl.load( + q_weight_ptr + cols, mask=col_mask, other=0.0, eviction_policy="evict_last" + ).to(tl.float32) + if ADD_UNIT_OFFSET: + qw = qw + 1.0 + else: + qw = tl.full((RBLOCK,), 1.0, tl.float32) kw = tl.load( k_weight_ptr + cols, mask=col_mask, other=0.0, eviction_policy="evict_last" ).to(tl.float32) if ADD_UNIT_OFFSET: - qw = qw + 1.0 kw = kw + 1.0 w = tl.where(is_q, qw, kw) @@ -704,7 +803,7 @@ def _fused_qk_norm_single_kernel( def fused_qk_norm( q: torch.Tensor, k: torch.Tensor, - q_weight: torch.Tensor, + q_weight: Optional[torch.Tensor], k_weight: torch.Tensor, eps: float, add_unit_offset: bool = False, @@ -714,11 +813,18 @@ def fused_qk_norm( Args: q: [num_tokens, num_heads, head_dim] k: [num_tokens, num_kv_heads, head_dim] - q_weight, k_weight: [head_dim] norm weights + q_weight: [head_dim] norm weight, or None for an implicit ones weight + (skips the Q-side weight load — for callers whose Q norm + is the identity, e.g. V4's per-head Q normalization). + k_weight: [head_dim] norm weight (always required) eps: epsilon for numerical stability add_unit_offset: True for GemmaRMSNorm (w+1), False for standard """ - head_dim = q_weight.shape[0] + head_dim = k_weight.shape[0] + if q_weight is not None: + assert ( + q_weight.shape[0] == head_dim + ), f"q_weight head_dim {q_weight.shape[0]} != k_weight {head_dim}" num_tokens = q.shape[0] num_q_heads = q.shape[1] num_k_heads = k.shape[1] @@ -735,12 +841,15 @@ def fused_qk_norm( # num_warps=1 is universally optimal for head_dim=256 workloads on MI355X. XBLOCK = 2 if total_rows > 8192 else 1 NUM_WARPS = 1 + # When q_weight is None pass k_weight as a placeholder pointer (the + # kernel won't load from it — Q_HAS_WEIGHT=False gates the load). + q_weight_arg = q_weight if q_weight is not None else k_weight _fused_qk_norm_single_kernel[((total_rows + XBLOCK - 1) // XBLOCK,)]( q, k, q_out, k_out, - q_weight, + q_weight_arg, k_weight, eps, num_tokens, @@ -752,6 +861,7 @@ def fused_qk_norm( num_q_heads, num_k_heads, ADD_UNIT_OFFSET=add_unit_offset, + Q_HAS_WEIGHT=q_weight is not None, RBLOCK=RBLOCK, XBLOCK=XBLOCK, num_warps=NUM_WARPS, @@ -763,6 +873,12 @@ class DualRMSNorm: """Fused Q/K RMSNorm — single Triton kernel launch. Not an nn.Module. References existing q_norm/k_norm for weights. + + Q-side weightless mode: when `q_norm.weight is None`, the Q-side norm + is treated as the identity (implicit ones weight). The kernel skips + the q_weight load entirely (Q_HAS_WEIGHT=False). Use this when the + checkpoint does not ship a Q weight (e.g. V4's per-head Q normalization, + equivalent to the prior `_rmsnorm_nw` helper). """ def __init__( @@ -779,7 +895,22 @@ def __init__( self.num_q_heads = num_q_heads self.num_kv_heads = num_kv_heads self.head_dim = head_dim - self.add_unit_offset = isinstance(q_norm, GemmaRMSNorm) + # add_unit_offset only applies when q has a real weight; the kernel's + # weightless path (q_norm.weight is None) emits qw=1.0 directly with + # no offset, so the GemmaRMSNorm test is gated on weight existence. + self.add_unit_offset = getattr( + q_norm, "weight", None + ) is not None and isinstance(q_norm, GemmaRMSNorm) + # Resolve eps once. Different RMSNorm implementations name the + # attribute differently — `variance_epsilon` (HF/Gemma style) or + # `eps` (ATOM RMSNorm). Cache to avoid the lookup per forward call. + self._eps = getattr(q_norm, "variance_epsilon", None) or getattr( + q_norm, "eps", None + ) + assert self._eps is not None, ( + f"q_norm {type(q_norm).__name__} must expose `eps` or " + f"`variance_epsilon`" + ) self.prefix = prefix @mark_trace @@ -793,12 +924,15 @@ def __call__( Returns: (q_normed, k_normed) same shapes as input """ + # `self.q_norm.weight` is None for weightless q_norm (e.g. V4 + # `q_norm2`); fused_qk_norm forwards that None to the kernel which + # takes the Q_HAS_WEIGHT=False fast path. q, k = fused_qk_norm( q.view(-1, self.num_q_heads, self.head_dim), k.view(-1, self.num_kv_heads, self.head_dim), self.q_norm.weight, self.k_norm.weight, - self.q_norm.variance_epsilon, + self._eps, add_unit_offset=self.add_unit_offset, ) return ( diff --git a/atom/model_ops/linear.py b/atom/model_ops/linear.py index a0cb99179e..e016de7032 100644 --- a/atom/model_ops/linear.py +++ b/atom/model_ops/linear.py @@ -13,6 +13,7 @@ gemm_a8w8, gemm_a8w8_blockscale_bpreshuffle, gemm_a8w8_bpreshuffle, + gemm_a8w8_blockscale, get_hip_quant, ) @@ -317,11 +318,19 @@ def weight_loader_process( loaded_weight: torch.Tensor, post_process_func: Callable = lambda a: a, ): - if ( - param.data.dtype != loaded_weight.dtype - and param.data.element_size() == loaded_weight.element_size() - ): - param.data = param.data.view(loaded_weight.dtype) + if param.data.dtype != loaded_weight.dtype: + if param.data.element_size() == loaded_weight.element_size(): + # Same byte-width: use view for raw-bit-compatible pairs + # (e.g. fp8 variants) but convert for semantically different + # formats (float16 ↔ bfloat16) where bit reinterpretation + # would corrupt values. + incompatible = {torch.float16, torch.bfloat16} + if {param.data.dtype, loaded_weight.dtype} == incompatible: + loaded_weight = loaded_weight.to(param.data.dtype) + else: + param.data = param.data.view(loaded_weight.dtype) + else: + loaded_weight = loaded_weight.to(param.data.dtype) loaded_weight = post_process_func(loaded_weight) if ( loaded_weight.shape != param.data.shape @@ -379,10 +388,14 @@ def process_weights_after_loading(self): shuffle_weights(self.weight) # self.weight_scale.data = fp4_utils.e8m0_shuffle(self.weight_scale.data) else: - if ( + need_shuffle = ( self.quant_type == QuantType.per_Token and self.params_dtype == dtypes.fp8 - ) or (self.quant_type in [QuantType.per_1x32, QuantType.per_1x128]): + ) or self.quant_type == QuantType.per_1x32 + # per_1x128 only needs shuffle when using the preshuffle GEMM path + if not need_shuffle and self.quant_type == QuantType.per_1x128: + need_shuffle = envs.ATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE + if need_shuffle: if self.weight.dim() == 2: shuffle_weights(self.weight) # self.weight_scale.data = fp4_utils.e8m0_shuffle(self.weight_scale.data) @@ -405,8 +418,11 @@ def forward( if x_scale is None: quant_func = self.quant_func if self.quant_type.value == QuantType.per_1x128.value: + # preshuffle GEMM expects column-major x_scale; + # non-preshuffle GEMM expects row-major x_scale quant_func = functools_partial( - self.quant_func, transpose_scale=True + self.quant_func, + transpose_scale=envs.ATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE, ) if self.quant_type.value != QuantType.per_1x32.value: x, x_scale = quant_func( @@ -444,14 +460,23 @@ def forward( if self.bias is not None: y += self.bias elif self.quant_type.value == QuantType.per_1x128.value: - y = gemm_a8w8_blockscale_preshuffle_impl( - x, - self.weight, - x_scale, - self.weight_scale, - dtype=otype, - prefix=self.prefix, - ) + if envs.ATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE: + y = gemm_a8w8_blockscale_preshuffle_impl( + x, + self.weight, + x_scale, + self.weight_scale, + dtype=otype, + prefix=self.prefix, + ) + else: + y = gemm_a8w8_blockscale( + x, + self.weight, + x_scale, + self.weight_scale, + dtype=otype, + ) if self.bias is not None: y += self.bias elif self.quant_type.value == QuantType.per_1x32.value: @@ -635,6 +660,17 @@ def weight_loader( class QKVZBAParallelLinear(ColumnParallelLinear): + """Fused QKVZBA linear with deinterleaved output layout. + + Output layout is always ``[q | k | v | z | b | a]`` contiguous on dim-0, + so the caller can use ``torch.split`` for zero-copy views. + + The weight_loader deinterleaves the Qwen3-Next interleaved per-k-head-group + checkpoint layout during loading. Qwen3.5 checkpoints (which have separate + per-component weights) are loaded via individual shard_ids and work without + deinterleaving. + """ + def __init__( self, input_size: int, @@ -655,11 +691,15 @@ def __init__( tp_size = get_tp_group().world_size self.num_k_heads = divide(self.num_k_heads, tp_size) self.num_v_heads = divide(self.num_v_heads, tp_size) - output_sizes = [ - (2 * head_k_dim * self.num_k_heads + 2 * head_v_dim * self.num_v_heads) - * tp_size, - 2 * self.num_v_heads * tp_size, - ] + # Output layout: [q_all | k_all | v_all | z_all | b_all | a_all] + q_size = self.num_k_heads * head_k_dim + k_size = self.num_k_heads * head_k_dim + v_size = self.num_v_heads * head_v_dim + z_size = self.num_v_heads * head_v_dim + b_size = self.num_v_heads + a_size = self.num_v_heads + self._section_sizes = [q_size, k_size, v_size, z_size, b_size, a_size] + output_sizes = [s * tp_size for s in self._section_sizes] super().__init__( input_size, output_sizes, @@ -669,71 +709,304 @@ def __init__( prefix=prefix, ) + # -- helpers for deinterleaving during weight loading -- + + def _deinterleave_qkvz(self, param_data: torch.Tensor, loaded_weight: torch.Tensor): + """Scatter interleaved qkvz checkpoint rows into [q|k|v|z] regions. + + Checkpoint layout per k-head group (``QKVZ_DIM_SIZE`` rows): + [q(head_k_dim) | k(head_k_dim) | v0..vR(R*head_v_dim) | z0..zR(R*head_v_dim)] + where R = num_v_heads / num_k_heads (KV_HEAD_RATIO). + """ + nk = self.num_k_heads + hk = self.head_k_dim + hv = self.head_v_dim + R = self.num_v_heads // nk # KV_HEAD_RATIO + group_size = 2 * hk + 2 * hv * R + + q_total = nk * hk + k_total = nk * hk + v_total = self.num_v_heads * hv + + # TP shard the source + src = loaded_weight.narrow( + self.tp_dim, self.tp_rank * nk * group_size, nk * group_size + ) + + for g in range(nk): + base = g * group_size + # q rows + param_data[g * hk : (g + 1) * hk] = src[base : base + hk] + # k rows + param_data[q_total + g * hk : q_total + (g + 1) * hk] = src[ + base + hk : base + 2 * hk + ] + # v sub-heads + for s in range(R): + v_src_start = base + 2 * hk + s * hv + v_dst_start = q_total + k_total + (g * R + s) * hv + param_data[v_dst_start : v_dst_start + hv] = src[ + v_src_start : v_src_start + hv + ] + # z sub-heads + for s in range(R): + z_src_start = base + 2 * hk + R * hv + s * hv + z_dst_start = q_total + k_total + v_total + (g * R + s) * hv + param_data[z_dst_start : z_dst_start + hv] = src[ + z_src_start : z_src_start + hv + ] + + def _deinterleave_ba(self, param_data: torch.Tensor, loaded_weight: torch.Tensor): + """Scatter interleaved ba checkpoint rows into [b|a] regions. + + Checkpoint layout per k-head group (2*R elements): + [b_sub0, b_sub1, ..., b_subR-1, a_sub0, a_sub1, ..., a_subR-1] + where R = num_v_heads / num_k_heads. + """ + nk = self.num_k_heads + nv = self.num_v_heads + R = nv // nk + ba_total = 2 * nv + + # TP shard the source + src = loaded_weight.narrow(self.tp_dim, self.tp_rank * ba_total, ba_total) + + qkvz_total = sum(self._section_sizes[:4]) + b_offset = qkvz_total + a_offset = qkvz_total + nv + + for g in range(nk): + group_base = g * 2 * R + # b sub-heads + for s in range(R): + param_data[b_offset + g * R + s] = src[group_base + s] + # a sub-heads + for s in range(R): + param_data[a_offset + g * R + s] = src[group_base + R + s] + def weight_loader( self, param: nn.Parameter, loaded_weight: torch.Tensor, loaded_shard_id: str ): param_data = param.data assert loaded_shard_id in ["qkvz", "ba", "qkv", "z", "b", "a"] + + is_scale = param is getattr(self, "weight_scale", None) or param is getattr( + self, "input_scale", None + ) + + # For interleaved checkpoint shards ("qkvz", "ba"), deinterleave + # weight rows so output is [q|k|v|z|b|a] contiguous. + if loaded_shard_id == "qkvz" and not is_scale: + self._deinterleave_qkvz(param_data, loaded_weight) + return + if loaded_shard_id == "ba" and not is_scale: + self._deinterleave_ba(param_data, loaded_weight) + return + + # For individual shard_ids or scale params, use offset-based loading. + q_size, k_size, v_size, z_size, b_size, a_size = self._section_sizes if loaded_shard_id == "qkvz": - shard_size = ( - 2 * self.num_k_heads * self.head_k_dim - + 2 * self.num_v_heads * self.head_v_dim - ) + shard_size = q_size + k_size + v_size + z_size shard_offset = 0 - shard_rank = self.tp_rank elif loaded_shard_id == "qkv": - shard_size = ( - 2 * self.num_k_heads * self.head_k_dim - + self.num_v_heads * self.head_v_dim - ) + shard_size = q_size + k_size + v_size shard_offset = 0 - shard_rank = self.tp_rank elif loaded_shard_id == "z": - shard_size = self.num_v_heads * self.head_v_dim - shard_offset = ( - 2 * self.num_k_heads * self.head_k_dim - + self.num_v_heads * self.head_v_dim - ) - shard_rank = self.tp_rank + shard_size = z_size + shard_offset = q_size + k_size + v_size elif loaded_shard_id == "ba": - shard_size = 2 * self.num_v_heads - shard_offset = ( - 2 * self.num_k_heads * self.head_k_dim - + 2 * self.num_v_heads * self.head_v_dim - ) - shard_rank = self.tp_rank + shard_size = b_size + a_size + shard_offset = q_size + k_size + v_size + z_size elif loaded_shard_id == "b": - shard_size = self.num_v_heads - shard_offset = ( - 2 * self.num_k_heads * self.head_k_dim - + 2 * self.num_v_heads * self.head_v_dim - ) - shard_rank = self.tp_rank + shard_size = b_size + shard_offset = q_size + k_size + v_size + z_size elif loaded_shard_id == "a": - shard_size = self.num_v_heads - shard_offset = ( - 2 * self.num_k_heads * self.head_k_dim - + 2 * self.num_v_heads * self.head_v_dim - + self.num_v_heads - ) - shard_rank = self.tp_rank + shard_size = a_size + shard_offset = q_size + k_size + v_size + z_size + b_size - if param is getattr(self, "weight_scale", None) or param is getattr( - self, "input_scale", None - ): + if is_scale: if self.quant_type == QuantType.per_1x128: shard_offset = (shard_offset + 127) // 128 shard_size = (shard_size + 127) // 128 elif self.quant_type == QuantType.per_Tensor: loaded_weight = loaded_weight.view(1, 1).repeat(self.tp_size, 1) - shard_offset = ["qkvz", "ba"].index(loaded_shard_id) + shard_offset = ["qkvz", "ba", "qkv", "z", "b", "a"].index( + loaded_shard_id + ) shard_size = 1 - start_idx = shard_rank * shard_size + start_idx = self.tp_rank * shard_size param_data = param_data.narrow(self.tp_dim, shard_offset, shard_size) loaded_weight = loaded_weight.narrow(self.tp_dim, start_idx, shard_size) param.weight_loader_process(param_data, loaded_weight) +class QKVZParallelLinear(ColumnParallelLinear): + """Deinterleaving linear for the separate ``in_proj_qkvz`` checkpoint. + + The Qwen3-Next checkpoint stores qkvz in per-k-head-group interleaved + layout:: + + [q(hk) | k(hk) | v0..vR(R*hv) | z0..zR(R*hv)] per group + + This class deinterleaves both weight rows **and** per-1x128 block scales + during loading so the output layout is ``[q_all | k_all | v_all | z_all]`` + contiguous, enabling zero-copy ``torch.split`` at runtime. + """ + + def __init__( + self, + input_size: int, + head_k_dim: int, + head_v_dim: int, + num_k_heads: int, + num_v_heads: int, + bias: bool = False, + quant_config: Optional[QuantizationConfig] = None, + source_quant_dtype: torch.dtype = None, + prefix: str = "", + **kwargs, + ): + self.head_k_dim = head_k_dim + self.head_v_dim = head_v_dim + tp_size = get_tp_group().world_size + self.nk = divide(num_k_heads, tp_size) + self.nv = divide(num_v_heads, tp_size) + self.R = num_v_heads // num_k_heads # v-heads per k-head group + output_size = 2 * num_k_heads * head_k_dim + 2 * num_v_heads * head_v_dim + super().__init__( + input_size, + output_size, + bias=bias, + quant_config=quant_config, + source_quant_dtype=source_quant_dtype, + prefix=prefix, + ) + + @staticmethod + def _deinterleave(param_data, src, nk, R, hk, hv): + """Scatter interleaved rows into [q|k|v|z] regions.""" + group_size = 2 * hk + 2 * hv * R + q_total = nk * hk + k_total = nk * hk + v_total = nk * R * hv + + for g in range(nk): + base = g * group_size + # q + param_data[g * hk : (g + 1) * hk] = src[base : base + hk] + # k + param_data[q_total + g * hk : q_total + (g + 1) * hk] = src[ + base + hk : base + 2 * hk + ] + # v sub-heads + for s in range(R): + v_src = base + 2 * hk + s * hv + v_dst = q_total + k_total + (g * R + s) * hv + param_data[v_dst : v_dst + hv] = src[v_src : v_src + hv] + # z sub-heads + for s in range(R): + z_src = base + 2 * hk + R * hv + s * hv + z_dst = q_total + k_total + v_total + (g * R + s) * hv + param_data[z_dst : z_dst + hv] = src[z_src : z_src + hv] + + @staticmethod + def _match_dtype(param_data, loaded_weight): + """View param_data as loaded_weight's dtype if they differ but share element size. + + This mirrors ``weight_loader_process`` behaviour for FP8 on ROCm where + the param is ``float8_e4m3fnuz`` but the checkpoint stores + ``float8_e4m3fn``. The normalisation happens later in + ``process_weights_after_loading``. + """ + if ( + param_data.dtype != loaded_weight.dtype + and param_data.element_size() == loaded_weight.element_size() + ): + return param_data.view(loaded_weight.dtype) + return param_data + + def weight_loader(self, param: nn.Parameter, loaded_weight: torch.Tensor): + # Load like ColumnParallelLinear (no deinterleave at load time). + super().weight_loader(param, loaded_weight) + + def process_weights_after_loading(self): + nk, R = self.nk, self.R + hk, hv = self.head_k_dim, self.head_v_dim + + # Deinterleave weight rows: interleaved → [q|k|v|z] + w = self.weight.data + dw = torch.empty_like(w) + self._deinterleave(dw, w, nk, R, hk, hv) + self.weight.data = dw + + # Deinterleave weight_scale rows (per_1x128 block scale) + ws = getattr(self, "weight_scale", None) + if ws is not None: + hk_s, hv_s = hk // 128, hv // 128 + s = ws.data + ds = torch.empty_like(s) + self._deinterleave(ds, s, nk, R, hk_s, hv_s) + self.weight_scale.data = ds + + super().process_weights_after_loading() + + +class BAParallelLinear(ColumnParallelLinear): + """Deinterleaving linear for the separate ``in_proj_ba`` checkpoint. + + The Qwen3-Next checkpoint stores ba in per-k-head-group layout:: + + [b_s0..b_sR | a_s0..a_sR] per group (R = num_v_heads / num_k_heads) + + This class deinterleaves during loading so the output layout is + ``[b_all | a_all]`` contiguous, enabling zero-copy ``torch.split``. + + ``in_proj_ba`` is always BF16 (listed in ``modules_to_not_convert``), + so no weight-scale handling is needed. + """ + + def __init__( + self, + input_size: int, + num_k_heads: int, + num_v_heads: int, + bias: bool = False, + quant_config: Optional[QuantizationConfig] = None, + source_quant_dtype: torch.dtype = None, + prefix: str = "", + **kwargs, + ): + tp_size = get_tp_group().world_size + self.nk = divide(num_k_heads, tp_size) + self.R = num_v_heads // num_k_heads + self.nv = self.nk * self.R + output_size = 2 * num_v_heads + super().__init__( + input_size, + output_size, + bias=bias, + quant_config=quant_config, + source_quant_dtype=source_quant_dtype, + prefix=prefix, + ) + + def weight_loader(self, param: nn.Parameter, loaded_weight: torch.Tensor): + # Load like ColumnParallelLinear (no deinterleave at load time). + super().weight_loader(param, loaded_weight) + + def process_weights_after_loading(self): + nk, R, nv = self.nk, self.R, self.nv + w = self.weight.data + dw = torch.empty_like(w) + group_size = 2 * R + for g in range(nk): + base = g * group_size + dw[g * R : (g + 1) * R] = w[base : base + R] + dw[nv + g * R : nv + (g + 1) * R] = w[base + R : base + 2 * R] + self.weight.data = dw + super().process_weights_after_loading() + + class QKVGParallelLinear(ColumnParallelLinear): """QKV + output-Gate parallel linear. diff --git a/atom/model_ops/mamba_ops/causal_conv1d.py b/atom/model_ops/mamba_ops/causal_conv1d.py index 94c74a1447..bb60a40664 100644 --- a/atom/model_ops/mamba_ops/causal_conv1d.py +++ b/atom/model_ops/mamba_ops/causal_conv1d.py @@ -91,33 +91,16 @@ def _causal_conv1d_fwd_kernel( # continuous batching # BLOCK_N elements along the feature-dimension (channel) idx_feats = tl.program_id(1) * BLOCK_N + tl.arange(0, BLOCK_N) - idx_feats_start = tl.program_id(1) * BLOCK_N - o_ptr = x_ptr - stride_o_token = stride_x_token - stride_o_dim = stride_x_dim - idx_feats_o = idx_feats - dim_limits = dim - # Assume the idx_feats can be fully divided by BLOCK_N - if idx_feats_start < k_start_dim: - o_ptr = query_ptr - stride_o_token = query_token_stride - stride_o_dim = query_dim_stride - idx_feats_o = idx_feats - dim_limits = k_dim_size - elif idx_feats_start < v_start_dim: - o_ptr = key_ptr - stride_o_token = key_token_stride - stride_o_dim = key_dim_stride - idx_feats_o = idx_feats - k_start_dim - dim_limits = k_dim_size - elif idx_feats_start < dim: - o_ptr = value_ptr - stride_o_token = value_token_stride - stride_o_dim = value_dim_stride - idx_feats_o = idx_feats - v_start_dim - dim_limits = v_dim_size - else: - return + + # Pre-compute per-output feature indices and block-level masks. + # BLOCK_N divides k_dim_size evenly, so each program block falls entirely + # within one of q/k/v — only one mask is all-true per block. + q_feat_idx = idx_feats + k_feat_idx = idx_feats - k_start_dim + v_feat_idx = idx_feats - v_start_dim + is_q_block = idx_feats < k_start_dim + is_k_block = (idx_feats >= k_start_dim) & (idx_feats < v_start_dim) + is_v_block = idx_feats >= v_start_dim if idx_seq == pad_slot_id: return @@ -505,17 +488,25 @@ def _causal_conv1d_fwd_kernel( # continuous batching if SILU_ACTIVATION: acc = acc / (1 + tl.exp(-acc)) - mask_1d = (idx_token < segment_len) & ( - idx_feats_o < dim_limits - ) # token-index # feature-index + token_pos = sequence_start_index + token_offset + idx_token + mask_token = idx_token < segment_len - o_ptrs = ( - o_ptr - + (sequence_start_index + token_offset + idx_token) * stride_o_token - + (idx_feats_o * stride_o_dim) + # Store to q/k/v via separate tl.store calls (avoids pointer + # reassignment that crashes AMD TritonAMDGPUCanonicalizePointers). + # Each block falls entirely within one output, so only one store + # per block writes real data. + q_ptrs = ( + query_ptr + token_pos * query_token_stride + q_feat_idx * query_dim_stride ) + tl.store(q_ptrs, acc, mask=mask_token & is_q_block) - tl.store(o_ptrs, acc, mask=mask_1d) + k_ptrs = key_ptr + token_pos * key_token_stride + k_feat_idx * key_dim_stride + tl.store(k_ptrs, acc, mask=mask_token & is_k_block) + + v_ptrs = ( + value_ptr + token_pos * value_token_stride + v_feat_idx * value_dim_stride + ) + tl.store(v_ptrs, acc, mask=mask_token & is_v_block) def causal_conv1d_fn( @@ -596,7 +587,6 @@ def causal_conv1d_fn( # Store original dtype to cast back at the end original_x_dtype = x.dtype x = x.to(conv_states.dtype) - out = torch.empty_like(x) if metadata is not None: nums_dict = metadata.nums_dict args = nums_dict @@ -656,15 +646,7 @@ def causal_conv1d_fn( stride_istate_seq = conv_states.stride(0) stride_istate_dim = conv_states.stride(1) stride_istate_token = conv_states.stride(2) - # Keep this aligned with upstream: the Triton kernel consumes the - # runtime conv_state strides directly instead of requiring dim-stride 1. - # assert stride_istate_dim == 1 - if out.dim() == 2: - stride_o_dim = out.stride(0) - stride_o_token = out.stride(1) - else: - stride_o_dim = out.stride(1) - stride_o_token = out.stride(2) + stride_cache_indices = cache_indices.stride(0) if cache_indices is not None else 0 if validate_data: @@ -884,42 +866,18 @@ def _causal_conv1d_update_kernel( # [BLOCK_N,] elements along the feature-dimension (channel) idx_feats = tl.program_id(1) * BLOCK_N + tl.arange(0, BLOCK_N) - idx_feats_start = tl.program_id(1) * BLOCK_N k_start_point: tl.constexpr = k_dim_size v_start_point: tl.constexpr = k_dim_size * 2 - stride_o_seq = stride_q_seq - - o_ptr = x_ptr - stride_o_token = stride_x_token - stride_o_dim = stride_x_dim - idx_feats_o = idx_feats - dim_limits = dim - if idx_feats_start < k_start_point: - o_ptr = query_ptr - # In VARLEN mode, tokens are indexed along dim-0 of the output tensor, - # so the token stride is stride_q_seq (not stride_q_token which is the - # last-dim stride of the [num_tokens, k_dim, 1] output). - stride_o_token = stride_q_seq if IS_VARLEN else stride_q_token - stride_o_dim = stride_q_dim - idx_feats_o = idx_feats - dim_limits = k_dim_size - stride_o_seq = stride_q_seq - elif idx_feats_start < v_start_point: - o_ptr = key_ptr - stride_o_token = stride_k_seq if IS_VARLEN else stride_k_token - stride_o_dim = stride_k_dim - idx_feats_o = idx_feats - k_start_point - dim_limits = k_dim_size - stride_o_seq = stride_k_seq - elif idx_feats_start < dim: - o_ptr = value_ptr - stride_o_token = stride_v_seq if IS_VARLEN else stride_v_token - stride_o_dim = stride_v_dim - idx_feats_o = idx_feats - v_start_point - dim_limits = v_dim_size - stride_o_seq = stride_v_seq - else: - return + + # Pre-compute per-output feature indices and block-level masks. + # BLOCK_N divides k_dim_size evenly, so each program block falls entirely + # within one of q/k/v — only one mask is all-true per block. + q_feat_idx = idx_feats + k_feat_idx = idx_feats - k_start_point + v_feat_idx = idx_feats - v_start_point + is_q_block = idx_feats < k_start_point + is_k_block = (idx_feats >= k_start_point) & (idx_feats < v_start_point) + is_v_block = (idx_feats >= v_start_point) & (idx_feats < dim) if IS_APC_ENABLED: # Get the state from the initial_state_idx @@ -946,12 +904,25 @@ def _causal_conv1d_update_kernel( state_len = state_len - (seqlen - (query_end_index - query_start_index)) seqlen = query_end_index - query_start_index x_offset = query_start_index * stride_x_token - o_offset = query_start_index * stride_o_token + # Varlen: output is [num_tokens, split_dim, 1], token_pos indexes dim-0 + # so the "token" stride is stride_q_seq (the dim-0 stride). + q_base = query_start_index * stride_q_seq + k_base = query_start_index * stride_k_seq + v_base = query_start_index * stride_v_seq + q_tok_stride = stride_q_seq + k_tok_stride = stride_k_seq + v_tok_stride = stride_v_seq else: query_start_index = idx_seq * seqlen query_end_index = query_start_index + seqlen x_offset = idx_seq * stride_x_seq - o_offset = idx_seq * stride_o_seq + # Non-varlen: output is [batch, split_dim, seqlen] + q_base = idx_seq * stride_q_seq + k_base = idx_seq * stride_k_seq + v_base = idx_seq * stride_v_seq + q_tok_stride = stride_q_token + k_tok_stride = stride_k_token + v_tok_stride = stride_v_token if query_start_index == query_end_index: return @@ -1184,14 +1155,24 @@ def _causal_conv1d_update_kernel( if SILU_ACTIVATION: acc = acc / (1 + tl.exp(-acc)) - mask_1d = (idx_token < seqlen) & ( - idx_feats_o < dim_limits - ) # token-index # feature-index - o_ptrs = ( - o_ptr + o_offset + idx_token * stride_o_token + (idx_feats_o * stride_o_dim) + mask_token = idx_token < seqlen + + # Store to q/k/v via separate tl.store calls (avoids pointer + # reassignment that crashes AMD TritonAMDGPUCanonicalizePointers). + # Each block falls entirely within one output, so only one store + # per block writes real data. + q_ptrs = ( + query_ptr + q_base + idx_token * q_tok_stride + q_feat_idx * stride_q_dim ) + tl.store(q_ptrs, acc, mask=mask_token & is_q_block) - tl.store(o_ptrs, acc, mask=mask_1d) + k_ptrs = key_ptr + k_base + idx_token * k_tok_stride + k_feat_idx * stride_k_dim + tl.store(k_ptrs, acc, mask=mask_token & is_k_block) + + v_ptrs = ( + value_ptr + v_base + idx_token * v_tok_stride + v_feat_idx * stride_v_dim + ) + tl.store(v_ptrs, acc, mask=mask_token & is_v_block) def causal_conv1d_update( @@ -1308,19 +1289,15 @@ def causal_conv1d_update( stride_k_seq, stride_k_dim, stride_k_token = key.stride() stride_v_seq, stride_v_dim, stride_v_token = value.stride() - out = x stride_w_dim, stride_w_width = weight.stride() if query_start_loc is None: # X (batch, dim, seqlen) stride_x_seq, stride_x_dim, stride_x_token = x.stride() - stride_o_seq, stride_o_dim, stride_o_token = out.stride() else: - # X (dim, cu_seqlen) + # X (num_tokens, dim) stride_x_token, stride_x_dim = x.stride() stride_x_seq = 0 - stride_o_token, stride_o_dim = out.stride() - stride_o_seq = 0 stride_istate_seq, stride_istate_dim, stride_istate_token = conv_state.stride() stride_state_indices = ( @@ -1392,8 +1369,6 @@ def grid(META): USE_PAD_SLOT=pad_slot_id is not None, BLOCK_N=256, ) - if unsqueeze: - out = out.squeeze(-1) query = query.squeeze(-1) key = key.squeeze(-1) value = value.squeeze(-1) diff --git a/atom/model_ops/moe.py b/atom/model_ops/moe.py index 9f6e0a0f43..b53aa73371 100644 --- a/atom/model_ops/moe.py +++ b/atom/model_ops/moe.py @@ -8,18 +8,18 @@ from typing import Callable, List, Optional, Tuple import torch -from aiter import ActivationType, QuantType, dtypes, get_hip_quant, topk_softplus +from aiter import ActivationType, QuantType, dtypes, get_hip_quant, topk_gating from aiter.dist.parallel_state import get_dp_group, get_tp_group from aiter.fused_moe import fused_moe from aiter.jit.utils.chip_info import get_gfx from aiter.jit.utils.torch_guard import torch_compile_guard -from aiter.ops.shuffle import shuffle_scale_a16w4, shuffle_weight_a16w4 -from aiter.utility import fp4_utils +from aiter.ops.shuffle import shuffle_weight, shuffle_scale from atom.config import ( Config, QuantizationConfig, get_current_atom_config, ) +from aiter.ops.flydsl.moe_common import GateMode from atom.quant_spec import LayerQuantConfig from atom.model_loader.weight_utils import set_weight_attrs from atom.model_ops.base_config import QuantizeMethodBase @@ -681,16 +681,20 @@ def __init__(self, quant_config: LayerQuantConfig, moe: FusedMoEConfig): self.quant_dtype = quant_config.quant_dtype self.quant_method = quant_config.quant_method or "" self.static_input_scales = not quant_config.is_dynamic + self.is_guinterleave = envs.ATOM_MOE_GU_ITLV self.block_quant = ( self.quant_type == QuantType.per_1x128 or self.quant_type == QuantType.per_1x32 ) gfx = get_gfx() - self.use_triton = ( - gfx.startswith("gfx94") - or (gfx.startswith("gfx95") and envs.ATOM_USE_TRITON_GEMM) - or os.environ.get("ATOM_USE_TRITON_MOE") == "1" - ) + if envs.is_set("ATOM_USE_TRITON_MOE"): + self.use_triton = envs.ATOM_USE_TRITON_MOE + else: + self.use_triton = ( + gfx.startswith("gfx94") + or gfx.startswith("gfx12") + or (gfx.startswith("gfx95") and envs.ATOM_USE_TRITON_GEMM) + ) if self.use_triton: from atom.model_ops.utils import has_triton_kernels @@ -752,7 +756,7 @@ def create_weights( if layer.has_bias: w13_bias = atom_parameter( - torch.empty( + torch.zeros( num_experts, 2 * intermediate_size_per_partition_after_pad, dtype=torch.bfloat16, @@ -789,7 +793,7 @@ def create_weights( if layer.has_bias: w2_bias = atom_parameter( - torch.empty( + torch.zeros( num_experts, hidden_size, dtype=torch.bfloat16, @@ -853,62 +857,33 @@ def process_weights_after_loading(self, layer): layer.w13_weight_scale = None layer.w2_weight_scale = None return - elif layer.activation == ActivationType.Swiglu: - e, n, k = layer.w13_weight.shape - layer.w13_weight.view(torch.uint8).copy_( - layer.w13_weight.data.view(torch.uint8) - .view(e, n // 2, 2, k) - .permute(0, 2, 1, 3) - .contiguous() - .view(e, n, k) - ) - layer.w13_weight_scale.data = ( - layer.w13_weight_scale.data.view(e, n // 2, 2, -1) - .permute(0, 2, 1, 3) - .contiguous() - .view(e, n, -1) - ) - layer.w13_weight.data = shuffle_weight_a16w4(layer.w13_weight, 16, True) - shuffled_w13_scale = shuffle_scale_a16w4( - layer.w13_weight_scale.view(-1, layer.w13_weight_scale.shape[-1]), - self.num_experts, - True, - ) - layer.w2_weight.data = shuffle_weight_a16w4(layer.w2_weight, 16, False) - shuffled_w2_scale = shuffle_scale_a16w4( - layer.w2_weight_scale.view(-1, layer.w2_weight_scale.shape[-1]), - self.num_experts, - False, - ) - if layer.w13_bias is not None: - layer.w13_bias.data = ( - layer.w13_bias.data.view(-1, n // 2, 2) - .permute(0, 2, 1) - .contiguous() - .view(-1, n) - ) - # quark method for moe, split it out? - elif self.quant_method == "quark": - shuffle_weights(layer.w13_weight, layer.w2_weight) - s0, s1, _ = layer.w13_weight_scale.shape - w13_weight_scale = layer.w13_weight_scale.view(s0 * s1, -1) - w13_weight_scale = fp4_utils.e8m0_shuffle(w13_weight_scale) - layer.w13_weight_scale.data = w13_weight_scale.view(s0, s1, -1) - - s0, s1, _ = layer.w2_weight_scale.shape - w2_weight_scale = layer.w2_weight_scale.view(s0 * s1, -1) - w2_weight_scale = fp4_utils.e8m0_shuffle(w2_weight_scale) - layer.w2_weight_scale.data = w2_weight_scale.view(s0, s1, -1) - return - else: - shuffle_weights(layer.w13_weight, layer.w2_weight) - shuffled_w13_scale = fp4_utils.e8m0_shuffle( - layer.w13_weight_scale.view(self.num_experts, -1) - ) - shuffled_w2_scale = fp4_utils.e8m0_shuffle( - layer.w2_weight_scale.view(self.num_experts, -1) - ) + # shuffle weight + layer.w13_weight.data = shuffle_weight( + layer.w13_weight, + is_guinterleave=self.is_guinterleave, + gate_up=True, + ) + layer.w2_weight.data = shuffle_weight( + layer.w2_weight, + is_guinterleave=self.is_guinterleave, + gate_up=False, + ) + layer.w13_weight.is_shuffled = True + layer.w2_weight.is_shuffled = True + + # shuffle scale + w13_scale_2d = layer.w13_weight_scale.reshape( + -1, layer.w13_weight_scale.shape[-1] + ) + w2_scale_2d = layer.w2_weight_scale.reshape(-1, layer.w2_weight_scale.shape[-1]) + + shuffled_w13_scale = shuffle_scale( + w13_scale_2d, self.num_experts, self.is_guinterleave, True + ) + shuffled_w2_scale = shuffle_scale( + w2_scale_2d, self.num_experts, self.is_guinterleave, False + ) layer.w13_weight_scale = atom_parameter(shuffled_w13_scale) layer.w2_weight_scale = atom_parameter(shuffled_w2_scale) @@ -946,7 +921,7 @@ def apply( from atom.model_ops.fused_moe_triton import ( triton_kernel_moe_forward, triton_kernel_fused_experts, - routing_from_topk, + fused_routing_from_topk_triton, ) # Check if the model needs custom routing that triton routing() @@ -975,13 +950,15 @@ def apply( num_fused_shared_experts=layer.num_fused_shared_experts, routed_scaling_factor=layer.routed_scaling_factor, ) + n_expts_act = topk_weights.shape[1] # Convert to triton routing data structures n_expts_tot = router_logits.shape[-1] if global_num_experts > 0: n_expts_tot = global_num_experts + n_expts_tot = n_expts_tot + layer.num_fused_shared_experts - routing_data, gather_idx, scatter_idx = routing_from_topk( + routing_data, gather_idx, scatter_idx = fused_routing_from_topk_triton( topk_weights, topk_ids, n_expts_tot ) @@ -994,7 +971,7 @@ def apply( routing_data, gather_idx, scatter_idx, - topk=top_k, + topk=n_expts_act, activation=activation, w13_precision_config=self.w13_precision_config, w2_precision_config=self.w2_precision_config, @@ -1002,7 +979,7 @@ def apply( w2_bias=layer.w2_bias, swiglu_limit=getattr(layer, "swiglu_limit", 0.0), apply_router_weight_on_input=layer.apply_router_weight_on_input, - global_num_experts=global_num_experts, + global_num_experts=n_expts_tot, expert_map=expert_map, ) return _moe_result @@ -1065,6 +1042,12 @@ def apply( intermediate_pad=self.intermediate_pad, bias1=layer.w13_bias, bias2=layer.w2_bias, + swiglu_limit=getattr(layer, "swiglu_limit", 0.0), + gate_mode=( + GateMode.INTERLEAVE.value + if self.is_guinterleave + else GateMode.SEPARATED.value + ), ) return self.fused_experts( hidden_states=x, @@ -1091,7 +1074,6 @@ def apply( # Refer to CompressedTensorsW8A8Fp8MoEMethod in vllm class CompressedTensorsFp8MoEMethod(FusedMoEMethodBase): - def __init__(self, quant_config: LayerQuantConfig, moe: FusedMoEConfig): super().__init__(moe) self.quant_config = quant_config @@ -1289,12 +1271,14 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: w2_input_scale = getattr(layer, "w2_input_scale", None) if self.need_normalize_e4m3fn_to_e4m3fnuz: - w13.data, w13_scale.data, w13_input_scale_data = ( - normalize_e4m3fn_to_e4m3fnuz( - w13.data, - w13_scale.data, - w13_input_scale.data if w13_input_scale is not None else None, - ) + ( + w13.data, + w13_scale.data, + w13_input_scale_data, + ) = normalize_e4m3fn_to_e4m3fnuz( + w13.data, + w13_scale.data, + w13_input_scale.data if w13_input_scale is not None else None, ) if w13_input_scale is not None and w13_input_scale_data is not None: w13_input_scale.data = w13_input_scale_data @@ -1351,13 +1335,22 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: # Update scale to single max scale per expert [E] layer.w13_weight_scale = atom_parameter(max_w13_scales) - # Shuffle weights for asm moe (moved from inference to load time for better performance) - if w13.dtype in [ - torch.int8, - torch.uint8, - torch.float8_e4m3fnuz, - torch.float8_e4m3fn, - ]: + # Shuffle weights for asm moe (moved from inference to load time for better performance). + # For per_1x128 blockscale (block_quant), only shuffle when the preshuffle GEMM + # path is enabled — the non-preshuffle kernel expects the un-shuffled layout. + skip_shuffle_for_block = ( + self.block_quant and not envs.ATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE + ) + if ( + w13.dtype + in [ + torch.int8, + torch.uint8, + torch.float8_e4m3fnuz, + torch.float8_e4m3fn, + ] + and not skip_shuffle_for_block + ): from aiter.ops.shuffle import shuffle_weight w13.data = shuffle_weight(w13.data) @@ -1667,7 +1660,11 @@ def _process_block_quant(self, layer: nn.Module) -> None: layer.w2_weight = atom_parameter(layer.w2_weight.data) layer.w2_weight_scale = atom_parameter(layer.w2_weight_scale.data) - shuffle_weights(layer.w13_weight, layer.w2_weight) + # per_1x128 blockscale MoE only needs weight bpreshuffle when the + # preshuffle GEMM path is enabled. Skip it to match the non-preshuffle + # kernel's expected weight layout. + if envs.ATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE: + shuffle_weights(layer.w13_weight, layer.w2_weight) def _process_channel_quant(self, layer: nn.Module) -> None: """PTPTC""" @@ -1713,9 +1710,10 @@ def _process_tensor_quant(self, layer: nn.Module) -> None: layer.w13_weight_scale[expert_id][shard_id], ) quant_func = get_hip_quant(self.quant_type) - layer.w13_weight[expert_id][start : start + shard_size, :], _ = ( - quant_func(dq_weight, max_w13_scales[expert_id]) - ) + ( + layer.w13_weight[expert_id][start : start + shard_size, :], + _, + ) = quant_func(dq_weight, max_w13_scales[expert_id]) start += shard_size shuffle_weights(layer.w13_weight, layer.w2_weight) @@ -2646,13 +2644,14 @@ def select_experts( topk_weights = torch.empty( tokens_num, top_k, dtype=torch.float32, device=router_logits.device ) - topk_softplus( + topk_gating( topk_weights, topk_ids, router_logits, e_score_correction_bias, renormalize, routed_scaling_factor, + score_func="sqrtsoftplus", ) else: raise ValueError( diff --git a/atom/model_ops/sampler.py b/atom/model_ops/sampler.py index 14276c3c1c..a1f52d47c4 100644 --- a/atom/model_ops/sampler.py +++ b/atom/model_ops/sampler.py @@ -2,7 +2,6 @@ # Copyright (C) 2024-2025, Advanced Micro Devices, Inc. All rights reserved. import warnings -from functools import lru_cache import torch from aiter import mixed_sample_outer_exponential @@ -32,7 +31,6 @@ SAMPLER_EPS = 1e-10 -@lru_cache(maxsize=1) def get_per_token_exponential(vocab_size: int, device) -> torch.Tensor: """Returns a tensor of shape (1, vocab_size) filled with exponential random values. This is key to deterministic inference, as it ensures that the same random values are used for each token across different runs. diff --git a/atom/model_ops/triton_fused_sigmoid_mul_quant.py b/atom/model_ops/triton_fused_sigmoid_mul_quant.py index 7829540521..baf9c6e1cf 100644 --- a/atom/model_ops/triton_fused_sigmoid_mul_quant.py +++ b/atom/model_ops/triton_fused_sigmoid_mul_quant.py @@ -120,6 +120,7 @@ def fused_sigmoid_mul_fp8_quant( attn_output: Tensor, gate: Tensor, group_size: int = 128, + transpose_scale: bool | None = None, ) -> tuple[Tensor, Tensor]: """Fused sigmoid(gate) * attn_output + FP8 per-group quantization. @@ -127,13 +128,20 @@ def fused_sigmoid_mul_fp8_quant( attn_output: [M, N] bf16/fp16 — attention output tensor. gate: [M, N] bf16/fp16 — gating tensor (pre-sigmoid). group_size: Quantization group size (default 128, matching per_1x128). + transpose_scale: If True, produce column-major x_scale (for preshuffle GEMM). + If False, produce row-major x_scale (for non-preshuffle GEMM). + If None (default), follows ATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE env var. Returns: (x_fp8, x_scale): x_fp8: [M, N] FP8 — quantized sigmoid(gate) * attn_output. - x_scale: [M, N // group_size] float32 — per-group scales in - transpose_scale layout (column-major storage). + x_scale: [M, N // group_size] float32 — per-group scales. """ + if transpose_scale is None: + from atom.utils import envs + + transpose_scale = envs.ATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE + M, N = attn_output.shape assert ( N % group_size == 0 @@ -144,11 +152,21 @@ def fused_sigmoid_mul_fp8_quant( out_fp8 = torch.empty((M, N), dtype=fp8_dtype, device=attn_output.device) num_scale_cols = N // group_size - # transpose_scale layout: allocate as (num_scale_cols, M) contiguous, - # then view as (M, num_scale_cols) after kernel - out_scale = torch.empty( - (num_scale_cols, M), dtype=torch.float32, device=attn_output.device - ) + if transpose_scale: + # column-major: allocate as (num_scale_cols, M) contiguous, + # then view as (M, num_scale_cols) after kernel + out_scale = torch.empty( + (num_scale_cols, M), dtype=torch.float32, device=attn_output.device + ) + stride_scale_m = out_scale.stride(1) + stride_scale_n = out_scale.stride(0) + else: + # row-major: allocate as (M, num_scale_cols) contiguous + out_scale = torch.empty( + (M, num_scale_cols), dtype=torch.float32, device=attn_output.device + ) + stride_scale_m = out_scale.stride(0) + stride_scale_n = out_scale.stride(1) grid = (M, N // BLOCK_SIZE_N) _fused_sigmoid_mul_fp8_group_quant_kernel[grid]( @@ -162,9 +180,8 @@ def fused_sigmoid_mul_fp8_quant( gate.stride(1), out_fp8.stride(0), out_fp8.stride(1), - # transpose_scale: swap strides so kernel writes in column-major - out_scale.stride(1), # stride_scale_m = stride along dim 1 (M dim) - out_scale.stride(0), # stride_scale_n = stride along dim 0 (scale_col dim) + stride_scale_m, + stride_scale_n, N=N, BLOCK_SIZE_N=BLOCK_SIZE_N, QUANT_BLOCK_SIZE=group_size, @@ -172,7 +189,8 @@ def fused_sigmoid_mul_fp8_quant( FP8_MIN=DTYPE_MIN, ) - # View transposed buffer back to (M, num_scale_cols) shape - out_scale = out_scale.view(M, num_scale_cols) + if transpose_scale: + # View transposed buffer back to (M, num_scale_cols) shape + out_scale = out_scale.view(M, num_scale_cols) return out_fp8, out_scale diff --git a/atom/model_ops/v4_kernels/__init__.py b/atom/model_ops/v4_kernels/__init__.py index 46d356da1d..da0c9b08c3 100644 --- a/atom/model_ops/v4_kernels/__init__.py +++ b/atom/model_ops/v4_kernels/__init__.py @@ -31,6 +31,10 @@ sparse_attn_v4_paged_prefill_reference, ) from atom.model_ops.v4_kernels.inverse_rope import inverse_rope_inplace +from atom.model_ops.v4_kernels.paged_decode_indices import ( + write_v4_paged_decode_indices, + write_v4_paged_decode_indices_reference, +) from atom.model_ops.v4_kernels.state_writes import update_compressor_states, swa_write __all__ = [ @@ -48,4 +52,6 @@ "make_compress_plans", "inverse_rope_inplace", "scale_indexer_weights", + "write_v4_paged_decode_indices", + "write_v4_paged_decode_indices_reference", ] diff --git a/atom/model_ops/v4_kernels/compress_plan.py b/atom/model_ops/v4_kernels/compress_plan.py index b7cde3b423..8e7fe2b4b2 100644 --- a/atom/model_ops/v4_kernels/compress_plan.py +++ b/atom/model_ops/v4_kernels/compress_plan.py @@ -57,54 +57,77 @@ class CompressPlan: def make_compress_plans( extend_lens_cpu: np.ndarray, - seq_lens_cpu: np.ndarray, + context_lens_cpu: np.ndarray, unique_ratios_overlap: Iterable[Tuple[int, bool]], - device: torch.device, *, - plan_buffers: dict | None = None, + plan_buffers: dict, + decode_capacity_per_ratio: dict[int, int] | None = None, ) -> dict[int, CompressPlan]: """Build a CompressPlan per (ratio, overlap) variant. Args: - extend_lens_cpu: np[bs] int — number of tokens this fwd processes per seq. - seq_lens_cpu: np[bs] int — absolute seq_len per seq (= prefix + extend). + extend_lens_cpu: np[bs] int — number of tokens this fwd processes per seq. + context_lens_cpu: np[bs] int — absolute seq_len per seq AFTER the new + extend tokens (= prefix + extend). Internally + reconstructs prefix via `context_lens - extend_lens`. unique_ratios_overlap: iterable of (ratio, is_overlap) pairs. Typically {(4, True), (128, False)} for V4-Pro; a subset for models with only CSA or only HCA layers. - device: target GPU device for the GPU-side plan tensors. - plan_buffers: optional dict[ratio] -> {"compress": CpuGpuBuffer, + plan_buffers: required dict[ratio] -> {"compress": CpuGpuBuffer, "write": CpuGpuBuffer} of pre-allocated fixed-capacity - plan buffers (CUDAGraph path). When provided, the - function writes into the existing buffers and sentinel- - fills (-1) the trailing rows beyond the actual count, so - the returned `compress_plan_gpu` / `write_plan_gpu` views - have stable data pointers across calls. When None, the - legacy fresh-allocation path is used (eager only). + plan buffers. Function writes into the existing buffers + and sentinel-fills (-1) trailing rows beyond the actual + count, so the returned `compress_plan_gpu` / + `write_plan_gpu` views have stable data pointers across + calls (CUDAGraph requirement). Fresh per-call alloc is + not supported — that pattern caused allocator-churn + races (see `write_v4_paged_decode_indices` docstring). + decode_capacity_per_ratio: optional dict[ratio] -> int. Slice length + for the returned `compress_plan_gpu`. When PROVIDED: + the slice is exactly that fixed value (independent of + `n_compress`) — required by the decode CUDAGraph path + so capture and replay produce kernel calls with + identical tensor shapes. Caller must size this to the + decode worst case (`max_decode_tokens // ratio + bs`). + When NONE: the slice is exactly `n_compress` (tight, + eager-only) — gives the smallest possible kernel grid. + The slice is contiguous-from-base so the data pointer + is stable; only `shape[0]` shrinks. The underlying + buffer is always sized to the prefill worst case. Returns: - dict[ratio] -> CompressPlan. Empty dict if `extend_lens_cpu.sum() == 0` - AND no plan_buffers are provided. With plan_buffers, an empty fwd still - returns CompressPlans pointing at the pre-allocated buffers (fully - sentinel-filled), so capture-time addresses match replay-time addresses - even on a zero-token fwd. + dict[ratio] -> CompressPlan. On empty fwd (`extend_lens_cpu.sum() == 0`) + still returns CompressPlans pointing at the pre-allocated buffers + (fully sentinel-filled), so capture-time addresses match replay-time + addresses even on a zero-token fwd. """ bs = len(extend_lens_cpu) extend_lens_cpu = np.ascontiguousarray(extend_lens_cpu, dtype=np.int32) - seq_lens_cpu = np.ascontiguousarray(seq_lens_cpu, dtype=np.int32) + context_lens_cpu = np.ascontiguousarray(context_lens_cpu, dtype=np.int32) total = int(extend_lens_cpu.sum()) out: dict[int, CompressPlan] = {} if total == 0 or bs == 0: - if plan_buffers is None: - return out - # Capture path with empty fwd: fully sentinel-fill the buffers and - # return CompressPlans so addresses are stable. Skipped via num_*=0. + # Empty fwd: produce CompressPlans pointing at the pre-allocated + # buffers so capture-time addresses match replay-time addresses + # even on a zero-token fwd. Skipped via num_*=0. + # + # CG path (decode_capacity_per_ratio provided): slice = fixed cap. + # Eager path (None): slice = 0 — kernel grid is empty, no work + # dispatched, no sentinel fill required. for ratio, _ in unique_ratios_overlap: cbuf = plan_buffers[ratio]["compress"] wbuf = plan_buffers[ratio]["write"] - cbuf.np[:].fill(-1) + cap = ( + decode_capacity_per_ratio.get(ratio) + if decode_capacity_per_ratio is not None + else 0 + ) + if cap > 0: + cbuf.np[:cap].fill(-1) + compress_plan_gpu = cbuf.copy_to_gpu(cap if cap > 0 else 0) wbuf.np[:].fill(-1) out[ratio] = CompressPlan( - compress_plan_gpu=cbuf.copy_to_gpu(), + compress_plan_gpu=compress_plan_gpu, write_plan_gpu=wbuf.copy_to_gpu(), num_compress=0, num_write=0, @@ -120,7 +143,7 @@ def make_compress_plans( cu_extend[0] = 0 np.cumsum(extend_lens_cpu, out=cu_extend[1:]) j_in_seq = ragged_ids - cu_extend[batch_ids] - prefix_lens = seq_lens_cpu - extend_lens_cpu + prefix_lens = context_lens_cpu - extend_lens_cpu positions = prefix_lens[batch_ids] + j_in_seq for ratio, is_overlap in unique_ratios_overlap: @@ -147,45 +170,50 @@ def make_compress_plans( # write: tokens whose absolute position falls in the per-seq # "last STATE_SIZE positions" window. STATE_SIZE = K. - # write_start[i] = max(0, seq_lens[i] - K) — uniform across overlap/non-overlap; + # write_start[i] = max(0, context_lens[i] - K) — uniform across overlap/non-overlap; # the SGLang formula `(seq_len // ratio) * ratio - (ratio if overlap else 0)` # is a stricter bound that includes only ratio-aligned writes; the looser - # `seq_len - K` is what ATOM's update_compressor_states already uses + # `context_len - K` is what ATOM's update_compressor_states already uses # (state_writes.py:152-154 docstring) and what the fused kernel's # state-cache reader expects. - write_starts = np.maximum(0, seq_lens_cpu - K).astype(np.int32) + write_starts = np.maximum(0, context_lens_cpu - K).astype(np.int32) write_mask = positions >= write_starts[batch_ids] write_plan = plan_rows[write_mask] n_compress = int(compress_plan.shape[0]) n_write = int(write_plan.shape[0]) - if plan_buffers is not None: - cbuf = plan_buffers[ratio]["compress"] - wbuf = plan_buffers[ratio]["write"] - assert n_compress <= cbuf.np.shape[0], ( - f"ratio={ratio} num_compress={n_compress} exceeds buffer " - f"capacity {cbuf.np.shape[0]}; bump in builder __init__." - ) - assert n_write <= wbuf.np.shape[0], ( - f"ratio={ratio} num_write={n_write} exceeds buffer " - f"capacity {wbuf.np.shape[0]}; bump in builder __init__." - ) - if n_compress > 0: - cbuf.np[:n_compress] = compress_plan - cbuf.np[n_compress:].fill(-1) # sentinel - if n_write > 0: - wbuf.np[:n_write] = write_plan - wbuf.np[n_write:].fill(-1) # sentinel - compress_plan_gpu = cbuf.copy_to_gpu() - write_plan_gpu = wbuf.copy_to_gpu() - else: - compress_plan_gpu = torch.from_numpy( - np.ascontiguousarray(compress_plan) - ).to(device, non_blocking=True) - write_plan_gpu = torch.from_numpy(np.ascontiguousarray(write_plan)).to( - device, non_blocking=True - ) + cbuf = plan_buffers[ratio]["compress"] + wbuf = plan_buffers[ratio]["write"] + full_cap = cbuf.np.shape[0] + # CG path: fixed slice (capture / replay must match). Eager + # path: slice = n_compress (smallest possible kernel grid). + cap = ( + decode_capacity_per_ratio.get(ratio) + if decode_capacity_per_ratio is not None + else None + ) + slice_cap = n_compress if cap is None else cap + assert n_compress <= slice_cap <= full_cap, ( + f"ratio={ratio} num_compress={n_compress}, slice={slice_cap}, " + f"buffer={full_cap}: invariant violated. CG path requires " + f"n_compress ≤ decode_cap; eager path uses n_compress as slice." + ) + assert n_write <= wbuf.np.shape[0], ( + f"ratio={ratio} num_write={n_write} exceeds buffer " + f"capacity {wbuf.np.shape[0]}; bump in builder __init__." + ) + if n_compress > 0: + cbuf.np[:n_compress] = compress_plan + # Sentinel only within the slice we hand to the kernel; rows + # beyond `slice_cap` are unreachable from this launch. + if slice_cap > n_compress: + cbuf.np[n_compress:slice_cap].fill(-1) + if n_write > 0: + wbuf.np[:n_write] = write_plan + wbuf.np[n_write:].fill(-1) # sentinel + compress_plan_gpu = cbuf.copy_to_gpu(slice_cap) + write_plan_gpu = wbuf.copy_to_gpu() out[ratio] = CompressPlan( compress_plan_gpu=compress_plan_gpu, diff --git a/atom/model_ops/v4_kernels/fused_compress.py b/atom/model_ops/v4_kernels/fused_compress.py index 0e26b7eb8a..3005ab6b0f 100644 --- a/atom/model_ops/v4_kernels/fused_compress.py +++ b/atom/model_ops/v4_kernels/fused_compress.py @@ -10,12 +10,14 @@ Each compression boundary across the entire fwd is one row in `compress_plan_gpu` — a packed `[num_compress, 4] int32` tensor where each row is `[ragged_id, batch_id, position, window_len]`. The kernel grid is - `num_compress` (zero waste, no early-exit), and each program does ONE 4×i32 - load to get all the metadata it needs: + the caller-supplied slice length (decode CG: `_decode_compress_cap[ratio]` + / eager prefill: `n_compress`); inactive plan rows are sentinel-marked + (`position == -1`) and bail at the top of the kernel. Each program does + ONE 4×i32 load to get all the metadata it needs: ragged_id → row index in the ragged kv_in / score_in stream batch_id → seq index → state_slot_mapping[batch_id], block_table[batch_id] - position → absolute token position (drives RoPE + paged scatter) + position → absolute token position (drives RoPE + paged scatter); -1 = skip window_len → number of leading K-loop iterations that read state cache (instead of the ragged input). K = STATE_SIZE. @@ -32,18 +34,21 @@ the caller MUST invoke this kernel BEFORE `update_compressor_states` runs (which would overwrite the historic positions this kernel needs). -Output: - Returns `[num_compress, head_dim]` BF16 tensor. Rows are in plan order - (= ragged_id ascending = per-seq grouped). Caller slices per seq via - `plan.cu_compress_cpu[i]:plan.cu_compress_cpu[i+1]`. - - Also writes valid rows into the paged kv_cache at compressed index - `position // RATIO` (when `block_table` is provided). - -TODO: FP8/FP4 quant fusion. Currently stores raw BF16 (no act_quant). The -`scale_fmt="ue8m0"` round-to-even f32→e8m0 path requires porting aiter's -`f32_to_e8m0` bit manipulation into Triton (~20 lines per quant block). -Skipping for this PR; follow-up PR will add it. +Quant modes (constexpr-selected by the `quant` arg of the Python wrapper): + - quant=False (CSA Main / HCA Main): writes BF16 rows into the paged BF16 + `kv_cache` at compressed slot `position // ratio`. + - quant=True (CSA Indexer-inner): per-row amax → ue8m0 (or raw) scale → + fp8 cast → preshuffled (MFMA 16x16 tile) write into the FP8 `kv_cache`, + plus fp32 scale into the per-block scale region (`cache_scale` — a + strided view of the same allocation built by the V4 builder). Bit-exact + match for `indexer_k_quant_and_cache` / + `cp_gather_indexer_k_quant_cache` (cache_kernels.cu:1145+). + +Output: side-effecting only — cache scatter IS the only output. The earlier +caller-visible `[num_compress, head_dim]` BF16 return tensor was vestigial +(paged_decode/paged_prefill read the scattered compress entries directly +from `unified_kv` (Main) or the FP8 indexer pool, not from the kernel +return). """ from typing import Optional @@ -83,15 +88,14 @@ def _fused_compress_attn_kernel( sin_cache_ptr, cos_sin_pos_stride, # = rope_head_dim // 2 # ── KV cache scatter (paged) ──────────────────────────────────────── - kv_cache_ptr, # [num_blocks, k_per_block, head_dim] bf16 + kv_cache_ptr, # bf16: [NB, k_per_block, head_dim] / fp8: [NB, k_per_block, head_dim] kv_cache_block_stride, kv_cache_token_stride, + cache_scale_ptr, # fp32 [NB, k_per_block] (QUANT path only; dummy ptr otherwise) + cache_scale_block_stride, # fp32 elements per block (= k_per_block * cache_stride / 4) block_table_ptr, # [bs, max_blocks_per_seq] int32 block_table_seq_stride, # row stride k_per_block, - # ── output to caller (post norm+rope BF16 for sparse_attn input) ──── - out_ptr, # [num_compress, head_dim] bf16 - out_token_stride, head_dim, rope_head_dim, # ── constexpr ─────────────────────────────────────────────────────── @@ -99,13 +103,25 @@ def _fused_compress_attn_kernel( HALF_ROPE: tl.constexpr, # = rope_head_dim // 2 OVERLAP: tl.constexpr, RATIO: tl.constexpr, - STATE_SIZE: tl.constexpr, # = 2*RATIO if OVERLAP else RATIO - K: tl.constexpr, # = STATE_SIZE (softmax-pool reduce dim) + STATE_SIZE: tl.constexpr, # ring buffer modulo = kv_state.shape[1] (≥ K_pool; + # for spec decode is K_pool + max_spec_steps + 1 to avoid R+1 re-commit + # borrow-reads of prior round's reject K/V; for non-spec decode is K_pool) + K: tl.constexpr, # pool-window reduce dim (= 2*RATIO if OVERLAP else RATIO); + # ≤ STATE_SIZE; used for `s = position - K + 1 + k_static` loop bound HAS_BLOCK_TABLE: tl.constexpr, + QUANT: tl.constexpr, # 0 = raw BF16 (CSA/HCA Main), 1 = FP8 e4m3 + ue8m0 scale (Indexer) + USE_UE8M0: tl.constexpr, # round scale to power-of-2 (only when QUANT == 1) + PRESHUFFLE: tl.constexpr, # MFMA 16x16 preshuffled FP8 layout (only when QUANT == 1) + # E4M3 max (=448 for E4M3FN, =240 for E4M3FNUZ). Constexpr so the clamp + # bounds and reciprocal fold to compile-time constants. Ignored when + # QUANT == 0 (caller passes 1.0 as a placeholder). + FP8_MAX: tl.constexpr = 1.0, ): - """One program per boundary in the plan. Grid = plan capacity (CUDAGraph- - safe fixed grid); inactive rows are sentinel-marked (position == -1) and - bail before any load/store/scatter.""" + """One program per boundary in the plan. Grid = caller-supplied slice + length (decode CG: `_decode_compress_cap[ratio]` for capture/replay + address stability; eager prefill: tight `n_compress`). Inactive rows + are sentinel-marked (position == -1) and bail before any load / + store / scatter.""" pid = tl.program_id(0) plan_base = plan_ptr + pid * 4 ragged_id = tl.load(plan_base + 0) @@ -122,57 +138,41 @@ def _fused_compress_attn_kernel( d_mask = d < head_dim # ── 1. Per-source-position load + online softmax-pool ────────────── + # Two-phase split (vs. the older single masked loop): for each program + # `is_input = k_static >= window_len` partitions the K iterations into + # a leading state-cache-only run and a trailing input-only run. Issuing + # only the live side's loads (rather than masked-off both) cuts HBM + # bandwidth ~40% on AMD CDNA where masked tl.load still issues the LD + # instruction (predicate only suppresses register write-back). + # + # Padding invariant: padding (`s < 0` ⟺ `k_static < K-1-position`) lies + # entirely within `k_static < window_len` because + # `window_len = K - min(j_in_seq+1, K) ≥ K - 1 - j_in_seq ≥ K - 1 - position` + # (`position = prefix_len + j_in_seq`, `prefix_len ≥ 0`). The input phase + # therefore needs no padding mask. NEG_INF: tl.constexpr = float("-inf") m_acc = tl.full([BLOCK_D], NEG_INF, tl.float32) kv_acc = tl.zeros([BLOCK_D], tl.float32) w_acc = tl.zeros([BLOCK_D], tl.float32) - # Dynamic loop (NOT unrolled) — K=128 (HCA) would otherwise blow up hsaco. - for k_static in tl.range(K): + # ── Phase 1: state cache (k_static ∈ [0, window_len)) ── + # Dynamic bound (window_len is per-program, not constexpr) — Triton + # cannot static-unroll this. Loop body issues only state-side loads. + for k_static in tl.range(0, window_len): s = position - K + 1 + k_static is_padding = s < 0 - is_input = k_static >= window_len - # is_state = (~is_input) & (~is_padding) - - # B-side (k < RATIO): cols [:head_dim] (= col_off=0) - # A-side (k >= RATIO): cols [head_dim:] (= col_off=head_dim) - # HCA (no overlap, K=RATIO): col_off=0 always (k_static < RATIO). + # B-side (k >= RATIO): cols [head_dim:]; A-side (k < RATIO): cols [:head_dim]. + # HCA (no overlap, K=RATIO): col_off=0 always. col_off = (k_static >= RATIO) * head_dim if OVERLAP else 0 - ape_row = k_static % RATIO # [0, RATIO) — same for B/A sides - # Input source: row index in ragged stream. - # k_static = K-1 corresponds to s = position (the boundary token itself, - # = ragged_id row). Earlier k_static values map to earlier ragged rows. - in_row = ragged_id - (K - 1 - k_static) - kv_a = tl.load( - kv_in_ptr + in_row * kv_in_row_stride + col_off + d, - mask=is_input & d_mask, - other=0.0, - ) - - # State cache source: ring slot indexed by absolute s. s_safe = tl.maximum(s, 0) ring = s_safe % STATE_SIZE - kv_b = tl.load( - kv_state_ptr - + slot * kv_state_slot_stride - + ring * kv_state_pos_stride - + col_off - + d, - mask=(~is_input) & (~is_padding) & d_mask, - other=0.0, - ) - kv_k = kv_a + kv_b # exactly one path active per source pos - - # ── score_k load ── - score_a = tl.load( - score_in_ptr + in_row * score_in_row_stride + col_off + d, - mask=is_input & d_mask, - other=NEG_INF, + state_row_off = ( + slot * kv_state_slot_stride + ring * kv_state_pos_stride + col_off ) - ape_v = tl.load( - ape_ptr + ape_row * dim_full + col_off + d, - mask=is_input & d_mask, + kv_b = tl.load( + kv_state_ptr + state_row_off + d, + mask=(~is_padding) & d_mask, other=0.0, ) score_b = tl.load( @@ -181,16 +181,55 @@ def _fused_compress_attn_kernel( + ring * score_state_pos_stride + col_off + d, - mask=(~is_input) & (~is_padding) & d_mask, + mask=(~is_padding) & d_mask, other=NEG_INF, ) - score_k = tl.where(is_input, score_a + ape_v, score_b) - # ── Online softmax-pool accumulate ── + m_new = tl.maximum(m_acc, score_b) + scale = tl.where(m_acc == NEG_INF, 0.0, tl.exp(m_acc - m_new)) + # Padding lanes have score_b = NEG_INF → w_k = 0, contributes nothing. + w_k = tl.where(score_b == NEG_INF, 0.0, tl.exp(score_b - m_new)) + kv_acc = kv_acc * scale + w_k * kv_b + w_acc = w_acc * scale + w_k + m_acc = m_new + + # ── Phase 2: ragged input (k_static ∈ [window_len, K)) ── + # No padding here (per invariant above). All loads unconditional in + # the position dimension; only the head_dim mask remains. + for k_static in tl.range(window_len, K): + col_off = (k_static >= RATIO) * head_dim if OVERLAP else 0 + ape_row = k_static % RATIO + # k_static = K-1 → s = position (the boundary token itself, + # = ragged_id row). Earlier k_static map to earlier ragged rows. + in_row = ragged_id - (K - 1 - k_static) + + # kv_in / score_in: single-use per program → evict_first to keep + # state-cache lines (small, possibly shared across programs) hot. + kv_a = tl.load( + kv_in_ptr + in_row * kv_in_row_stride + col_off + d, + mask=d_mask, + other=0.0, + eviction_policy="evict_first", + ).to(tl.float32) + score_a = tl.load( + score_in_ptr + in_row * score_in_row_stride + col_off + d, + mask=d_mask, + other=0.0, + eviction_policy="evict_first", + ).to(tl.float32) + ape_v = tl.load( + ape_ptr + ape_row * dim_full + col_off + d, + mask=d_mask, + other=0.0, + eviction_policy="evict_last", + ) + score_k = score_a + ape_v + m_new = tl.maximum(m_acc, score_k) scale = tl.where(m_acc == NEG_INF, 0.0, tl.exp(m_acc - m_new)) - w_k = tl.where(score_k == NEG_INF, 0.0, tl.exp(score_k - m_new)) - kv_acc = kv_acc * scale + w_k * kv_k + # score_k always finite in input phase → no NEG_INF guard needed. + w_k = tl.exp(score_k - m_new) + kv_acc = kv_acc * scale + w_k * kv_a w_acc = w_acc * scale + w_k m_acc = m_new @@ -231,26 +270,80 @@ def _fused_compress_attn_kernel( new_odd = odd_v * cos_per_pair + even_v * sin_per_pair rotated = tl.interleave(new_even, new_odd) # [BLOCK_D] fp32 - # ── 4. Cast to BF16 + store ──────────────────────────────────────── - rotated_bf16 = rotated.to(tl.bfloat16) - - # Output: row pid (plan order = per-seq grouped, caller uses cu_compress). - tl.store(out_ptr + pid * out_token_stride + d, rotated_bf16, mask=d_mask) - - # KV cache scatter (paged): block_table[batch_id, ci // k_per_block][ci % k_per_block] + # ── 4. Cache scatter (paged) ─────────────────────────────────────── + # The Compressor's BF16 return value was historically consumed by sparse + # attention but is now vestigial — paged_decode/paged_prefill read the + # scattered compress entries directly from `unified_kv` (Main) or the FP8 + # indexer pool. So no caller-visible `out` write; the cache scatter IS + # the only output. if HAS_BLOCK_TABLE: + # block_table[batch_id, ci // k_per_block] → physical_block. + # slot_in_block = ci % k_per_block. Same address resolution as + # `indexer_k_quant_and_cache` (which used a pre-flattened slot_mapping + # = physical_block * k_per_block + slot_in_block). ci = position // RATIO block_in_seq = ci // k_per_block slot_in_block = ci % k_per_block physical_block = tl.load( block_table_ptr + batch_id * block_table_seq_stride + block_in_seq ).to(tl.int64) - cache_addr = ( - physical_block * kv_cache_block_stride - + slot_in_block * kv_cache_token_stride - + d - ) - tl.store(kv_cache_ptr + cache_addr, rotated_bf16, mask=d_mask) + + if QUANT: + # FP8 e4m3 quantization: per-row amax → ue8m0 (or raw) scale → + # clamp+cast to fp8 → write preshuffled (MFMA 16x16 tile) or + # linear layout into the FP8 region; write the fp32 scale into + # the per-block scale region (separate `cache_scale_ptr` view). + # Bit-exact match for `indexer_k_quant_and_cache` / + # `cp_gather_indexer_k_quant_cache` (cache_kernels.cu:1145+). + # + # Quant pattern follows aiter's `_fp8_quant_op` / + # `_fused_rms_gated_fp8_group_quant_kernel` — the explicit + # `fp_downcast_rounding="rtne"` is intentionally NOT used here: + # on AMD it forces a slow software-RTNE path and bypasses the + # `v_cvt_pk_fp8_f32` HW intrinsic (which already rounds RTNE). + # The pre-cast `tl.clamp` saturates intermediate overflow and + # mirrors aiter's hand-tuned HIP `aiter::scaled_cast`. + rotated_for_amax = tl.where(d_mask, tl.abs(rotated), 0.0) + amax = tl.max(rotated_for_amax, axis=0) # scalar (per row) + scale = tl.maximum(amax, 1e-4) * (1.0 / FP8_MAX) + if USE_UE8M0: + scale = tl.exp2(tl.ceil(tl.log2(scale))) + inv_scale = 1.0 / scale + scaled = tl.clamp(rotated * inv_scale, -FP8_MAX, FP8_MAX) + fp8_val = scaled.to(kv_cache_ptr.dtype.element_ty) + if PRESHUFFLE: + TILE: tl.constexpr = 16 + token_tile_id = slot_in_block // TILE + token_in_tile = slot_in_block % TILE + col_tile_id = d // TILE + col_in_tile = d % TILE + fp8_offset = ( + physical_block * kv_cache_block_stride + + token_tile_id * (TILE * head_dim) + + col_tile_id * (TILE * TILE) + + token_in_tile * TILE + + col_in_tile + ) + else: + fp8_offset = ( + physical_block * kv_cache_block_stride + + slot_in_block * head_dim + + d + ) + # Streaming write — these slots aren't reread inside this kernel. + tl.store( + kv_cache_ptr + fp8_offset, fp8_val, mask=d_mask, cache_modifier=".cs" + ) + # Scale: one fp32 per row, packed at end of block. + scale_offset = physical_block * cache_scale_block_stride + slot_in_block + tl.store(cache_scale_ptr + scale_offset, scale, cache_modifier=".cs") + else: + cache_addr = ( + physical_block * kv_cache_block_stride + + slot_in_block * kv_cache_token_stride + + d + ) + tl.store(kv_cache_ptr + cache_addr, rotated.to(tl.bfloat16), mask=d_mask) def fused_compress_attn( @@ -270,7 +363,9 @@ def fused_compress_attn( cos_cache: torch.Tensor, # [max_seq, ..., rope_head_dim/2] bf16/fp16 sin_cache: torch.Tensor, # same shape # KV cache scatter - kv_cache: Optional[torch.Tensor], # [num_blocks, k_per_block, head_dim] bf16 + kv_cache: Optional[ + torch.Tensor + ], # bf16: [NB, k_per_block, head_dim] / fp8: same shape, fp8 block_tables: Optional[torch.Tensor], # [bs, max_blocks_per_seq] int32 k_per_block: int, # Geometry @@ -278,41 +373,56 @@ def fused_compress_attn( ratio: int, head_dim: int, rope_head_dim: int, - out_dtype: torch.dtype = torch.bfloat16, - out: Optional[torch.Tensor] = None, -) -> Optional[torch.Tensor]: - """Batched fused per-source-position pool + RMSNorm + RoPE + bf16 kv_cache - scatter, dispatched via SGLang-style packed plan. - - Returns `[num_compress, head_dim]` BF16 tensor in plan order (= ragged_id - ascending = per-seq grouped). Caller uses `plan.cu_compress_cpu[i:i+2]` to - slice per-seq chunks. - - Returns None if `plan.num_compress == 0` AND no `out` buffer is provided. - When `out` is supplied (CUDAGraph path), kernel ALWAYS launches at full - plan capacity — inactive rows are sentinel-skipped inside the kernel — - and `out` is returned unchanged when `num_compress == 0`. + # FP8 quant fusion (Indexer-inner Compressor path) + quant: bool = False, + cache_scale: Optional[ + torch.Tensor + ] = None, # fp32 [NB, k_per_block]; required when quant=True + use_ue8m0: bool = True, # round scale to power-of-2 (UE8M0); only when quant=True + preshuffle: bool = True, # MFMA 16x16 preshuffled FP8 layout; only when quant=True + fp8_max: Optional[float] = None, # E4M3 max; required when quant=True +) -> None: + """Batched fused per-source-position pool + RMSNorm + RoPE + cache scatter, + dispatched via SGLang-style packed plan. + + Two scatter modes (constexpr-selected by `quant`): + + - quant=False (CSA Main / HCA Main): writes BF16 rows into the paged + BF16 kv_cache (compressed slot = `position // ratio`). + + - quant=True (CSA Indexer-inner): per-row amax → ue8m0 scale → fp8 cast + → preshuffled (MFMA 16x16 tile) write into the FP8 kv_cache, plus + fp32 scale into `cache_scale` (the per-block scale region of the + same allocation). Bit-exact match for `indexer_k_quant_and_cache` / + `cp_gather_indexer_k_quant_cache` (cache_kernels.cu:1145+). + + Side-effecting: cache scatter IS the only output (Main path's BF16 + return tensor was vestigial — paged_decode/paged_prefill read directly + from `unified_kv` and the indexer FP8 pool, not from the kernel return). + Grid is always `plan_capacity` (CUDAGraph-safe); inactive plan rows are + sentinel-skipped (`position == -1`) inside the kernel. Caller MUST invoke BEFORE `update_compressor_states` (state cache reads must see previous-fwd data). """ - device = kv_in.device - num_compress = plan.num_compress plan_capacity = plan.compress_plan_gpu.shape[0] + num_compress = plan.num_compress if plan_capacity == 0: - return out # nothing to do; out (or None) returned as-is. - if num_compress == 0 and out is None: - # Eager path: nothing to do, no caller buffer to fill. - return None + return # nothing to do — no plan rows ever populated. # Validate shapes dim_full = (2 if overlap else 1) * head_dim - state_size = (2 if overlap else 1) * ratio + K_pool = (2 if overlap else 1) * ratio # pool window size (algorithm-defined) + state_size = kv_state.shape[ + 1 + ] # ring buffer modulo (≥ K_pool; spec path enlarges to K_pool + max_spec_steps + 1) assert ( kv_in.dim() == 2 and kv_in.shape[1] == dim_full ), f"kv_in {kv_in.shape}, expected [*, {dim_full}]" assert score_in.shape == kv_in.shape - assert kv_state.shape[1] == state_size and kv_state.shape[2] == dim_full + assert ( + state_size >= K_pool and kv_state.shape[2] == dim_full + ), f"kv_state {kv_state.shape}, expected [*, ≥{K_pool}, {dim_full}]" assert score_state.shape == kv_state.shape assert ape.shape == (ratio, dim_full) assert rms_weight.shape == (head_dim,) @@ -345,30 +455,45 @@ def fused_compress_attn( else: bt_seq_stride = 0 - if out is None: - # Eager path: allocate output sized to the actual num_compress so the - # returned tensor has the legacy [num_compress, head_dim] shape. - out = torch.empty(num_compress, head_dim, dtype=out_dtype, device=device) - else: - # CUDAGraph path: caller-provided buffer of capacity ≥ plan_capacity. - # Validate shape; kernel will write into rows [0:num_compress) and - # leave [num_compress:plan_capacity) untouched (those plan rows are - # sentinel-skipped). Inactive output rows therefore carry stale data - # but no consumer reads them (caller slices via cu_compress_cpu). + # Quant validation. quant=True is FP8 cache write (Indexer-inner path); + # requires a valid block_tables (slot resolution) AND a paired fp32 scale + # view of the same allocation (see Compressor.forward for how to slice it). + if quant: + assert has_bt, "quant=True requires block_tables for slot resolution" assert ( - out.shape[0] >= plan_capacity and out.shape[1] == head_dim - ), f"out {tuple(out.shape)}, expected ≥ ({plan_capacity}, {head_dim})" - assert out.dtype == out_dtype - assert out.is_contiguous() or out.stride(1) == 1 + kv_cache.dtype != torch.bfloat16 + ), f"quant=True expects an FP8/uint8 kv_cache; got {kv_cache.dtype}" + assert ( + cache_scale is not None and cache_scale.dtype == torch.float32 + ), "quant=True requires `cache_scale` (fp32 [NB, k_per_block])" + assert cache_scale.dim() == 2 and cache_scale.shape[0] == kv_cache.shape[0] + assert fp8_max is not None and fp8_max > 0 + if preshuffle: + assert ( + head_dim % 16 == 0 + ), f"preshuffle requires head_dim%16==0, got {head_dim}" + assert ( + k_per_block % 16 == 0 + ), f"preshuffle requires k_per_block%16==0, got {k_per_block}" BLOCK_D = triton.next_power_of_2(head_dim) HALF_ROPE = rope_head_dim // 2 - K = state_size + K = K_pool # pool window reduce-dim (constexpr; not equal to ring modulo) + + # Cache-scale args (only consumed by the quant path; pass placeholders + # otherwise so the constexpr branch is never taken). + if quant: + cache_scale_arg = cache_scale + cache_scale_block_stride_arg = cache_scale.stride(0) + fp8_max_arg = float(fp8_max) + else: + cache_scale_arg = state_slot_mapping # placeholder int32 ptr (unused) + cache_scale_block_stride_arg = 0 + fp8_max_arg = 1.0 # placeholder; FP8_MAX is constexpr, must be > 0 # Fixed grid for CUDAGraph compat: launch one program per plan row; - # sentinel rows (position=-1) skipped inside the kernel. When out is - # eager-allocated above, grid still shrinks to num_compress (no pad). - grid = (plan_capacity if out.shape[0] >= plan_capacity else num_compress,) + # sentinel rows (position=-1) skip inside the kernel. + grid = (plan_capacity,) _fused_compress_attn_kernel[grid]( kv_in, kv_in.stride(0), @@ -392,11 +517,11 @@ def fused_compress_attn( kv_cache if has_bt else cos_cache, # placeholder when no scatter kv_cache.stride(0) if has_bt else 0, kv_cache.stride(1) if has_bt else 0, + cache_scale_arg, + cache_scale_block_stride_arg, block_tables if has_bt else state_slot_mapping, # placeholder bt_seq_stride, k_per_block, - out, - out.stride(0), head_dim, rope_head_dim, BLOCK_D=BLOCK_D, @@ -406,10 +531,12 @@ def fused_compress_attn( STATE_SIZE=state_size, K=K, HAS_BLOCK_TABLE=int(has_bt), + QUANT=int(quant), + FP8_MAX=fp8_max_arg, + USE_UE8M0=int(use_ue8m0), + PRESHUFFLE=int(preshuffle), ) - return out # [num_compress, head_dim] - def fused_compress_attn_reference( *, @@ -440,8 +567,8 @@ def fused_compress_attn_reference( if plan.num_compress == 0: return None device = kv_in.device - K = (2 if overlap else 1) * ratio - state_size = K + K = (2 if overlap else 1) * ratio # pool window + state_size = kv_state.shape[1] # ring buffer modulo (≥ K) plan_cpu = plan.compress_plan_gpu.detach().cpu() slot_map_cpu = state_slot_mapping.detach().cpu() if block_tables is not None: diff --git a/atom/model_ops/v4_kernels/paged_decode_indices.py b/atom/model_ops/v4_kernels/paged_decode_indices.py new file mode 100644 index 0000000000..1403acbe86 --- /dev/null +++ b/atom/model_ops/v4_kernels/paged_decode_indices.py @@ -0,0 +1,235 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""V4 paged-decode index scatter — single Triton kernel writes SWA window- +prefix paged offsets into the three ragged-packed destination buffers +(`kv_indices_swa` / `kv_indices_csa` / `kv_indices_hca`). + +Replaces the prior chain (numpy `_build_window_topk_np` + `index_copy_`): + + window_topk_np = _build_window_topk_np(positions, win, cs) # [T, win] + swa_paged_2d = torch.where(window_topk >= 0, slot * cs + topk, -1) + swa_paged_flat = swa_paged_2d.reshape(-1) + swa_indices_gpu[:T*win].copy_(swa_paged_flat) + csa_indices_gpu.index_copy_(0, csa_win_pos, swa_paged_flat) + hca_indices_gpu.index_copy_(0, hca_win_pos, swa_paged_flat) + +Two simplifications vs the prior implementation (see plan +`sequential-noodling-turing.md` for details): + +1. The ring-index formula `ring = (pos - win + 1 + w) % cs` is computed + inline inside the kernel from `positions[t]`. The `[T, win]` + `window_topk` intermediate buffer (mnbt·win·4 = 4 MB at typical config) + is gone; no separate CPU build + H2D copy. +2. The destination layout is now ragged-packed (same as prefill): each + token's SWA prefix segment has length `n = min(positions[t]+1, win)` + (NOT a fixed `win` padded with `-1` sentinels). The caller's + `swa_indptr` / `csa_indptr` / `hca_indptr` reflect this ragged sizing. + +Bytewise correctness: for tokens with `position >= win-1` (all `n == win`), +the output is identical to the prior implementation. For shorter +positions, the prior layout wrote `(win - n)` leading `-1` entries that +the sparse-attention kernel masked out; the new layout omits those slots +entirely, saving sparse-attn loop iterations. + +Caller contract: +- Grid = T (one program per token). +- `batch_id_per_token[:T]` may carry `-1` sentinels in the CG-padded tail — + kernel checks and bails (matches `_attach_v4_per_fwd_meta` convention). +- `swa_indptr` / `csa_indptr` / `hca_indptr` must reflect the ragged-packed + sizing: per-token slot count = `min(positions[t]+1, win) + n_compress[t]` + where `n_compress[t]` is 0 for SWA, `min(n_committed_csa, index_topk)` + for CSA, `n_committed_hca` for HCA. +- `swa_indices` / `csa_indices` / `hca_indices` capacity ≥ corresponding + indptr[T]; this kernel only writes the SWA-prefix segment + `[indptr[t], indptr[t] + n)` per token. The compress-tail is filled + elsewhere (HCA: numpy fill in caller, CSA: `csa_translate_pack` per layer). +""" + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _v4_paged_decode_indices_kernel( + state_slot_per_seq_ptr, # [bs] int32 + batch_id_per_token_ptr, # [T+pad] int — sentinel -1 in pad tail + positions_ptr, # [T+pad] int — global token position + swa_indptr_ptr, # [T+1] int32 — ragged SWA-prefix cumsum + csa_indptr_ptr, # [T+1] int32 — ragged (SWA + CSA topk) + hca_indptr_ptr, # [T+1] int32 — ragged (SWA + HCA committed) + swa_indices_ptr, # [swa_total] int32, output + csa_indices_ptr, # [csa_total] int32, output (writes SWA-prefix segment only) + hca_indices_ptr, # [hca_total] int32, output (writes SWA-prefix segment only) + cs, # win_with_spec — stride into unified_kv SWA region (paper §3.6.1) + win: tl.constexpr, # window_size — max SWA prefix slots + BLOCK_N: tl.constexpr, # next_pow2(win) +): + """One program per token. Writes `n = min(positions[t]+1, win)` paged + offsets to the SWA prefix segment of each of SWA/CSA/HCA index buffers. + + For token `t`: + bid = batch_id_per_token[t] # bail if -1 (CG pad) + slot = state_slot_per_seq[bid] + pos = positions[t] + n = min(pos + 1, win) + # Old -1 sentinels were at the leading `win - n` cols; reparameterize + # to skip them: i in [0, n) → abs_pos = pos - n + 1 + i ∈ [0, pos]. + for i in range(n): + abs_pos = pos - n + 1 + i + ring = abs_pos % cs + paged = slot * cs + ring + swa_indices[swa_indptr[t] + i] = paged + csa_indices[csa_indptr[t] + i] = paged + hca_indices[hca_indptr[t] + i] = paged + """ + t = tl.program_id(0) + bid = tl.load(batch_id_per_token_ptr + t) + if bid < 0: + return # CG-padded sentinel — leave outputs untouched + + slot = tl.load(state_slot_per_seq_ptr + bid) + pos = tl.load(positions_ptr + t) + # `n` = actual valid SWA prefix count. Cast to match `win` (compile-time + # int) — pos is i32/i64 from positions buffer. + n = tl.minimum(pos + 1, win) + swa_base = tl.load(swa_indptr_ptr + t) + csa_base = tl.load(csa_indptr_ptr + t) + hca_base = tl.load(hca_indptr_ptr + t) + + i = tl.arange(0, BLOCK_N) + mask = i < n + abs_pos = pos - n + 1 + i # ∈ [0, pos] for valid i + ring_idx = abs_pos % cs + paged = slot * cs + ring_idx + + tl.store(swa_indices_ptr + swa_base + i, paged, mask=mask) + tl.store(csa_indices_ptr + csa_base + i, paged, mask=mask) + tl.store(hca_indices_ptr + hca_base + i, paged, mask=mask) + + +def write_v4_paged_decode_indices( + *, + state_slot_per_seq: torch.Tensor, + batch_id_per_token: torch.Tensor, + positions: torch.Tensor, + swa_indptr: torch.Tensor, + csa_indptr: torch.Tensor, + hca_indptr: torch.Tensor, + swa_indices: torch.Tensor, + csa_indices: torch.Tensor, + hca_indices: torch.Tensor, + T: int, + win: int, + cs: int, +) -> None: + """In-place fill SWA / CSA / HCA window-prefix offsets via a single + Triton kernel. Replaces the prior `_build_window_topk_np` (CPU O(T·win)) + + `index_copy_` chain. All inputs are persistent forward_vars buffers — + no allocator churn. + + Args (all GPU tensors except T/win/cs): + state_slot_per_seq: [bs] int32 — per-seq state cache slot. + batch_id_per_token: [>=T] int — token→seq map; -1 sentinel skipped. + positions: [>=T] int — global token position + (forward_vars["positions"]); used to derive + `n = min(pos+1, win)` per token + the ring + index `(pos - n + 1 + i) % cs`. + swa_indptr: [>=T+1] int32 — ragged SWA-prefix cumsum, where + `swa_indptr[t+1] - swa_indptr[t] = + min(positions[t]+1, win)`. + csa_indptr: [>=T+1] int32 — ragged CSA buffer indptr (SWA + prefix + CSA topk per token). + hca_indptr: [>=T+1] int32 — ragged HCA buffer indptr (SWA + prefix + HCA committed per token). + swa_indices: [>=swa_indptr[T]] int32 OUT — fully written by + this kernel (no other source). + csa_indices: [>=csa_indptr[T]] int32 OUT — window-prefix + `[csa_indptr[t], +n)` written here; CSA + topk tail filled per-layer by + `csa_translate_pack`. + hca_indices: [>=hca_indptr[T]] int32 OUT — same semantics; HCA + compress tail filled in the caller via + numpy fill. + T: int — number of real tokens (grid size). + win: int — SWA window size (typically 128 for V4-Pro). + cs: int — `win_with_spec = window_size + max_spec_steps`, + stride into unified_kv SWA region per slot + AND modulo for ring-index wrap. + """ + if T == 0: + return + assert state_slot_per_seq.dim() == 1 + assert batch_id_per_token.dim() == 1 and batch_id_per_token.shape[0] >= T + assert positions.dim() == 1 and positions.shape[0] >= T + assert swa_indptr.dim() == 1 and swa_indptr.shape[0] >= T + 1 + assert csa_indptr.dim() == 1 and csa_indptr.shape[0] >= T + 1 + assert hca_indptr.dim() == 1 and hca_indptr.shape[0] >= T + 1 + assert swa_indices.dim() == 1 + assert csa_indices.dim() == 1 + assert hca_indices.dim() == 1 + + BLOCK_N = triton.next_power_of_2(win) + _v4_paged_decode_indices_kernel[(T,)]( + state_slot_per_seq, + batch_id_per_token, + positions, + swa_indptr, + csa_indptr, + hca_indptr, + swa_indices, + csa_indices, + hca_indices, + cs, + win=win, + BLOCK_N=BLOCK_N, + ) + + +def write_v4_paged_decode_indices_reference( + *, + state_slot_per_seq: torch.Tensor, + batch_id_per_token: torch.Tensor, + positions: torch.Tensor, + swa_indptr: torch.Tensor, + csa_indptr: torch.Tensor, + hca_indptr: torch.Tensor, + swa_indices: torch.Tensor, + csa_indices: torch.Tensor, + hca_indices: torch.Tensor, + T: int, + win: int, + cs: int, +) -> None: + """Pure-PyTorch reference equivalent of `write_v4_paged_decode_indices`. + For unit tests and bisect verification. Mirrors the kernel exactly: + per-token ragged-packed write, no -1 sentinels in output. + """ + if T == 0: + return + bid = batch_id_per_token[:T].long() + pos_t = positions[:T].long() + valid = bid >= 0 + # n = min(pos+1, win) per token; clamp invalid rows to 0 to skip writes. + n_per_tok = torch.minimum(pos_t + 1, torch.full_like(pos_t, win)) + n_per_tok = torch.where(valid, n_per_tok, torch.zeros_like(n_per_tok)) + slot = torch.where( + valid, state_slot_per_seq[bid.clamp(min=0)].long(), torch.zeros_like(bid) + ) + for t in range(T): + n = int(n_per_tok[t].item()) + if n == 0: + continue + p = int(pos_t[t].item()) + s = int(slot[t].item()) + i_arr = torch.arange(n, device=positions.device, dtype=torch.long) + abs_pos = p - n + 1 + i_arr # [n] + ring = abs_pos % cs + paged = (s * cs + ring).to(torch.int32) + swa_base = int(swa_indptr[t].item()) + csa_base = int(csa_indptr[t].item()) + hca_base = int(hca_indptr[t].item()) + swa_indices[swa_base : swa_base + n] = paged + csa_indices[csa_base : csa_base + n] = paged + hca_indices[hca_base : hca_base + n] = paged diff --git a/atom/model_ops/v4_kernels/state_writes.py b/atom/model_ops/v4_kernels/state_writes.py index 03d9f25cca..2ebfbc4577 100644 --- a/atom/model_ops/v4_kernels/state_writes.py +++ b/atom/model_ops/v4_kernels/state_writes.py @@ -9,13 +9,16 @@ Currently implemented: - `swa_write`: writes `swa_kv[state_slot_per_seq[batch_id_per_token[t]], - positions[t] % win, :] = kv[t, :]` for each src row id `t` selected by - `write_indices`. The kernel does ALL gathers (kv row, position, batch id, - state slot) itself — caller passes only stable forward_vars buffers (full - `kv`, full `positions`, full `batch_id_per_token`, per-seq `state_slot`). - Race-free for long prefill via the `write_indices` filter (only the last - `win` tokens per seq selected); CUDAGraph-safe via sentinel skip - (`write_indices[pid] < 0` → bail). + positions[t] % cache_size, :] = kv[t, :]` for each src row id `t` selected + by `write_indices`. The kernel does ALL gathers (kv row, position, batch + id, state slot) itself — caller passes only stable forward_vars buffers + (full `kv`, full `positions`, full `batch_id_per_token`, per-seq + `state_slot`). Race-free for long prefill via the `write_indices` filter + (only the last `cache_size` tokens per seq selected); CUDAGraph-safe via + sentinel skip (`write_indices[pid] < 0` → bail). `cache_size` is + `window_size + max_spec_steps` — for non-MTP this reduces to `window_size` + (no behavioral change); for MTP-k draft tokens get their own ring slots + separate from the verified token's slot. - `update_compressor_states`: unified in-place update of Compressor's per-request `kv_state` + `score_state` ring buffers, covering both prefill (B-side overlap context + tail) and decode (every token at `pos % STATE_SIZE` @@ -31,17 +34,19 @@ - `kv` [T, head_dim] flat — full per-fwd KV (forward_vars). - `write_indices` [W] int — src row ids into `kv` / `positions` / `batch_id_per_token`. Sentinel = -1 → kernel skips. - For long prefill (`seqlen > win`) builder pre-filters - to the last `win` rows per seq to avoid `pos % win` - collisions. For decode/MTP every row is written. + For long prefill (`seqlen > cache_size`) builder + pre-filters to the last `cache_size` rows per seq to + avoid `pos % cache_size` collisions. For decode/MTP + every row is written. - `positions` [T] int — full positions buffer (forward_vars). - `batch_id_per_token` [T] int — Phase-B `v4_batch_id_per_token` mapping; kernel does `state_slot_per_seq[batch_id]` for the per-seq state-cache slot. Single per-token mapping principle (no per-token slot alias). - `state_slot_per_seq` [bs] int — `state_slot_mapping_gpu_i32`. -- `swa_kv` [num_slots, win, head_dim] in-place buffer. -- `win` int sliding-window size (e.g. 128). +- `swa_kv` [num_slots, cache_size, head_dim] in-place buffer. +- `cache_size` int ring-slot count = `window_size + max_spec_steps` + (e.g. 128 + 0 = 128 non-MTP; 128 + 1 = 129 MTP-1). Grid = `write_indices.shape[0]`; each program processes one src row id. """ @@ -58,11 +63,11 @@ def _swa_write_kernel( positions_ptr, # [T] int — full positions batch_id_per_token_ptr, # [T] int — v4_batch_id_per_token state_slot_per_seq_ptr, # [bs] int — state_slot_mapping_gpu_i32 - swa_kv_ptr, # [num_slots, win, head_dim] - swa_kv_slot_stride, # = win * head_dim + swa_kv_ptr, # [num_slots, cache_size, head_dim] + swa_kv_slot_stride, # = cache_size * head_dim swa_kv_pos_stride, # = head_dim head_dim, - win, + cache_size, BLOCK_D: tl.constexpr, ): """One program per write_indices entry. Sentinel skip via write_indices < 0. @@ -78,7 +83,7 @@ def _swa_write_kernel( pos = tl.load(positions_ptr + src_id) bid = tl.load(batch_id_per_token_ptr + src_id) slot = tl.load(state_slot_per_seq_ptr + bid) - ring_idx = pos % win + ring_idx = pos % cache_size d_offsets = tl.arange(0, BLOCK_D) d_mask = d_offsets < head_dim @@ -103,10 +108,11 @@ def swa_write( batch_id_per_token: torch.Tensor, state_slot_per_seq: torch.Tensor, swa_kv: torch.Tensor, - win: int, + cache_size: int, ) -> None: - """In-place write `swa_kv[state_slot_per_seq[bid], pos % win, :] = kv[r, :]` - for each `r = write_indices[pid]` (skip pid where `write_indices[pid] < 0`). + """In-place write + `swa_kv[state_slot_per_seq[bid], pos % cache_size, :] = kv[r, :]` for + each `r = write_indices[pid]` (skip pid where `write_indices[pid] < 0`). Per-token quantities (`pos`, `bid`) are gathered inside the kernel via `positions[r]` / `batch_id_per_token[r]`; per-seq `state_slot_per_seq` @@ -122,21 +128,25 @@ def swa_write( positions: [T] int — full forward_vars["positions"]. batch_id_per_token: [T] int — v4_batch_id_per_token mapping. state_slot_per_seq: [bs] int — per-seq state cache slot. - swa_kv: [num_slots, win, head_dim] in-place ring buffer. - win: sliding-window size. + swa_kv: [num_slots, cache_size, head_dim] in-place ring buffer. + cache_size: ring-slot count = `window_size + max_spec_steps`. + For non-MTP this equals `window_size` and the kernel is bytewise + identical to the pre-MTP behavior. """ assert kv.dim() == 2, f"kv must be [T, D], got {kv.shape}" assert write_indices.dim() == 1 assert positions.dim() == 1 assert batch_id_per_token.dim() == 1 assert state_slot_per_seq.dim() == 1 - assert swa_kv.dim() == 3, f"swa_kv must be [S, W, D], got {swa_kv.shape}" + assert swa_kv.dim() == 3, f"swa_kv must be [S, C, D], got {swa_kv.shape}" T, head_dim = kv.shape assert positions.shape[0] >= T, f"positions {positions.shape[0]} < kv T={T}" assert ( batch_id_per_token.shape[0] >= T ), f"batch_id_per_token {batch_id_per_token.shape[0]} < kv T={T}" - assert swa_kv.shape[1] == win + assert ( + swa_kv.shape[1] == cache_size + ), f"swa_kv ring dim {swa_kv.shape[1]} != cache_size {cache_size}" assert swa_kv.shape[2] == head_dim assert kv.is_contiguous() and swa_kv.is_contiguous() @@ -148,6 +158,7 @@ def swa_write( # block per token covers it. Round up to the next power of two for tl. BLOCK_D = triton.next_power_of_2(head_dim) grid = (W,) + _swa_write_kernel[grid]( kv, write_indices, @@ -158,7 +169,7 @@ def swa_write( swa_kv.stride(0), swa_kv.stride(1), head_dim, - win, + cache_size, BLOCK_D=BLOCK_D, ) @@ -170,7 +181,7 @@ def swa_write_reference( batch_id_per_token: torch.Tensor, state_slot_per_seq: torch.Tensor, swa_kv: torch.Tensor, - win: int, + cache_size: int, ) -> None: """Pure-PyTorch reference equivalent of `swa_write`. For tests / dump-bisect. @@ -185,7 +196,7 @@ def swa_write_reference( src_pos = positions[src_ids] bids = batch_id_per_token[src_ids].long() slots = state_slot_per_seq[bids].long() - ring_idx = src_pos % win + ring_idx = src_pos % cache_size swa_kv[slots, ring_idx] = src_kv @@ -319,10 +330,11 @@ def update_compressor_states( assert kv.dim() == 2 and score.dim() == 2 assert kv.shape == score.shape, f"{kv.shape} vs {score.shape}" assert ape.dim() == 2 and ape.shape[0] == ratio - state_size = (2 if overlap else 1) * ratio + K_pool = (2 if overlap else 1) * ratio # pool window (lower bound) + state_size = kv_state.shape[1] # ring buffer modulo (≥ K_pool) assert ( - kv_state.shape[1] == state_size - ), f"kv_state.shape[1]={kv_state.shape[1]}, expected {state_size}" + state_size >= K_pool + ), f"kv_state.shape[1]={state_size}, must be ≥ K_pool={K_pool}" dim = kv.shape[1] assert write_plan.dim() == 2 and write_plan.shape[1] == 4 assert write_plan.dtype == torch.int32 @@ -379,7 +391,7 @@ def update_compressor_states_reference( `write_plan[i] = (ragged_id, batch_id, position, _)` — each row is one token to write. No mask (host filtered). """ - state_size = (2 if overlap else 1) * ratio + state_size = kv_state.shape[1] # ring buffer modulo (≥ (1+overlap)*ratio) plan_cpu = write_plan.detach().cpu() slot_map_cpu = state_slot_mapping.detach().cpu() for i in range(plan_cpu.shape[0]): diff --git a/atom/models/deepseek_mtp.py b/atom/models/deepseek_mtp.py index 924608403d..952c4d1039 100644 --- a/atom/models/deepseek_mtp.py +++ b/atom/models/deepseek_mtp.py @@ -8,6 +8,7 @@ from atom.config import Config, QuantizationConfig from atom.model_ops.embed_head import ParallelLMHead, VocabParallelEmbedding from atom.model_ops.layernorm import RMSNorm +from atom.model_ops.linear import ReplicatedLinear from atom.model_ops.moe import FusedMoE from atom.model_ops.topK import is_rocm_aiter_fusion_shared_expert_enabled from atom.models.utils import IntermediateTensors @@ -16,7 +17,7 @@ from transformers import DeepseekV2Config, DeepseekV3Config, PretrainedConfig from .deepseek_v2 import DeepseekV2DecoderLayer -from .utils import maybe_prefix +from .utils import ckpt_has_tensor_suffix, maybe_prefix class SharedHead(nn.Module): @@ -40,7 +41,13 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: class DeepSeekMultiTokenPredictorLayer(nn.Module): - def __init__(self, atom_config: Config, prefix: str, layer_idx: int) -> None: + def __init__( + self, + atom_config: Config, + prefix: str, + layer_idx: int, + topk_indices_buffer: Optional[torch.Tensor] = None, + ) -> None: super().__init__() config = atom_config.hf_config @@ -48,7 +55,27 @@ def __init__(self, atom_config: Config, prefix: str, layer_idx: int) -> None: self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.eh_proj = nn.Linear(config.hidden_size * 2, config.hidden_size, bias=False) + self.eh_proj = ReplicatedLinear( + config.hidden_size * 2, + config.hidden_size, + bias=False, + quant_config=atom_config.quant_config, + prefix=maybe_prefix(prefix, "eh_proj"), + ) + + if hasattr(config, "index_topk"): + max_num_batched_tokens = getattr( + atom_config, "max_num_batched_tokens", atom_config.max_num_seqs + ) + buffer_device = getattr(atom_config, "device", "cuda") + topk_indices_buffer = torch.empty( + max_num_batched_tokens, + config.index_topk, + dtype=torch.int32, + device=buffer_device, + ) + else: + topk_indices_buffer = None self.shared_head = SharedHead( config=config, prefix=prefix, quant_config=atom_config.quant_config @@ -63,6 +90,7 @@ def __init__(self, atom_config: Config, prefix: str, layer_idx: int) -> None: quant_config=quant_config, layer_num=layer_idx, is_mtp_block=True, + topk_indices_buffer=topk_indices_buffer, ) def forward( @@ -74,9 +102,6 @@ def forward( spec_step_index: int = 0, ) -> torch.Tensor: assert inputs_embeds is not None - # masked_inputs_embeds = torch.where( - # positions.unsqueeze(-1) == 0, 0, inputs_embeds - # ) masked_inputs_embeds = inputs_embeds inputs_embeds = self.enorm(masked_inputs_embeds) previous_hidden_states = self.hnorm(previous_hidden_states) @@ -95,7 +120,13 @@ def forward( class DeepSeekMultiTokenPredictor(nn.Module): - def __init__(self, *, atom_config: Config, prefix: str = ""): + def __init__( + self, + *, + atom_config: Config, + prefix: str = "", + topk_indices_buffer: Optional[torch.Tensor] = None, + ): super().__init__() config = atom_config.hf_config self.mtp_start_layer_idx = config.num_hidden_layers @@ -104,7 +135,10 @@ def __init__(self, *, atom_config: Config, prefix: str = ""): self.layers = torch.nn.ModuleDict( { str(idx): DeepSeekMultiTokenPredictorLayer( - atom_config, f"{prefix}.layers.{idx}", layer_idx=idx + atom_config, + f"{prefix}.layers.{idx}", + layer_idx=idx, + topk_indices_buffer=topk_indices_buffer, ) for idx in range( self.mtp_start_layer_idx, @@ -154,6 +188,20 @@ def __init__(self, atom_config: Config, prefix: str = ""): super().__init__() self.config = atom_config.hf_config + # Several MTP checkpoints (DeepSeek R1/V3/V3.2 FP8 + the Quark mixed + # MXFP4/FP8 variants) store eh_proj as BF16 with no weight_scale even + # though their HF quantization_config does not list eh_proj in the + # exclude set. Without this guard ReplicatedLinear is built with the + # global FP8/MXFP4 spec, the BF16 weight is cast into the FP8 slot + # against an uninitialized weight_scale, and MTP accept rate collapses. + # GLM-FP8 ckpts already list eh_proj explicitly (this becomes a no-op); + # GLM-5.1-MXFP4 truly quantizes eh_proj and ships weight_scale on disk + # so the check below leaves the global spec in effect. + if atom_config.quant_config is not None and not ckpt_has_tensor_suffix( + atom_config.model, "eh_proj.weight_scale" + ): + atom_config.quant_config.apply_default_exclude_layers(["*.eh_proj"]) + if hasattr(self.config, "q_lora_rank") and self.config.q_lora_rank is not None: self.packed_modules_mapping = { "q_a_proj": ("fused_qkv_a_proj", 0), @@ -167,8 +215,20 @@ def __init__(self, atom_config: Config, prefix: str = ""): "up_proj": ("gate_up_proj", 1), } + if hasattr(self.config, "index_topk"): + topk_indices_buffer = torch.empty( + atom_config.max_num_batched_tokens, + self.config.index_topk, + dtype=torch.int32, + device="cuda", + ) + else: + topk_indices_buffer = None + self.model = DeepSeekMultiTokenPredictor( - atom_config=atom_config, prefix=maybe_prefix(prefix, "model") + atom_config=atom_config, + prefix=maybe_prefix(prefix, "model"), + topk_indices_buffer=topk_indices_buffer, ) def remap_mtp_weight_name(self, name: str) -> str | None: diff --git a/atom/models/deepseek_v2.py b/atom/models/deepseek_v2.py index c9ad8e53f4..92145c1cf8 100644 --- a/atom/models/deepseek_v2.py +++ b/atom/models/deepseek_v2.py @@ -35,6 +35,7 @@ gemm_a8w8_blockscale_bpreshuffle, get_hip_quant, indexer_k_quant_and_cache, + indexer_qk_rope_quant_and_cache, top_k_per_row_decode, top_k_per_row_prefill, ) @@ -116,6 +117,48 @@ ENABLE_DS_QKNORM_FUSION = envs.ATOM_ENABLE_DS_QKNORM_FUSION ENABLE_ALLREDUCE_RMSNORM_FUSION = envs.ATOM_ENABLE_ALLREDUCE_RMSNORM_FUSION ENABLE_DS_INPUT_RMSNORM_QUANT_FUSION = envs.ATOM_ENABLE_DS_INPUT_RMSNORM_QUANT_FUSION +ENABLE_DS_INDEXER_QK_ROPE_CACHE_FUSION = ( + envs.ATOM_ENABLE_DS_INDEXER_QK_ROPE_CACHE_FUSION +) +_FP8_DTYPES = tuple( + dtype + for dtype in ( + getattr(torch, "float8_e4m3fn", None), + getattr(torch, "float8_e4m3fnuz", None), + ) + if dtype is not None +) + + +def _supports_fused_indexer_kernel_config(config: PretrainedConfig) -> bool: + if not hasattr(config, "index_topk"): + return False + if getattr(config, "model_type", None) == "glm_moe_dsa": + return False + return ( + getattr(config, "index_head_dim", None) == 128 + and getattr(config, "qk_rope_head_dim", None) == 64 + ) + + +def _can_fuse_indexer_wk_weights_proj( + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig], + indexer_prefixes: list[str], +) -> bool: + if not _supports_fused_indexer_kernel_config(config): + return False + if quant_config is None: + return True + + for indexer_prefix in indexer_prefixes: + wk_quant_config = quant_config.get_layer_quant_config(f"{indexer_prefix}.wk") + if ( + wk_quant_config.quant_type != QuantType.No + and wk_quant_config.quant_dtype != dtypes.fp8 + ): + return False + return True def _fuse_rmsnorm_fp4_quant_fake( @@ -160,7 +203,7 @@ def _fuse_rmsnorm_fp4_quant_fake( return out1_quantized, out1_bs, out1_unquantized, out2, out_res1 -def _fused_rms_fp8_group_quant_fake( +def _fused_rms_fp8_quant_fake( x1: torch.Tensor, x1_weight: torch.Tensor, x1_epsilon: float, @@ -170,6 +213,7 @@ def _fused_rms_fp8_group_quant_fake( res1: Optional[torch.Tensor] = None, dtype_quant: torch.dtype = dtypes.fp8, group_size: int = 128, + quant_type: Optional[int] = None, output_unquantized_inp1: bool = False, transpose_scale: bool = False, ) -> Tuple[ @@ -180,9 +224,18 @@ def _fused_rms_fp8_group_quant_fake( torch.Tensor, ]: m, n1 = x1.shape - out1_quantized = torch.empty((m, n1), dtype=dtype_quant, device=x1.device) - num_bs_cols = (n1 + group_size - 1) // group_size - out1_bs = torch.empty((m, num_bs_cols), dtype=torch.float32, device=x1.device) + no_quant = quant_type is None or quant_type == QuantType.No.value + if not no_quant: + out1_quantized = torch.empty((m, n1), dtype=dtype_quant, device=x1.device) + else: + out1_quantized = torch.empty_like(x1) + if no_quant: + out1_bs = None + elif quant_type == QuantType.per_Token.value: + out1_bs = torch.empty((m, 1), dtype=torch.float32, device=x1.device) + else: + num_bs_cols = (n1 + group_size - 1) // group_size + out1_bs = torch.empty((m, num_bs_cols), dtype=torch.float32, device=x1.device) out1_unquantized = torch.empty_like(x1) if output_unquantized_inp1 else None out2 = None if x2 is not None: @@ -236,8 +289,8 @@ def _fuse_rmsnorm_fp4_quant( return out1_quantized, out1_bs, out1_unquantized, out2, out_res1 -@torch_compile_guard(gen_fake=_fused_rms_fp8_group_quant_fake) -def _fused_rms_fp8_group_quant( +@torch_compile_guard(gen_fake=_fused_rms_fp8_quant_fake) +def _fused_rms_fp8_quant( x1: torch.Tensor, x1_weight: torch.Tensor, x1_epsilon: float, @@ -247,6 +300,7 @@ def _fused_rms_fp8_group_quant( res1: Optional[torch.Tensor] = None, dtype_quant: torch.dtype = dtypes.fp8, group_size: int = 128, + quant_type: Optional[int] = None, output_unquantized_inp1: bool = False, transpose_scale: bool = False, ) -> Tuple[ @@ -257,7 +311,7 @@ def _fused_rms_fp8_group_quant( torch.Tensor, ]: out1_quantized, out1_bs, out1_unquantized, out2, out_res1 = ( - _fused_rms_fp8_group_quant_fake( + _fused_rms_fp8_quant_fake( x1, x1_weight, x1_epsilon, @@ -267,14 +321,18 @@ def _fused_rms_fp8_group_quant( res1, dtype_quant, group_size, + quant_type, output_unquantized_inp1, transpose_scale, ) ) - from aiter.ops.fused_qk_rmsnorm_group_quant import fused_qk_rmsnorm_group_quant + if quant_type is None: + quant_type = QuantType.No + else: + quant_type = QuantType(quant_type) - fused_qk_rmsnorm_group_quant( + fused_qk_rmsnorm( q_out_quantized=out1_quantized, q_out_scale=out1_bs, q=x1, @@ -287,6 +345,7 @@ def _fused_rms_fp8_group_quant( k_weight=x2_weight, k_epsilon=x2_epsilon, q_residual=res1, + quant_type=quant_type, group_size=group_size, transpose_scale=transpose_scale, ) @@ -306,6 +365,7 @@ def _fuse_rmsnorm_quant( shuffle: bool = True, scale_shuffle_padding: bool = False, group_size: int = 128, + quant_type: Optional[int] = None, output_unquantized_inp1: bool = False, transpose_scale: bool = False, ): @@ -324,9 +384,9 @@ def _fuse_rmsnorm_quant( output_unquantized_inp1, ) ) - elif dtype_quant == dtypes.fp8: + elif dtype_quant == dtypes.fp8 or dtype_quant == torch.bfloat16: out1_quantized, out1_bs, out1_unquantized, out2, out_res1 = ( - _fused_rms_fp8_group_quant( + _fused_rms_fp8_quant( x1, x1_weight, x1_epsilon, @@ -334,10 +394,11 @@ def _fuse_rmsnorm_quant( x2_weight, x2_epsilon, res1, - dtype_quant, - group_size, - output_unquantized_inp1, - transpose_scale, + dtype_quant=dtype_quant, + group_size=group_size, + quant_type=quant_type, + output_unquantized_inp1=output_unquantized_inp1, + transpose_scale=transpose_scale, ) ) else: @@ -703,38 +764,7 @@ def _fuse_qkv_a_proj_reduce_rmsnorm_quant( return q_c, q_c_scale, kv_c_normed, k_pe -def _fused_qk_rmsnorm_fake( - q_c: torch.Tensor, - q_a_layernorm_weight: torch.Tensor, - q_a_layernorm_variance_epsilon: float, - kv_c: torch.Tensor, - kv_a_layernorm_weight: torch.Tensor, - kv_a_layernorm_variance_epsilon: float, -) -> Tuple[torch.Tensor, torch.Tensor]: - return torch.empty_like(q_c), torch.empty_like(kv_c) - - -@torch_compile_guard(gen_fake=_fused_qk_rmsnorm_fake) -def _fused_qk_rmsnorm( - q_c: torch.Tensor, - q_a_layernorm_weight: torch.Tensor, - q_a_layernorm_variance_epsilon: float, - kv_c: torch.Tensor, - kv_a_layernorm_weight: torch.Tensor, - kv_a_layernorm_variance_epsilon: float, -) -> Tuple[torch.Tensor, torch.Tensor]: - return fused_qk_rmsnorm( - q_c, - q_a_layernorm_weight, - q_a_layernorm_variance_epsilon, - kv_c, - kv_a_layernorm_weight, - kv_a_layernorm_variance_epsilon, - ) - - class DeepseekV2MLP(nn.Module): - def __init__( self, hidden_size: int, @@ -762,8 +792,7 @@ def __init__( ) if hidden_act != "silu": raise ValueError( - f"Unsupported activation: {hidden_act}. " - "Only silu is supported for now." + f"Unsupported activation: {hidden_act}. Only silu is supported for now." ) self.act_fn = SiluAndMul() @@ -775,7 +804,6 @@ def forward(self, x): class DeepseekV2MoE(nn.Module): - def __init__( self, config: PretrainedConfig, @@ -944,7 +972,6 @@ def yarn_get_mscale(scale: float = 1, mscale: float = 1) -> float: @DeepseekV32IndexerCacheDecoratorForPluginMode class DeepseekV32IndexerCache(nn.Module): - def __init__( self, head_dim: int, dtype: torch.dtype, prefix: str, cache_config: str ): @@ -960,7 +987,7 @@ def sparse_attn_indexer( hidden_states: torch.Tensor, k_cache_prefix: str, kv_cache: torch.Tensor, - q_fp8: torch.Tensor, + q_input: torch.Tensor, k: torch.Tensor, weights: torch.Tensor, quant_block_size: int, @@ -970,6 +997,15 @@ def sparse_attn_indexer( max_model_len: int, total_seq_lens: int, topk_indices_buffer: torch.Tensor, + k_norm_weight: torch.Tensor, + k_norm_bias: torch.Tensor, + k_norm_eps: float, + positions: torch.Tensor, + cos_cache: torch.Tensor, + sin_cache: torch.Tensor, + weights_scale: float, + is_neox_style: bool, + use_qk_rope_cache_fusion: bool, ) -> torch.Tensor: # careful! this will be None in dummy run forward_context = get_forward_context() @@ -979,15 +1015,50 @@ def sparse_attn_indexer( # Skip for dummy runs to avoid corrupting KV cache if forward_context.context.is_dummy_run: # dummy runner - return weights - num_decode_tokens = context.batch_size if not context.is_prefill else 0 - indexer_k_quant_and_cache( - k, - kv_cache, - slot_mapping, - quant_block_size, - scale_fmt, + return torch.zeros_like(weights, dtype=torch.float32) + # For MTP verify decode, max_seqlen_q > 1 so total decode tokens = batch_size * max_seqlen_q + num_decode_tokens = ( + context.batch_size * attn_metadata.max_seqlen_q if not context.is_prefill else 0 ) + runner_block_size = get_current_atom_config().kv_cache_block_size + kv_cache = kv_cache.view(-1, runner_block_size, kv_cache.shape[-1]) + if use_qk_rope_cache_fusion: + q_bf16 = q_input + q_fp8 = torch.empty_like(q_bf16, dtype=dtypes.fp8) + weights_out = torch.empty( + weights.shape, device=weights.device, dtype=torch.float32 + ) + indexer_qk_rope_quant_and_cache( + q_bf16, + q_fp8, + weights, + weights_out, + k, + kv_cache, + slot_mapping, + k_norm_weight, + k_norm_bias, + positions, + cos_cache, + sin_cache, + k_norm_eps, + quant_block_size, + scale_fmt, + weights_scale, + preshuffle=True, + is_neox=is_neox_style, + ) + weights = weights_out + else: + q_fp8 = q_input + indexer_k_quant_and_cache( + k, + kv_cache, + slot_mapping, + quant_block_size, + scale_fmt, + preshuffle=True, + ) if context.is_prefill: if attn_metadata.max_seqlen_k <= topk_indices_buffer.shape[1]: return weights @@ -1018,6 +1089,7 @@ def sparse_attn_indexer( if prefill_metadata.has_cached else prefill_metadata.cu_seqlens_q ), + preshuffle=True, ) cu_seqlen_ks = prefill_metadata.cu_seqlen_ks cu_seqlen_ke = prefill_metadata.cu_seqlen_ke @@ -1069,6 +1141,8 @@ def sparse_attn_indexer( decode_metadata.context_lens, attn_metadata.block_tables, max_model_len, + KVBlockSize=runner_block_size, + Preshuffle=True, ) num_rows = logits.shape[0] assert topk_tokens == 2048, "top_k_per_row assumes size 2048" @@ -1089,7 +1163,7 @@ def sparse_attn_indexer_fake( hidden_states: torch.Tensor, k_cache_prefix: str, kv_cache: torch.Tensor, - q_fp8: torch.Tensor, + q_input: torch.Tensor, k: torch.Tensor, weights: torch.Tensor, quant_block_size: int, @@ -1099,6 +1173,15 @@ def sparse_attn_indexer_fake( max_model_len: int, total_seq_lens: int, topk_indices_buffer: torch.Tensor, + k_norm_weight: torch.Tensor, + k_norm_bias: torch.Tensor, + k_norm_eps: float, + positions: torch.Tensor, + cos_cache: torch.Tensor, + sin_cache: torch.Tensor, + weights_scale: float, + is_neox_style: bool, + use_qk_rope_cache_fusion: bool, ) -> torch.Tensor: # profile run # NOTE(Chen): create the max possible flattened_kv. So that @@ -1108,7 +1191,7 @@ def sparse_attn_indexer_fake( ) _k_fp8 = _flattened_kv[..., :head_dim].view(torch.float8_e4m3fn).contiguous() _k_scale = _flattened_kv[..., head_dim:].view(torch.float32).contiguous() - return weights + return torch.empty(weights.shape, device=weights.device, dtype=torch.float32) direct_register_custom_op( @@ -1119,9 +1202,120 @@ def sparse_attn_indexer_fake( ) +def _dequant_fp8_block_to_bf16( + weight_fp8: torch.Tensor, + scale: torch.Tensor, + block_size: int = 128, +) -> torch.Tensor: + """Dequantize block-scaled FP8 weights to BF16 for BF16-only fused GEMMs.""" + out_dim, in_dim = weight_fp8.shape + if out_dim % block_size != 0 or in_dim % block_size != 0: + raise ValueError( + "FP8 block dequant expects dimensions divisible by " + f"{block_size}, got {tuple(weight_fp8.shape)}" + ) + weight = ( + weight_fp8.unflatten(0, (-1, block_size)) + .unflatten(-1, (-1, block_size)) + .float() + ) + scale = scale.float() + return (weight * scale[:, None, :, None]).flatten(2, 3).flatten(0, 1).bfloat16() + + +class IndexerWkWeightsProjLinear(MergedReplicatedLinear): + """Fused Indexer wk + weights projection with FP8 wk load support.""" + + def __init__( + self, + hidden_size: int, + head_dim: int, + n_head: int, + prefix: str = "", + ): + self._wk_pending_weight: Optional[torch.Tensor] = None + self._wk_pending_scale: Optional[torch.Tensor] = None + self._wk_loaded = False + super().__init__( + hidden_size, + [head_dim, n_head], + bias=False, + quant_config=None, + prefix=prefix, + ) + # Checkpoints may store indexer.wk as FP8 plus weight_scale_inv. The + # fused GEMM runs in BF16, so this parameter only receives the scale + # during loading and is not consumed in forward. + self.weight_scale = atom_parameter( + torch.empty( + ((head_dim + 127) // 128, (hidden_size + 127) // 128), + dtype=torch.float32, + ) + ) + self.weight_scale.weight_loader_process = self.weight_loader_process + self.weight_scale.weight_loader = self.weight_loader + + def _maybe_load_pending_wk(self) -> None: + if self._wk_pending_weight is None or self._wk_pending_scale is None: + return + wk_weight_fp8 = self._wk_pending_weight + if wk_weight_fp8.device != self._wk_pending_scale.device: + wk_weight_fp8 = wk_weight_fp8.to(self._wk_pending_scale.device) + wk_weight = _dequant_fp8_block_to_bf16( + wk_weight_fp8, + self._wk_pending_scale, + ) + super().weight_loader(self.weight, wk_weight, 0) + self._wk_pending_weight = None + self._wk_pending_scale = None + self._wk_loaded = True + + def weight_loader( + self, + param: nn.Parameter, + loaded_weight: torch.Tensor, + loaded_shard_id: Optional[int] = None, + ): + if param is self.weight_scale: + if loaded_shard_id == 0: + param.weight_loader_process(param.data, loaded_weight) + self._wk_pending_scale = param.data.detach().clone() + self._maybe_load_pending_wk() + return + + if ( + param is self.weight + and loaded_shard_id == 0 + and loaded_weight.dtype in _FP8_DTYPES + ): + self._wk_pending_weight = loaded_weight.detach().clone() + self._maybe_load_pending_wk() + return + + if param is self.weight and loaded_shard_id == 0: + self._wk_pending_weight = None + self._wk_pending_scale = None + self._wk_loaded = True + + super().weight_loader(param, loaded_weight, loaded_shard_id) + + def process_weights_after_loading(self): + if self._wk_pending_weight is not None or ( + self._wk_pending_scale is not None and not self._wk_loaded + ): + raise RuntimeError( + "Incomplete FP8 indexer.wk load: both weight and weight_scale " + "are required before building wk_weights_proj." + ) + if not self._wk_loaded: + raise RuntimeError( + "Missing indexer.wk load before building wk_weights_proj." + ) + super().process_weights_after_loading() + + @IndexerDecoratorForPluginMode class Indexer(nn.Module): - def __init__( self, atom_config: Config, @@ -1131,6 +1325,7 @@ def __init__( quant_config: Optional[QuantizationConfig], cache_config: str, topk_indices_buffer: Optional[torch.Tensor], + use_wk_weights_proj_fusion: bool = True, prefix: str = "", ): super().__init__() @@ -1150,25 +1345,41 @@ def __init__( quant_config=quant_config, prefix=f"{prefix}.wq_b", ) - self.wk = ReplicatedLinear( - hidden_size, - self.head_dim, - bias=False, - quant_config=quant_config, - prefix=f"{prefix}.wk", - ) + self.use_wk_weights_proj_fusion = use_wk_weights_proj_fusion + if self.use_wk_weights_proj_fusion: + self.wk_weights_proj = IndexerWkWeightsProjLinear( + hidden_size, + self.head_dim, + self.n_head, + prefix=f"{prefix}.wk_weights_proj", + ) + else: + self.wk = ReplicatedLinear( + hidden_size, + self.head_dim, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.wk", + ) + self.weights_proj = ReplicatedLinear( + hidden_size, + self.n_head, + quant_config=None, + prefix=f"{prefix}.weights_proj", + ) self.k_norm = LayerNorm(self.head_dim, eps=1e-6) - self.weights_proj = ReplicatedLinear( - hidden_size, - self.n_head, - quant_config=quant_config, - prefix=f"{prefix}.weights_proj", - ) self.softmax_scale = self.head_dim**-0.5 + self._weights_scale = self.softmax_scale * self.n_head**-0.5 self.scale_fmt = "ue8m0" self.quant_func = get_hip_quant(QuantType.per_1x128) self.quant_block_size = 128 # TODO: get from config + self.use_qk_rope_cache_fusion = ( + ENABLE_DS_INDEXER_QK_ROPE_CACHE_FUSION + and _supports_fused_indexer_kernel_config(config) + and self.head_dim == self.quant_block_size + and self.rope_dim == self.head_dim // 2 + ) self.topk_indices_buffer = topk_indices_buffer # TODO (zyongye) change dim to fp8 later to (self.head_dim + 4) @@ -1195,36 +1406,45 @@ def forward( ) -> torch.Tensor: q = self.wq_b(qr, qr_scale) q = q.view(-1, self.n_head, self.head_dim) - q_pe, q_nope = torch.split( - q, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1 - ) - - k = self.wk(hidden_states) - k = self.k_norm(k) - k_pe, k_nope = torch.split( - k, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1 - ) - - q_pe, k_pe = rotary_emb(positions, q_pe, k_pe) - - # we only quant q here since k quant is fused with cache insertion - q = q.view(-1, self.head_dim) - - q_fp8, q_scale = self.quant_func(q, quant_dtype=dtypes.fp8) - q_fp8 = q_fp8.view(-1, self.n_head, self.head_dim) - q_scale = q_scale.view(-1, self.n_head, 1) - weights = self.weights_proj(hidden_states) - weights = ( - weights.unsqueeze(-1) * q_scale * self.softmax_scale * self.n_head**-0.5 - ) - weights = weights.squeeze(-1) + if self.use_wk_weights_proj_fusion: + k, weights = torch.split( + self.wk_weights_proj(hidden_states), + [self.head_dim, self.n_head], + dim=-1, + ) + else: + k = self.wk(hidden_states) + weights = self.weights_proj(hidden_states) + k = k.contiguous() + weights = weights.contiguous() + + if not self.use_qk_rope_cache_fusion: + q_pe, _ = torch.split( + q, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1 + ) + k = self.k_norm(k) + k_pe, _ = torch.split( + k, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1 + ) + q_pe, k_pe = rotary_emb(positions, q_pe, k_pe) + + q = q.view(-1, self.head_dim) + q_fp8, q_scale = self.quant_func(q, quant_dtype=dtypes.fp8) + q_fp8 = q_fp8.view(-1, self.n_head, self.head_dim) + q_scale = q_scale.view(-1, self.n_head, 1) + weights = (weights.unsqueeze(-1) * q_scale * self._weights_scale).squeeze( + -1 + ) + q_input = q_fp8 + else: + q_input = q return self.sparse_attn_indexer_impl( hidden_states, self.k_cache.prefix, self.k_cache.kv_cache[0], - q_fp8, + q_input, k, weights, self.quant_block_size, @@ -1234,6 +1454,15 @@ def forward( self.max_model_len, self.max_total_seq_len, self.topk_indices_buffer, + self.k_norm.weight, + self.k_norm.bias, + self.k_norm.eps, + positions, + rotary_emb.cos_cache.squeeze(-2).squeeze(-2), + rotary_emb.sin_cache.squeeze(-2).squeeze(-2), + self._weights_scale, + rotary_emb.is_neox_style, + self.use_qk_rope_cache_fusion, ) @@ -1261,6 +1490,7 @@ def __init__( prefix: str = "", layer_num: int = 0, topk_indices_buffer: Optional[torch.Tensor] = None, + use_indexer_wk_weights_proj_fusion: Optional[bool] = None, ) -> None: super().__init__() self.hidden_size = hidden_size @@ -1271,6 +1501,7 @@ def __init__( self.q_lora_rank = q_lora_rank self.kv_lora_rank = kv_lora_rank + model_quant_config = quant_config self.num_heads = num_heads tp_size = get_tensor_model_parallel_world_size() @@ -1289,6 +1520,14 @@ def __init__( layer_quant_dtype = quant_config.get_layer_quant_config( f"{prefix}.{q_a_proj_name}" ).quant_dtype + layer_quant_type = quant_config.get_layer_quant_config( + f"{prefix}.{q_a_proj_name}" + ).quant_type + # Keep a plain int on the module so Dynamo does not guard on a + # QuantType enum attribute. + layer_quant_type_value = ( + None if layer_quant_type is None else layer_quant_type.value + ) if layer_quant_dtype == dtypes.fp4x2: if not use_triton_gemm(): source_quant_dtype = None @@ -1424,6 +1663,15 @@ def __init__( base_quant_config, cache_config, topk_indices_buffer, + ( + _can_fuse_indexer_wk_weights_proj( + config, + model_quant_config, + [f"{prefix}.indexer"], + ) + if use_indexer_wk_weights_proj_fusion is None + else use_indexer_wk_weights_proj_fusion + ), f"{prefix}.indexer", ) else: @@ -1464,7 +1712,8 @@ def __init__( # When ATOM_ENABLE_DS_QKNORM_QUANT_FUSION is turned on, self.fuse_qknorm_quant is turned on only if FP8 or (use_triton_gemm() and FP4), self.prefix = prefix - self.quant_dtype = None + self.quant_dtype = layer_quant_dtype + self.qknorm_quant_type = layer_quant_type_value self.fuse_qknorm_quant = False # always fuse qknorm self.fuse_qknorm = ENABLE_DS_QKNORM_FUSION @@ -1472,7 +1721,6 @@ def __init__( if layer_quant_dtype == dtypes.fp8 or ( layer_quant_dtype == dtypes.fp4x2 and use_triton_gemm() ): - self.quant_dtype = layer_quant_dtype self.fuse_qknorm_quant = True def forward( @@ -1518,7 +1766,7 @@ def forward( dim=-1, ) # fuse q_c norm + kv_c norm + quant of hidden_states_or_q_c - if self.fuse_qknorm_quant: + if self.fuse_qknorm_quant or self.fuse_qknorm: ( (hidden_states_or_q_c, hidden_states_or_q_c_scale), _, @@ -1536,19 +1784,10 @@ def forward( shuffle=False, scale_shuffle_padding=False, group_size=128, + quant_type=self.qknorm_quant_type, output_unquantized_inp1=False, transpose_scale=True, ) - elif self.fuse_qknorm: - hidden_states_or_q_c, kv_c_normed = _fused_qk_rmsnorm( - q_c, - self.q_a_layernorm.weight, - self.q_a_layernorm.eps, - kv_c, - self.kv_a_layernorm.weight, - self.kv_a_layernorm.eps, - ) - hidden_states_or_q_c_scale = None else: hidden_states_or_q_c = self.q_a_layernorm(q_c) else: @@ -1580,7 +1819,6 @@ def forward( class DeepseekV2DecoderLayer(nn.Module): - def __init__( self, config: PretrainedConfig, @@ -1591,6 +1829,7 @@ def __init__( layer_num: int = 0, is_mtp_block: bool = False, alt_stream: Optional[torch.cuda.Stream] = None, + use_indexer_wk_weights_proj_fusion: Optional[bool] = None, ) -> None: super().__init__() self.hidden_size = config.hidden_size @@ -1615,6 +1854,7 @@ def __init__( prefix=f"{prefix}.self_attn", layer_num=layer_num, topk_indices_buffer=topk_indices_buffer, + use_indexer_wk_weights_proj_fusion=use_indexer_wk_weights_proj_fusion, ) # When ATOM_ENABLE_DS_INPUT_RMSNORM_QUANT_FUSION is turned on self.fuse_input_norm_quant is turned on only if use_triton_gemm and (FP8 or FP4), @@ -1627,6 +1867,11 @@ def __init__( if quant_config is None else quant_config.get_layer_quant_config(prefix).quant_dtype ) + self.input_norm_quant_type = ( + None + if quant_config is None + else quant_config.get_layer_quant_config(prefix).quant_type.value + ) self.fuse_input_norm_quant = False self.fuse_ar_input_norm = ENABLE_ALLREDUCE_RMSNORM_FUSION if quant_config is not None and ENABLE_DS_INPUT_RMSNORM_QUANT_FUSION: @@ -1710,6 +1955,7 @@ def forward( shuffle=True, scale_shuffle_padding=True, group_size=128, + quant_type=self.input_norm_quant_type, output_unquantized_inp1=False, transpose_scale=True, ) @@ -1728,6 +1974,7 @@ def forward( shuffle=True, scale_shuffle_padding=True, group_size=128, + quant_type=self.input_norm_quant_type, output_unquantized_inp1=False, transpose_scale=True, ) @@ -1779,6 +2026,7 @@ def __init__( atom_config: Config, prefix: str = "", layer_type: type[nn.Module] = DeepseekV2DecoderLayer, + use_indexer_wk_weights_proj_fusion: Optional[bool] = None, ): super().__init__() @@ -1823,6 +2071,7 @@ def __init__( quant_config=quant_config, layer_num=layer_num, alt_stream=_alt_stream, + use_indexer_wk_weights_proj_fusion=use_indexer_wk_weights_proj_fusion, ), prefix=f"{prefix}.layers", layer_num_offset=0, @@ -1837,6 +2086,8 @@ def __init__( ) else: self.norm = PPMissingLayer() + self.aux_hidden_state_layers: tuple[int, ...] = tuple() + self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( ["hidden_states", "residual"], config.hidden_size ) @@ -1850,7 +2101,9 @@ def forward( positions: torch.Tensor, intermediate_tensors: Optional[IntermediateTensors], inputs_embeds: Optional[torch.Tensor] = None, - ) -> Union[torch.Tensor, IntermediateTensors]: + ) -> Union[ + torch.Tensor, IntermediateTensors, Tuple[torch.Tensor, list[torch.Tensor]] + ]: if get_pp_group().is_first_rank: if inputs_embeds is not None: hidden_states = inputs_embeds @@ -1862,7 +2115,13 @@ def forward( hidden_states = intermediate_tensors["hidden_states"] residual = intermediate_tensors["residual"] - for layer in self.layers[self.start_layer : self.end_layer]: + aux_hidden_states = [] + for idx in range(self.start_layer, self.end_layer): + layer = self.layers[idx] + if idx in self.aux_hidden_state_layers: + aux_hidden_states.append( + hidden_states if residual is None else hidden_states + residual + ) hidden_states, residual = layer(positions, hidden_states, residual) if not get_pp_group().is_last_rank: @@ -1871,6 +2130,9 @@ def forward( ) hidden_states, _ = self.norm(hidden_states, residual) + + if aux_hidden_states: + return hidden_states, aux_hidden_states return hidden_states def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: @@ -1890,7 +2152,6 @@ def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: class DeepseekV2ForCausalLM(nn.Module): - def __init__( self, atom_config: Config, @@ -1903,6 +2164,23 @@ def __init__( self.config = config self.quant_config = quant_config + model_prefix = maybe_prefix(prefix, "model") + attn_module_list_cfg = getattr(config, "attn_module_list_cfg", None) + indexer_prefixes = [] + if isinstance(attn_module_list_cfg, (list, tuple)): + indexer_prefixes = [ + f"{model_prefix}.layers.{layer_idx}.self_attn.indexer" + for layer_idx, layer_cfg in enumerate(attn_module_list_cfg) + if isinstance(layer_cfg, dict) + and layer_cfg.get("attn_index") is not None + ] + if not indexer_prefixes: + indexer_prefixes = [f"{model_prefix}.layers.0.self_attn.indexer"] + use_indexer_wk_weights_proj_fusion = _can_fuse_indexer_wk_weights_proj( + config, + quant_config, + indexer_prefixes, + ) if hasattr(config, "q_lora_rank") and config.q_lora_rank is not None: self.packed_modules_mapping = { "q_a_proj": ("fused_qkv_a_proj", 0), @@ -1915,11 +2193,19 @@ def __init__( "gate_proj": ("gate_up_proj", 0), "up_proj": ("gate_up_proj", 1), } + if use_indexer_wk_weights_proj_fusion: + self.packed_modules_mapping.update( + { + "indexer.wk": ("indexer.wk_weights_proj", 0), + "indexer.weights_proj": ("indexer.wk_weights_proj", 1), + } + ) self.model = DeepseekV2Model( atom_config=atom_config, - prefix=maybe_prefix(prefix, "model"), + prefix=model_prefix, layer_type=layer_type, + use_indexer_wk_weights_proj_fusion=use_indexer_wk_weights_proj_fusion, ) if get_pp_group().is_last_rank: self.lm_head = ParallelLMHead( @@ -1971,22 +2257,36 @@ def make_empty_intermediate_tensors( } ) + def set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None: + self.model.aux_hidden_state_layers = layers + + def get_eagle3_aux_hidden_state_layers(self) -> tuple[int, ...]: + """Default Eagle3 aux hidden-state layer ids: early / middle / late + of the target model. Aligned with vLLM's default (see + vllm/model_executor/models/deepseek_v2.py). + """ + num_layers = len(self.model.layers) + return (2, num_layers // 2, num_layers - 3) + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: return self.model.get_expert_mapping() class DeepseekV3ForCausalLM(DeepseekV2ForCausalLM): - # DeepSeek-V3.2's indexer weights_proj are not quantized, but they are not listed in - # model's quant exclude mapping. So we add it to quant_default_exclude_layers. - quant_default_exclude_layers: list[str] = ["*.indexer.weights_proj"] + # DeepSeek-V3.2's indexer weights projection is BF16. Keep the original + # checkpoint path and the fused ATOM path excluded from default quantization. + quant_default_exclude_layers: list[str] = [ + "*.indexer.weights_proj", + "*.indexer.wk_weights_proj", + ] class GlmMoeDsaForCausalLM(DeepseekV2ForCausalLM): """GLM 5.0 MoE (structurally similar to DeepSeek v3.2). Reuses DeepseekV2 implementation.""" # GLM-5's HF quant config uses `indexers_proj` in modules_to_not_convert, but - # the ATOM module path is `indexer.weights_proj`. Declaring the mapping here - # keeps the translation co-located with the model and out of config.py. + # the unfused ATOM module path is `indexer.weights_proj`. Keep that path + # excluded so FP4/MXFP4 fallback does not quantize the BF16 projection. quant_exclude_name_mapping: dict[str, str] = { # HF quant config uses "indexers_proj" but the ATOM module path is # "indexer.weights_proj". str.replace translates each exclude entry. diff --git a/atom/models/deepseek_v4.py b/atom/models/deepseek_v4.py index 002511913c..7bb2389416 100644 --- a/atom/models/deepseek_v4.py +++ b/atom/models/deepseek_v4.py @@ -33,12 +33,23 @@ cp_gather_indexer_k_quant_cache, dtypes, get_hip_quant, - indexer_k_quant_and_cache, + silu_and_mul as aiter_silu_and_mul, +) +from aiter.jit.utils.torch_guard import torch_compile_guard +from aiter.dist.communication_op import ( + tensor_model_parallel_all_reduce, +) +from aiter.dist.parallel_state import ( + get_tensor_model_parallel_world_size, ) -from aiter.dist.communication_op import tensor_model_parallel_all_reduce -from aiter.dist.parallel_state import get_tensor_model_parallel_world_size from aiter.ops.topk import top_k_per_row_decode, top_k_per_row_prefill from aiter.ops.triton.fp8_mqa_logits import fp8_mqa_logits +from aiter.ops.triton.fusions.fused_clamp_act_mul import ( + fused_clamp_act_mul, +) +from aiter.ops.triton.fusions.fused_reduce_qk_norm_rope_swa_write import ( + fused_reduce_qk_norm_rope_swa_write, +) from aiter.ops.triton.pa_mqa_logits import deepgemm_fp8_paged_mqa_logits from atom.config import ( Config, @@ -48,8 +59,8 @@ get_current_atom_config, ) from atom.model_loader.loader import WeightsMapper -from atom.model_ops.embed_head import VocabParallelEmbedding -from atom.model_ops.layernorm import RMSNorm, rmsnorm2d_fwd_ +from atom.model_ops.embed_head import ParallelLMHead, VocabParallelEmbedding +from atom.model_ops.layernorm import DualRMSNorm, RMSNorm, rmsnorm2d_fwd_ from atom.model_ops.triton_rmsnorm_nw import rmsnorm_nw from atom.model_ops.linear import ( ColumnParallelLinear, @@ -66,7 +77,7 @@ hc_split_sinkhorn, ) from atom.model_ops.utils import atom_parameter -from atom.utils import envs +from atom.utils import envs, mark_spliting_op # Side-effect import: registers `torch.ops.aiter.maybe_dual_stream_forward` # (shared with deepseek_v2) and `torch.ops.aiter.indexer_score_topk` (V4-only). @@ -86,6 +97,7 @@ update_compressor_states, ) from atom.utils.forward_context import get_forward_context +from atom.utils.decorators import support_torch_compile # --------------------------------------------------------------------------- # Classical KV cache scatter / gather helpers (PR3-pre2c-B). @@ -111,6 +123,9 @@ # forward burns syscalls (V4-Pro: 64 layers × multiple sites per call). _V4_FORCE_UE8M0_QUANT = os.environ.get("V4_FORCE_UE8M0_QUANT", "0") == "1" _V4_USE_REF_QUANT = os.environ.get("V4_USE_REF_QUANT", "0") == "1" +# Fused-kernel switches. Default off; flip via env to A/B against the eager path. +_V4_USE_TRITON_FUSION = os.environ.get("ATOM_V4_USE_TRITON_FUSION", "0") == "1" +ENABLE_DS_QKNORM_QUANT_FUSION = envs.ATOM_ENABLE_DS_QKNORM_QUANT_FUSION def _rmsnorm_nw(x: torch.Tensor, eps: float, dim: int) -> torch.Tensor: @@ -120,6 +135,174 @@ def _rmsnorm_nw(x: torch.Tensor, eps: float, dim: int) -> torch.Tensor: return rmsnorm2d_fwd_(x, ones, eps, dim) +def _fused_qk_norm_rope_swa_write_fake( + q: torch.Tensor, + kv: torch.Tensor, + cos_cache: torch.Tensor, + sin_cache: torch.Tensor, + positions: torch.Tensor, + n_local_heads: int, + head_dim: int, + rope_head_dim: int, + kv_weight: torch.Tensor, + eps: float, + win: int, + swa_write_indices: Optional[torch.Tensor] = None, + batch_id_per_token: Optional[torch.Tensor] = None, + state_slot_mapping: Optional[torch.Tensor] = None, + swa_kv: Optional[torch.Tensor] = None, +) -> torch.Tensor: + M = q.shape[0] + return torch.empty( + (M, n_local_heads, head_dim), + dtype=torch.bfloat16, + device=q.device, + ) + + +@torch_compile_guard(gen_fake=_fused_qk_norm_rope_swa_write_fake) +def fused_qk_norm_rope_swa_write( + q: torch.Tensor, + kv: torch.Tensor, + cos_cache: torch.Tensor, + sin_cache: torch.Tensor, + positions: torch.Tensor, + n_local_heads: int, + head_dim: int, + rope_head_dim: int, + kv_weight: torch.Tensor, + eps: float, + win: int, + swa_write_indices: Optional[torch.Tensor] = None, + batch_id_per_token: Optional[torch.Tensor] = None, + state_slot_mapping: Optional[torch.Tensor] = None, + swa_kv: Optional[torch.Tensor] = None, +) -> torch.Tensor: + """Fused wq_b GEMM (a8w8 1x128 blockscale) + per-head RMSNorm-nw + RoPE on + q/kv tail (+ SWA write) in a single triton kernel. + + `kv` must already be kv_norm-applied; the kernel does not weight-norm kv. + `kv` is RoPE-mutated in place. When all SWA tensors are provided, the + kernel also writes the windowed kv into `swa_kv`; the caller must then + skip the standalone `swa_write` for those rows. + """ + num_tokens = q.shape[0] + if num_tokens <= 64: + q_out = torch.empty( + (num_tokens, n_local_heads, head_dim), + dtype=torch.bfloat16, + device=q.device, + ) + fused_reduce_qk_norm_rope_swa_write( + q, + kv, + None, + kv_weight, + eps, + eps, + rope_head_dim, + cos_cache, + sin_cache, + positions, + q_out=q_out, + is_neox=False, + dtype=torch.bfloat16, + write_indices=swa_write_indices, + batch_id_per_token=batch_id_per_token, + state_slot_mapping=state_slot_mapping, + swa_kv=swa_kv, + win=win, + ) + else: + q = q.view(num_tokens, n_local_heads, head_dim) + q_out = _rmsnorm_nw(q, eps, head_dim) + kv = rmsnorm2d_fwd_(kv, kv_weight, eps, kv.shape[-1]) + aiter.rope_cached_positions_2c_fwd_inplace( + q_out[..., -rope_head_dim:].view(1, num_tokens, -1, rope_head_dim), + kv[..., -rope_head_dim:].view(1, num_tokens, -1, rope_head_dim), + cos_cache, + sin_cache, + positions.view(1, num_tokens), + 1, + reuse_freqs_front_part=True, + nope_first=False, + ) + if swa_write_indices is not None: + swa_write( + kv, + swa_write_indices, + positions, + batch_id_per_token, + state_slot_mapping, + swa_kv, + win, + ) + return q_out + + +def _make_weightless_rmsnorm(dim: int, eps: float) -> RMSNorm: + """Build an `RMSNorm(dim, eps)` whose `.weight` is `None`. + + Drops the learnable Parameter so `state_dict()` is empty for this + submodule and `load_model` doesn't expect a weight from disk (no + "unloaded parameter" warning). `DualRMSNorm` recognizes the + sentinel `weight is None` and instructs the fused kernel to skip + both the q_weight load and the multiply (Q_HAS_WEIGHT=False), + matching the prior `_rmsnorm_nw(x, eps, dim)` behavior. + """ + norm = RMSNorm(dim, eps) + del norm.weight # remove Parameter from `_parameters` + norm.weight = None # sentinel for downstream consumers + return norm + + +def _hc_head_reduce_fake( + x: torch.Tensor, + hc_fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + norm_eps: float, + hc_eps: float, +) -> torch.Tensor: + return torch.empty(x.shape[0], x.shape[-1], dtype=x.dtype, device=x.device) + + +@torch_compile_guard(gen_fake=_hc_head_reduce_fake, mutates_args=[]) +def _hc_head_reduce( + x: torch.Tensor, + hc_fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + norm_eps: float, + hc_eps: float, +) -> torch.Tensor: + x_flat = x.flatten(-2) + x_normed = _rmsnorm_nw(x_flat, norm_eps, x_flat.shape[-1]) + mixes = F.linear(x_normed.float(), hc_fn) + pre = torch.sigmoid(mixes * hc_scale + hc_base) + hc_eps + y = torch.sum(pre.unsqueeze(-1) * x, dim=-2) + return y.to(x.dtype) + + +def _v4_attention_fake( + x: torch.Tensor, + positions: torch.Tensor, + layer_name: str, +) -> torch.Tensor: + return torch.empty_like(x) + + +@mark_spliting_op(is_custom=True, gen_fake=_v4_attention_fake, mutates_args=[]) +def v4_attention_with_output( + x: torch.Tensor, + positions: torch.Tensor, + layer_name: str, +) -> torch.Tensor: + atom_config = get_current_atom_config() + self = atom_config.compilation_config.static_forward_context[layer_name] + return self.forward_impl(x, positions) + + # --------------------------------------------------------------------------- # Config wrapper # --------------------------------------------------------------------------- @@ -615,17 +798,15 @@ def __init__( # External tensors — assigned by the owning Attention / Indexer at first forward. self.kv_cache: Optional[torch.Tensor] = None self.rotary_emb: Optional[_V4RoPE] = None - # Shared compress-output buffer (set by V4 builder.build_kv_cache_tensor - # to a per-kind torch.empty pre-allocation in `forward_vars`). When - # present, fused_compress_attn writes into it via `out=` so the GPU - # pointer stays stable across captures (CUDAGraph requirement); when - # None, falls back to the eager fresh-allocation path. - self.compress_out: Optional[torch.Tensor] = None + # FP8 quant path only: strided fp32 view of the per-block scale region + # of `self.kv_cache`. Bound by the V4 builder when `kv_cache.dtype` is + # FP8 (Indexer-inner Compressor); None for BF16 cache (Main path). + self.cache_scale: Optional[torch.Tensor] = None # State cache (per paper §3.6.1 "uncompressed tail + B-side overlap # window" portion). Indexed as a single ring buffer of size - # `coff * compress_ratio` by `pos % STATE_SIZE` per token — no - # segment switching, no roll. The `forward` softmax-pool consumer + # `ring_size` (≥ coff * compress_ratio) by `pos % ring_size` per token + # — no segment switching, no roll. The `forward` softmax-pool consumer # resolves A-side (current block) vs B-side (previous block) by # block-id parity (`comp_id % 2`). # @@ -633,9 +814,12 @@ def __init__( # runs before allocate_kv_cache → build_kv_cache_tensor) sees a # valid tensor; afterwards `DeepseekV4AttentionMetadataBuilder. # build_kv_cache_tensor` setattr-replaces these attributes with - # views of the per-request cache pool (shape - # `[max_num_seqs, coff*ratio, coff*head_dim]`). The 1-slot init - # buffers (≈9 MB total across all layers) are GC'd once replaced. + # views of the per-request cache pool whose second dim is the real + # ring_size = coff*ratio + max_spec_steps + 1 (spec) or coff*ratio + # (non-spec). The 1-slot init buffers (≈9 MB total across all layers) + # are GC'd once replaced before any real kernel call, so the + # placeholder's smaller second dim never actually flows through the + # kernel's `state_size >= K_pool` assertion. self.register_buffer( "kv_state", torch.zeros( @@ -662,45 +846,41 @@ def forward( plan: "CompressPlan", state_slot_mapping: torch.Tensor, # [bs] int32 block_tables: Optional[torch.Tensor] = None, # [bs, max_blocks_per_seq] int32 - out_dtype: Optional[torch.dtype] = None, - idx_slot_mapping: Optional[torch.Tensor] = None, # [num_compress] int64 - ) -> Optional[torch.Tensor]: # [num_compress, head_dim] BF16 (None if plan empty) + ) -> None: """Batched plan-style compress: one fused kernel call for the whole fwd's batch (across all seqs). - Single fused Triton kernel does pool + RMSNorm + RoPE + bf16 kv_cache - scatter in one launch. Each compression boundary across the batch is - one row in `plan.compress_plan_gpu`. State cache update fires after - (write order critical — fused kernel reads state-cache-as-of-previous- - fwd; update_compressor_states overwrites for next fwd). - - TODO: quant (FP8 ue8m0 for CSA Main, FP4 + Hadamard for Indexer) is - NOT applied for the Main path; see class docstring. The Indexer-inner - path *does* apply FP8 quantization via `indexer_k_quant_and_cache` - when `idx_slot_mapping` is provided. + Single fused Triton kernel does pool + RMSNorm + RoPE + cache scatter + in one launch. Each compression boundary across the batch is one row + in `plan.compress_plan_gpu`. State cache update fires after (write + order critical — fused kernel reads state-cache-as-of-previous-fwd; + `update_compressor_states` overwrites for next fwd). + + Quant mode is auto-selected by `self.kv_cache.dtype`: + - BF16 cache (CSA Main / HCA Main): raw BF16 row write into + `self.kv_cache` (consumed by paged_decode/paged_prefill via + `unified_kv` per-fwd indices). + - FP8 cache (Indexer-inner): per-row amax → ue8m0 scale → fp8 cast + → preshuffled (MFMA 16x16 tile) write into `self.kv_cache`, plus + fp32 scale into `self.cache_scale` (a strided view of the same + allocation built by the V4 builder). Bit-exact with + `indexer_k_quant_and_cache` / `cp_gather_indexer_k_quant_cache` + (cache_kernels.cu:1145+). + + Side-effecting only — no return value (cache scatter IS the output). + + TODO: QAT for the BF16 Main path (FP8 round-trip per Compressor + docstring) is not yet fused. End-to-end accuracy unaffected today + because the input act_quant simulation is applied upstream. Args: x: [num_tokens, dim] flat ragged batch hidden state. plan: CompressPlan from attn_metadata.compress_plans[ratio] (or a synthetic bs=1 plan during warmup). - state_slot_mapping: [bs] int32 — per-seq state cache slot - (attn_metadata.state_slot_mapping). + state_slot_mapping: [bs] int32 — per-seq state cache slot. block_tables: [bs, max_blocks_per_seq] int32 — physical block IDs per seq; None during warmup (skips kv_cache scatter). - out_dtype: override kernel output dtype. Defaults to bf16. - idx_slot_mapping: [num_compress] int64 — global flat slot in the - FP8 indexer cache pool, one per compress row (kernel - signature requires int64). When provided, the fused - kernel skips its BF16 scatter and we instead call - `indexer_k_quant_and_cache` to FP8-quantize+write each - row into `self.kv_cache` (interleaved [head_dim] FP8 + - 4-byte fp32 scale per row, dtypes.fp8 storage). - Only set by the CSA-inner-Compressor caller. - - Returns: - Compressed KV `[num_compress, head_dim]` post-norm post-rope BF16, - in plan order (= per-seq grouped by `plan.cu_compress_cpu`). - Returns None if `plan.num_compress == 0`. + Required for the Indexer FP8 path (slot resolution). """ assert self.rotary_emb is not None, "compressor.rotary_emb must be set by owner" assert ( @@ -710,15 +890,15 @@ def forward( overlap = self.overlap d = self.head_dim rd = self.rope_head_dim - if out_dtype is None: - out_dtype = torch.bfloat16 if x.dtype == torch.float32 else x.dtype - # Single fused BF16 GEMM via tgemm; otype=fp32 keeps softmax-pool - # accumulator in fp32. torch.split returns zero-copy strided views; - # downstream kernels (fused_compress_attn, update_compressor_states) - # accept strided kv/score (only inner stride must be 1). + # Single fused BF16 GEMM via tgemm. (Probing whether dropping the + # otype=fp32 upcast — relying on fused_compress_attn's internal fp32 + # accumulator instead — is accuracy-neutral.) torch.split returns + # zero-copy strided views; downstream kernels (fused_compress_attn, + # update_compressor_states) accept strided kv/score (only inner + # stride must be 1). coff_d = (1 + overlap) * d - combined = self.wkv_gate(x, otype=torch.float32) + combined = self.wkv_gate(x) kv, score = torch.split(combined, [coff_d, coff_d], dim=-1) # ====== Unified fused kernel path (CSA + Indexer) ====== @@ -726,29 +906,21 @@ def forward( # PREVIOUS-fwd. `update_compressor_states` overwrites them with this # fwd's data for the NEXT fwd's overlap — must run AFTER the fused # kernel. - # - # The kernel does pool + RMSNorm + RoPE + (optional) BF16 store to - # kv_cache. For the Indexer-inner path (`idx_slot_mapping is not None`) - # we bypass the kernel scatter and instead call - # `indexer_k_quant_and_cache` separately, which FP8-quantizes the - # post-process K row and writes the (FP8 + scale) interleaved layout - # into the uint8 indexer pool. cos_cache, sin_cache = self.rotary_emb._caches(x.device) - is_idx_fp8 = idx_slot_mapping is not None - # Warmup runs through the same path: kv_cache is None until the V4 - # builder binds it post-warmup, so the kernel skips the scatter. - # Indexer-FP8 path always skips the kernel scatter (separate write). - scatter_kv_cache = ( - None - if is_idx_fp8 - else (self.kv_cache if block_tables is not None else None) - ) - scatter_block_tables = ( - None - if is_idx_fp8 - else (block_tables if scatter_kv_cache is not None else None) - ) - out = fused_compress_attn( + # Quant path triggers when the bound cache is FP8 (Indexer-inner). + # `self.cache_scale` is bound alongside `self.kv_cache` by the V4 + # builder when the cache is FP8 (strided fp32 view of the per-block + # scale region). + is_quant = self.kv_cache is not None and self.kv_cache.dtype != torch.bfloat16 + # Skip the kernel's cache scatter during warmup (kv_cache/block_tables + # not yet bound). + if block_tables is None or self.kv_cache is None: + scatter_kv_cache = None + scatter_block_tables = None + else: + scatter_kv_cache = self.kv_cache + scatter_block_tables = block_tables + fused_compress_attn( kv_in=kv, score_in=score, kv_state=self.kv_state, @@ -767,21 +939,12 @@ def forward( ratio=ratio, head_dim=d, rope_head_dim=rd, - out_dtype=out_dtype, - out=self.compress_out, + quant=is_quant, + cache_scale=self.cache_scale if is_quant else None, + use_ue8m0=(self.scale_fmt == "ue8m0"), + preshuffle=True, + fp8_max=(torch.finfo(self.kv_cache.dtype).max if is_quant else None), ) - # Indexer-FP8 path: quantize the post-RoPE BF16 K rows into the uint8 - # FP8+scale layout of `self.kv_cache`. `out` is [num_compress, head_dim] - # BF16 in plan order (matches `idx_slot_mapping` row order). - if is_idx_fp8 and out is not None and self.kv_cache is not None: - indexer_k_quant_and_cache( - out, - self.kv_cache, - idx_slot_mapping, - d, # quant_block_size = head_dim (one fp32 scale per row) - self.scale_fmt, - preshuffle=True, - ) update_compressor_states( kv, score, @@ -794,7 +957,6 @@ def forward( ratio=ratio, overlap=overlap, ) - return out class Indexer(nn.Module): @@ -882,8 +1044,11 @@ def __init__(self, args: DeepseekV4Args, compress_ratio: int = 4, prefix: str = def forward_batched( self, x_full: torch.Tensor, # [total_tokens, dim] - qr_full: torch.Tensor, # [total_tokens, q_lora_rank] + qr_full: torch.Tensor, # [total_tokens, q_lora_rank] — fp8 when qr_full_scale given positions: torch.Tensor, # [total_tokens] + qr_full_scale: Optional[ + torch.Tensor + ] = None, # per_1x128 scale paired with qr_full ) -> torch.Tensor: """Q proj + RoPE + FP8-quant + weights compute (have module state), then dispatch to `torch.ops.aiter.indexer_score_topk`, which calls @@ -908,7 +1073,9 @@ def forward_batched( # Q proj + RoPE + rotate (batched). rotary_emb internally reshapes # to (1, num_tokens, -1, rotary_dim) so the input doesn't need an # explicit batch dim. rotate_activation is last-dim-only. - q = self.wq_b(qr_full).view(total_tokens, self.n_heads, self.head_dim) + q = self.wq_b(qr_full, x_scale=qr_full_scale).view( + total_tokens, self.n_heads, self.head_dim + ) # self.rotary_emb(positions, q[..., -rd:]) # q = rotate_activation(q) cos, sin = self.rotary_emb._caches(q.device) @@ -1085,13 +1252,18 @@ def _score_topk_decode( kv_cache_4d = self.kv_cache.unsqueeze( -2 ) # [num_blocks, k1_csa, 1, head_dim+scale_dim] uint8 - # Pre-allocated decode logits buffer (sized [max_bs, max_model_len_idx] - # in builder, sliced here to [total_tokens, max_model_len_idx]). - # No `fill_(-inf)` needed — `top_k_per_row_decode` bounds each row - # by `n_committed_per_seq[batch]` so unwritten cols are never picked. - logits = indexer_meta["decode_logits_gpu"][ - :total_tokens - ] # [total_tokens, max_model_len_idx] fp32 + # Per-fwd write-once GPU scratch — no CPU mirror, no cross-fwd state. + # Under CUDAGraph capture, torch allocates from the graph's private + # memory pool and the address is stable across replays at this + # captured `total_tokens`. No `fill_(-inf)` needed — + # `top_k_per_row_decode` bounds each row by `n_committed_per_seq[batch]` + # so unwritten cols are never picked. + logits = torch.empty( + total_tokens, + self._max_model_len_idx, + dtype=torch.float32, + device=q_fp8.device, + ) deepgemm_fp8_paged_mqa_logits( q_4d, kv_cache_4d, @@ -1103,13 +1275,12 @@ def _score_topk_decode( KVBlockSize=self.kv_cache.size(1), # k1_csa = 32 Preshuffle=True, ) - # Pre-allocated indices buffer — fixed [max_bs, index_topk] int32 in - # builder, sliced to [total_tokens, index_topk] for this fwd. Kernel - # writes exactly `index_topk` ints per row (valid seq-local indices - # then -1 sentinels), no overflow risk. - topk_local = indexer_meta["decode_topk_indices_gpu"][ - :total_tokens - ] # [total_tokens, index_topk] int32 + # Per-fwd write-once int32 scratch. Kernel writes exactly `index_topk` + # ints per row (valid seq-local indices then -1 sentinels). CG-safe + # for the same reason as `logits` above. + topk_local = torch.empty( + total_tokens, self.index_topk, dtype=torch.int32, device=q_fp8.device + ) top_k_per_row_decode( logits, next_n, @@ -1148,7 +1319,14 @@ class DeepseekV4Attention(nn.Module): - attn_sink: per-head learnable logit added only to softmax denominator """ - def __init__(self, layer_id: int, args: DeepseekV4Args, prefix: str = ""): + def __init__( + self, + layer_id: int, + args: DeepseekV4Args, + prefix: str = "", + alt_stream: Optional[torch.cuda.Stream] = None, + compress_stream: Optional[torch.cuda.Stream] = None, + ): super().__init__() self.layer_id = layer_id self.dim = args.dim @@ -1193,8 +1371,22 @@ def __init__(self, layer_id: int, args: DeepseekV4Args, prefix: str = ""): quant_config=qc, prefix=f"{p}.wqkv_a", ) - self.q_norm = RMSNorm(self.q_lora_rank, self.eps) - self.q_norm2 = RMSNorm(self.head_dim, self.eps) + # Fuse q_norm + per_1x128 FP8 quant: kernel emits (qr_fp8, qr_scale) + # in one launch, both wq_b consumers (outer ColumnParallel + Indexer + # ReplicatedLinear) skip their own input quant. + self.q_norm = RMSNorm( + self.q_lora_rank, + self.eps, + fused_quant=True, + quant_config=qc, + prefix=f"{p}.q_norm", + ) + # q_norm2: per-head Q normalization. The checkpoint has no + # `q_norm2.weight` entry, so the module is constructed with + # `weight=None` (no learnable Parameter) — `DualRMSNorm` reads the + # sentinel and tells the kernel to skip the q_weight load entirely. + # Behavior is identical to the prior `_rmsnorm_nw(q, eps, head_dim)`. + self.q_norm2 = _make_weightless_rmsnorm(self.head_dim, self.eps) self.wq_b = ColumnParallelLinear( self.q_lora_rank, self.n_heads * self.head_dim, @@ -1203,6 +1395,19 @@ def __init__(self, layer_id: int, args: DeepseekV4Args, prefix: str = ""): prefix=f"{p}.wq_b", ) self.kv_norm = RMSNorm(self.head_dim, self.eps) + # Fused per-head Q RMSNorm (identity weight) + KV RMSNorm — both have + # feature dim = head_dim, single Triton kernel via DualRMSNorm. + # `q_norm2.weight is None` tells the kernel's Q_HAS_WEIGHT=False + # path to skip the q_weight load (Q side has no learnable weight; + # equivalent to the prior `_rmsnorm_nw` helper). + self.qk_norm = DualRMSNorm( + self.q_norm2, + self.kv_norm, + num_q_heads=self.n_local_heads, + num_kv_heads=1, + head_dim=self.head_dim, + prefix=f"{p}.qk_norm", + ) # wo_a: grouped LoRA — V4QuantConfig forces this BF16 even though disk is FP8. # The grouped einsum (`bsgd,grd->bsgr`) needs BF16 weights; aiter has no FP8 einsum. self.wo_a = ColumnParallelLinear( @@ -1291,6 +1496,18 @@ def __init__(self, layer_id: int, args: DeepseekV4Args, prefix: str = ""): self.indexer.rotary_emb = self.rotary_emb self.indexer.compressor.rotary_emb = self.rotary_emb + self.alt_stream = alt_stream + self.compress_stream = compress_stream + self._use_async_compress = ( + self.alt_stream is not None and self.compressor is not None + ) + + self.use_fuse_qk_norm_rope_swa_write = _V4_USE_TRITON_FUSION + + self.layer_name = prefix + atom_config = get_current_atom_config() + atom_config.compilation_config.static_forward_context[self.layer_name] = self + def process_weights_after_loading(self) -> None: """Dequant wo_a (FP8 + e8m0 block scale) → BF16 in place. @@ -1341,7 +1558,43 @@ def process_weights_after_loading(self) -> None: # batched-FP8 kernel. See attention_mla.py:211 for reference. self.wo_a.quant_type = QuantType.No + def _launch_compressors_async(self, x, plan, state_slot_mapping, block_tables): + """Fire Compressor(s) on side streams, return immediately. + + Main Compressor → alt_stream (CSA + HCA). + Indexer Compressor → compress_stream (CSA only). + Waits resolve instantly: side streams ~25us, main Q/KV chain ~87us.""" + current_stream = get_forward_context().main_stream + if self.compressor is not None and self.alt_stream is not None: + self.alt_stream.wait_stream(current_stream) + if self.indexer is not None and self.compress_stream is not None: + self.compress_stream.wait_stream(current_stream) + + if self.compressor is not None and self.alt_stream is not None: + with torch.cuda.stream(self.alt_stream): + self.compressor( + x, + plan=plan, + state_slot_mapping=state_slot_mapping, + block_tables=block_tables, + ) + if self.indexer is not None and self.compress_stream is not None: + with torch.cuda.stream(self.compress_stream): + self.indexer.compressor( + x, + plan=plan, + state_slot_mapping=state_slot_mapping, + block_tables=block_tables, + ) + def forward( + self, + x: torch.Tensor, + positions: torch.Tensor, + ) -> torch.Tensor: + return torch.ops.aiter.v4_attention_with_output(x, positions, self.layer_name) + + def forward_impl( self, x: torch.Tensor, # [num_tokens, dim] flat ragged-batch hidden state positions: torch.Tensor, # [num_tokens] int absolute token positions @@ -1366,32 +1619,36 @@ def forward( if get_forward_context().context.is_dummy_run: return torch.zeros_like(x) num_tokens = x.size(0) - win = self.window_size + # Async-compress (alt_stream main Compressor + compress_stream + # indexer.compressor) is only safe inside CUDAGraph capture: graph + # records the fork-join edges and replay re-uses the same stream + # layout. In eager mode the side-stream launches accumulate across + # 60 layers and deadlock the hipStream queue when the first splitk + # GEMM kernel allocates its workspace — verified hang on both + # small (~800-token) and large (>2k-token) prefill batches in + # eager. Replay does not re-execute this Python code so the flag + # doesn't matter then. + use_async_compress = ( + self._use_async_compress and get_forward_context().in_hipgraph + ) + # SWA ring-slot count per req (= window_size + max_spec_steps for + # MTP-aware cache). Sourced from the bound cache to avoid threading + # `max_spec_steps` through V4Args; for non-MTP this equals + # `self.window_size`. Used as the modulo in `swa_write` (and matches + # the per-row case_c modulo in window_topk). + cache_size = self.swa_kv.shape[1] ratio = self.compress_ratio rd = self.rope_head_dim - # ----- Batched ops on full flat tensors ----- - # `_V4_FORCE_UE8M0_QUANT` (module-level): round-trip x/qr to ue8m0-FP8 - # to mirror the reference's `act_quant(scale_fmt="ue8m0")` Linear-input - # quantization. EXPERIMENT only. - if _V4_FORCE_UE8M0_QUANT: - x = x.clone() - act_quant_inplace(x, 128, "ue8m0") - # Single fused FP8 GEMM for [wq_a; wkv]; torch.split returns zero-copy views. - qkv_a = self.wqkv_a(x) - q_lora, kv_pre = torch.split(qkv_a, [self.q_lora_rank, self.head_dim], dim=-1) - qr = self.q_norm(q_lora) # [num_tokens, q_lora_rank] shared with Indexer - if _V4_FORCE_UE8M0_QUANT: - qr = qr.clone() - act_quant_inplace(qr, 128, "ue8m0") - q = self.wq_b(qr).view(num_tokens, self.n_local_heads, self.head_dim) - q = _rmsnorm_nw(q, self.eps, self.head_dim) - # q [S, H, D] / kv [S, head_dim] — rotary_emb internally unsqueezes - # to (1, num_tokens, ...) for aiter's per-position rope kernel. - kv = self.kv_norm(kv_pre) - self.rotary_emb(positions, q[..., -rd:], kv[..., -rd:]) - if _V4_USE_REF_QUANT: - act_quant_inplace(kv[..., :-rd], 64, self.scale_fmt) + # Idempotent one-time plumb of rotary_emb into compressor / indexer + # (and the indexer's inner compressor). `rotary_emb` is set by the + # owning layer after __init__, so this can't move into __init__. + if self.compress_ratio and self.compressor.rotary_emb is None: + self.compressor.rotary_emb = self.rotary_emb + if self.indexer is not None: + self.indexer.rotary_emb = self.rotary_emb + self.indexer.compressor.rotary_emb = self.rotary_emb + # ===== Per-fwd metadata (built once in prepare_prefill/decode). ===== # All per-fwd state read once. Production prepare_decode/prefill # always populates these; warmup goes through the same path @@ -1400,83 +1657,156 @@ def forward( # compress_plans, swa_write_indices, ...) is well-typed for pyright. attn_md = cast("AttentionMetaData_DSV4", get_forward_context().attn_metadata) compress_plans = attn_md.compress_plans - v4_indexer_meta = attn_md.indexer_meta swa_write_indices = attn_md.swa_write_indices v4_batch_id_per_token = attn_md.batch_id_per_token block_tables_gpu = attn_md.block_tables state_slot_mapping = attn_md.state_slot_mapping - # ===== Batched compressor + Indexer (ONCE per layer) ===== - # State cache reset on fresh prefill is redundant: SWA's start_pos==0 - # path uses raw seq_kv, and fused_compress's state-cache reads are - # already masked by `s < 0` for fresh prefill (compress_plan.py:88 + - # fused_compress.py:124-127). - # Compressor's return value is unused — paged_decode/paged_prefill read - # the scattered compress entries from `unified_kv` via per-fwd indices. - # Called purely for its side-effect (scatter into self.kv_cache). + # ----- Batched ops on full flat tensors ----- + # `_V4_FORCE_UE8M0_QUANT` (module-level): round-trip x/qr to ue8m0-FP8 + # to mirror the reference's `act_quant(scale_fmt="ue8m0")` Linear-input + # quantization. EXPERIMENT only. + if _V4_FORCE_UE8M0_QUANT: + x = x.clone() + act_quant_inplace(x, 128, "ue8m0") + + attn_md = cast("AttentionMetaData_DSV4", get_forward_context().attn_metadata) + compress_plans = attn_md.compress_plans + swa_write_indices = attn_md.swa_write_indices + v4_batch_id_per_token = attn_md.batch_id_per_token + block_tables_gpu = attn_md.block_tables + state_slot_mapping = attn_md.state_slot_mapping plan_for_layer = compress_plans[ratio] if ratio else None - if ratio: - self.compressor( - x, - plan=plan_for_layer, - state_slot_mapping=state_slot_mapping, - block_tables=block_tables_gpu, + + # ===== Triple-stream: Q/KV path + both Compressors in parallel ===== + current_stream = get_forward_context().main_stream + if use_async_compress: + self._launch_compressors_async( + x, plan_for_layer, state_slot_mapping, block_tables_gpu ) - if self.indexer is not None: # i.e. CSA layer, compress_ratio == 4 - # Indexer's inner compressor populates the indexer kv_cache (FP8 - # via indexer_k_quant_and_cache when idx_slot_mapping is set); - # the outer `forward_batched` reads `v4_indexer_meta` (built once - # per fwd) so it has zero per-layer H2D / CPU index math. - idx_slot_mapping = v4_indexer_meta.get("compress_slot_mapping_gpu") - self.indexer.compressor( - x, - plan=plan_for_layer, - state_slot_mapping=state_slot_mapping, - block_tables=block_tables_gpu, - idx_slot_mapping=idx_slot_mapping, + + # ----- Q/KV projections (main stream) ----- + qkv_a = self.wqkv_a(x) + q_lora, kv_pre = torch.split(qkv_a, [self.q_lora_rank, self.head_dim], dim=-1) + assert ( + not _V4_FORCE_UE8M0_QUANT + ), "_V4_FORCE_UE8M0_QUANT incompatible with fused q_norm quant (qr is already FP8)" + qr, qr_scale = self.q_norm(q_lora) + q = self.wq_b(qr, x_scale=qr_scale) + if attn_md.is_pure_decode and self.use_fuse_qk_norm_rope_swa_write: + # Fused: wq_b GEMM (a8w8 1x128 blockscale) + per-head RMSNorm-nw + # + RoPE on q tail + RoPE on kv tail (+ SWA write) in one triton + # kernel. KV RMSNorm stays out (kernel doesn't apply weighted norm + # to kv); the standalone swa_write below is gated off when this + # path runs since the kernel already wrote the window slots. + cos_cache, sin_cache = self.rotary_emb._caches(x.device) + if swa_write_indices is not None: + write_indices = swa_write_indices + batch_id = v4_batch_id_per_token + slot_map = state_slot_mapping + swa_kv_buf = self.swa_kv + else: + write_indices = None + batch_id = None + slot_map = None + swa_kv_buf = None + q = fused_qk_norm_rope_swa_write( + q, + kv_pre, + cos_cache, + sin_cache, + positions, + self.n_local_heads, + self.head_dim, + rd, + self.kv_norm.weight, + self.eps, + cache_size, + swa_write_indices=write_indices, + batch_id_per_token=batch_id, + state_slot_mapping=slot_map, + swa_kv=swa_kv_buf, ) + else: + # Flat q_flat is [num_tokens, n_local_heads * head_dim]; DualRMSNorm + # views internally to per-head shape and returns flat. Single Triton + # launch fuses per-head Q RMSNorm (weightless) + KV RMSNorm (both + # head_dim=128), replacing the prior `_rmsnorm_nw + kv_norm` pair. + q_flat, kv = self.qk_norm(q, kv_pre) + q = q_flat.view(num_tokens, self.n_local_heads, self.head_dim) + # q [S, H, D] / kv [S, head_dim] — rotary_emb internally unsqueezes + # to (1, num_tokens, ...) for aiter's per-position rope kernel. + self.rotary_emb(positions, q[..., -rd:], kv[..., -rd:]) + if _V4_USE_REF_QUANT: + act_quant_inplace(kv[..., :-rd], 64, self.scale_fmt) + + # ===== Compressor + Indexer ===== + if not use_async_compress: + if self.compressor is not None: + self.compressor( + x, + plan=plan_for_layer, + state_slot_mapping=state_slot_mapping, + block_tables=block_tables_gpu, + ) + if self.indexer is not None: + self.indexer.compressor( + x, + plan=plan_for_layer, + state_slot_mapping=state_slot_mapping, + block_tables=block_tables_gpu, + ) + if self.indexer is not None: + if use_async_compress: + current_stream.wait_stream(self.compress_stream) + current_stream.wait_stream(self.alt_stream) indexer_topk_batched = self.indexer.forward_batched( x_full=x, qr_full=qr, + qr_full_scale=qr_scale, positions=positions, - ) # raw seq-local [T, index_topk] int32, -1 sentinels in tail cols + ) # Translate seq-local topk → physical paged offsets and write into # the CSA section of either: # - decode buffer `kv_indices_csa` (is_pure_decode) # - prefill buffer `kv_indices_prefix_csa` (otherwise) # `_fill_csa_paged_compress` dispatches internally on is_pure_decode. self._fill_csa_paged_compress(attn_md, indexer_topk_batched, num_tokens) + elif use_async_compress: + current_stream.wait_stream(self.alt_stream) # ===== Sparse attention dispatch ===== # Two paths over the unified KV pool. The order of `swa_write` vs # `sparse_attn` differs because the two kernels read SWA differently: # # decode (is_pure_decode==True): - # `paged_decode` reads SWA from `unified_kv` ring slot `pos % win`. - # The current decode token's K must be present in the ring before - # attn fires, otherwise the token can't see its own K. So: + # `paged_decode` reads SWA from `unified_kv` ring slot + # `pos % cache_size`. The current decode token's K must be present + # in the ring before attn fires, otherwise the token can't see its + # own K. So: # swa_write → paged_decode # # prefill / mixed (is_pure_decode==False): # `paged_prefill` reads in-chunk K from per-fwd `kv` tensor (extend # region) and prior-chunk K from `unified_kv` ring (prefix region). - # `swa_write` writes the LAST `win` tokens of THIS fwd into ring - # slots `pos % win`, which OVERLAP with prior-chunk slots that the - # prefix SWA region wants to read (chunked prefill). So: + # `swa_write` writes the LAST `cache_size` tokens of THIS fwd into + # ring slots `pos % cache_size`, which OVERLAP with prior-chunk + # slots that the prefix SWA region wants to read (chunked prefill): # paged_prefill → swa_write # Pure prefill (chunk_start==0) has prefix_swa_count==0 so no prior # ring read; swa_write order is irrelevant in that subcase. q_sa = q.contiguous() if attn_md.is_pure_decode: - swa_write( - kv, - swa_write_indices, - positions, - v4_batch_id_per_token, - state_slot_mapping, - self.swa_kv, - win, - ) + if not self.use_fuse_qk_norm_rope_swa_write: + swa_write( + kv, + swa_write_indices, + positions, + v4_batch_id_per_token, + state_slot_mapping, + self.swa_kv, + cache_size, + ) if ratio == 0: kv_indices = attn_md.kv_indices_swa kv_indptr = attn_md.kv_indptr_swa @@ -1520,7 +1850,7 @@ def forward( ) # [S, H, head_dim] # swa_write AFTER attn so chunked-prefill prefix SWA reads see # prior-chunk's ring contents (current swa_write would overwrite - # ring slots `pos % win` for positions in this chunk's tail). + # ring slots `pos % cache_size` for positions in this chunk's tail). swa_write( kv, swa_write_indices, @@ -1528,7 +1858,7 @@ def forward( v4_batch_id_per_token, state_slot_mapping, self.swa_kv, - win, + cache_size, ) # Inverse RoPE on output's rope dims to remove absolute-position @@ -1639,24 +1969,44 @@ def __init__( prefix=f"{prefix}.w2", ) self.swiglu_limit = swiglu_limit + # Switch: route clamp + silu(gate)*up [+ weights] + per-token FP8 1x128 + # quant through a single aiter triton kernel. The fused kernel emits + # FP8 + scale; w2 accepts `x_scale` and skips its own quant step. + self.use_fused_clamp_act_mul = _V4_USE_TRITON_FUSION def forward( self, x: torch.Tensor, # [num_tokens, dim] weights: Optional[torch.Tensor] = None, # [num_tokens, 1] optional gate ) -> torch.Tensor: # [num_tokens, dim] + dtype = x.dtype - # Single fused GEMM, then chunk(2) along last dim so TP-sharded per-rank - # output (inter_dim_per_tp * 2) is split correctly regardless of tp_size. - combined = self.gate_up_proj(x).float() # [num_tokens, 2*inter_dim_per_tp] - gate, up = combined.chunk(2, dim=-1) # each [num_tokens, inter_dim_per_tp] - if self.swiglu_limit > 0: - up = torch.clamp(up, min=-self.swiglu_limit, max=self.swiglu_limit) - gate = torch.clamp(gate, max=self.swiglu_limit) - x = F.silu(gate) * up # [num_tokens, inter_dim_per_tp] + # Single fused GEMM. Layout is [gate | up] concat on last dim — matches + # aiter silu_and_mul's split([d, d], dim=-1) contract. The kernel does + # silu/clamp/mul in fp32 internally regardless of input dtype, so we + # feed the bf16 GEMM output directly. + combined = self.gate_up_proj(x) # [num_tokens, 2*inter_dim_per_tp] + if self.use_fused_clamp_act_mul: + x_fp8, x_scale = fused_clamp_act_mul( + combined, + swiglu_limit=self.swiglu_limit, + activation="silu", + weights=weights, + dtype_quant=dtypes.fp8, + transpose_scale=True, + ) + return self.w2(x_fp8, x_scale=x_scale) + out = torch.empty( + (combined.shape[0], combined.shape[-1] // 2), + dtype=dtype, + device=combined.device, + ) + # limit > 0 enables in-kernel clamp (gate≤limit, up∈[-limit,limit]) via + # ROCm v_med3_f32 — same semantics as the prior torch.clamp pair. + aiter_silu_and_mul(out, combined, self.swiglu_limit) if weights is not None: - x = weights * x - return self.w2(x.to(dtype)) # [num_tokens, dim] + out = weights.to(dtype) * out + return self.w2(out) # [num_tokens, dim] class MoE(nn.Module): @@ -1723,10 +2073,25 @@ def __init__( # (set by ModelRunner). torch.compile silently drops NNModule # attribute mutation across the compile boundary, so stashing on # `self.foo` from inside forward is a no-op at runtime. - + assert args.n_shared_experts == 1 + # aiter can fuse shared expert into the routed FusedMoE kernel ONLY + # when shared and routed have the same quant dtype. + # `is_rocm_aiter_fusion_shared_expert_enabled` compares shared vs + # GLOBAL (=base dtype), but V4-Pro has routed=FP4 (per-layer override + # for `.ffn.experts`) and shared=FP8 (=global). The function returns + # True for V4 because shared matches global, missing the routed-shared + # mismatch. Direct comparison: get routed and shared specs and compare. + routed_dtype = qc.get_layer_quant_config(f"{prefix}.experts").quant_dtype + shared_dtype = qc.get_layer_quant_config(f"{prefix}.shared_experts").quant_dtype + self._fuse_shared_into_routed = ( + routed_dtype == shared_dtype + and is_rocm_aiter_fusion_shared_expert_enabled() + ) moe_cfg = SimpleNamespace( routed_scaling_factor=self.routed_scaling_factor, - n_shared_experts=args.n_shared_experts, + n_shared_experts=( + args.n_shared_experts if self._fuse_shared_into_routed else 0 + ), ) self.experts = FusedMoE( num_experts=self.n_routed_experts, @@ -1743,21 +2108,9 @@ def __init__( config=moe_cfg, ) self.experts.swiglu_limit = args.swiglu_limit - assert args.n_shared_experts == 1 - # aiter can fuse shared expert into the routed FusedMoE kernel ONLY - # when shared and routed have the same quant dtype. - # `is_rocm_aiter_fusion_shared_expert_enabled` compares shared vs - # GLOBAL (=base dtype), but V4-Pro has routed=FP4 (per-layer override - # for `.ffn.experts`) and shared=FP8 (=global). The function returns - # True for V4 because shared matches global, missing the routed-shared - # mismatch. Direct comparison: get routed and shared specs and compare. - routed_dtype = qc.get_layer_quant_config(f"{prefix}.experts").quant_dtype - shared_dtype = qc.get_layer_quant_config(f"{prefix}.shared_experts").quant_dtype - self._fuse_shared_into_routed = ( - routed_dtype == shared_dtype - and is_rocm_aiter_fusion_shared_expert_enabled() - ) + if not self._fuse_shared_into_routed: + # self.experts.num_fused_shared_experts = 0 self.shared_experts = Expert( args.dim, args.moe_inter_dim, @@ -1813,6 +2166,23 @@ def _hash_topk( fwd_input_ids is not None ), "forward_context.context.input_ids is None — caller must invoke DeepseekV4ForCausalLM.forward, not DeepseekV4Model.forward directly." ids = fwd_input_ids.flatten() + # Under DP-attention gather/scatter, forward_impl_graph gathers + # hidden_states and gating_output across DP ranks before calling + # select_experts → _hash_topk. input_ids is still local (not + # gathered), so topk_ids would be [local_bs, topk] while the + # gathered gating_output is [local_bs*dp, n_experts]. Gather + # input_ids to match. + num_tokens = gating_output.shape[0] + if ids.shape[0] < num_tokens: + from atom.model_ops.moe import pad_for_all_gather + + ids_2d = ids.unsqueeze(-1) + ids_2d, _ = pad_for_all_gather(ids_2d) + from aiter.dist.parallel_state import get_dp_group + + ids_2d = get_dp_group().all_gather(ids_2d, dim=0) + ids = ids_2d[:num_tokens].flatten() + ids = ids.clamp(0, self.gate.tid2eid.shape[0] - 1) topk_ids = self.gate.tid2eid[ids].to(torch.int32) # [N, topk] scores = torch.nn.functional.softplus(gating_output.float()).sqrt() topk_weights = scores.gather(dim=-1, index=topk_ids.long()) @@ -1833,9 +2203,7 @@ def routed_expert_forward( `forward_context.context.input_ids` before each forward, and `_hash_topk` (FusedMoE's custom_routing_function) reads it there. """ - router_logits = self.gate( - x, otype=torch.float32 - ) # [num_tokens, n_routed_experts] + router_logits = self.gate(x) # [num_tokens, n_routed_experts] return self.experts(hidden_states=x, router_logits=router_logits) def combine_outputs( @@ -1868,11 +2236,11 @@ def dual_stream_moe_forward( independent; main stream waits on alt_stream's completion before combining. """ - current_stream = torch.cuda.current_stream() + current_stream = get_forward_context().main_stream self.alt_stream.wait_stream(current_stream) + routed = self.routed_expert_forward(x) with torch.cuda.stream(self.alt_stream): shared = self.shared_experts(x) - routed = self.routed_expert_forward(x) current_stream.wait_stream(self.alt_stream) return self.combine_outputs(routed, shared) @@ -1916,11 +2284,18 @@ def __init__( args: DeepseekV4Args, prefix: str = "", alt_stream: Optional[torch.cuda.Stream] = None, + compress_stream: Optional[torch.cuda.Stream] = None, ): super().__init__() self.layer_id = layer_id self.norm_eps = args.norm_eps - self.attn = DeepseekV4Attention(layer_id, args, prefix=f"{prefix}.attn") + self.attn = DeepseekV4Attention( + layer_id, + args, + prefix=f"{prefix}.attn", + alt_stream=alt_stream, + compress_stream=compress_stream, + ) self.ffn = MoE(layer_id, args, prefix=f"{prefix}.ffn", alt_stream=alt_stream) self.attn_norm = RMSNorm(args.dim, self.norm_eps) self.ffn_norm = RMSNorm(args.dim, self.norm_eps) @@ -2044,10 +2419,12 @@ def forward( ) -> torch.Tensor: # [num_tokens, hc, dim] updated residual stream # ----- Attention sub-layer with mHC mixing ----- residual = x # [num_tokens, hc, dim] - x, post, comb = ( - self.hc_pre( # [num_tokens, dim], [num_tokens, hc], [num_tokens, hc, hc] - x, self.hc_attn_fn, self.hc_attn_scale, self.hc_attn_base - ) + ( + x, + post, + comb, + ) = self.hc_pre( # [num_tokens, dim], [num_tokens, hc], [num_tokens, hc, hc] + x, self.hc_attn_fn, self.hc_attn_scale, self.hc_attn_base ) x = self.attn_norm(x) # [num_tokens, dim] x = self.attn(x, positions) # [num_tokens, dim] @@ -2065,58 +2442,43 @@ def forward( return x -class ParallelHead(nn.Module): - """LM head with mHC reduction. +class ParallelHead(ParallelLMHead): + """V4 LM head with mHC reduction; vocab-parallel sharded across TP ranks. + + Port of inference/model.py:704-736. Inherits from `ParallelLMHead` so the + vocab-axis sharding, `weight_loader`, last-token slicing, bf16 a16w16 GEMM + (`tgemm.mm`), and TP all-gather come for free. V4 only adds: - Port of inference/model.py:704-736. Unlike `Block.hc_pre` (which uses - Sinkhorn projection on the combination matrix), `hc_head` uses simple - `Sigmoid(mix*scale + base) + eps` weights to reduce the [num_tokens, hc, dim] - residual to [num_tokens, dim] before applying the LM head linear projection. + - `forward(...)` taking the mHC residual + hc_head params + final norm + - `get_logits(...)` so `compute_logits` can call it directly on the + hidden-state output of `model.forward` (CUDAGraph contract) + - `hc_head(...)` Sigmoid-gated mHC reduction (vs `Block.hc_pre`'s Sinkhorn) - `get_logits` projects only the last token of each sequence (decode mode); - for prefill the caller should slice the desired positions before passing - through. + Note on weight dtype: the V4 reference (model.py:713-714) keeps the LM head + in fp32 because the disk weight is bf16; on AMD CDNA3/CDNA4 the bf16 MFMA + instruction accumulates in fp32 natively, so a bf16 GEMM with the + bf16-on-disk weight has the same effective precision as the reference's + fp32 path while halving VRAM and using the faster a16w16 kernel. """ def __init__( self, vocab_size: int, dim: int, norm_eps: float = 1e-6, hc_eps: float = 1e-6 ): - super().__init__() - self.vocab_size = vocab_size + super().__init__(vocab_size, dim, bias=False) self.dim = dim self.norm_eps = norm_eps self.hc_eps = hc_eps - # PR1 single-rank: full vocab on this rank. - self.weight = atom_parameter( - torch.empty(self.vocab_size, self.dim, dtype=torch.float32) - ) def get_logits( self, x: torch.Tensor # [num_tokens, dim] ) -> torch.Tensor: # [bs, vocab] - """Project the last-token-per-seq slice of `x` to vocab logits. - - Picks the last token of each sequence using `cu_seqlens_q` from the - forward_context. Falls back to `x[-1:]` (single-seq) when no - forward_context is set (warmup / standalone). + """Project to vocab logits via the inherited `ParallelLMHead.forward`, + which handles last-token slicing (prefill) + tgemm.mm + all-gather. """ assert ( x.dim() == 2 and x.shape[-1] == self.dim ), f"get_logits expects [num_tokens, {self.dim}], got {tuple(x.shape)}" - ctx = get_forward_context() - cu_seqlens_q = ( - ctx.attn_metadata.cu_seqlens_q - if ctx is not None and ctx.attn_metadata is not None - else None - ) - if cu_seqlens_q is not None and cu_seqlens_q.numel() >= 2: - # Last-token positions per seq: cu_seqlens_q[1:bs+1] - 1. - # block_tables tells us the actual scheduled bs (cu_seqlens_q - # may have trailing zeros from the CpuGpuBuffer pool). - bs = ctx.attn_metadata.block_tables.size(0) - last_idx = cu_seqlens_q[1 : bs + 1] - 1 - return F.linear(x.index_select(0, last_idx.long()).float(), self.weight) - return F.linear(x[-1:].float(), self.weight) + return super().forward(x) def hc_head( self, @@ -2128,10 +2490,14 @@ def hc_head( """Reduce mHC residual `[num_tokens, hc, dim]` → `[num_tokens, dim]` via Sigmoid-gated weighted sum (vs Block.hc_pre's Sinkhorn variant). """ - _, _, y = aiter.mhc_pre( - x, hc_fn, hc_scale, hc_base, self.norm_eps, self.hc_eps, sinkhorn_repeat=0 + return _hc_head_reduce( + x, + hc_fn, + hc_scale, + hc_base, + self.norm_eps, + self.hc_eps, ) - return y def forward( self, @@ -2142,98 +2508,25 @@ def forward( norm: nn.Module, ) -> torch.Tensor: # [bs, vocab] x = self.hc_head(x, hc_fn, hc_scale, hc_base) # [num_tokens, dim] - logits = self.get_logits(norm(x)) # [bs, vocab] - # PR1 single-rank: skip all_gather - return logits - - -class MTPBlock(Block): - """MTP block: V4 dense block + e_proj/h_proj/enorm/hnorm + own hc_head params + LM head. - - Port of inference/model.py:739-767. Subclass of Block reusing all HC + Attention + FFN - machinery; adds a token-embed projection (`e_proj`), a hidden-state projection - (`h_proj`), per-input RMSNorms, and its own `hc_head_fn/base/scale` parameters - for the final LM head reduction. - - `embed` and `head` are assigned externally by `DeepseekV4Model` (shared with - the main model's embedding and LM head). - """ - - def __init__(self, layer_id: int, args: DeepseekV4Args, prefix: str = ""): - super().__init__(layer_id, args, prefix=prefix) - # e_proj / h_proj are FP8 on disk per index; ATOM Linear with V4QuantConfig - # picks per_1x128 automatically. nn.Linear at construction works for the - # toy/dummy path; for real-checkpoint loading, switch to ReplicatedLinear. - qc = args.quant_config - if qc is None: - self.e_proj = nn.Linear(args.dim, args.dim, bias=False) - self.h_proj = nn.Linear(args.dim, args.dim, bias=False) - else: - self.e_proj = ReplicatedLinear( - args.dim, - args.dim, - bias=False, - quant_config=qc, - prefix=f"{prefix}.e_proj", - ) - self.h_proj = ReplicatedLinear( - args.dim, - args.dim, - bias=False, - quant_config=qc, - prefix=f"{prefix}.h_proj", - ) - self.enorm = RMSNorm(args.dim, args.norm_eps) - self.hnorm = RMSNorm(args.dim, args.norm_eps) - self.norm = RMSNorm(args.dim, args.norm_eps) - # Per-MTP hc_head params (distinct from Block's hc_attn/hc_ffn params). - hc_mult = args.hc_mult - hc_dim = hc_mult * args.dim - self.hc_head_fn = atom_parameter( - torch.empty(hc_mult, hc_dim, dtype=torch.float32) - ) - self.hc_head_base = atom_parameter(torch.empty(hc_mult, dtype=torch.float32)) - self.hc_head_scale = atom_parameter(torch.empty(1, dtype=torch.float32)) - # Externally-assigned by DeepseekV4Model (shared with main model). - self.embed: Optional[nn.Module] = None - self.head: Optional[ParallelHead] = None - - def forward( - self, - x: torch.Tensor, # [num_tokens, hc, dim] residual stream from main model - positions: torch.Tensor, # [num_tokens] int absolute positions - input_ids: torch.Tensor, # [num_tokens] int - ) -> torch.Tensor: # [bs, vocab] last-token logits via self.head - assert ( - self.embed is not None and self.head is not None - ), "MTPBlock requires .embed and .head to be assigned by the parent model" - e = self.enorm(self.embed(input_ids)) # [num_tokens, dim] - x = self.hnorm(x) # [num_tokens, hc, dim] - # Mix token-embed + hidden into a fresh residual. e_proj output is - # [num_tokens, dim]; unsqueeze adds the hc axis so it broadcasts - # against h_proj(x) [num_tokens, hc, dim]. - x = self.e_proj(e).unsqueeze(-2) + self.h_proj(x) # [num_tokens, hc, dim] - x = super().forward(x, positions) # [num_tokens, hc, dim] - return self.head( - x, self.hc_head_fn, self.hc_head_scale, self.hc_head_base, self.norm - ) + # get_logits handles the per-rank vocab shard + all-gather internally. + return self.get_logits(norm(x)) # [bs, vocab] +@support_torch_compile class DeepseekV4Model(nn.Module): """Full model: embed -> expand to hc_mult copies -> N blocks -> hc_head -> logits. - Port of inference/model.py:Transformer (770-810). MTP blocks are constructed - and have their `.embed` and `.head` linked to the main model's, but they are - NOT called from the main forward path — PR5 will integrate them into ATOM's - EagleProposer via `self.mtp[k].forward(...)` from outside. - - PR1 single-rank: uses plain `nn.Embedding` for `self.embed` (state_dict-compatible - with reference's `ParallelEmbedding` since both store a single `weight` parameter). + Port of inference/model.py:Transformer (770-810). MTP blocks live in the + EagleProposer wrapper (`atom.models.deepseek_v4_mtp.DeepseekV4MTP`), not + on this target. The ckpt's `mtp.*` weights are filtered out at + `loader.load_model(spec_decode=False)` via the auto-detected + `need_load_mtp` flag (target has no `mtp.*` params). """ def __init__( self, *, + atom_config: Config, args: DeepseekV4Args, ): super().__init__() @@ -2247,12 +2540,16 @@ def __init__( # equals nn.Embedding's [vocab_size, dim] so dummy state_dicts load # directly. At TP>1 each rank holds vocab_size/tp rows. self.embed = VocabParallelEmbedding(args.vocab_size, args.dim) - # alt_stream for dual-stream MoE (shared_experts // routed_experts). - # Allocated once and shared across all blocks; MoE.__init__ decides - # per-layer whether to actually use it (env-gated, mismatched-dtype-aware). + # alt_stream: dual-stream MoE (shared_experts // routed_experts) AND + # Main Compressor overlap. compress_stream: Indexer Compressor overlap. + # Both allocated once, shared across all blocks. Attention runs before + # MoE in each block, so attn and MoE never contend for alt_stream. self.alt_stream: Optional[torch.cuda.Stream] = ( torch.cuda.Stream() if torch.cuda.is_available() else None ) + self.compress_stream: Optional[torch.cuda.Stream] = ( + torch.cuda.Stream() if torch.cuda.is_available() else None + ) self.layers = nn.ModuleList( [ Block( @@ -2260,6 +2557,7 @@ def __init__( args, prefix=f"layers.{layer_id}", alt_stream=self.alt_stream, + compress_stream=self.compress_stream, ) for layer_id in range(args.n_layers) ] @@ -2267,14 +2565,6 @@ def __init__( self.norm = RMSNorm(args.dim, self.norm_eps) self.head = ParallelHead(args.vocab_size, args.dim, self.norm_eps, self.hc_eps) - # MTP blocks: constructed and linked, but only invoked externally (PR5). - self.mtp = nn.ModuleList() - for layer_id in range(args.n_mtp_layers): - blk = MTPBlock(args.n_layers + layer_id, args, prefix=f"mtp.{layer_id}") - blk.embed = self.embed - blk.head = self.head - self.mtp.append(blk) - # Top-level hc_head params used to reduce the final hc_mult residual stack # before the LM head linear projection. hc_mult = args.hc_mult @@ -2285,19 +2575,17 @@ def __init__( self.hc_head_base = atom_parameter(torch.empty(hc_mult, dtype=torch.float32)) self.hc_head_scale = atom_parameter(torch.empty(1, dtype=torch.float32)) - @torch.inference_mode() def forward( self, input_ids: torch.Tensor, # [num_tokens] int flat ragged-batch token ids positions: torch.Tensor, # [num_tokens] int abs positions (required) - ) -> ( - torch.Tensor - ): # [num_tokens, dim] hidden_states (vocab head deferred to compute_logits) + ) -> torch.Tensor: # [num_tokens, hc, dim] pre-hc_head residual stream """Forward over `num_tokens` flat ragged-batch tokens. - Returns hidden_states `[num_tokens, dim]` — the vocab projection is - deferred to `compute_logits` to satisfy ATOM's CUDAGraph contract - (capture buffer is sized to hidden_size, not vocab). + Returns the mHC residual stack `[num_tokens, hc, dim]` BEFORE hc_head + reduction — `hc_head + RMSNorm + LM head` are all deferred to + `compute_logits`. Returning the hc-shaped residual lets the (future) + MTP draft consume it without re-expanding from a dim-reduced state. """ assert input_ids.dim() == 1, f"input_ids must be 1D, got {input_ids.shape}" h = self.embed(input_ids) # [num_tokens, dim] @@ -2307,12 +2595,7 @@ def forward( for layer in self.layers: h = layer(h, positions) # [num_tokens, hc, dim] - # Reduce the mHC residual stack (hc_head + final RMSNorm); leave the - # vocab projection to compute_logits. - x_hc = self.head.hc_head( # [num_tokens, dim] - h, self.hc_head_fn, self.hc_head_scale, self.hc_head_base - ) - return self.norm(x_hc) + return h class DeepseekV4ForCausalLM(nn.Module): @@ -2355,7 +2638,6 @@ class DeepseekV4ForCausalLM(nn.Module): "norm.weight": "model.norm.weight", "head.weight": "model.head.weight", "hc_head_": "model.hc_head_", - "mtp.": "model.mtp.", } ) weights_mapping = { @@ -2382,7 +2664,12 @@ def __init__(self, config: Config, prefix: str = "") -> None: # config lacks `quantization_config` (e.g. dummy / toy validation), # this still works — base spec is QuantType.No. self.args.quant_config = make_v4_quant_config(self.hf_config) - self.model = DeepseekV4Model(args=self.args) + self.model = DeepseekV4Model(atom_config=config, args=self.args) + # Tell ModelRunner to size the CG outputs buffer as + # [max_num_batched_tokens, hc_mult, hidden_size] instead of the + # default [max_num_batched_tokens, hidden_size]. forward returns + # the un-reduced mHC residual stack [N, hc, dim]. + self.extra_output_dims: tuple[int, ...] = (self.args.hc_mult,) def forward( self, @@ -2401,12 +2688,21 @@ def forward( def compute_logits( self, - hidden_states: torch.Tensor, # [num_tokens, dim] post hc_head + final norm + hidden_states: torch.Tensor, # [num_tokens, hc, dim] pre-hc_head residual ) -> torch.Tensor: # [bs, vocab] - # Vocab projection is split off from `model.forward` so the latter - # returns hidden_size-shaped tensors — required by ATOM's CUDAGraph - # capture contract (outputs buffer is sized to hidden_size, not vocab). - return self.model.head.get_logits(hidden_states) + # mHC reduce + final RMSNorm + LM head are all here so `model.forward` + # can return the un-reduced [N, hc, dim] residual stream — the future + # MTP draft consumes it directly without re-expanding from a dim-reduced + # state. CG output buffer is sized [N, hc, dim] in ModelRunner via the + # `extra_output_dims = (hc_mult,)` hook on this class. + x = self.model.head.hc_head( + hidden_states, + self.model.hc_head_fn, + self.model.hc_head_scale, + self.model.hc_head_base, + ) + x = self.model.norm(x) + return self.model.head.get_logits(x) def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: """Return (param_name, weight_name, expert_id, shard_id) tuples for FusedMoE. diff --git a/atom/models/deepseek_v4_mtp.py b/atom/models/deepseek_v4_mtp.py new file mode 100644 index 0000000000..79eaadbb0c --- /dev/null +++ b/atom/models/deepseek_v4_mtp.py @@ -0,0 +1,267 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""DeepSeek-V4 MTP wrapper for ATOM's EagleProposer. + +Mirrors the V2/V3 + Qwen MTP convention: + +1. The target (`DeepseekV4ForCausalLM`) constructs and loads ONLY the main + stack — no MTP modules, no MTP weights. + +2. This wrapper owns the MTP block(s) and loads `mtp.{i}.*` ckpt entries + into them via the standard `load_model` path (with `spec_decode=True`). + `weights_mapper`, `packed_modules_mapping`, `get_expert_mapping`, and + `remap_mtp_weight_name` are declared so the loader recognizes V4 MTP + ckpt naming + filters out target-only entries silently. + +3. `EagleProposer.load_model` then calls `share_with_target(...)` which + rebinds `self.model.{embed,head}` to point at the target's already-loaded + instances, and propagates them onto each MTPBlock (which requires + `embed`/`head` set externally). No second weight load, no double KV. + +4. V4's MTP block consumes the **un-reduced mHC residual stack** `[N, hc, dim]` + from the target — enabled by deferring `hc_head + RMSNorm + LM head` from + `DeepseekV4Model.forward` to `DeepseekV4ForCausalLM.compute_logits`. The + wrapper's `forward` contract matches ATOM's EagleProposer: + - input `hidden_states` is the target's `[N, hc, dim]` residual + - output is `[N, dim]` post-(MTP block + its own hc_head + norm) +""" + +from typing import Optional + +import torch +from torch import nn + +from atom.config import Config +from atom.model_loader.loader import WeightsMapper +from atom.model_ops.layernorm import RMSNorm +from atom.model_ops.linear import ReplicatedLinear +from atom.model_ops.moe import FusedMoE +from atom.model_ops.utils import atom_parameter +from atom.utils.forward_context import get_forward_context + +from .deepseek_v4 import ( + Block, + DeepseekV4Args, + ParallelHead, + make_v4_quant_config, +) + + +class MTPBlock(Block): + """MTP block: V4 dense block + e_proj/h_proj/enorm/hnorm + own hc_head params + LM head. + + Port of inference/model.py:739-767. Subclass of Block reusing all HC + Attention + FFN + machinery; adds a token-embed projection (`e_proj`), a hidden-state projection + (`h_proj`), per-input RMSNorms, and its own `hc_head_fn/base/scale` parameters + for the final LM head reduction. + + `embed` and `head` are assigned externally by the wrapper (shared with + the target's embedding and LM head via `share_with_target`). + """ + + def __init__(self, layer_id: int, args: DeepseekV4Args, prefix: str = ""): + super().__init__(layer_id, args, prefix=prefix) + # e_proj / h_proj are FP8 on disk per index; ATOM Linear with V4QuantConfig + # picks per_1x128 automatically. nn.Linear at construction works for the + # toy/dummy path; for real-checkpoint loading, switch to ReplicatedLinear. + qc = args.quant_config + if qc is None: + self.e_proj = nn.Linear(args.dim, args.dim, bias=False) + self.h_proj = nn.Linear(args.dim, args.dim, bias=False) + else: + self.e_proj = ReplicatedLinear( + args.dim, + args.dim, + bias=False, + quant_config=qc, + prefix=f"{prefix}.e_proj", + ) + self.h_proj = ReplicatedLinear( + args.dim, + args.dim, + bias=False, + quant_config=qc, + prefix=f"{prefix}.h_proj", + ) + self.enorm = RMSNorm(args.dim, args.norm_eps) + self.hnorm = RMSNorm(args.dim, args.norm_eps) + self.norm = RMSNorm(args.dim, args.norm_eps) + # Per-MTP hc_head params (distinct from Block's hc_attn/hc_ffn params). + hc_mult = args.hc_mult + hc_dim = hc_mult * args.dim + self.hc_head_fn = atom_parameter( + torch.empty(hc_mult, hc_dim, dtype=torch.float32) + ) + self.hc_head_base = atom_parameter(torch.empty(hc_mult, dtype=torch.float32)) + self.hc_head_scale = atom_parameter(torch.empty(1, dtype=torch.float32)) + # Externally-assigned by the wrapper (shared with the target). + self.embed: Optional[nn.Module] = None + self.head: Optional[ParallelHead] = None + + def forward( + self, + x: torch.Tensor, # [num_tokens, hc, dim] residual stream from main model + positions: torch.Tensor, # [num_tokens] int absolute positions + input_ids: torch.Tensor, # [num_tokens] int + ) -> torch.Tensor: # [num_tokens, hc, dim] pre-hc_head residual + """Run one MTP step. Returns the un-reduced mHC residual stack + `[num_tokens, hc, dim]` — same shape contract as + `DeepseekV4Model.forward`. The hc_head reduction + RMSNorm + LM head + are all deferred to `DeepseekV4MTP.compute_logits` so the wrapper's + forward output can be fed back in as `x` for the NEXT MTP draft + step (mtp_k > 1) without re-expanding from a `[N, dim]` post-reduction + state. + """ + assert ( + self.embed is not None + ), "MTPBlock requires .embed to be assigned by the wrapper" + e = self.enorm(self.embed(input_ids)) # [num_tokens, dim] + x = self.hnorm(x) # [num_tokens, hc, dim] + # Mix token-embed + hidden into a fresh residual. h_proj is FP8 + # (V4QuantConfig → ReplicatedLinear over `gemm_a8w8_blockscale_preshuffle`), + # which only supports 2D `[M, K]` input. Flatten the hc axis into the + # batch dim and reshape back after the projection. e_proj's input is + # already 2D [num_tokens, dim]; unsqueeze adds the hc axis so it + # broadcasts over h_proj_out's hc dim. + n_tok, hc, d = x.shape + h_proj_out = self.h_proj(x.reshape(n_tok * hc, d)).reshape(n_tok, hc, d) + x = self.e_proj(e).unsqueeze(-2) + h_proj_out # [num_tokens, hc, dim] + return super().forward(x, positions) # [num_tokens, hc, dim] + + +class DeepseekV4MTPModel(nn.Module): + """V4 MTP inner model: owns the MTP blocks. Each block's `embed` / `head` + are set externally by `DeepseekV4MTP.share_with_target` to point at the + already-loaded target instances; this module itself holds no embed/head. + """ + + def __init__(self, atom_config: Config, args: DeepseekV4Args) -> None: + super().__init__() + # ModelRunner reads `drafter.model.model.mtp_start_layer_idx` to bind + # the draft attention to KV slots `[n_layers, n_layers + n_mtp)`. + self.mtp_start_layer_idx = args.n_layers + # Real MTP blocks loaded via the standard load_model path with + # spec_decode=True. Each block's `.embed` / `.head` are set by + # share_with_target after load. + self.mtp = nn.ModuleList( + [ + MTPBlock(args.n_layers + i, args, prefix=f"mtp.{i}") + for i in range(args.n_mtp_layers) + ] + ) + + def forward( + self, + input_ids: torch.Tensor, # [num_tokens] int + positions: torch.Tensor, # [num_tokens] int + hidden_states: torch.Tensor, # [num_tokens, hc, dim] pre-hc_head from target + spec_step_idx: int = 0, + ) -> torch.Tensor: # [num_tokens, hc, dim] pre-hc_head residual + """Returns the un-reduced mHC residual `[N, hc, dim]` from the + selected MTP block — same shape contract as `DeepseekV4Model.forward`. + EagleProposer feeds this back in as `hidden_states` for the NEXT + draft step (mtp_k > 1).""" + idx = spec_step_idx % len(self.mtp) + return self.mtp[idx](hidden_states, positions, input_ids) + + def compute_logits( + self, + hidden_states: torch.Tensor, # [num_tokens, hc, dim] pre-hc_head residual + spec_step_idx: int = 0, + ) -> torch.Tensor: # [bs, vocab] + """Mirror `DeepseekV4ForCausalLM.compute_logits`: hc_head + RMSNorm + + LM head all here so MTPBlock.forward can return the un-reduced + residual stack. Each MTP block has its own `hc_head_fn/base/scale` + and own `norm` (the LM head is target-shared via `share_with_target`).""" + idx = spec_step_idx % len(self.mtp) + blk = self.mtp[idx] + x = blk.head.hc_head( + hidden_states, blk.hc_head_fn, blk.hc_head_scale, blk.hc_head_base + ) # [num_tokens, dim] + x = blk.norm(x) + return blk.head.get_logits(x) + + +class DeepseekV4MTP(nn.Module): + """Top-level V4 MTP wrapper. Owns its own MTP weights (loaded via the + standard `load_model` path with `spec_decode=True`) and shares only + `embed` / `head` with the target through `share_with_target`. + """ + + # Disk `mtp.{i}.*` -> wrapper param `model.mtp.{i}.*`. Prefix-anchored to + # avoid the `mtp.` substring colliding with anything inside MTPBlock subtree. + weights_mapper = WeightsMapper(orig_to_new_prefix={"mtp.": "model.mtp."}) + # Same on-disk -> internal name conventions as the target (V4 ckpt quirks). + weights_mapping = { + ".gate.bias": ".gate.e_score_correction_bias", + ".scale": ".weight_scale_inv", + } + # Same packed-module fusions as the target — MTPBlock subclasses Block, so + # the same attention / shared-experts fused param layouts apply. + packed_modules_mapping = { + "attn.wq_a": ("attn.wqkv_a", 0), + "attn.wkv": ("attn.wqkv_a", 1), + "compressor.wkv": ("compressor.wkv_gate", 0), + "compressor.wgate": ("compressor.wkv_gate", 1), + "shared_experts.w1": ("shared_experts.gate_up_proj", 0), + "shared_experts.w3": ("shared_experts.gate_up_proj", 1), + } + + def __init__(self, config: Config, prefix: str = "") -> None: + super().__init__() + self.atom_config = config + self.hf_config = config.hf_config + self.args = DeepseekV4Args.from_hf_config(self.hf_config) + self.args.quant_config = make_v4_quant_config(self.hf_config) + self.model = DeepseekV4MTPModel(atom_config=config, args=self.args) + + def remap_mtp_weight_name(self, name: str) -> str | None: + """Filter loader input to MTP-only weights. + + Called per ckpt entry AFTER `weights_mapper` rewrites `mtp.` -> + `model.mtp.`. Target-only entries (`embed.weight`, `layers.X.*`, + `head.weight`, `hc_head_*`, `norm.weight`) pass through unchanged + and have no matching wrapper param, which would otherwise generate + loud `dropped_ckpt_keys` warnings. Returning None drops them silently. + """ + return name if "mtp." in name else None + + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: + """FusedMoE expert param mapping for MTPBlock's MoE layer. Same + ckpt convention as the target (`ffn.experts.{e}.w{1,2,3}`). + """ + return FusedMoE.make_expert_params_mapping( + ckpt_gate_proj_name="w1", + ckpt_down_proj_name="w2", + ckpt_up_proj_name="w3", + num_experts=self.args.n_routed_experts, + ) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + hidden_states: torch.Tensor, + spec_step_idx: int = 0, + ) -> torch.Tensor: + # Hash MoE routing inside MTPBlock's MoE looks at this — same as + # `DeepseekV4ForCausalLM.forward`. + get_forward_context().context.input_ids = input_ids + return self.model(input_ids, positions, hidden_states, spec_step_idx) + + def compute_logits( + self, + hidden_states: torch.Tensor, + spec_step_idx: int = 0, + ) -> Optional[torch.Tensor]: + return self.model.compute_logits(hidden_states, spec_step_idx) + + def share_with_target(self, target_base: nn.Module, loaded: set[str]) -> None: + """Bind embed/head on each MTPBlock to the already-loaded target's + instances. MTPBlock requires `.embed` / `.head` set externally before + its first forward; `compute_logits` then reaches them via `mtp[0].head`. + """ + for blk in self.model.mtp: + blk.embed = target_base.model.embed + blk.head = target_base.model.head diff --git a/atom/models/eagle3_llama.py b/atom/models/eagle3_llama.py new file mode 100644 index 0000000000..67837907b4 --- /dev/null +++ b/atom/models/eagle3_llama.py @@ -0,0 +1,301 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2025, Advanced Micro Devices, Inc. All rights reserved. + +"""Eagle3 draft model (Llama full-attention) for speculative decoding. + +Implements the Eagle3 draft model matching the lightseekorg/kimi-k2.5-eagle3 +checkpoint layout: + + embed_tokens.weight — independent embedding + fc.weight — aux fusion projection (hidden*3 -> hidden) + midlayer.* — single decoder layer (dual-norm, wide QKV) + norm.weight — final RMSNorm + lm_head.weight — independent lm_head + +Weight keys map directly to model attribute paths; no key rewriting needed. +""" + +import torch +from aiter.dist.parallel_state import get_tensor_model_parallel_world_size +from aiter.rotary_embedding import get_rope +from atom.config import Config +from atom.model_ops.activation import SiluAndMul +from atom.model_ops.base_attention import Attention +from atom.model_ops.embed_head import ParallelLMHead, VocabParallelEmbedding +from atom.model_ops.layernorm import RMSNorm +from atom.model_ops.linear import ( + MergedColumnParallelLinear, + QKVParallelLinear, + ReplicatedLinear, + RowParallelLinear, +) +from atom.utils.decorators import support_torch_compile +from torch import nn + + +class Eagle3LlamaAttention(nn.Module): + """Llama full-attention with input_size = hidden_size * 2. + + The QKV projection accepts the concatenation of normalized embeddings + and fc output, hence input_size is doubled compared to standard Llama. + """ + + def __init__( + self, + config, + hidden_size: int, + num_heads: int, + num_kv_heads: int, + cache_config: str = "bf16", + prefix: str = "", + layer_num: int = 0, + ) -> None: + super().__init__() + self.hidden_size = hidden_size + tp_size = get_tensor_model_parallel_world_size() + self.total_num_heads = num_heads + assert self.total_num_heads % tp_size == 0 + self.num_heads = self.total_num_heads // tp_size + self.total_num_kv_heads = num_kv_heads + if self.total_num_kv_heads >= tp_size: + assert self.total_num_kv_heads % tp_size == 0 + else: + assert tp_size % self.total_num_kv_heads == 0 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) + head_dim = getattr(config, "head_dim", None) + if head_dim is None: + head_dim = hidden_size // self.total_num_heads + self.head_dim = head_dim + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + self.scaling = self.head_dim**-0.5 + + # QKV input_size = hidden_size * 2 (concat of embed + fc_output) + attn_input_size = hidden_size * 2 + self.qkv_proj = QKVParallelLinear( + hidden_size=attn_input_size, + head_size=self.head_dim, + total_num_heads=self.total_num_heads, + total_num_kv_heads=self.total_num_kv_heads, + bias=False, + prefix=f"{prefix}.qkv_proj", + ) + + self.o_proj = RowParallelLinear( + input_size=self.total_num_heads * self.head_dim, + output_size=hidden_size, + bias=False, + prefix=f"{prefix}.o_proj", + ) + + rope_theta = getattr(config, "rope_theta", 10000) + max_position_embeddings = getattr(config, "max_position_embeddings", 8192) + self.rotary_emb = get_rope( + self.head_dim, + rotary_dim=self.head_dim, + max_position=max_position_embeddings, + base=rope_theta, + is_neox_style=True, + ) + + sliding_window = -1 + if getattr(config, "use_sliding_window", False) and getattr( + config, "sliding_window", None + ): + sliding_window = config.sliding_window + self.attn = Attention( + self.num_heads, + self.head_dim, + self.scaling, + num_kv_heads=self.num_kv_heads, + kv_cache_dtype=cache_config, + layer_num=layer_num, + prefix=f"{prefix}.attn", + rotary_emb=self.rotary_emb, + per_layer_sliding_window=sliding_window, + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + ) -> torch.Tensor: + qkv = self.qkv_proj(hidden_states) + q, k, v = torch.split(qkv, [self.q_size, self.kv_size, self.kv_size], dim=-1) + attn_output = self.attn(q, k, v, positions) + output = self.o_proj(attn_output) + return output + + +class Eagle3LlamaDecoderLayer(nn.Module): + """Single decoder layer for Eagle3 with dual-norm input. + + Unlike standard LlamaDecoderLayer, this layer has: + - input_layernorm: normalizes the embedding input + - hidden_norm: normalizes the fc output (projected aux hidden states) + - Attention input is concat(normed_embed, normed_hidden) -> [N, hidden*2] + """ + + def __init__( + self, + config, + cache_config: str = "bf16", + prefix: str = "", + layer_num: int = 0, + ) -> None: + super().__init__() + self.hidden_size = config.hidden_size + + self.self_attn = Eagle3LlamaAttention( + config=config, + hidden_size=self.hidden_size, + num_heads=config.num_attention_heads, + num_kv_heads=getattr( + config, "num_key_value_heads", config.num_attention_heads + ), + cache_config=cache_config, + prefix=f"{prefix}.self_attn", + layer_num=layer_num, + ) + + self.mlp = Eagle3LlamaMLP( + hidden_size=self.hidden_size, + intermediate_size=config.intermediate_size, + prefix=f"{prefix}.mlp", + ) + + # Dual norms matching checkpoint keys: midlayer.input_layernorm, midlayer.hidden_norm + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.hidden_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + + def forward( + self, + positions: torch.Tensor, + embeds: torch.Tensor, + hidden_states: torch.Tensor, + ) -> torch.Tensor: + normed_embeds = self.input_layernorm(embeds) + normed_hidden = self.hidden_norm(hidden_states) + # Concat for attention input: [N, hidden*2] + attn_input = torch.cat([normed_embeds, normed_hidden], dim=-1) + attn_output = self.self_attn(positions, attn_input) + # Residual connection on hidden_states + hidden_states = hidden_states + attn_output + # MLP with pre-norm + residual + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + return hidden_states + + +class Eagle3LlamaMLP(nn.Module): + """Simple Llama MLP (gate+up fused, silu activation, down projection).""" + + def __init__( + self, + hidden_size: int, + intermediate_size: int, + prefix: str = "", + ) -> None: + super().__init__() + self.gate_up_proj = MergedColumnParallelLinear( + input_size=hidden_size, + output_sizes=[intermediate_size] * 2, + bias=False, + prefix=f"{prefix}.gate_up_proj", + ) + self.down_proj = RowParallelLinear( + input_size=intermediate_size, + output_size=hidden_size, + bias=False, + prefix=f"{prefix}.down_proj", + ) + self.act_fn = SiluAndMul() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + gate_up = self.gate_up_proj(x) + x = self.act_fn(gate_up) + x = self.down_proj(x) + return x + + +@support_torch_compile +class Eagle3LlamaModel(nn.Module): + """Eagle3 draft model (Llama full-attention, single decoder layer). + + Matches the lightseekorg/kimi-k2.5-eagle3 checkpoint layout: + embed_tokens.weight [163840, 7168] independent embedding + fc.weight [7168, 21504] aux fusion (hidden*3 -> hidden) + midlayer.* single decoder layer + norm.weight final RMSNorm + lm_head.weight [163840, 7168] independent lm_head + """ + + packed_modules_mapping = { + "q_proj": ("qkv_proj", "q"), + "k_proj": ("qkv_proj", "k"), + "v_proj": ("qkv_proj", "v"), + "gate_proj": ("gate_up_proj", 0), + "up_proj": ("gate_up_proj", 1), + } + + def __init__(self, atom_config: Config, prefix: str = "", layer_offset: int = 0): + super().__init__() + config = atom_config.hf_config + cache_config = atom_config.kv_cache_dtype + self.config = config + + # Independent embedding (vocab matches target model) + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, config.hidden_size + ) + + # Aux fusion: concatenated aux hidden states [N, hidden*3] -> [N, hidden] + self.fc = ReplicatedLinear( + config.hidden_size * 3, config.hidden_size, bias=False + ) + + # Draft attention layer_num must start from the target model's layer + # count so kv_cache_data["layer_N"] maps to the correct cache entry. + self.midlayer = Eagle3LlamaDecoderLayer( + config=config, + cache_config=cache_config, + prefix="midlayer", + layer_num=layer_offset, + ) + + # Final norm + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + # Independent lm_head (not shared with target model) + self.lm_head = ParallelLMHead(config.vocab_size, config.hidden_size) + + def combine_hidden_states(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Project concatenated aux hidden states through fc. + + Args: + hidden_states: [N, hidden_size * 3] (3 aux layers concatenated) + + Returns: + [N, hidden_size] projected hidden states + """ + return self.fc(hidden_states) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + hidden_states: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + embeds = self.embed_tokens(input_ids) + hidden_states = self.midlayer(positions, embeds, hidden_states) + hidden_states_prenorm = hidden_states + hidden_states = self.norm(hidden_states) + return hidden_states, hidden_states_prenorm + + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.lm_head(hidden_states) diff --git a/atom/models/glm4_moe_mtp.py b/atom/models/glm4_moe_mtp.py new file mode 100644 index 0000000000..b546b28323 --- /dev/null +++ b/atom/models/glm4_moe_mtp.py @@ -0,0 +1,205 @@ +from typing import Optional + +import torch +import torch.nn as nn +from atom.config import Config, QuantizationConfig +from atom.model_ops.embed_head import ParallelLMHead, VocabParallelEmbedding +from atom.model_ops.layernorm import RMSNorm +from atom.model_ops.moe import FusedMoE +from atom.model_ops.topK import is_rocm_aiter_fusion_shared_expert_enabled +from atom.models.utils import IntermediateTensors +from atom.utils.decorators import support_torch_compile +from transformers import PretrainedConfig + +from .deepseek_mtp import rewrite_spec_layer_name + +from .glm4_moe import Glm4MoeDecoderLayer, get_spec_layer_idx_from_weight_name +from .utils import maybe_prefix + + +class SharedHead(nn.Module): + def __init__( + self, + config: PretrainedConfig, + prefix: str, + quant_config: Optional[QuantizationConfig] = None, + ) -> None: + super().__init__() + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "head"), + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.norm(hidden_states) + + +class Glm4MoeMultiTokenPredictorLayer(nn.Module): + def __init__(self, atom_config: Config, prefix: str) -> None: + super().__init__() + + config = atom_config.hf_config + self.config = config + + self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.eh_proj = nn.Linear(config.hidden_size * 2, config.hidden_size, bias=False) + + self.shared_head = SharedHead( + config=config, prefix=prefix, quant_config=atom_config.quant_config + ) + + self.mtp_block = Glm4MoeDecoderLayer( + config=config, + atom_config=atom_config, + prefix=prefix, + ) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + previous_hidden_states: torch.Tensor, + inputs_embeds: torch.Tensor, + spec_step_index: int = 0, + ) -> torch.Tensor: + assert inputs_embeds is not None + masked_inputs_embeds = inputs_embeds + inputs_embeds = self.enorm(masked_inputs_embeds) + previous_hidden_states = self.hnorm(previous_hidden_states) + + hidden_states = self.eh_proj( + torch.cat([inputs_embeds, previous_hidden_states], dim=-1) + ) + + hidden_states, residual = self.mtp_block( + positions=positions, hidden_states=hidden_states, residual=None + ) + hidden_states = residual + hidden_states + return hidden_states + + +class Glm4MoeMultiTokenPredictor(nn.Module): + def __init__(self, *, atom_config: Config, prefix: str = ""): + super().__init__() + config = atom_config.hf_config + self.mtp_start_layer_idx = config.num_hidden_layers + self.num_mtp_layers = config.num_nextn_predict_layers + # to map the exact layer index from weights + self.layers = torch.nn.ModuleDict( + { + str(idx): Glm4MoeMultiTokenPredictorLayer( + atom_config, f"{prefix}.layers.{idx}" + ) + for idx in range( + self.mtp_start_layer_idx, + self.mtp_start_layer_idx + self.num_mtp_layers, + ) + } + ) + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + ) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_tokens(input_ids) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + previous_hidden_states: torch.Tensor, + inputs_embeds: torch.Tensor | None = None, + spec_step_idx: int = 0, + ) -> torch.Tensor: + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + current_step_idx = spec_step_idx % self.num_mtp_layers + return self.layers[str(self.mtp_start_layer_idx + current_step_idx)]( + input_ids, + positions, + previous_hidden_states, + inputs_embeds, + current_step_idx, + ) + + def compute_logits( + self, + hidden_states: torch.Tensor, + spec_step_idx: int = 0, + ) -> torch.Tensor: + current_step_idx = spec_step_idx % self.num_mtp_layers + mtp_layer = self.layers[str(self.mtp_start_layer_idx + current_step_idx)] + logits = mtp_layer.shared_head.head(mtp_layer.shared_head(hidden_states)) + return logits + + +@support_torch_compile +class Glm4MoeMTP(nn.Module): + # ATOM format: checkpoint_weight_name -> (model_param_name, shard_id) + packed_modules_mapping = { + "q_proj": ("qkv_proj", "q"), + "k_proj": ("qkv_proj", "k"), + "v_proj": ("qkv_proj", "v"), + "gate_proj": ("gate_up_proj", 0), + "up_proj": ("gate_up_proj", 1), + } + + def __init__(self, atom_config: Config, prefix: str = ""): + super().__init__() + self.config = atom_config.hf_config + + self.model = Glm4MoeMultiTokenPredictor( + atom_config=atom_config, prefix=maybe_prefix(prefix, "model") + ) + + def remap_mtp_weight_name(self, name: str) -> str | None: + # GLM-4 MoE MTP shares the rewrite rules with DeepSeek MTP: + # - shared scalars (embed_tokens) → top-level model.* + # - per-layer scalars (enorm/hnorm/eh_proj/shared_head) → kept verbatim + # - decoder block weights (self_attn/mlp/...) → model.layers.{i}.mtp_block.* + spec_layer = get_spec_layer_idx_from_weight_name(self.config, name) + if spec_layer is None: + return None + return rewrite_spec_layer_name(spec_layer, name) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + hidden_states: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + spec_step_idx: int = 0, + ) -> torch.Tensor: + hidden_states = self.model( + input_ids, positions, hidden_states, inputs_embeds, spec_step_idx + ) + return hidden_states + + def compute_logits( + self, + hidden_states: torch.Tensor, + spec_step_idx: int = 0, + ) -> torch.Tensor | None: + return self.model.compute_logits(hidden_states, spec_step_idx) + + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: + # Params for weights, fp8 weight scales, fp8 activation scales + # (param_name, weight_name, expert_id, shard_id) + num_experts = self.config.n_routed_experts + if ( + is_rocm_aiter_fusion_shared_expert_enabled() + and self.config.n_shared_experts + ): + num_experts += self.config.n_shared_experts + return FusedMoE.make_expert_params_mapping( + ckpt_gate_proj_name="gate_proj", + ckpt_down_proj_name="down_proj", + ckpt_up_proj_name="up_proj", + num_experts=num_experts, + ) diff --git a/atom/models/gpt_oss.py b/atom/models/gpt_oss.py index eca04cb62f..66b2b22b0f 100644 --- a/atom/models/gpt_oss.py +++ b/atom/models/gpt_oss.py @@ -153,6 +153,37 @@ def forward( return output +def _interleave_swiglu_weights(experts: FusedMoE): + """Interleave gate/up weights, scales, and biases for Swiglu activation. + + Must run before Mxfp4MoEMethod.process_weights_after_loading (shuffle). + The loader calls module.process_weights_after_loading() before + quant_method.process_weights_after_loading(module), so this ordering + is guaranteed. + """ + e, n, k = experts.w13_weight.shape + experts.w13_weight.view(torch.uint8).copy_( + experts.w13_weight.data.view(torch.uint8) + .view(e, n // 2, 2, k) + .permute(0, 2, 1, 3) + .contiguous() + .view(e, n, k) + ) + experts.w13_weight_scale.data = ( + experts.w13_weight_scale.data.view(e, n // 2, 2, -1) + .permute(0, 2, 1, 3) + .contiguous() + .view(e, n, -1) + ) + if experts.w13_bias is not None: + experts.w13_bias.data = ( + experts.w13_bias.data.view(-1, n // 2, 2) + .permute(0, 2, 1) + .contiguous() + .view(-1, n) + ) + + class MLPBlock(torch.nn.Module): def __init__( self, @@ -202,6 +233,9 @@ def __init__( else: self.moe_hidden_pad = 0 + def process_weights_after_loading(self): + _interleave_swiglu_weights(self.experts) + def forward(self, x: torch.Tensor) -> torch.Tensor: num_tokens = x.shape[0] diff --git a/atom/models/kimi_k25.py b/atom/models/kimi_k25.py index eb4aaafb61..06c3b1412f 100644 --- a/atom/models/kimi_k25.py +++ b/atom/models/kimi_k25.py @@ -40,6 +40,10 @@ class KimiK25ForCausalLM(nn.Module): "vision_tower.", "mm_projector.", ] + quant_exclude_name_mapping = { + "language_model.model.": "model.", + "language_model.lm_head": "lm_head", + } def __init__( self, @@ -88,5 +92,11 @@ def compute_logits( ) -> Optional[torch.Tensor]: return self.language_model.compute_logits(hidden_states) + def set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None: + self.language_model.set_aux_hidden_state_layers(layers) + + def get_eagle3_aux_hidden_state_layers(self) -> tuple[int, ...]: + return self.language_model.get_eagle3_aux_hidden_state_layers() + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: return self.language_model.get_expert_mapping() diff --git a/atom/models/mimo_v2_flash_mtp.py b/atom/models/mimo_v2_flash_mtp.py index 1da461eca5..fef4b7a8f5 100644 --- a/atom/models/mimo_v2_flash_mtp.py +++ b/atom/models/mimo_v2_flash_mtp.py @@ -10,7 +10,8 @@ from atom.config import Config from atom.model_ops.embed_head import ParallelLMHead, VocabParallelEmbedding from atom.model_ops.layernorm import RMSNorm -from atom.models.utils import IntermediateTensors, maybe_prefix +from atom.model_ops.linear import ReplicatedLinear +from atom.models.utils import IntermediateTensors, ckpt_has_tensor_suffix, maybe_prefix from atom.utils.decorators import support_torch_compile from .mimo_v2_flash import MiMoV2Attention, MiMoV2MLP @@ -113,7 +114,13 @@ def __init__(self, atom_config: Config, prefix: str, layer_idx: int) -> None: self.enorm = RMSNorm(config.hidden_size, eps=config.layernorm_epsilon) self.hnorm = RMSNorm(config.hidden_size, eps=config.layernorm_epsilon) - self.eh_proj = nn.Linear(config.hidden_size * 2, config.hidden_size, bias=False) + self.eh_proj = ReplicatedLinear( + config.hidden_size * 2, + config.hidden_size, + bias=False, + quant_config=atom_config.quant_config, + prefix=maybe_prefix(prefix, "eh_proj"), + ) self.mtp_block = MiMoV2FlashMTPLayer( atom_config=atom_config, @@ -213,6 +220,40 @@ def __init__(self, atom_config: Config, prefix: str = ""): super().__init__() self.config = atom_config.hf_config self._draft_hf_config = atom_config.speculative_config.draft_model_hf_config + num_spec = atom_config.speculative_config.num_speculative_tokens + assert num_spec == 1, ( + f"MiMo-V2-Flash MTP only supports --num-speculative-tokens=1 now " + f"(got {num_spec})." + ) + + # See deepseek_mtp.DeepSeekMTP for the full rationale: MTP eh_proj is + # commonly stored as BF16 with no weight_scale even when the model's + # global quant_config is FP8/MXFP4. Skip quantization for eh_proj only + # when the checkpoint actually has no scale tensor for it. + if atom_config.quant_config is not None and not ckpt_has_tensor_suffix( + atom_config.model, "eh_proj.weight_scale" + ): + atom_config.quant_config.apply_default_exclude_layers(["*.eh_proj"]) + # MiMo additionally keeps every self_attn.o_proj as BF16. Its HF + # `ignored_layers` covers all 48 base-model o_proj entries but + # forgets the MTP-layer ones (model.mtp.layers.0..2.self_attn.o_proj). + # Without this exclude, MTP o_proj falls back to the global FP8 spec, + # weight_scale stays at torch.empty junk memory, and accept rate + # collapses — same corruption chain as the eh_proj case above. + # Detect from disk: if the ckpt has no o_proj.weight_scale[_inv] for + # any layer, exclude *.self_attn.o_proj globally (HF entries dedup). + if ( + atom_config.quant_config is not None + and not ckpt_has_tensor_suffix( + atom_config.model, "self_attn.o_proj.weight_scale_inv" + ) + and not ckpt_has_tensor_suffix( + atom_config.model, "self_attn.o_proj.weight_scale" + ) + ): + atom_config.quant_config.apply_default_exclude_layers( + ["*.self_attn.o_proj"] + ) self.model = MiMoV2FlashMultiTokenPredictor( atom_config=atom_config, prefix=maybe_prefix(prefix, "model") diff --git a/atom/models/minimax_m2.py b/atom/models/minimax_m2.py index 7d960edd0d..d928e54840 100644 --- a/atom/models/minimax_m2.py +++ b/atom/models/minimax_m2.py @@ -1,7 +1,10 @@ from typing import Optional, Union import torch -from aiter.dist.communication_op import tensor_model_parallel_all_reduce +from aiter.dist.communication_op import ( + tensor_model_parallel_all_reduce, + tensor_model_parallel_fused_qknorm_allreduce, +) from aiter.dist.parallel_state import ( get_pp_group, get_tensor_model_parallel_rank, @@ -224,31 +227,42 @@ def forward( hidden_states: torch.Tensor, ) -> torch.Tensor: qkv = self.qkv_proj(hidden_states) - q, k, v = torch.split(qkv, [self.q_size, self.kv_size, self.kv_size], dim=-1) if self.use_qk_norm: # TP-aware RMSNorm: all-reduce variance across TP ranks so # normalization uses the global variance (over 6144/1024 dims) # rather than per-rank variance (768/128 dims). - orig_dtype = q.dtype - q = q.to(torch.float32) - k = k.to(torch.float32) - q_var = q.pow(2).mean(dim=-1, keepdim=True) - k_var = k.pow(2).mean(dim=-1, keepdim=True) - if self.tp_size > 1: - qk_var = torch.cat([q_var, k_var], dim=-1) - qk_var = tensor_model_parallel_all_reduce(qk_var) / self.tp_size - q_var, k_var = qk_var.chunk(2, dim=-1) - q = (q * torch.rsqrt(q_var + self.rms_norm_eps) * self.q_norm.weight).to( - orig_dtype - ) - k = (k * torch.rsqrt(k_var + self.rms_norm_eps) * self.k_norm.weight).to( - orig_dtype + if qkv.shape[0] <= 256 and self.tp_size > 1: + q, k, v = tensor_model_parallel_fused_qknorm_allreduce( + qkv, self.q_norm.weight, self.k_norm.weight, self.rms_norm_eps + ) + else: + q, k, v = torch.split( + qkv, [self.q_size, self.kv_size, self.kv_size], dim=-1 + ) + orig_dtype = q.dtype + q = q.to(torch.float32) + k = k.to(torch.float32) + q_var = q.pow(2).mean(dim=-1, keepdim=True) + k_var = k.pow(2).mean(dim=-1, keepdim=True) + if self.tp_size > 1: + qk_var = torch.cat([q_var, k_var], dim=-1) + qk_var = tensor_model_parallel_all_reduce(qk_var) / self.tp_size + q_var, k_var = qk_var.chunk(2, dim=-1) + q = ( + q * torch.rsqrt(q_var + self.rms_norm_eps) * self.q_norm.weight + ).to(orig_dtype) + k = ( + k * torch.rsqrt(k_var + self.rms_norm_eps) * self.k_norm.weight + ).to(orig_dtype) + else: + q, k, v = torch.split( + qkv, [self.q_size, self.kv_size, self.kv_size], dim=-1 ) - attn_output = self.attn( - query=q, key=k, value=v, positions=positions, q_scale=None, qkv=qkv - ) + attn_output = self.attn( + query=q, key=k, value=v, positions=positions, q_scale=None, qkv=qkv + ) output = self.o_proj(attn_output) return output diff --git a/atom/models/qwen3_5.py b/atom/models/qwen3_5.py index 008e8877f7..88282ac739 100644 --- a/atom/models/qwen3_5.py +++ b/atom/models/qwen3_5.py @@ -40,10 +40,6 @@ maybe_prefix, extract_layer_index, ) -from atom.model_ops.split_chunk import ( - fused_split_chunk_zeros, - fused_split_chunk_zeros_qwen3_5_qkvzba, -) if is_vllm(): from vllm.model_executor.layers.mamba.mamba_utils import ( @@ -251,35 +247,35 @@ def forward( 1. Input projection 2. Core attention (custom op) """ + num_tokens = hidden_states.size(0) # ============================================================ # Part 1: Input Projection # ============================================================ + v_heads_tp = self.num_v_heads // self.tp_size + qkv_size = self.conv_dim // self.tp_size + z_size = v_heads_tp * self.head_v_dim + b_size = v_heads_tp + a_size = v_heads_tp + if hasattr(self, "in_proj_qkvzba"): qkvzba = self.in_proj_qkvzba(hidden_states) - k_heads_after_tp = self.num_k_heads // self.tp_size - v_heads_after_tp = self.num_v_heads // self.tp_size - mixed_qkv, z, b, a, core_attn_out = fused_split_chunk_zeros_qwen3_5_qkvzba( - qkvzba, - k_heads_after_tp, - v_heads_after_tp, - self.head_k_dim, - self.head_v_dim, + # Qwen3.5 layout is already contiguous [q|k|v|z|b|a] + mixed_qkv, z_flat, b, a = torch.split( + qkvzba, [qkv_size, z_size, b_size, a_size], dim=-1 ) else: if x_fp8 is not None: mixed_qkvz = self.in_proj_qkvz(x_fp8, x_scale=x_scale) else: mixed_qkvz = self.in_proj_qkvz(hidden_states) - ba = self.in_proj_ba(hidden_states) - - qkv_size = (self.key_dim * 2 + self.value_dim) // self.tp_size - z_size = self.value_dim // self.tp_size - num_v_heads_tp = self.num_v_heads // self.tp_size + projected_ba = self.in_proj_ba(hidden_states) + # Qwen3.5 layout is already contiguous [q|k|v|z] and [b|a] + mixed_qkv, z_flat = torch.split(mixed_qkvz, [qkv_size, z_size], dim=-1) + b, a = torch.split(projected_ba, [b_size, a_size], dim=-1) - mixed_qkv, z, b, a, core_attn_out = fused_split_chunk_zeros( - mixed_qkvz, ba, qkv_size, z_size, self.head_v_dim, num_v_heads_tp - ) + z = z_flat.view(num_tokens, v_heads_tp, self.head_v_dim) + core_attn_out = torch.empty_like(z) # ============================================================ # Part 2: Core Attention (Custom Op) @@ -494,6 +490,25 @@ def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: ) +_BF16_IN_PROJ_MAPPING = { + "in_proj_qkv": ("in_proj_qkvzba", (0, 1, 2)), + "in_proj_z": ("in_proj_qkvzba", 3), + "in_proj_b": ("in_proj_qkvzba", 4), + "in_proj_a": ("in_proj_qkvzba", 5), +} + + +def _apply_bf16_in_proj_mapping(mapping: dict, atom_config: Config) -> dict: + if atom_config.quant_config.global_quant_config.quant_dtype != torch.bfloat16: + return mapping + + mapping.pop("in_proj_qkvz", None) + mapping.pop("in_proj_ba", None) + mapping["in_proj_qkvzba"] = ("in_proj_qkvzba", None) + mapping.update(_BF16_IN_PROJ_MAPPING) + return mapping + + class Qwen3_5ForConditionalGenerationTextOnly(nn.Module): packed_modules_mapping = { "q_proj": ("qkv_proj", "q"), @@ -521,6 +536,9 @@ class Qwen3_5ForConditionalGenerationTextOnly(nn.Module): def __init__(self, atom_config: Config, prefix: str = ""): super().__init__() self.config = atom_config.hf_config + self.packed_modules_mapping = _apply_bf16_in_proj_mapping( + dict(self.packed_modules_mapping), atom_config + ) self.visual = PPMissingLayer() self.language_model = Qwen3_5ForCausalLM(atom_config=atom_config, prefix="") self.make_empty_intermediate_tensors = ( @@ -558,6 +576,9 @@ class Qwen3_5MoeForConditionalGenerationTextOnly( def __init__(self, atom_config: Config, prefix: str = ""): nn.Module.__init__(self) self.config = atom_config.hf_config + self.packed_modules_mapping = _apply_bf16_in_proj_mapping( + dict(self.packed_modules_mapping), atom_config + ) self.visual = PPMissingLayer() self.language_model = Qwen3_5MoeForCausalLM(atom_config=atom_config, prefix="") self.make_empty_intermediate_tensors = ( diff --git a/atom/models/qwen3_5_mtp.py b/atom/models/qwen3_5_mtp.py index 1ec4df8a46..73e8b39116 100644 --- a/atom/models/qwen3_5_mtp.py +++ b/atom/models/qwen3_5_mtp.py @@ -136,8 +136,6 @@ def remap_mtp_weight_name(self, name: str) -> str | None: def __init__(self, atom_config: Config, prefix: str = ""): super().__init__() config = get_qwen3_5_text_config(atom_config) - if atom_config.enable_prefix_caching: - raise ValueError("Qwen3_5MTP currently does not support prefix caching") self.config = config # Reindex MTP exclude entries on a copy: checkpoint uses 0-based diff --git a/atom/models/qwen3_next.py b/atom/models/qwen3_next.py index 0361602119..40ad2380a3 100644 --- a/atom/models/qwen3_next.py +++ b/atom/models/qwen3_next.py @@ -20,19 +20,17 @@ from atom.model_ops.layernorm import GemmaRMSNorm as Qwen3NextRMSNorm from atom.model_ops.layernorm import RMSNormGated from atom.model_ops.linear import ( + BAParallelLinear, ColumnParallelLinear, MergedColumnParallelLinear, MergedReplicatedLinear, QKVGParallelLinear, QKVParallelLinear, QKVZBAParallelLinear, + QKVZParallelLinear, RowParallelLinear, -) +) # noqa: F401 from atom.model_ops.moe import FusedMoE -from atom.model_ops.split_chunk import ( - fused_split_chunk_qwen_next_qkvz_ba, - fused_split_chunk_qwen_next_qkvzba, -) from atom.model_ops.topK import is_rocm_aiter_fusion_shared_expert_enabled from atom.model_ops.utils import atom_parameter from atom.models.utils import ( @@ -487,7 +485,8 @@ def __init__( self.config = config self.quant_config = quant_config - self.speculative_config = speculative_config + + self.speculative_config = speculative_config or atom_config.speculative_config self.num_spec = ( self.speculative_config.num_speculative_tokens if self.speculative_config @@ -594,103 +593,27 @@ def create_qkvzba_proj(self, quant_config, prefix): prefix=f"{self.prefix}.in_proj_qkvzba", ) else: - self.in_proj_qkvz = self.create_qkvz_proj( - hidden_size=self.hidden_size, - key_dim=self.key_dim, - value_dim=self.value_dim, + # Non-fused path (FP8): deinterleave weights + scales at load + # time so the forward can use zero-copy torch.split. + self.in_proj_qkvz = QKVZParallelLinear( + input_size=self.hidden_size, + head_k_dim=self.head_k_dim, + head_v_dim=self.head_v_dim, + num_k_heads=self.num_k_heads, + num_v_heads=self.num_v_heads, + bias=False, quant_config=quant_config, - prefix=prefix, + prefix=f"{prefix}.in_proj_qkvz", ) - - self.in_proj_ba = self.create_ba_proj( - hidden_size=self.hidden_size, + self.in_proj_ba = BAParallelLinear( + input_size=self.hidden_size, + num_k_heads=self.num_k_heads, num_v_heads=self.num_v_heads, + bias=False, quant_config=quant_config, - prefix=prefix, + prefix=f"{prefix}.in_proj_ba", ) - def create_qkvz_proj( - self, - hidden_size: int, - key_dim: int, - value_dim: int, - quant_config: QuantizationConfig | None, - prefix: str, - ) -> MergedColumnParallelLinear: - return MergedColumnParallelLinear( - input_size=hidden_size, - output_sizes=[sum((key_dim, key_dim, value_dim, value_dim))], - bias=False, - quant_config=quant_config, - prefix=f"{prefix}.in_proj_qkvz", - ) - - def create_ba_proj( - self, - hidden_size: int, - num_v_heads: int, - quant_config: QuantizationConfig | None, - prefix: str, - ) -> MergedColumnParallelLinear: - return MergedColumnParallelLinear( - input_size=hidden_size, - output_sizes=[num_v_heads, num_v_heads], - bias=False, - quant_config=quant_config, - prefix=f"{prefix}.in_proj_ba", - ) - - def fix_query_key_value_ordering( - self, - mixed_qkvz: torch.Tensor, - mixed_ba: torch.Tensor, - ): - """ - Derives `query`, `key` and `value` tensors from `mixed_qkvzba`. - """ - new_tensor_shape_qkvz = mixed_qkvz.size()[:-1] + ( - self.num_k_heads // self.tp_size, - ( - self.head_k_dim - + self.head_k_dim - + (self.head_v_dim + self.head_v_dim) - * self.num_v_heads - // self.num_k_heads - ), - ) - new_tensor_shape_ba = mixed_qkvz.size()[:-1] + ( - self.num_k_heads // self.tp_size, - 2 * self.num_v_heads // self.num_k_heads, - ) - - mixed_qkvz = mixed_qkvz.view(*new_tensor_shape_qkvz) - mixed_ba = mixed_ba.view(*new_tensor_shape_ba) - - split_arg_list_qkvz = [ - self.head_k_dim, - self.head_k_dim, - (self.num_v_heads // self.num_k_heads * self.head_v_dim), - (self.num_v_heads // self.num_k_heads * self.head_v_dim), - ] - split_arg_list_ba = [ - self.num_v_heads // self.num_k_heads, - self.num_v_heads // self.num_k_heads, - ] - - # [b, sq, ng, (hn + hn + np/ng * hn + np/ng + np/ng)] - # --> [b, sq, ng, hn], [b, sq, ng, hn], [b, sq, ng, np/ng * hn], - # [b, sq, ng, np/ng * hn], [b, sq, ng, np/ng], [b, sq, ng, np/ng] - query, key, value, z = torch.split(mixed_qkvz, split_arg_list_qkvz, dim=2) - b, a = torch.split(mixed_ba, split_arg_list_ba, dim=2) - - # [b, sq, ng, np/ng * hn] -> [b, sq, np, hn] - value = value.reshape(value.size(0), -1, self.head_v_dim) - z = z.reshape(z.size(0), -1, self.head_v_dim) - b = b.reshape(b.size(0), self.num_v_heads // self.tp_size) - a = a.reshape(a.size(0), self.num_v_heads // self.tp_size) - - return query, key, value, z, b, a - def rearrange_mixed_qkv(self, mixed_qkv): if mixed_qkv is None: return None, None, None @@ -723,43 +646,42 @@ def forward( 3. Output projection """ + num_tokens = hidden_states.shape[0] # ============================================================ # Part 1: Input Projection # ============================================================ + v_heads_tp = self.num_v_heads // self.tp_size + qkv_size = self.conv_dim // self.tp_size + z_size = v_heads_tp * self.head_v_dim + b_size = v_heads_tp + a_size = v_heads_tp + if hasattr(self, "in_proj_qkvzba"): - projected_states_qkvzba = self.in_proj_qkvzba(hidden_states) - k_heads_after_tp = self.num_k_heads // self.tp_size - v_heads_after_tp = self.num_v_heads // self.tp_size - mixed_qkv, z, b, a, core_attn_out = fused_split_chunk_qwen_next_qkvzba( - projected_states_qkvzba, - k_heads_after_tp, - v_heads_after_tp, - self.head_k_dim, - self.head_v_dim, + projected = self.in_proj_qkvzba(hidden_states) + # Output layout is [q|k|v|z|b|a] contiguous (deinterleaved at load) + mixed_qkv, z_flat, b, a = torch.split( + projected, [qkv_size, z_size, b_size, a_size], dim=-1 ) + z = z_flat.view(num_tokens, v_heads_tp, self.head_v_dim) else: if x_fp8 is not None: - projected_states_qkvz = self.in_proj_qkvz(x_fp8, x_scale=x_scale) + projected_qkvz = self.in_proj_qkvz(x_fp8, x_scale=x_scale) else: - projected_states_qkvz = self.in_proj_qkvz(hidden_states) - projected_states_ba = self.in_proj_ba(hidden_states) # always BF16 - # Use Triton kernel to process qkvz and ba - num_k_heads_tp = self.num_k_heads // self.tp_size - num_v_heads_tp = self.num_v_heads // self.tp_size - mixed_qkv, z, b, a, core_attn_out = fused_split_chunk_qwen_next_qkvz_ba( - projected_states_qkvz, - projected_states_ba, - num_k_heads_tp, - num_v_heads_tp, - self.head_k_dim, - self.head_v_dim, + projected_qkvz = self.in_proj_qkvz(hidden_states) + projected_ba = self.in_proj_ba(hidden_states) # always BF16 + # Weights deinterleaved at load → output is [q|k|v|z] and [b|a] + mixed_qkv = projected_qkvz[:, :qkv_size] + z = projected_qkvz[:, qkv_size:].view( + num_tokens, v_heads_tp, self.head_v_dim ) + b = projected_ba[:, :b_size] + a = projected_ba[:, b_size:] + + core_attn_out = torch.empty(z.shape, dtype=z.dtype, device=z.device) # ============================================================ # Part 2: Core Attention (Custom Op) # ============================================================ - # Note: we should not use torch.empty here like other attention backends, - # see discussions in https://github.com/vllm-project/vllm/pull/28182 core_attn_out = self.attn(mixed_qkv, b, a, core_attn_out) @@ -942,7 +864,6 @@ def forward( residual: torch.Tensor | None, ) -> tuple[torch.Tensor, torch.Tensor]: # Self Attention - if self.input_layernorm.use_fused_quant: if residual is None: residual = hidden_states @@ -1138,6 +1059,11 @@ def __init__( if self.config.tie_word_embeddings: self.lm_head.weight = self.model.embed_tokens.weight + # Expose embed_tokens at this level for vLLM MTP embedding sharing. + # vLLM's proposer accesses target_wrapper.model.embed_tokens, where + # target_wrapper.model = this class (Qwen3NextForCausalLM). + self.embed_tokens = self.model.embed_tokens + self.make_empty_intermediate_tensors = ( self.model.make_empty_intermediate_tensors ) @@ -1211,6 +1137,7 @@ def get_mamba_state_shape_from_config( if vllm_config.speculative_config else 0 ) + return MambaStateShapeCalculator.gated_delta_net_state_shape( tp_size, hf_config.linear_num_key_heads, diff --git a/atom/models/qwen3_next_mtp.py b/atom/models/qwen3_next_mtp.py index 8fe57e7408..f865aa2dbc 100644 --- a/atom/models/qwen3_next_mtp.py +++ b/atom/models/qwen3_next_mtp.py @@ -43,6 +43,7 @@ def __init__(self, atom_config: Config, prefix: str = ""): self.config.hidden_size, bias=False, quant_config=quant_config, + prefix=f"{prefix}.fc", ) self.layers = torch.nn.ModuleList( @@ -107,6 +108,8 @@ class Qwen3NextMTP(nn.Module): "v_proj": ("qkv_proj", "v"), "gate_proj": ("gate_up_proj", 0), "up_proj": ("gate_up_proj", 1), + ".gate.": (".gate.", 0), + "shared_expert_gate": ("gate", 1), } weights_mapping = {"mtp.": "model."} @@ -128,8 +131,6 @@ def remap_mtp_weight_name(self, name: str) -> str | None: def __init__(self, atom_config: Config, prefix: str = ""): super().__init__() config = atom_config.hf_config - if atom_config.enable_prefix_caching: - raise ValueError("Qwen3NextMTP currently does not support prefix caching") self.config = config self.model = Qwen3NextMultiTokenPredictor( atom_config=atom_config, prefix=maybe_prefix(prefix, "mtp") @@ -168,9 +169,19 @@ def compute_logits( def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: # Params for weights, fp8 weight scales, fp8 activation scales # (param_name, weight_name, expert_id, shard_id) + # Mirror target's get_expert_mapping: when shared-expert fusion is on, + # the loader rewrites `mlp.shared_expert.*` to `mlp.experts.{N}.*` + # (where N == n_routed_experts), so the expert_mapping must include + # an extra slot for that fused shared-expert. Without this, MTP's + # shared_expert weights get silently dropped during loading. + from atom.model_ops.topK import is_rocm_aiter_fusion_shared_expert_enabled + + n_routed = getattr(self.config, "n_routed_experts", self.config.num_experts) + n_shared = getattr(self.config, "n_shared_experts", 1) return FusedMoE.make_expert_params_mapping( ckpt_gate_proj_name="gate_proj", ckpt_down_proj_name="down_proj", ckpt_up_proj_name="up_proj", - num_experts=self.config.num_experts, + num_experts=n_routed + + (n_shared if is_rocm_aiter_fusion_shared_expert_enabled() else 0), ) diff --git a/atom/models/utils.py b/atom/models/utils.py index e5d834af3a..afb5f0ea82 100644 --- a/atom/models/utils.py +++ b/atom/models/utils.py @@ -241,6 +241,134 @@ def maybe_prefix(prefix: str, name: str) -> str: return name if not prefix else f"{prefix}.{name}" +def ckpt_has_tensor_suffix(model_path: str, suffix: str) -> bool: + """Return True if the checkpoint at ``model_path`` contains a tensor whose + key ends with ``suffix``. + + Used by MTP modules to decide whether the entry projection (eh_proj/fc) is + actually quantized on disk: when the HF quantization_config's exclude list + is incomplete, falling back to the global quant spec creates an + uninitialized weight_scale that silently corrupts inference. Checking the + safetensors index for a sibling ``*.weight_scale`` is the only reliable + signal. + """ + import glob + import json + import os + + if not model_path or not os.path.isdir(model_path): + logger.warning( + "ckpt_has_tensor_suffix: model_path %r is not a directory; " + "assuming suffix %r is absent", + model_path, + suffix, + ) + return False + + index_files = glob.glob(os.path.join(model_path, "*.safetensors.index.json")) + if index_files: + try: + with open(index_files[0]) as f: + weight_map = json.load(f)["weight_map"] + except (OSError, json.JSONDecodeError, KeyError) as e: + logger.warning( + "ckpt_has_tensor_suffix: failed to read %s (%s); " + "assuming suffix %r is absent", + index_files[0], + e, + suffix, + ) + return False + return any(k.endswith(suffix) for k in weight_map) + + safetensors_files = glob.glob(os.path.join(model_path, "*.safetensors")) + if safetensors_files: + try: + from safetensors import safe_open + + with safe_open(safetensors_files[0], framework="pt") as sf: + return any(k.endswith(suffix) for k in sf.keys()) + except Exception as e: + logger.warning( + "ckpt_has_tensor_suffix: failed to read %s (%s); " + "assuming suffix %r is absent", + safetensors_files[0], + e, + suffix, + ) + return False + + return False + + +def ckpt_shared_expert_count(model_path: str | None) -> int: + """Return the number of shared experts present in the local checkpoint. + + 0 means no `shared_expert` weights were found. 1 covers the flat-block + layout used by every MTP checkpoint shipped today (a single shared block + written without a numeric index, e.g. `...mlp.shared_experts.gate_proj`). + A value >1 is parsed from indexed keys (e.g. `shared_experts.0.`, + `shared_experts.1.`) and emits a warning, since no released model uses + that layout — the path exists to surface a mismatch instead of silently + loading the wrong count. + + Used by SpeculativeConfig to synthesize n_shared_experts when the HF + config doesn't carry it. Non-local paths or unreadable indexes return 0 + (leaving the field unset, safer than fabricating one). + """ + import glob + import json + import os + import re + + if not model_path or not os.path.isdir(model_path): + return 0 + + keys: list[str] | None = None + index_files = glob.glob(os.path.join(model_path, "*.safetensors.index.json")) + if index_files: + try: + with open(index_files[0]) as f: + keys = list(json.load(f)["weight_map"].keys()) + except (OSError, json.JSONDecodeError, KeyError): + return 0 + else: + safetensors_files = glob.glob(os.path.join(model_path, "*.safetensors")) + if safetensors_files: + try: + from safetensors import safe_open + + with safe_open(safetensors_files[0], framework="pt") as sf: + keys = list(sf.keys()) + except Exception: + return 0 + + if not keys: + return 0 + + shared_keys = [k for k in keys if "shared_expert" in k] + if not shared_keys: + return 0 + + indexed = re.compile(r"shared_experts?\.(\d+)\.") + max_idx = -1 + for k in shared_keys: + m = indexed.search(k) + if m: + max_idx = max(max_idx, int(m.group(1))) + if max_idx >= 0: + count = max_idx + 1 + logger.warning( + "Checkpoint at %r carries indexed shared_experts (parsed " + "n_shared_experts=%d). This layout is unprecedented in shipped " + "models; verify against the model definition before trusting it.", + model_path, + count, + ) + return count + return 1 + + def cast_overflow_tensors( tensors: torch.Tensor, offset: float = 1000, diff --git a/atom/plugin/attention.py b/atom/plugin/attention.py index d6aaf4f11e..e46441af37 100644 --- a/atom/plugin/attention.py +++ b/atom/plugin/attention.py @@ -22,6 +22,17 @@ disable_vllm_plugin_attention = envs.ATOM_DISABLE_VLLM_PLUGIN_ATTENTION +def _get_aiter_kv_cache_dtype(config) -> torch.dtype: + kv_cache_dtype = config.cache_config.cache_dtype + if kv_cache_dtype == "auto": + kv_cache_dtype = "bf16" + elif kv_cache_dtype == "bfloat16": + kv_cache_dtype = "bf16" + elif kv_cache_dtype == "float16": + kv_cache_dtype = "fp16" + return dtypes.d_dtypes[kv_cache_dtype] + + @dataclass class AiterFlashAttentionPhaseMetadata: max_query_len: int @@ -283,6 +294,28 @@ def init_method_under_plugin_mode( i64_kwargs = {"dtype": torch.int64, "device": device} self.positions = CpuGpuBuffer(max_num_batched_tokens, **i64_kwargs) + # Bump reorder_batch_threshold so multi-token spec-decode requests + # (MTP / EAGLE) are routed through the decode path. Mirrors vLLM's + # AttentionMetadataBuilder._init_reorder_batch_threshold(supports_spec_as_decode=True). + speculative_config = getattr(config, "speculative_config", None) + if ( + getattr(self, "reorder_batch_threshold", None) is not None + and speculative_config is not None + and getattr(speculative_config, "num_speculative_tokens", None) is not None + ): + parallel_drafting = getattr(speculative_config, "parallel_drafting", False) + max_num_queries_for_spec = 1 + (2 if parallel_drafting else 1) * ( + speculative_config.num_speculative_tokens + ) + self.reorder_batch_threshold = max( + self.reorder_batch_threshold, max_num_queries_for_spec + ) + logger.info( + "Spec decode: bumped reorder_batch_threshold to %d (num_spec_tokens=%d)", + self.reorder_batch_threshold, + speculative_config.num_speculative_tokens, + ) + return init_method_under_plugin_mode @@ -300,7 +333,7 @@ def setup_attn_metadata_builder_base_class_and_attributes(class_dict: dict): needs_generic = True # align with vllm rocm aiter fa - class_dict["_cudagraph_support"] = AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE + class_dict["_cudagraph_support"] = AttentionCGSupport.UNIFORM_BATCH class_dict["reorder_batch_threshold"] = 1 return base_class, generic_base, needs_generic, class_dict @@ -324,9 +357,12 @@ def build( from vllm.v1.attention.backends.utils import split_decodes_prefills_and_extends - # here assume the decode num token is 1 per request + # decode_threshold tracks reorder_batch_threshold so MTP/EAGLE + # multi-token verification (query_len > 1) routes through decode. + decode_threshold = getattr(self, "reorder_batch_threshold", 1) or 1 split_ret = split_decodes_prefills_and_extends( - common_attn_metadata=common_attn_metadata, decode_threshold=1 + common_attn_metadata=common_attn_metadata, + decode_threshold=decode_threshold, ) ( @@ -350,7 +386,15 @@ def build( query_lens_cpu = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1] + # The spec-decode draft path invalidates it between proposal steps when + # num_speculative_tokens > 1. + # Fall back to seq_lens - query_lens computed on already-CPU tensors. num_computed_tokens_cpu = common_attn_metadata._num_computed_tokens_cpu + # In async spec-decode mode (auto-enabled for MTP/EAGLE), vLLM sets + # _num_computed_tokens_cpu to None because the GPU seq_lens is the + # authoritative source. Reconstruct from CPU tensors we already have. + if num_computed_tokens_cpu is None: + num_computed_tokens_cpu = seq_lens - query_lens_cpu prefill_max_query_len = decode_max_query_len = ( common_attn_metadata.max_query_len @@ -840,12 +884,8 @@ def _build_decode( ) paged_kv_indptr = self.paged_kv_indptr[: 1 + num_reqs] - max_qo_len = ( - (query_start_loc_cpu[-1] - query_start_loc_cpu[-2]).item() - if query_start_loc_cpu.numel() > 1 - else 1 - ) - + qo_len = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1] + max_qo_len = qo_len.max().item() kv_indices_generate_triton( block_table_tensor, self.paged_kv_indices, @@ -882,7 +922,9 @@ def _build_decode( self.qo_indptr[1 + num_reqs :] = num_decode_tokens qo_indptr = self.qo_indptr[: 1 + num_reqs] - ctx_mla_ps = self._set_mla_persistent_worker_buffers(num_reqs, qo_indptr, 1) + ctx_mla_ps = self._set_mla_persistent_worker_buffers( + num_reqs, qo_indptr, max_qo_len + ) self.mla_persistent_metadata.update(ctx_mla_ps) attn_metadata = AiterMLADecodeMetadataForPluginMode( @@ -936,7 +978,6 @@ def build( query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu seq_lens = common_attn_metadata.seq_lens dcp_local_seq_lens = common_attn_metadata.dcp_local_seq_lens - num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = ( split_decodes_and_prefills( common_attn_metadata, @@ -1269,7 +1310,7 @@ def init_method_under_plugin_mode( 1, self.padded_num_attention_heads, torch.bfloat16, - dtypes.d_dtypes[config.cache_config.cache_dtype], + _get_aiter_kv_cache_dtype(config), is_sparse=False, # TODO: support sparse fast_mode=True, ) @@ -1747,11 +1788,20 @@ class AiterMLASparseMetadataForPluginMode: paged_kv_last_page_len: torch.Tensor paged_kv_indices: torch.Tensor paged_kv_indptr: torch.Tensor - paged_kv_indptr_rest: torch.Tensor + attn_out_dtype: torch.dtype block_size: int = 1 topk_tokens: int = 2048 + # Persistent MLA metadata (only populated when persistent mode is enabled, + # i.e. when the aiter sparse decode kernel supports work-stealing splits). + work_meta_data: torch.Tensor | None = None + work_indptr: torch.Tensor | None = None + work_info_set: torch.Tensor | None = None + reduce_indptr: torch.Tensor | None = None + reduce_final_map: torch.Tensor | None = None + reduce_partial_map: torch.Tensor | None = None + class vllmMLASparseAttentionMetadataBuilderMethods: def __init__(self): @@ -1769,18 +1819,61 @@ def build(self, common_prefix_len, common_attn_metadata, fast_build=False): ) # Zero-fill for cudagraphs self.req_id_per_token_buffer.fill_(0) + self.paged_kv_indices.fill_(0) + self.paged_kv_indptr.fill_(0) self.req_id_per_token_buffer[: req_id_per_token.shape[0]].copy_( req_id_per_token, non_blocking=True ) - self.paged_kv_indices.fill_(0) - self.paged_kv_indptr.fill_(0) + query_lens = ( + common_attn_metadata.query_start_loc[1:] + - common_attn_metadata.query_start_loc[:-1] + ) + seq_lens = common_attn_metadata.seq_lens + + from atom.plugin.attention_mla_sparse import generate_sparse_seqlen_triton + + sparse_seqlen = generate_sparse_seqlen_triton( + query_lens, + seq_lens, + common_attn_metadata.query_start_loc, + self.topk_tokens, + num_tokens, + common_attn_metadata.max_query_len, + ) + + torch.cumsum(sparse_seqlen, dim=0, out=self.paged_kv_indptr[1 : num_tokens + 1]) + self.paged_kv_indptr[num_tokens + 1 :].fill_(self.paged_kv_indptr[num_tokens]) req_id_per_token = self.req_id_per_token_buffer[:num_tokens] qo_indptr = self.qo_indptr[: num_tokens + 1] paged_kv_last_page_len = self.paged_kv_last_page_len[:num_tokens] paged_kv_indices = self.paged_kv_indices[: num_tokens * self.topk_tokens] paged_kv_indptr = self.paged_kv_indptr[: num_tokens + 1] - paged_kv_indptr_rest = self.paged_kv_indptr[num_tokens + 1 :] + + # ----- Compute persistent MLA metadata ----- + # The aiter sparse decode kernel uses qseqlen=1 (each query token is + # treated as its own batch entry), so persistent metadata can always + # be precomputed here. The kernel switches to the persistent + # work-stealing path automatically when work_meta_data is non-None. + get_mla_metadata_v1( + qo_indptr, + paged_kv_indptr, + paged_kv_last_page_len, + self.padded_num_heads, + 1, + True, + self._mla_work_meta_data, + self._mla_work_info_set, + self._mla_work_indptr, + self._mla_reduce_indptr, + self._mla_reduce_final_map, + self._mla_reduce_partial_map, + page_size=1, + kv_granularity=16, + max_seqlen_qo=1, + uni_seqlen_qo=1, + fast_mode=True, + ) attn_metadata_for_plugin_mode = AiterMLASparseMetadataForPluginMode( num_reqs=common_attn_metadata.num_reqs, @@ -1793,12 +1886,18 @@ def build(self, common_prefix_len, common_attn_metadata, fast_build=False): block_table=common_attn_metadata.block_table_tensor, req_id_per_token=req_id_per_token, block_size=self.kv_cache_spec.block_size, + attn_out_dtype=self.model_dtype, topk_tokens=self.topk_tokens, qo_indptr=qo_indptr, paged_kv_last_page_len=paged_kv_last_page_len, paged_kv_indices=paged_kv_indices, paged_kv_indptr=paged_kv_indptr, - paged_kv_indptr_rest=paged_kv_indptr_rest, + work_meta_data=self._mla_work_meta_data, + work_indptr=self._mla_work_indptr, + work_info_set=self._mla_work_info_set, + reduce_indptr=self._mla_reduce_indptr, + reduce_final_map=self._mla_reduce_final_map, + reduce_partial_map=self._mla_reduce_partial_map, ) attn_metadata = AttentionMetaData( @@ -2066,14 +2165,12 @@ def __init__(self): @staticmethod def get_supported_kernel_block_sizes(): - return [1] + return [1, 64] @classmethod def get_preferred_block_size(cls, default_block_size: int) -> int: - # ATOM's sparse MLA plugin path assumes kernel/page block size == 1 - # throughout metadata construction and KV-cache indexing, so force that - # value when vLLM 0.19 negotiates the backend-specific cache block size. - return 1 + # Prefer block_size == 64 so the indexer's preshuffled path is taken + return 64 @staticmethod def get_kv_cache_shape( @@ -2136,6 +2233,7 @@ def init_method_under_plugin_mode( from vllm.utils.platform_utils import num_compute_units except ImportError: from vllm.utils.platform_utils import get_cu_count as num_compute_units + from atom.models.utils import extract_layer_index from vllm.v1.worker.cp_utils import get_total_cp_world_size from vllm.utils.math_utils import cdiv @@ -2146,16 +2244,30 @@ def init_method_under_plugin_mode( self.kv_cache_spec = kv_cache_spec self.device = device max_num_batched_tokens = config.scheduler_config.max_num_batched_tokens + self._init_reorder_batch_threshold(1, supports_spec_as_decode=True) self.max_prefill_buffer_size = get_max_prefill_buffer_size( self.model_config.max_model_len ) - self.num_speculative_tokens = ( - self.vllm_config.speculative_config.num_speculative_tokens - if self.vllm_config.speculative_config - else 0 - ) - self.reorder_batch_threshold += self.num_speculative_tokens + # Determine if this builder is for draft model layers (MTP). + # Draft model layers have layer indices >= num_hidden_layers. + # The draft model itself does not do speculative decoding, so + # num_speculative_tokens should be 0 for its builders. + _is_draft_layer = False + if layer_names: + num_hidden_layers = config.model_config.hf_config.num_hidden_layers + layer_indices = [ + extract_layer_index(layer_name) for layer_name in layer_names + ] + _is_draft_layer = all(idx >= num_hidden_layers for idx in layer_indices) + if _is_draft_layer: + self.num_speculative_tokens = 0 + else: + self.num_speculative_tokens = ( + self.vllm_config.speculative_config.num_speculative_tokens + if self.vllm_config.speculative_config + else 0 + ) sm_count = num_compute_units(self.device.index) self.num_sms = sm_count @@ -2219,18 +2331,17 @@ def init_method_under_plugin_mode( self.vllm_config = config self.model_config = config.model_config + self.model_dtype = self.model_config.dtype self.kv_cache_spec = kv_cache_spec self.device = device max_num_batched_tokens = config.scheduler_config.max_num_batched_tokens + self._init_reorder_batch_threshold(1, supports_spec_as_decode=True) parallel_config = config.parallel_config self.num_heads = self.model_config.get_num_attention_heads(parallel_config) self.padded_num_heads = max(self.num_heads, _MLA_MIN_HEADS) self.mla_dims = get_mla_dims(self.model_config) self.topk_tokens = config.model_config.hf_config.index_topk - self.topk_tokens_tensor = torch.tensor( - [self.topk_tokens], device=device, dtype=torch.int32 - ) self.max_model_len_tensor = torch.tensor( [self.model_config.max_model_len], device=device, dtype=torch.int32 ) @@ -2262,6 +2373,52 @@ def init_method_under_plugin_mode( [max_num_batched_tokens + 1], dtype=torch.int32, device=device ) + # ----- Persistent MLA metadata buffers ----- + # The aiter sparse decode kernel supports a "persistent" path that + # uses precomputed work-splitting metadata for better load balancing + # across CUs. Mirrors the approach used in rocm_aiter_mla.py. + # + # In the sparse case each query token is its own "batch" entry in the + # qo_indptr (qo_indptr = [0, 1, 2, ..., num_tokens]) and max_qo_len=1. + # We pad get_mla_metadata_info_v1's batch_size to max_num_batched_tokens + # so the buffers are large enough for any decode shape we might see. + ( + (work_meta_data_size, work_meta_data_type), + (work_indptr_size, work_indptr_type), + (work_info_set_size, work_info_set_type), + (reduce_indptr_size, reduce_indptr_type), + (reduce_final_map_size, reduce_final_map_type), + (reduce_partial_map_size, reduce_partial_map_type), + ) = get_mla_metadata_info_v1( + max_num_batched_tokens, + 1, + self.padded_num_heads, + torch.bfloat16, + _get_aiter_kv_cache_dtype(config), + is_sparse=True, + fast_mode=True, + ) + self._mla_work_meta_data = torch.empty( + work_meta_data_size, dtype=work_meta_data_type, device=device + ) + self._mla_work_indptr = torch.empty( + work_indptr_size, dtype=work_indptr_type, device=device + ) + self._mla_work_info_set = torch.empty( + work_info_set_size, dtype=work_info_set_type, device=device + ) + self._mla_reduce_indptr = torch.empty( + reduce_indptr_size, dtype=reduce_indptr_type, device=device + ) + self._mla_reduce_final_map = torch.empty( + reduce_final_map_size, dtype=reduce_final_map_type, device=device + ) + self._mla_reduce_partial_map = torch.empty( + reduce_partial_map_size, + dtype=reduce_partial_map_type, + device=device, + ) + return init_method_under_plugin_mode @@ -2275,8 +2432,9 @@ def setup_mla_sparse_attn_metadata_builder_base_class_and_attributes(class_dict: generic_base = AttentionMetadataBuilder needs_generic = True - # align with vllm ROCMAiterMLASparseMetadataBuilder - class_dict["_cudagraph_support"] = AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE + # The plugin sparse MLA path builds per-token ragged metadata, so uniform + # spec-decode batches from the main model can use full decode cudagraphs. + class_dict["_cudagraph_support"] = AttentionCGSupport.UNIFORM_BATCH # For indexer metadata, not present in vllm sparse MLA class_dict["reorder_batch_threshold"] = 1 @@ -2297,12 +2455,16 @@ def unified_attention_with_output_base_for_plugin_mode( use_mla: bool, qkv: torch.Tensor, ) -> torch.Tensor: - atom_config = get_current_atom_config() + current_atom_config = get_current_atom_config() + static_forward_context = ( + current_atom_config.compilation_config.static_forward_context + ) + if use_mla: # raise NotImplementedError("MLA is not supported for plugin mode for now") kv_c_normed = k k_pe = v - self = atom_config.compilation_config.static_forward_context[layer_name] + self = static_forward_context[layer_name] q = self.q_proj(q, q_scale) q = q.view(-1, self.num_heads, self.qk_head_dim) # Add head dim of 1 to k_pe @@ -2321,8 +2483,7 @@ def unified_attention_with_output_base_for_plugin_mode( ) return self.o_proj(output) else: - self = atom_config.compilation_config.static_forward_context[layer_name] - + self = static_forward_context[layer_name] # here is the standard vllm attention impl interface # when using fusion, we need to pass the qkv and positions through the q,k,v # [watch out] accept_output_buffer must be False for plugin mode @@ -2330,7 +2491,14 @@ def unified_attention_with_output_base_for_plugin_mode( # ATOM needs to handle all of the buffer here # Positions for compiled unified_attention are provided via vLLM ForwardContext # (atom_positions) in ATOMModelBase.forward, not via this Python arg. - if envs.ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION: + # + # RoPE single-source-of-truth in plugin mode: when ATOM's + # PagedAttentionImpl is active, it applies RoPE inside + # rope_cache_plugin_mode. Applying RoPE here too would double apply + # it for any model hitting the "else" branch in rope_cache_plugin_mode when + # ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION is not set. + # if envs.ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION: + if current_atom_config.plugin_config.vllm_use_atom_attention: output = self.attn(q, k, v) else: # calculate the q and k with rotary embedding diff --git a/atom/plugin/attention_mha.py b/atom/plugin/attention_mha.py index ba2073d794..100c492aba 100644 --- a/atom/plugin/attention_mha.py +++ b/atom/plugin/attention_mha.py @@ -234,22 +234,37 @@ def paged_attention_triton_plugin_mode( v_cache: torch.Tensor, k_scale: torch.Tensor, v_scale: torch.Tensor, + num_decodes: int, out: torch.Tensor, attn_metadata: "AttentionMetaData", ps: bool = True, ): - o = out - num_seqs, num_q_heads_total, head_size = q.shape + # q.shape[0] == num_decodes * max_query_len for MTP (one row per decode + # token, query_len > 1). For non-MTP it equals num_decodes (query_len = 1). + # pa_decode_gluon handles multi-token causal masking internally when + # `query_length > 1` is passed; intermediate buffers must be sized + # `num_decodes` (not q.shape[0]) and `query_group_size` must include + # the max_qlen multiplier — mirroring server-mode `paged_attention_triton`. + _, num_q_heads_total, head_size = q.shape num_blocks, num_kv_heads, _, block_size, _ = k_cache.shape - query_group_size = num_q_heads_total // num_kv_heads + decode_metadata = attn_metadata.plugin_metadata.decode_metadata + max_qlen = decode_metadata.max_query_len if decode_metadata is not None else 1 assert num_q_heads_total % num_kv_heads == 0 + + seq_lens = attn_metadata.plugin_metadata.seq_lens[:num_decodes] + block_tables = attn_metadata.plugin_metadata.block_table[:num_decodes] + + query_group_size = max_qlen * (num_q_heads_total // num_kv_heads) context_partition_size = 256 - use_ps = self.adopt_persistent_kernel( - head_size, num_kv_heads, num_q_heads_total - ) + # use_ps = self.adopt_persistent_kernel( + # head_size, num_kv_heads, num_q_heads_total + # ) + use_ps = True if use_ps: - max_context_partition_num = get_recommended_splits(num_seqs, num_kv_heads) + max_context_partition_num = get_recommended_splits( + num_decodes, num_kv_heads + ) else: max_context_partition_num = _NO_PS_FIXED_SPLITS @@ -257,9 +272,8 @@ def paged_attention_triton_plugin_mode( max_context_partition_num = 1 context_partition_size = 128 - # Output buffers (same as Triton) intermediate_shape = ( - num_seqs, + num_decodes, num_kv_heads, max_context_partition_num, query_group_size, @@ -282,21 +296,19 @@ def paged_attention_triton_plugin_mode( k_scale = k_scale.unsqueeze(-1) v_scale = v_scale.unsqueeze(-1) - num_decode_seqs = q.shape[0] - seq_lens_decode = attn_metadata.plugin_metadata.seq_lens[:num_decode_seqs] - block_tables_decode = attn_metadata.plugin_metadata.block_table[ - :num_decode_seqs - ] - + # Kernel takes natural q layout [batch * query_length, num_q_heads, head_size]. + # Internally it derives batch_size = q.shape[0] // query_length and reshapes + # to [batch, query_length, num_kv_heads, group, head_size]. See + # aiter/aiter/ops/triton/gluon/pa_decode_gluon.py:5371-5377 and 5542-5544. torch.ops.aiter.pa_decode_gluon( - o, + out, q, k_cache, v_cache, - seq_lens_decode, - block_tables_decode, + seq_lens, + block_tables, self.scale, - 1, # query_lenth + max_qlen, # query_length — handles multi-token causal mask internally max_context_partition_num, context_partition_size, compute_type, @@ -311,8 +323,7 @@ def paged_attention_triton_plugin_mode( sliding_window=self.sliding_window, ps=use_ps, ) - - return o + return out def paged_attention_asm_plugin_mode( self, @@ -326,6 +337,11 @@ def paged_attention_asm_plugin_mode( attn_metadata: "AttentionMetaData", out: torch.Tensor, ): + decode_metadata = attn_metadata.plugin_metadata.decode_metadata + max_qlen = decode_metadata.max_query_len if decode_metadata is not None else 1 + qo_indptr = ( + decode_metadata.query_start_loc if decode_metadata is not None else None + ) aiter.pa_fwd_asm( Q=q, K=k_cache, @@ -335,9 +351,11 @@ def paged_attention_asm_plugin_mode( block_tables_stride0=attn_metadata.plugin_metadata.block_table[ :num_decodes ].stride(0), + max_qlen=max_qlen, K_QScale=k_scale, V_QScale=v_scale, out_=out[:num_decode_tokens], + qo_indptr=qo_indptr, high_precision=0, ) @@ -705,12 +723,13 @@ def forward_impl_plugin_mode( extend_tokens_slice = slice( num_decode_tokens, num_decode_tokens + num_extend_tokens ) + extend_reqs_slice = slice(num_decodes, num_decodes + num_extends) extend_querys = query[extend_tokens_slice] extend_keys = key[extend_tokens_slice] extend_values = value[extend_tokens_slice] extend_outputs = output[extend_tokens_slice] extend_block_table = attn_metadata.plugin_metadata.block_table[ - extend_tokens_slice + extend_reqs_slice ] extend_slot_mapping = attn_metadata.plugin_metadata.slot_mapping[ extend_tokens_slice @@ -744,6 +763,7 @@ def forward_impl_plugin_mode( v_cache=new_value_cache, k_scale=k_scale, v_scale=v_scale, + num_decodes=num_decodes, out=output_actual_tokens[:num_decode_tokens], attn_metadata=attn_metadata, ) @@ -756,6 +776,7 @@ def forward_impl_plugin_mode( v_cache=new_value_cache, k_scale=k_scale, v_scale=v_scale, + num_decodes=num_decodes, out=output_actual_tokens[:num_decode_tokens], attn_metadata=attn_metadata, ) diff --git a/atom/plugin/attention_mla_sparse.py b/atom/plugin/attention_mla_sparse.py index 6a3e01635c..c63dd09fa0 100644 --- a/atom/plugin/attention_mla_sparse.py +++ b/atom/plugin/attention_mla_sparse.py @@ -26,8 +26,8 @@ cp_gather_indexer_k_quant_cache, dtypes, indexer_k_quant_and_cache, + indexer_qk_rope_quant_and_cache, top_k_per_row_decode, - top_k_per_row_prefill, ) from aiter.ops.triton.fp8_mqa_logits import fp8_mqa_logits from aiter.ops.triton.pa_mqa_logits import deepgemm_fp8_paged_mqa_logits @@ -46,6 +46,188 @@ logger = logging.getLogger("atom") +@triton.jit +def _convert_req_index_to_global_index_kernel( + req_id_ptr, # int32 [num_tokens] + block_table_ptr, # int32 [num_requests, max_num_blocks_per_req] + token_indices_ptr, # int32 [num_tokens, NUM_TOPK_TOKENS] + cu_seqlens_ptr, # int32 [num_tokens + 1] + out_ptr, # int32 [num_tokens, NUM_TOPK_TOKENS] + # shapes (compile-time where possible) + max_num_blocks_per_req: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + BLOCK_N: tl.constexpr, # tile width along columns + # strides (in elements) + bt_stride0, + bt_stride1, + ti_stride0, + ti_stride1, +): + # program_id(0) -> token_id (row) + # program_id(1) -> tile index along columns + token_id = tl.program_id(0) + tile_id = tl.program_id(1) + + # Each program covers BLOCK_N consecutive columns + indice_id = tile_id * BLOCK_N + tl.arange(0, BLOCK_N) + + # Load request id for this token (no mask: grid is exact) + req = tl.load(req_id_ptr + token_id) + + # Load cumulative sequence lengths to get starting index of this request + seq_start = tl.load(cu_seqlens_ptr + token_id) + seq_end = tl.load(cu_seqlens_ptr + token_id + 1) + + if tile_id * BLOCK_N + seq_start >= seq_end: + return + + # Load token indices for this tile + ti_ptr = token_indices_ptr + token_id * ti_stride0 + indice_id * ti_stride1 + tok = tl.load(ti_ptr) # int32 + + # Only token == -1 should propagate as -1 + is_invalid_tok = tok < 0 + + # Compute block id and in-block offset + block_id = tok // BLOCK_SIZE + inblock_off = tok % BLOCK_SIZE + + # Guard block_table access + valid_block = (block_id < max_num_blocks_per_req) & (block_id >= 0) + bt_ptr = block_table_ptr + req * bt_stride0 + block_id * bt_stride1 + base = tl.load(bt_ptr, mask=valid_block, other=0) + + # # If token == -1 OR block_id OOB, output 0; else base * BLOCK_SIZE + offset + out_val = tl.where( + is_invalid_tok | (~valid_block), 0, base * BLOCK_SIZE + inblock_off + ) + out_ptr_ij = out_ptr + seq_start + indice_id + out_ptr_ij_mask = (seq_start + indice_id) < seq_end + + # store the results with mask + tl.store(out_ptr_ij, out_val, mask=out_ptr_ij_mask) + + +def triton_convert_req_index_to_global_index( + req_id: torch.Tensor, # int32 [num_tokens] + block_table: torch.Tensor, # int32 [num_requests, max_num_blocks_per_req] + token_indices: torch.Tensor, # int32 [num_tokens, NUM_TOPK_TOKENS] + cu_seqlens: torch.Tensor, # int32 [num_tokens + 1] + paged_kv_indices: torch.Tensor, # int32 [num_tokens * topk] out_buffer + BLOCK_SIZE: int = 64, + NUM_TOPK_TOKENS: int = 2048, + BLOCK_N: int = 128, # tile width along columns +): + """ + out[token_id, indice_id] = + block_table[req_id[token_id], + token_indices[token_id, indice_id] // BLOCK_SIZE] * BLOCK_SIZE + + token_indices[token_id, indice_id] % BLOCK_SIZE + + Only when token_indices[token_id, indice_id] == -1 do we output -1. + For safety, we also output -1 if the derived block_id would be + out-of-bounds. + """ + assert req_id.dtype == torch.int32 + assert block_table.dtype == torch.int32 + assert token_indices.dtype == torch.int32 + assert token_indices.shape[1] == NUM_TOPK_TOKENS + assert ( + NUM_TOPK_TOKENS % BLOCK_N == 0 + ), f"NUM_TOPK_TOKENS ({NUM_TOPK_TOKENS}) must be divisible byBLOCK_N ({BLOCK_N})" + # print("req_id: ", req_id, flush=True) + num_tokens = req_id.shape[0] + _, max_num_blocks_per_req = block_table.shape + tiles_per_row = NUM_TOPK_TOKENS // BLOCK_N + + # Ensure contiguous tensors on the same device + req_id_c = req_id.contiguous() + block_table_c = block_table.contiguous() + token_indices_c = token_indices.contiguous() + + # Strides in elements + bt_stride0, bt_stride1 = block_table_c.stride() + ti_stride0, ti_stride1 = token_indices_c.stride() + + # Exact 2D grid: tokens × column tiles + grid = (num_tokens, tiles_per_row) + + _convert_req_index_to_global_index_kernel[grid]( + req_id_c, + block_table_c, + token_indices_c, + cu_seqlens, + paged_kv_indices, + # shapes / constexprs + max_num_blocks_per_req, + BLOCK_SIZE, + BLOCK_N, + # strides + bt_stride0, + bt_stride1, + ti_stride0, + ti_stride1, + ) + return + + +@triton.jit +def generate_sparse_seqlen_kernel( + seq_len_ptr, # [num_seq] + cu_query_lens_ptr, # [num_seq] + out_ptr, # [num_query_tokens] + topk_token: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + seq_id = tl.program_id(0) + query_offset = tl.program_id(1) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + query_start = tl.load(cu_query_lens_ptr + seq_id) + query_end = tl.load(cu_query_lens_ptr + seq_id + 1) + if query_start + tl.program_id(1) * BLOCK_SIZE > query_end: + return + query_len = query_end - query_start + query_mask = query_offset + query_start < query_end + seq_len = tl.load(seq_len_ptr + seq_id) + # Just return since the out_ptr is zero initialized. + if seq_len == 0: + return + context_start_point = seq_len - query_len + sparse_seqlen = context_start_point + query_offset + sparse_seqlen_masked = tl.where( + sparse_seqlen + 1 < topk_token, sparse_seqlen + 1, topk_token + ) + tl.store( + out_ptr + query_start + query_offset, sparse_seqlen_masked, mask=query_mask + ) + + +def generate_sparse_seqlen_triton( + query_lens: torch.Tensor, + seq_lens: torch.Tensor, + cu_query_lens: torch.Tensor, + topk_token: int, + num_tokens: int, + max_query_len: int, +): + num_seqs = query_lens.size(0) + # zero initialize the tensor to make sure invalid positions will be zero + out = torch.zeros([num_tokens], dtype=torch.int32, device=query_lens.device) + block_size = 64 + num_block_per_row = triton.cdiv(max_query_len, block_size) + grid = ( + num_seqs, + num_block_per_row, + ) + generate_sparse_seqlen_kernel[grid]( + seq_lens, + cu_query_lens, + out, + topk_token, + block_size, + ) + return out + + @triton.jit def fetch_id_to_ragged_kernel( in_tensor_ptr, # [num_seq, topk] @@ -101,11 +283,10 @@ def __init__(self): "It is only used as a method container for the decorator." ) - def _forward_sparse_bf16_kv( + def _forward_sparse_mla( self, q: torch.Tensor, # [sq, heads, d_qk] kv_cache: torch.Tensor, # [blocks, heads, d_qk] - topk_indices_global: torch.Tensor, # [sq, topk] attn_metadata, layer, ) -> torch.Tensor: @@ -115,25 +296,12 @@ def _forward_sparse_bf16_kv( output = torch.empty( [num_tokens, self.padded_num_heads, self.kv_lora_rank], # dtype=q.dtype, - dtype=torch.bfloat16, + dtype=sparse_meta.attn_out_dtype, device=q.device, ) - seq_len = (topk_indices_global != -1).sum(dim=-1) - torch.cumsum(seq_len, dim=0, out=sparse_meta.paged_kv_indptr[1:]) - sparse_meta.paged_kv_indptr_rest.fill_(sparse_meta.paged_kv_indptr[-1]) - fetch_id_to_ragged_triton( - topk_indices_global, - sparse_meta.paged_kv_indptr, - sparse_meta.paged_kv_indices, - sparse_meta.topk_tokens, - ) - kv_buffer = kv_cache.unsqueeze(2) - # TODO: Currently, persistent mode has memory access issues when input context is long - # in fp8 kv cache settings, so it is not used for now. - # Re-enable persistent mode when the issue is fixed. mla_decode_fwd( q, kv_buffer.view(-1, 1, 1, q.shape[-1]), @@ -147,6 +315,12 @@ def _forward_sparse_bf16_kv( q_scale=layer._q_scale, kv_scale=layer._k_scale, page_size=1, + work_meta_data=sparse_meta.work_meta_data, + work_indptr=sparse_meta.work_indptr, + work_info_set=sparse_meta.work_info_set, + reduce_indptr=sparse_meta.reduce_indptr, + reduce_final_map=sparse_meta.reduce_final_map, + reduce_partial_map=sparse_meta.reduce_partial_map, ) if self.head_repeat_factor > 1: @@ -288,22 +462,15 @@ def forward_impl_sparse_plugin_mode( assert self.topk_indices_buffer is not None topk_indices = self.topk_indices_buffer[:num_actual_toks] - try: - from vllm.v1.attention.backends.mla.sparse_utils import ( - triton_convert_req_index_to_global_index, - ) - except ImportError: - from vllm.v1.attention.backends.mla.flashmla_sparse import ( - triton_convert_req_index_to_global_index, - ) - req_id_i32 = sparse_meta.req_id_per_token.to(dtype=torch.int32) block_table_i32 = sparse_meta.block_table.to(dtype=torch.int32) topk_indices_i32 = topk_indices.to(dtype=torch.int32) - topk_indices_global = triton_convert_req_index_to_global_index( - req_id_i32, # sparse_meta.req_id_per_token, - block_table_i32, # sparse_meta.block_table, - topk_indices_i32, # topk_indices, + triton_convert_req_index_to_global_index( + req_id_i32, + block_table_i32, + topk_indices_i32, + sparse_meta.paged_kv_indptr, + sparse_meta.paged_kv_indices, BLOCK_SIZE=sparse_meta.block_size, NUM_TOPK_TOKENS=sparse_meta.topk_tokens, ) @@ -316,9 +483,7 @@ def forward_impl_sparse_plugin_mode( layer._q_scale, ) q_out = q_flat.reshape(q_out.shape) - attn_out = self._forward_sparse_bf16_kv( - q_out, kv_cache, topk_indices_global, attn_metadata, layer - ) + attn_out = self._forward_sparse_mla(q_out, kv_cache, attn_metadata, layer) # V up-projection self._v_up_proj(attn_out, out=output[:num_actual_toks]) @@ -355,13 +520,13 @@ def MLASparseAttentionImplDecoratorForPluginMode(cls): Decorator that injects sparse MLA methods into the MLAAttentionImpl class. Applied alongside the regular MLAAttentionImplDecoratorForPluginMode. - Injects forward_impl_sparse_plugin_mode and _forward_sparse_bf16_kv. + Injects forward_impl_sparse_plugin_mode and _forward_sparse_mla. The patched forward_impl in register.py calls forward_impl_plugin_mode, which dispatches to forward_impl_sparse_plugin_mode when topk_indices_buffer is set. """ sparse_method_names = [ - "_forward_sparse_bf16_kv", + "_forward_sparse_mla", "forward_impl_sparse_plugin_mode", ] @@ -393,7 +558,7 @@ def sparse_attn_indexer_plugin_mode( hidden_states: torch.Tensor, k_cache_prefix: str, kv_cache: torch.Tensor, - q_fp8: torch.Tensor, + q_input: torch.Tensor, k: torch.Tensor, weights: torch.Tensor, quant_block_size: int, @@ -403,6 +568,15 @@ def sparse_attn_indexer_plugin_mode( max_model_len: int, total_seq_lens: int, topk_indices_buffer: torch.Tensor, + k_norm_weight: torch.Tensor, + k_norm_bias: torch.Tensor, + k_norm_eps: float, + positions: torch.Tensor, + cos_cache: torch.Tensor, + sin_cache: torch.Tensor, + weights_scale: float, + is_neox_style: bool, + use_qk_rope_cache_fusion: bool, ) -> torch.Tensor: try: from vllm.forward_context import ( @@ -419,12 +593,12 @@ def sparse_attn_indexer_plugin_mode( # During profile/dummy run the metadata dict may not contain # our layer or may be None. if attn_metadata_dict is None: - return weights + return torch.zeros_like(weights, dtype=torch.float32) if k_cache_prefix not in attn_metadata_dict: - return weights + return torch.zeros_like(weights, dtype=torch.float32) layer_meta = attn_metadata_dict[k_cache_prefix] if layer_meta is None: - return weights + return torch.zeros_like(weights, dtype=torch.float32) # In plugin mode, plugin_metadata is vllmDeepseekV32IndexerMetadata from # AiterMLASparseIndexerMetadataBuilder. @@ -434,14 +608,46 @@ def sparse_attn_indexer_plugin_mode( has_decode = indexer_meta.num_decodes > 0 has_prefill = indexer_meta.num_prefills > 0 num_decode_tokens = indexer_meta.num_decode_tokens - - indexer_k_quant_and_cache( - k, - kv_cache, - slot_mapping, - quant_block_size, - scale_fmt, - ) + kv_block_size = kv_cache.shape[1] + preshuffle_cache = kv_block_size != 1 + + if use_qk_rope_cache_fusion: + q_bf16 = q_input + q_fp8 = torch.empty_like(q_bf16, dtype=dtypes.fp8) + weights_out = torch.empty( + weights.shape, device=weights.device, dtype=torch.float32 + ) + indexer_qk_rope_quant_and_cache( + q_bf16, + q_fp8, + weights, + weights_out, + k, + kv_cache, + slot_mapping, + k_norm_weight, + k_norm_bias, + positions, + cos_cache, + sin_cache, + k_norm_eps, + quant_block_size, + scale_fmt, + weights_scale, + preshuffle=preshuffle_cache, + is_neox=is_neox_style, + ) + weights = weights_out + else: + q_fp8 = q_input + indexer_k_quant_and_cache( + k, + kv_cache, + slot_mapping, + quant_block_size, + scale_fmt, + preshuffle=preshuffle_cache, + ) topk_indices_buffer[: hidden_states.shape[0]] = -1 # topk_indices_buffer[: num_actual_tokens] = -1 @@ -465,6 +671,7 @@ def sparse_attn_indexer_plugin_mode( k_scale.view(dtypes.fp8), chunk.block_table, chunk.cu_seq_lens, + preshuffle=preshuffle_cache, ) logits = fp8_mqa_logits( @@ -530,6 +737,10 @@ def sparse_attn_indexer_plugin_mode( decode_metadata.seq_lens, decode_metadata.block_table, max_model_len, + ChunkK=256, + Preshuffle=preshuffle_cache, + KVBlockSize=kv_block_size, + WavePerEU=2, ) num_rows = logits.shape[0] @@ -565,7 +776,7 @@ def sparse_attn_indexer_fake( hidden_states: torch.Tensor, k_cache_prefix: str, kv_cache: torch.Tensor, - q_fp8: torch.Tensor, + q_input: torch.Tensor, k: torch.Tensor, weights: torch.Tensor, quant_block_size: int, @@ -575,6 +786,15 @@ def sparse_attn_indexer_fake( max_model_len: int, total_seq_lens: int, topk_indices_buffer: torch.Tensor, + k_norm_weight: torch.Tensor, + k_norm_bias: torch.Tensor, + k_norm_eps: float, + positions: torch.Tensor, + cos_cache: torch.Tensor, + sin_cache: torch.Tensor, + weights_scale: float, + is_neox_style: bool, + use_qk_rope_cache_fusion: bool, ) -> torch.Tensor: # profile run # NOTE(Chen): create the max possible flattened_kv. So that @@ -584,7 +804,7 @@ def sparse_attn_indexer_fake( ) _k_fp8 = _flattened_kv[..., :head_dim].view(torch.float8_e4m3fn).contiguous() _k_scale = _flattened_kv[..., head_dim:].view(torch.float32).contiguous() - return weights + return torch.empty(weights.shape, device=weights.device, dtype=torch.float32) direct_register_custom_op( @@ -616,8 +836,15 @@ def new_init(self, *args, **kwargs): def _deepseek_v32_indexer_get_kv_cache_spec(self, vllm_config): from vllm.v1.kv_cache_interface import MLAAttentionSpec + # Use the negotiated cache_config.block_size so the indexer cache and the + # main MLA cache share a single block_size. With a uniform block_size, + # vLLM groups the two layer types via `UniformTypeKVCacheSpecs.from_specs` + # and allocates a separate KVCacheTensor per layer sized to its own + # page_size_bytes (576B/token MLA vs 132B/token indexer); otherwise the + # smaller indexer page does not divide the MLA page and + # `unify_kv_cache_spec_page_size` raises NotImplementedError. return MLAAttentionSpec( - block_size=1, # block_size = 1 for indexer on ROCm + block_size=vllm_config.cache_config.block_size, num_kv_heads=1, head_size=self.head_dim, dtype=self.dtype, diff --git a/atom/plugin/config.py b/atom/plugin/config.py index 07aafa4a52..575948aab3 100644 --- a/atom/plugin/config.py +++ b/atom/plugin/config.py @@ -1,3 +1,4 @@ +import copy from typing import Any, Optional from dataclasses import dataclass @@ -31,6 +32,7 @@ class PluginConfig: sglang_enable_torch_compile: bool = False sglang_disable_cuda_graph: bool = False sglang_enable_dp_attention: bool = False + sglang_aiter_rank_id: int = 0 sglang_dist_init_addr: Optional[str] = None sglang_port_args: Any = None @@ -71,6 +73,45 @@ def _normalize_sglang_parallel_config( return tp_size, 1, 0, tp_rank +def _build_atom_speculative_config_from_vllm(vllm_spec_config: Any): + """Translate vLLM's SpeculativeConfig into ATOM's SpeculativeConfig. + + Reuses vLLM's already-loaded draft hf_config (skips a second disk fetch + in ATOM SpeculativeConfig.__post_init__) but still runs ATOM's + hf_config_override on it — so MTP model_type remap, n_routed_experts + backfill (Qwen families), and architecture rewrite all land on the + draft config in one place. Mirrors how standalone ATOM MTP exposes + the draft hf_config via atom_config.speculative_config. + + The draft hf_config is deepcopied first because hf_config_override + mutates `architectures` to ATOM's standalone naming (e.g. + "Qwen3NextMTPModel"), which differs from vLLM's registry name + ("Qwen3NextMTP"). Mutating in place would make vLLM's later draft + architecture lookup fail. + """ + if vllm_spec_config is None: + return None + + from atom.config import SpeculativeConfig + + draft_model_config = getattr(vllm_spec_config, "draft_model_config", None) + draft_hf_config = getattr(draft_model_config, "hf_config", None) + if draft_hf_config is not None: + draft_hf_config = copy.deepcopy(draft_hf_config) + model_path = getattr(draft_model_config, "model", None) or getattr( + vllm_spec_config, "model", None + ) + + return SpeculativeConfig( + method=getattr(vllm_spec_config, "method", "") or "", + model=model_path, + num_speculative_tokens=getattr( + vllm_spec_config, "num_speculative_tokens", None + ), + draft_model_hf_config=draft_hf_config, + ) + + def _generate_atom_config_from_vllm_config(config: Any) -> PluginConfig: from atom.config import Config, CompilationConfig @@ -117,6 +158,10 @@ def _generate_atom_config_from_vllm_config(config: Any) -> PluginConfig: max_num_batched_tokens = vllm_scheduler_config.max_num_batched_tokens + atom_speculative_config = _build_atom_speculative_config_from_vllm( + getattr(config, "speculative_config", None) + ) + return Config( model=vllm_model_config.model, trust_remote_code=getattr(vllm_model_config, "trust_remote_code", False), @@ -140,6 +185,7 @@ def _generate_atom_config_from_vllm_config(config: Any) -> PluginConfig: master_addr=None, enable_dp_attention=False, plugin_config=plugin_config, + speculative_config=atom_speculative_config, ) @@ -148,6 +194,7 @@ def _generate_atom_config_from_sglang_config(config: Any): from sglang.srt.server_args import ( get_global_server_args, PortArgs, + ZMQ_TCP_PORT_DELTA, ) from sglang.srt.configs.model_config import ModelConfig as SglangModelConfig from sglang.srt.configs.modelopt_config import ModelOptConfig @@ -212,7 +259,9 @@ def _generate_atom_config_from_sglang_config(config: Any): # sglang uses the atom parallel config sgl_parallel_config = ParallelConfig( data_parallel_size=atom_data_parallel_size, + data_parallel_size_local=atom_data_parallel_size, data_parallel_rank=atom_data_parallel_rank, + data_parallel_rank_local=atom_data_parallel_rank, ) # use sglang torch compile policy and cuda graph policy @@ -224,6 +273,26 @@ def _generate_atom_config_from_sglang_config(config: Any): cudagraph_mode=None, ) + sglang_dist_init_addr = server_args.dist_init_addr + # In single-node DP attention, synthesize the same TCP base address that + # SGLang uses for its DP-attention TCP port family. The primary purpose is + # to avoid calling PortArgs.init_new() again in ATOM plugin mode, because a + # second call would probe that fixed TCP range again and conflict with + # SGLang's existing allocation. In the current plugin path, this value + # should be treated as a compatibility/fallback hint rather than a + # guaranteed representation of the runtime default torch.distributed world + # rendezvous endpoint. + if ( + sglang_dist_init_addr is None + and server_args.enable_dp_attention + and server_args.nnodes == 1 + ): + sglang_dist_init_addr = f"127.0.0.1:{server_args.port + ZMQ_TCP_PORT_DELTA}" + + sglang_port_args = None + if sglang_dist_init_addr is None: + sglang_port_args = PortArgs.init_new(server_args) + plugin_config = PluginConfig( # common config model_config=sgl_model_config, @@ -237,8 +306,9 @@ def _generate_atom_config_from_sglang_config(config: Any): sglang_enable_torch_compile=server_args.enable_torch_compile, sglang_disable_cuda_graph=server_args.disable_cuda_graph, sglang_enable_dp_attention=server_args.enable_dp_attention, - sglang_dist_init_addr=server_args.dist_init_addr, - sglang_port_args=PortArgs.init_new(server_args), + sglang_aiter_rank_id=sglang_aiter_rank_id, + sglang_dist_init_addr=sglang_dist_init_addr, + sglang_port_args=sglang_port_args, ) # force max num batched tokens to 16K because sgl doesn't have diff --git a/atom/plugin/prepare.py b/atom/plugin/prepare.py index 297d9faa95..3c7b722d29 100644 --- a/atom/plugin/prepare.py +++ b/atom/plugin/prepare.py @@ -84,13 +84,7 @@ def prepare_model(config: Any, engine: str): apply_prepare_model_adaptations(atom_config, model_arch) - # Qwen3-Next and Qwen3.5 series models keep the upstream attention backend path. - if model_arch not in { - "Qwen3NextForCausalLM", - "Qwen3_5ForConditionalGeneration", - "Qwen3_5MoeForConditionalGeneration", - }: - register_ops_to_sglang(atom_config=atom_config) + register_ops_to_sglang(atom_config=atom_config) set_attn_cls() # init aiter dist for using aiter custom collective ops diff --git a/atom/plugin/register.py b/atom/plugin/register.py index 1af8811e50..043f14133a 100644 --- a/atom/plugin/register.py +++ b/atom/plugin/register.py @@ -40,21 +40,28 @@ def _register_custom_attention_to_sglang() -> None: sglang only accepts pre-registered backend names, so we reuse the "aiter" name to inject ATOMAttnBackendForSgl without modifying sglang source. """ + import sglang.srt.layers.attention.aiter_backend as sglang_aiter_backend + from sglang.srt.layers.attention.attention_registry import ( register_attention_backend, ) + from atom.plugin.sglang.attention_backend.full_attention.full_attention_backend import ( + ATOMAttnBackendForSgl, + ) # here register the custom attention backend with the name "aiter" # as sglang defines the fixed attention backend choices, which must be # in-tree logger.info("Register custom attention backend ATOMAttnBackendForSgl to SGLang") + # Speculative draft paths instantiate AiterAttnBackend directly inside + # AiterMultiStepDraftBackend, bypassing the attention registry. Rebind the + # module symbol as well so both registry lookup and direct construction use + # the plugin backend. + sglang_aiter_backend.AiterAttnBackend = ATOMAttnBackendForSgl + @register_attention_backend("aiter") def create_atom_backend(runner): - from atom.plugin.sglang.attention_backend.sgl_attn_backend import ( - ATOMAttnBackendForSgl, - ) - return ATOMAttnBackendForSgl(runner) @@ -94,6 +101,8 @@ def init_aiter_dist(config: Config) -> None: ) rank = config.plugin_config.rank + if getattr(config.plugin_config, "is_sglang", False): + rank = getattr(config.plugin_config, "sglang_aiter_rank_id", rank) tensor_parallel_size = config.tensor_parallel_size assert ( diff --git a/atom/plugin/sglang/attention_backend/attention_gdn.py b/atom/plugin/sglang/attention_backend/attention_gdn.py index 2403efaa9e..f29a43b1b7 100644 --- a/atom/plugin/sglang/attention_backend/attention_gdn.py +++ b/atom/plugin/sglang/attention_backend/attention_gdn.py @@ -164,7 +164,7 @@ def _build_gdn_metadata( def build( cls, forward_batch_or_metadata: Any ) -> Optional["SGLangGDNForwardContext"]: - from atom.plugin.sglang.models.base_model_wrapper import ( + from atom.plugin.sglang.runtime import ( SGLangForwardBatchMetadata, ) diff --git a/atom/plugin/sglang/attention_backend/full_attention/__init__.py b/atom/plugin/sglang/attention_backend/full_attention/__init__.py new file mode 100644 index 0000000000..2b9f6f2726 --- /dev/null +++ b/atom/plugin/sglang/attention_backend/full_attention/__init__.py @@ -0,0 +1,8 @@ +from .radix_attention import RadixAttention +from .full_attention_backend import ATOMAttnBackendForSgl, ForwardMetadata + +__all__ = [ + "RadixAttention", + "ATOMAttnBackendForSgl", + "ForwardMetadata", +] diff --git a/atom/plugin/sglang/attention_backend/sgl_attn_backend.py b/atom/plugin/sglang/attention_backend/full_attention/full_attention_backend.py similarity index 58% rename from atom/plugin/sglang/attention_backend/sgl_attn_backend.py rename to atom/plugin/sglang/attention_backend/full_attention/full_attention_backend.py index ad6723570d..2b58f6fed5 100644 --- a/atom/plugin/sglang/attention_backend/sgl_attn_backend.py +++ b/atom/plugin/sglang/attention_backend/full_attention/full_attention_backend.py @@ -1,10 +1,10 @@ from __future__ import annotations -# sglang-specific attention backend replacing sglang's built-in AiterAttnBackend. -# Shared by ALL models (DeepSeek, Qwen3, etc.) — handles KV cache writes, -# page-table fixup, pa_persistent_fwd decode path, and MLA prefill kernels. -# Sits at the lowest layer of the attention stack: sglang's RadixAttention -# delegates the actual kernel dispatch here. +# SGLang full-attention backend replacing sglang's built-in AiterAttnBackend. +# Shared by ALL full-attention models (DeepSeek, Qwen3, etc.) — handles KV +# cache writes, page-table fixup, pa_persistent_fwd decode path, and MLA +# prefill kernels. Sits at the lowest layer of the attention stack: +# sglang's RadixAttention delegates the actual kernel dispatch here. # # TODO: rewrite this file once sglang's attention flow is unified into ATOM's # attention layer — KV cache management and attention kernel dispatch will then @@ -20,7 +20,11 @@ import sglang.srt.layers.attention.aiter_backend as _sglang_aiter from sglang.srt.layers.attention.aiter_backend import AiterAttnBackend -from sglang.srt.layers.attention.utils import create_flashinfer_kv_indices_triton +from sglang.srt.layers.attention.utils import ( + create_flashinfer_kv_indices_triton, + launch_reshape_and_cache_flash, + pad_sequence_with_mask, +) from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode from sglang.srt.utils import get_bool_env_var @@ -35,13 +39,15 @@ dtypes, get_pa_metadata_info_v1, get_pa_metadata_v1, + mha_batch_prefill_func, pa_fwd_asm, pa_persistent_fwd, + paged_attention_ragged, ) except ImportError as e: raise ImportError( "Failed to import 'aiter', which provides AMD-specific attention kernels " - "required by sgl_attn_backend. Please ensure 'aiter' is installed and " + "required by full_attention_backend. Please ensure 'aiter' is installed and " f"available on your AMD system. Original import error: {e}" ) from e @@ -190,6 +196,7 @@ class ForwardMetadata: reduce_partial_map: Optional[torch.Tensor] = None fp8_prefill_kv_indices: Optional[torch.Tensor] = None num_kv_splits: Optional[int] = None + run_graph: Optional[bool] = True # PA metadata for pa_persistent_fwd (only used in decode mode, non-MLA) pa_metadata_qo_indptr: Optional[torch.Tensor] = None pa_metadata_pages_kv_indptr: Optional[torch.Tensor] = None @@ -212,8 +219,9 @@ def __init__( model_runner: ModelRunner, skip_prefill: bool = False, kv_indptr_buf: Optional[torch.Tensor] = None, + topk: int = 1, ): - super().__init__(model_runner, skip_prefill, kv_indptr_buf) + super().__init__(model_runner, skip_prefill, kv_indptr_buf, topk) mapping = getattr( model_runner.token_to_kv_pool, "full_attention_layer_id_mapping", None ) @@ -283,6 +291,10 @@ def init_forward_metadata(self, forward_batch: ForwardBatch): """Init auxiliary variables for triton attention backend.""" if forward_batch.forward_mode.is_decode_or_idle(): self._init_forward_metadata_decode(forward_batch) + elif self.use_mla and forward_batch.forward_mode.is_draft_extend(): + self._init_draft_extend_mla(forward_batch.batch_size, forward_batch) + elif self.use_mla and forward_batch.forward_mode.is_target_verify(): + self._init_target_verify_mla(forward_batch.batch_size, forward_batch) else: self._init_forward_metadata_extend(forward_batch) self._fixup_page_table(forward_batch) @@ -380,20 +392,11 @@ def _init_decode_mha(self, bs, kv_indptr, kv_indices, forward_batch): if seq_lens_cpu is None: seq_lens_cpu = forward_batch.seq_lens.cpu() - page_table_persistent = self.page_table - seq_lens_persistent = self.seq_lens - seq_lens_persistent.fill_(0) - page_table_persistent.fill_(0) - seq_lens_persistent[:bs].copy_(forward_batch.seq_lens, non_blocking=True) - max_seq_pages = ( - seq_lens_cpu.max().item() + self.page_size - 1 - ) // self.page_size + 1 - page_table = self.req_to_token[ - forward_batch.req_pool_indices[:, None], - self.strided_indices[:max_seq_pages][None, :], - ] - page_table_persistent[:bs, :max_seq_pages].copy_( - page_table // self.page_size, non_blocking=True + page_table, seq_lens = self._update_decode_page_table( + bs, + forward_batch.req_pool_indices, + forward_batch.seq_lens, + seq_lens_cpu=seq_lens_cpu, ) self.forward_metadata = ForwardMetadata( kv_indptr, @@ -402,8 +405,8 @@ def _init_decode_mha(self, bs, kv_indptr, kv_indices, forward_batch): None, 1, None, - page_table_persistent[:bs, :max_seq_pages], - seq_lens_persistent[:bs], + page_table, + seq_lens, ) self._build_pa_metadata_for_decode(bs, tp_q_head_num=self.num_head) else: @@ -429,6 +432,186 @@ def _init_forward_metadata_extend(self, forward_batch: ForwardBatch): else: self._init_extend_mha(bs, forward_batch) + def _init_draft_extend_mla(self, bs, forward_batch): + """Init MLA metadata for speculative draft_extend.""" + spec_info = forward_batch.spec_info + if spec_info is None: + raise RuntimeError("MLA draft_extend requires speculative metadata") + + kv_indices, kv_indptr, qo_indptr, _ = spec_info.generate_attn_arg_prefill( + forward_batch.req_pool_indices, + forward_batch.seq_lens, + forward_batch.seq_lens_sum, + self.req_to_token, + ) + + extend_seq_lens_cpu = forward_batch.extend_seq_lens_cpu + if extend_seq_lens_cpu is not None: + max_q_len = ( + int(extend_seq_lens_cpu.max().item()) + if isinstance(extend_seq_lens_cpu, torch.Tensor) + else max(extend_seq_lens_cpu) + ) + elif forward_batch.extend_seq_lens is not None: + max_q_len = int(forward_batch.extend_seq_lens.max().item()) + elif getattr(spec_info, "accept_length", None) is not None: + max_q_len = int(spec_info.accept_length.max().item()) + else: + raise RuntimeError("MLA draft_extend is missing extend sequence lengths") + + seq_lens_cpu = forward_batch.seq_lens_cpu + max_kv_len = ( + ( + int(seq_lens_cpu.max().item()) + if isinstance(seq_lens_cpu, torch.Tensor) + else max(seq_lens_cpu) + ) + if seq_lens_cpu is not None + else int(forward_batch.seq_lens.max().item()) + ) + + work_metadata = None + work_indptr = None + work_info_set = None + reduce_indptr = None + reduce_final_map = None + reduce_partial_map = None + num_kv_splits = None + + if _sglang_aiter._use_mla_ps_kernel: + ( + work_metadata, + work_indptr, + work_info_set, + reduce_indptr, + reduce_final_map, + reduce_partial_map, + ) = self.make_mla_decode_meta_data_buffer(max_q_len, bs) + num_kv_splits = self.max_split_per_batch + self.make_mla_meta_data( + qo_indptr, + kv_indptr, + self.kv_last_page_len[:bs], + work_metadata, + work_info_set, + work_indptr, + reduce_indptr, + reduce_final_map, + reduce_partial_map, + max_q_len, + fast_mode=_sglang_aiter.fast_mode, + max_split_per_batch=num_kv_splits, + intra_batch_mode=_sglang_aiter.intra_batch_mode, + ) + + self.forward_metadata = ForwardMetadata( + kv_indptr, + kv_indices, + qo_indptr, + self.kv_last_page_len[:bs], + max_q_len, + max_kv_len, + None, + None, + work_metadata=work_metadata, + work_info_set=work_info_set, + work_indptr=work_indptr, + reduce_indptr=reduce_indptr, + reduce_final_map=reduce_final_map, + reduce_partial_map=reduce_partial_map, + num_kv_splits=num_kv_splits, + run_graph=False, + ) + + def _init_target_verify_mla(self, bs, forward_batch): + """Init MLA metadata for speculative target_verify.""" + spec_info = forward_batch.spec_info + if spec_info is None: + raise RuntimeError("MLA target_verify requires speculative metadata") + + draft_num = spec_info.draft_token_num + kv_lens = forward_batch.seq_lens + draft_num + kv_lens_sum = forward_batch.seq_lens_sum + draft_num * bs + device = forward_batch.seq_lens.device + + qo_indptr = torch.arange( + 0, + (1 + bs) * draft_num, + step=draft_num, + dtype=torch.int32, + device=device, + ) + kv_indptr = self.kv_indptr + kv_indptr[1 : bs + 1] = torch.cumsum(kv_lens, dim=0) + kv_indptr = kv_indptr[: bs + 1] + kv_indices = torch.empty( + kv_lens_sum, + dtype=torch.int32, + device=device, + ) + create_flashinfer_kv_indices_triton[(bs,)]( + self.req_to_token, + forward_batch.req_pool_indices, + kv_lens, + kv_indptr, + None, + kv_indices, + self.req_to_token.stride(0), + ) + + work_metadata = None + work_indptr = None + work_info_set = None + reduce_indptr = None + reduce_final_map = None + reduce_partial_map = None + num_kv_splits = None + + if _sglang_aiter._use_mla_ps_kernel: + ( + work_metadata, + work_indptr, + work_info_set, + reduce_indptr, + reduce_final_map, + reduce_partial_map, + ) = self.make_mla_decode_meta_data_buffer(draft_num, bs) + num_kv_splits = self.max_split_per_batch + self.make_mla_meta_data( + qo_indptr, + kv_indptr, + self.kv_last_page_len[:bs], + work_metadata, + work_info_set, + work_indptr, + reduce_indptr, + reduce_final_map, + reduce_partial_map, + draft_num, + fast_mode=_sglang_aiter.fast_mode, + max_split_per_batch=num_kv_splits, + intra_batch_mode=_sglang_aiter.intra_batch_mode, + ) + + self.forward_metadata = ForwardMetadata( + kv_indptr, + kv_indices, + qo_indptr, + self.kv_last_page_len[:bs], + draft_num, + None, + None, + None, + work_metadata=work_metadata, + work_info_set=work_info_set, + work_indptr=work_indptr, + reduce_indptr=reduce_indptr, + reduce_final_map=reduce_final_map, + reduce_partial_map=reduce_partial_map, + num_kv_splits=num_kv_splits, + run_graph=False, + ) + def _init_extend_mla(self, bs, forward_batch): self.mla_indices_updater_prefill.update( forward_batch.req_pool_indices, @@ -519,8 +702,11 @@ def _init_extend_mha(self, bs, forward_batch): self.indices_updater_prefill.kv_indices, self.qo_indptr[: bs + 1], None, - self.indices_updater_prefill.max_q_len, - self.indices_updater_prefill.max_kv_len, + self._max_len( + forward_batch.extend_seq_lens_cpu, + forward_batch.extend_seq_lens, + ), + self._max_len(forward_batch.seq_lens_cpu, forward_batch.seq_lens), None, forward_batch.seq_lens, ) @@ -711,6 +897,11 @@ def init_cuda_graph_state( self.cuda_graph_kv_last_page_len = torch.ones( max_bs, dtype=torch.int, device=self.device ) + assert self.cuda_graph_kv_last_page_len.is_cuda, ( + "ATOMAttnBackendForSgl.init_cuda_graph_state created " + f"non-CUDA cuda_graph_kv_last_page_len on {self.cuda_graph_kv_last_page_len.device}, " + f"backend={type(self)}" + ) if kv_indices_buf is None: self.cuda_graph_kv_indices = torch.zeros( (max_bs * self.max_context_len), @@ -838,27 +1029,244 @@ def init_forward_metadata_capture_cuda_graph( forward_mode: ForwardMode, spec_info: Optional[SpecInput], ): - if not forward_mode.is_decode_or_idle(): - raise ValueError(f"Invalid mode: {forward_mode=}") + num_kv_splits = None + work_metadata = None + work_info_set = None + work_indptr = None + reduce_indptr = None + reduce_final_map = None + reduce_partial_map = None - if self.use_mla: - self._init_mla_cuda_graph_metadata(bs, req_pool_indices, seq_lens) - else: - page_table = self.page_table[:bs, :] - self.seq_lens[:bs].copy_(seq_lens, non_blocking=True) - seq_lens_persistent = self.seq_lens[:bs] - self.forward_metadata = ForwardMetadata( - None, - None, - None, + if forward_mode.is_decode_or_idle(): + if self.use_mla: + self._init_mla_cuda_graph_metadata(bs, req_pool_indices, seq_lens) + else: + kv_indptr = self.kv_indptr + kv_indptr[1 : bs + 1] = torch.cumsum(seq_lens, dim=0) + kv_indptr = kv_indptr[: bs + 1] + kv_indices = self.cuda_graph_kv_indices + create_flashinfer_kv_indices_triton[(bs,)]( + self.req_to_token, + req_pool_indices, + seq_lens, + kv_indptr, + None, + kv_indices, + self.req_to_token.stride(0), + ) + page_table, seq_lens_persistent = self._update_decode_page_table( + bs, + req_pool_indices, + seq_lens, + static_columns=True, + ) + self.forward_metadata = ForwardMetadata( + kv_indptr, + kv_indices, + None, + None, + 1, + None, + page_table, + seq_lens_persistent, + ) + if self.decode_using_pa_ps: + self._build_pa_metadata_for_decode(bs, tp_q_head_num=self.num_head) + elif forward_mode.is_target_verify(): + qo_indptr = self.qo_indptr[: bs + 1] + qo_indptr[: bs + 1] = torch.arange( + 0, + (1 + bs) * self.num_draft_tokens, + step=self.num_draft_tokens, + dtype=torch.int32, + device=self.device, + ) + kv_lens = seq_lens + self.num_draft_tokens if self.use_mla else seq_lens + kv_indptr = self.kv_indptr[: bs + 1] + kv_indptr[1 : bs + 1] = torch.cumsum(kv_lens, dim=0) + kv_indices = self.cuda_graph_kv_indices + create_flashinfer_kv_indices_triton[(bs,)]( + self.req_to_token, + req_pool_indices, + kv_lens, + kv_indptr, None, - 1, + kv_indices, + self.req_to_token.stride(0), + ) + kv_last_page_len = self.cuda_graph_kv_last_page_len[:bs] + max_q_len = self.num_draft_tokens + + if self.use_mla: + if _sglang_aiter._use_mla_ps_kernel: + num_kv_splits = self.max_split_per_batch + self.make_mla_meta_data( + qo_indptr, + kv_indptr, + kv_last_page_len, + self.work_metadata, + self.work_info_set, + self.work_indptr, + self.reduce_indptr, + self.reduce_final_map, + self.reduce_partial_map, + max_q_len, + fast_mode=_sglang_aiter.fast_mode, + max_split_per_batch=num_kv_splits, + intra_batch_mode=_sglang_aiter.intra_batch_mode, + ) + work_metadata = self.work_metadata + work_info_set = self.work_info_set + work_indptr = self.work_indptr + reduce_indptr = self.reduce_indptr + reduce_final_map = self.reduce_final_map + reduce_partial_map = self.reduce_partial_map + + self.forward_metadata = ForwardMetadata( + kv_indptr, + kv_indices, + qo_indptr, + kv_last_page_len, + max_q_len, + kv_indptr[-1].item(), + None, + None, + work_metadata=work_metadata, + work_info_set=work_info_set, + work_indptr=work_indptr, + reduce_indptr=reduce_indptr, + reduce_final_map=reduce_final_map, + reduce_partial_map=reduce_partial_map, + num_kv_splits=num_kv_splits, + ) + assert ( + self.forward_metadata.kv_last_page_len is None + or self.forward_metadata.kv_last_page_len.is_cuda + ), ( + "capture_cuda_graph TARGET_VERIFY produced non-CUDA kv_last_page_len: " + f"{self.forward_metadata.kv_last_page_len.device}, " + f"backend={type(self)}, metadata_backend={type(self.forward_metadata)}" + ) + else: + custom_mask = self.cuda_graph_custom_mask + assert spec_info is not None and spec_info.custom_mask is not None + custom_mask[: spec_info.custom_mask.shape[0]] = spec_info.custom_mask + seq_mask_len = max_q_len * (seq_lens + max_q_len) + mask_indptr = self.mask_indptr + mask_indptr[1 : bs + 1] = torch.cumsum(seq_mask_len[:bs], dim=0) + mask_indptr = mask_indptr[: bs + 1] + self.forward_metadata = ForwardMetadata( + kv_indptr, + kv_indices, + qo_indptr, + kv_last_page_len, + max_q_len, + kv_indptr[-1].item(), + None, + None, + custom_mask=custom_mask, + mask_indptr=mask_indptr, + max_extend_len=max_q_len, + ) + assert ( + self.forward_metadata.kv_last_page_len is None + or self.forward_metadata.kv_last_page_len.is_cuda + ), ( + "capture_cuda_graph TARGET_VERIFY(non-MLA) produced non-CUDA kv_last_page_len: " + f"{self.forward_metadata.kv_last_page_len.device}, " + f"backend={type(self)}" + ) + elif forward_mode.is_draft_extend(): + num_tokens_per_bs = self.speculative_num_steps + 1 + qo_indptr = self.qo_indptr[: bs + 1] + qo_indptr[: bs + 1] = torch.arange( + 0, + bs * num_tokens_per_bs + 1, + step=num_tokens_per_bs, + dtype=torch.int32, + device=self.device, + ) + kv_indptr = self.kv_indptr[: bs + 1] + kv_indptr[1 : bs + 1] = torch.cumsum(seq_lens, dim=0) + kv_indices = self.cuda_graph_kv_indices + create_flashinfer_kv_indices_triton[(bs,)]( + self.req_to_token, + req_pool_indices, + seq_lens, + kv_indptr, None, - page_table, - seq_lens_persistent, + kv_indices, + self.req_to_token.stride(0), ) - if self.decode_using_pa_ps: - self._build_pa_metadata_for_decode(bs, tp_q_head_num=self.num_head) + + if self.use_mla: + kv_last_page_len = self.cuda_graph_kv_last_page_len[:bs] + max_q_len = num_tokens_per_bs + if _sglang_aiter._use_mla_ps_kernel: + num_kv_splits = self.max_split_per_batch + self.make_mla_meta_data( + qo_indptr, + kv_indptr, + kv_last_page_len, + self.work_metadata, + self.work_info_set, + self.work_indptr, + self.reduce_indptr, + self.reduce_final_map, + self.reduce_partial_map, + max_q_len, + fast_mode=_sglang_aiter.fast_mode, + max_split_per_batch=num_kv_splits, + intra_batch_mode=_sglang_aiter.intra_batch_mode, + ) + work_metadata = self.work_metadata + work_info_set = self.work_info_set + work_indptr = self.work_indptr + reduce_indptr = self.reduce_indptr + reduce_final_map = self.reduce_final_map + reduce_partial_map = self.reduce_partial_map + + self.forward_metadata = ForwardMetadata( + kv_indptr, + kv_indices, + qo_indptr, + kv_last_page_len, + max_q_len, + kv_indptr[-1].item(), + None, + None, + work_metadata=work_metadata, + work_info_set=work_info_set, + work_indptr=work_indptr, + reduce_indptr=reduce_indptr, + reduce_final_map=reduce_final_map, + reduce_partial_map=reduce_partial_map, + num_kv_splits=num_kv_splits, + ) + assert ( + self.forward_metadata.kv_last_page_len is None + or self.forward_metadata.kv_last_page_len.is_cuda + ), ( + "capture_cuda_graph DRAFT_EXTEND produced non-CUDA kv_last_page_len: " + f"{self.forward_metadata.kv_last_page_len.device}, " + f"backend={type(self)}" + ) + else: + self.forward_metadata = ForwardMetadata( + kv_indptr, + kv_indices, + qo_indptr, + None, + num_tokens_per_bs, + None, + None, + None, + custom_mask=None, + mask_indptr=None, + max_extend_len=num_tokens_per_bs, + ) + else: + raise ValueError(f"Invalid mode: {forward_mode=}") def init_forward_metadata_replay_cuda_graph( self, @@ -872,40 +1280,331 @@ def init_forward_metadata_replay_cuda_graph( seq_lens_cpu: Optional[torch.Tensor], out_cache_loc: Optional[torch.Tensor] = None, ): - if not forward_mode.is_decode_or_idle(): - raise ValueError("Invalid forward mode") + num_kv_splits = None + work_metadata = None + work_info_set = None + work_indptr = None + reduce_indptr = None + reduce_final_map = None + reduce_partial_map = None - if self.use_mla: - self._init_mla_cuda_graph_metadata(bs, req_pool_indices, seq_lens) - else: - page_table_persistent = self.page_table - seq_lens_persistent = self.seq_lens - seq_lens_persistent.fill_(0) - page_table_persistent.fill_(0) - seq_lens_persistent[:bs].copy_(seq_lens, non_blocking=True) - max_seq_pages = ( - seq_lens_cpu.max().item() + self.page_size - 1 - ) // self.page_size + 1 - page_table = self.req_to_token[ - req_pool_indices[:, None], - self.strided_indices[:max_seq_pages][None, :], - ] - page_table_persistent[:bs, :max_seq_pages].copy_( - page_table // self.page_size, non_blocking=True - ) + if forward_mode.is_decode_or_idle(): + if self.use_mla: + self._init_mla_cuda_graph_metadata(bs, req_pool_indices, seq_lens) + else: + kv_indptr = self.kv_indptr + kv_indptr[1 : bs + 1] = torch.cumsum(seq_lens, dim=0) + kv_indptr = kv_indptr[: bs + 1] + kv_indices = self.cuda_graph_kv_indices + create_flashinfer_kv_indices_triton[(bs,)]( + self.req_to_token, + req_pool_indices, + seq_lens, + kv_indptr, + None, + kv_indices, + self.req_to_token.stride(0), + ) + page_table, seq_lens_persistent = self._update_decode_page_table( + bs, + req_pool_indices, + seq_lens, + seq_lens_cpu=seq_lens_cpu, + static_columns=True, + ) - self.forward_metadata = ForwardMetadata( - None, - None, - None, + self.forward_metadata = ForwardMetadata( + kv_indptr, + kv_indices, + None, + None, + 1, + None, + page_table, + seq_lens_persistent[:bs], + ) + if self.decode_using_pa_ps: + self._build_pa_metadata_for_decode(bs, tp_q_head_num=self.num_head) + elif forward_mode.is_target_verify(): + bs = len(req_pool_indices) + qo_indptr = self.qo_indptr[: bs + 1] + qo_indptr[: bs + 1] = torch.arange( + 0, + (1 + bs) * self.num_draft_tokens, + step=self.num_draft_tokens, + dtype=torch.int32, + device=self.device, + ) + kv_lens = seq_lens + self.num_draft_tokens if self.use_mla else seq_lens + kv_indptr = self.kv_indptr[: bs + 1] + kv_indptr[1 : bs + 1] = torch.cumsum(kv_lens, dim=0) + kv_indices = self.cuda_graph_kv_indices + create_flashinfer_kv_indices_triton[(bs,)]( + self.req_to_token, + req_pool_indices, + kv_lens, + kv_indptr, None, - 1, + kv_indices, + self.req_to_token.stride(0), + ) + kv_last_page_len = self.cuda_graph_kv_last_page_len[:bs] + max_q_len = self.num_draft_tokens + + if self.use_mla: + if _sglang_aiter._use_mla_ps_kernel: + num_kv_splits = self.max_split_per_batch + self.make_mla_meta_data( + qo_indptr, + kv_indptr, + kv_last_page_len, + self.work_metadata, + self.work_info_set, + self.work_indptr, + self.reduce_indptr, + self.reduce_final_map, + self.reduce_partial_map, + max_q_len, + fast_mode=_sglang_aiter.fast_mode, + max_split_per_batch=num_kv_splits, + intra_batch_mode=_sglang_aiter.intra_batch_mode, + ) + work_metadata = self.work_metadata + work_info_set = self.work_info_set + work_indptr = self.work_indptr + reduce_indptr = self.reduce_indptr + reduce_final_map = self.reduce_final_map + reduce_partial_map = self.reduce_partial_map + + self.forward_metadata = ForwardMetadata( + kv_indptr, + kv_indices, + qo_indptr, + kv_last_page_len, + max_q_len, + kv_indptr[-1].item(), + None, + None, + work_metadata=work_metadata, + work_info_set=work_info_set, + work_indptr=work_indptr, + reduce_indptr=reduce_indptr, + reduce_final_map=reduce_final_map, + reduce_partial_map=reduce_partial_map, + num_kv_splits=num_kv_splits, + ) + assert ( + self.forward_metadata.kv_last_page_len is None + or self.forward_metadata.kv_last_page_len.is_cuda + ), ( + "replay_cuda_graph TARGET_VERIFY produced non-CUDA kv_last_page_len: " + f"{self.forward_metadata.kv_last_page_len.device}, " + f"backend={type(self)}, metadata_backend={type(self.forward_metadata)}" + ) + else: + custom_mask = self.cuda_graph_custom_mask + assert spec_info is not None and spec_info.custom_mask is not None + custom_mask[: spec_info.custom_mask.shape[0]] = spec_info.custom_mask + seq_mask_len = max_q_len * (seq_lens + max_q_len) + mask_indptr = self.mask_indptr + mask_indptr[1 : bs + 1] = torch.cumsum(seq_mask_len[:bs], dim=0) + mask_indptr = mask_indptr[: bs + 1] + self.forward_metadata = ForwardMetadata( + kv_indptr, + kv_indices, + qo_indptr, + kv_last_page_len, + max_q_len, + kv_indptr[-1].item(), + None, + None, + custom_mask=custom_mask, + mask_indptr=mask_indptr, + max_extend_len=max_q_len, + ) + assert ( + self.forward_metadata.kv_last_page_len is None + or self.forward_metadata.kv_last_page_len.is_cuda + ), ( + "replay_cuda_graph TARGET_VERIFY(non-MLA) produced non-CUDA kv_last_page_len: " + f"{self.forward_metadata.kv_last_page_len.device}, " + f"backend={type(self)}" + ) + elif forward_mode.is_draft_extend(): + num_tokens_per_bs = self.speculative_num_steps + 1 + seq_lens = seq_lens[:bs] + accept_lens = spec_info.accept_length[:bs] + qo_indptr = self.qo_indptr[: bs + 1] + qo_indptr[1 : bs + 1] = torch.cumsum(accept_lens, dim=0) + kv_indptr = self.kv_indptr[: bs + 1] + kv_indptr[1 : bs + 1] = torch.cumsum(seq_lens, dim=0) + kv_indices = self.cuda_graph_kv_indices + create_flashinfer_kv_indices_triton[(bs,)]( + self.req_to_token, + req_pool_indices, + seq_lens, + kv_indptr, None, - page_table_persistent[:bs, :max_seq_pages], - seq_lens_persistent[:bs], + kv_indices, + self.req_to_token.stride(0), ) - if self.decode_using_pa_ps: - self._build_pa_metadata_for_decode(bs, tp_q_head_num=self.num_head) + + if self.use_mla: + kv_last_page_len = self.cuda_graph_kv_last_page_len[:bs] + max_q_len = num_tokens_per_bs + if _sglang_aiter._use_mla_ps_kernel: + num_kv_splits = self.max_split_per_batch + self.make_mla_meta_data( + qo_indptr, + kv_indptr, + kv_last_page_len, + self.work_metadata, + self.work_info_set, + self.work_indptr, + self.reduce_indptr, + self.reduce_final_map, + self.reduce_partial_map, + max_q_len, + fast_mode=_sglang_aiter.fast_mode, + max_split_per_batch=num_kv_splits, + intra_batch_mode=_sglang_aiter.intra_batch_mode, + ) + work_metadata = self.work_metadata + work_info_set = self.work_info_set + work_indptr = self.work_indptr + reduce_indptr = self.reduce_indptr + reduce_final_map = self.reduce_final_map + reduce_partial_map = self.reduce_partial_map + + self.forward_metadata = ForwardMetadata( + kv_indptr, + kv_indices, + qo_indptr, + kv_last_page_len, + max_q_len, + kv_indptr[-1].item(), + None, + None, + work_metadata=work_metadata, + work_info_set=work_info_set, + work_indptr=work_indptr, + reduce_indptr=reduce_indptr, + reduce_final_map=reduce_final_map, + reduce_partial_map=reduce_partial_map, + num_kv_splits=num_kv_splits, + ) + assert ( + self.forward_metadata.kv_last_page_len is None + or self.forward_metadata.kv_last_page_len.is_cuda + ), ( + "replay_cuda_graph DRAFT_EXTEND produced non-CUDA kv_last_page_len: " + f"{self.forward_metadata.kv_last_page_len.device}, " + f"backend={type(self)}" + ) + else: + self.forward_metadata = ForwardMetadata( + kv_indptr, + kv_indices, + qo_indptr, + None, + num_tokens_per_bs, + None, + None, + None, + custom_mask=None, + mask_indptr=None, + max_extend_len=num_tokens_per_bs, + ) + else: + raise ValueError(f"Invalid mode: {forward_mode=}") + + @staticmethod + def _max_len(cpu_values, device_values) -> int: + values = cpu_values if cpu_values is not None else device_values + if isinstance(values, torch.Tensor): + return int(values.max().item()) + return int(max(values)) + + def _update_decode_page_table( + self, + bs: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + seq_lens_cpu: Optional[torch.Tensor] = None, + static_columns: bool = False, + ): + page_table_persistent = self.page_table + seq_lens_persistent = self.seq_lens + seq_lens_persistent.fill_(0) + page_table_persistent.fill_(0) + seq_lens_persistent[:bs].copy_(seq_lens, non_blocking=True) + + if seq_lens_cpu is None: + max_seq_len = int(seq_lens.max().item()) + elif isinstance(seq_lens_cpu, torch.Tensor): + max_seq_len = int(seq_lens_cpu.max().item()) + else: + max_seq_len = int(max(seq_lens_cpu)) + + max_seq_pages = (max_seq_len + self.page_size - 1) // self.page_size + 1 + max_seq_pages = min(max_seq_pages, page_table_persistent.shape[1]) + page_table = self.req_to_token[ + req_pool_indices[:, None], + self.strided_indices[:max_seq_pages][None, :], + ] + page_table_persistent[:bs, :max_seq_pages].copy_( + page_table // self.page_size, non_blocking=True + ) + + if static_columns: + return page_table_persistent[:bs, :], seq_lens_persistent[:bs] + return page_table_persistent[:bs, :max_seq_pages], seq_lens_persistent[:bs] + + def _should_use_native_dense_mha(self, layer) -> bool: + sliding_window_size = getattr(layer, "sliding_window_size", None) + return ( + not self.use_mla + and not layer.is_cross_attention + and layer.head_dim == 256 + and layer.qk_head_dim == 256 + and layer.v_head_dim == 256 + and (sliding_window_size is None or sliding_window_size <= -1) + ) + + def _kv_descales(self, layer): + if self.kv_cache_dtype != dtypes.fp8: + return None, None + k_descale = layer.k_scale if layer.k_scale is not None else self.k_scale + v_descale = layer.v_scale if layer.v_scale is not None else self.v_scale + return k_descale, v_descale + + def _get_aiter_paged_ragged_kv_cache_dtype(self) -> str: + if self.kv_cache_dtype != dtypes.fp8: + return "auto" + return "fp8_e4m3" + + def _set_kv_buffer_native_dense(self, layer, cache_loc, k, v, forward_batch): + k_descale, v_descale = self._kv_descales(layer) + if self.kv_cache_dtype == dtypes.fp8: + k_cache, v_cache = forward_batch.token_to_kv_pool.get_kv_buffer( + layer.layer_id + ) + launch_reshape_and_cache_flash( + k.view(-1, layer.tp_k_head_num, layer.qk_head_dim), + v.view(-1, layer.tp_v_head_num, layer.v_head_dim), + k_cache.view( + -1, self.page_size, layer.tp_k_head_num, layer.qk_head_dim + ), + v_cache.view(-1, self.page_size, layer.tp_v_head_num, layer.v_head_dim), + cache_loc, + k_scale=k_descale, + v_scale=v_descale, + ) + return + + forward_batch.token_to_kv_pool.set_kv_buffer( + layer, cache_loc, k, v, k_descale, v_descale + ) def set_kv_buffer_with_layout_shuffle( self, @@ -947,11 +1646,16 @@ def forward_extend(self, q, k, v, layer, forward_batch, save_kv_cache=True): if not layer.is_cross_attention else forward_batch.encoder_out_cache_loc ) + use_native_dense_mha = self._should_use_native_dense_mha(layer) if k is not None: assert v is not None if save_kv_cache: - if self.use_mla: + if use_native_dense_mha: + self._set_kv_buffer_native_dense( + layer, cache_loc, k, v, forward_batch + ) + elif self.use_mla: forward_batch.token_to_kv_pool.set_kv_buffer(layer, cache_loc, k, v) else: k_buffer, v_buffer = forward_batch.token_to_kv_pool.get_kv_buffer( @@ -970,9 +1674,44 @@ def forward_extend(self, q, k, v, layer, forward_batch, save_kv_cache=True): if self.use_mla: return self._forward_extend_mla(q, k, v, layer, forward_batch) + if use_native_dense_mha: + return self._forward_extend_native_dense_mha(q, layer, forward_batch) else: return self._forward_extend_mha(q, k, v, layer, forward_batch) + def _forward_extend_native_dense_mha(self, q, layer, forward_batch): + k_cache, v_cache = forward_batch.token_to_kv_pool.get_kv_buffer(layer.layer_id) + q_descale, k_descale, v_descale = None, None, None + + if self.kv_cache_dtype == dtypes.fp8: + q = q.to(dtypes.fp8) + q_descale = layer.k_scale if layer.k_scale is not None else self.k_scale + k_descale = layer.k_scale if layer.k_scale is not None else self.k_scale + v_descale = layer.v_scale if layer.v_scale is not None else self.v_scale + + bs0 = forward_batch.batch_size + 1 + o = mha_batch_prefill_func( + q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim), + k_cache, + v_cache, + self.forward_metadata.qo_indptr[:bs0], + self.forward_metadata.kv_indptr[:bs0], + self.forward_metadata.kv_indices, + self.forward_metadata.max_q_len, + self.forward_metadata.max_kv_len, + causal=True, + logits_soft_cap=layer.logit_cap, + alibi_slopes=None, + return_lse=False, + return_attn_probs=False, + window_size=(-1, -1), + sink_ptr=None, + q_descale=q_descale, + k_descale=k_descale, + v_descale=v_descale, + ) + return o.view(-1, layer.tp_q_head_num * layer.head_dim) + def _forward_extend_mha(self, q, k, v, layer, forward_batch): """Non-MLA extend path: standard MHA with flash_attn_varlen_func.""" seqlens_in_batch = forward_batch.seq_lens @@ -1019,6 +1758,7 @@ def _forward_extend_mla(self, q, k, v, layer, forward_batch): layer, K_Buffer, qo_indptr, + forward_batch, ) if not forward_batch.forward_mode.is_extend(): raise ValueError( @@ -1356,14 +2096,51 @@ def _call_mla_decode_fwd(self, q, k_buffer, o, layer): num_kv_splits=md.num_kv_splits, ) - def _forward_extend_mla_speculative(self, q, layer, K_Buffer, qo_indptr): + def _forward_extend_mla_speculative( + self, q, layer, K_Buffer, qo_indptr, forward_batch + ): """MLA speculative path (target_verify / draft_extend).""" - o = q.new_empty( - (q.shape[0], layer.tp_q_head_num, layer.v_head_dim), - dtype=self.input_dtype, + md = self.forward_metadata + + if forward_batch.forward_mode.is_target_verify(): + o = q.new_empty( + (q.shape[0], layer.tp_q_head_num, layer.v_head_dim), + dtype=self.input_dtype, + ) + self._call_mla_decode_fwd(q, K_Buffer, o, layer) + return o + + if forward_batch.forward_mode.is_draft_extend(): + if md.run_graph is not True: + bs, q_pad, _ = pad_sequence_with_mask( + q.view(q.shape[0], -1), + qo_indptr[:-1], + forward_batch.extend_seq_lens, + md.max_q_len, + ) + o = q.new_empty( + (bs * md.max_q_len, layer.tp_q_head_num, layer.v_head_dim), + dtype=self.input_dtype, + ) + self._call_mla_decode_fwd( + q_pad.view(-1, layer.tp_q_head_num, layer.qk_head_dim), + K_Buffer, + o, + layer, + ) + total_valid_q = int(qo_indptr[-1].item()) + return o[:total_valid_q] + + o = q.new_empty( + (q.shape[0], layer.tp_q_head_num, layer.v_head_dim), + dtype=self.input_dtype, + ) + self._call_mla_decode_fwd(q, K_Buffer, o, layer) + return o + + raise ValueError( + f"Invalid forward mode for MLA speculative path: {forward_batch.forward_mode=}" ) - self._call_mla_decode_fwd(q, K_Buffer, o, layer) - return o def forward_decode( self, @@ -1401,6 +2178,14 @@ def forward_decode( return o # Non-MLA decode paths + use_native_dense_mha = self._should_use_native_dense_mha(layer) + if use_native_dense_mha: + if save_kv_cache: + self._set_kv_buffer_native_dense( + layer, forward_batch.out_cache_loc, k, v, forward_batch + ) + return self._forward_decode_native_dense_mha(q, layer, forward_batch) + o = q.new_empty((batch_size, layer.tp_q_head_num, head_dim_out)) if save_kv_cache: @@ -1476,3 +2261,31 @@ def forward_decode( ) return o.view(-1, layer.tp_q_head_num * head_dim_out) + + def _forward_decode_native_dense_mha(self, q, layer, forward_batch): + k_cache, v_cache = forward_batch.token_to_kv_pool.get_kv_buffer(layer.layer_id) + aiter_kv_str = self._get_aiter_paged_ragged_kv_cache_dtype() + + o = torch.empty_like(q, dtype=self.input_dtype) + paged_attention_ragged( + o.view(-1, layer.tp_q_head_num, layer.head_dim), + self.workspace_buffer, + q.view(-1, layer.tp_q_head_num, layer.head_dim), + k_cache.view(-1, 1, layer.tp_k_head_num, layer.head_dim), + v_cache.view(-1, 1, layer.tp_v_head_num, layer.v_head_dim), + layer.scaling, + self.forward_metadata.kv_indptr, + self.forward_metadata.kv_indices, + self.kv_last_page_len, + 1, + self.max_num_partitions, + None, + aiter_kv_str, + "NHD", + layer.logit_cap, + self.k_scale, + self.v_scale, + None, + getattr(_sglang_aiter, "_AITER_PARTITION_SIZE_ROCM", 256), + ) + return o diff --git a/atom/plugin/sglang/attention_backend/radix_attention.py b/atom/plugin/sglang/attention_backend/full_attention/radix_attention.py similarity index 93% rename from atom/plugin/sglang/attention_backend/radix_attention.py rename to atom/plugin/sglang/attention_backend/full_attention/radix_attention.py index ba46eb4c79..613458ee37 100644 --- a/atom/plugin/sglang/attention_backend/radix_attention.py +++ b/atom/plugin/sglang/attention_backend/full_attention/radix_attention.py @@ -88,10 +88,18 @@ def __init__( self.attn.k_scale = atom_parameter( torch.tensor([1.0], dtype=torch.float32, device="cuda") ) + elif not self.attn.k_scale.is_cuda: + self.attn.k_scale = atom_parameter( + self.attn.k_scale.detach().to(device="cuda") + ) if self.attn.v_scale is None: self.attn.v_scale = atom_parameter( torch.tensor([1.0], dtype=torch.float32, device="cuda") ) + elif not self.attn.v_scale.is_cuda: + self.attn.v_scale = atom_parameter( + self.attn.v_scale.detach().to(device="cuda") + ) # Some SGLang attention backends consume the host-side float scales # directly. Keep them in sync with the device-side defaults so the # plugin path works even when checkpoint loading never populates them. @@ -121,7 +129,7 @@ def forward_impl_plugin_mode( # for sglang, forward_batch is required forward_batch = kwargs.get("forward_batch", None) if forward_batch is None: - from atom.plugin.sglang.models.base_model_wrapper import ( + from atom.plugin.sglang.runtime import ( get_current_forward_batch, ) diff --git a/atom/plugin/sglang/attention_backend/sgl_attention_mla.py b/atom/plugin/sglang/attention_backend/sgl_attention_mla.py deleted file mode 100644 index d0c2b3b07a..0000000000 --- a/atom/plugin/sglang/attention_backend/sgl_attention_mla.py +++ /dev/null @@ -1,1187 +0,0 @@ -"""Sglang-specific MLA forward and weight processing for DeepseekV2/V3. - -DeepSeek MLA (Multi-Latent Attention) forward logic for sglang plugin mode: -absorbed BMM computation, MHA/MLA path dispatch (prefill -> MHA, decode -> MLA), -kv_b_proj weight splitting (w_kc/w_vc), and monkey-patch setup via -setup_deepseek_for_sglang(). - -This module is lazily imported from base_model_wrapper.py only when running in -sglang plugin mode (``is_sglang() == True``). Keeping all sglang-dependent -imports here avoids crashing when sglang is not installed. - -TODO: rewrite this file once sglang's attention flow is unified into ATOM's -attention layer — the MLA absorbed path and MHA dispatch will then be handled -natively by ATOM's attention ops, making this sglang-specific module unnecessary. -""" - -from __future__ import annotations - -from typing import TYPE_CHECKING, Any, NamedTuple, Optional - -import torch -from aiter import dtypes -from aiter.dist.parallel_state import get_tensor_model_parallel_world_size, get_tp_group -from aiter.ops.triton.fused_fp8_quant import fused_rms_fp8_group_quant -from atom.model_ops.base_attention import Attention -from atom.model_ops.attention_mla import ( - dynamic_per_batched_tensor_quant, - fused_qk_rope_concat_and_cache_mla, -) -from atom.models.utils import maybe_prefix - -# sglang imports -from sglang.srt.layers.communicator import AttentionInputs, get_attn_tp_context -from sglang.srt.layers.attention.nsa.utils import nsa_use_prefill_cp -from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode -from sglang.srt.models.deepseek_common.utils import ( - _use_aiter_gfx95, - _is_hip, - _is_cpu, - _is_cpu_amx_available, - _is_cuda, - _is_fp8_fnuz, - _is_npu, - awq_dequantize_func, -) -from sglang.srt.layers.quantization.rocm_mxfp4_utils import ( - batched_gemm_afp4wfp4_pre_quant, -) -from aiter.ops.triton.batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant import ( - batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant, -) -from sglang.srt.layers.quantization.fp8_kernel import ( - per_tensor_quant_mla_fp8, - per_token_group_quant_mla_deep_gemm_masked_fp8, -) -from sglang.srt.utils import bind_or_assign, get_bool_env_var - -if TYPE_CHECKING: - from atom.models.deepseek_v2 import DeepseekV2MLAAttention - - -# bmm_fp8 custom-op wrapper (adapted from sglang forward_mla.py) -if _is_cuda: - from sgl_kernel import bmm_fp8 as _raw_bmm_fp8 - from sglang.srt.utils.custom_op import register_custom_op - - @register_custom_op(mutates_args=["out"]) - def _bmm_fp8_op( - A: torch.Tensor, - B: torch.Tensor, - out: torch.Tensor, - A_scale: torch.Tensor, - B_scale: torch.Tensor, - ) -> None: - _raw_bmm_fp8(A, B, A_scale, B_scale, out.dtype, out) - - def bmm_fp8(A, B, A_scale, B_scale, dtype, out=None): - if out is None: - out = torch.empty( - (A.shape[0], A.shape[1], B.shape[2]), - device=A.device, - dtype=dtype, - ) - _bmm_fp8_op(A, B, out, A_scale, B_scale) - return out - -else: - - def bmm_fp8(A, B, A_scale, B_scale, dtype, out=None): - raise RuntimeError("bmm_fp8 requires CUDA (sgl_kernel)") - - -# NamedTuple for prepare → core data flow -class SglPrepareResult(NamedTuple): - q_pe: torch.Tensor - k_pe: torch.Tensor - q_nope_out: torch.Tensor - k_nope: torch.Tensor - forward_batch: Any - zero_allocator: Any - positions: torch.Tensor - topk_indices: Optional[torch.Tensor] - llama_4_scaling: Optional[Any] - - -class SglMhaPrepareResult(NamedTuple): - q: torch.Tensor - k: torch.Tensor - v: torch.Tensor - forward_batch: Any - - -def _unwrap_linear_output(output: Any) -> torch.Tensor: - """Normalize ATOM/public-SGLang linear outputs to a tensor.""" - if isinstance(output, tuple): - return output[0] - return output - - -def _fuse_qk_rmsnorm_and_q_quant( - attn: DeepseekV2MLAAttention, - q: torch.Tensor, - k_nope: torch.Tensor, - *, - output_unquantized_q: bool = False, -) -> tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor], torch.Tensor]: - """Fuse q/k RMSNorm and q quant using ATOM's DeepSeek-V2 path.""" - from atom.models.deepseek_v2 import _fuse_rmsnorm_quant - - (q_quantized, q_scale), q_normed, k_nope_normed, _ = _fuse_rmsnorm_quant( - q, - attn.q_a_layernorm.weight, - attn.q_a_layernorm.eps, - k_nope, - attn.kv_a_layernorm.weight, - attn.kv_a_layernorm.eps, - None, - dtype_quant=attn.quant_dtype, - shuffle=False, - scale_shuffle_padding=False, - group_size=128, - output_unquantized_inp1=output_unquantized_q, - transpose_scale=True, - ) - return q_quantized, q_scale, q_normed, k_nope_normed - - -def _fuse_qk_rmsnorm( - attn: DeepseekV2MLAAttention, - q: torch.Tensor, - k_nope: torch.Tensor, -) -> tuple[torch.Tensor, torch.Tensor]: - """Fuse q/k RMSNorm without quantizing q.""" - from atom.models.deepseek_v2 import _fused_qk_rmsnorm - - return _fused_qk_rmsnorm( - q, - attn.q_a_layernorm.weight, - attn.q_a_layernorm.eps, - k_nope, - attn.kv_a_layernorm.weight, - attn.kv_a_layernorm.eps, - ) - - -def _prepare_weight_for_bmm( - weight: torch.Tensor, in_dim: int, out_dim: int -) -> torch.Tensor: - """Normalize absorbed weight layout for torch.bmm fallback.""" - if weight.shape[1] == in_dim and weight.shape[2] == out_dim: - return weight - if weight.shape[1] == out_dim and weight.shape[2] == in_dim: - return weight.transpose(-2, -1) - raise RuntimeError( - "Unexpected absorbed weight shape for bmm fallback: " - f"{tuple(weight.shape)} with in_dim={in_dim}, out_dim={out_dim}" - ) - - -# Init helpers -def init_sgl_attrs( - attn: DeepseekV2MLAAttention, - config, - kv_cache_dtype: str = "bf16", -) -> None: - """Initialise sglang-only attributes on DeepseekV2MLAAttention.""" - from sglang.srt.configs.model_config import is_deepseek_nsa - - attn.use_nsa = is_deepseek_nsa(config) - attn.use_deep_gemm_bmm = False - attn.alt_stream = None - attn.kv_cache_dtype = kv_cache_dtype - attn.use_fused_qk_rope_concat_and_cache_mla = _use_aiter_gfx95 - attn.current_sgl_plugin_attn_path = None - attn.w_kc, attn.w_vc = None, None - attn.w_scale = None - attn.w_scale_k = None - attn.w_scale_v = None - attn.attn_mha = Attention( - num_heads=attn.num_local_heads, - head_dim=attn.qk_head_dim, - scale=attn.scaling, - num_kv_heads=attn.num_local_heads, - kv_cache_dtype=kv_cache_dtype, - layer_num=attn.layer_num, - use_mla=False, - v_head_dim=attn.v_head_dim, - prefix=maybe_prefix(attn.prefix, "attn_mha"), - ) - if hasattr(attn.attn_mha, "attn"): - attn.attn_mha.attn.kv_b_proj = None - - -# Absorbed batched-matmul (shared by prepare and core) -def mla_absorbed_bmm( - attn: DeepseekV2MLAAttention, - inp: torch.Tensor, - weight: torch.Tensor, - weight_scale: Optional[torch.Tensor], - weight_scale_k: Optional[torch.Tensor], - out_dim: int, -) -> torch.Tensor: - """Batched matmul for MLA absorbed weights (w_kc / w_vc). - - Handles deep_gemm, mxfp4, fp8-triton, fp8-cublas, and bf16 fallback paths. - inp: (num_tokens, num_heads, in_dim) — token-major - Returns: (num_tokens, num_heads, out_dim) — token-major - """ - effective_weight_scale = ( - weight_scale_k if weight_scale_k is not None else weight_scale - ) - - if attn.use_deep_gemm_bmm: - from sglang.srt.layers import deep_gemm_wrapper - - val, scale, masked_m, expected_m, aligned_m = ( - per_token_group_quant_mla_deep_gemm_masked_fp8(inp.transpose(0, 1)) - ) - out = inp.new_empty((attn.num_local_heads, aligned_m, out_dim)) - deep_gemm_wrapper.grouped_gemm_nt_f8f8bf16_masked( - (val, scale), - (weight, weight_scale_k), - out, - masked_m, - expected_m, - ) - return out[:, :expected_m, :].transpose(0, 1) - - if _is_hip: - if _use_aiter_gfx95 and weight.dtype == torch.uint8: - x = inp.transpose(0, 1) - out = torch.empty( - x.shape[0], - x.shape[1], - weight.shape[2], - device=x.device, - dtype=torch.bfloat16, - ) - batched_gemm_afp4wfp4_pre_quant( - x, - weight.transpose(-2, -1), - weight_scale_k.transpose(-2, -1), - torch.bfloat16, - out, - ) - return out.transpose(0, 1) - - if (_use_aiter_gfx95 and weight.dtype == torch.float8_e4m3fn) or ( - get_is_capture_mode() and weight.dtype == torch.float8_e4m3fnuz - ): - x = inp.transpose(0, 1) - out = ( - batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant( - X=x, - WQ=weight, - w_scale=effective_weight_scale, - group_size=128, - YQ=None, - transpose_bm=True, - transpose_bm_in=False, - dtype=torch.bfloat16, - ) - ) - return out - - w_bf16 = _prepare_weight_for_bmm(weight, inp.shape[-1], out_dim).to( - torch.bfloat16 - ) - if effective_weight_scale is not None: - w_bf16 = w_bf16 * effective_weight_scale - out = torch.bmm( - inp.to(torch.bfloat16).transpose(0, 1), - w_bf16, - ) - return out.transpose(0, 1) - - # CUDA fp8 path - if weight.dtype == torch.float8_e4m3fn: - val, scale = per_tensor_quant_mla_fp8( - inp.transpose(0, 1), - torch.zeros((1,), dtype=torch.float32, device=inp.device), - ) - out = bmm_fp8(val, weight, scale, effective_weight_scale, torch.bfloat16) - return out.transpose(0, 1) - - # bf16 fallback - return torch.bmm(inp.transpose(0, 1), weight).transpose(0, 1) - - -def mla_v_up_proj( - attn: DeepseekV2MLAAttention, - inp: torch.Tensor, - weight: torch.Tensor, - weight_scale: Optional[torch.Tensor], - weight_scale_k: Optional[torch.Tensor], - out_dim: int, -) -> torch.Tensor: - """Project MLA decode output to a flat o_proj input.""" - effective_weight_scale = ( - weight_scale_k if weight_scale_k is not None else weight_scale - ) - if _is_hip and ( - (_use_aiter_gfx95 and weight.dtype == torch.float8_e4m3fn) - or (get_is_capture_mode() and weight.dtype == torch.float8_e4m3fnuz) - ): - x = inp.transpose(0, 1) - out = torch.empty( - (inp.shape[0], attn.num_local_heads * out_dim), - device=inp.device, - dtype=torch.bfloat16, - ) - out_3d = out.view(inp.shape[0], attn.num_local_heads, out_dim) - batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant( - X=x, - WQ=weight, - w_scale=effective_weight_scale, - group_size=128, - YQ=out_3d, - transpose_bm=True, - transpose_bm_in=False, - dtype=torch.bfloat16, - ) - return out - - return mla_absorbed_bmm( - attn, inp, weight, weight_scale, weight_scale_k, out_dim - ).flatten(1, 2) - - -# Forward: prepare → core -def forward_sgl_prepare( - attn: DeepseekV2MLAAttention, - positions: torch.Tensor, - hidden_states: torch.Tensor, - **model_kwargs, -) -> SglPrepareResult: - """Prepare QKV for sglang MLA attention (adapted from sglang forward_absorb_prepare).""" - hidden_states_scale = None - if isinstance(hidden_states, tuple): - hidden_states, hidden_states_scale = hidden_states - - forward_batch = model_kwargs.get("forward_batch", None) - zero_allocator = model_kwargs.get("zero_allocator", None) - llama_4_scaling = model_kwargs.get("llama_4_scaling", None) - q_lora = None - topk_indices = None - - if attn.q_lora_rank is not None: - q, latent_cache = ( - get_attn_tp_context() - .fetch_qkv_latent() - .split( - [attn.q_lora_rank, attn.kv_lora_rank + attn.qk_rope_head_dim], - dim=-1, - ) - ) - - if ( - q.shape[0] != positions.shape[0] - and get_tensor_model_parallel_world_size() > 1 - ): - qkv_lora = torch.cat([q, latent_cache], dim=-1) - qkv_lora = get_tp_group().all_gather(qkv_lora, dim=0) - if qkv_lora.shape[0] < positions.shape[0]: - raise RuntimeError( - f"qkv_lora gather mismatch: got {qkv_lora.shape[0]}, " - f"expected {positions.shape[0]}" - ) - qkv_lora = qkv_lora[: positions.shape[0]] - q, latent_cache = torch.split( - qkv_lora, - [attn.q_lora_rank, attn.kv_lora_rank + attn.qk_rope_head_dim], - dim=-1, - ) - - k_nope = latent_cache[..., : attn.kv_lora_rank] - q_scale = None - - # Reuse native ATOM gating for q/k RMSNorm fusion. Quant fusion is used - # when DeepSeek enables qknorm-quant; otherwise keep the non-quant fused - # path aligned with native ATOM before falling back to plain layernorm. - if getattr(attn, "fuse_qknorm_quant", False): - q, q_scale, q_lora, k_nope = _fuse_qk_rmsnorm_and_q_quant( - attn, - q, - k_nope, - output_unquantized_q=attn.use_nsa, - ) - elif getattr(attn, "fuse_qknorm", False): - q, k_nope = _fuse_qk_rmsnorm(attn, q, k_nope) - # Otherwise keep the original overlap path for unfused qk norm. - elif attn.alt_stream is not None and get_is_capture_mode(): - current_stream = torch.cuda.current_stream() - attn.alt_stream.wait_stream(current_stream) - q = attn.q_a_layernorm(q) - with torch.cuda.stream(attn.alt_stream): - k_nope = attn.kv_a_layernorm(k_nope) - current_stream.wait_stream(attn.alt_stream) - else: - q = attn.q_a_layernorm(q) - k_nope = attn.kv_a_layernorm(k_nope) - - if attn.use_nsa: - if q_lora is None: - q_lora = q - - # overlap q_b_proj and indexer during decode - if ( - attn.alt_stream is not None - and get_is_capture_mode() - and forward_batch.forward_mode.is_decode_or_idle() - and q_lora is not None - ): - current_stream = torch.cuda.current_stream() - attn.alt_stream.wait_stream(current_stream) - with torch.cuda.stream(attn.alt_stream): - k_nope = k_nope.unsqueeze(1) - q = _unwrap_linear_output( - attn.q_b_proj(q, q_scale) - if q_scale is not None - else attn.q_b_proj(q) - ).view(-1, attn.num_local_heads, attn.qk_head_dim) - topk_indices = attn.indexer( - x=hidden_states, - q_lora=q_lora, - positions=positions, - forward_batch=forward_batch, - layer_id=attn.layer_num, - ) - current_stream.wait_stream(attn.alt_stream) - else: - k_nope = k_nope.unsqueeze(1) - q = _unwrap_linear_output( - attn.q_b_proj(q, q_scale) if q_scale is not None else attn.q_b_proj(q) - ).view(-1, attn.num_local_heads, attn.qk_head_dim) - if q_lora is not None: - topk_indices = attn.indexer( - x=hidden_states, - q_lora=q_lora, - positions=positions, - forward_batch=forward_batch, - layer_id=attn.layer_num, - ) - else: - q = _unwrap_linear_output(attn.q_proj(hidden_states)).view( - -1, attn.num_local_heads, attn.qk_head_dim - ) - latent_cache = _unwrap_linear_output(attn.kv_a_proj_with_mqa(hidden_states)) - k_nope = latent_cache[..., : attn.kv_lora_rank] - k_nope = attn.kv_a_layernorm(k_nope).unsqueeze(1) - - q_nope, q_pe = q.split([attn.qk_nope_head_dim, attn.qk_rope_head_dim], dim=-1) - k_pe = latent_cache[..., attn.kv_lora_rank :].unsqueeze(1) - - q_nope_out = mla_absorbed_bmm( - attn, q_nope, attn.w_kc, attn.w_scale, attn.w_scale_k, attn.kv_lora_rank - ) - - if attn.rotary_emb is not None and not attn.use_fused_qk_rope_concat_and_cache_mla: - q_pe, k_pe = attn.rotary_emb(positions, q_pe, k_pe) - - if nsa_use_prefill_cp(forward_batch): - k_nope, k_pe = attn.rebuild_cp_kv_cache( - latent_cache, forward_batch, k_nope, k_pe - ) - - return SglPrepareResult( - q_pe=q_pe, - k_pe=k_pe, - q_nope_out=q_nope_out, - k_nope=k_nope, - forward_batch=forward_batch, - zero_allocator=zero_allocator, - positions=positions, - topk_indices=topk_indices, - llama_4_scaling=llama_4_scaling, - ) - - -def forward_sgl_core( - attn: DeepseekV2MLAAttention, - prepared: SglPrepareResult, -) -> torch.Tensor: - """Core MLA attention computation for sglang (adapted from sglang forward_absorb_core).""" - save_kv_cache = True - - if attn.use_fused_qk_rope_concat_and_cache_mla: - mla_attn = _get_sglang_radix_attn(attn.mla_attn) - kv_cache = prepared.forward_batch.token_to_kv_pool.get_key_buffer( - mla_attn.layer_id - ) - q_out_dtype = ( - dtypes.fp8 - if attn.kv_cache_dtype == "fp8_e4m3" - else prepared.q_nope_out.dtype - ) - q = torch.empty( - ( - prepared.q_nope_out.shape[0], - attn.num_local_heads, - attn.kv_lora_rank + attn.qk_rope_head_dim, - ), - dtype=q_out_dtype, - device=prepared.q_nope_out.device, - ) - - fused_qk_rope_concat_and_cache_mla( - prepared.q_nope_out, - prepared.q_pe, - prepared.k_nope, - prepared.k_pe, - kv_cache, - q, - prepared.forward_batch.out_cache_loc, - mla_attn.k_scale, - mla_attn.k_scale, - prepared.positions, - attn.rotary_emb.cos_cache, - attn.rotary_emb.sin_cache, - is_neox=attn.rotary_emb.is_neox_style, - is_nope_first=True, - ) - # Decode/speculative MLA consumes q plus packed MLA cache directly. - k = None - v = None - save_kv_cache = False - else: - q = torch.cat([prepared.q_nope_out, prepared.q_pe], dim=-1) - k = torch.cat([prepared.k_nope, prepared.k_pe], dim=-1) - v = prepared.k_nope - - if prepared.llama_4_scaling is not None: - q = q * prepared.llama_4_scaling - - extra_kwargs = {} - if prepared.topk_indices is not None: - extra_kwargs["topk_indices"] = prepared.topk_indices - - attn_output = attn.mla_attn( - q, - k, - v, - forward_batch=prepared.forward_batch, - save_kv_cache=save_kv_cache, - **extra_kwargs, - ) - attn_output = attn_output.view(-1, attn.num_local_heads, attn.kv_lora_rank) - - # up-proj by w_vc - attn_bmm_output = mla_v_up_proj( - attn, attn_output, attn.w_vc, attn.w_scale, attn.w_scale_v, attn.v_head_dim - ) - - return attn.o_proj(attn_bmm_output) - - -def _dispatch_sgl_plugin_attn_path(forward_batch) -> str: - """Decide the attention algorithm for this batch based on forward_mode. - - Returns "mha" for extend/prefill (uses standard Q×K×V with flash_attn) - or "mla" for decode (uses absorbed weights + mla_decode_fwd). - - This is the per-batch *routing* decision, distinct from - ``_can_run_sgl_mha_now`` which is a *capability* gate checking whether - the model configuration supports the MHA path at all. - """ - if forward_batch.forward_mode.is_extend_without_speculative(): - return "mha" - return "mla" - - -def forward_sgl_plugin_mode_mla( - attn: DeepseekV2MLAAttention, - positions: torch.Tensor, - hidden_states: torch.Tensor, - **model_kwargs, -) -> torch.Tensor: - prepared = forward_sgl_prepare(attn, positions, hidden_states, **model_kwargs) - return forward_sgl_core(attn, prepared) - - -def _get_sglang_radix_attn(attn_module): - return attn_module.attn if hasattr(attn_module, "attn") else attn_module - - -def _set_mla_kv_buffer_for_mha( - attn: DeepseekV2MLAAttention, - kv_a: torch.Tensor, - k_pe: torch.Tensor, - forward_batch, -) -> None: - attn_mha = _get_sglang_radix_attn(attn.attn_mha) - cache_k = torch.cat([kv_a.unsqueeze(1), k_pe], dim=-1) - forward_batch.token_to_kv_pool.set_kv_buffer( - attn_mha, - forward_batch.out_cache_loc, - cache_k, - cache_k, - ) - - -def _can_run_sgl_mha_now(attn: DeepseekV2MLAAttention, forward_batch) -> bool: - """Check if the model configuration supports the MHA attention path. - - This is a *capability* gate — NSA models and MXFP4-quantised weights - (uint8) cannot use the MHA path. Distinct from - ``_dispatch_sgl_plugin_attn_path`` which routes each batch. - """ - del forward_batch - if attn.use_nsa: - return False - if attn.kv_b_proj.weight.dtype == torch.uint8: - return False - return True - - -def forward_sgl_mha_prepare( - attn: DeepseekV2MLAAttention, - positions: torch.Tensor, - hidden_states: torch.Tensor, - **model_kwargs, -) -> SglMhaPrepareResult: - forward_batch = model_kwargs.get("forward_batch", None) - if forward_batch is None: - raise RuntimeError("forward_batch is required in forward_sgl_mha_prepare") - - hidden_states_scale = None - if isinstance(hidden_states, tuple): - hidden_states, hidden_states_scale = hidden_states - - attn_mha = _get_sglang_radix_attn(attn.attn_mha) - if getattr(attn_mha, "kv_b_proj", None) is None: - attn_mha.kv_b_proj = attn.kv_b_proj - - if attn.q_lora_rank is not None: - q, latent_cache = ( - get_attn_tp_context() - .fetch_qkv_latent() - .split( - [attn.q_lora_rank, attn.kv_lora_rank + attn.qk_rope_head_dim], - dim=-1, - ) - ) - - if ( - q.shape[0] != positions.shape[0] - and get_tensor_model_parallel_world_size() > 1 - ): - qkv_lora = torch.cat([q, latent_cache], dim=-1) - qkv_lora = get_tp_group().all_gather(qkv_lora, dim=0) - if qkv_lora.shape[0] < positions.shape[0]: - raise RuntimeError( - f"qkv_lora gather mismatch: got {qkv_lora.shape[0]}, " - f"expected {positions.shape[0]}" - ) - qkv_lora = qkv_lora[: positions.shape[0]] - q, latent_cache = torch.split( - qkv_lora, - [attn.q_lora_rank, attn.kv_lora_rank + attn.qk_rope_head_dim], - dim=-1, - ) - - if _use_aiter_gfx95 and attn.q_b_proj.weight.dtype == torch.float8_e4m3fn: - (q, q_scale), _, _, _ = fused_rms_fp8_group_quant( - q, - attn.q_a_layernorm.weight, - attn.q_a_layernorm.eps, - None, - None, - None, - group_size=128, - dtype_quant=torch.float8_e4m3fn, - res1=None, - output_unquantized_inp1=False, - transpose_scale=True, - ) - q = _unwrap_linear_output(attn.q_b_proj(q, q_scale)).view( - -1, attn.num_local_heads, attn.qk_head_dim - ) - else: - q = attn.q_a_layernorm(q) - q = _unwrap_linear_output(attn.q_b_proj(q)).view( - -1, attn.num_local_heads, attn.qk_head_dim - ) - else: - q = _unwrap_linear_output(attn.q_proj(hidden_states, hidden_states_scale)).view( - -1, attn.num_local_heads, attn.qk_head_dim - ) - latent_cache = _unwrap_linear_output( - attn.kv_a_proj_with_mqa(hidden_states, hidden_states_scale) - ) - - _, q_pe = q.split([attn.qk_nope_head_dim, attn.qk_rope_head_dim], dim=-1) - kv_a, _ = latent_cache.split([attn.kv_lora_rank, attn.qk_rope_head_dim], dim=-1) - latent_cache = latent_cache.unsqueeze(1) - - if _use_aiter_gfx95 and attn.kv_b_proj.weight.dtype == torch.float8_e4m3fn: - (kv_a_quanted, kv_a_quanted_scale), kv_a, _, _ = fused_rms_fp8_group_quant( - kv_a, - attn.kv_a_layernorm.weight, - attn.kv_a_layernorm.eps, - None, - None, - None, - group_size=128, - dtype_quant=torch.float8_e4m3fn, - res1=None, - output_unquantized_inp1=True, - transpose_scale=True, - ) - else: - kv_a_quanted = None - kv_a = attn.kv_a_layernorm(kv_a) - - k_pe = latent_cache[:, :, attn.kv_lora_rank :] - if attn.rotary_emb is not None: - q_pe, k_pe = attn.rotary_emb(positions, q_pe, k_pe) - q[..., attn.qk_nope_head_dim :] = q_pe - - _set_mla_kv_buffer_for_mha(attn, kv_a, k_pe, forward_batch) - - if kv_a_quanted is not None: - kv = _unwrap_linear_output(attn.kv_b_proj(kv_a_quanted, kv_a_quanted_scale)) - else: - kv = _unwrap_linear_output(attn.kv_b_proj(kv_a)) - kv = kv.view(-1, attn.num_local_heads, attn.qk_nope_head_dim + attn.v_head_dim) - k_nope = kv[..., : attn.qk_nope_head_dim] - v = kv[..., attn.qk_nope_head_dim :] - k = torch.cat( - [k_nope, k_pe.expand(-1, attn.num_local_heads, -1)], - dim=-1, - ) - return SglMhaPrepareResult(q=q, k=k, v=v, forward_batch=forward_batch) - - -def forward_sgl_mha_core( - attn: DeepseekV2MLAAttention, - prepared: SglMhaPrepareResult, -) -> torch.Tensor: - attn_output = attn.attn_mha( - prepared.q, - prepared.k, - prepared.v, - forward_batch=prepared.forward_batch, - save_kv_cache=False, - ) - attn_output = attn_output.reshape(-1, attn.num_local_heads * attn.v_head_dim) - return attn.o_proj(attn_output) - - -def forward_sgl_plugin_mode_mha( - attn: DeepseekV2MLAAttention, - positions: torch.Tensor, - hidden_states: torch.Tensor, - **model_kwargs, -) -> torch.Tensor: - forward_batch = model_kwargs.get("forward_batch", None) - if forward_batch is None: - raise RuntimeError("forward_batch is required in forward_sgl_plugin_mode_mha") - if not _can_run_sgl_mha_now(attn, forward_batch): - attn.current_sgl_plugin_attn_path = "mla_fallback" - return forward_sgl_plugin_mode_mla( - attn, - positions, - hidden_states, - **model_kwargs, - ) - prepared = forward_sgl_mha_prepare(attn, positions, hidden_states, **model_kwargs) - return forward_sgl_mha_core(attn, prepared) - - -def prepare_qkv_latent( - attn: DeepseekV2MLAAttention, - hidden_states: torch.Tensor, - forward_batch, -) -> torch.Tensor: - """Prepare QKV latent tensor for the sglang communicator.""" - assert attn.q_lora_rank is not None - hidden_states_scale = None - if isinstance(hidden_states, tuple): - hidden_states, hidden_states_scale = hidden_states - qkv_lora = attn.fused_qkv_a_proj(hidden_states, hidden_states_scale) - - # Fallback: when communicator does not enable input_scattered gather, - # force qkv latent token dimension to align with positions. - expected_tokens = 0 - if hasattr(forward_batch, "positions") and forward_batch.positions is not None: - expected_tokens = int(forward_batch.positions.shape[0]) - if expected_tokens <= 0: - expected_tokens = int(getattr(forward_batch, "seq_lens_sum", 0) or 0) - - if ( - expected_tokens > 0 - and qkv_lora.shape[0] != expected_tokens - and get_tensor_model_parallel_world_size() > 1 - ): - qkv_lora = get_tp_group().all_gather(qkv_lora, dim=0) - if qkv_lora.shape[0] > expected_tokens: - qkv_lora = qkv_lora[:expected_tokens] - elif qkv_lora.shape[0] < expected_tokens: - raise RuntimeError( - f"prepare_qkv_latent gather mismatch: got {qkv_lora.shape[0]}, " - f"expected {expected_tokens}" - ) - return qkv_lora - - -# Top-level forward entry point -def forward_sgl_plugin_mode( - attn: DeepseekV2MLAAttention, - positions: torch.Tensor, - hidden_states: torch.Tensor, - **model_kwargs, -) -> torch.Tensor: - """Full MLA forward in sglang plugin mode.""" - forward_batch = model_kwargs.get("forward_batch", None) - if forward_batch is None: - raise RuntimeError("forward_batch is required in forward_sgl_plugin_mode") - - attn_tp_context = get_attn_tp_context() - with attn_tp_context.maybe_input_scattered(forward_batch): - if attn.q_lora_rank is not None: - attn_tp_context.set_attn_inputs( - AttentionInputs( - hidden_states, - forward_batch, - lambda hs, fb: prepare_qkv_latent(attn, hs, fb), - ) - ) - attn_path = _dispatch_sgl_plugin_attn_path(forward_batch) - attn.current_sgl_plugin_attn_path = attn_path - if attn_path == "mha": - return forward_sgl_plugin_mode_mha( - attn, - positions, - hidden_states, - **model_kwargs, - ) - if attn_path == "mla": - return forward_sgl_plugin_mode_mla( - attn, - positions, - hidden_states, - **model_kwargs, - ) - raise ValueError(f"Unsupported plugin attention path: {attn_path}") - - -# Weight post-processing: decomposed into sub-functions -def _read_kv_b_proj_weight(attn: DeepseekV2MLAAttention) -> torch.Tensor: - """Read kv_b_proj weight, handling AWQ and fnuz dtypes.""" - if hasattr(attn.kv_b_proj, "qweight"): - awq_dequant = awq_dequantize_func() - if awq_dequant is None: - raise ValueError( - "AWQ dequantize function is not supported for current device" - ) - w = awq_dequant( - attn.kv_b_proj.qweight, - attn.kv_b_proj.scales, - attn.kv_b_proj.qzeros, - ).T - else: - w = attn.kv_b_proj.weight - - # On ROCm, ATOM creates parameters with fnuz dtype but loads fn bytes. - # View-cast back to fn so the normalize path works correctly. - if _is_fp8_fnuz and w.dtype == torch.float8_e4m3fnuz: - w = w.view(torch.float8_e4m3fn) - - return w - - -def _get_weight_block_size(attn: DeepseekV2MLAAttention) -> Optional[list[int]]: - """Derive weight_block_size from ATOM's quant_type system.""" - from aiter import QuantType as _AiterQuantType - - qt = getattr(attn.kv_b_proj, "quant_type", None) - if qt == _AiterQuantType.per_1x128: - return [128, 128] - elif qt == _AiterQuantType.per_1x32: - return [1, 32] - return None - - -def _process_fp8_weight( - attn: DeepseekV2MLAAttention, - w: torch.Tensor, - weight_block_size: Optional[list[int]], -) -> tuple[torch.Tensor, bool, Optional[torch.Tensor]]: - """Process FP8 weights for kv_b_proj. - - Returns (w, use_deep_gemm_bmm, block_scale). - """ - from atom.model_ops.utils import normalize_e4m3fn_to_e4m3fnuz - from sglang.srt.layers.quantization.fp8_utils import ( - block_quant_dequant, - block_quant_to_tensor_quant, - channel_quant_to_tensor_quant, - inverse_transform_scale_ue8m0, - ) - from sglang.srt.layers.deep_gemm_wrapper import ( - ENABLE_JIT_DEEPGEMM, - DEEPGEMM_BLACKWELL, - ) - from sglang.srt.model_loader.utils import should_deepgemm_weight_requant_ue8m0 - - use_deep_gemm_bmm = False - block_scale = None - - if weight_block_size is not None: - assert hasattr(attn.kv_b_proj, "weight_scale_inv") or hasattr( - attn.kv_b_proj, "weight_scale" - ) - weight_scale = ( - attn.kv_b_proj.weight_scale - if hasattr(attn.kv_b_proj, "weight_scale") - else attn.kv_b_proj.weight_scale_inv - ) - - if _is_fp8_fnuz and w.dtype == torch.float8_e4m3fn: - weight, weight_scale, _ = normalize_e4m3fn_to_e4m3fnuz( - weight=w, weight_scale=weight_scale, input_scale=None - ) - else: - weight = w - - if should_deepgemm_weight_requant_ue8m0( - weight_block_size=weight_block_size - ) and getattr(weight_scale, "format_ue8m0", False): - weight_scale = inverse_transform_scale_ue8m0( - weight_scale, mn=weight.shape[-2] - ) - - if _is_cuda and weight_block_size[0] == 128 and weight_block_size[1] == 128: - if ( - ENABLE_JIT_DEEPGEMM - and not DEEPGEMM_BLACKWELL - and get_bool_env_var("SGL_USE_DEEPGEMM_BMM", "false") - ): - block_scale = weight_scale - use_deep_gemm_bmm = True - else: - w = block_quant_dequant( - weight, weight_scale, weight_block_size, torch.bfloat16 - ) - else: - w, scale = block_quant_to_tensor_quant( - weight, weight_scale, weight_block_size - ) - attn.w_scale = scale - else: - if w.dtype == torch.float8_e4m3fn and _is_fp8_fnuz: - weight, weight_scale, _ = normalize_e4m3fn_to_e4m3fnuz( - weight=w, weight_scale=attn.kv_b_proj.weight_scale, input_scale=None - ) - else: - weight = w - weight_scale = attn.kv_b_proj.weight_scale - - w, scale = channel_quant_to_tensor_quant(weight, weight_scale) - attn.w_scale = scale - - return w, use_deep_gemm_bmm, block_scale - - -def _process_int8_weight( - attn: DeepseekV2MLAAttention, - w: torch.Tensor, - weight_block_size: Optional[list[int]], -) -> torch.Tensor: - """Process INT8 weights for kv_b_proj.""" - from sglang.srt.layers.quantization.int8_utils import ( - block_dequant as int8_block_dequant, - ) - - if weight_block_size is not None: - assert hasattr(attn.kv_b_proj, "weight_scale_inv") - return int8_block_dequant( - w, attn.kv_b_proj.weight_scale_inv, weight_block_size - ).to(torch.bfloat16) - else: - return w.to(torch.bfloat16) * attn.kv_b_proj.weight_scale.to(torch.bfloat16) - - -def _split_kc_vc_like_vllm( - attn: DeepseekV2MLAAttention, w: torch.Tensor -) -> tuple[torch.Tensor, torch.Tensor]: - """Split kv_b_proj weight using vLLM's transpose-first layout.""" - kv_b_proj_weight = w.T - assert kv_b_proj_weight.shape == ( - attn.kv_lora_rank, - attn.num_local_heads * (attn.qk_nope_head_dim + attn.v_head_dim), - ), ( - f"{kv_b_proj_weight.shape=}, " - f"{attn.kv_lora_rank=}, " - f"{attn.num_local_heads=}, " - f"{attn.qk_nope_head_dim=}, " - f"{attn.v_head_dim=}" - ) - kv_b_proj_weight = kv_b_proj_weight.view( - attn.kv_lora_rank, - attn.num_local_heads, - attn.qk_nope_head_dim + attn.v_head_dim, - ) - w_uk, w_uv = kv_b_proj_weight.split( - [attn.qk_nope_head_dim, attn.v_head_dim], dim=-1 - ) - return w_uk.transpose(0, 1).contiguous(), w_uv.permute(1, 2, 0).contiguous() - - -def _split_and_assign_kc_vc( - attn: DeepseekV2MLAAttention, - w: torch.Tensor, - use_deep_gemm_bmm: bool, - block_scale: Optional[torch.Tensor], - weight_block_size: Optional[list[int]], -) -> None: - """Split weight into kc/vc and assign to attn.""" - from atom.model_ops.utils import quark_post_load_weights - - w_kc, w_vc = w.unflatten(0, (-1, attn.qk_nope_head_dim + attn.v_head_dim)).split( - [attn.qk_nope_head_dim, attn.v_head_dim], dim=1 - ) - - # quark fp4 special path - quant_method = getattr(attn.kv_b_proj, "quant_method", None) - quant_config = getattr(quant_method, "quant_config", None) - if ( - _use_aiter_gfx95 - and quant_config is not None - and quant_config.get_name() == "quark" - ): - w_kc, attn.w_scale_k, w_vc, attn.w_scale_v = quark_post_load_weights( - attn, w, "mxfp4" - ) - - if not use_deep_gemm_bmm: - use_vllm_weight_layout = _is_hip and not ( - quant_config is not None and quant_config.get_name() == "quark" - ) - - if use_vllm_weight_layout: - w_kc, w_vc = _split_kc_vc_like_vllm(attn, w) - else: - w_kc = w_kc.transpose(1, 2).contiguous().transpose(1, 2) - w_vc = w_vc.contiguous().transpose(1, 2) - - # Align bf16 kv_b_proj post-load handling with vLLM: split first, then - # quantize kc/vc independently for the fp8 BMM path. - if w.dtype == torch.bfloat16 and (_is_hip or _is_cuda): - w_kc, w_scale_k = dynamic_per_batched_tensor_quant(w_kc, dtype=dtypes.fp8) - w_vc, w_scale_v = dynamic_per_batched_tensor_quant(w_vc, dtype=dtypes.fp8) - attn.w_scale_k = bind_or_assign(attn.w_scale_k, w_scale_k) - attn.w_scale_v = bind_or_assign(attn.w_scale_v, w_scale_v) - - attn.w_kc = bind_or_assign(attn.w_kc, w_kc) - if _is_npu: - w_vc = w_vc.contiguous() - attn.w_vc = bind_or_assign(attn.w_vc, w_vc) - - if _is_cpu and _is_cpu_amx_available and w.dtype == torch.float8_e4m3fn: - attn.w_kc = attn.w_kc.to(torch.bfloat16) * attn.w_scale - attn.w_vc = attn.w_vc.to(torch.bfloat16) * attn.w_scale - else: - num_tiles_k = attn.qk_nope_head_dim // weight_block_size[1] - num_tiles_n = attn.v_head_dim // weight_block_size[0] - ws_kc, ws_vc = block_scale.unflatten( - 0, (-1, (num_tiles_k + num_tiles_n)) - ).split([num_tiles_k, num_tiles_n], dim=1) - - attn.w_scale_k = bind_or_assign( - attn.w_scale_k, ws_kc.transpose(1, 2).contiguous() - ) - attn.w_scale_v = bind_or_assign(attn.w_scale_v, ws_vc.contiguous()) - attn.w_kc = bind_or_assign(attn.w_kc, w_kc.transpose(1, 2).contiguous()) - attn.w_vc = bind_or_assign(attn.w_vc, w_vc.contiguous()) - attn.use_deep_gemm_bmm = True - - -def process_mla_kv_b_proj_after_loading(attn: DeepseekV2MLAAttention) -> None: - """Process kv_b_proj weights after loading for sglang MLA mode. - - Orchestrates reading, quantization handling, and splitting of - kv_b_proj into absorbed w_kc / w_vc weights. - """ - w = _read_kv_b_proj_weight(attn) - weight_block_size = _get_weight_block_size(attn) - - use_deep_gemm_bmm = False - block_scale = None - - # fp8 path - if w.dtype in (torch.float8_e4m3fn, torch.float8_e4m3fnuz): - w, use_deep_gemm_bmm, block_scale = _process_fp8_weight( - attn, w, weight_block_size - ) - - # int8 path - if w.dtype == torch.int8: - w = _process_int8_weight(attn, w, weight_block_size) - - # split and assign kc/vc - _split_and_assign_kc_vc(attn, w, use_deep_gemm_bmm, block_scale, weight_block_size) - - -# One-time model setup (called from base_model_wrapper.py) -def setup_deepseek_for_sglang(model) -> None: - """Patch a DeepseekV2/V3 model for sglang plugin mode. - - - Initialises sglang TP context - - Patches each MLAAttention.forward to dispatch to the sglang MLA path - - Registers process_weights_after_loading hooks - - Stores atom_config on the model - """ - config = model.config - - # Store atom_config (needed by load_weights in the OOT wrapper) - if not hasattr(model, "atom_config"): - from atom.config import get_current_atom_config - - model.atom_config = get_current_atom_config() - - kv_cache_dtype = model.atom_config.kv_cache_dtype - - # Initialise sglang TP context for MLA gather/scatter - from sglang.srt.configs.model_config import is_deepseek_nsa - from sglang.srt.layers.communicator import get_attn_tp_context - - get_attn_tp_context().init_context(config.q_lora_rank, is_deepseek_nsa(config)) - - # Patch each MLAAttention instance - from atom.models.deepseek_v2 import DeepseekV2MLAAttention - - for module in model.modules(): - if isinstance(module, DeepseekV2MLAAttention): - _patch_mla_attention_for_sglang(module, config, kv_cache_dtype) - - -def _patch_mla_attention_for_sglang(attn, config, kv_cache_dtype: str = "bf16") -> None: - """Patch a single DeepseekV2MLAAttention for sglang plugin mode. - - We patch attn.forward (rather than relying solely on ops.Attention = - RadixAttention) because MLA's absorbed-weight forward path replaces the - *entire* forward method — including RoPE, and absorbed - BMM — not just the attention backend. ops.Attention = RadixAttention - handles the backend layer (flash_attn / paged_attn dispatch) and is - already set via set_attn_cls(); this patch sits above that layer. - """ - init_sgl_attrs(attn, config, kv_cache_dtype) - - def patched_forward( - positions: torch.Tensor, - hidden_states: torch.Tensor, - **kwargs, - ) -> torch.Tensor: - from atom.plugin.sglang.models.base_model_wrapper import ( - get_current_forward_batch, - ) - - kwargs["forward_batch"] = get_current_forward_batch() - return forward_sgl_plugin_mode(attn, positions, hidden_states, **kwargs) - - attn.forward = patched_forward - attn.process_weights_after_loading = lambda: process_mla_kv_b_proj_after_loading( - attn - ) diff --git a/atom/plugin/sglang/models/base_model_wrapper.py b/atom/plugin/sglang/models/base_model_wrapper.py index afeaf6b5da..85bc6608e1 100644 --- a/atom/plugin/sglang/models/base_model_wrapper.py +++ b/atom/plugin/sglang/models/base_model_wrapper.py @@ -2,13 +2,12 @@ Registers model architecture classes via SGLANG_EXTERNAL_MODEL_PACKAGE, replacing sglang's built-in implementations with ATOM-optimized versions. + +To add a new model, append its architecture class name to _MODEL_NAMES. """ import logging -from contextlib import contextmanager -from contextvars import ContextVar -from dataclasses import dataclass -from typing import Any, ClassVar, Iterable, Optional, Tuple, Union +from typing import Any, Iterable, Optional, Tuple, Union import torch from torch import nn @@ -18,96 +17,26 @@ from sglang.srt.layers.quantization.base_config import QuantizationConfig from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors -logger = logging.getLogger("atom.plugin.sglang.models") - - -_current_forward_batch: ContextVar[Optional[ForwardBatch]] = ContextVar( - "atom_sglang_current_forward_batch", default=None +from atom.plugin.sglang.runtime import ( + MODEL_ARCH_SPECS, + SGLangForwardBatchMetadata, + SGLangPluginRuntime, + bind_current_forward_batch, + get_current_forward_batch, + get_model_arch_spec, + plugin_runtime_scope, ) +logger = logging.getLogger("atom.plugin.sglang.models") -def get_current_forward_batch(): - return _current_forward_batch.get() - - -@dataclass(frozen=True) -class SGLangForwardBatchMetadata: - """Small context object for one SGLang model forward.""" - - forward_batch: Optional[ForwardBatch] - pp_proxy_tensors: Optional[PPProxyTensors] = None - save_kv_cache: bool = True - _current: ClassVar[ContextVar[Optional["SGLangForwardBatchMetadata"]]] = ContextVar( - "atom_sglang_current_forward_batch_metadata", - default=None, - ) - - @classmethod - def current(cls) -> Optional["SGLangForwardBatchMetadata"]: - return cls._current.get() - - @classmethod - def build( - cls, - forward_batch: Optional[ - Union[ForwardBatch, "SGLangForwardBatchMetadata"] - ] = None, - *, - pp_proxy_tensors: Optional[PPProxyTensors] = None, - save_kv_cache: Optional[bool] = None, - ) -> Optional["SGLangForwardBatchMetadata"]: - if isinstance(forward_batch, cls): - return forward_batch - if forward_batch is None and pp_proxy_tensors is None and save_kv_cache is None: - return cls.current() - return cls( - forward_batch=forward_batch, - pp_proxy_tensors=pp_proxy_tensors, - save_kv_cache=True if save_kv_cache is None else save_kv_cache, - ) - - @classmethod - @contextmanager - def bind(cls, metadata: Optional["SGLangForwardBatchMetadata"]): - meta_token = cls._current.set(metadata) - batch_token = _current_forward_batch.set( - None if metadata is None else metadata.forward_batch - ) - try: - yield metadata - finally: - _current_forward_batch.reset(batch_token) - cls._current.reset(meta_token) - - @staticmethod - def to_intermediate_tensors( - intermediate_tensors, - metadata: Optional["SGLangForwardBatchMetadata"], - ): - if intermediate_tensors is not None or metadata is None: - return intermediate_tensors - pp_proxy_tensors = metadata.pp_proxy_tensors - if pp_proxy_tensors is None: - return intermediate_tensors - tensors = getattr(pp_proxy_tensors, "tensors", None) - if tensors is None: - return intermediate_tensors - from atom.models.utils import IntermediateTensors - - return IntermediateTensors(dict(tensors)) - - -@dataclass(frozen=True) -class ModelArchSpec: - wrapper_binds_gdn_context: bool = False - apply_deepseek_patch: bool = False - - -_MODEL_ARCH_SPECS = { - "DeepseekV3ForCausalLM": ModelArchSpec(apply_deepseek_patch=True), - "Qwen3MoeForCausalLM": ModelArchSpec(), - "Qwen3NextForCausalLM": ModelArchSpec(wrapper_binds_gdn_context=True), -} +__all__ = [ + "EntryClass", + "SGLangForwardBatchMetadata", + "SGLangPluginRuntime", + "bind_current_forward_batch", + "get_current_forward_batch", + "plugin_runtime_scope", +] class _AtomCausalLMBaseForSglang(nn.Module): @@ -133,26 +62,89 @@ def __init__( self.vocab_size = config.vocab_size self.unpadded_vocab_size = config.vocab_size self.model_arch = getattr(config, "architectures", [""])[0] - self.model_arch_spec = _MODEL_ARCH_SPECS.get(self.model_arch, ModelArchSpec()) + self.model_arch_spec = get_model_arch_spec(self.model_arch) import atom - self.model = atom.prepare_model(config=config, engine="sglang") + # TODO: prepare_model() currently handles model construction, config + # generation, attention backend registration, and distributed init. + # Refactor so this wrapper only dispatches the attention backend + # (register_ops_to_sglang + set_attn_cls), and let sglang handle + # model construction directly + with plugin_runtime_scope(framework="sglang"): + from atom.config import get_current_atom_config + + self.model = atom.prepare_model(config=config, engine="sglang") + self.atom_config = getattr(self.model, "atom_config", None) + if self.atom_config is None: + self.atom_config = get_current_atom_config() + self.model.atom_config = self.atom_config if self.model is None: raise ValueError( f"ATOM failed to create model for architecture {self.model_arch}" ) - self.logits_processor = LogitsProcessor(config) + # Under SGLang dp-attention, ATOM runtime interprets non-MoE modules + # like lm_head with tp=1 semantics, so plugin logits must not perform + # an extra TP all-gather after local lm_head matmul. + plugin_skip_all_gather = bool(self.model.atom_config.enable_dp_attention) + self.logits_processor = LogitsProcessor( + config, skip_all_gather=plugin_skip_all_gather + ) # Apply ds model-specific sglang patches (attn dispatch, weight hooks, etc.) # TODO: will remove this after sglang supports atom attention backend if self.model_arch_spec.apply_deepseek_patch: - from atom.plugin.sglang.attention_backend.sgl_attention_mla import ( + from atom.plugin.sglang.models.deepseek_mla import ( setup_deepseek_for_sglang, ) - setup_deepseek_for_sglang(self.model) + with plugin_runtime_scope(framework="sglang", atom_config=self.atom_config): + setup_deepseek_for_sglang(self.model) + + def get_embed_and_head(self): + if hasattr(self.model, "get_embed_and_head"): + return self.model.get_embed_and_head() + + embed_owner = ( + self.model.model + if hasattr(self.model, "model") + and hasattr(self.model.model, "embed_tokens") + else self.model + ) + return embed_owner.embed_tokens.weight, self.model.lm_head.weight + + def set_embed_and_head(self, embed, head): + if hasattr(self.model, "set_embed_and_head"): + return self.model.set_embed_and_head(embed, head) + + embed_owner = ( + self.model.model + if hasattr(self.model, "model") + and hasattr(self.model.model, "embed_tokens") + else self.model + ) + del embed_owner.embed_tokens.weight + del self.model.lm_head.weight + embed_owner.embed_tokens.weight = embed + self.model.lm_head.weight = head + torch.cuda.empty_cache() + torch.cuda.synchronize() + + def set_embed(self, embed): + if hasattr(self.model, "set_embed"): + return self.model.set_embed(embed) + + embed_owner = ( + self.model.model + if hasattr(self.model, "model") + and hasattr(self.model.model, "embed_tokens") + else self.model + ) + del embed_owner.embed_tokens.weight + embed_owner.embed_tokens.weight = embed + torch.cuda.empty_cache() + torch.cuda.synchronize() @torch.no_grad() def forward( @@ -165,38 +157,61 @@ def forward( pp_proxy_tensors: Optional[PPProxyTensors] = None, **model_kwargs: Any, ) -> Union[LogitsProcessorOutput, PPProxyTensors]: - metadata = SGLangForwardBatchMetadata.build( - forward_batch, - pp_proxy_tensors=pp_proxy_tensors, - save_kv_cache=model_kwargs.get("save_kv_cache"), - ) - model_inputs = dict( - input_ids=input_ids, - positions=positions, - intermediate_tensors=SGLangForwardBatchMetadata.to_intermediate_tensors( - pp_proxy_tensors, metadata - ), - inputs_embeds=input_embeds, - ) - with SGLangForwardBatchMetadata.bind(metadata): - if self.model_arch_spec.wrapper_binds_gdn_context: - from atom.plugin.sglang.attention_backend.attention_gdn import ( - SGLangGDNForwardContext, + with plugin_runtime_scope(framework="sglang", atom_config=self.atom_config): + with SGLangPluginRuntime( + atom_config=self.atom_config, + forward_batch=forward_batch, + positions=positions, + input_ids=input_ids, + input_embeds=input_embeds, + set_forward_context=not self.model_arch_spec.wrapper_binds_gdn_context, + ) as runtime: + metadata = SGLangForwardBatchMetadata.build( + runtime.forward_batch, + pp_proxy_tensors=pp_proxy_tensors, + save_kv_cache=model_kwargs.get("save_kv_cache"), ) - - with SGLangGDNForwardContext.bind(metadata): - hidden_states = self.model(**model_inputs) - else: - hidden_states = self.model(**model_inputs) - - if self.pp_group.is_last_rank: - return self.logits_processor( - input_ids, - hidden_states, - self.model.lm_head, - forward_batch, - ) - return hidden_states + model_inputs = dict( + input_ids=runtime.input_ids, + positions=runtime.positions, + intermediate_tensors=SGLangForwardBatchMetadata.to_intermediate_tensors( + pp_proxy_tensors, metadata + ), + inputs_embeds=runtime.input_embeds, + ) + uses_context_only_forward = ( + self.model_arch_spec.apply_deepseek_patch + or self.model_arch_spec.wrapper_binds_gdn_context + ) + with SGLangForwardBatchMetadata.bind(metadata): + if self.model_arch_spec.wrapper_binds_gdn_context: + from atom.plugin.sglang.attention_backend.attention_gdn import ( + SGLangGDNForwardContext, + ) + + with SGLangGDNForwardContext.bind(metadata): + hidden_states = self.model(**model_inputs) + elif uses_context_only_forward: + hidden_states = self.model(**model_inputs) + else: + hidden_states = self.model( + **model_inputs, + forward_batch=runtime.forward_batch, + get_embedding=get_embedding, + pp_proxy_tensors=pp_proxy_tensors, + **model_kwargs, + ) + + hidden_states = runtime.trim_output(hidden_states) + + if self.pp_group.is_last_rank: + return self.logits_processor( + input_ids, + hidden_states, + self.model.lm_head, + forward_batch, + ) + return hidden_states def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): # The passed `weights` iterable from sglang is ignored because ATOM @@ -205,13 +220,14 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): # sglang's default weight iterator. from atom.model_loader.loader import load_model_in_plugin_mode - return load_model_in_plugin_mode( - model=self.model, config=self.model.atom_config, prefix="model." - ) + with plugin_runtime_scope(framework="sglang", atom_config=self.atom_config): + return load_model_in_plugin_mode( + model=self.model, config=self.atom_config, prefix="model." + ) EntryClass = [] -for _name in _MODEL_ARCH_SPECS: +for _name in MODEL_ARCH_SPECS: _cls = type(_name, (_AtomCausalLMBaseForSglang,), {}) globals()[_name] = _cls EntryClass.append(_cls) diff --git a/atom/plugin/sglang/models/deepseek_mla.py b/atom/plugin/sglang/models/deepseek_mla.py new file mode 100644 index 0000000000..a7165b7112 --- /dev/null +++ b/atom/plugin/sglang/models/deepseek_mla.py @@ -0,0 +1,63 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2025, Advanced Micro Devices, Inc. All rights reserved. + +"""Model-level DeepSeek MLA patching for SGLang plugin mode. + +This module owns the install-time hooks that adapt DeepSeek MLA models to +SGLang plugin mode. The heavy DeepSeek-specific runtime helpers live in +`atom.plugin.sglang.models.deepseek_mla_forward`. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +from atom.plugin.sglang.models.deepseek_mla_attention import ( + SGLangDeepseekMLAAttention, +) +from atom.plugin.sglang.models.deepseek_mla_forward import ( + init_sgl_attrs, + process_mla_kv_b_proj_after_loading, +) + +if TYPE_CHECKING: + from atom.models.deepseek_v2 import DeepseekV2MLAAttention + + +def setup_deepseek_for_sglang(model) -> None: + """Patch a DeepSeek V2/V3 model for SGLang plugin mode.""" + config = model.config + + # Store atom_config for the OOT wrapper before install-time hooks run. + if not hasattr(model, "atom_config"): + from atom.config import get_current_atom_config + + model.atom_config = get_current_atom_config() + + kv_cache_dtype = model.atom_config.kv_cache_dtype + + # Initialise SGLang's MLA TP context before patching per-layer forwards. + from sglang.srt.configs.model_config import is_deepseek_nsa + from sglang.srt.layers.communicator import get_attn_tp_context + + get_attn_tp_context().init_context(config.q_lora_rank, is_deepseek_nsa(config)) + + from atom.models.deepseek_v2 import DeepseekV2MLAAttention + + for module in model.modules(): + if isinstance(module, DeepseekV2MLAAttention): + _patch_mla_attention_for_sglang(module, config, kv_cache_dtype) + + +def _patch_mla_attention_for_sglang( + attn: "DeepseekV2MLAAttention", + config: Any, + kv_cache_dtype: str = "bf16", +) -> None: + """Patch one DeepSeek MLA layer for SGLang plugin mode.""" + init_sgl_attrs(attn, config, kv_cache_dtype) + if not isinstance(attn.mla_attn, SGLangDeepseekMLAAttention): + attn.mla_attn = SGLangDeepseekMLAAttention(attn, attn.mla_attn) + attn.process_weights_after_loading = lambda: process_mla_kv_b_proj_after_loading( + attn + ) diff --git a/atom/plugin/sglang/models/deepseek_mla_attention.py b/atom/plugin/sglang/models/deepseek_mla_attention.py new file mode 100644 index 0000000000..bd89870125 --- /dev/null +++ b/atom/plugin/sglang/models/deepseek_mla_attention.py @@ -0,0 +1,332 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2025, Advanced Micro Devices, Inc. All rights reserved. + +"""DeepSeek MLA wrapper for SGLang plugin mode. + +This adapter keeps the model-side entry at ``self.mla_attn(...)`` and owns the +SGLang-specific runtime dispatch for DeepSeek MLA. It is intentionally shaped +closer to the vLLM plugin path than the older model-side monkey-patched entry. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +import torch +from torch import nn + +if TYPE_CHECKING: + from atom.models.deepseek_v2 import DeepseekV2MLAAttention + + +class SGLangDeepseekMLAAttention(nn.Module): + """Enter SGLang DeepSeek MLA runtime through ``self.mla_attn(...)``.""" + + def __init__( + self, + owner_attn: "DeepseekV2MLAAttention", + base_attn: nn.Module, + ) -> None: + super().__init__() + self.owner_attn = owner_attn + self.base_attn = base_attn + + @property + def attn(self): + return getattr(self.base_attn, "attn", self.base_attn) + + def _get_forward_batch(self, kwargs: dict[str, Any]): + forward_batch = kwargs.get("forward_batch", None) + if forward_batch is None: + from atom.plugin.sglang.runtime import ( + get_current_forward_batch, + ) + + forward_batch = get_current_forward_batch() + kwargs["forward_batch"] = forward_batch + if forward_batch is None: + raise RuntimeError( + "forward_batch is required for SGLang DeepSeek MLA wrapper" + ) + return forward_batch + + def _infer_total_tokens(self, forward_batch, tensor: torch.Tensor) -> int: + if hasattr(forward_batch, "input_ids") and forward_batch.input_ids is not None: + return int(forward_batch.input_ids.shape[0]) + if hasattr(forward_batch, "positions") and forward_batch.positions is not None: + return int(forward_batch.positions.shape[0]) + if hasattr(forward_batch, "seq_lens_sum"): + return int(forward_batch.seq_lens_sum) + return int(tensor.shape[0]) + + def _maybe_all_gather( + self, + tensor: torch.Tensor | None, + *, + total_tokens: int, + input_scattered: bool, + ): + if tensor is None or not input_scattered: + return tensor + from sglang.srt.distributed import get_tp_group + + output = tensor.new_empty((total_tokens, *tensor.shape[1:])) + get_tp_group().all_gather_into_tensor(output, tensor) + return output + + def _gather_runtime_inputs( + self, + q_input: torch.Tensor, + kv_c_normed: torch.Tensor, + k_pe: torch.Tensor, + positions: torch.Tensor, + q_scale: torch.Tensor | None, + *, + forward_batch, + input_scattered: bool, + ): + total_tokens = self._infer_total_tokens(forward_batch, q_input) + q_input = self._maybe_all_gather( + q_input, + total_tokens=total_tokens, + input_scattered=input_scattered, + ) + kv_c_normed = self._maybe_all_gather( + kv_c_normed, + total_tokens=total_tokens, + input_scattered=input_scattered, + ) + k_pe = self._maybe_all_gather( + k_pe, + total_tokens=total_tokens, + input_scattered=input_scattered, + ) + positions = self._maybe_all_gather( + positions, + total_tokens=total_tokens, + input_scattered=input_scattered, + ) + q_scale = self._maybe_all_gather( + q_scale, + total_tokens=total_tokens, + input_scattered=input_scattered, + ) + return q_input, kv_c_normed, k_pe, positions, q_scale + + def _project_q( + self, + q_input: torch.Tensor, + q_scale: torch.Tensor | None, + ) -> torch.Tensor: + attn = self.owner_attn + from atom.plugin.sglang.models.deepseek_mla_forward import _unwrap_linear_output + + if attn.q_lora_rank is not None: + q = ( + attn.q_b_proj(q_input, q_scale) + if q_scale is not None + else attn.q_b_proj(q_input) + ) + else: + q = ( + attn.q_proj(q_input, q_scale) + if q_scale is not None + else attn.q_proj(q_input) + ) + return _unwrap_linear_output(q).view( + -1, attn.num_local_heads, attn.qk_head_dim + ) + + def _forward_absorbed( + self, + q_input: torch.Tensor, + kv_c_normed: torch.Tensor, + k_pe: torch.Tensor, + positions: torch.Tensor, + q_scale: torch.Tensor | None, + *, + forward_batch, + ) -> torch.Tensor: + attn = self.owner_attn + from aiter import dtypes + from atom.model_ops.attention_mla import fused_qk_rope_concat_and_cache_mla + from atom.plugin.sglang.models.deepseek_mla_forward import ( + _get_sglang_radix_attn, + mla_absorbed_bmm, + mla_v_up_proj, + ) + from sglang.srt.layers.attention.nsa.utils import nsa_use_prefill_cp + + q = self._project_q(q_input, q_scale) + k_nope = kv_c_normed.unsqueeze(1) + k_pe = k_pe.unsqueeze(1) + q_nope, q_pe = q.split( + [attn.qk_nope_head_dim, attn.qk_rope_head_dim], dim=-1 + ) + q_nope_out = mla_absorbed_bmm( + attn, q_nope, attn.w_kc, attn.w_scale, attn.w_scale_k, attn.kv_lora_rank + ) + + if attn.rotary_emb is not None and not attn.use_fused_qk_rope_concat_and_cache_mla: + q_pe, k_pe = attn.rotary_emb(positions, q_pe, k_pe) + + if nsa_use_prefill_cp(forward_batch): + latent_cache = torch.cat([k_nope.squeeze(1), k_pe.squeeze(1)], dim=-1) + k_nope, k_pe = attn.rebuild_cp_kv_cache( + latent_cache, forward_batch, k_nope, k_pe + ) + + save_kv_cache = True + if attn.use_fused_qk_rope_concat_and_cache_mla: + mla_attn = _get_sglang_radix_attn(self.base_attn) + kv_cache = forward_batch.token_to_kv_pool.get_key_buffer(mla_attn.layer_id) + q_out_dtype = ( + dtypes.fp8 + if attn.kv_cache_dtype == "fp8_e4m3" + else q_nope_out.dtype + ) + q = torch.empty( + ( + q_nope_out.shape[0], + attn.num_local_heads, + attn.kv_lora_rank + attn.qk_rope_head_dim, + ), + dtype=q_out_dtype, + device=q_nope_out.device, + ) + fused_qk_rope_concat_and_cache_mla( + q_nope_out, + q_pe, + k_nope, + k_pe, + kv_cache, + q, + forward_batch.out_cache_loc, + mla_attn.k_scale, + mla_attn.k_scale, + positions, + attn.rotary_emb.cos_cache, + attn.rotary_emb.sin_cache, + is_neox=attn.rotary_emb.is_neox_style, + is_nope_first=True, + ) + k = None + v = None + save_kv_cache = False + else: + q = torch.cat([q_nope_out, q_pe], dim=-1) + k = torch.cat([k_nope, k_pe], dim=-1) + v = k_nope + + attn_output = self.base_attn( + q, + k, + v, + forward_batch=forward_batch, + save_kv_cache=save_kv_cache, + ) + attn_output = attn_output.view(-1, attn.num_local_heads, attn.kv_lora_rank) + attn_bmm_output = mla_v_up_proj( + attn, attn_output, attn.w_vc, attn.w_scale, attn.w_scale_v, attn.v_head_dim + ) + return attn.o_proj(attn_bmm_output) + + def _forward_non_absorbed( + self, + q_input: torch.Tensor, + kv_c_normed: torch.Tensor, + k_pe: torch.Tensor, + positions: torch.Tensor, + q_scale: torch.Tensor | None, + *, + forward_batch, + ) -> torch.Tensor: + attn = self.owner_attn + from atom.plugin.sglang.models.deepseek_mla_forward import ( + _set_mla_kv_buffer_for_non_absorbed, + _unwrap_linear_output, + ) + + q = self._project_q(q_input, q_scale) + _, q_pe = q.split([attn.qk_nope_head_dim, attn.qk_rope_head_dim], dim=-1) + + kv_a = kv_c_normed + k_pe = k_pe.unsqueeze(1) + if attn.rotary_emb is not None: + q_pe, k_pe = attn.rotary_emb(positions, q_pe, k_pe) + q[..., attn.qk_nope_head_dim :] = q_pe + + _set_mla_kv_buffer_for_non_absorbed(attn, kv_a, k_pe, forward_batch) + + kv = _unwrap_linear_output(attn.kv_b_proj(kv_a)).view( + -1, attn.num_local_heads, attn.qk_nope_head_dim + attn.v_head_dim + ) + k_nope = kv[..., : attn.qk_nope_head_dim] + v = kv[..., attn.qk_nope_head_dim :] + k = torch.cat( + [k_nope, k_pe.expand(-1, attn.num_local_heads, -1)], + dim=-1, + ) + + attn_output = attn.attn_non_absorbed( + q, + k, + v, + forward_batch=forward_batch, + save_kv_cache=False, + ) + attn_output = attn_output.reshape(-1, attn.num_local_heads * attn.v_head_dim) + return attn.o_proj(attn_output) + + def forward( + self, + q_input: torch.Tensor, + kv_c_normed: torch.Tensor, + k_pe: torch.Tensor, + positions: torch.Tensor, + q_scale: torch.Tensor | None = None, + **kwargs: Any, + ) -> torch.Tensor: + attn = self.owner_attn + forward_batch = self._get_forward_batch(kwargs) + + from atom.plugin.sglang.models.deepseek_mla_forward import ( + _can_run_non_absorbed_mla_now, + ) + from sglang.srt.layers.communicator import get_attn_tp_context + + attn_tp_context = get_attn_tp_context() + with attn_tp_context.maybe_input_scattered(forward_batch): + q_input, kv_c_normed, k_pe, positions, q_scale = self._gather_runtime_inputs( + q_input, + kv_c_normed, + k_pe, + positions, + q_scale, + forward_batch=forward_batch, + input_scattered=attn_tp_context.input_scattered, + ) + + if forward_batch.forward_mode.is_extend_without_speculative(): + if _can_run_non_absorbed_mla_now(attn, forward_batch): + attn.current_sgl_plugin_attn_path = "non_absorbed" + return self._forward_non_absorbed( + q_input, + kv_c_normed, + k_pe, + positions, + q_scale, + forward_batch=forward_batch, + ) + attn.current_sgl_plugin_attn_path = "absorbed_fallback" + else: + attn.current_sgl_plugin_attn_path = "absorbed" + + return self._forward_absorbed( + q_input, + kv_c_normed, + k_pe, + positions, + q_scale, + forward_batch=forward_batch, + ) diff --git a/atom/plugin/sglang/models/deepseek_mla_forward.py b/atom/plugin/sglang/models/deepseek_mla_forward.py new file mode 100644 index 0000000000..f88aa762f8 --- /dev/null +++ b/atom/plugin/sglang/models/deepseek_mla_forward.py @@ -0,0 +1,533 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2025, Advanced Micro Devices, Inc. All rights reserved. + +"""Helper functions for DeepSeek MLA in SGLang plugin mode. + +This module now contains only the low-level helpers that are still shared by +the SGLang DeepSeek MLA wrapper and the install-time weight hooks: +absorbed BMM math, small utility helpers, non-absorbed cache staging, and +kv_b_proj post-load processing. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Optional + +import torch +from aiter import dtypes +from atom.model_ops.base_attention import Attention +from atom.model_ops.attention_mla import ( + dynamic_per_batched_tensor_quant, +) +from atom.models.utils import maybe_prefix + +from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode +from sglang.srt.models.deepseek_common.utils import ( + _use_aiter_gfx95, + _is_hip, + _is_cpu, + _is_cpu_amx_available, + _is_cuda, + _is_fp8_fnuz, + _is_npu, + awq_dequantize_func, +) +from sglang.srt.layers.quantization.rocm_mxfp4_utils import ( + batched_gemm_afp4wfp4_pre_quant, +) +from aiter.ops.triton.batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant import ( + batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant, +) +from sglang.srt.layers.quantization.fp8_kernel import ( + per_tensor_quant_mla_fp8, + per_token_group_quant_mla_deep_gemm_masked_fp8, +) +from sglang.srt.utils import bind_or_assign, get_bool_env_var + +if TYPE_CHECKING: + from atom.models.deepseek_v2 import DeepseekV2MLAAttention + + +if _is_cuda: + from sgl_kernel import bmm_fp8 as _raw_bmm_fp8 + from sglang.srt.utils.custom_op import register_custom_op + + @register_custom_op(mutates_args=["out"]) + def _bmm_fp8_op( + A: torch.Tensor, + B: torch.Tensor, + out: torch.Tensor, + A_scale: torch.Tensor, + B_scale: torch.Tensor, + ) -> None: + _raw_bmm_fp8(A, B, A_scale, B_scale, out.dtype, out) + + def bmm_fp8(A, B, A_scale, B_scale, dtype, out=None): + if out is None: + out = torch.empty( + (A.shape[0], A.shape[1], B.shape[2]), + device=A.device, + dtype=dtype, + ) + _bmm_fp8_op(A, B, out, A_scale, B_scale) + return out + +else: + + def bmm_fp8(A, B, A_scale, B_scale, dtype, out=None): + raise RuntimeError("bmm_fp8 requires CUDA (sgl_kernel)") + + +def _unwrap_linear_output(output: Any) -> torch.Tensor: + """Normalize ATOM/public-SGLang linear outputs to a tensor.""" + if isinstance(output, tuple): + return output[0] + return output + + +def _prepare_weight_for_bmm( + weight: torch.Tensor, in_dim: int, out_dim: int +) -> torch.Tensor: + """Normalize absorbed weight layout for torch.bmm fallback.""" + if weight.shape[1] == in_dim and weight.shape[2] == out_dim: + return weight + if weight.shape[1] == out_dim and weight.shape[2] == in_dim: + return weight.transpose(-2, -1) + raise RuntimeError( + "Unexpected absorbed weight shape for bmm fallback: " + f"{tuple(weight.shape)} with in_dim={in_dim}, out_dim={out_dim}" + ) + + +def init_sgl_attrs( + attn: DeepseekV2MLAAttention, + config, + kv_cache_dtype: str = "bf16", +) -> None: + """Initialise sglang-only attributes on DeepseekV2MLAAttention.""" + from sglang.srt.configs.model_config import is_deepseek_nsa + + attn.use_nsa = is_deepseek_nsa(config) + attn.use_deep_gemm_bmm = False + attn.alt_stream = None + attn.kv_cache_dtype = kv_cache_dtype + attn.use_fused_qk_rope_concat_and_cache_mla = _use_aiter_gfx95 + attn.current_sgl_plugin_attn_path = None + attn.w_kc, attn.w_vc = None, None + attn.w_scale = None + attn.w_scale_k = None + attn.w_scale_v = None + attn.attn_non_absorbed = Attention( + num_heads=attn.num_local_heads, + head_dim=attn.qk_head_dim, + scale=attn.scaling, + num_kv_heads=attn.num_local_heads, + kv_cache_dtype=kv_cache_dtype, + layer_num=attn.layer_num, + use_mla=False, + v_head_dim=attn.v_head_dim, + prefix=maybe_prefix(attn.prefix, "attn_non_absorbed"), + ) + if hasattr(attn.attn_non_absorbed, "attn"): + attn.attn_non_absorbed.attn.kv_b_proj = None + + +def mla_absorbed_bmm( + attn: DeepseekV2MLAAttention, + inp: torch.Tensor, + weight: torch.Tensor, + weight_scale: Optional[torch.Tensor], + weight_scale_k: Optional[torch.Tensor], + out_dim: int, +) -> torch.Tensor: + """Batched matmul for MLA absorbed weights (w_kc / w_vc).""" + effective_weight_scale = ( + weight_scale_k if weight_scale_k is not None else weight_scale + ) + + if attn.use_deep_gemm_bmm: + from sglang.srt.layers import deep_gemm_wrapper + + val, scale, masked_m, expected_m, aligned_m = ( + per_token_group_quant_mla_deep_gemm_masked_fp8(inp.transpose(0, 1)) + ) + out = inp.new_empty((attn.num_local_heads, aligned_m, out_dim)) + deep_gemm_wrapper.grouped_gemm_nt_f8f8bf16_masked( + (val, scale), + (weight, weight_scale_k), + out, + masked_m, + expected_m, + ) + return out[:, :expected_m, :].transpose(0, 1) + + if _is_hip: + if _use_aiter_gfx95 and weight.dtype == torch.uint8: + x = inp.transpose(0, 1) + out = torch.empty( + x.shape[0], + x.shape[1], + weight.shape[2], + device=x.device, + dtype=torch.bfloat16, + ) + batched_gemm_afp4wfp4_pre_quant( + x, + weight.transpose(-2, -1), + weight_scale_k.transpose(-2, -1), + torch.bfloat16, + out, + ) + return out.transpose(0, 1) + + if (_use_aiter_gfx95 and weight.dtype == torch.float8_e4m3fn) or ( + get_is_capture_mode() and weight.dtype == torch.float8_e4m3fnuz + ): + x = inp.transpose(0, 1) + out = ( + batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant( + X=x, + WQ=weight, + w_scale=effective_weight_scale, + group_size=128, + YQ=None, + transpose_bm=True, + transpose_bm_in=False, + dtype=torch.bfloat16, + ) + ) + return out + + w_bf16 = _prepare_weight_for_bmm(weight, inp.shape[-1], out_dim).to( + torch.bfloat16 + ) + if effective_weight_scale is not None: + w_bf16 = w_bf16 * effective_weight_scale + out = torch.bmm( + inp.to(torch.bfloat16).transpose(0, 1), + w_bf16, + ) + return out.transpose(0, 1) + + if weight.dtype == torch.float8_e4m3fn: + val, scale = per_tensor_quant_mla_fp8( + inp.transpose(0, 1), + torch.zeros((1,), dtype=torch.float32, device=inp.device), + ) + out = bmm_fp8(val, weight, scale, effective_weight_scale, torch.bfloat16) + return out.transpose(0, 1) + + return torch.bmm(inp.transpose(0, 1), weight).transpose(0, 1) + + +def mla_v_up_proj( + attn: DeepseekV2MLAAttention, + inp: torch.Tensor, + weight: torch.Tensor, + weight_scale: Optional[torch.Tensor], + weight_scale_k: Optional[torch.Tensor], + out_dim: int, +) -> torch.Tensor: + """Project MLA decode output to a flat o_proj input.""" + effective_weight_scale = ( + weight_scale_k if weight_scale_k is not None else weight_scale + ) + if _is_hip and ( + (_use_aiter_gfx95 and weight.dtype == torch.float8_e4m3fn) + or (get_is_capture_mode() and weight.dtype == torch.float8_e4m3fnuz) + ): + x = inp.transpose(0, 1) + out = torch.empty( + (inp.shape[0], attn.num_local_heads * out_dim), + device=inp.device, + dtype=torch.bfloat16, + ) + out_3d = out.view(inp.shape[0], attn.num_local_heads, out_dim) + batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant( + X=x, + WQ=weight, + w_scale=effective_weight_scale, + group_size=128, + YQ=out_3d, + transpose_bm=True, + transpose_bm_in=False, + dtype=torch.bfloat16, + ) + return out + + return mla_absorbed_bmm( + attn, inp, weight, weight_scale, weight_scale_k, out_dim + ).flatten(1, 2) +def _get_sglang_radix_attn(attn_module): + return attn_module.attn if hasattr(attn_module, "attn") else attn_module + + +def _set_mla_kv_buffer_for_non_absorbed( + attn: DeepseekV2MLAAttention, + kv_a: torch.Tensor, + k_pe: torch.Tensor, + forward_batch, +) -> None: + attn_non_absorbed = _get_sglang_radix_attn(attn.attn_non_absorbed) + cache_k = torch.cat([kv_a.unsqueeze(1), k_pe], dim=-1) + forward_batch.token_to_kv_pool.set_kv_buffer( + attn_non_absorbed, + forward_batch.out_cache_loc, + cache_k, + cache_k, + ) + + +def _can_run_non_absorbed_mla_now( + attn: DeepseekV2MLAAttention, + forward_batch, +) -> bool: + """Check if the non-absorbed MLA path is allowed for this batch.""" + del forward_batch + if attn.use_nsa: + return False + if attn.kv_b_proj.weight.dtype == torch.uint8: + return False + return True + + +def _read_kv_b_proj_weight(attn: DeepseekV2MLAAttention) -> torch.Tensor: + """Read kv_b_proj weight, handling AWQ and fnuz dtypes.""" + if hasattr(attn.kv_b_proj, "qweight"): + awq_dequant = awq_dequantize_func() + if awq_dequant is None: + raise ValueError( + "AWQ dequantize function is not supported for current device" + ) + w = awq_dequant( + attn.kv_b_proj.qweight, + attn.kv_b_proj.scales, + attn.kv_b_proj.qzeros, + ).T + else: + w = attn.kv_b_proj.weight + + if _is_fp8_fnuz and w.dtype == torch.float8_e4m3fnuz: + w = w.view(torch.float8_e4m3fn) + + return w + + +def _get_weight_block_size(attn: DeepseekV2MLAAttention) -> Optional[list[int]]: + """Derive weight_block_size from ATOM's quant_type system.""" + from aiter import QuantType as _AiterQuantType + + qt = getattr(attn.kv_b_proj, "quant_type", None) + if qt == _AiterQuantType.per_1x128: + return [128, 128] + elif qt == _AiterQuantType.per_1x32: + return [1, 32] + return None + + +def _process_fp8_weight( + attn: DeepseekV2MLAAttention, + w: torch.Tensor, + weight_block_size: Optional[list[int]], +) -> tuple[torch.Tensor, bool, Optional[torch.Tensor]]: + """Process FP8 weights for kv_b_proj.""" + from atom.model_ops.utils import normalize_e4m3fn_to_e4m3fnuz + from sglang.srt.layers.quantization.fp8_utils import ( + block_quant_dequant, + block_quant_to_tensor_quant, + channel_quant_to_tensor_quant, + inverse_transform_scale_ue8m0, + ) + from sglang.srt.layers.deep_gemm_wrapper import ( + ENABLE_JIT_DEEPGEMM, + DEEPGEMM_BLACKWELL, + ) + from sglang.srt.model_loader.utils import should_deepgemm_weight_requant_ue8m0 + + use_deep_gemm_bmm = False + block_scale = None + + if weight_block_size is not None: + assert hasattr(attn.kv_b_proj, "weight_scale_inv") or hasattr( + attn.kv_b_proj, "weight_scale" + ) + weight_scale = ( + attn.kv_b_proj.weight_scale + if hasattr(attn.kv_b_proj, "weight_scale") + else attn.kv_b_proj.weight_scale_inv + ) + + if _is_fp8_fnuz and w.dtype == torch.float8_e4m3fn: + weight, weight_scale, _ = normalize_e4m3fn_to_e4m3fnuz( + weight=w, weight_scale=weight_scale, input_scale=None + ) + else: + weight = w + + if should_deepgemm_weight_requant_ue8m0( + weight_block_size=weight_block_size + ) and getattr(weight_scale, "format_ue8m0", False): + weight_scale = inverse_transform_scale_ue8m0( + weight_scale, mn=weight.shape[-2] + ) + + if _is_cuda and weight_block_size[0] == 128 and weight_block_size[1] == 128: + if ( + ENABLE_JIT_DEEPGEMM + and not DEEPGEMM_BLACKWELL + and get_bool_env_var("SGL_USE_DEEPGEMM_BMM", "false") + ): + block_scale = weight_scale + use_deep_gemm_bmm = True + else: + w = block_quant_dequant( + weight, weight_scale, weight_block_size, torch.bfloat16 + ) + else: + w, scale = block_quant_to_tensor_quant( + weight, weight_scale, weight_block_size + ) + attn.w_scale = scale + else: + if w.dtype == torch.float8_e4m3fn and _is_fp8_fnuz: + weight, weight_scale, _ = normalize_e4m3fn_to_e4m3fnuz( + weight=w, weight_scale=attn.kv_b_proj.weight_scale, input_scale=None + ) + else: + weight = w + weight_scale = attn.kv_b_proj.weight_scale + + w, scale = channel_quant_to_tensor_quant(weight, weight_scale) + attn.w_scale = scale + + return w, use_deep_gemm_bmm, block_scale + + +def _process_int8_weight( + attn: DeepseekV2MLAAttention, + w: torch.Tensor, + weight_block_size: Optional[list[int]], +) -> torch.Tensor: + """Process INT8 weights for kv_b_proj.""" + from sglang.srt.layers.quantization.int8_utils import ( + block_dequant as int8_block_dequant, + ) + + if weight_block_size is not None: + assert hasattr(attn.kv_b_proj, "weight_scale_inv") + return int8_block_dequant( + w, attn.kv_b_proj.weight_scale_inv, weight_block_size + ).to(torch.bfloat16) + else: + return w.to(torch.bfloat16) * attn.kv_b_proj.weight_scale.to(torch.bfloat16) + + +def _split_kc_vc_like_vllm( + attn: DeepseekV2MLAAttention, w: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor]: + """Split kv_b_proj weight using vLLM's transpose-first layout.""" + kv_b_proj_weight = w.T + assert kv_b_proj_weight.shape == ( + attn.kv_lora_rank, + attn.num_local_heads * (attn.qk_nope_head_dim + attn.v_head_dim), + ), ( + f"{kv_b_proj_weight.shape=}, " + f"{attn.kv_lora_rank=}, " + f"{attn.num_local_heads=}, " + f"{attn.qk_nope_head_dim=}, " + f"{attn.v_head_dim=}" + ) + kv_b_proj_weight = kv_b_proj_weight.view( + attn.kv_lora_rank, + attn.num_local_heads, + attn.qk_nope_head_dim + attn.v_head_dim, + ) + w_uk, w_uv = kv_b_proj_weight.split( + [attn.qk_nope_head_dim, attn.v_head_dim], dim=-1 + ) + return w_uk.transpose(0, 1).contiguous(), w_uv.permute(1, 2, 0).contiguous() + + +def _split_and_assign_kc_vc( + attn: DeepseekV2MLAAttention, + w: torch.Tensor, + use_deep_gemm_bmm: bool, + block_scale: Optional[torch.Tensor], + weight_block_size: Optional[list[int]], +) -> None: + """Split weight into kc/vc and assign to attn.""" + from atom.model_ops.utils import quark_post_load_weights + + w_kc, w_vc = w.unflatten(0, (-1, attn.qk_nope_head_dim + attn.v_head_dim)).split( + [attn.qk_nope_head_dim, attn.v_head_dim], dim=1 + ) + + quant_method = getattr(attn.kv_b_proj, "quant_method", None) + quant_config = getattr(quant_method, "quant_config", None) + if ( + _use_aiter_gfx95 + and quant_config is not None + and quant_config.get_name() == "quark" + ): + w_kc, attn.w_scale_k, w_vc, attn.w_scale_v = quark_post_load_weights( + attn, w, "mxfp4" + ) + + if not use_deep_gemm_bmm: + use_vllm_weight_layout = _is_hip and not ( + quant_config is not None and quant_config.get_name() == "quark" + ) + + if use_vllm_weight_layout: + w_kc, w_vc = _split_kc_vc_like_vllm(attn, w) + else: + w_kc = w_kc.transpose(1, 2).contiguous().transpose(1, 2) + w_vc = w_vc.contiguous().transpose(1, 2) + + if w.dtype == torch.bfloat16 and (_is_hip or _is_cuda): + w_kc, w_scale_k = dynamic_per_batched_tensor_quant(w_kc, dtype=dtypes.fp8) + w_vc, w_scale_v = dynamic_per_batched_tensor_quant(w_vc, dtype=dtypes.fp8) + attn.w_scale_k = bind_or_assign(attn.w_scale_k, w_scale_k) + attn.w_scale_v = bind_or_assign(attn.w_scale_v, w_scale_v) + + attn.w_kc = bind_or_assign(attn.w_kc, w_kc) + if _is_npu: + w_vc = w_vc.contiguous() + attn.w_vc = bind_or_assign(attn.w_vc, w_vc) + + if _is_cpu and _is_cpu_amx_available and w.dtype == torch.float8_e4m3fn: + attn.w_kc = attn.w_kc.to(torch.bfloat16) * attn.w_scale + attn.w_vc = attn.w_vc.to(torch.bfloat16) * attn.w_scale + else: + num_tiles_k = attn.qk_nope_head_dim // weight_block_size[1] + num_tiles_n = attn.v_head_dim // weight_block_size[0] + ws_kc, ws_vc = block_scale.unflatten( + 0, (-1, (num_tiles_k + num_tiles_n)) + ).split([num_tiles_k, num_tiles_n], dim=1) + + attn.w_scale_k = bind_or_assign( + attn.w_scale_k, ws_kc.transpose(1, 2).contiguous() + ) + attn.w_scale_v = bind_or_assign(attn.w_scale_v, ws_vc.contiguous()) + attn.w_kc = bind_or_assign(attn.w_kc, w_kc.transpose(1, 2).contiguous()) + attn.w_vc = bind_or_assign(attn.w_vc, w_vc.contiguous()) + attn.use_deep_gemm_bmm = True + + +def process_mla_kv_b_proj_after_loading(attn: DeepseekV2MLAAttention) -> None: + """Process kv_b_proj weights after loading for sglang MLA mode.""" + w = _read_kv_b_proj_weight(attn) + weight_block_size = _get_weight_block_size(attn) + + use_deep_gemm_bmm = False + block_scale = None + + if w.dtype in (torch.float8_e4m3fn, torch.float8_e4m3fnuz): + w, use_deep_gemm_bmm, block_scale = _process_fp8_weight( + attn, w, weight_block_size + ) + + if w.dtype == torch.int8: + w = _process_int8_weight(attn, w, weight_block_size) + + _split_and_assign_kc_vc(attn, w, use_deep_gemm_bmm, block_scale, weight_block_size) diff --git a/atom/plugin/sglang/models/deepseek_nextn_wrapper.py b/atom/plugin/sglang/models/deepseek_nextn_wrapper.py new file mode 100644 index 0000000000..730355aaff --- /dev/null +++ b/atom/plugin/sglang/models/deepseek_nextn_wrapper.py @@ -0,0 +1,202 @@ +"""ATOM DeepSeek NextN wrapper for SGLang external loading. + +This keeps SGLang's draft architecture name (`DeepseekV3ForCausalLMNextN`) +so ModelRegistry can override the upstream implementation, but delegates the +actual draft core to ATOM's `DeepSeekMTP`. +""" + +import logging +from typing import Iterable, Optional, Tuple + +import torch +from torch import nn + +from sglang.srt.distributed import get_pp_group +from sglang.srt.layers.logits_processor import LogitsProcessor +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.server_args import get_global_server_args + +from atom.config import SpeculativeConfig +from atom.plugin.config import generate_atom_config_for_plugin_mode +from atom.plugin.sglang.models.deepseek_mla import ( + setup_deepseek_for_sglang, +) +from atom.plugin.sglang.runtime import ( + SGLangPluginRuntime, + plugin_runtime_scope, +) + +logger = logging.getLogger("atom.plugin.sglang.models") + + +def _sync_replaced_weights() -> None: + if torch.cuda.is_available(): + torch.cuda.empty_cache() + torch.cuda.synchronize() + + +def _replace_weight(module: nn.Module, attr_name: str, weight) -> None: + if hasattr(module, attr_name): + delattr(module, attr_name) + setattr(module, attr_name, weight) + + +def _set_runtime_layer_id(layer_module: nn.Module, layer_id: int) -> None: + if hasattr(layer_module, "layer_id"): + layer_module.layer_id = layer_id + if hasattr(layer_module, "layer_num"): + layer_module.layer_num = layer_id + + +def _retag_mtp_runtime_layer_ids(model: nn.Module) -> None: + """Retag MTP runtime layer ids to draft-local indices. + + ATOM's DeepSeekMTP keeps checkpoint/global layer numbering (e.g. 61, 62...) + in module prefixes so weight remapping still works. SGLang's draft KV cache, + however, allocates layers using draft-local indices (0..num_nextn_layers-1). + Rebind only the runtime ids used by the attention/KV-cache path. + """ + + for local_layer_id, mtp_layer in enumerate(model.model.layers.values()): + mtp_block = mtp_layer.mtp_block + self_attn = mtp_block.self_attn + + _set_runtime_layer_id(self_attn, local_layer_id) + + for attr_name in ("mla_attn", "attn_mha"): + attn_obj = getattr(self_attn, attr_name, None) + if attn_obj is None: + continue + _set_runtime_layer_id(attn_obj, local_layer_id) + nested_attn = getattr(attn_obj, "attn", None) + if nested_attn is not None: + _set_runtime_layer_id(nested_attn, local_layer_id) + + +class DeepseekV3ForCausalLMNextN(nn.Module): + """SGLang-compatible draft wrapper backed by ATOM's `DeepSeekMTP`.""" + + def __init__( + self, + config, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + del prefix + super().__init__() + + logger.info("Initializing ATOM backend for %s", self.__class__.__name__) + + self.pp_group = get_pp_group() + self.quant_config = quant_config + self.config = config + self.vocab_size = config.vocab_size + self.unpadded_vocab_size = config.vocab_size + + with plugin_runtime_scope(framework="sglang"): + self.atom_config = generate_atom_config_for_plugin_mode(config) + + # Draft workers need ATOM's MTP-specific config semantics rather than the + # default target-model translation used by the generic plugin wrapper. + SpeculativeConfig.hf_config_override(self.atom_config.hf_config) + + with plugin_runtime_scope(framework="sglang", atom_config=self.atom_config): + from atom.plugin.register import ( + init_aiter_dist, + register_ops_to_sglang, + set_attn_cls, + ) + from atom.models.deepseek_mtp import DeepSeekMTP + + register_ops_to_sglang(atom_config=self.atom_config) + set_attn_cls() + init_aiter_dist(config=self.atom_config) + + self.model = DeepSeekMTP(atom_config=self.atom_config) + self.model.atom_config = self.atom_config + setup_deepseek_for_sglang(self.model) + _retag_mtp_runtime_layer_ids(self.model) + + self.logits_processor = LogitsProcessor(config) + self.lm_head = self._first_mtp_layer().shared_head.head + + def _mtp_layers(self): + return list(self.model.model.layers.values()) + + def _first_mtp_layer(self): + layers = self._mtp_layers() + if not layers: + raise ValueError("DeepSeekMTP does not contain any draft layers") + return layers[0] + + def get_embed_and_head(self): + return self.model.model.embed_tokens.weight, self.lm_head.weight + + def set_embed_and_head(self, embed, head): + self.set_embed(embed) + for mtp_layer in self._mtp_layers(): + _replace_weight(mtp_layer.shared_head.head, "weight", head) + self.lm_head = self._first_mtp_layer().shared_head.head + _sync_replaced_weights() + + def set_embed(self, embed): + _replace_weight(self.model.model.embed_tokens, "weight", embed) + _sync_replaced_weights() + + @torch.no_grad() + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: torch.Tensor = None, + **kwargs, + ): + if forward_batch.spec_info is None: + raise ValueError("DeepSeek MTP draft forward requires speculative info") + + with plugin_runtime_scope(framework="sglang", atom_config=self.atom_config): + with SGLangPluginRuntime( + atom_config=self.atom_config, + forward_batch=forward_batch, + positions=positions, + input_ids=input_ids, + input_embeds=input_embeds, + ): + hidden_states = self.model( + input_ids=input_ids, + positions=positions, + hidden_states=forward_batch.spec_info.hidden_states, + inputs_embeds=input_embeds, + ) + + if self.pp_group.is_last_rank: + return self.logits_processor( + input_ids, + hidden_states, + self.lm_head, + forward_batch, + ) + return hidden_states + + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): + del weights + from atom.model_loader.loader import load_model + + server_args = get_global_server_args() + draft_model_path = ( + server_args.speculative_draft_model_path or server_args.model_path + ) + self.atom_config.model = draft_model_path + with plugin_runtime_scope(framework="sglang", atom_config=self.atom_config): + return load_model( + model=self.model, + model_name_or_path=draft_model_path, + hf_config=self.atom_config.hf_config, + load_dummy=self.atom_config.load_dummy, + spec_decode=True, + ) + + +EntryClass = [DeepseekV3ForCausalLMNextN] diff --git a/atom/plugin/sglang/models/qwen3_5.py b/atom/plugin/sglang/models/qwen3_5.py index 22294b8548..c0b6b0ca0b 100644 --- a/atom/plugin/sglang/models/qwen3_5.py +++ b/atom/plugin/sglang/models/qwen3_5.py @@ -37,7 +37,7 @@ from atom.plugin.sglang.attention_backend.attention_gdn import ( SGLangGDNForwardContext, ) -from atom.plugin.sglang.models.base_model_wrapper import ( +from atom.plugin.sglang.runtime import ( SGLangForwardBatchMetadata, ) @@ -448,5 +448,5 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: # SGLang discovers these multimodal wrappers from this module's `EntryClass`. # They are not covered by `base_model_wrapper.py`, whose generated entries only -# handle the plain causal-LM architectures in `_MODEL_ARCH_SPECS`. +# handle the plain causal-LM architectures in `MODEL_ARCH_SPECS`. EntryClass = [Qwen3_5ForConditionalGeneration, Qwen3_5MoeForConditionalGeneration] diff --git a/atom/plugin/sglang/runtime/__init__.py b/atom/plugin/sglang/runtime/__init__.py new file mode 100644 index 0000000000..f63890418f --- /dev/null +++ b/atom/plugin/sglang/runtime/__init__.py @@ -0,0 +1,20 @@ +"""Runtime utilities for ATOM's SGLang plugin integration.""" + +from atom.plugin.sglang.runtime.context import ( + SGLangForwardBatchMetadata, + bind_current_forward_batch, + get_current_forward_batch, + plugin_runtime_scope, +) +from atom.plugin.sglang.runtime.forward_context import SGLangPluginRuntime +from atom.plugin.sglang.runtime.model_arch import MODEL_ARCH_SPECS, get_model_arch_spec + +__all__ = [ + "MODEL_ARCH_SPECS", + "SGLangForwardBatchMetadata", + "SGLangPluginRuntime", + "bind_current_forward_batch", + "get_current_forward_batch", + "get_model_arch_spec", + "plugin_runtime_scope", +] diff --git a/atom/plugin/sglang/runtime/context.py b/atom/plugin/sglang/runtime/context.py new file mode 100644 index 0000000000..1fffae92f2 --- /dev/null +++ b/atom/plugin/sglang/runtime/context.py @@ -0,0 +1,126 @@ +"""Runtime context helpers for ATOM's SGLang plugin path.""" + +from __future__ import annotations + +from contextlib import contextmanager +from contextvars import ContextVar +from dataclasses import dataclass +from typing import ClassVar, Optional, Union + +from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors + +_RUNTIME_SENTINEL = object() +_current_forward_batch: ContextVar[Optional[ForwardBatch]] = ContextVar( + "atom_sglang_current_forward_batch", default=None +) + + +def get_current_forward_batch(): + return _current_forward_batch.get() + + +@contextmanager +def bind_current_forward_batch(forward_batch: Optional[ForwardBatch]): + token = _current_forward_batch.set(forward_batch) + try: + yield + finally: + _current_forward_batch.reset(token) + + +@contextmanager +def plugin_runtime_scope( + *, + framework: Optional[str] = None, + atom_config=_RUNTIME_SENTINEL, +): + """Temporarily bind process-global ATOM plugin runtime state. + + SGLang target/draft wrappers can coexist during speculative decoding, while + ATOM core still reads process-global framework/config state in some paths. + Keep those globals scoped to one wrapper call and restore them afterwards. + """ + + import atom.config as atom_config_module + import atom.plugin.prepare as plugin_prepare + + prev_framework = plugin_prepare._CURRENT_FRAMEWORK + prev_atom_config = getattr(atom_config_module, "_current_atom_config", None) + + if framework is not None: + plugin_prepare._set_framework_backbone(framework) + if atom_config is not _RUNTIME_SENTINEL: + atom_config_module._current_atom_config = atom_config + + try: + yield + finally: + plugin_prepare._CURRENT_FRAMEWORK = prev_framework + atom_config_module._current_atom_config = prev_atom_config + + +@dataclass(frozen=True) +class SGLangForwardBatchMetadata: + """Small context object for one SGLang model forward.""" + + forward_batch: Optional[ForwardBatch] + pp_proxy_tensors: Optional[PPProxyTensors] = None + save_kv_cache: bool = True + _current: ClassVar[ContextVar[Optional["SGLangForwardBatchMetadata"]]] = ContextVar( + "atom_sglang_current_forward_batch_metadata", + default=None, + ) + + @classmethod + def current(cls) -> Optional["SGLangForwardBatchMetadata"]: + return cls._current.get() + + @classmethod + def build( + cls, + forward_batch: Optional[ + Union[ForwardBatch, "SGLangForwardBatchMetadata"] + ] = None, + *, + pp_proxy_tensors: Optional[PPProxyTensors] = None, + save_kv_cache: Optional[bool] = None, + ) -> Optional["SGLangForwardBatchMetadata"]: + if isinstance(forward_batch, cls): + return forward_batch + if forward_batch is None and pp_proxy_tensors is None and save_kv_cache is None: + return cls.current() + return cls( + forward_batch=forward_batch, + pp_proxy_tensors=pp_proxy_tensors, + save_kv_cache=True if save_kv_cache is None else save_kv_cache, + ) + + @classmethod + @contextmanager + def bind(cls, metadata: Optional["SGLangForwardBatchMetadata"]): + meta_token = cls._current.set(metadata) + batch_token = _current_forward_batch.set( + None if metadata is None else metadata.forward_batch + ) + try: + yield metadata + finally: + _current_forward_batch.reset(batch_token) + cls._current.reset(meta_token) + + @staticmethod + def to_intermediate_tensors( + intermediate_tensors, + metadata: Optional["SGLangForwardBatchMetadata"], + ): + if intermediate_tensors is not None or metadata is None: + return intermediate_tensors + pp_proxy_tensors = metadata.pp_proxy_tensors + if pp_proxy_tensors is None: + return intermediate_tensors + tensors = getattr(pp_proxy_tensors, "tensors", None) + if tensors is None: + return intermediate_tensors + from atom.models.utils import IntermediateTensors + + return IntermediateTensors(dict(tensors)) diff --git a/atom/plugin/sglang/runtime/forward_context.py b/atom/plugin/sglang/runtime/forward_context.py new file mode 100644 index 0000000000..05939cc7cb --- /dev/null +++ b/atom/plugin/sglang/runtime/forward_context.py @@ -0,0 +1,234 @@ +"""Scoped runtime adapter from SGLang batches to ATOM core.""" + +from __future__ import annotations + +import copy +from contextlib import ExitStack +from dataclasses import dataclass, field +from typing import Any, Optional + +import torch +from sglang.srt.model_executor.forward_batch_info import ForwardBatch + +from atom.plugin.sglang.runtime.context import bind_current_forward_batch + + +def _is_dummy_forward(forward_batch: ForwardBatch) -> bool: + """Return whether an SGLang batch represents an empty/idle dummy run.""" + + forward_mode = getattr(forward_batch, "forward_mode", None) + return bool( + forward_mode is not None + and hasattr(forward_mode, "is_idle") + and forward_mode.is_idle() + ) + + +def _pad_dummy_like( + tensor: Optional[torch.Tensor], + *, + length: int, + fill_value: int | float = 0, +) -> Optional[torch.Tensor]: + if tensor is None: + return None + shape = (length, *tensor.shape[1:]) + return torch.full(shape, fill_value, dtype=tensor.dtype, device=tensor.device) + + +def _materialize_atom_dummy_forward( + input_ids: Optional[torch.Tensor], + positions: Optional[torch.Tensor], + input_embeds: Optional[torch.Tensor], + forward_batch: ForwardBatch, +) -> tuple[ + Optional[torch.Tensor], + Optional[torch.Tensor], + Optional[torch.Tensor], + ForwardBatch, +]: + """Convert an empty SGLang IDLE batch into ATOM-style dummy inputs.""" + + if positions is None: + raise RuntimeError("SGLang dummy forward materialization requires positions") + if input_ids is None: + raise RuntimeError("SGLang dummy forward materialization requires input_ids") + + dummy_positions = positions.new_zeros((1,)) + dummy_input_ids = input_ids.new_zeros((1,)) + dummy_input_embeds = _pad_dummy_like(input_embeds, length=1, fill_value=0) + + model_forward_batch = copy.copy(forward_batch) + model_forward_batch.positions = dummy_positions + model_forward_batch.batch_size = 1 + model_forward_batch.seq_lens_sum = 1 + model_forward_batch.seq_lens = forward_batch.seq_lens.new_ones((1,)) + model_forward_batch.seq_lens_cpu = forward_batch.seq_lens_cpu.new_ones((1,)) + + return dummy_input_ids, dummy_positions, dummy_input_embeds, model_forward_batch + + +def _trim_hidden_states_for_output(hidden_states, num_tokens: int): + if torch.is_tensor(hidden_states): + return hidden_states[:num_tokens] + if isinstance(hidden_states, tuple): + return tuple( + tensor[:num_tokens] if torch.is_tensor(tensor) else tensor + for tensor in hidden_states + ) + return hidden_states + + +def _resolve_num_tokens_across_dp( + atom_config: Any, + forward_batch: ForwardBatch, + num_tokens: int, + is_dummy_run: bool, +) -> torch.Tensor: + """Resolve per-DP token counts for ATOM's CPU-side DPMetadata.""" + + global_num_tokens_cpu = getattr(forward_batch, "global_num_tokens_cpu", None) + if global_num_tokens_cpu is not None: + num_tokens_across_dp = torch.tensor( + global_num_tokens_cpu, dtype=torch.int32, device="cpu" + ) + else: + dp_size = atom_config.parallel_config.data_parallel_size + global_num_tokens_gpu = getattr(forward_batch, "global_num_tokens_gpu", None) + global_dp_buffer_len = getattr(forward_batch, "global_dp_buffer_len", None) + is_static_same_shape_batch = ( + global_num_tokens_gpu is not None + and global_dp_buffer_len == num_tokens * dp_size + ) + if not is_static_same_shape_batch: + raise RuntimeError( + "[SGL+ATOM] SGLang dp-attention requires " + "forward_batch.global_num_tokens_cpu unless the batch uses static " + "same-shape DP metadata." + ) + + # Static batches, such as CUDA graph capture batches, may only keep + # global token counts on GPU. Avoid GPU-to-CPU reads here and mirror + # their same-shape layout directly for ATOM's CPU DPMetadata. + num_tokens_across_dp = torch.full( + (dp_size,), num_tokens, dtype=torch.int32, device="cpu" + ) + + if is_dummy_run: + # SGLang reports idle ranks as 0 tokens, but ATOM materializes them + # as one local dummy token so collectives and DPMetadata stay aligned. + dp_rank = atom_config.parallel_config.data_parallel_rank + num_tokens_across_dp[dp_rank] = num_tokens + return num_tokens_across_dp + + +def _set_atom_forward_context( + atom_config: Any, + forward_batch: ForwardBatch, + positions: torch.Tensor, +) -> None: + """Bridge SGLang batch metadata into ATOM's global forward context.""" + + from atom.utils.forward_context import ( + AttentionMetaData, + Context, + set_forward_context, + ) + + forward_mode = forward_batch.forward_mode + # This value is only used by ATOM-side MoE padding in the SGLang wrapper. + max_seqlen_q = 1 if forward_mode.is_decode_or_idle() else 0 + attn_metadata = AttentionMetaData(max_seqlen_q=max_seqlen_q) + batch_size = int(forward_batch.batch_size) + is_dummy_run = _is_dummy_forward(forward_batch) + is_prefill = forward_mode.is_prefill() + num_tokens = int(positions.shape[0]) + + if bool(atom_config.enable_dp_attention): + num_tokens_across_dp = _resolve_num_tokens_across_dp( + atom_config, forward_batch, num_tokens, is_dummy_run + ) + graph_bs = int(torch.max(num_tokens_across_dp).item()) + else: + num_tokens_across_dp = None + graph_bs = num_tokens if is_prefill else batch_size + + context = Context( + positions=positions, + is_prefill=is_prefill, + is_dummy_run=is_dummy_run, + batch_size=batch_size, + graph_bs=graph_bs, + ) + set_forward_context( + attn_metadata=attn_metadata, + atom_config=atom_config, + context=context, + num_tokens=num_tokens, + num_tokens_across_dp=num_tokens_across_dp, + ) + + +def _reset_atom_forward_context() -> None: + from atom.utils.forward_context import reset_forward_context + + reset_forward_context() + + +@dataclass +class SGLangPluginRuntime: + """Scoped adapter for running ATOM model code under SGLang plugin runtime. + + The adapter owns the temporary translation from SGLang's ``ForwardBatch`` to + ATOM's process-local runtime state. Callers should use the normalized + ``input_ids``, ``positions``, ``input_embeds``, and ``forward_batch`` exposed + by this object while inside the context. + """ + + atom_config: Any + forward_batch: ForwardBatch + positions: torch.Tensor + input_ids: Optional[torch.Tensor] = None + input_embeds: Optional[torch.Tensor] = None + set_forward_context: bool = True + _original_forward_batch: ForwardBatch = field(init=False, repr=False) + _is_dummy_run: bool = field(init=False, default=False) + _exit_stack: ExitStack = field(init=False, repr=False) + + def __enter__(self) -> "SGLangPluginRuntime": + self._original_forward_batch = self.forward_batch + self._is_dummy_run = _is_dummy_forward(self.forward_batch) + + if self._is_dummy_run: + ( + self.input_ids, + self.positions, + self.input_embeds, + self.forward_batch, + ) = _materialize_atom_dummy_forward( + self.input_ids, + self.positions, + self.input_embeds, + self.forward_batch, + ) + + self._exit_stack = ExitStack() + self._exit_stack.enter_context(bind_current_forward_batch(self.forward_batch)) + if self.set_forward_context: + _set_atom_forward_context( + self.atom_config, + self.forward_batch, + self.positions, + ) + self._exit_stack.callback(_reset_atom_forward_context) + return self + + def __exit__(self, exc_type, exc, tb) -> None: + self._exit_stack.close() + + def trim_output(self, hidden_states): + """Map ATOM-visible outputs back to SGLang-visible token count.""" + + if self._is_dummy_run: + return _trim_hidden_states_for_output(hidden_states, 0) + return hidden_states diff --git a/atom/plugin/sglang/runtime/model_arch.py b/atom/plugin/sglang/runtime/model_arch.py new file mode 100644 index 0000000000..1ddd3772c9 --- /dev/null +++ b/atom/plugin/sglang/runtime/model_arch.py @@ -0,0 +1,22 @@ +"""Runtime behavior flags for SGLang plugin model wrappers.""" + +from __future__ import annotations + +from dataclasses import dataclass + + +@dataclass(frozen=True) +class ModelArchSpec: + wrapper_binds_gdn_context: bool = False + apply_deepseek_patch: bool = False + + +MODEL_ARCH_SPECS = { + "DeepseekV3ForCausalLM": ModelArchSpec(apply_deepseek_patch=True), + "Qwen3MoeForCausalLM": ModelArchSpec(), + "Qwen3NextForCausalLM": ModelArchSpec(wrapper_binds_gdn_context=True), +} + + +def get_model_arch_spec(model_arch: str) -> ModelArchSpec: + return MODEL_ARCH_SPECS.get(model_arch, ModelArchSpec()) diff --git a/atom/plugin/vllm/attention_backend/attention_gdn.py b/atom/plugin/vllm/attention_backend/attention_gdn.py index bff5370215..87a2f2f9f9 100644 --- a/atom/plugin/vllm/attention_backend/attention_gdn.py +++ b/atom/plugin/vllm/attention_backend/attention_gdn.py @@ -22,6 +22,7 @@ from atom.model_ops.fla_ops.fused_sigmoid_gating import ( fused_sigmoid_gating_delta_rule_update, ) + from atom.utils import envs from torch import nn @@ -75,17 +76,19 @@ def fused_gdn_gating_kernel( dt_bias, seq_len, NUM_HEADS: tl.constexpr, + stride_a_batch, + stride_b_batch, beta: tl.constexpr, threshold: tl.constexpr, BLK_HEADS: tl.constexpr, ): i_b, i_s, i_d = tl.program_id(0), tl.program_id(1), tl.program_id(2) head_off = i_d * BLK_HEADS + tl.arange(0, BLK_HEADS) - off = i_b * seq_len * NUM_HEADS + i_s * NUM_HEADS + head_off + out_off = i_b * seq_len * NUM_HEADS + i_s * NUM_HEADS + head_off mask = head_off < NUM_HEADS blk_A_log = tl.load(A_log + head_off, mask=mask) - blk_a = tl.load(a + off, mask=mask) - blk_b = tl.load(b + off, mask=mask) + blk_a = tl.load(a + i_b * stride_a_batch + head_off, mask=mask) + blk_b = tl.load(b + i_b * stride_b_batch + head_off, mask=mask) blk_bias = tl.load(dt_bias + head_off, mask=mask) # If the model is loaded in fp16, without the .float() here, A might be -inf x = blk_a.to(tl.float32) + blk_bias.to(tl.float32) @@ -93,11 +96,13 @@ def fused_gdn_gating_kernel( beta * x <= threshold, (1 / beta) * tl.log(1 + tl.exp(beta * x)), x ) blk_g = -tl.exp(blk_A_log.to(tl.float32)) * softplus_x - tl.store(g + off, blk_g.to(g.dtype.element_ty), mask=mask) + tl.store(g + out_off, blk_g.to(g.dtype.element_ty), mask=mask) # compute beta_output = sigmoid(b) blk_beta_output = tl.sigmoid(blk_b.to(tl.float32)) tl.store( - beta_output + off, blk_beta_output.to(beta_output.dtype.element_ty), mask=mask + beta_output + out_off, + blk_beta_output.to(beta_output.dtype.element_ty), + mask=mask, ) @@ -129,6 +134,8 @@ def fused_gdn_gating( dt_bias, seq_len, num_heads, + a.stride(0), + b.stride(0), beta, threshold, 8, @@ -206,6 +213,7 @@ def forward( if attn_metadata is None: # V1 profile run + core_attn_out.zero_() return core_attn_out assert isinstance(attn_metadata, dict) @@ -378,7 +386,13 @@ def forward( ssm_state[non_spec_state_indices_tensor] = last_recurrent_state.to( ssm_state.dtype ) - core_attn_out[:num_actual_tokens] = core_attn_out_non_spec.squeeze(0) + # Only write directly when there are no spec tokens. With spec + # decode active, mixed_qkv was index_select'd by non_spec_token_indx + # so core_attn_out_non_spec has fewer rows than num_actual_tokens. + # The merge below (index_copy_) handles the scatter back to the + # correct slot positions. + if spec_sequence_masks is None: + core_attn_out[:num_actual_tokens] = core_attn_out_non_spec.squeeze(0) elif attn_metadata.num_decodes > 0: o = core_attn_out[: attn_metadata.num_decode_tokens] if USE_FLYDSL_GDR: @@ -437,4 +451,8 @@ def forward( elif spec_sequence_masks is not None: core_attn_out[:num_actual_tokens] = core_attn_out_spec.squeeze(0) + # Zero padding tail for CUDA graph replay safety + if num_actual_tokens < core_attn_out.shape[0]: + core_attn_out[num_actual_tokens:].zero_() + return core_attn_out diff --git a/atom/plugin/vllm/attention_backend/mla_sparse.py b/atom/plugin/vllm/attention_backend/mla_sparse.py index 595dbad965..d8d8e27b9f 100644 --- a/atom/plugin/vllm/attention_backend/mla_sparse.py +++ b/atom/plugin/vllm/attention_backend/mla_sparse.py @@ -23,7 +23,7 @@ class AiterMLASparseBackend(AttentionBackend): @staticmethod def get_name() -> str: - return "ROCM_AITER_MLA_SPARSE" if not is_plugin_mode() else "CUSTOM" + return "ROCM_AITER_MLA_SPARSE" if not is_plugin_mode() else "CUSTOM_MLA_SPARSE" @staticmethod def get_builder_cls() -> Type["AiterMLASparseMetadataBuilder"]: @@ -59,7 +59,11 @@ class AiterMLASparseIndexerBackend(AttentionBackend): @staticmethod def get_name() -> str: - return "ROCM_AITER_MLA_SPARSE_INDEXER" if not is_plugin_mode() else "CUSTOM" + return ( + "ROCM_AITER_MLA_SPARSE_INDEXER" + if not is_plugin_mode() + else "CUSTOM_MLA_SPARSE_INDEXER" + ) @staticmethod def get_builder_cls() -> Type["AiterMLASparseIndexerMetadataBuilder"]: diff --git a/atom/plugin/vllm/model_wrapper.py b/atom/plugin/vllm/model_wrapper.py index 1eeee3ddb9..fbdd5e94dd 100644 --- a/atom/plugin/vllm/model_wrapper.py +++ b/atom/plugin/vllm/model_wrapper.py @@ -1,6 +1,8 @@ from collections.abc import Iterable +import functools import importlib +import types import torch import torch.nn as nn from aiter.dist.parallel_state import ( @@ -33,7 +35,9 @@ logger = logging.getLogger("atom") - +_MTP_MASK_INPUT_ARCH: set[str] = { + "DeepSeekMTPModel", +} _ATOM_MODEL_CLASSES: dict[str, str] = { "LlamaForCausalLM": "atom.models.llama:LlamaForCausalLM", "Qwen3ForCausalLM": "atom.models.qwen3:Qwen3ForCausalLM", @@ -43,7 +47,10 @@ "DeepseekV32ForCausalLM": "atom.models.deepseek_v2:DeepseekV3ForCausalLM", "Glm4MoeForCausalLM": "atom.models.glm4_moe:Glm4MoeForCausalLM", "GlmMoeDsaForCausalLM": "atom.models.deepseek_v2:GlmMoeDsaForCausalLM", + "DeepSeekMTPModel": "atom.models.deepseek_mtp:DeepSeekMTP", + "Glm4MoeMTPModel": "atom.models.glm4_moe_mtp:Glm4MoeMTP", "Qwen3NextForCausalLM": "atom.models.qwen3_next:Qwen3NextForCausalLM", + "Qwen3NextMTP": "atom.models.qwen3_next_mtp:Qwen3NextMTP", "Qwen3_5MoeForConditionalGeneration": "atom.models.qwen3_5:Qwen3_5MoeForConditionalGeneration_", "Qwen3_5ForConditionalGeneration": "atom.models.qwen3_5:Qwen3_5ForConditionalGeneration_", "KimiK25ForConditionalGeneration": "atom.plugin.vllm.models.kimi_k25:KimiK25ForConditionalGeneration_", @@ -73,12 +80,52 @@ def _prepare_env(atom_config) -> None: init_aiter_dist(config=atom_config) +def _safe_get_first_arch(config_like) -> str | None: + if config_like is None: + return None + architectures = getattr(config_like, "architectures", None) + if isinstance(architectures, list) and len(architectures) > 0: + return architectures[0] + return None + + +def _select_model_arch(vllm_config: VllmConfig) -> str: + model_arch = _safe_get_first_arch(getattr(vllm_config, "model_config", None)) + if model_arch is None: + raise ValueError("Cannot determine model architecture from vLLM model_config") + speculative_config = getattr(vllm_config, "speculative_config", None) + draft_model_config = getattr(speculative_config, "draft_model_config", None) + draft_arch = _safe_get_first_arch(draft_model_config) + if draft_arch is None: + return model_arch + model_tag = None + try: + from vllm.compilation import backends as vllm_backends + + model_tag = getattr(vllm_backends, "model_tag", None) + except Exception: + pass + if model_tag is None: + model_tag = getattr( + getattr(vllm_config, "compilation_config", None), "model_tag", None + ) + if model_tag in {"eagle_head", "draft_model", "drafter"}: + logger.info( + f"Use draft model architecture {draft_arch} for speculative tag {model_tag}" + ) + return draft_arch + return model_arch + + class ATOMModelBase(nn.Module, VllmModel, SupportsQuant, SupportsPP): + # forced_model_arch: str | None = None + def __init_subclass__(cls, *args, **kwargs): super().__init_subclass__(*args, **kwargs) def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() + from atom.config import get_current_atom_config _set_framework_backbone("vllm") @@ -98,11 +145,21 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.ignore_unexpected_suffixes: list[str] = [] self.vllm_config = vllm_config - self.atom_config = generate_atom_config_for_plugin_mode(vllm_config) - + self.is_mtp = False + speculative_config = getattr(vllm_config, "speculative_config", None) + if speculative_config is not None: + spec_method = speculative_config.method + self.is_mtp = spec_method == "mtp" + + main_model_arch = vllm_config.model_config.architectures[0] + model_arch = _select_model_arch(vllm_config) + self.is_mtp_draft_model = self.is_mtp and model_arch != main_model_arch + if self.is_mtp_draft_model: + self.atom_config = get_current_atom_config() + else: + self.atom_config = generate_atom_config_for_plugin_mode(vllm_config) + self.model_arch = model_arch _prepare_env(atom_config=self.atom_config) - - model_arch = vllm_config.model_config.architectures[0] model_cls = _get_atom_model_cls(model_arch) module_remapping = getattr(model_cls, "packed_modules_mapping", {}) weights_mapper = getattr(model_cls, "hf_to_atom_mapper", {}) @@ -116,6 +173,13 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): # to ATOM's format. It is invoked in ATOM's model_runner initialization, but # lacks correspondences in vLLM. So we invoke the translation here for vLLM OOT. exclude_mapping = getattr(model_cls, "quant_exclude_name_mapping", {}) + # add exclude mapping for mtp layer of GLM5. + if model_arch != main_model_arch and main_model_arch == "GlmMoeDsaForCausalLM": + exclude_mapping.update( + { + "indexers_proj": "indexer.weights_proj", + } + ) if exclude_mapping and self.atom_config.quant_config is not None: self.atom_config.quant_config.apply_exclude_name_mapping(exclude_mapping) @@ -126,13 +190,17 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): logger.info(f"Construct ATOM model {model_arch} for vLLM plugin mode") self.model = model_cls(self.atom_config) + if model_arch in _MTP_MASK_INPUT_ARCH: + self._adapt_mtp_layers_for_vllm() + # Mirror nested attributes required by vLLM speculative decoding. + self._expose_spec_decode_attrs() + # For sparse MLA, register the Indexer's DeepseekV32IndexerCache as # a virtual subclass of vLLM's AttentionLayerBase so vLLM can discover # it and allocate KV cache. self._register_indexer_caches_with_vllm() if self.model is None: - model_arch = vllm_config.model_config.architectures[0] raise ValueError( f"The model {model_arch} is not supported by model impl backend atom" ) @@ -141,6 +209,58 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.pp_group = get_pp_group() self.tp_group = get_tp_group() + # Attributes whose writes on the outer model must propagate to the + # inner model so vLLM's weight-sharing reaches the forward path. + _WEIGHT_SHARED_ATTRS = frozenset({"embed_tokens", "embedding", "lm_head"}) + + def _expose_spec_decode_attrs(self) -> None: + """Bridge the extra nesting level between vLLM and ATOM for spec decode. + + ATOM wraps the HF model with one extra level: + vLLM sees: wrapper.model (DeepSeekMTP) + forward uses: .model (DeepSeekMultiTokenPredictor) + + vLLM's EagleSpeculator reads/writes embed_tokens, lm_head, layers on + the outer model. The forward path reads them from the inner model. + + We need two things: + 1. Mirror inner → outer so vLLM can discover the attrs. + 2. When vLLM later *replaces* embed_tokens / lm_head with shared + target-model weights, propagate the write to the inner model + so the forward path picks up the shared tensor. + """ + model = self.model + inner = getattr(model, "model", None) + if inner is None: + if hasattr(model, "lm_head") and not hasattr(self, "lm_head"): + self.lm_head = model.lm_head + return + + # (1) Mirror: make attrs visible on the outer model for vLLM discovery. + for attr in (*self._WEIGHT_SHARED_ATTRS, "layers"): + if not hasattr(model, attr) and hasattr(inner, attr): + setattr(model, attr, getattr(inner, attr)) + + if not hasattr(self, "lm_head") and hasattr(model, "lm_head"): + self.lm_head = model.lm_head + + # (2) Propagate: future writes on the outer model sync to the inner + # model. We create a one-off subclass so the hook only affects + # this particular draft-model instance, not the base class. + shared = self._WEIGHT_SHARED_ATTRS + base_setattr = model.__class__.__setattr__ + + def _syncing_setattr(self_model, name, value): + base_setattr(self_model, name, value) + if name in shared and hasattr(inner, name): + base_setattr(inner, name, value) + + model.__class__ = type( + model.__class__.__name__, + (model.__class__,), + {"__setattr__": _syncing_setattr}, + ) + def _register_indexer_caches_with_vllm(self): """Register DeepseekV32IndexerCache instances with vLLM so that: 1. vLLM discovers them via isinstance(AttentionLayerBase) for KV cache @@ -197,8 +317,7 @@ def _register_indexer_caches_with_vllm(self): if prefix not in vllm_sfc: vllm_sfc[prefix] = module logger.info( - f"Registered indexer cache in vLLM static_forward_context: " - f"{prefix}" + f"Registered indexer cache in vLLM static_forward_context: {prefix}" ) else: logger.warning( @@ -206,6 +325,51 @@ def _register_indexer_caches_with_vllm(self): f"static_forward_context, skipping" ) + def _adapt_mtp_layers_for_vllm(self) -> None: + """Install vLLM-only MTP input masking without changing model code.""" + if not self.is_mtp_draft_model: + return + + inner_model = getattr(self.model, "model", None) + layers = ( + getattr(inner_model, "layers", None) if inner_model is not None else None + ) + if layers is None: + return + + layer_iter = layers.values() if isinstance(layers, nn.ModuleDict) else layers + for layer in layer_iter: + if getattr(layer, "_atom_vllm_mtp_masked", False): + continue + + layer.forward = types.MethodType( + self._make_vllm_mtp_layer_forward(layer.forward), + layer, + ) + layer._atom_vllm_mtp_masked = True + + @staticmethod + def _make_vllm_mtp_layer_forward(original_forward): + @functools.wraps(original_forward) + def masked_forward( + self_layer, + input_ids, + positions, + previous_hidden_states, + inputs_embeds, + spec_step_index=0, + ): + inputs_embeds = torch.where(positions.unsqueeze(-1) == 0, 0, inputs_embeds) + return original_forward( + input_ids, + positions, + previous_hidden_states, + inputs_embeds, + spec_step_index, + ) + + return masked_forward + def forward( self, input_ids: torch.Tensor | None, @@ -221,7 +385,12 @@ def forward( # pass positions from vLLM to OOT execution path via vLLM's per-forward context if is_forward_context_available(): - get_vllm_forward_context().additional_kwargs["atom_positions"] = positions + forward_context = get_vllm_forward_context() + forward_context.additional_kwargs["atom_positions"] = positions + # set atom_config into vLLM forward_context in order to + # make sure main model and draft model can get their specific + # static_forward_context from their own atom_config + forward_context.additional_kwargs["atom_config"] = self.atom_config elif "positions" in self.atom_config.compilation_config.static_forward_context: buf = self.atom_config.compilation_config.static_forward_context[ "positions" @@ -235,7 +404,6 @@ def forward( inputs_embeds=inputs_embeds, **model_kwargs, ) - if not self.pp_group.is_last_rank: return IntermediateTensors({"hidden_states": hidden_states}) @@ -248,8 +416,29 @@ def load_weights( # prevent circular import from atom.model_loader.loader import load_model_in_plugin_mode + is_mtp_draft_model = self.model_arch in { + "DeepSeekMTPModel", + "Qwen3NextMTP", + "Glm4MoeMTPModel", + } + draft_hf_config = None + if is_mtp_draft_model: + draft_model_config = getattr( + getattr(self.atom_config, "speculative_config", None), + "draft_model_config", + None, + ) + if draft_model_config is not None: + draft_hf_config = getattr( + draft_model_config, "hf_config", draft_model_config + ) + loaded_weights_record = load_model_in_plugin_mode( - model=self.model, config=self.atom_config, prefix="model." + model=self.model, + config=self.atom_config, + prefix="model.", + spec_decode=is_mtp_draft_model, + hf_config_override=draft_hf_config, ) return loaded_weights_record @@ -270,7 +459,6 @@ class ATOMMoEForCausalLM(ATOMModelBase, VllmModelForTextGeneration): ... class ATOMForConditionalGeneration( ATOMModelBase, VllmModelForTextGeneration, SupportsMultiModal, SupportsMRoPE ): - @classmethod def get_placeholder_str(cls, modality: str, i: int) -> str | None: """ diff --git a/atom/plugin/vllm/models/kimi_k25.py b/atom/plugin/vllm/models/kimi_k25.py index 47cfb9ce22..f3b16dd641 100644 --- a/atom/plugin/vllm/models/kimi_k25.py +++ b/atom/plugin/vllm/models/kimi_k25.py @@ -80,6 +80,7 @@ def __init__( ) else: self.norm = PPMissingLayer() + self.aux_hidden_state_layers: tuple[int, ...] = tuple() self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( ["hidden_states", "residual"], config.hidden_size ) @@ -173,6 +174,10 @@ class KimiK25ForConditionalGeneration_(vLLMKimiK25): "gate_proj": ("gate_up_proj", 0), "up_proj": ("gate_up_proj", 1), } + quant_exclude_name_mapping = { + "language_model.model.": "model.language_model.model.", + "language_model.lm_head": "model.language_model.lm_head", + } hf_to_atom_mapper = WeightsMapper( orig_to_new_prefix={ "model.visual.": "visual.", @@ -187,7 +192,11 @@ class KimiK25ForConditionalGeneration_(vLLMKimiK25): def __init__(self, atom_config: Config, prefix: str = "model"): # protocols have not __init__ method, so we need to use nn.Module.__init__ nn.Module.__init__(self) - config: KimiK25Config = atom_config.hf_config + hf_config = getattr(atom_config, "hf_config", None) + assert hf_config is not None, "hf_config is not found in atom_config" + vision_config = getattr(hf_config, "vision_config", None) + text_config = getattr(hf_config, "text_config", None) + config = KimiK25Config(vision_config, text_config) vllm_config = atom_config.plugin_config.vllm_config # quant_config from vLLM ignores exclude_layers in model's quantization config diff --git a/atom/plugin/vllm/register.py b/atom/plugin/vllm/register.py index 32308c1e4a..25c75bba55 100644 --- a/atom/plugin/vllm/register.py +++ b/atom/plugin/vllm/register.py @@ -5,6 +5,7 @@ from atom.plugin.prepare import _set_framework_backbone from atom.utils import envs from atom.plugin.vllm.mla_patch import patch_vllm_mla_attention +from atom.plugin.vllm.spec_decode_patch import apply_vllm_spec_decode_patch logger = logging.getLogger("atom") @@ -27,7 +28,10 @@ "DeepseekV32ForCausalLM": ATOM_MOE_CAUSAL_LM_MODEL_WRAPPER, "Glm4MoeForCausalLM": ATOM_MOE_CAUSAL_LM_MODEL_WRAPPER, "GlmMoeDsaForCausalLM": ATOM_MOE_CAUSAL_LM_MODEL_WRAPPER, + "DeepSeekMTPModel": ATOM_MOE_CAUSAL_LM_MODEL_WRAPPER, + "Glm4MoeMTPModel": ATOM_MOE_CAUSAL_LM_MODEL_WRAPPER, "Qwen3NextForCausalLM": "atom.models.qwen3_next:Qwen3NextForCausalLMVllm", + "Qwen3NextMTP": ATOM_MOE_CAUSAL_LM_MODEL_WRAPPER, "Qwen3_5ForConditionalGeneration": "atom.models.qwen3_5:Qwen3_5ForConditionalGeneration", "Qwen3_5MoeForConditionalGeneration": "atom.models.qwen3_5:Qwen3_5MoeForConditionalGeneration", "KimiK25ForConditionalGeneration": "atom.plugin.vllm.models.kimi_k25:KimiK25ForConditionalGeneration", @@ -126,6 +130,7 @@ def register_model() -> None: _patch_vllm_attention_process_weights_after_loading(Attention) _patch_vllm_attention_process_weights_after_loading(MLAAttention) + apply_vllm_spec_decode_patch() # Patch vLLM graph_capture to also enter aiter's ca_comm.capture(), # avoiding hipMemcpyAsync in fused_allreduce_rmsnorm when model uses aiter collectives diff --git a/atom/plugin/vllm/spec_decode_patch.py b/atom/plugin/vllm/spec_decode_patch.py new file mode 100644 index 0000000000..857af6874e --- /dev/null +++ b/atom/plugin/vllm/spec_decode_patch.py @@ -0,0 +1,35 @@ +import functools +import logging + +logger = logging.getLogger("atom") + + +def apply_vllm_spec_decode_patch() -> None: + """Patch vLLM speculative decoding for ATOM metadata compatibility.""" + try: + from atom.utils.forward_context import ( + AttentionMetaData as AtomAttentionMetaData, + ) + from vllm.v1.spec_decode.eagle import SpecDecodeBaseProposer + except Exception as exc: + logger.debug("ATOM spec-decode patch skipped: %s", exc) + return + + original_init = SpecDecodeBaseProposer.__init__ + if getattr(original_init, "_atom_allowed_attn_types_patched", False): + return + + @functools.wraps(original_init) + def wrapped_init(self, *args, **kwargs): + original_init(self, *args, **kwargs) + allowed = getattr(self, "allowed_attn_types", None) + if allowed is not None and AtomAttentionMetaData not in allowed: + self.allowed_attn_types = (*allowed, AtomAttentionMetaData) + + setattr(wrapped_init, "_atom_allowed_attn_types_patched", True) + SpecDecodeBaseProposer.__init__ = wrapped_init + + logger.info( + "ATOM plugin: patched vLLM speculative decoder for " + "ATOM attention-metadata compatibility." + ) diff --git a/atom/spec_decode/eagle.py b/atom/spec_decode/eagle.py index aa3aa4c8b1..5d86c62fd0 100644 --- a/atom/spec_decode/eagle.py +++ b/atom/spec_decode/eagle.py @@ -1,10 +1,13 @@ +import copy import logging +from typing import Optional import numpy as np import torch import torch.nn as nn +from aiter import dtypes from aiter.dist.parallel_state import get_pp_group -from atom.config import CompilationLevel, Config +from atom.config import CompilationLevel, Config, KVCacheTensor from atom.model_loader.loader import load_model from atom.utils import CpuGpuBuffer, resolve_obj_by_qualname from atom.utils.forward_context import SpecDecodeMetadata, get_forward_context @@ -15,12 +18,141 @@ support_eagle_model_arch_dict = { "DeepSeekMTPModel": "atom.models.deepseek_mtp.DeepSeekMTP", + "DeepseekV4MTPModel": "atom.models.deepseek_v4_mtp.DeepseekV4MTP", "Qwen3NextMTPModel": "atom.models.qwen3_next_mtp.Qwen3NextMTP", "MiMoV2FlashMTPModel": "atom.models.mimo_v2_flash_mtp.MiMoV2FlashMTP", "Qwen3_5MTPModel": "atom.models.qwen3_5_mtp.Qwen3_5MTP", + "Eagle3LlamaModel": "atom.models.eagle3_llama.Eagle3LlamaModel", } +class Eagle3DraftBuilder: + """KV cache subsystem for an Eagle3 MHA draft alongside a non-MHA target. + + Implements the same subset of `AttentionMetadataBuilder` hooks that + ModelRunner consults during KV pool sizing and per-module binding — + `compute_block_bytes`, `allocate_kv_cache_tensors`, and + `build_kv_cache_tensor` — so the draft's independent non-MLA cache + fits the post-#659 builder protocol without leaking into the target's + builder. The draft does NOT drive prepare_decode/prepare_prefill; + it piggybacks on the target builder's metadata flow during propose. + """ + + def __init__(self, model_runner, draft_hf): + self.model_runner = model_runner + self.draft_hf = draft_hf + self.block_size = model_runner.block_size + self.num_kv_heads = draft_hf.num_key_value_heads // model_runner.world_size + self.num_layers = draft_hf.num_hidden_layers + self.head_dim = draft_hf.head_dim + self._next_layer_id = 0 # consumed by build_kv_cache_tensor + self.num_blocks = 0 # set in allocate_kv_cache_tensors + + def compute_block_bytes(self) -> int: + """Per-block bytes for the draft's independent non-MLA KV cache.""" + kv_dtype_size = dtypes.d_dtypes[ + self.model_runner.config.kv_cache_dtype + ].itemsize + bb = ( + 2 + * self.num_layers + * self.block_size + * self.num_kv_heads + * self.head_dim + * kv_dtype_size + ) + if self.model_runner.config.kv_cache_dtype == "fp8": + # fp8 KV cache needs an extra per-(layer, block, kv_head) scale + # tensor (one fp32 per element) to dequantize fp8 → bf16 at + # attention time. Reserve that space alongside the cache. + bb += ( + 2 + * self.num_layers + * self.block_size + * self.num_kv_heads + * dtypes.fp32.itemsize + ) + return bb + + def allocate_kv_cache_tensors(self, num_kv_heads, num_draft_layers) -> dict: + """Allocate the draft's [2, L, blocks, block_size, kv_heads, head_dim] + cache and matching fp32 scale; ModelRunner setattr's both onto itself + under namespaced keys so they don't collide with the target builder's + `kv_cache` / `kv_scale`. + """ + runner = self.model_runner + config = runner.config + # Draft's block budget scales with the target pool: same total token + # capacity, just paged at the draft's own block size. + self.num_blocks = ( + config.num_kvcache_blocks * runner.block_size // self.block_size + ) + cache = torch.zeros( + 2, + self.num_layers, + self.num_blocks, + self.block_size, + self.num_kv_heads, + self.head_dim, + dtype=dtypes.d_dtypes[config.kv_cache_dtype], + device="cuda", + ) + scale = torch.zeros( + 2, + self.num_layers, + self.num_blocks, + self.num_kv_heads, + self.block_size, + dtype=dtypes.fp32, + device="cuda", + ) + logger.info(f"Allocated Eagle3 draft KV cache: {cache.shape}") + return {"eagle3_kv_cache": cache, "eagle3_kv_scale": scale} + + def build_kv_cache_tensor(self, layer_id: int, module): + """Bind one Eagle3 draft attention module to its slice of the + independent draft KV cache. Returns None for non-MHA modules so + ModelRunner falls through to the target builder. + """ + if not ( + hasattr(module, "base_attention") + and hasattr(module, "use_mla") + and not module.use_mla + ): + return None + runner = self.model_runner + idx = self._next_layer_id + self._next_layer_id += 1 + cache = runner.eagle3_kv_cache + x = 16 // cache.element_size() + k_cache = cache[0, idx].view( + self.num_blocks, + self.num_kv_heads, + self.head_dim // x, + self.block_size, + x, + ) + v_cache = cache[1, idx].view( + self.num_blocks, + self.num_kv_heads, + self.head_dim, + self.block_size, + ) + module.max_model_len = runner.config.max_model_len + if runner.config.kv_cache_dtype == "fp8": + module.k_scale = runner.eagle3_kv_scale[0, idx] + module.v_scale = runner.eagle3_kv_scale[1, idx] + module.k_cache = k_cache + module.v_cache = v_cache + return KVCacheTensor( + layer_num=layer_id, + k_cache=k_cache, + v_cache=v_cache, + k_scale=getattr(module, "k_scale", None), + v_scale=getattr(module, "v_scale", None), + ) + + class EagleProposer: def __init__( @@ -49,7 +181,30 @@ def __init__( self.device = device draft_model_hf_config = self.speculative_config.draft_model_hf_config model_class = resolve_obj_by_qualname(support_eagle_model_arch_dict[draft_model_hf_config.architectures[0]]) # type: ignore - self.model = model_class(self.config) + + if self.speculative_config.method == "eagle3": + # Eagle3 draft model has its own architecture (Llama, not MLA), + # so it must be constructed with the draft model's hf_config. + # Also disable torch.compile for the draft model to avoid + # Dynamo tracing issues with the separate KV cache binding. + draft_atom_config = copy.deepcopy(atom_config) + draft_atom_config.hf_config = draft_model_hf_config + draft_atom_config.compilation_config.level = CompilationLevel.NO_COMPILATION + # Draft attention layer_num must continue from the target model's + # layer count so it maps to the correct kv_cache_data entry. + self.model = model_class( + draft_atom_config, + layer_offset=atom_config.hf_config.num_hidden_layers, + ) + # Attach the draft's KV-cache builder to the runner. ModelRunner + # consults `runner.eagle3_draft_builder` from `_compute_block_bytes` + # / `allocate_kv_cache` to size + allocate + bind the draft's + # independent non-MLA cache through the standard builder protocol. + runner.eagle3_draft_builder = Eagle3DraftBuilder( + runner, draft_model_hf_config + ) + else: + self.model = model_class(self.config) i32_kwargs = {"dtype": torch.int32, "device": self.device} i64_kwargs = {"dtype": torch.int64, "device": self.device} @@ -78,6 +233,23 @@ def _share_if_not_loaded( setattr(owner, attr, source) def load_model(self, target_model: nn.Module) -> None: + if self.speculative_config.method == "eagle3": + # Eagle3: load from a separate draft model checkpoint with + # independent embed_tokens and lm_head (no sharing). + load_model( + self.model, + self.speculative_config.model, + self.speculative_config.draft_model_hf_config, + self.config.load_dummy, + False, + ) + logger.info( + "Eagle3 draft model loaded from %s (independent embed/lm_head)", + self.speculative_config.model, + ) + return + + # MTP: load from the target model checkpoint and share embeddings/lm_head. loaded = load_model( self.model, self.config.model, @@ -89,6 +261,15 @@ def load_model(self, target_model: nn.Module) -> None: # Resolve the base model (unwrap multimodal wrapper if present) target_base = getattr(target_model, "language_model", target_model) + # Model-specific share hook escape valve. Models whose embed/lm_head + # naming doesn't match the standard `model.embed_tokens` / + # `lm_head` convention (e.g. DeepSeek-V4 uses `model.embed` / + # `model.head`) implement `share_with_target(target_base)` to do + # their own setattr-rebinding and short-circuit the default path. + if hasattr(self.model, "share_with_target"): + self.model.share_with_target(target_base, loaded) + return + # Share embed_tokens with the target model if ( get_pp_group().world_size == 1 @@ -101,11 +282,6 @@ def load_model(self, target_model: nn.Module) -> None: ) del self.model.model.embed_tokens self.model.model.embed_tokens = target_base.model.embed_tokens - else: - logger.info( - "The EAGLE head's vocab embedding will be loaded separately" - " from the target model." - ) # Share lm_head from target if not loaded from checkpoint. # Case 1: per-layer shared_head.head (DeepSeek MTP) @@ -148,6 +324,7 @@ def propose( num_reject_tokens: torch.Tensor, next_token_ids: torch.Tensor, last_token_indices: torch.Tensor, + aux_hidden_states: Optional[list[torch.Tensor]] = None, ) -> torch.Tensor: forward_context = get_forward_context() @@ -161,21 +338,55 @@ def propose( # input_ids[last_token_indices] = next_token_ids input_ids.scatter_(0, last_token_indices, next_token_ids) positions = target_positions + 1 - hidden_states = target_hidden_states + + # Eagle3: project concatenated aux hidden states through fc + if aux_hidden_states is not None: + concat_aux = torch.cat(aux_hidden_states, dim=-1) + hidden_states = self.model.combine_hidden_states(concat_aux) + else: + hidden_states = target_hidden_states draft_token_ids = torch.empty( bs, self.mtp_k, dtype=next_token_ids.dtype, device=next_token_ids.device ) - # return draft_token_ids.fill_(1) # for debug var = self.runner.forward_vars - use_mla = self.runner.use_mla + target_uses_mla = self.runner.use_mla + # Eaale3 only support mha currently + draft_uses_mha = hasattr(self.runner, "eagle3_draft_builder") + + # Eagle3 MLA: re-slice slot_mapping to len(input_ids). + # Target's MLA prepare_decode sized it + # to bs*max_q_len; after rejection len(input_ids) may be smaller, + # and the MHA cache-write kernel asserts slot_mapping <= q. + # Other fields (block_tables, context_lens, slot_mapping values + # themselves) are already in a draft-compatible format because + # MLA's prepare_decode uses the same runner.block_size as the draft. + if draft_uses_mha: + attn_metadata.slot_mapping = var["slot_mapping"].gpu[: len(input_ids)] + + # Backends that expose flat per-seq kv_indices/kv_indptr (MLA, MHA) + # wire them through eagle's mid-step block; V4 has block_tables + + # context_lens instead (its v4_kv_indices_{swa,csa,hca} are per-token + # non-equivalent). Hoisted out of the loop so the value is bound for + # every iteration (used at i>=1 too, even though i==0 sets it). + has_flat_kv = "kv_indices" in var + for i in range(self.mtp_k): with record_function(f"draft[{i}/{self.mtp_k} bs={bs}]"): - ret_hidden_states = self.model( + model_output = self.model( input_ids=input_ids, positions=positions, hidden_states=hidden_states, ) + # Eagle3 draft (the only draft_uses_mha case under narrow + # semantics) returns (post_norm, pre_norm); MTP drafts return + # a single hidden tensor. + if draft_uses_mha: + ret_hidden_states, ret_hidden_prenorm = model_output + else: + ret_hidden_states = model_output + ret_hidden_prenorm = None + sample_hidden_states = ( torch.index_select(ret_hidden_states, 0, last_token_indices) if i == 0 @@ -194,25 +405,30 @@ def propose( if i == 0: i0_max_seqlen_q = attn_metadata.max_seqlen_q attn_metadata.max_seqlen_q = 1 - kv_indptr = var["kv_indptr"].gpu[: bs + 1] - kv_indices = var["kv_indices"].gpu slot_mapping = var["slot_mapping"].gpu[ : bs * attn_metadata.max_seqlen_q ] cu_seqlens_q = var["cu_seqlens_q"].gpu[: bs + 1] - attn_metadata.kv_indptr = kv_indptr - attn_metadata.kv_indices = kv_indices attn_metadata.cu_seqlens_q = cu_seqlens_q attn_metadata.slot_mapping = slot_mapping - if use_mla: + if has_flat_kv: + kv_indptr = var["kv_indptr"].gpu[: bs + 1] + kv_indices = var["kv_indices"].gpu + attn_metadata.kv_indptr = kv_indptr + attn_metadata.kv_indices = kv_indices + if target_uses_mla: kv_last_page_lens = var["kv_last_page_lens"].gpu[:bs] attn_metadata.kv_last_page_lens = kv_last_page_lens - else: - # MHA needs block_tables and context_lens - attn_metadata.block_tables = var["block_tables"].gpu[:bs] - attn_metadata.context_lens = var["context_lens"].gpu[:bs] + # block_tables, context_lens, and sparse_kv_indptr are + # needed by both MHA and MLA+sparse attention + attn_metadata.block_tables = var["block_tables"].gpu[:bs] + attn_metadata.context_lens = var["context_lens"].gpu[:bs] + if "sparse_kv_indptr" in var: + attn_metadata.sparse_kv_indptr = var[ + "sparse_kv_indptr" + ].gpu[: bs + 1] cu_seqlens_q[: bs + 1] = self.arrange_bs[: bs + 1] - if use_mla: + if target_uses_mla and has_flat_kv: # MLA: block_size=1, kv_indptr tracks tokens kv_indptr[1 : bs + 1] -= torch.cumsum( num_reject_tokens, dim=0 @@ -222,9 +438,9 @@ def propose( # update metadata attn_metadata.max_seqlen_k += 1 - if not use_mla: - # MHA: update context_lens for this draft step - attn_metadata.context_lens[:bs] += 1 + # Update context_lens for each draft step (needed by both + # MHA attention and MLA+sparse indexer) + attn_metadata.context_lens[:bs] += 1 workinfos = self.runner.attn_metadata_builder.prepare_mtp_decode( bs, ( @@ -238,10 +454,25 @@ def propose( ) for k, v in workinfos.items(): attn_metadata.__dict__[k] = v - slot_mapping[:] = kv_indices[kv_indptr[1 : bs + 1] - 1] + if has_flat_kv: + # MLA/MHA path: slot derived from flat kv_indices. + # V4 doesn't expose flat kv_indices and its kernels + # don't read attn_metadata.slot_mapping (state-ring + + # swa_write_indices instead), so the update is skipped. + slot_mapping[:] = kv_indices[kv_indptr[1 : bs + 1] - 1] + input_ids = new_draft_ids positions += 1 - hidden_states = sample_hidden_states + if ret_hidden_prenorm is not None: + hidden_states = ( + torch.index_select( + ret_hidden_prenorm, 0, last_token_indices + ) + if i == 0 + else ret_hidden_prenorm + ) + else: + hidden_states = sample_hidden_states # self.runner.debug(f"final {draft_token_ids=}") # [batch_size, mtp_k] diff --git a/atom/utils/decorators.py b/atom/utils/decorators.py index eb16cd2463..e262030321 100644 --- a/atom/utils/decorators.py +++ b/atom/utils/decorators.py @@ -1,7 +1,14 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from typing import Callable, Optional, TypeVar, Union +from typing import ( + Callable, + Optional, + ParamSpec, + TypeVar, + Union, + overload as _overload, +) import inspect import os import sys @@ -23,6 +30,8 @@ # from atom.utils import start_monitoring_torch_compile _T = TypeVar("_T", bound=type[nn.Module]) +_P = ParamSpec("_P") +_R = TypeVar("_R") context_manager = None torch_compile_start_time: float = 0.0 @@ -167,6 +176,18 @@ def wrapped(*args, **kwargs): return wrapped +# Overloads preserve the decorated callable's signature for pyright. +# Without them, `@mark_trace` instances were inferred as plain `Callable`, +# and class instances whose `__call__` was wrapped (DualRMSNorm, LinearBase, +# ...) were flagged as "not callable" at the call site. +@_overload +def mark_trace(func: Callable[_P, _R], /) -> Callable[_P, _R]: ... +@_overload +def mark_trace( + *, + torch_compile: bool = ..., + prefix: Optional[str] = ..., +) -> Callable[[Callable[_P, _R]], Callable[_P, _R]]: ... def mark_trace( func: Optional[Callable] = None, *, diff --git a/atom/utils/envs.py b/atom/utils/envs.py index 269fa0c20a..105f109749 100644 --- a/atom/utils/envs.py +++ b/atom/utils/envs.py @@ -33,11 +33,13 @@ "ATOM_USE_TRITON_MXFP4_BMM": lambda: ( os.getenv("ATOM_USE_TRITON_MXFP4_BMM", "0") == "1" ), + "ATOM_USE_TRITON_MLA": lambda: os.getenv("ATOM_USE_TRITON_MLA", "0") == "1", + "ATOM_USE_TRITON_MOE": lambda: os.getenv("ATOM_USE_TRITON_MOE", "0") == "1", # --- Kernel Fusion Toggles --- # QK-norm-rope-cache-quant fusion for Qwen3-MoE; disabled by default. # Enable for Qwen3-MoE to get better performance. "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION": lambda: ( - os.getenv("ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION", "0") == "1" + os.getenv("ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION", "1") == "1" ), "ATOM_ENABLE_DS_INPUT_RMSNORM_QUANT_FUSION": lambda: ( os.getenv("ATOM_ENABLE_DS_INPUT_RMSNORM_QUANT_FUSION", "1") == "1" @@ -48,6 +50,9 @@ "ATOM_ENABLE_DS_QKNORM_FUSION": lambda: ( os.getenv("ATOM_ENABLE_DS_QKNORM_FUSION", "1") == "1" ), + "ATOM_ENABLE_DS_INDEXER_QK_ROPE_CACHE_FUSION": lambda: ( + os.getenv("ATOM_ENABLE_DS_INDEXER_QK_ROPE_CACHE_FUSION", "1") == "1" + ), "ATOM_ENABLE_ALLREDUCE_RMSNORM_FUSION": lambda: ( os.getenv("ATOM_ENABLE_ALLREDUCE_RMSNORM_FUSION", "1") == "1" ), @@ -69,6 +74,15 @@ "ATOM_DISABLE_MMAP": lambda: ( os.getenv("ATOM_DISABLE_MMAP", "false").lower() == "true" ), + # Use a thread pool for weight loading instead of main-process sequential I/O. + # Set to 0 to disable if the thread pool causes hangs (e.g. on gfx1250). + "ATOM_LOADER_USE_THREADPOOL": lambda: ( + os.getenv("ATOM_LOADER_USE_THREADPOOL", "1") == "1" + ), + # --- Attention Backend --- + # Use unified_attention (flash-style) for MHA paged/prefill attention instead + # of pa_decode_gluon. Set to 1 to enable the unified_attention path. + "ATOM_USE_UNIFIED_ATTN": lambda: os.getenv("ATOM_USE_UNIFIED_ATTN", "0") == "1", # --- Plugin Mode --- "ATOM_DISABLE_VLLM_PLUGIN": lambda: ( os.getenv("ATOM_DISABLE_VLLM_PLUGIN", "0").lower() == "1" @@ -85,6 +99,9 @@ "ATOM_DUAL_STREAM_MOE_TOKEN_THRESHOLD": lambda: int( os.getenv("ATOM_DUAL_STREAM_MOE_TOKEN_THRESHOLD", "1024") ), + # Gate/Up interleave mode for MoE weight preshuffle and kernel gate_mode. + # "0" (default) = SEPARATED layout; "1" = INTERLEAVE layout. + "ATOM_MOE_GU_ITLV": lambda: os.getenv("ATOM_MOE_GU_ITLV", "0") == "1", # --- MTP (relaxed mtp for quantized mtp) --- "ATOM_ENABLE_RELAXED_MTP": lambda: ( os.getenv("ATOM_ENABLE_RELAXED_MTP", "0").lower() == "1" @@ -93,6 +110,11 @@ # Enable gradient tracking on model parameters. Default "0" (disabled) # is correct for inference; set to "1" only for training / fine-tuning. "ATOM_REQUIRES_GRAD": lambda: os.getenv("ATOM_REQUIRES_GRAD", "0") == "1", + # --- Bpreshuffle for weight --- + # Preshuffle weight. Default "1" (enabled) + "ATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE": lambda: ( + os.getenv("ATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE", "1") == "1" + ), # --- V4 Attention Backend Refactor (PR-A: kill .item(), unlock CUDAGraph) --- # `legacy` (default) keeps the per-seq Python dispatch loop with .item() # syncs in deepseek_v4.py. `new` routes through V4AttentionBackend with diff --git a/atom/utils/forward_context.py b/atom/utils/forward_context.py index 77b28fa490..cdac102e2d 100644 --- a/atom/utils/forward_context.py +++ b/atom/utils/forward_context.py @@ -336,6 +336,28 @@ class ForwardContext: ubatch_slices: Optional[list[Any]] = None + # Cached current_stream() captured at set_forward_context() time, so + # downstream code (V4 attention / MoE / metadata builder) doesn't have + # to query torch.cuda.current_stream() repeatedly during a forward — + # multiple call sites caching independent Stream handles was widening + # the hipStream handle pool and complicating reasoning about which + # logical stream each wait_stream() refers to. CG capture / TBO + # threads each call set_forward_context() inside their own stream + # context, so the cached value is correct for the captured graph or + # active thread. + main_stream: Optional[torch.cuda.Stream] = None + + # True only while the model forward runs inside a CUDAGraph capture + # block (model_runner.capture_model loop). Components that gate + # multi-stream side-launches (V4 main Compressor on alt_stream, + # indexer.compressor on compress_stream) check this flag: side-stream + # work is safe to emit inside a captured graph (graph records the + # fork-join edges and replay re-uses the same stream layout) but + # racy in eager mode where launches accumulate across layers and + # deadlock the hipStream queue. Replay does not re-execute Python + # forward, so it ignores the flag entirely. + in_hipgraph: bool = False + def __post_init__(self): if not hasattr(self, "no_compile_layers") or self.no_compile_layers is None: self.no_compile_layers = {} @@ -346,6 +368,10 @@ def __post_init__(self): _forward_context: Optional[ForwardContext] = ForwardContext() _forward_kv_cache_context: Optional[ForwardContext] = ForwardContext() +# Cached once at module import — CUDA availability does not change at +# runtime, so we don't pay torch.cuda.is_available() per set_forward_context(). +_CUDA_AVAILABLE: bool = torch.cuda.is_available() + # Thread-local storage for TBO dual-thread execution _forward_context_local = threading.local() @@ -373,6 +399,7 @@ def set_forward_context( num_tokens_across_dp: Optional[torch.Tensor] = None, spec_decode_metadata: Optional[SpecDecodeMetadata] = None, ubatch_slices: Optional[list[Any]] = None, + in_hipgraph: bool = False, ) -> None: global _forward_context dp_metadata: Optional[DPMetadata] = None @@ -392,6 +419,8 @@ def set_forward_context( dp_metadata=dp_metadata, spec_decode_metadata=spec_decode_metadata, ubatch_slices=ubatch_slices, + main_stream=(torch.cuda.current_stream() if _CUDA_AVAILABLE else None), + in_hipgraph=in_hipgraph, ) # _forward_context.attn_metadata = attn_metadata # _forward_context.no_compile_layers = atom_config.compilation_config.static_forward_context # _forward_context = ForwardContext(no_compile_layers=atom_config.compilation_config.static_forward_context, attn_metadata=attn_metadata) diff --git a/atom/utils/selector.py b/atom/utils/selector.py index e87b1f8194..bb9b1c7d65 100644 --- a/atom/utils/selector.py +++ b/atom/utils/selector.py @@ -7,6 +7,7 @@ from atom.model_ops.attentions.backends import AttentionBackend from atom.utils import resolve_obj_by_qualname from atom.plugin.prepare import is_sglang, is_vllm +from atom.utils import envs def get_attn_backend( @@ -51,13 +52,9 @@ def get_attn_backend_cls( if use_v4: return "atom.model_ops.attentions.deepseek_v4_attn.DeepseekV4Backend" if use_mla: - # if block_size == 1: - return "atom.model_ops.attentions.aiter_mla.AiterMLABackend" # noqa: E501 - # else: - # raise ValueError( - # f" The selected backend" - # f"does not support block size {block_size}." - # "(currently only supports block size 1)") + if envs.ATOM_USE_TRITON_MLA: + return "atom.model_ops.attentions.triton_mla.TritonMLABackend" + return "atom.model_ops.attentions.aiter_mla.AiterMLABackend" if use_gdn: if use_vllm: return "atom.plugin.vllm.attention_backend.gdn_attn.GDNAttentionBackend" @@ -66,4 +63,6 @@ def get_attn_backend_cls( "atom.plugin.sglang.attention_backend.attention_gdn.GDNAttentionBackend" ) return "atom.model_ops.attentions.gdn_attn.GDNAttentionBackend" + if envs.ATOM_USE_UNIFIED_ATTN: + return "atom.model_ops.attentions.triton_mha.TritonMHABackend" return "atom.model_ops.attentions.aiter_attention.AiterBackend" # noqa: E501 diff --git a/docker/Dockerfile b/docker/Dockerfile index 38f93a2152..964168bd6e 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -155,6 +155,14 @@ RUN echo "========== [OOT] Restore base image triton ==========" && \ rm -f /tmp/triton_pip_show.txt && \ "${VENV_PYTHON}" -c "import triton; print(f'triton.__version__ = {triton.__version__}')" +# Qwen3-next encounter accuracy issue with the triton from base image +# we install verified good version triton==3.6.0 here as a temporary fix +# this should be removed after the triton issue is resolved +RUN echo "========== [OOT] Reinstall known good triton ==========" && \ + pip uninstall -y triton amd-triton && \ + pip install triton==3.6.0 +RUN pip show triton || true + RUN echo "========== [vLLM-ATOM] Final transformers version ==========" && \ "${VENV_PYTHON}" -c "import transformers; print(f'transformers.__version__ = {transformers.__version__}')" && \ "${VENV_PYTHON}" -m pip show transformers || true @@ -390,6 +398,13 @@ RUN cd /app/aiter-test && pip install -e . --no-build-isolation && \ RUN pip install rocm-trace-lite && \ rtl --version || true +# Qwen3-next encounter accuracy issue with the triton from base image +# we install verified good version triton==3.6.0 here as a temporary fix +# this should be removed after the triton issue is resolved +RUN pip uninstall -y triton amd-triton && \ + pip install triton==3.6.0 +RUN pip show triton || true + # ATOM: lightweight install (no compilation needed) # CACHEBUST invalidates only this layer so parallel stages stay cached ARG CACHEBUST=1 diff --git a/recipes/GPT-OSS.md b/recipes/GPT-OSS.md index 29fae6e5c5..a65449dac2 100644 --- a/recipes/GPT-OSS.md +++ b/recipes/GPT-OSS.md @@ -20,7 +20,7 @@ GPT-OSS-120B fits on a single MI300X/MI355X GPU: python -m atom.entrypoints.openai_server \ --model openai/gpt-oss-120b \ --kv_cache_dtype fp8 \ - --gpu-memory-utilization 0.3 + --gpu-memory-utilization 0.5 ``` ### Multi-GPU with DP Attention + Expert Parallel @@ -32,12 +32,12 @@ python -m atom.entrypoints.openai_server \ --model openai/gpt-oss-120b \ --kv_cache_dtype fp8 -tp 2 \ --enable-dp-attention --enable-expert-parallel \ - --gpu-memory-utilization 0.3 + --gpu-memory-utilization 0.5 ``` Tips on server configuration: - Single GPU is sufficient for this model — no TP required. -- `--gpu-memory-utilization 0.3` is recommended to leave room for KV cache growth. +- `--gpu-memory-utilization 0.5` is recommended to leave room for KV cache growth. - DP attention + EP mode does not require `--trust-remote-code` since ATOM has built-in `GptOssForCausalLM`. - Sliding window attention is applied on even-indexed layers automatically. @@ -70,4 +70,4 @@ lm_eval \ --num_fewshot 3 ``` -CI accuracy threshold: `flexible-extract ≥ 0.38`. +CI accuracy threshold: `flexible-extract ≥ 0.88`. diff --git a/recipes/atom_sglang/DeepSeek-R1.md b/recipes/atom_sglang/DeepSeek-R1.md new file mode 100644 index 0000000000..917959a7c8 --- /dev/null +++ b/recipes/atom_sglang/DeepSeek-R1.md @@ -0,0 +1,146 @@ +# DeepSeek-R1 with ATOM SGLang Backend + +This recipe shows how to run `deepseek-ai/DeepSeek-R1-0528` or `amd/DeepSeek-R1-0528-MXFP4-MTP-MoEFP4` with the SGLang-ATOM backend. For background on the SGLang-ATOM integration, see [Introduce ATOM as external model package of SGLang](https://github.com/ROCm/ATOM/issues/359). + +## Step 1: Pull the SGLang-ATOM Docker + +```bash +docker pull rocm/atom-dev:sglang-latest +``` + +Launch a container from this image and run the remaining commands inside the container. + +## Step 2: Launch SGLang-ATOM Server + +The SGLang-ATOM backend keeps the standard SGLang CLI, server APIs, and general usage flow compatible with upstream SGLang. For general server options and API usage, users can refer to the [official SGLang documentation](https://docs.sglang.ai/). + +Before launching the server, export the same SGLang-ATOM settings used by the benchmark workflow: + +```bash +export AITER_QUICK_REDUCE_QUANTIZATION=INT4 +export SGLANG_AITER_FP8_PREFILL_ATTN=0 +export SGLANG_USE_AITER=1 +export ATOM_ENABLE_DS_QKNORM_QUANT_FUSION=1 +# Introduce ATOM as external model package of SGLang +export SGLANG_EXTERNAL_MODEL_PACKAGE=atom.plugin.sglang.models +``` + +### DeepSeek with FP8 (TP=8) + +Users can use this command to launch the FP8 server with the same settings as the SGLang benchmark workflow. + +```bash +python3 -m sglang.launch_server \ + --model-path deepseek-ai/DeepSeek-R1-0528 \ + --host localhost \ + --port 8000 \ + --trust-remote-code \ + --tensor-parallel-size 8 \ + --kv-cache-dtype fp8_e4m3 \ + --mem-fraction-static 0.8 \ + --page-size 1 \ + --disable-radix-cache +``` + +### DeepSeek with FP8 (TP=4) + +```bash +python3 -m sglang.launch_server \ + --model-path deepseek-ai/DeepSeek-R1-0528 \ + --host localhost \ + --port 8000 \ + --trust-remote-code \ + --tensor-parallel-size 4 \ + --kv-cache-dtype fp8_e4m3 \ + --mem-fraction-static 0.8 \ + --page-size 1 \ + --disable-radix-cache +``` + +### DeepSeek with MXFP4 (TP=8) + +AMD Instinct MI355X GPU supports MXFP4 computation instructions, and users can use the following command to launch the MXFP4 server on MI355X. For MXFP4 model weight, we suggest using the checkpoint quantized from AMD Quark. + +```bash +python3 -m sglang.launch_server \ + --model-path amd/DeepSeek-R1-0528-MXFP4-MTP-MoEFP4 \ + --host localhost \ + --port 8000 \ + --trust-remote-code \ + --tensor-parallel-size 8 \ + --kv-cache-dtype fp8_e4m3 \ + --mem-fraction-static 0.8 \ + --page-size 1 \ + --disable-radix-cache +``` + +### DeepSeek with MXFP4 (TP=4) + +```bash +python3 -m sglang.launch_server \ + --model-path amd/DeepSeek-R1-0528-MXFP4-MTP-MoEFP4 \ + --host localhost \ + --port 8000 \ + --trust-remote-code \ + --tensor-parallel-size 4 \ + --kv-cache-dtype fp8_e4m3 \ + --mem-fraction-static 0.8 \ + --page-size 1 \ + --disable-radix-cache +``` + +## Step 3: Performance Benchmark + +The SGLang benchmark workflow uses the `bench_serving` client for performance benchmarking. The following example matches the workflow command pattern for an MXFP4 TP4 case. + +```bash +git clone --depth 1 https://github.com/kimbochen/bench_serving.git /tmp/bench_serving + +ISL=8192 +OSL=1024 +CONC=64 +RANDOM_RANGE_RATIO=0.8 +RESULT_DIR=./benchmark-results +RESULT_FILENAME=deepseek-r1-fp4-tp4-${ISL}-${OSL}-${CONC}-${RANDOM_RANGE_RATIO}.json + +python3 /tmp/bench_serving/benchmark_serving.py \ + --model=amd/DeepSeek-R1-0528-MXFP4-MTP-MoEFP4 \ + --backend=sglang \ + --base-url=http://127.0.0.1:8000 \ + --dataset-name=random \ + --random-input-len="${ISL}" \ + --random-output-len="${OSL}" \ + --random-range-ratio "${RANDOM_RANGE_RATIO}" \ + --num-prompts="$(( CONC * 10 ))" \ + --max-concurrency="${CONC}" \ + --trust-remote-code \ + --num-warmups="$(( 2 * CONC ))" \ + --request-rate=inf \ + --ignore-eos \ + --save-result \ + --percentile-metrics="ttft,tpot,itl,e2el" \ + --result-dir="${RESULT_DIR}" \ + --result-filename="${RESULT_FILENAME}" +``` + +For FP8 or TP8 cases, keep the same benchmark command and replace `--model` with the checkpoint used in Step 2. + +### Optional: Enable Profiling +If you want to collect profiling trace, set the SGLang profiling environment variables before launching the server, and add `--profile` to the benchmark client command. + +```bash +export SGLANG_PROFILE_RECORD_SHAPES=1 +export SGLANG_PROFILE_WITH_STACK=1 +export SGLANG_TORCH_PROFILER_DIR=./profile_sglang/ +``` + +Then append `--profile` to the `benchmark_serving.py` command in Step 3. + +## Step 4: Accuracy Validation + +```bash +lm_eval --model local-completions \ + --model_args model=amd/DeepSeek-R1-0528-MXFP4-MTP-MoEFP4,base_url=http://localhost:8000/v1/completions,num_concurrent=65,max_retries=1,tokenized_requests=False,trust_remote_code=True \ + --tasks gsm8k \ + --num_fewshot 3 +``` diff --git a/recipes/atom_sglang/Qwen3_5.md b/recipes/atom_sglang/Qwen3_5.md index 6baf617816..f9c1d17c04 100644 --- a/recipes/atom_sglang/Qwen3_5.md +++ b/recipes/atom_sglang/Qwen3_5.md @@ -27,6 +27,8 @@ Users can use this command to launch the FP8 server with the same settings as th ```bash +export ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=0 + model_path=[Qwen/Qwen3.5-397B-A17B-FP8](https://huggingface.co/Qwen/Qwen3.5-397B-A17B-FP8) # [Qwen/Qwen3-Next-80B-A3B-Instruct-FP8](https://huggingface.co/Qwen/Qwen3-Next-80B-A3B-Instruct-FP8) tp=4 @@ -35,7 +37,6 @@ python3 -m sglang.launch_server \ --port 8000 \ --tensor-parallel-size ${tp} \ --mem-fraction-static 0.9 \ - --cuda-graph-max-bs 256 \ --reasoning-parser qwen3 \ --disable-radix-cache ``` @@ -86,10 +87,9 @@ Then append `--profile` to the `sglang.bench_serving` command in Step 3. ```bash lm_eval --model local-completions \ - --model_args model=${model_path},base_url=http://localhost:30000/v1/completions,num_concurrent=256,max_retries=2,tokenized_requests=False,trust_remote_code=True \ + --model_args model=${model_path},base_url=http://localhost:30000/v1/completions,num_concurrent=65,max_retries=1,tokenized_requests=False,trust_remote_code=True \ --tasks gsm8k \ - --batch_size auto \ - --num_fewshot 5 \ + --num_fewshot 3 \ --trust_remote_code ``` diff --git a/recipes/atom_vllm/GLM-5.md b/recipes/atom_vllm/GLM-5.md index 6ee61c6305..1a8720dd54 100644 --- a/recipes/atom_vllm/GLM-5.md +++ b/recipes/atom_vllm/GLM-5.md @@ -23,8 +23,7 @@ vllm serve zai-org/GLM-5-FP8 \ --gpu_memory_utilization 0.9 \ --async-scheduling \ --compilation-config '{"cudagraph_mode": "FULL_AND_PIECEWISE"}' \ - --no-enable-prefix-caching \ - --block-size 1 + --no-enable-prefix-caching ``` ## Step 3: Performance Benchmark diff --git a/recipes/atom_vllm/GPT-OSS.md b/recipes/atom_vllm/GPT-OSS.md index 61557c9a74..414b5abc39 100644 --- a/recipes/atom_vllm/GPT-OSS.md +++ b/recipes/atom_vllm/GPT-OSS.md @@ -15,13 +15,12 @@ The ATOM vLLM plugin backend keeps the standard vLLM CLI, server APIs, and gener GPT-OSS-120B is a single-GPU model, so `--tensor-parallel-size` defaults to 1 and can be omitted. ```bash -export ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1 - +ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1 \ vllm serve openai/gpt-oss-120b \ --host localhost \ --port 8000 \ --kv-cache-dtype fp8 \ - --gpu_memory_utilization 0.9 \ + --gpu_memory_utilization 0.5 \ --async-scheduling \ --compilation-config '{"cudagraph_mode": "FULL_AND_PIECEWISE"}' \ --no-enable-prefix-caching diff --git a/recipes/atom_vllm/Kimi-K2-Thinking.md b/recipes/atom_vllm/Kimi-K2-Thinking.md index 949955b30b..5118129c2c 100644 --- a/recipes/atom_vllm/Kimi-K2-Thinking.md +++ b/recipes/atom_vllm/Kimi-K2-Thinking.md @@ -14,12 +14,13 @@ docker pull rocm/atom-dev:vllm-latest ## Step 2: Launch vLLM Server The ATOM vLLM plugin backend keeps the standard vLLM CLI, server APIs, and general usage flow compatible with upstream vLLM. For general server options and API usage, refer to the [official vLLM documentation](https://docs.vllm.ai/en/latest/). +We adopt [amd/Kimi-K2-Thinking-MXFP4-AttnFP8](https://huggingface.co/amd/Kimi-K2-Thinking-MXFP4-AttnFP8) for better performance by leveraging FP8 weights for attention layers. ```bash # use quick allreduce to reduce TTFT export AITER_QUICK_REDUCE_QUANTIZATION=INT4 -vllm serve amd/Kimi-K2-Thinking-MXFP4 \ +vllm serve amd/Kimi-K2-Thinking-MXFP4-AttnFP8 \ --host localhost \ --port 8000 \ --trust-remote-code \ @@ -28,6 +29,8 @@ vllm serve amd/Kimi-K2-Thinking-MXFP4 \ --gpu_memory_utilization 0.9 \ --async-scheduling \ --compilation-config '{"cudagraph_mode": "FULL_AND_PIECEWISE"}' \ + --max-num-batched-tokens 16384 \ + --max-model-len 16384 \ --no-enable-prefix-caching ``` @@ -37,7 +40,7 @@ Users can use the default vllm bench command for performance benchmarking. vllm bench serve \ --host localhost \ --port 8000 \ - --model amd/Kimi-K2-Thinking-MXFP4 \ + --model amd/Kimi-K2-Thinking-MXFP4-AttnFP8 \ --dataset-name random \ --random-input-len 8000 \ --random-output-len 1000 \ @@ -52,7 +55,15 @@ vllm bench serve \ ```bash lm_eval --model local-completions \ - --model_args model=amd/Kimi-K2-Thinking-MXFP4,base_url=http://localhost:8000/v1/completions,num_concurrent=16,max_retries=3,tokenized_requests=False \ + --model_args model=amd/Kimi-K2-Thinking-MXFP4-AttnFP8,base_url=http://localhost:8000/v1/completions,num_concurrent=16,max_retries=3,tokenized_requests=False \ --tasks gsm8k \ --num_fewshot 3 ``` +The reference values of corresponding metrics: +```bash +local-completions ({'model': 'amd/Kimi-K2-Thinking-MXFP4-AttnFP8', 'base_url': 'http://localhost:8000/v1/completions', 'num_concurrent': 16, 'max_retries': 3, 'tokenized_requests': False}), gen_kwargs: ({}), limit: None, num_fewshot: 3, batch_size: 1 +|Tasks|Version| Filter |n-shot| Metric | |Value | |Stderr| +|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:| +|gsm8k| 3|flexible-extract| 3|exact_match|↑ |0.9340|± |0.0068| +| | |strict-match | 3|exact_match|↑ |0.9325|± |0.0069| +``` \ No newline at end of file diff --git a/recipes/atom_vllm/Kimi-K2.5.md b/recipes/atom_vllm/Kimi-K2.5.md index aeaaa6ad1d..f94a139d92 100644 --- a/recipes/atom_vllm/Kimi-K2.5.md +++ b/recipes/atom_vllm/Kimi-K2.5.md @@ -14,12 +14,13 @@ docker pull rocm/atom-dev:vllm-latest ## Step 2: Launch vLLM Server The ATOM vLLM plugin backend keeps the standard vLLM CLI, server APIs, and general usage flow compatible with upstream vLLM. For general server options and API usage, refer to the [official vLLM documentation](https://docs.vllm.ai/en/latest/). +We adopt [amd/Kimi-K2.5-MXFP4-AttnFP8](https://huggingface.co/amd/Kimi-K2.5-MXFP4-AttnFP8) for better performance by leveraging FP8 weights for attention layers. ```bash # use quick allreduce to reduce TTFT export AITER_QUICK_REDUCE_QUANTIZATION=INT4 -vllm serve amd/Kimi-K2.5-MXFP4 \ +vllm serve amd/Kimi-K2.5-MXFP4-AttnFP8 \ --host localhost \ --port 8000 \ --async-scheduling \ @@ -28,6 +29,8 @@ vllm serve amd/Kimi-K2.5-MXFP4 \ --gpu_memory_utilization 0.9 \ --compilation-config '{"cudagraph_mode": "FULL_AND_PIECEWISE"}' \ --kv-cache-dtype fp8 \ + --max-num-batched-tokens 16384 \ + --max-model-len 16384 \ --no-enable-prefix-caching ``` @@ -51,7 +54,7 @@ IMAGE_BASE64=$(base64 -w 0 ATOM/recipes/atom_vllm/dog.png) curl -X POST "http://localhost:8000/v1/chat/completions" \ -H "Content-Type: application/json" \ -d '{ - "model": "amd/Kimi-K2.5-MXFP4", + "model": "amd/Kimi-K2.5-MXFP4-AttnFP8", "messages": [ { "role": "user", @@ -82,7 +85,7 @@ The expected response: "id": "chatcmpl-941a3736cc5cce95", "object": "chat.completion", "created": 1774365813, - "model": "amd/Kimi-K2.5-MXFP4", + "model": "amd/Kimi-K2.5-MXFP4-AttnFP8", "choices": [ { "index": 0, @@ -122,7 +125,7 @@ Users can use the default vllm bench command for performance benchmarking. vllm bench serve \ --host localhost \ --port 8000 \ - --model amd/Kimi-K2.5-MXFP4 \ + --model amd/Kimi-K2.5-MXFP4-AttnFP8 \ --dataset-name random \ --random-input-len 8000 \ --random-output-len 1000 \ @@ -137,15 +140,15 @@ vllm bench serve \ The accuracy can be verified on gsm8k dataset with command: ```bash lm_eval --model local-completions \ - --model_args model=amd/Kimi-K2.5-MXFP4,base_url=http://localhost:8000/v1/completions,num_concurrent=64,max_retries=3,tokenized_requests=False \ + --model_args model=amd/Kimi-K2.5-MXFP4-AttnFP8,base_url=http://localhost:8000/v1/completions,num_concurrent=64,max_retries=3,tokenized_requests=False \ --tasks gsm8k \ --num_fewshot 3 ``` The reference values of corresponding metrics: ```bash -local-completions ({'model': '/workspace/shared/data/models/Kimi-K2.5-MXFP4', 'base_url': 'http://localhost:8000/v1/completions', 'num_concurrent': 64, 'max_retries': 3, 'tokenized_requests': False}), gen_kwargs: ({}), limit: None, num_fewshot: 3, batch_size: 1 +local-completions ({'model': 'amd/Kimi-K2.5-MXFP4-AttnFP8', 'base_url': 'http://localhost:8000/v1/completions', 'num_concurrent': 64, 'max_retries': 3, 'tokenized_requests': False}), gen_kwargs: ({}), limit: None, num_fewshot: 3, batch_size: 1 |Tasks|Version| Filter |n-shot| Metric | |Value | |Stderr| |-----|------:|----------------|-----:|-----------|---|-----:|---|-----:| -|gsm8k| 3|flexible-extract| 3|exact_match|↑ |0.9401|± |0.0065| -| | |strict-match | 3|exact_match|↑ |0.9386|± |0.0066| +|gsm8k| 3|flexible-extract| 3|exact_match|↑ |0.9325|± |0.0061| +| | |strict-match | 3|exact_match|↑ |0.9240|± |0.0062| ``` diff --git a/recipes/atom_vllm/Qwen3.5.md b/recipes/atom_vllm/Qwen3.5.md index 94a900e074..4e3b8c077f 100644 --- a/recipes/atom_vllm/Qwen3.5.md +++ b/recipes/atom_vllm/Qwen3.5.md @@ -18,6 +18,7 @@ The ATOM vLLM plugin backend keeps the standard vLLM CLI, server APIs, and gener export AITER_QUICK_REDUCE_QUANTIZATION=INT4 export ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1 export ATOM_USE_CUSTOM_ALL_GATHER=0 +export ATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE=0 vllm serve Qwen/Qwen3.5-35B-A3B-FP8 \ --host localhost \ @@ -37,6 +38,7 @@ vllm serve Qwen/Qwen3.5-35B-A3B-FP8 \ export AITER_QUICK_REDUCE_QUANTIZATION=INT4 export ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1 export ATOM_USE_CUSTOM_ALL_GATHER=0 +export ATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE=0 vllm serve Qwen/Qwen3.5-397B-A17B-FP8 \ --host localhost \ @@ -56,6 +58,7 @@ vllm serve Qwen/Qwen3.5-397B-A17B-FP8 \ export AITER_QUICK_REDUCE_QUANTIZATION=INT4 export ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1 export ATOM_USE_CUSTOM_ALL_GATHER=0 +export ATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE=0 vllm serve amd/Qwen3.5-397B-A17B-MXFP4 \ --host localhost \ @@ -69,10 +72,10 @@ vllm serve amd/Qwen3.5-397B-A17B-MXFP4 \ --no-enable-prefix-caching ``` -**Important**: The following three environment variables are required for Qwen3.5: +**Important**: The following environment variables are required for Qwen3.5: -- `ATOM_DISABLE_VLLM_PLUGIN_ATTENTION=1`: Disables ATOM attention plugin to use vLLM's implementation for full attention layers (required because Qwen3.5 uses a hybrid architecture with both linear attention (GatedDeltaNet) and full attention layers) - `ATOM_USE_CUSTOM_ALL_GATHER=0`: Disables custom all-gather for compatibility with Qwen3.5 model architecture +- `ATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE=0`: Disables FP8 blockscale weight preshuffle - `AITER_QUICK_REDUCE_QUANTIZATION=INT4`: **Performance optimization** - enables INT4 quantization for quick reduce operations, which can significantly improve TTFT (Time To First Token) performance. **Note**: This optimization may introduce a risk of accuracy degradation. For accuracy-critical workloads, consider validating with your specific use case. ## Step 3: Performance Benchmark @@ -133,8 +136,8 @@ Reference result (TP=4): ## Key Environment Variables -- `ATOM_DISABLE_VLLM_PLUGIN_ATTENTION=1`: **Required** - disables ATOM attention plugin to use vLLM's implementation for full attention layers - `ATOM_USE_CUSTOM_ALL_GATHER=0`: **Required** - disables custom all-gather for compatibility with Qwen3.5 model architecture +- `ATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE=0`: **Required** - disables FP8 blockscale weight preshuffle - `AITER_QUICK_REDUCE_QUANTIZATION=INT4`: **Performance optimization** - enables INT4 quantization for quick reduce operations - **Benefit**: Significantly improves TTFT (Time To First Token) performance by reducing communication overhead during tensor parallelism all-reduce operations - **Risk**: May cause slight accuracy degradation due to lower quantization precision diff --git a/recipes/atom_vllm/Qwen3Next.md b/recipes/atom_vllm/Qwen3Next.md index e22f80d1c1..019e297f65 100644 --- a/recipes/atom_vllm/Qwen3Next.md +++ b/recipes/atom_vllm/Qwen3Next.md @@ -17,6 +17,8 @@ The ATOM vLLM plugin backend keeps the standard vLLM CLI, server APIs, and gener ```bash export ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1 export ATOM_USE_CUSTOM_ALL_GATHER=0 +export AITER_QUICK_REDUCE_QUANTIZATION=INT4 +export ATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE=0 vllm serve Qwen/Qwen3-Next-80B-A3B-Instruct-FP8 \ --host localhost \ @@ -31,8 +33,26 @@ vllm serve Qwen/Qwen3-Next-80B-A3B-Instruct-FP8 \ --no-enable-prefix-caching ``` -**Important**: `ATOM_DISABLE_VLLM_PLUGIN_ATTENTION=1` is required for Qwen3-Next because it uses a hybrid architecture with both linear attention (GatedDeltaNet) and full attention layers. This env var ensures full attention layers use vLLM's default implementation. +### Qwen3-Next-80B-A3B-Instruct-FP8 MTP (TP=1, MI355X) +```bash +export ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1 +export ATOM_USE_CUSTOM_ALL_GATHER=0 +export AITER_QUICK_REDUCE_QUANTIZATION=INT4 +export ATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE=0 +vllm serve Qwen/Qwen3-Next-80B-A3B-Instruct-FP8 \ + --host localhost \ + --port 8000 \ + --tensor-parallel-size 1 \ + --kv-cache-dtype fp8 \ + --gpu_memory_utilization 0.9 \ + --async-scheduling \ + --compilation-config '{"cudagraph_mode": "FULL_AND_PIECEWISE"}' \ + --max-model-len 16384 \ + --max-num-batched-tokens 32768 \ + --speculative-config '{"num_speculative_tokens":1, "method": "mtp"}' \ + --no-enable-prefix-caching +``` ## Step 3: Performance Benchmark Users can use the default vllm bench commands for performance benchmarking. @@ -70,9 +90,6 @@ lm_eval --model local-completions \ --num_fewshot 3 ``` -## Key Environment Variables - -- `ATOM_DISABLE_VLLM_PLUGIN_ATTENTION=1`: **Required** - disables ATOM attention plugin to use vLLM's implementation for full attention layers ## Architecture Notes diff --git a/recipes/pd_disaggregation_guide.md b/recipes/pd_disaggregation_guide.md new file mode 100644 index 0000000000..8a21c7c661 --- /dev/null +++ b/recipes/pd_disaggregation_guide.md @@ -0,0 +1,177 @@ +# PD Disaggregation with Mooncake (RDMA) + +Prefill-Decode disaggregation splits inference into two stages on separate nodes: +- **Producer** (prefill): runs prompt prefill, pushes KV cache via RDMA +- **Consumer** (decode): receives KV cache, runs autoregressive decode + +## Prerequisites + +- Two nodes with AMD MI300X GPUs (8 GPUs each for TP=8) +- RDMA network connectivity between nodes (RoCE or InfiniBand) +- Mooncake package installed (`pip install mooncake`) +- Producer and consumer should be in the **same network partition** for best accuracy + +## Building Mooncake for ROCm + +If Mooncake is not pre-installed in your Docker image, build from source: + +### 1. System dependencies + +```bash +apt update && apt install -y \ + zip unzip wget gcc make libtool autoconf cmake \ + librdmacm-dev rdmacm-utils infiniband-diags ibverbs-utils perftest ethtool \ + libibverbs-dev rdma-core \ + openssh-server openmpi-bin openmpi-common libopenmpi-dev +``` + +### 2. Install Go 1.22.2 + +Mooncake's etcd wrapper requires Go. **Must use Go 1.22.2** — Go 1.23+ causes `mallocHeaderSize redeclared` errors. + +```bash +apt remove -y golang golang-go 2>/dev/null || true +rm -rf /usr/local/go +wget https://go.dev/dl/go1.22.2.linux-amd64.tar.gz +tar -C /usr/local -xzf go1.22.2.linux-amd64.tar.gz +rm go1.22.2.linux-amd64.tar.gz +export PATH=$PATH:/usr/local/go/bin +go version # expect: go1.22.2 +``` + +### 3. Clone and build + +Use the SGLang-validated commit: + +```bash +git clone https://github.com/kvcache-ai/Mooncake.git +cd Mooncake +git checkout b6a841dc78c707ec655a563453277d969fb8f38d +git submodule update --init --recursive + +# Install C++ dependencies (etcd, protobuf, gRPC, abseil) +bash dependencies.sh -y + +# CMake build +mkdir -p build && cd build +cmake .. -DUSE_HIP=ON -DUSE_ETCD=ON +make -j$(nproc) +make install +``` + +Key flags: `-DUSE_HIP=ON` for ROCm (default is CUDA), `-DUSE_ETCD=ON` for metadata store. + +### 4. Install Python package + +```bash +cd ../mooncake-wheel +pip install . + +# Copy compiled .so to Python package +MOONCAKE_DIR=$(python -c "import mooncake; print(mooncake.__path__[0])") +cp ../build/mooncake-integration/engine.cpython-*-linux-gnu.so "$MOONCAKE_DIR/" + +ldconfig +``` + +### 5. Verify + +```bash +python -c "from mooncake.engine import TransferEngine; print('Mooncake ROCm OK')" +``` + +If you see `libglog.so` or other missing library errors, install them (`apt install -y libgoogle-glog-dev`) and re-run `ldconfig`. + +> **Important**: Producer and consumer nodes must use the **same Mooncake build version**. Mismatched versions cause `Corrupted segment descriptor` errors due to incompatible metadata serialization formats. + +## Quick Start + +### Step 0: Find local IP + +On each node, find the network IP (not loopback): + +```bash +export LOCAL_IP=$(ip addr show | grep "inet " | grep -v 127.0.0.1 | awk '{print $2}' | cut -d/ -f1 | head -1) +echo "Local IP: ${LOCAL_IP}" +``` + +### Step 1: Start Proxy (on producer node) + +The proxy handles routing between producer and consumer: + +```bash +python -m atom.kv_transfer.disaggregation.proxy --port 10001 +``` + +### Step 2: Start Producer (prefill node) + +```bash +ATOM_DISABLE_MMAP=true \ +NCCL_SOCKET_IFNAME=lo \ +AITER_LOG_LEVEL=WARNING \ +python -m atom.entrypoints.openai_server \ + --model /data/models/DeepSeek-R1/ \ + --kv_cache_dtype fp8 \ + -tp 8 \ + --server-port 8003 \ + --kv-transfer-config '{ + "kv_role": "kv_producer", + "kv_connector": "mooncake", + "proxy_ip": "'"${LOCAL_IP}"'", + "proxy_ping_port": 36367, + "http_port": 8003 + }' \ + 2>&1 | tee producer.log +``` + +### Step 3: Start Consumer (decode node) + +Replace `PRODUCER_IP` with the producer node's IP: + +```bash +export PRODUCER_IP= + +ATOM_DISABLE_MMAP=true \ +NCCL_SOCKET_IFNAME=lo \ +AITER_LOG_LEVEL=WARNING \ +python -m atom.entrypoints.openai_server \ + --model /data/models/DeepSeek-R1/ \ + --kv_cache_dtype fp8 \ + -tp 8 \ + --server-port 8004 \ + --kv-transfer-config '{ + "kv_role": "kv_consumer", + "kv_connector": "mooncake", + "proxy_ip": "'"${PRODUCER_IP}"'", + "proxy_ping_port": 36367, + "http_port": 8004 + }' \ + 2>&1 | tee consumer.log +``` + +### Step 4: Validate Accuracy + +Run GSM8K evaluation against the consumer endpoint: + +```bash +lm_eval --model local-chat-completions \ + --model_args "model=DeepSeek-R1,base_url=http://${CONSUMER_IP}:8004/v1,tokenizer_backend=huggingface,pretrained=/data/models/DeepSeek-R1/" \ + --tasks gsm8k_cot \ + --batch_size 1 \ + --limit 100 \ + --apply_chat_template \ + --fewshot_as_multiturn \ + --predict_only \ + --log_samples \ + --output_path results/ \ + --gen_kwargs "max_tokens=8192,temperature=0.6" +``` + +Expected accuracy: ~0.95-0.96 (matching non-PD baseline). +``` +ewshot: None, batch_size: 1 +|Tasks|Version| Filter |n-shot| Metric | |Value| |Stderr| +|-----|------:|----------------|-----:|-----------|---|----:|---|-----:| +|gsm8k| 3|flexible-extract| 5|exact_match|↑ | 0.96|± | 0.028| +| | |strict-match | 5|exact_match|↑ | 0.96|± | 0.028| +``` \ No newline at end of file diff --git a/scripts/run_benchmark.sh b/scripts/run_benchmark.sh new file mode 100755 index 0000000000..8d9174669e --- /dev/null +++ b/scripts/run_benchmark.sh @@ -0,0 +1,92 @@ +#!/bin/bash +# Performance benchmark for ATOM serving +# Usage: bash run_benchmark.sh [MODEL] [PORT] [ISL] [OSL] [CONCURRENCY] [PROMPT_MULTIPLIER] [PROFILE] +# +# Examples: +# bash run_benchmark.sh /data/DeepSeek-R1-0528 8000 1024 1024 128 # normal +# bash run_benchmark.sh /data/DeepSeek-R1-0528 8000 1024 1024 64 2 1 # profile trace (conc*1 requests) +# bash run_benchmark.sh /data/DeepSeek-R1-0528 8000 1024 8192 32 + +set -euo pipefail + +MODEL="${1:-/data/DeepSeek-R1-0528}" +PORT="${2:-8000}" +ISL="${3:-1024}" +OSL="${4:-1024}" +CONCURRENCY="${5:-128}" +PROMPT_MULTIPLIER="${6:-10}" +PROFILE="${7:-0}" +_n=$(( $# < 7 ? $# : 7 )) +shift "$_n" 2>/dev/null || true +EXTRA_ARGS="$*" +NUM_PROMPTS=$(( CONCURRENCY * PROMPT_MULTIPLIER )) +RANDOM_RANGE_RATIO="${RANDOM_RANGE_RATIO:-0.8}" +OUTPUT_DIR="/app/logs_claude" + +echo "========================================" +echo " ATOM Performance Benchmark" +echo "========================================" +echo " Model: $MODEL" +echo " Base URL: http://localhost:${PORT}" +echo " ISL/OSL: ${ISL}/${OSL}" +echo " Concurrency: $CONCURRENCY" +echo " Num Prompts: $NUM_PROMPTS (x${PROMPT_MULTIPLIER})" +echo " Range Ratio: $RANDOM_RANGE_RATIO" +echo " Profile: $PROFILE" +echo " Extra args: ${EXTRA_ARGS:-none}" +echo "========================================" + +# Wait for server ready (dual check: HTTP + GPU VRAM) +echo "Waiting for server to be ready..." +for i in $(seq 1 120); do + if curl -sf "http://localhost:${PORT}/v1/models" > /dev/null 2>&1; then + VRAM_COUNT=$(rocm-smi --showmemuse 2>/dev/null | grep "VRAM%" | awk '{print $NF}' | awk '$1 > 0' | wc -l) + if [ "$VRAM_COUNT" -gt 0 ]; then + echo "Server is ready! (GPU VRAM loaded on $VRAM_COUNT GPUs)" + break + fi + fi + if [ "$i" -eq 120 ]; then + echo "ERROR: Server not ready after 120s" + exit 1 + fi + [ $((i % 10)) -eq 0 ] && echo " Waiting... (${i}s)" + sleep 1 +done + +LOG_FILE="${OUTPUT_DIR}/benchmark.log" +RESULT_FILE="${OUTPUT_DIR}/benchmark.json" + +# Write config header to log +{ +echo "========================================" +echo " ATOM Performance Benchmark" +echo "========================================" +echo " Model: $MODEL" +echo " Base URL: http://localhost:${PORT}" +echo " ISL/OSL: ${ISL}/${OSL}" +echo " Concurrency: $CONCURRENCY" +echo " Num Prompts: $NUM_PROMPTS (x${PROMPT_MULTIPLIER})" +echo " Range Ratio: $RANDOM_RANGE_RATIO" +echo " Profile: $PROFILE" +echo " Extra args: ${EXTRA_ARGS:-none}" +echo " Date: $(date)" +echo "========================================" +} > "$LOG_FILE" + +python -m atom.benchmarks.benchmark_serving \ + --model="$MODEL" --backend=vllm --base-url="http://localhost:${PORT}" \ + --dataset-name=random \ + --random-input-len="$ISL" --random-output-len="$OSL" \ + --random-range-ratio="$RANDOM_RANGE_RATIO" \ + --max-concurrency="$CONCURRENCY" \ + --num-prompts="$NUM_PROMPTS" \ + --trust-remote-code \ + --num-warmups=$((CONCURRENCY * 2)) \ + --request-rate=inf --ignore-eos \ + --save-result \ + --result-filename="$RESULT_FILE" \ + --percentile-metrics="ttft,tpot,itl,e2el" \ + $([ "$PROFILE" = "1" ] && echo "--profile") \ + ${EXTRA_ARGS:-} \ + 2>&1 | tee -a "$LOG_FILE" diff --git a/scripts/run_benchmark_sweep.sh b/scripts/run_benchmark_sweep.sh new file mode 100755 index 0000000000..5f54809a45 --- /dev/null +++ b/scripts/run_benchmark_sweep.sh @@ -0,0 +1,35 @@ +#!/bin/bash +# Run benchmark across multiple concurrency levels +# Usage: bash run_benchmark_sweep.sh MODEL PORT ISL OSL [CONC_LIST] +# +# Examples: +# bash run_benchmark_sweep.sh /data/Kimi-K2.5-MXFP4 8000 1024 1024 +# bash run_benchmark_sweep.sh /data/Kimi-K2.5-MXFP4 8000 1024 1024 "4 8 16 32 64 128" +# bash run_benchmark_sweep.sh /data/DeepSeek-R1-0528 8000 1024 1024 "1 2 4 8 16 32 64 128 256" + +set -uo pipefail + +MODEL="${1:?Usage: $0 MODEL PORT ISL OSL [CONC_LIST]}" +PORT="${2:-8000}" +ISL="${3:-1024}" +OSL="${4:-1024}" +CONC_LIST="${5:-4 8 16 32 64 128}" + +echo "========================================" +echo " Benchmark Sweep" +echo "========================================" +echo " Model: $MODEL" +echo " ISL/OSL: ${ISL}/${OSL}" +echo " Conc: $CONC_LIST" +echo "========================================" + +for CONC in $CONC_LIST; do + echo "" + echo "================================================================" + echo "=== CONC=$CONC ===" + echo "================================================================" + bash /app/ATOM/scripts/run_benchmark.sh "$MODEL" "$PORT" "$ISL" "$OSL" "$CONC" +done + +echo "" +echo "=== Sweep complete ===" diff --git a/scripts/run_debug_agent.sh b/scripts/run_debug_agent.sh new file mode 100755 index 0000000000..67bd81e628 --- /dev/null +++ b/scripts/run_debug_agent.sh @@ -0,0 +1,82 @@ +#!/bin/bash +# Launch ATOM under rocm-debug-agent so a GPU memory fault dumps wave state + +# code-object disassembly. Intended for debugging async kernel races (Memory +# access fault, MEMORY_VIOLATION, ASSERT_TRAP, silent GPU livelocks). +# +# Two modes: +# +# Server mode (default): +# bash scripts/run_debug_agent.sh [MODEL] [TP] [PORT] [EXTRA_ARGS...] +# MODEL default /data/DeepSeek-V4-Pro +# TP default 8 +# PORT default 8000 +# EXTRA_ARGS forwarded to start_atom_server.sh +# +# Offline simple_inference mode: +# bash scripts/run_debug_agent.sh --simple [MODEL] [TP] [EXTRA_ARGS...] +# MODEL default /data/DeepSeek-V4-Pro +# TP default 8 +# (no PORT — simple_inference is offline) +# EXTRA_ARGS forwarded to start_simple_inference.sh +# Default log path: /app/logs_claude/simple_inference_debug_agent.log +# (override with LOG_FILE=...). +# +# REQUIRED in EXTRA_ARGS for either mode: --enforce-eager --level 0 +# (graph mode + Inductor are incompatible with the agent's no-caching +# allocator). The script does NOT inject these — pass them via EXTRA_ARGS so +# the model-specific launch matches your repro. +# +# Output: the agent prints wave dumps to stderr; both launchers redirect into +# /app/logs_claude/. Code objects (one per faulting kernel, +# ~4 MB each) land in /app/logs_claude/debug_run/. +# +# After fault triggers in server mode, the agent keeps running. Stop with: +# bash scripts/stop_atom_server.sh +# In simple mode the launcher exits when generation finishes (or faults). + +set -uo pipefail + +MODE="server" +if [ "${1:-}" = "--simple" ]; then + MODE="simple" + shift +fi + +MODEL="${1:-/data/DeepSeek-V4-Pro}" +TP="${2:-8}" + +if [ "$MODE" = "server" ]; then + PORT="${3:-8000}" + shift 3 2>/dev/null || true +else + shift 2 2>/dev/null || true +fi + +# 1) no gpucore files (a single ROCm fault dumps 30-50 GB per rank) +ulimit -c 0 + +# 2) clean cwd for code-object dumps; --save-code-objects writes +# `memory____offset__size_` files into the process cwd. +DEBUG_DIR="/app/logs_claude/debug_run" +mkdir -p "$DEBUG_DIR" +cd "$DEBUG_DIR" +rm -f memory_* *.s + +# 3) debug agent env (must be exported so spawn workers inherit) +export HSA_TOOLS_LIB=/opt/rocm/lib/librocm-debug-agent.so.2 +export HSA_ENABLE_DEBUG=1 +export ROCM_DEBUG_AGENT_OPTIONS="--save-code-objects" + +# 4) ATOM-side default; caller may override via env before invocation +export AITER_LOG_LEVEL="${AITER_LOG_LEVEL:-WARNING}" + +# 5) launch (model-specific env like ATOM_USE_TRITON_MOE must come from caller) +SCRIPT_DIR="$(dirname "$0")" +if [ "$MODE" = "server" ]; then + exec bash "$SCRIPT_DIR/start_atom_server.sh" \ + "$MODEL" "$TP" "$PORT" "$@" +else + export LOG_FILE="${LOG_FILE:-/app/logs_claude/simple_inference_debug_agent.log}" + exec bash "$SCRIPT_DIR/start_simple_inference.sh" \ + "$MODEL" "$TP" "$@" +fi diff --git a/scripts/run_gsm8k_eval.sh b/scripts/run_gsm8k_eval.sh new file mode 100755 index 0000000000..79a06609ad --- /dev/null +++ b/scripts/run_gsm8k_eval.sh @@ -0,0 +1,56 @@ +#!/bin/bash +# GSM8K 精度测试 (使用 lm_eval) +# 用法: bash run_gsm8k_eval.sh [MODEL_NAME] [PORT] [NUM_FEWSHOT] +# +# 示例: +# bash run_gsm8k_eval.sh /data/DeepSeek-R1-0528 8000 5 +# bash run_gsm8k_eval.sh meta-llama/Meta-Llama-3-8B 8000 5 + +set -euo pipefail + +MODEL="${1:-/data/DeepSeek-R1-0528}" +PORT="${2:-8000}" +NUM_FEWSHOT="${3:-5}" +NUM_CONCURRENT="${NUM_CONCURRENT:-65}" +LIMIT="${LIMIT:-}" # set LIMIT=50 to run only first 50 samples +BASE_URL="http://localhost:${PORT}/v1/completions" +OUTPUT_DIR="/app/logs_claude" +LOG_FILE="${OUTPUT_DIR}/gsm8k_eval.log" + +{ +echo "========================================" +echo " GSM8K Accuracy Evaluation (lm_eval)" +echo "========================================" +echo " Model: $MODEL" +echo " Base URL: $BASE_URL" +echo " Few-shot: $NUM_FEWSHOT" +echo " Concurrency: $NUM_CONCURRENT" +echo "========================================" +} > "$LOG_FILE" + +# 等待服务就绪 +echo "Waiting for server to be ready..." +for i in $(seq 1 120); do + if curl -sf "http://localhost:${PORT}/v1/models" > /dev/null 2>&1; then + echo "Server is ready!" + break + fi + if [ "$i" -eq 120 ]; then + echo "ERROR: Server not ready after 120 retries (20 min)" + exit 1 + fi + echo " Waiting... ($i/120)" + sleep 10 +done + +LIMIT_ARG=() +if [ -n "$LIMIT" ]; then + LIMIT_ARG=(--limit "$LIMIT") +fi + +lm_eval --model local-completions \ + --model_args "model=${MODEL},base_url=${BASE_URL},num_concurrent=${NUM_CONCURRENT},max_retries=3,tokenized_requests=False,trust_remote_code=True" \ + --tasks gsm8k \ + --num_fewshot "$NUM_FEWSHOT" \ + "${LIMIT_ARG[@]}" \ + 2>&1 | tee -a "${LOG_FILE}" diff --git a/scripts/start_atom_server.sh b/scripts/start_atom_server.sh new file mode 100755 index 0000000000..197c040dd1 --- /dev/null +++ b/scripts/start_atom_server.sh @@ -0,0 +1,115 @@ +#!/bin/bash +# Start ATOM OpenAI-compatible server +# Usage: bash start_atom_server.sh [MODEL_PATH] [TP_SIZE] [PORT] [EXTRA_ARGS...] +# +# Examples: +# bash start_atom_server.sh # DeepSeek-R1-0528, tp=8, port=8000 +# bash start_atom_server.sh /data/Llama-3.1-8B-Instruct-FP8-KV 1 8000 +# bash start_atom_server.sh /data/DeepSeek-R1-0528 8 8000 --method mtp --num-speculative-tokens 3 + +set -euo pipefail + +MODEL_PATH="${1:-/data/DeepSeek-R1-0528}" +TP_SIZE="${2:-8}" +PORT="${3:-8000}" +shift 3 2>/dev/null || true +EXTRA_ARGS="$*" +KV_CACHE_DTYPE="${KV_CACHE_DTYPE:-fp8}" +MAX_NUM_SEQS="${MAX_NUM_SEQS:-256}" +GPU_MEM_UTIL="${GPU_MEM_UTIL:-0.9}" +LOG_FILE="/app/logs_claude/atom_server.log" + +export AITER_LOG_LEVEL="${AITER_LOG_LEVEL:-INFO}" +export KINETO_CONFIG="/home/ljin1/dk/libkineto.conf" + +# === Pre-flight: ensure GPU is clean === +echo "Pre-flight: cleaning up processes and GPU memory..." + +# 1. Kill atom server processes +pkill -f 'atom.entrypoints' 2>/dev/null || true +sleep 2 + +# 2. Kill orphaned multiprocessing spawn/tracker (these hold GPU memory after server dies) +pkill -9 -f 'multiprocessing.spawn' 2>/dev/null || true +pkill -9 -f 'multiprocessing.resource_tracker' 2>/dev/null || true +sleep 3 + +# 3. Verify GPU memory is actually free +MAX_WAIT=30 +for i in $(seq 1 $MAX_WAIT); do + USED_GPUS=$(rocm-smi --showmemuse 2>/dev/null | grep "VRAM%" | awk '{print $NF}' | awk '$1 > 0' | wc -l) + if [ "$USED_GPUS" -eq 0 ]; then + echo "GPU memory clear after ${i}s" + break + fi + if [ "$i" -eq "$MAX_WAIT" ]; then + echo "WARNING: GPU memory still in use after ${MAX_WAIT}s. Dumping GPU process info:" + rocm-smi --showpidgpus 2>&1 | grep "PID.*is using" | grep -v "0 DRM" || true + echo "Attempting force kill of GPU-holding processes..." + rocm-smi --showpidgpus 2>&1 | grep -oP 'PID \K\d+' | while read pid; do + kill -9 "$pid" 2>/dev/null || true + done + sleep 5 + fi + sleep 1 +done + +# 4. Clear stale compile cache +rm -rf ~/.cache/atom/* +rm -rf ./gpucore.* + +# Write config header to log (truncates old content). +# Inherited env vars are dumped explicitly so you never have to wonder +# whether ATOM_USE_TRITON_MOE / V4_USE_REF_QUANT / etc. were set. +{ +echo "========================================" +echo " ATOM Server Launcher" +echo "========================================" +echo " Model: $MODEL_PATH" +echo " TP Size: $TP_SIZE" +echo " Port: $PORT" +echo " KV Cache dtype: $KV_CACHE_DTYPE" +echo " Max num seqs: $MAX_NUM_SEQS" +echo " GPU mem util: $GPU_MEM_UTIL" +echo " Extra args: ${EXTRA_ARGS:-none}" +echo " Date: $(date)" +echo "----------------------------------------" +echo " Inherited env vars (ATOM_*, V4_*, AITER_*, HSA_*, AMD_*, HIP_*):" +env | grep -E '^(ATOM_|V4_|AITER_|HSA_|AMD_|HIP_|KV_CACHE|MAX_NUM_SEQS|MAX_MODEL_LEN|MAX_BATCHED_TOKENS|GPU_MEM_UTIL)' \ + | sort | sed 's/^/ /' || echo " (none set)" +echo "========================================" +} | tee "$LOG_FILE" + +python -m atom.entrypoints.openai_server \ + --model "$MODEL_PATH" \ + --kv_cache_dtype "$KV_CACHE_DTYPE" \ + -tp "$TP_SIZE" \ + --max-num-seqs "$MAX_NUM_SEQS" \ + --gpu-memory-utilization "$GPU_MEM_UTIL" \ + --server-port "$PORT" \ + $EXTRA_ARGS \ + >> "$LOG_FILE" 2>&1 & + +SERVER_PID=$! +echo "Server started in background (PID: $SERVER_PID)" + +# Wait for server ready with GPU verification +echo "Waiting for server to be ready..." +for i in $(seq 1 120); do + if curl -sf "http://localhost:${PORT}/v1/models" > /dev/null 2>&1; then + VRAM_COUNT=$(rocm-smi --showmemuse 2>/dev/null | grep "VRAM%" | awk '{print $NF}' | awk '$1 > 0' | wc -l) + if [ "$VRAM_COUNT" -gt 0 ]; then + echo "Server is ready! (PID: $SERVER_PID, GPU VRAM loaded)" + exit 0 + fi + fi + if ! kill -0 "$SERVER_PID" 2>/dev/null; then + echo "ERROR: Server process died. Check $LOG_FILE" + exit 1 + fi + if [ "$i" -eq 600 ]; then + echo "ERROR: Server not ready after 600s (10 min)" + exit 1 + fi + sleep 1 +done diff --git a/scripts/start_simple_inference.sh b/scripts/start_simple_inference.sh new file mode 100755 index 0000000000..1627b34311 --- /dev/null +++ b/scripts/start_simple_inference.sh @@ -0,0 +1,132 @@ +#!/bin/bash +# Run ATOM offline simple_inference for any registered model. +# Usage: bash start_simple_inference.sh [MODEL_PATH] [TP_SIZE] [EXTRA_ARGS...] +# +# All engine flags (kv_cache_dtype, max_num_seqs, max_model_len, gpu mem util, +# spec method, …) go in EXTRA_ARGS as native CLI args — no env-var translation. +# Only LOG_FILE and true env vars (ATOM_*, AITER_*, HSA_*, HIP_*) remain +# overridable via env. +# +# Examples: +# bash start_simple_inference.sh # default DeepSeek-R1-0528 tp=8 +# bash start_simple_inference.sh /data/Llama-3.1-8B-Instruct-FP8-KV 1 --kv_cache_dtype fp8 +# bash start_simple_inference.sh /data/DeepSeek-V4-Pro 8 --kv_cache_dtype fp8 --enforce-eager --level 0 +# bash start_simple_inference.sh /data/DeepSeek-V4-Pro 8 \ +# --kv_cache_dtype fp8 --enforce-eager --level 0 \ +# --method mtp --num-speculative-tokens 3 +# MAX_MODEL_LEN=1024 ... # ← no longer supported; pass --max-model-len 1024 instead +# +# DeepSeek-V4 must also export ATOM_USE_TRITON_MOE=1 (real env var, not a CLI flag). + +set -euo pipefail + +MODEL_PATH="${1:-/data/DeepSeek-R1-0528}" +TP_SIZE="${2:-8}" +shift 2 2>/dev/null || true +EXTRA_ARGS=("$@") + +LOG_FILE="${LOG_FILE:-/app/logs_claude/simple_inference.log}" + +export AITER_LOG_LEVEL="${AITER_LOG_LEVEL:-INFO}" + +# === Pre-flight: ensure GPU is clean === +echo "Pre-flight: cleaning up processes and GPU memory..." + +pkill -f 'atom.entrypoints' 2>/dev/null || true +pkill -f 'atom.examples.simple_inference' 2>/dev/null || true +sleep 2 + +pkill -9 -f 'multiprocessing.spawn' 2>/dev/null || true +pkill -9 -f 'multiprocessing.resource_tracker' 2>/dev/null || true +sleep 3 + +# Verify GPU memory is actually free +MAX_WAIT=30 +for i in $(seq 1 $MAX_WAIT); do + USED_GPUS=$(rocm-smi --showmemuse 2>/dev/null | grep "VRAM%" | awk '{print $NF}' | awk '$1 > 0' | wc -l) + if [ "$USED_GPUS" -eq 0 ]; then + echo " All GPUs free." + break + fi + if [ "$i" -eq "$MAX_WAIT" ]; then + echo " WARN: $USED_GPUS GPU(s) still showing memory after $MAX_WAIT s — proceeding anyway" + break + fi + sleep 1 +done + +# Clear ATOM compile cache (stale cache after code change causes silent fails) +rm -rf /root/.cache/atom/* 2>/dev/null || true + +# Clear ROCm core dumps (each ~85GB; HSA exceptions can fill the disk fast). +# Find from CWD + a few common roots — gpucore.PID files dropped in process CWD. +GPUCORE_DELETED=0 +for d in . /app/ATOM /app /root /tmp /app/logs_claude; do + [ -d "$d" ] || continue + for f in "$d"/gpucore.* "$d"/core.[0-9]*; do + if [ -e "$f" ]; then + sz=$(du -BG --apparent-size "$f" 2>/dev/null | awk '{print $1}') + rm -f "$f" 2>/dev/null && GPUCORE_DELETED=$((GPUCORE_DELETED+1)) + echo " removed $f ($sz)" + fi + done +done +[ "$GPUCORE_DELETED" -gt 0 ] && echo " Cleared $GPUCORE_DELETED ROCm core dump(s)" + +# Disk-space guard — refuse to launch if root fs has < 50 GB free +# (each crashed worker can dump 85+GB; better to fail-fast than fill disk again). +AVAIL_GB=$(df -BG / 2>/dev/null | awk 'NR==2 {gsub("G",""); print $4}') +if [ -n "$AVAIL_GB" ] && [ "$AVAIL_GB" -lt 50 ]; then + echo "ERROR: only ${AVAIL_GB}GB free on / — refusing to launch (would fill disk on crash)." + echo " Run: rm -f /app/ATOM/gpucore.* /app/logs_claude/**/gpucore.*" + exit 2 +fi + +# Disable ROCm core dumps for this run (HSA crashes won't write 85GB files). +# Set HSA_ENABLE_COREDUMP_FILTER to 0x0 to skip all core regions. +export HSA_ENABLE_COREDUMP=0 +export AMD_LOG_LEVEL=0 # quiet HSA error spam + +# === Header to log === +# Inherited env vars are dumped explicitly so missing ATOM_USE_TRITON_MOE etc. +# is visible at a glance instead of producing silent accuracy regressions. +{ +echo "========================================" +echo " ATOM simple_inference Launcher" +echo "========================================" +echo " Model: $MODEL_PATH" +echo " TP Size: $TP_SIZE" +echo " Extra args: ${EXTRA_ARGS[*]:-none}" +echo " Date: $(date)" +echo "----------------------------------------" +echo " Inherited env vars (ATOM_*, V4_*, AITER_*, HSA_*, AMD_*, HIP_*):" +env | grep -E '^(ATOM_|V4_|AITER_|HSA_|AMD_|HIP_)' \ + | sort | sed 's/^/ /' || echo " (none set)" +echo "========================================" +} | tee "$LOG_FILE" + +set +e +python -m atom.examples.simple_inference \ + --model "$MODEL_PATH" \ + -tp "$TP_SIZE" \ + "${EXTRA_ARGS[@]}" \ + >> "$LOG_FILE" 2>&1 +EXIT_CODE=$? +set -e + +echo "" +echo "========================================" +echo " Exit code: $EXIT_CODE" +echo " Log: $LOG_FILE" +echo "========================================" + +# Show first error if any +if [ "$EXIT_CODE" -ne 0 ]; then + echo "" + echo "First error in log:" + grep -nE "RuntimeError|AttributeError|TypeError|ValueError|^\[rank0\]:" "$LOG_FILE" \ + | grep -v "AsyncIOProc\|ModelRunner.*has no attribute" \ + | head -5 +fi + +exit $EXIT_CODE diff --git a/scripts/stop_atom_server.sh b/scripts/stop_atom_server.sh new file mode 100755 index 0000000000..1a3aee5c18 --- /dev/null +++ b/scripts/stop_atom_server.sh @@ -0,0 +1,92 @@ +#!/bin/bash +# Stop ATOM server and release GPU memory. +# Usage: bash stop_atom_server.sh +# +# Escalates SIGTERM -> SIGKILL on the main entrypoint process AND its +# multiprocessing children, then waits up to 60s for VRAM to drop to 0. +# If VRAM is still held, dumps the GPU PID list and force-kills each one. + +set -uo pipefail + +echo "=== Stopping ATOM server ===" + +# Patterns matching every process the server can spawn AND clients hammering +# it (lm_eval) so the next launch starts from a fully clean state. +PATTERNS=( + 'atom.entrypoints' + 'multiprocessing.spawn' + 'multiprocessing.resource_tracker' + 'atom.model_engine' + 'lm_eval' +) + +count_alive() { + local total=0 + for pat in "${PATTERNS[@]}"; do + local n + n=$(pgrep -f "$pat" 2>/dev/null | wc -l) + total=$((total + n)) + done + echo "$total" +} + +# 1. Graceful SIGTERM round. +if [ "$(count_alive)" -gt 0 ]; then + echo "Sending SIGTERM to ATOM processes..." + for pat in "${PATTERNS[@]}"; do + pkill -TERM -f "$pat" 2>/dev/null || true + done + for i in $(seq 1 5); do + sleep 1 + [ "$(count_alive)" -eq 0 ] && break + done +fi + +# 2. SIGKILL anything still alive — escalate without polite waiting. +if [ "$(count_alive)" -gt 0 ]; then + echo "SIGTERM did not finish cleanup; escalating to SIGKILL..." + for pat in "${PATTERNS[@]}"; do + pkill -KILL -f "$pat" 2>/dev/null || true + done + sleep 2 +fi + +# 3. Final sweep: any leftover that still references the server by name. +if [ "$(count_alive)" -gt 0 ]; then + echo "Stubborn processes still alive — final SIGKILL pass:" + for pat in "${PATTERNS[@]}"; do + pgrep -af "$pat" 2>/dev/null || true + done + for pat in "${PATTERNS[@]}"; do + pkill -KILL -9 -f "$pat" 2>/dev/null || true + done + sleep 3 +fi + +# 4. Wait for GPU memory to release. +echo "Waiting for GPU memory to release..." +for i in $(seq 1 60); do + USED_GPUS=$(rocm-smi --showmemuse 2>/dev/null | grep "VRAM%" | awk '{print $NF}' | awk '$1 > 0' | wc -l) + if [ "$USED_GPUS" -eq 0 ]; then + echo "GPU memory released after ${i}s" + break + fi + if [ "$i" -eq 60 ]; then + echo "WARNING: GPU memory still in use after 60s, force-killing GPU PIDs..." + # Newer rocm-smi exposes per-GPU PIDs via --showpids + for pid in $(rocm-smi --showpids 2>/dev/null | awk '/^PID/{f=1; next} f && $1 ~ /^[0-9]+$/ {print $1}'); do + echo " kill -9 $pid" + kill -9 "$pid" 2>/dev/null || true + done + sleep 5 + USED_GPUS=$(rocm-smi --showmemuse 2>/dev/null | grep "VRAM%" | awk '{print $NF}' | awk '$1 > 0' | wc -l) + if [ "$USED_GPUS" -gt 0 ]; then + echo "ERROR: GPU memory still held after force-kill. Manual intervention required." + rocm-smi --showpids 2>&1 | head -40 + exit 1 + fi + fi + sleep 1 +done + +echo "Server stopped." diff --git a/scripts/wait_infer_drain.sh b/scripts/wait_infer_drain.sh new file mode 100755 index 0000000000..ba78d94c32 --- /dev/null +++ b/scripts/wait_infer_drain.sh @@ -0,0 +1,191 @@ +#!/bin/bash +# Watch an in-flight ATOM inference workload (server OR offline simple_inference). +# Exit 0 when the workload drains cleanly, exit 1 on hang, exit 2 on GPU fault. +# +# Workload modes (auto-detected by process pattern): +# - Server mode (atom.entrypoints): drains when eval client (lm_eval / curl +# / benchmark) is gone and a single poll shows no new output. +# Hang = STUCK_POLLS consecutive polls with no new "Engine +# Core: output send" while a client is still hammering. +# - Offline mode (atom.examples.simple_inference / atom.examples.benchmark): +# drains when the inference process exits naturally (no fault +# in log). Hang detection is N/A (no client concept; no +# Engine Core ZMQ progress). +# +# Server log auto-discovery: the canonical signal for "engine made progress" +# is `Engine Core: output send` in the server's stdout. Rather than asking +# the caller to know where the server log lives, the script `readlink`s +# `/proc//fd/1` of the atom.entrypoints process and uses THAT as the +# authoritative server log — path-independent across repo layouts. Falls +# back to the caller-supplied LOG_FILE only if /proc lookup fails. +# +# The caller-supplied LOG_FILE is still useful as a SECONDARY signal: +# - Fault grep runs on BOTH server log AND caller LOG_FILE (max coverage) +# - File-mtime growth on caller LOG_FILE counts as progress when the +# engine marker is absent (e.g. simple_inference offline, benchmark tqdm) +# +# Use this script for ANY blocking wait on an ATOM workload — don't fall back +# to `sleep + tail`. Offline path is supported in-script so you can wrap +# simple_inference the same way as server runs. +# +# Liveness: pgrep `$SERVER_PATTERN` (NOT curl /v1/models — under heavy load +# the HTTP endpoint can fail to respond within reasonable timeout even when +# the server is alive, producing false-positive "dead" reports). +# +# Usage: bash scripts/wait_infer_drain.sh [PORT] [MAX_MIN] [POLL_SEC] [LOG_FILE] [STUCK_POLLS] +# PORT default 8000 (kept for API symmetry with wait_server_ready.sh; unused in offline mode) +# MAX_MIN default 30 (full GSM8K 1319 typically 5-15 min on V4-Pro) +# POLL_SEC default 10 (fast poll for quick hang detection) +# LOG_FILE default empty (server log auto-discovered via /proc). Pass a +# client/workload log (e.g. benchmark.log, simple_inference.log) +# to add it as a secondary signal source for fault grep and +# mtime-based progress detection. +# STUCK_POLLS default 6 (6 × 10s = 1 min of no progress → declare hang) +# +# Exit codes: +# 0 — workload drained cleanly (server: client gone + no pending output; +# offline: process exited without fault) +# 1 — hang detected (server only: no progress while client running) +# 2 — fault detected (MEMORY_VIOLATION / ASSERT_TRAP / Memory access fault / proc died) +# 4 — max wait elapsed without resolution + +set -uo pipefail + +PORT="${1:-8000}" +MAX_MIN="${2:-30}" +POLL="${3:-10}" +LOG_FILE="${4:-}" +STUCK_POLLS="${5:-6}" +ITERS=$(( MAX_MIN * 60 / POLL )) + +# Match anything that indicates the eval driver is still hammering the +# server. Extend if you use a different client. +CLIENT_PATTERN='lm_eval|curl.*v1/(completions|chat)|atom\.examples\.benchmark|atom\.benchmarks\.benchmark' +# Server- or offline-mode workload process. simple_inference and +# atom.examples.benchmark also count — process-exit + no-fault = drain. +SERVER_PATTERN='atom\.entrypoints|atom\.examples\.simple_inference' + +# Resolve the authoritative server/workload log by reading the workload +# process's stdout fd. Re-run every poll because the workload PID can change +# (start_atom_server has just spawned it; CG capture forks; etc.) and we +# don't want a stale path. Returns empty string on failure. +discover_server_log() { + local pid log + pid=$(pgrep -f "$SERVER_PATTERN" 2>/dev/null | head -1) + [ -z "$pid" ] && return + log=$(readlink "/proc/$pid/fd/1" 2>/dev/null) + # Reject /dev/pts/, pipes, sockets, /dev/null — only real files count. + case "$log" in + /*) [ -f "$log" ] && echo "$log" ;; + esac +} + +# Grep helper that tolerates missing/empty file and never returns a count +# polluted with grep's stderr. +count_in() { + local pat=$1 file=$2 c + [ -z "$file" ] || [ ! -r "$file" ] && { echo 0; return; } + c=$(grep -cE "$pat" "$file" 2>/dev/null | head -1) + echo "${c:-0}" +} + +mtime_of() { + local file=$1 + [ -z "$file" ] || [ ! -r "$file" ] && { echo 0; return; } + stat -c %Y "$file" 2>/dev/null || echo 0 +} + +FAULT_PATTERN='stopped, reason|MEMORY_VIOLATION|ASSERT_TRAP|proc died unexpectedly|Memory access fault by GPU' + +prev_outputs=0 +prev_mtime=0 +stuck=0 +server_log="" + +for ((i=1; i<=ITERS; i++)); do + sleep "$POLL" + + # Refresh discovered server log (PID may have just appeared). + new_log=$(discover_server_log) + if [ -n "$new_log" ] && [ "$new_log" != "$server_log" ]; then + server_log="$new_log" + echo "[t=$((i*POLL))s] server log auto-discovered: $server_log" + fi + + # GPU fault? Scan BOTH server log AND caller LOG_FILE (max coverage). + # Check BEFORE process-exit so that fault-then-exit is attributed to + # fault (exit 2), not normal drain. + fault_server=$(count_in "$FAULT_PATTERN" "$server_log") + fault_caller=$(count_in "$FAULT_PATTERN" "$LOG_FILE") + fault_total=$(( fault_server + fault_caller )) + if [ "$fault_total" -gt 0 ]; then + echo "[t=$((i*POLL))s] GPU fault detected ($fault_total signals) — exiting 2" + for f in "$server_log" "$LOG_FILE"; do + [ -n "$f" ] && [ -r "$f" ] && grep -E "$FAULT_PATTERN" "$f" 2>/dev/null | head -3 + done + exit 2 + fi + + # Workload process alive? Only by process presence — no curl (HTTP can + # false-negative under heavy load). + if ! pgrep -f "$SERVER_PATTERN" >/dev/null 2>&1; then + # No fault grep matched above + process gone = clean exit (drain). + # Covers both stop_atom_server (server mode) and simple_inference + # finishing its prompts (offline mode). + echo "[t=$((i*POLL))s] workload process exited cleanly — DRAINED" + exit 0 + fi + + # Engine progress? Two signals — either resets the stuck counter: + # 1. "Engine Core: output send" count rising in server log + # (auto-discovered; authoritative engine progress). + # 2. Caller LOG_FILE mtime advancing (covers client logs / offline + # stdout where engine marker is absent: benchmark tqdm, + # simple_inference, etc.). + cur_outputs=$(count_in "Engine Core: output send" "$server_log") + delta_out=$(( cur_outputs - prev_outputs )) + cur_mtime=$(mtime_of "$LOG_FILE") + delta_mtime=$(( cur_mtime - prev_mtime )) + + # Client still running? + client_alive=$(pgrep -af "$CLIENT_PATTERN" 2>/dev/null | grep -v grep | wc -l) + + echo "[t=$((i*POLL))s] outputs=${cur_outputs} (+${delta_out}) mtime+${delta_mtime}s clients=${client_alive} stuck=${stuck}/${STUCK_POLLS}" + + if [ "$delta_out" -eq 0 ] && [ "$delta_mtime" -eq 0 ]; then + stuck=$(( stuck + 1 )) + else + stuck=0 + prev_outputs=$cur_outputs + prev_mtime=$cur_mtime + fi + + # Drained cleanly: client gone AND no new output this poll → done. + # No need to wait STUCK_POLLS — once the client is gone no new requests + # can arrive, so a single quiet poll is definitive. + if [ "$client_alive" -eq 0 ] && [ "$stuck" -ge 1 ]; then + echo "Engine DRAINED (client gone + no pending output)" + exit 0 + fi + + # Hung: no progress AND clients still alive → declare hang + if [ "$stuck" -ge "$STUCK_POLLS" ] && [ "$client_alive" -gt 0 ]; then + echo "HANG detected (no progress for ${stuck} polls while ${client_alive} client(s) still running)" + for f in "$server_log" "$LOG_FILE"; do + if [ -n "$f" ] && [ -r "$f" ]; then + echo "--- last 20 lines of $f ---" + tail -20 "$f" 2>/dev/null + fi + done + exit 1 + fi +done + +echo "MAX_WAIT ${MAX_MIN} min elapsed without resolution" +for f in "$server_log" "$LOG_FILE"; do + if [ -n "$f" ] && [ -r "$f" ]; then + echo "--- last 20 lines of $f ---" + tail -20 "$f" 2>/dev/null + fi +done +exit 4 diff --git a/scripts/wait_server_ready.sh b/scripts/wait_server_ready.sh new file mode 100755 index 0000000000..1f489a9191 --- /dev/null +++ b/scripts/wait_server_ready.sh @@ -0,0 +1,41 @@ +#!/bin/bash +# Poll http://localhost:$PORT/v1/models until the ATOM server is ready or +# the max wait elapses. Exit 0 = ready, 1 = timeout / startup error. +# +# Usage: bash scripts/wait_server_ready.sh [PORT] [MAX_MIN] [POLL_SEC] [LOG_FILE] +# PORT default 8000 +# MAX_MIN default 6 (server warmup typically 1-3 min; under debug agent 3-5 min) +# POLL_SEC default 30 +# LOG_FILE default /app/logs_claude/atom_server.log (used to detect startup errors) +# +# Side effect: every poll prints "[t=Ns] ready=...". Last line shows pass/fail. + +set -uo pipefail + +PORT="${1:-8000}" +MAX_MIN="${2:-6}" +POLL="${3:-30}" +LOG_FILE="${4:-/app/logs_claude/atom_server.log}" +ITERS=$(( MAX_MIN * 60 / POLL )) + +for ((i=1; i<=ITERS; i++)); do + sleep "$POLL" + READY=$(curl -s -m 3 "http://localhost:${PORT}/v1/models" 2>/dev/null | head -c 60) + ERR=$(grep -c "cluster_dims\|InductorError\|SHUTDOWN signal\|proc died" \ + "$LOG_FILE" 2>/dev/null | head -1) + ERR="${ERR:-0}" + echo "[t=$((i*POLL))s] ready=${READY:-(empty)} err=$ERR" + if [ -n "$READY" ]; then + echo "Server READY" + exit 0 + fi + if [ "$ERR" -gt 0 ]; then + echo "Server FAILED to start (errors detected)" + tail -30 "$LOG_FILE" + exit 1 + fi +done + +echo "Server NOT ready after ${MAX_MIN} min" +tail -30 "$LOG_FILE" +exit 1 diff --git a/tests/conftest.py b/tests/conftest.py index 18ae73642b..16d0cac657 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -118,6 +118,7 @@ def __init__(self, **overrides): enable_prefix_caching=False, max_num_seqs=4, max_num_batched_tokens=64, + max_model_len=64, bos_token_id=1, eos_token_id=2, stop_token_ids=[], diff --git a/tests/entrypoints/test_reasoning.py b/tests/entrypoints/test_reasoning.py index 0539528198..0961fe488e 100644 --- a/tests/entrypoints/test_reasoning.py +++ b/tests/entrypoints/test_reasoning.py @@ -82,6 +82,20 @@ def test_whitespace_after_thinking(self): reasoning, content = separate_reasoning(text) assert content == "The answer." + def test_no_think_start_tag(self): + """MiniMax M2.7 pattern: model doesn't generate , only . + The chat template injects as part of the prompt.""" + text = "The user wants hello world...\n\n\nprint('Hello')" + reasoning, content = separate_reasoning(text) + assert reasoning == "The user wants hello world..." + assert content == "print('Hello')" + + def test_no_think_start_tag_empty_content(self): + text = "Reasoning only\n" + reasoning, content = separate_reasoning(text) + assert reasoning == "Reasoning only" + assert content == "" + # ============================================================================ # ReasoningFilter (Streaming) Tests diff --git a/tests/plugin/test_sglang_model_wrapper.py b/tests/plugin/test_sglang_model_wrapper.py index e4015ed9dc..20d0e07923 100644 --- a/tests/plugin/test_sglang_model_wrapper.py +++ b/tests/plugin/test_sglang_model_wrapper.py @@ -54,8 +54,7 @@ def _make_fake_modules(*, is_last_rank: bool, setup_hook=None) -> dict[str, Modu forward_batch_mod.ForwardBatch = object forward_batch_mod.PPProxyTensors = object - attn_backend_pkg = _package("atom.plugin.sglang.attention_backend") - mla_mod = ModuleType("atom.plugin.sglang.attention_backend.sgl_attention_mla") + mla_mod = ModuleType("atom.plugin.sglang.models.deepseek_mla") mla_mod.setup_deepseek_for_sglang = setup_hook or (lambda model: None) return { @@ -68,8 +67,7 @@ def _make_fake_modules(*, is_last_rank: bool, setup_hook=None) -> dict[str, Modu "sglang.srt.layers.quantization.base_config": quant_base_mod, "sglang.srt.model_executor": model_executor_pkg, "sglang.srt.model_executor.forward_batch_info": forward_batch_mod, - "atom.plugin.sglang.attention_backend": attn_backend_pkg, - "atom.plugin.sglang.attention_backend.sgl_attention_mla": mla_mod, + "atom.plugin.sglang.models.deepseek_mla": mla_mod, } diff --git a/tests/plugin/test_sglang_register.py b/tests/plugin/test_sglang_register.py index 562aaf8bc8..1adb20fb42 100644 --- a/tests/plugin/test_sglang_register.py +++ b/tests/plugin/test_sglang_register.py @@ -324,10 +324,13 @@ def __init__(self, runner): "atom.models.qwen3_moe": ModuleType("atom.models.qwen3_moe"), "atom.models.glm4_moe": ModuleType("atom.models.glm4_moe"), "atom.models.deepseek_v2": ModuleType("atom.models.deepseek_v2"), + "atom.models.minimax_m2": ModuleType("atom.models.minimax_m2"), + "atom.models.qwen3_next": ModuleType("atom.models.qwen3_next"), + "atom.models.qwen3_5": ModuleType("atom.models.qwen3_5"), "atom.config": ModuleType("atom.config"), "atom.plugin.prepare": fake_prepare_mod, - "atom.plugin.sglang.attention_backend.sgl_attn_backend": ModuleType( - "atom.plugin.sglang.attention_backend.sgl_attn_backend" + "atom.plugin.sglang.attention_backend.full_attention.full_attention_backend": ModuleType( + "atom.plugin.sglang.attention_backend.full_attention.full_attention_backend" ), } fake_modules["atom.models.qwen3"].Qwen3ForCausalLM = type( @@ -342,9 +345,21 @@ def __init__(self, runner): fake_modules["atom.models.deepseek_v2"].DeepseekV3ForCausalLM = type( "DeepseekV3ForCausalLM", (), {} ) + fake_modules["atom.models.minimax_m2"].MiniMaxM2ForCausalLM = type( + "MiniMaxM2ForCausalLM", (), {} + ) + fake_modules["atom.models.qwen3_next"].Qwen3NextForCausalLM = type( + "Qwen3NextForCausalLM", (), {} + ) + fake_modules["atom.models.qwen3_5"].Qwen3_5ForCausalLM = type( + "Qwen3_5ForCausalLM", (), {} + ) + fake_modules["atom.models.qwen3_5"].Qwen3_5MoeForCausalLM = type( + "Qwen3_5MoeForCausalLM", (), {} + ) fake_modules["atom.config"].Config = type("Config", (), {}) fake_modules[ - "atom.plugin.sglang.attention_backend.sgl_attn_backend" + "atom.plugin.sglang.attention_backend.full_attention.full_attention_backend" ].ATOMAttnBackendForSgl = _FakeBackend with patch.dict(sys.modules, fake_modules): diff --git a/tests/test_block_manager.py b/tests/test_block_manager.py index 60a6c10612..9ad059ff2d 100644 --- a/tests/test_block_manager.py +++ b/tests/test_block_manager.py @@ -301,10 +301,11 @@ def test_preempt_and_reschedule_reuses_cache(self, seq_factory): assert s1.num_cached_tokens == 0 assert s1.block_table == [] - # Re-allocate — should get cache hits on both blocks + # Re-allocate — first block is a cache hit; the last full block is + # force-recomputed so prefill has at least one token to forward. s1_retry = seq_factory([1, 2, 3, 4, 5, 6, 7, 8]) bm.allocate(s1_retry) - assert s1_retry.num_cached_tokens == 8 # both blocks cached + assert s1_retry.num_cached_tokens == 4 # ── Edge cases ─────────────────────────────────────────────────────────── @@ -325,8 +326,9 @@ def test_single_token_no_cache(self, seq_factory): # Partial block → hash is -1 → no caching assert s2.num_cached_tokens == 0 - def test_exact_block_size_fully_cached(self, seq_factory): - """Sequence with exactly block_size tokens — fully cached on reuse.""" + def test_exact_block_size_last_block_recomputed(self, seq_factory): + """Single-block prompt: last full block is force-recomputed on reuse so + prefill has at least one token to forward and produce logits.""" cfg = MockConfig( num_kvcache_blocks=4, kv_cache_block_size=4, enable_prefix_caching=True ) @@ -336,7 +338,7 @@ def test_exact_block_size_fully_cached(self, seq_factory): bm.deallocate(s1) s2 = seq_factory([1, 2, 3, 4]) bm.allocate(s2) - assert s2.num_cached_tokens == 4 + assert s2.num_cached_tokens == 0 def test_free_block_ids_set_consistent(self, block_manager, seq_factory): """free_block_ids_set stays consistent through allocate/deallocate.""" diff --git a/test/kv_transfer/disaggregation/test_kv_aggregator.py b/tests/test_kv_aggregator.py similarity index 100% rename from test/kv_transfer/disaggregation/test_kv_aggregator.py rename to tests/test_kv_aggregator.py diff --git a/test/kv_transfer/disaggregation/test_kv_connector_scheduler.py b/tests/test_kv_connector_scheduler.py similarity index 100% rename from test/kv_transfer/disaggregation/test_kv_connector_scheduler.py rename to tests/test_kv_connector_scheduler.py diff --git a/tests/test_mxfp4_moe_has_bias.py b/tests/test_mxfp4_moe_has_bias.py index a9255bfddc..a195392c6c 100644 --- a/tests/test_mxfp4_moe_has_bias.py +++ b/tests/test_mxfp4_moe_has_bias.py @@ -19,10 +19,28 @@ import sys import unittest -# Clear cached atom modules (conftest.py stubs) -for mod_name in list(sys.modules): - if mod_name.startswith("atom"): - del sys.modules[mod_name] +# This test needs to inspect real atom source (not conftest.py stubs), so it +# wipes any cached `atom.*` modules at module-import time. Previously this +# also wiped the conftest stubs and never restored them, polluting later +# tests (test_arg_utils_spec / test_scheduler / test_sequence) that depend +# on the stubs. setUpModule / tearDownModule now snapshots-and-restores +# sys.modules so this file's effect is local to its own collection. +_saved_atom_modules: dict[str, object] = {} + + +def setUpModule(): + global _saved_atom_modules + _saved_atom_modules = { + name: mod for name, mod in sys.modules.items() if name.startswith("atom") + } + for name in list(_saved_atom_modules): + del sys.modules[name] + + +def tearDownModule(): + for name in [n for n in sys.modules if n.startswith("atom")]: + del sys.modules[name] + sys.modules.update(_saved_atom_modules) class TestFusedMoEDefaultHasBias(unittest.TestCase): @@ -201,6 +219,15 @@ class TestSwiGLUInterleavingWithoutBias(unittest.TestCase): and guard only the bias interleaving on ``layer.w13_bias is not None``. """ + @unittest.skip( + "Obsolete: Mxfp4MoEMethod.process_weights_after_loading no longer " + "branches on `layer.activation == ActivationType.Swiglu`. The function " + "now routes via `use_triton` (Triton swizzle) vs. the AITER shuffle " + "path, with bias cast handled unconditionally up top. The original " + "regression this guard was added for — the SwiGLU branch being " + "incorrectly gated on `w13_bias is not None` — cannot recur in the " + "current structure. Re-evaluate or delete when revisiting Mxfp4 MoE." + ) def test_swiglu_branch_condition_no_bias_check(self): """The SwiGLU branch must NOT require bias to be present.""" import inspect diff --git a/tests/test_per_req_cache_decoupling.py b/tests/test_per_req_cache_decoupling.py index 81166817ac..07ef2000e6 100644 --- a/tests/test_per_req_cache_decoupling.py +++ b/tests/test_per_req_cache_decoupling.py @@ -22,6 +22,7 @@ def per_req_cache_config(**overrides): enable_prefix_caching=False, max_num_seqs=8, max_num_batched_tokens=256, + max_model_len=256, bos_token_id=1, eos_token_id=2, stop_token_ids=[], diff --git a/test/kv_transfer/disaggregation/test_proxy.py b/tests/test_proxy.py similarity index 100% rename from test/kv_transfer/disaggregation/test_proxy.py rename to tests/test_proxy.py diff --git a/test/kv_transfer/disaggregation/test_transfer_engine.py b/tests/test_transfer_engine.py similarity index 100% rename from test/kv_transfer/disaggregation/test_transfer_engine.py rename to tests/test_transfer_engine.py diff --git a/test/kv_transfer/disaggregation/test_types.py b/tests/test_types.py similarity index 100% rename from test/kv_transfer/disaggregation/test_types.py rename to tests/test_types.py diff --git a/tools/backfill_dashboard_gpu_name.py b/tools/backfill_dashboard_gpu_name.py new file mode 100644 index 0000000000..d9268fd3e4 --- /dev/null +++ b/tools/backfill_dashboard_gpu_name.py @@ -0,0 +1,119 @@ +#!/usr/bin/env python3 +"""Backfill historical `AMD Radeon Graphics` labels in the dashboard data.js. + +`benchmark-action/github-action-benchmark@v1` (the action that powers +https://rocm.github.io/ATOM/benchmark-dashboard/) is append-only: every nightly +run pushes a new `{commit, date, benches}` entry to `benchmark-dashboard/data.js` +on the `gh-pages` branch and never rewrites history. + +That meant the entries written before PR #692 — and the entries written by +accuracy-validation workflows that use `linux-atom-mi35x-*` runners which a +follow-up patch fixed — are permanently stuck with the mis-detected +`GPU: AMD Radeon Graphics` label, even though every ATOM CI machine is +physically an MI355X. + +This script does a one-time, in-place rewrite of `data.js` to relabel those +entries. It is intentionally a literal text substitution (not a JSON +round-trip) so the action's exact formatting is preserved and the diff stays +inspectable. + +Assumptions (verified against the live data.js on 2026-05-11): + +* The dashboard only ever stores two GPU labels: `AMD Radeon Graphics` (6357 + occurrences, all mis-detections) and `AMD Instinct MI355X` (4321 correct). +* No MI300X / MI250X / MI210 entries have ever been written — the ATOM project + has only run on MI355X-class hardware since the dashboard was created. + +If either assumption stops holding (e.g. someone adds an MI300X runner), this +script needs to be tightened to gate the substitution on the runner hint. + +Usage: + # Inspect what would change (default — no writes): + python tools/backfill_dashboard_gpu_name.py path/to/data.js + + # Apply the rewrite in place: + python tools/backfill_dashboard_gpu_name.py path/to/data.js --write + + # Full end-to-end against the live gh-pages branch: + git fetch origin gh-pages + git worktree add /tmp/atom-gh-pages gh-pages + python tools/backfill_dashboard_gpu_name.py /tmp/atom-gh-pages/benchmark-dashboard/data.js --write + cd /tmp/atom-gh-pages + git add benchmark-dashboard/data.js + git commit -m "ci(dashboard): backfill MI355X label on historical Radeon entries" + git push origin gh-pages +""" + +from __future__ import annotations + +import argparse +import re +import sys +from pathlib import Path + +WRONG = "GPU: AMD Radeon Graphics" +RIGHT = "GPU: AMD Instinct MI355X" + +# Defensive: only touch occurrences that appear inside a benchmark `extra` +# string. Every `extra` value the action writes starts with `Run: https://...`, +# so anchoring on `"extra": "Run:` makes accidental matches in commit messages +# or other free-form fields essentially impossible. +EXTRA_LINE_RE = re.compile(r'("extra":\s*"Run: [^"]*?)' + re.escape(WRONG)) + + +def rewrite(text: str) -> tuple[str, int]: + """Return (new_text, replacement_count).""" + new_text, count = EXTRA_LINE_RE.subn(r"\1" + RIGHT, text) + return new_text, count + + +def main() -> int: + p = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + p.add_argument("data_js", type=Path, help="Path to benchmark-dashboard/data.js") + p.add_argument( + "--write", + action="store_true", + help="Rewrite the file in place (default is dry-run).", + ) + args = p.parse_args() + + if not args.data_js.is_file(): + print(f"error: {args.data_js} not found", file=sys.stderr) + return 2 + + original = args.data_js.read_text(encoding="utf-8") + pre_wrong = original.count(WRONG) + pre_right = original.count(RIGHT) + + new_text, n = rewrite(original) + post_wrong = new_text.count(WRONG) + post_right = new_text.count(RIGHT) + + print(f"file: {args.data_js}") + print(f"size: {len(original):,} bytes") + print(f"before: {pre_wrong:>6} Radeon, {pre_right:>6} MI355X") + print(f"after: {post_wrong:>6} Radeon, {post_right:>6} MI355X") + print(f"replacements made: {n:,}") + + if post_wrong != 0: + print( + f"warning: {post_wrong} Radeon mention(s) remain outside of " + f'`"extra": "Run: ..."` strings; left untouched.', + file=sys.stderr, + ) + + if not args.write: + print("\nDry-run only. Re-run with --write to apply.") + return 0 + + if n == 0: + print("\nNothing to write.") + return 0 + + args.data_js.write_text(new_text, encoding="utf-8") + print(f"\nWrote {args.data_js}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main())