[Perf][GLM-5.3-Flash] Dense/masked-MHA sparse prefill for the NoPE (256, 0, 256) layout + skip the NoPE K concat

.buildkite/scripts/hardware_ci/run-cpu-test.sh

@@ -57,7 +57,7 @@ prune_if_disk_pressure
 # omits `--pull` so it already prefers a cached image over the network. A
 # single attempt only -- the retry loop below covers a miss here too.
 echo "--- :docker: Pre-fetching base image"
-docker pull ubuntu:22.04 || true
+docker pull ubuntu:25.04 || true
 
 # building the docker image
 echo "--- :docker: Building Docker image"

.buildkite/test_areas/disaggregated_mooncake.yaml

@@ -19,3 +19,57 @@ steps:
   commands:
     - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
     - bash v1/kv_connector/mooncake_integration/config_sweep_accuracy_test.sh
+
+- label: ":nvidia: (L4) Mooncake EC TCP E2E"
+  key: mooncake-ec-tcp-e2e-2-gpus
+  timeout_in_minutes: 30
+  working_dir: "/vllm-workspace"
+  device: l4
+  num_devices: 2
+  source_file_dependencies:
+    - vllm/distributed/ec_transfer/
+    - vllm/config/ec_transfer.py
+    - vllm/config/multimodal.py
+    - vllm/config/vllm.py
+    - vllm/multimodal/
+    - vllm/v1/core/encoder_cache_manager.py
+    - vllm/v1/core/sched/
+    - vllm/v1/engine/
+    - vllm/v1/worker/ec_connector_model_runner_mixin.py
+    - vllm/v1/worker/gpu_model_runner.py
+    - vllm/v1/worker/gpu_worker.py
+    - vllm/v1/worker/gpu/ec_connector.py
+    - vllm/v1/worker/gpu/model_runner.py
+    - vllm/v1/worker/gpu/mm/
+    - examples/disaggregated/disaggregated_encoder/
+    - tests/v1/ec_connector/
+    - requirements/kv_connectors.txt
+  env:
+    PYTHON_BIN: "/vllm-workspace/.venv/bin/python"
+    PYTHONUNBUFFERED: "1"
+    VLLM_HOST_IP: "127.0.0.1"
+    VLLM_USE_V2_MODEL_RUNNER: "1"
+    MOONCAKE_EC_PROTOCOL: "tcp"
+    USE_MM_PROMPTS: "1"
+    SKIP_BASELINE: "0"
+    CONCURRENCY: "3"
+    REPEAT: "2"
+    LOG_PATH: "/tmp/mooncake-ec-e2e"
+    BASELINE_FILE: "/tmp/mooncake-ec-e2e/baseline.json"
+  commands:
+    - uv venv --system-site-packages --python 3.12 .venv
+    - |
+      uv pip install --python .venv/bin/python "$(.venv/bin/python - <<'PY'
+      from pathlib import Path
+      import torch
+
+      requirement = next(
+          line for line in Path("requirements/kv_connectors.txt").read_text().splitlines()
+          if line.startswith("mooncake-transfer-engine ")
+      )
+      if torch.version.cuda.split(".")[0] == "13":
+          requirement = requirement.replace("mooncake-transfer-engine", "mooncake-transfer-engine-cuda13")
+      print(requirement)
+      PY
+      )"
+    - bash tests/v1/ec_connector/integration/run_epd_mooncake_ec_full_pipeline.sh

.github/CODEOWNERS

@@ -218,6 +218,19 @@ mkdocs.yaml @hmellor
 /vllm/config/pooler.py @noooop
 /vllm/model_executor/layers/pooler @noooop
 
+# Reinforcement learning
+/vllm/distributed/weight_transfer @aoshen02
+/vllm/entrypoints/serve/dev/rlhf @njhill @aoshen02
+/tests/distributed/test_weight_transfer.py @aoshen02
+/tests/entrypoints/weight_transfer @DarkLight1337 @robertgshaw2-redhat @aarnphm @NickLucche @AndreasKaratzas @aoshen02
+/tests/entrypoints/serve/dev/rlhf @DarkLight1337 @robertgshaw2-redhat @aarnphm @NickLucche @AndreasKaratzas @aoshen02
+/tests/v1/worker/test_gpu_worker_weight_transfer.py @aoshen02
+/examples/rl @aoshen02
+/examples/features/pause_resume @aoshen02
+/docs/training/rlhf.md @aoshen02
+/docs/training/async_rl.md @aoshen02
+/docs/training/weight_transfer @aoshen02
+
 # Security guide and policies
 /docs/usage/security.md @russellb @jperezdealgaba
 /SECURITY.md @russellb @jperezdealgaba

CMakeLists.txt

@@ -1122,7 +1122,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
     cuda_archs_loose_intersection(FUSED_KDA_DECODE_ARCHS
       "9.0a;10.0f;12.0f" "${CUDA_ARCHS}")
     cuda_archs_loose_intersection(FUSED_GDN_DECODE_ARCHS
-      "8.0;8.6;8.9;9.0a;10.0f;12.0f" "${CUDA_ARCHS}")
+      "8.0;8.6;8.9;9.0a;10.0f;11.0f;12.0f" "${CUDA_ARCHS}")
   endif()
   if(FUSED_KDA_DECODE_ARCHS)
     set(FUSED_KDA_DECODE_SRC

benchmarks/kernels/benchmark_hy_v4_ihc.py

@@ -0,0 +1,199 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Benchmark the HY V4 Triton iHC pre/post kernels against eager PyTorch."""
+
+import os
+import subprocess
+from functools import partial
+from importlib.metadata import version
+from statistics import median
+
+import torch
+import torch.nn.functional as F
+
+import vllm
+from vllm.models.hy_v4.nvidia.triton_ihc import (
+    triton_ihc_post,
+    triton_ihc_pre,
+)
+from vllm.triton_utils import triton
+from vllm.utils.argparse_utils import FlexibleArgumentParser
+
+
+def eager_pre(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    scale: torch.Tensor,
+    base: torch.Tensor,
+    magnitude: float,
+    hc_eps: float,
+    norm_eps: float,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    num_tokens, hc_mult, hidden_size = x.shape
+    x_flat = x.flatten(1).float()
+    reciprocal_rms = torch.rsqrt(x_flat.square().mean(-1, keepdim=True) + norm_eps)
+    mixes = F.linear(x_flat, weight) * reciprocal_rms
+    pre = torch.sigmoid(mixes[:, :hc_mult] * scale[0] + base[:hc_mult])
+    post = torch.sigmoid(mixes[:, hc_mult:] * scale[1] + base[hc_mult:])
+    pre = pre + hc_eps
+    post = magnitude * post + hc_eps
+    output = torch.sum(pre.unsqueeze(-1) * x.float(), dim=1)
+    return output.to(x.dtype).reshape(num_tokens, hidden_size), post
+
+
+def eager_post(
+    x: torch.Tensor,
+    residual: torch.Tensor,
+    post: torch.Tensor,
+) -> torch.Tensor:
+    return (post.float().unsqueeze(-1) * x.float().unsqueeze(-2) + residual.float()).to(
+        x.dtype
+    )
+
+
+def _timer(method: str):
+    if method == "cupti":
+        from flashinfer.testing import bench_gpu_time_with_cupti
+
+        return partial(
+            bench_gpu_time_with_cupti,
+            use_cuda_graph=True,
+            cold_l2_cache=True,
+        )
+    if method == "cudagraph":
+        from flashinfer.testing import bench_gpu_time_with_cudagraph
+
+        return partial(bench_gpu_time_with_cudagraph, cold_l2_cache=True)
+    raise ValueError(f"unknown timing method: {method}")
+
+
+@torch.inference_mode()
+def run_benchmark(
+    token_counts: list[int],
+    hidden_size: int,
+    dtype: torch.dtype,
+    method: str,
+) -> None:
+    hc_mult = 4
+    device = torch.device("cuda")
+    torch.manual_seed(0)
+    weight = torch.randn(
+        2 * hc_mult,
+        hc_mult * hidden_size,
+        device=device,
+        dtype=torch.float32,
+    )
+    scale = torch.randn(2, device=device, dtype=torch.float32) * 0.01
+    base = torch.randn(2 * hc_mult, device=device, dtype=torch.float32)
+    timer = _timer(method)
+
+    properties = torch.cuda.get_device_properties(device)
+    git_branch = subprocess.run(
+        ["git", "branch", "--show-current"],
+        check=False,
+        capture_output=True,
+        text=True,
+    ).stdout.strip()
+    git_commit = subprocess.run(
+        ["git", "rev-parse", "--short", "HEAD"],
+        check=False,
+        capture_output=True,
+        text=True,
+    ).stdout.strip()
+    print(f"device: {properties.name}")
+    print(f"branch: {git_branch}; commit: {git_commit}")
+    print(
+        f"vllm: {vllm.__version__}; torch: {torch.__version__}; "
+        f"CUDA: {torch.version.cuda}"
+    )
+    print(
+        f"triton: {triton.__version__}; flashinfer: {version('flashinfer-python')}; "
+        f"dtype: {dtype}; method: {method}; cache: cold L2"
+    )
+    print(
+        "env: VLLM_ENABLE_HPC_OPS="
+        f"{os.getenv('VLLM_ENABLE_HPC_OPS', '<unset>')}; "
+        "VLLM_BATCH_INVARIANT="
+        f"{os.getenv('VLLM_BATCH_INVARIANT', '<unset>')}"
+    )
+    print(f"hidden_size: {hidden_size}; hc_mult: {hc_mult}")
+    print(
+        f"{'tokens':>8} {'op':>6} {'eager (us)':>12} "
+        f"{'triton (us)':>12} {'speedup':>9} {'GiB':>8} "
+        f"{'eager GB/s':>12} {'triton GB/s':>13}"
+    )
+
+    for num_tokens in token_counts:
+        x = torch.randn(
+            num_tokens,
+            hc_mult,
+            hidden_size,
+            device=device,
+            dtype=dtype,
+        )
+        block_output = torch.randn(num_tokens, hidden_size, device=device, dtype=dtype)
+        residual = torch.randn_like(x)
+        eager_output, post = eager_pre(x, weight, scale, base, 2.0, 1e-6, 1e-5)
+        triton_output, triton_post = triton_ihc_pre(
+            x, weight, scale, base, 2.0, 1e-6, 1e-5
+        )
+        torch.testing.assert_close(triton_output, eager_output, atol=2e-2, rtol=1e-2)
+        torch.testing.assert_close(triton_post, post, atol=2e-5, rtol=1e-5)
+        torch.testing.assert_close(
+            triton_ihc_post(block_output, residual, post),
+            eager_post(block_output, residual, post),
+            atol=0,
+            rtol=0,
+        )
+
+        benchmarks = (
+            (
+                "pre",
+                partial(eager_pre, x, weight, scale, base, 2.0, 1e-6, 1e-5),
+                partial(triton_ihc_pre, x, weight, scale, base, 2.0, 1e-6, 1e-5),
+                x.nbytes
+                + weight.nbytes
+                + scale.nbytes
+                + base.nbytes
+                + eager_output.nbytes
+                + post.nbytes,
+            ),
+            (
+                "post",
+                partial(eager_post, block_output, residual, post),
+                partial(triton_ihc_post, block_output, residual, post),
+                block_output.nbytes + residual.nbytes + post.nbytes + residual.nbytes,
+            ),
+        )
+        for op_name, eager_fn, triton_fn, logical_bytes in benchmarks:
+            eager_us = median(timer(eager_fn)) * 1e3
+            triton_us = median(timer(triton_fn)) * 1e3
+            speedup = eager_us / triton_us
+            logical_gib = logical_bytes / 2**30
+            eager_gbps = logical_bytes / eager_us / 1e3
+            triton_gbps = logical_bytes / triton_us / 1e3
+            print(
+                f"{num_tokens:>8} {op_name:>6} {eager_us:>12.1f} "
+                f"{triton_us:>12.1f} {speedup:>8.2f}x {logical_gib:>8.3f} "
+                f"{eager_gbps:>12.1f} {triton_gbps:>13.1f}"
+            )
+
+
+if __name__ == "__main__":
+    parser = FlexibleArgumentParser(description=__doc__)
+    parser.add_argument(
+        "--token-counts",
+        type=int,
+        nargs="+",
+        default=[1, 2, 4, 8, 16, 64, 256, 1024, 4096, 8192],
+    )
+    parser.add_argument("--hidden-size", type=int, choices=[4096, 6144], default=6144)
+    parser.add_argument("--dtype", choices=["float16", "bfloat16"], default="bfloat16")
+    parser.add_argument("--method", choices=["cupti", "cudagraph"], default="cupti")
+    args = parser.parse_args()
+    run_benchmark(
+        args.token_counts,
+        args.hidden_size,
+        getattr(torch, args.dtype),
+        args.method,
+    )