[Draft][MoE][Offload] Expert pool: shared GPU bank of NVFP4 Marlin expert rows with a device-side planner (moe_expert_pool_rows)

benchmarks/expert_pool/LICENSE.prompts

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025.
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

benchmarks/expert_pool/README.md

@@ -0,0 +1,127 @@
+# Expert重みのGPUキャッシュ：速度測定の再現手順
+
+Qwen3.8 FlashNext NVFP4に、準備用の要求と速度測定用の要求を順番に送ります。
+各要求ではソフトウェアの不具合報告を読み、修正案と検証計画を生成します。
+
+## ファイルと要求の内容
+
+| ファイル | 内容 |
+| --- | --- |
+| [benchmark.py](benchmark.py) | HTTP要求、ストリーム保存、生成速度の集計 |
+| [pair.json](pair.json) | 実測に使用したプロンプト全文、固定トークン列、出典 |
+| [prompts.md](prompts.md) | プロンプト全文の読みやすい表示 |
+| [provenance.json](provenance.json) | 実測コードの版、元ファイルのSHA256、移植差分 |
+| [LICENSE.prompts](LICENSE.prompts) | プロンプトの出典に付属するMITライセンス |
+
+| 送信順 | 用途 | 不具合報告 | 入力トークン数 |
+| --- | --- | --- | --- |
+| 1回目 | 準備用（結果のroleはwarmup） | ファイル選択ボタンの「Choose File」を「Choose file」に修正する課題（28096_836） | 1070 |
+| 2回目 | 速度測定用（roleはmeasure） | 言語設定を変更したとき「Link sent!」表示も更新する課題（18827_741） | 753 |
+
+出典は[OpenAI frontier-evals](https://github.com/openai/frontier-evals/tree/51052cede8cc608f95bb00346635e03759013e5a)のSWE-Lancerです。
+既存の動作確認用課題から選んだ2件で、測定対象は修正案の文章生成です。
+実際のコード編集・公式採点を行う品質試験は別の手順です。
+
+## 実測に使用したコードと環境
+
+| 用途 | コードの版 |
+| --- | --- |
+| 比較の土台となるvLLM main | `a97dacb7106ee49f39f3d1fc6ae1800ff724e01d` |
+| Expert重みのGPUキャッシュの実装（[fork PR #48](https://github.com/01554/vllm/pull/48)） | `5fbc240ba5ddec82a10362340ac77339a1c24017` |
+| PLEの読み出しと生成計算を重ねる実装（[fork PR #46](https://github.com/01554/vllm/pull/46)、[上流PR #54129](https://github.com/vllm-project/vllm/pull/54129)を前提とする差分） | `4f859de9d0f55760b50358aee4834e6966e13bc8` |
+| 上記機能を組み合わせ、以下の速度を測定した版 | `7dedc6d8d9b178b60f6a5b32f03d677145982441` |
+
+サーバーは実測版のPythonソース、上記mainからビルドしたwheel、別途ビルドした
+`_ple_memops`拡張を組み合わせて動かしました。このディレクトリは測定後に追加した
+クライアント用ファイルです。`pair.json`は実測ファイルのbyteコピーです。
+
+RTX 6000 Adaの48GBに収める構成の検証を目的として、手元の
+RTX PRO 6000 Blackwell Max-Q（96GiB）上で別プロセスにGPUメモリを確保させ、
+サーバーに利用可能な容量を48GiBにして測定しました。以下はこのGPU上の実測値です。
+ホスト側のコンテナのメモリ上限は100GiBでした。
+
+GPUキャッシュは48層それぞれ258 expert行、約32GiBです。
+容量制限は外部プロセスで設定します。下記の`gpu-memory-utilization`は物理GPU容量に
+対するvLLMの予算比率です。クライアント実行前に容量とサーバー起動状態を確認してください。
+
+## サーバーの起動設定
+
+上の実測版と同じ機能をビルドした環境で、チェックポイントのパスを指定します。
+`VLLM_USE_BREAKABLE_CUDAGRAPH`は未設定（自動選択）で測定しました。
+
+```bash
+export CHECKPOINT=/data/models/Qwen3.8-Flash-Next-NVFP4-nvidia
+unset VLLM_USE_BREAKABLE_CUDAGRAPH
+export VLLM_DEBUG_WORKSPACE=1 VLLM_LOGGING_LEVEL=DEBUG
+export PYTORCH_ALLOC_CONF=pinned_max_round_threshold_mb:1,pinned_max_cached_size_mb:1
+export OMP_NUM_THREADS=2 OPENBLAS_NUM_THREADS=2
+export VLLM_USE_V2_MODEL_RUNNER=1
+export VLLM_PLE_MMAP=1 VLLM_PLE_MMAP_DEFERRED=1
+export VLLM_PLE_MMAP_PREWARM=0 VLLM_PLE_MMAP_PINNED=0
+export HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1
+.venv/bin/python -m vllm.entrypoints.openai.api_server \
+  --model "$CHECKPOINT" --served-model-name flashnext \
+  --host 0.0.0.0 --port 8000 --tensor-parallel-size 1 \
+  --quantization modelopt --dtype bfloat16 --moe-backend marlin \
+  --moe-expert-pool-rows 258 --language-model-only \
+  --max-model-len 4096 --max-num-seqs 1 --max-num-batched-tokens 512 \
+  --compilation-config '{"cudagraph_mode":"FULL_DECODE_ONLY"}' \
+  --no-enable-flashinfer-autotune --gpu-memory-utilization 0.4548806288994517 \
+  --safetensors-load-strategy lazy \
+  --default-chat-template-kwargs '{"enable_thinking":false}' \
+  --reasoning-parser qwen3 --generation-config vllm
+```
+
+速度測定のコンテキスト上限は4096です。別途実施した品質試験では32768を使いました。
+
+## クライアントの実行
+
+サーバーの準備完了後、次を1回実行します。クライアントはPython標準ライブラリで動きます。
+
+```bash
+.venv/bin/python benchmarks/expert_pool/benchmark.py \
+  --base-url http://127.0.0.1:8000 --model flashnext \
+  --label fresh-0 --output results/fresh-0.jsonl
+```
+
+送信先は`/v1/completions`です。`pair.json`の固定トークン列を送信するので、
+チェックポイントのtokenizerが`pair.json`の`tokenization`に記録したSHA256と
+一致することを確認してください。要求には`temperature=0`、`top_p=1`、`seed=0`、
+`max_tokens=2048`を指定し、thinkingを無効にしたチャットテンプレートのトークン列を使います。
+`tokenization.pair_sha256`はトークン列追加前の資料のハッシュで、ファイル全体のハッシュは
+`provenance.json`にあります。
+
+3回の測定では、**毎回サーバーを終了して新しいプロセスで起動し、準備用→測定用を1組送信**します。
+出力名を`fresh-0.jsonl`、`fresh-1.jsonl`、`fresh-2.jsonl`と変えます。
+実測はこの順で3組を実行し、2回目の要求の速度3値から中央値を求めました。
+クライアントは既存の出力ファイルを保護し、HTTPエラーや不完全なストリームを保存して停止します。
+失敗した測定も結果として保持してください。
+
+## 指標と実測値
+
+生成速度は、usageの生成トークン数とクライアント側の受信時刻から計算します。
+
+```text
+decode_tok_s = (completion_tokens - 1) / (最後の本文受信時刻 - 最初の本文受信時刻)
+```
+
+`first_token_s`は要求開始から最初の本文受信までの時間です。
+`e2e_tok_s`は要求全体の所要時間あたりの生成トークン数です。
+ストリームの1イベントに複数トークンが入ることがあるため、いずれもクライアント観測の値です。
+生のSSE、送信body、usage、finish reason、本文とSHA256も同じJSONLに保存します。
+
+| 新しいサーバーでの実行 | 2回目の要求の生成速度 |
+| --- | --- |
+| 1回目 | 63.1882 tok/s |
+| 2回目 | 62.7087 tok/s |
+| 3回目 | 63.6519 tok/s |
+| 中央値 | **63.1882 tok/s** |
+
+これは表の実測版で各機能を組み合わせた値です。各要求の終了理由は`stop`でした。
+容量は起動中・生成中の標本で記録し、プロセス終了コード0とOOMKilled=falseを確認しました。
+
+## クライアントの動作確認
+
+```bash
+.venv/bin/python -m unittest discover -s benchmarks/expert_pool -p 'test_*.py'
+```

benchmarks/expert_pool/benchmark.py

@@ -0,0 +1,189 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Send one frozen warmup request, then one measured request to a fresh server.
+
+This is issue-text generation timing, not the SWE-Lancer correctness evaluation.
+See README.md for the server configuration and the historical script provenance.
+"""
+
+import argparse
+import datetime
+import hashlib
+import json
+import math
+import time
+import urllib.error
+import urllib.request
+from pathlib import Path
+
+
+def main():
+    p = argparse.ArgumentParser(description=__doc__)
+    p.add_argument("--base-url", required=True)
+    p.add_argument("--label", required=True)
+    p.add_argument("--pair", type=Path, default=Path(__file__).with_name("pair.json"))
+    p.add_argument("--output", required=True)
+    p.add_argument("--model", default="flashnext")
+    p.add_argument("--max-tokens", type=int, default=2048)
+    p.add_argument("--timeout", type=int, default=1200)
+    a = p.parse_args()
+    source = Path(a.pair).read_bytes()
+    pair = json.loads(source)
+    tasks = pair["tasks"]
+    if len(tasks) != 2 or [t["role"] for t in tasks] != ["warmup", "measure"]:
+        raise ValueError(
+            "Expected exactly one warmup task followed by one measure task"
+        )
+    sequence = [(tasks[0], "warmup"), (tasks[1], "measure")]
+    output = Path(a.output)
+    output.parent.mkdir(parents=True, exist_ok=True)
+    # Refuse to overwrite earlier measurements.
+    with output.open("x") as out:
+        for index, (task, role) in enumerate(sequence):
+            prompt = task["prompt_token_ids"]
+            request_body = dict(
+                model=a.model,
+                prompt=prompt,
+                max_tokens=a.max_tokens,
+                temperature=0,
+                top_p=1,
+                stream=True,
+                stream_options={"include_usage": True},
+                seed=0,
+            )
+            request = urllib.request.Request(
+                a.base_url.rstrip("/") + "/v1/completions",
+                data=json.dumps(request_body).encode(),
+                headers={"Content-Type": "application/json"},
+            )
+            result = dict(
+                schema_version=1,
+                timestamp=datetime.datetime.now(datetime.timezone.utc).isoformat(),
+                label=a.label,
+                sequence_index=index,
+                role=role,
+                task_id=task["id"],
+                pair_sha256=hashlib.sha256(source).hexdigest(),
+                prompt_sha256=task["prompt_sha256"],
+                expected_prompt_token_ids=task["prompt_token_ids"],
+                request=request_body,
+                content="",
+                usage=None,
+                finish_reason=None,
+                first_token_s=None,
+                last_token_s=None,
+                stream_events=0,
+            )
+            print(
+                f"Starting {a.label}: {role} {task['id']} "
+                f"({len(task['prompt_token_ids'])} input tokens)",
+                flush=True,
+            )
+            result["raw_sse_events"] = []
+            result["sse_done"] = False
+            start = time.perf_counter()  # Includes HTTP request and server scheduling.
+            try:
+                with urllib.request.urlopen(request, timeout=a.timeout) as response:
+                    for raw in response:
+                        line = raw.decode().strip()
+                        if not line.startswith("data:"):
+                            continue
+                        data = line[5:].strip()
+                        if data == "[DONE]":
+                            result["sse_done"] = True
+                            break
+                        result["raw_sse_events"].append(data)
+                        event = json.loads(data)
+                        if event.get("error"):
+                            raise RuntimeError(event["error"])
+                        if event.get("usage"):
+                            result["usage"] = event["usage"]
+                        for choice in event.get("choices", []):
+                            content = choice.get("text") or ""
+                            if content:
+                                elapsed = time.perf_counter() - start
+                                if result["first_token_s"] is None:
+                                    result["first_token_s"] = elapsed
+                                result["last_token_s"] = elapsed
+                                result["content"] += content
+                                result["stream_events"] += 1
+                            if choice.get("finish_reason"):
+                                result["finish_reason"] = choice["finish_reason"]
+                result["elapsed_s"] = time.perf_counter() - start
+                if not result["sse_done"]:
+                    raise RuntimeError("SSE ended without DONE")
+                completion_tokens = (result["usage"] or {}).get("completion_tokens")
+                if (
+                    type(completion_tokens) is not int
+                    or completion_tokens < 1
+                    or not result["content"]
+                ):
+                    raise RuntimeError(
+                        "Missing completion token usage or nonempty output"
+                    )
+                if result["finish_reason"] not in ("stop", "length"):
+                    raise RuntimeError(
+                        "Incomplete or unexpected finish reason: "
+                        f"{result['finish_reason']}"
+                    )
+                if (result["usage"] or {}).get("prompt_tokens") != len(
+                    task["prompt_token_ids"]
+                ):
+                    raise RuntimeError(
+                        "Server input token count does not match frozen prompt"
+                    )
+                result["completion_tokens"] = completion_tokens
+                result["e2e_tok_s"] = completion_tokens / result["elapsed_s"]
+                decode_duration = result["last_token_s"] - result["first_token_s"]
+                result["decode_tok_s"] = (
+                    (completion_tokens - 1) / decode_duration
+                    if completion_tokens > 1 and decode_duration > 0
+                    else None
+                )
+                for metric in ("e2e_tok_s", "decode_tok_s"):
+                    value = result[metric]
+                    if value is not None and not math.isfinite(value):
+                        raise RuntimeError(f"Non-finite {metric}")
+                result["decode_metric_note"] = (
+                    "Approximate client-observed "
+                    "(completion_tokens-1)/(last-text-event-first-text-event)"
+                )
+                result["output_sha256"] = hashlib.sha256(
+                    result["content"].encode()
+                ).hexdigest()
+                result["truncated"] = result["finish_reason"] == "length"
+                result["correctness"] = "Not graded: issue-text inference workload only"
+            except Exception as exc:
+                result["elapsed_s"] = time.perf_counter() - start
+                result["error"] = repr(exc)
+                if isinstance(exc, urllib.error.HTTPError):
+                    result["error_body"] = exc.read().decode(errors="replace")
+            out.write(json.dumps(result, ensure_ascii=False) + "\n")
+            out.flush()
+            print(
+                json.dumps(
+                    {
+                        k: result.get(k)
+                        for k in (
+                            "label",
+                            "role",
+                            "task_id",
+                            "completion_tokens",
+                            "first_token_s",
+                            "elapsed_s",
+                            "decode_tok_s",
+                            "e2e_tok_s",
+                            "finish_reason",
+                            "error",
+                        )
+                    },
+                    ensure_ascii=False,
+                ),
+                flush=True,
+            )
+            if result.get("error"):
+                raise RuntimeError(result["error"])
+
+
+if __name__ == "__main__":
+    main()