[GG] fix MLA query BMM cuBLAS read-ahead without query copies

CMakeLists.txt

@@ -406,6 +406,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
     "csrc/libtorch_stable/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu"
     "csrc/libtorch_stable/quantization/fused_kernels/fused_silu_mul_block_quant.cu"
     "csrc/libtorch_stable/attention/merge_attn_states.cu"
+    "csrc/libtorch_stable/attention/mla/safe_query_bmm.cu"
     "csrc/libtorch_stable/sampler.cu"
     "csrc/libtorch_stable/topk.cu"
     "csrc/libtorch_stable/mamba/selective_scan_fwd.cu"
@@ -1106,6 +1107,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)

csrc/libtorch_stable/attention/mla/safe_query_bmm.cu

@@ -0,0 +1,102 @@
+#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,
+          CUBLAS_COMPUTE_32F_PEDANTIC, CUBLAS_GEMM_DEFAULT),
+      "cublasGemmStridedBatchedEx");
+}
+
+STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
+  m.impl("safe_mla_query_bmm", TORCH_BOX(&safe_mla_query_bmm));
+}

csrc/libtorch_stable/ops.h

@@ -198,6 +198,12 @@ void merge_attn_states(
     const std::optional<int64_t> prefill_tokens_with_context,
     const std::optional<torch::stable::Tensor>& output_scale = std::nullopt);
 
+// BF16 MLA query absorption BMM that avoids cuBLAS read-ahead on tight
+// DCP/custom-allocation inputs without materializing the query operand.
+void safe_mla_query_bmm(torch::stable::Tensor const& query,
+                        torch::stable::Tensor const& weight,
+                        torch::stable::Tensor& output);
+
 torch::stable::Tensor hadacore_transform(torch::stable::Tensor& x,
                                          bool inplace);
 

csrc/libtorch_stable/torch_bindings.cpp

@@ -326,6 +326,11 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
   ops.def(
       "dsv3_fused_a_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()");
 
+  // BF16 MLA query absorption BMM. The implementation selects a cuBLAS
+  // compute mode that is safe for tight DCP/custom-allocation inputs.
+  ops.def(
+      "safe_mla_query_bmm(Tensor query, Tensor weight, Tensor! output) -> ()");
+
   // BF16/FP32 x FP32 -> FP32 router GEMM for H=3072, E=256, M<=32 (SM90+).
   // conditionally compiled so impl registration is in source file
   ops.def("fp32_router_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()");

tests/kernels/test_safe_mla_query_bmm.py

@@ -0,0 +1,68 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+
+
+def _require_safe_mla_query_bmm():
+    if not hasattr(torch.ops._C, "safe_mla_query_bmm"):
+        pytest.skip("safe_mla_query_bmm is not built")
+
+
+@pytest.mark.parametrize("heads,tokens", [(8, 1), (8, 6), (8, 11), (11, 6), (16, 6)])
+def test_safe_mla_query_bmm_matches_torch_bmm(heads: int, tokens: int):
+    _require_safe_mla_query_bmm()
+    device = torch.device("cuda")
+    q_dim = 512
+    rope_dim = 64
+    latent_dim = 512
+    torch.manual_seed(0)
+
+    query_storage = torch.randn(
+        tokens, heads, q_dim + rope_dim, dtype=torch.bfloat16, device=device
+    )
+    query = query_storage[..., :q_dim].transpose(0, 1)
+    weight = torch.randn(heads, q_dim, latent_dim, dtype=torch.bfloat16, device=device)
+    output = torch.empty(heads, tokens, latent_dim, dtype=torch.bfloat16, device=device)
+
+    assert not query.is_contiguous()
+    torch.ops._C.safe_mla_query_bmm(query, weight, output)
+    expected = torch.bmm(query.contiguous(), weight)
+
+    torch.testing.assert_close(output.float(), expected.float(), rtol=5e-2, atol=5e-2)
+
+
+def test_safe_mla_query_bmm_cuda_graph_replay():
+    _require_safe_mla_query_bmm()
+    device = torch.device("cuda")
+    heads = 8
+    tokens = 6
+    q_dim = 512
+    latent_dim = 512
+    torch.manual_seed(1)
+
+    query_storage = torch.randn(
+        tokens, heads, q_dim + 64, dtype=torch.bfloat16, device=device
+    )
+    query = query_storage[..., :q_dim].transpose(0, 1)
+    weight = torch.randn(heads, q_dim, latent_dim, dtype=torch.bfloat16, device=device)
+    output = torch.empty(heads, tokens, latent_dim, dtype=torch.bfloat16, device=device)
+
+    stream = torch.cuda.Stream()
+    stream.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(stream):
+        torch.ops._C.safe_mla_query_bmm(query, weight, output)
+    torch.cuda.current_stream().wait_stream(stream)
+
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        torch.ops._C.safe_mla_query_bmm(query, weight, output)
+
+    graph.replay()
+    graph.replay()
+    expected = torch.bmm(query.contiguous(), weight)
+
+    torch.testing.assert_close(output.float(), expected.float(), rtol=5e-2, atol=5e-2)

tests/v1/attention/test_mla_backends.py

@@ -163,6 +163,75 @@ def test_mla_kv_cache_spec_uses_layer_cache_dtype(
         assert spec.page_size_bytes == 64 * 656
 
 
+@pytest.mark.cpu_test
+def test_mla_init_propagates_safe_query_bmm_contract(monkeypatch):
+    class FakeImpl:
+        is_sparse = True
+        supports_mha_prefill = False
+
+        def __init__(self, **kwargs):
+            self.use_safe_mla_query_bmm = True
+
+    class FakeBackend:
+        @staticmethod
+        def is_mla():
+            return True
+
+        @staticmethod
+        def get_name():
+            return "TEST_MLA"
+
+        @staticmethod
+        def get_impl_cls():
+            return FakeImpl
+
+    config = SimpleNamespace(
+        compilation_config=SimpleNamespace(
+            static_forward_context={}, cudagraph_capture_sizes=[]
+        ),
+        attention_config=SimpleNamespace(mla_prefill_backend=None),
+        parallel_config=SimpleNamespace(
+            decode_context_parallel_size=1,
+            dcp_comm_backend="a2a",
+        ),
+        scheduler_config=SimpleNamespace(max_num_batched_tokens=128),
+    )
+    monkeypatch.setattr(mla_attention_module, "get_current_vllm_config", lambda: config)
+    monkeypatch.setattr(
+        mla_attention_module, "get_current_vllm_config_or_none", lambda: config
+    )
+    monkeypatch.setattr(mla_attention_module, "_init_kv_cache_quant", lambda *a: None)
+    monkeypatch.setattr(
+        mla_attention_module,
+        "get_mla_prefill_backend",
+        lambda _: (_ for _ in ()).throw(ValueError),
+    )
+    monkeypatch.setattr(mla_attention_module, "_DecodeConcatQuantFP8", lambda **_: None)
+    monkeypatch.setattr(mla_attention_module, "QuantFP8", lambda **_: None)
+    monkeypatch.setattr(
+        mla_attention_module.rocm_aiter_ops, "is_fp8bmm_enabled", lambda: False
+    )
+    monkeypatch.setattr(
+        mla_attention_module.rocm_aiter_ops, "is_fp4bmm_enabled", lambda: False
+    )
+
+    layer = MLAAttention(
+        num_heads=8,
+        scale=1.0,
+        qk_nope_head_dim=4,
+        qk_rope_head_dim=2,
+        v_head_dim=3,
+        q_lora_rank=None,
+        kv_lora_rank=4,
+        kv_b_proj=SimpleNamespace(),
+        prefix="test",
+        attn_backend=FakeBackend,
+        use_sparse=True,
+    )
+
+    assert layer.use_safe_mla_query_bmm
+
+
 # Remove sm100 backends from the list if not using sm100
 if not torch.cuda.is_available() or torch.cuda.get_device_properties(0).major < 10:
     BACKENDS_TO_TEST.remove(AttentionBackendEnum.CUTLASS_MLA)
@@ -778,6 +847,80 @@ def guarded_empty(*args, **kwargs):
     assert torch.equal(output, torch.zeros_like(output))
 
 
+@pytest.mark.cpu_test
+def test_mla_query_absorb_safe_bmm_fallback_materializes_input(monkeypatch):
+    class FakeSparseImpl(SparseMLACommonImpl):
+        supports_quant_query_input = False
+
+        def __init__(self):
+            self.dcp_world_size = 1
+
+        def forward_mqa(self, q, kv_cache, attn_metadata, layer):
+            assert isinstance(q, tuple)
+            return torch.ones(
+                (q[0].shape[0], q[0].shape[1], layer.kv_lora_rank),
+                dtype=q[0].dtype,
+            ), None
+
+    layer = object.__new__(MLAAttention)
+    layer.impl = FakeSparseImpl()
+    layer.kv_cache_dtype = "auto"
+    layer.num_heads = 2
+    layer.qk_nope_head_dim = 4
+    layer.qk_rope_head_dim = 2
+    layer.kv_lora_rank = 3
+    layer.v_head_dim = 3
+    layer.q_pad_num_heads = None
+    layer.use_safe_mla_query_bmm = True
+    layer.is_aiter_triton_fp4_bmm_enabled = False
+    layer.is_aiter_triton_fp8_bmm_enabled = False
+    layer.W_UK_T = torch.ones((2, 4, 3), dtype=torch.bfloat16)
+
+    def fake_v_up_proj(x, out):
+        out.copy_(x.reshape(out.shape))
+
+    layer._v_up_proj = fake_v_up_proj
+
+    real_bmm = torch.bmm
+    seen_input_layouts = []
+
+    def checked_bmm(input_tensor, mat2, *, out=None):
+        if mat2 is layer.W_UK_T:
+            seen_input_layouts.append(input_tensor.is_contiguous())
+        return real_bmm(input_tensor, mat2, out=out)
+
+    monkeypatch.setattr(torch, "bmm", checked_bmm)
+
+    num_tokens = 3
+    q = torch.ones((num_tokens, 2, 6), dtype=torch.bfloat16)
+    q_nope, _ = q.split((4, 2), dim=-1)
+    assert not q_nope.transpose(0, 1).is_contiguous()
+    kv_c = torch.empty((num_tokens, 3), dtype=torch.bfloat16)
+    k_pe = torch.empty((num_tokens, 1, 2), dtype=torch.bfloat16)
+    kv_cache = torch.empty((0,), dtype=torch.uint8)
+    output = torch.empty((num_tokens, 6), dtype=torch.bfloat16)
+    attn_metadata = SimpleNamespace(
+        num_actual_tokens=num_tokens,
+        num_decodes=num_tokens,
+        num_prefills=0,
+        num_decode_tokens=num_tokens,
+        max_query_len=1,
+    )
+
+    result = MLAAttention.forward_impl(
+        layer,
+        q,
+        kv_c,
+        k_pe,
+        kv_cache,
+        attn_metadata,
+        output,
+    )
+
+    assert result is output
+    assert seen_input_layouts == [True]
+
+
 def test_fp8_dcp_sparse_mla_uses_lse_gather_path(monkeypatch):
     from vllm.model_executor.layers.attention import mla_attention
 
@@ -837,6 +980,7 @@ def fake_lse_reduce(attn_out, lse, group, is_lse_base_on_e=True):
     layer.v_head_dim = 3
     layer.q_pad_num_heads = None
     layer.force_contiguous_mla_bmm_input = False
+    layer.use_safe_mla_query_bmm = False
     layer.force_contiguous_mla_bmm_output = False
     layer.is_aiter_triton_fp4_bmm_enabled = False
     layer.is_aiter_triton_fp8_bmm_enabled = False
@@ -910,6 +1054,7 @@ def forward_mqa(self, *args, **kwargs):
     layer.v_head_dim = 3
     layer.q_pad_num_heads = None
     layer.force_contiguous_mla_bmm_input = False
+    layer.use_safe_mla_query_bmm = False
     layer.is_aiter_triton_fp4_bmm_enabled = False
     layer.is_aiter_triton_fp8_bmm_enabled = False
     layer.W_UK_T = torch.ones((2, 4, 3), dtype=torch.bfloat16)