[Draft][MoE] Expert pool decode path: validation in the step kernel, one per-forward assertion, fewer launches

.buildkite/test_areas/kernels.yaml

@@ -231,13 +231,15 @@ steps:
   - csrc/quantization/cutlass_w8a8/moe/
   - csrc/moe/
   - tests/kernels/moe
+  - tests/kernels/expert_pool
   - vllm/model_executor/layers/fused_moe/
   - vllm/distributed/device_communicators/
   - vllm/envs.py
   - vllm/config
   commands:
     - pytest -v -s kernels/moe --ignore=kernels/moe/test_modular_oai_triton_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
     - pytest -v -s kernels/moe/test_modular_oai_triton_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
+    - pytest -v -s kernels/expert_pool --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
   parallelism: 5
   mirror:
     amd:
@@ -249,6 +251,7 @@ steps:
       - csrc/quantization/cutlass_w8a8/moe/
       - csrc/moe/
       - tests/kernels/moe
+      - tests/kernels/expert_pool
       - vllm/model_executor/layers/fused_moe/
       - vllm/distributed/device_communicators/
       - vllm/envs.py

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,142 @@
+# Expert pool: reproducing the generation-speed measurement
+
+Sends a warmup request and then a measured request to Qwen3.8-Flash-Next
+NVFP4. Each request is a software bug report; the model writes a fix
+proposal and a verification plan.
+
+## Files and requests
+
+| file | content |
+| --- | --- |
+| [benchmark.py](benchmark.py) | HTTP requests, stream capture, decode-speed computation |
+| [pair.json](pair.json) | the two prompts as measured: full text, frozen token ids, provenance |
+| [prompts.md](prompts.md) | the prompt texts, readable |
+| [provenance.json](provenance.json) | code versions measured, SHA256 of the original files, changes from the historical runner |
+| [LICENSE.prompts](LICENSE.prompts) | MIT license of the prompt source |
+
+| order | role | bug report | input tokens |
+| --- | --- | --- | --- |
+| 1st | warmup (`role` = `warmup`) | rename the "Choose File" button to "Choose file" (task 28096_836) | 1070 |
+| 2nd | measured (`role` = `measure`) | update the "Link sent!" message when the language setting changes (task 18827_741) | 753 |
+
+The prompts come from SWE-Lancer in
+[OpenAI frontier-evals](https://github.com/openai/frontier-evals/tree/51052cede8cc608f95bb00346635e03759013e5a),
+two of its existing sanity tasks. This measures fix-proposal text
+generation only; the code-editing and official-grading quality checks are a
+separate procedure.
+
+## Code and environment measured
+
+| purpose | version |
+| --- | --- |
+| vLLM main used as the base | `a97dacb7106ee49f39f3d1fc6ae1800ff724e01d` |
+| expert pool implementation (this PR, at the time of measurement) | `5fbc240ba5ddec82a10362340ac77339a1c24017` |
+| deferred PLE rows ([01554/vllm#46](https://github.com/01554/vllm/pull/46), a diff on top of upstream PR [#54129](https://github.com/vllm-project/vllm/pull/54129)) | `4f859de9d0f55760b50358aee4834e6966e13bc8` |
+| the combination of the above that produced the numbers below | `7dedc6d8d9b178b60f6a5b32f03d677145982441` |
+
+The server ran the Python sources of the measured combination over a wheel
+built from the base commit above, plus a separately built `_ple_memops`
+extension. This directory holds client-side files added after the
+measurement; `pair.json` is a byte copy of the file used.
+
+The goal was a configuration that fits an RTX 6000 Ada (48 GB). Measurements
+were taken on an RTX PRO 6000 Blackwell Max-Q (96 GiB) with a separate
+process holding GPU memory so that the server had 48 GiB available; the
+container's host memory limit was 100 GiB. The numbers below are from that GPU.
+
+The expert pool holds 258 expert rows per layer for 48 layers, about
+32 GiB. The memory limit is set by the external process; the
+`gpu-memory-utilization` below is vLLM's budget as a fraction of the
+physical GPU. Check the available memory and that the server is up before
+running the client.
+
+## Server launch
+
+Same features as the measured combination, checkpoint path as needed.
+`VLLM_USE_BREAKABLE_CUDAGRAPH` was unset (automatic selection).
+`VLLM_DEBUG_WORKSPACE`, `VLLM_LOGGING_LEVEL`, `PYTORCH_ALLOC_CONF` and the
+thread counts are the values used during measurement, not requirements.
+
+```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
+```
+
+The speed measurement uses a context limit of 4096. The separate quality
+checks used 32768.
+
+## Client
+
+Run once after the server is ready. The client needs only the Python
+standard library.
+
+```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
+```
+
+Requests go to `/v1/completions` with the frozen token ids from
+`pair.json`, so confirm that the checkpoint's tokenizer matches the SHA256
+recorded under `tokenization` in `pair.json`. Each request uses
+`temperature=0`, `top_p=1`, `seed=0`, `max_tokens=2048` and the
+chat-template token sequence with thinking disabled.
+`tokenization.pair_sha256` is the hash of the material before the token ids
+were added; the hash of the whole file is in `provenance.json`.
+
+For the three-run measurement, **stop the server and start a new process
+before each run, then send warmup -> measure once**, writing to
+`fresh-0.jsonl`, `fresh-1.jsonl`, `fresh-2.jsonl`. The reported value is the
+median of the three measured-request speeds. The client refuses to
+overwrite an existing output file, and saves HTTP errors and incomplete
+streams before stopping; keep failed runs as results too.
+
+## Metric and measured values
+
+Decode speed is computed from the usage completion-token count and the
+client-side receive times:
+
+```text
+decode_tok_s = (completion_tokens - 1) / (last text event time - first text event time)
+```
+
+`first_token_s` is the time from request start to the first text event;
+`e2e_tok_s` is completion tokens over the whole request time. A stream
+event may carry more than one token, so both are client-observed values.
+The raw SSE events, request body, usage, finish reason, text and its SHA256
+are stored in the same JSONL.
+
+| fresh server run | measured-request decode speed |
+| --- | --- |
+| 1 | 63.1882 tok/s |
+| 2 | 62.7087 tok/s |
+| 3 | 63.6519 tok/s |
+| median | **63.1882 tok/s** |
+
+Measured on the combination listed above. Every request finished with
+`stop`. Memory was sampled during startup and generation; the process
+exited 0 with OOMKilled=false.
+
+## Client self-check
+
+```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()