feat: EXL3 LoRA cartridge support for MSRT additive quantization

.gitignore

@@ -255,5 +255,16 @@ vllm/grpc/vllm_engine_pb2.py
 vllm/grpc/vllm_engine_pb2_grpc.py
 vllm/grpc/vllm_engine_pb2.pyi
 
+# DGX Spark local build products
+/.spark-artifacts/
+/.spark-build/
+
 # Ignore generated cpu headers 
 csrc/cpu/cpu_attn_dispatch_generated.h
+
+# Raw routing-stat dumps: tens of MB each, published to HF instead
+# (see results/k3-fq/CONVERGENCE-RESULTS.md for sha256 + location)
+research/fungible-quant/runs/m5-serve/results/*/stats*.jsonl
+
+# PoC virtualenv (local only, not for commit)
+.poc-venv/

CMakeLists.txt

@@ -433,6 +433,9 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
   endif()
 
   if(VLLM_GPU_LANG STREQUAL "CUDA")
+    list(APPEND VLLM_STABLE_EXT_SRC
+      "csrc/libtorch_stable/attention/mla/safe_query_bmm.cu")
+
     SET(CUTLASS_ENABLE_HEADERS_ONLY ON CACHE BOOL "Enable only the header library")
 
     # Set CUTLASS_REVISION. Used for FetchContent. Also fixes some bogus messages when building.
@@ -1106,6 +1109,15 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
   # Needed to use cuda/hip APIs from C-shim
   if(VLLM_GPU_LANG STREQUAL "CUDA")
     target_compile_definitions(_C_stable_libtorch PRIVATE USE_CUDA)
+    find_library(VLLM_CUBLAS_LIBRARY cublas
+      PATHS "/usr/local/cuda/lib64" "/usr/local/cuda/lib" NO_DEFAULT_PATH)
+    if(NOT VLLM_CUBLAS_LIBRARY)
+      find_library(VLLM_CUBLAS_LIBRARY cublas)
+    endif()
+    if(NOT VLLM_CUBLAS_LIBRARY)
+      message(FATAL_ERROR "safe_mla_query_bmm requires libcublas")
+    endif()
+    target_link_libraries(_C_stable_libtorch PRIVATE ${VLLM_CUBLAS_LIBRARY})
     if(COOPERATIVE_TOPK_ARCHS)
       target_compile_definitions(_C_stable_libtorch PRIVATE
         VLLM_ENABLE_COOPERATIVE_TOPK=1)

benchmarks/profile_dspark_sps_curve.py

@@ -0,0 +1,165 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Profile the engine step-rate curve for the DSpark prefix scheduler.
+
+Times the captured FULL cudagraph replays of the target verification step at
+every captured batch token count and emits ``dspark_sps_curve`` breakpoints,
+one per capture size. The scheduler linearly interpolates between
+breakpoints, which amortizes cudagraph padding smoothly instead of
+concentrating it into thresholds at capture-size boundaries.
+
+Example:
+    python benchmarks/profile_dspark_sps_curve.py <target-model> \\
+        --speculative-config '{"method": "dspark", "model": "...", ...}' \\
+        --engine-args '{"tensor_parallel_size": 4, "max_num_seqs": 32}'
+
+Paste the printed ``dspark_sps_curve`` entry into --speculative-config.
+
+Caveats: replays run on whatever (dummy) buffer state capture left behind, so
+data-dependent kernels (e.g. MoE routing) may be timed on unrepresentative
+inputs, and per-step CPU/draft overhead is modeled only through the constant
+``--overhead-ms``. Only the curve's shape matters to the scheduler.
+"""
+
+import argparse
+import json
+
+
+def _time_fullgraph_replays(worker, iters: int, warmup: int) -> dict[int, float]:
+    """Worker-side: time FULL graph replay per batch token count (ms/step).
+
+    Runs on every TP rank via collective_rpc so the collectives captured in
+    the graphs stay matched; every rank replays the same descs in the same
+    sorted order. Before timing each descriptor the input buffers are
+    refreshed into the same coherent dummy state capture used, so replays
+    never read stale metadata.
+    """
+    import torch
+
+    from vllm.v1.worker.gpu.cudagraph_utils import prepare_inputs_to_capture
+
+    runner = worker.model_runner
+    mgr = runner.cudagraph_manager
+    assert mgr is not None and mgr.graphs, (
+        "No FULL cudagraphs captured; run with a cudagraph_mode that captures "
+        "FULL decode graphs."
+    )
+    # Prefer varlen spec-decode descs; fall back to all captured graphs.
+    descs = [d for d in mgr.graphs if d.max_req_tokens is not None]
+    if not descs:
+        descs = list(mgr.graphs.keys())
+    # One desc per token count: the largest request count is the most
+    # representative shape under load.
+    by_tokens: dict[int, object] = {}
+    for d in descs:
+        cur = by_tokens.get(d.num_tokens)
+        if cur is None or (d.num_reqs or 0) > (cur.num_reqs or 0):
+            by_tokens[d.num_tokens] = d
+
+    results: dict[int, float] = {}
+    for num_tokens in sorted(by_tokens):
+        desc = by_tokens[num_tokens]
+        num_reqs = desc.num_reqs or min(num_tokens, mgr.max_num_reqs)
+        prepare_inputs_to_capture(
+            num_reqs,
+            num_tokens,
+            runner.model_state,
+            runner.input_buffers,
+            runner.block_tables,
+            runner.attn_groups,
+            runner.kv_cache_config,
+            max_req_tokens=desc.max_req_tokens,
+        )
+        graph = mgr.graphs[desc]
+        for _ in range(warmup):
+            graph.replay()
+        torch.accelerator.synchronize()
+        start = torch.Event(enable_timing=True)
+        end = torch.Event(enable_timing=True)
+        start.record()
+        for _ in range(iters):
+            graph.replay()
+        end.record()
+        torch.accelerator.synchronize()
+        results[num_tokens] = start.elapsed_time(end) / iters
+    return results
+
+
+def curve_breakpoints(
+    ms_per_step: dict[int, float], overhead_ms: float
+) -> list[list[float]]:
+    """Convert per-capture-size step times into ``dspark_sps_curve``
+    breakpoints, one per capture size. The scheduler's table linearly
+    interpolates between them (and clamps at the ends)."""
+    return [
+        [size, round(1000.0 / (ms_per_step[size] + overhead_ms), 3)]
+        for size in sorted(ms_per_step)
+    ]
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("model", help="Target model (path or HF id)")
+    parser.add_argument(
+        "--speculative-config",
+        required=True,
+        help="JSON speculative config (same value you pass to vllm serve)",
+    )
+    parser.add_argument(
+        "--engine-args",
+        default="{}",
+        help="JSON dict of extra vllm.LLM kwargs "
+        '(e.g. \'{"tensor_parallel_size": 4, "max_num_seqs": 32}\')',
+    )
+    parser.add_argument("--iters", type=int, default=50)
+    parser.add_argument("--warmup", type=int, default=5)
+    parser.add_argument(
+        "--overhead-ms",
+        type=float,
+        default=0.0,
+        help="Constant per-step overhead (draft forward, sampling, CPU gap) "
+        "added to every measured step time before converting to a rate.",
+    )
+    parser.add_argument("--output", help="Write the curve JSON to this file")
+    args = parser.parse_args()
+    if args.iters <= 0:
+        parser.error("--iters must be greater than zero")
+    if args.warmup < 0:
+        parser.error("--warmup must be non-negative")
+    if args.overhead_ms < 0:
+        parser.error("--overhead-ms must be non-negative")
+
+    # The timing callable is shipped to the workers via collective_rpc, which
+    # requires the pickle fallback. Local profiling tool, trusted input.
+    import os
+
+    os.environ.setdefault("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
+
+    from vllm import LLM
+
+    llm = LLM(
+        model=args.model,
+        speculative_config=json.loads(args.speculative_config),
+        **json.loads(args.engine_args),
+    )
+    per_rank = llm.collective_rpc(
+        _time_fullgraph_replays, kwargs={"iters": args.iters, "warmup": args.warmup}
+    )
+    ms_per_step = per_rank[0]
+
+    print("\nMeasured FULL-graph step times (rank 0):")
+    for size in sorted(ms_per_step):
+        print(f"  B={size:5d} tokens: {ms_per_step[size]:8.3f} ms/step")
+
+    curve = curve_breakpoints(ms_per_step, args.overhead_ms)
+    entry = {"dspark_sps_curve": curve}
+    print("\nAdd to --speculative-config:")
+    print(json.dumps(entry))
+    if args.output:
+        with open(args.output, "w") as f:
+            json.dump(entry, f, indent=2)
+        print(f"\nWritten to {args.output}")
+
+
+if __name__ == "__main__":
+    main()

cmake/external_projects/vllm_flash_attn.cmake

@@ -36,10 +36,15 @@ if(VLLM_FLASH_ATTN_SRC_DIR)
           BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn
   )
 else()
+  set(VLLM_FLASH_ATTN_GIT_SUBMODULES "")
+  if(VLLM_GPU_LANG STREQUAL "CUDA")
+    set(VLLM_FLASH_ATTN_GIT_SUBMODULES "csrc/cutlass")
+  endif()
   FetchContent_Declare(
           vllm-flash-attn
           GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
           GIT_TAG caaa4eb59845388a20b1f435ecaafb4bd9517ad8
+          GIT_SUBMODULES ${VLLM_FLASH_ATTN_GIT_SUBMODULES}
           GIT_PROGRESS TRUE
           # Don't share the vllm-flash-attn build between build types
           BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn

csrc/libtorch_stable/attention/mla/safe_query_bmm.cu

@@ -0,0 +1,106 @@
+#include <torch/csrc/stable/library.h>
+#include <torch/csrc/stable/tensor.h>
+#include <torch/headeronly/core/ScalarType.h>
+
+#include "core/registration.h"
+#include "libtorch_stable/torch_utils.h"
+
+#include <cublas_v2.h>
+#include <cuda_bf16.h>
+
+#include <limits>
+#include <sstream>
+
+namespace {
+
+void check_cublas(cublasStatus_t status, const char* operation) {
+  if (status == CUBLAS_STATUS_SUCCESS) {
+    return;
+  }
+  std::ostringstream error;
+  error << operation << " failed with cuBLAS status "
+        << static_cast<int>(status);
+  STD_TORCH_CHECK(false, error.str());
+}
+
+int checked_int(int64_t value, const char* name) {
+  STD_TORCH_CHECK(value > 0 && value <= std::numeric_limits<int>::max(), name,
+                  " is outside cuBLAS int range: ", value);
+  return static_cast<int>(value);
+}
+
+void check_bf16_cuda_3d(torch::stable::Tensor const& tensor,
+                        const char* name) {
+  STD_TORCH_CHECK(tensor.device().is_cuda(), name, " must be on CUDA");
+  STD_TORCH_CHECK(tensor.dim() == 3, name, " must be a 3D tensor");
+  STD_TORCH_CHECK(
+      tensor.scalar_type() == torch::headeronly::ScalarType::BFloat16, name,
+      " must be BF16");
+}
+
+}  // namespace
+
+void safe_mla_query_bmm(torch::stable::Tensor const& query,
+                        torch::stable::Tensor const& weight,
+                        torch::stable::Tensor& output) {
+  check_bf16_cuda_3d(query, "query");
+  check_bf16_cuda_3d(weight, "weight");
+  check_bf16_cuda_3d(output, "output");
+  STD_TORCH_CHECK(query.get_device_index() == weight.get_device_index() &&
+                      query.get_device_index() == output.get_device_index(),
+                  "query, weight, and output must be on the same CUDA device");
+
+  const int64_t heads_i64 = query.size(0);
+  const int64_t tokens_i64 = query.size(1);
+  const int64_t q_dim_i64 = query.size(2);
+  const int64_t latent_dim_i64 = weight.size(2);
+
+  STD_TORCH_CHECK(weight.size(0) == heads_i64 && weight.size(1) == q_dim_i64,
+                  "weight must have shape [heads, q_dim, latent_dim]");
+  STD_TORCH_CHECK(output.size(0) == heads_i64 &&
+                      output.size(1) == tokens_i64 &&
+                      output.size(2) == latent_dim_i64,
+                  "output must have shape [heads, tokens, latent_dim]");
+  STD_TORCH_CHECK(query.stride(2) == 1,
+                  "query q_dim must be contiguous for safe_mla_query_bmm");
+  STD_TORCH_CHECK(weight.stride(2) == 1,
+                  "weight latent_dim must be contiguous for safe_mla_query_bmm");
+  STD_TORCH_CHECK(output.stride(2) == 1,
+                  "output latent_dim must be contiguous for safe_mla_query_bmm");
+
+  const int heads = checked_int(heads_i64, "heads");
+  const int tokens = checked_int(tokens_i64, "tokens");
+  const int q_dim = checked_int(q_dim_i64, "q_dim");
+  const int latent_dim = checked_int(latent_dim_i64, "latent_dim");
+  const int query_ld = checked_int(query.stride(1), "query.stride(1)");
+  const int weight_ld = checked_int(weight.stride(1), "weight.stride(1)");
+  const int output_ld = checked_int(output.stride(1), "output.stride(1)");
+
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      query.get_device_index());
+  cublasHandle_t handle = get_current_cuda_blas_handle();
+  check_cublas(cublasSetStream(handle, get_current_cuda_stream(
+                                           query.get_device_index())),
+               "cublasSetStream");
+
+  const float alpha = 1.0f;
+  const float beta = 0.0f;
+  check_cublas(
+      cublasGemmStridedBatchedEx(
+          handle, CUBLAS_OP_N, CUBLAS_OP_N, latent_dim, tokens, q_dim, &alpha,
+          weight.const_data_ptr(), CUDA_R_16BF, weight_ld,
+          static_cast<long long>(weight.stride(0)), query.const_data_ptr(),
+          CUDA_R_16BF, query_ld, static_cast<long long>(query.stride(0)),
+          &beta, output.mutable_data_ptr(), CUDA_R_16BF, output_ld,
+          static_cast<long long>(output.stride(0)), heads,
+          // The explicit operand order and leading dimensions provide the
+          // tight-query contract. Regular FP32 accumulation keeps tensor-core
+          // kernels eligible; PEDANTIC forces a much slower fallback for
+          // production prefill shapes without improving that contract.
+          CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT),
+      "cublasGemmStridedBatchedEx");
+}
+
+STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
+  m.impl("safe_mla_query_bmm", TORCH_BOX(&safe_mla_query_bmm));
+}