[GG] fix(mla): tail-pad cuBLAS BMM inputs

tests/distributed/test_dcp_a2a.py

@@ -344,6 +344,7 @@ def fake_b12x(
         is_lse_base_on_e,
         max_batch_size,
         query_head_dim,
+        output_tail_padding_bytes,
     ):
         captured.update(
             output=cp_attn_out,
@@ -353,6 +354,7 @@ def fake_b12x(
             is_lse_base_on_e=is_lse_base_on_e,
             max_batch_size=max_batch_size,
             query_head_dim=query_head_dim,
+            output_tail_padding_bytes=output_tail_padding_bytes,
         )
         return expected
 
@@ -375,9 +377,61 @@ def fake_b12x(
         "is_lse_base_on_e": True,
         "max_batch_size": 8192,
         "query_head_dim": None,
+        "output_tail_padding_bytes": 0,
     }
 
 
+def test_cublas_tail_padding_preserves_values_and_storage_contract():
+    from vllm.utils.cublas import ensure_cublas_tail_padding, storage_tail_bytes
+
+    source = torch.arange(6 * 8 * 16, dtype=torch.bfloat16).view(6, 8, 16)
+    padded = ensure_cublas_tail_padding(source)
+
+    assert padded.shape == source.shape
+    assert padded.stride() == source.stride()
+    assert padded.data_ptr() != source.data_ptr()
+    assert storage_tail_bytes(padded) >= 64 * 1024
+    torch.testing.assert_close(padded, source)
+
+    reused = ensure_cublas_tail_padding(padded)
+    assert reused.data_ptr() == padded.data_ptr()
+
+
+def test_ag_rs_honors_tail_padding_bytes_for_mla_bmm(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    from vllm.utils.cublas import storage_tail_bytes
+    from vllm.v1.attention.ops import common
+
+    partial = torch.arange(2 * 4 * 8, dtype=torch.bfloat16).view(2, 4, 8)
+    lse = torch.zeros(2, 4, dtype=torch.float32)
+    reduced = partial[:, :2].clone()
+
+    class _Group:
+        world_size = 2
+
+        @staticmethod
+        def reduce_scatter(value, dim):
+            assert dim == 1
+            return reduced
+
+    monkeypatch.setattr(
+        common,
+        "_cp_lse_common",
+        lambda *args, **kwargs: (partial, lse),
+    )
+    actual = common.cp_lse_ag_out_rs(
+        partial,
+        lse,
+        _Group(),  # type: ignore[arg-type]
+        output_tail_padding_bytes=8 * 1024,
+    )
+
+    torch.testing.assert_close(actual, reduced)
+    assert actual.data_ptr() != reduced.data_ptr()
+    assert storage_tail_bytes(actual) >= 8 * 1024
+
+
 def test_packed_a2a_capture_buffers_stay_live_per_shape(
     monkeypatch: pytest.MonkeyPatch,
 ):
@@ -427,9 +481,7 @@ def fake_empty(*args, **kwargs):
     dcp_alltoall._DCP_A2A_GRAPH_BUFFERS.clear()
     monkeypatch.setattr(torch, "empty", fake_empty)
     monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False)
-    monkeypatch.setattr(
-        dcp_alltoall, "is_vllm_cudagraph_capture_active", lambda: True
-    )
+    monkeypatch.setattr(dcp_alltoall, "is_vllm_cudagraph_capture_active", lambda: True)
     device = torch.device("cuda:0")
 
     prewarm = dcp_alltoall._dcp_a2a_send_recv_buffers(
@@ -460,9 +512,7 @@ def fake_empty(*args, **kwargs):
     dcp_alltoall._DCP_A2A_GRAPH_BUFFERS.clear()
     monkeypatch.setattr(torch, "empty", fake_empty)
     monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False)
-    monkeypatch.setattr(
-        dcp_alltoall, "is_vllm_cudagraph_capture_active", lambda: False
-    )
+    monkeypatch.setattr(dcp_alltoall, "is_vllm_cudagraph_capture_active", lambda: False)
     device = torch.device("cuda:0")
 
     first = dcp_alltoall._dcp_a2a_send_recv_buffers(
@@ -850,16 +900,16 @@ def fake_get_pool(
 @pytest.mark.skipif(torch.accelerator.device_count() < 1, reason="CUDA is required.")
 def test_b12x_lse_reduce_makes_views_contiguous(monkeypatch: pytest.MonkeyPatch):
     """Head-sliced attention views must reach the PCIe pool contiguous."""
+    from vllm.utils.cublas import storage_tail_bytes
     from vllm.v1.attention.ops import dcp_alltoall
 
     monkeypatch.setenv("VLLM_USE_B12X_DCP_A2A", "1")
     received: dict[str, Any] = {}
-    sentinel = torch.zeros(1)
 
     class _FakePool:
-        def lse_reduce_scatter(self, out, lse, *, is_lse_base_on_e):
-            received.update(out=out, lse=lse)
-            return sentinel
+        def lse_reduce_scatter(self, partial_out, lse, out=None, *, is_lse_base_on_e):
+            received.update(partial_out=partial_out, lse=lse, output=out)
+            return out
 
     monkeypatch.setattr(
         dcp_alltoall,
@@ -884,10 +934,13 @@ def lse_reduce_scatter(self, out, lse, *, is_lse_base_on_e):
         is_lse_base_on_e=True,
         max_batch_size=8192,
         query_head_dim=64,
+        output_tail_padding_bytes=12 * 1024,
     )
-    assert result is sentinel
-    assert received["out"].is_contiguous()
+    assert result is received["output"]
+    assert received["partial_out"].is_contiguous()
     assert received["lse"].is_contiguous()
+    assert result is not None
+    assert storage_tail_bytes(result) >= 12 * 1024
 
 
 def test_b12x_query_gather_requires_env(monkeypatch: pytest.MonkeyPatch):
@@ -946,12 +999,15 @@ class TestPackedA2AKernels:
     @pytest.mark.parametrize("dtype_name", ["float16", "bfloat16", "float32"])
     @pytest.mark.parametrize("return_lse", [False, True])
     @pytest.mark.parametrize("is_lse_base_on_e", [False, True])
+    @pytest.mark.parametrize("output_tail_padding_bytes", [0, 8 * 1024])
     def test_pack_unpack_combine_matches_reference(
         self,
         dtype_name: str,
         return_lse: bool,
         is_lse_base_on_e: bool,
+        output_tail_padding_bytes: int,
     ):
+        from vllm.utils.cublas import storage_tail_bytes
         from vllm.v1.attention.ops.dcp_alltoall import (
             _dcp_a2a_lse_pack_dim,
             _dcp_a2a_pack_send,
@@ -982,7 +1038,12 @@ def test_pack_unpack_combine_matches_reference(
             lse_pack_dim,
         )
         actual = _dcp_a2a_unpack_combine(
-            send_buffer, D, lse_pack_dim, return_lse, is_lse_base_on_e
+            send_buffer,
+            D,
+            lse_pack_dim,
+            return_lse,
+            is_lse_base_on_e,
+            output_tail_padding_bytes,
         )
         expected_out, expected_lse = _packed_a2a_reference(
             cp_attn_out, cp_attn_lse, world_size, h_per_rank, is_lse_base_on_e
@@ -993,7 +1054,10 @@ def test_pack_unpack_combine_matches_reference(
             _assert_packed_a2a_close(actual_out, expected_out, dtype)
             torch.testing.assert_close(actual_lse, expected_lse, rtol=1e-4, atol=1e-4)
         else:
+            actual_out = actual
             _assert_packed_a2a_close(actual, expected_out, dtype)
+        if output_tail_padding_bytes:
+            assert storage_tail_bytes(actual_out) >= output_tail_padding_bytes
 
 
 def _distributed_packed_a2a_worker(env: dict[str, str]) -> None:

tests/model_executor/kernels/test_b12x_mxfp8_linear.py

@@ -304,9 +304,12 @@ def mxfp8_linear(
         bias: torch.Tensor | None = None,
         expected_m: int | None = None,
         stream: object = None,
+        tail_padding_bytes: int = 0,
     ) -> torch.Tensor:
         del stream
-        calls.append((source, packed_weight, bias, expected_m))
+        calls.append(
+            (source, packed_weight, bias, expected_m, tail_padding_bytes)
+        )
         return source.new_full((source.shape[0], packed_weight.out_features), 3.0)
 
     monkeypatch.setattr(
@@ -318,6 +321,7 @@ def mxfp8_linear(
     layer = torch.nn.Module()
     packed = types.SimpleNamespace(out_features=48)
     layer.b12x_mxfp8_packed_weight = packed
+    layer.output_tail_padding_bytes = 64 * 1024
     x = torch.empty((2, 3, 128), dtype=torch.bfloat16)
     bias = torch.empty((48,), dtype=torch.bfloat16)
     kernel = object.__new__(B12xMxfp8LinearKernel)
@@ -327,11 +331,12 @@ def mxfp8_linear(
     assert output.shape == (2, 3, 48)
     assert output.dtype == x.dtype
     assert len(calls) == 1
-    source, called_packed, called_bias, expected_m = calls[0]
+    source, called_packed, called_bias, expected_m, tail_padding_bytes = calls[0]
     assert source.shape == (6, 128)
     assert called_packed is packed
     assert called_bias is bias
     assert expected_m == 6
+    assert tail_padding_bytes == 64 * 1024
 
 
 def test_b12x_mxfp8_compile_path_uses_forward_context_custom_op(
@@ -384,9 +389,12 @@ def mxfp8_linear(
         bias: torch.Tensor | None = None,
         expected_m: int | None = None,
         stream: object = None,
+        tail_padding_bytes: int = 0,
     ) -> torch.Tensor:
         del stream
-        calls.append((source, packed_weight, bias, expected_m))
+        calls.append(
+            (source, packed_weight, bias, expected_m, tail_padding_bytes)
+        )
         return source.new_full((source.shape[0], packed_weight.out_features), 11.0)
 
     monkeypatch.setattr(
@@ -399,6 +407,7 @@ def mxfp8_linear(
     layer.prefix = "model.layers.2.mlp.down_proj"
     packed = types.SimpleNamespace(out_features=16)
     layer.b12x_mxfp8_packed_weight = packed
+    layer.output_tail_padding_bytes = 64 * 1024
     x = torch.empty((2, 3, 128), dtype=torch.bfloat16)
     bias = torch.empty((16,), dtype=torch.bfloat16)
     vllm_config = VllmConfig()
@@ -409,11 +418,12 @@ def mxfp8_linear(
 
     assert output.shape == (2, 3, 16)
     assert len(calls) == 1
-    source, called_packed, called_bias, expected_m = calls[0]
+    source, called_packed, called_bias, expected_m, tail_padding_bytes = calls[0]
     assert source.shape == (6, 128)
     assert called_packed is packed
     assert called_bias is bias
     assert expected_m == 6
+    assert tail_padding_bytes == 64 * 1024
     torch.testing.assert_close(output, torch.full_like(output, 11.0))
 
 

tests/model_executor/layers/test_unquantized_linear_tail_padding.py

@@ -0,0 +1,47 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+
+from vllm.model_executor.layers.linear import UnquantizedLinearMethod
+from vllm.utils.cublas import storage_tail_bytes
+
+
+def test_unquantized_linear_writes_directly_to_tail_padded_storage() -> None:
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    layer = torch.nn.Module().to(device)
+    layer.weight = torch.nn.Parameter(torch.randn(7, 5, device=device))
+    layer.output_tail_padding_bytes = 64 * 1024
+    source = torch.randn(3, 5, device=device)
+
+    with torch.inference_mode():
+        output = UnquantizedLinearMethod().apply(layer, source)
+
+    torch.testing.assert_close(output, torch.nn.functional.linear(source, layer.weight))
+    assert storage_tail_bytes(output) >= layer.output_tail_padding_bytes
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
+def test_unquantized_linear_tail_padding_survives_cuda_graph_replay() -> None:
+    layer = torch.nn.Module().cuda()
+    layer.weight = torch.nn.Parameter(
+        torch.randn(2048, 2048, device="cuda", dtype=torch.bfloat16),
+        requires_grad=False,
+    )
+    layer.output_tail_padding_bytes = 64 * 1024
+    source = torch.randn(6, 2048, device="cuda", dtype=torch.bfloat16)
+
+    with torch.inference_mode():
+        expected = UnquantizedLinearMethod().apply(layer, source).clone()
+        UnquantizedLinearMethod().apply(layer, source)
+        torch.cuda.synchronize()
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            output = UnquantizedLinearMethod().apply(layer, source)
+        for _ in range(3):
+            graph.replay()
+        torch.cuda.synchronize()
+
+    assert storage_tail_bytes(output) >= layer.output_tail_padding_bytes
+    torch.testing.assert_close(output, expected, rtol=0, atol=0)

vllm/model_executor/kernels/linear/mxfp8/b12x.py

@@ -137,6 +137,7 @@ def _apply_b12x_mxfp8_packed_linear(
         bias=bias,
         expected_m=_b12x_mxfp8_expected_m(int(input_2d.shape[0])),
         stream=current_stream().cuda_stream,
+        tail_padding_bytes=int(getattr(layer, "output_tail_padding_bytes", 0)),
     )
     return output.view(*output_shape)
 

vllm/model_executor/layers/attention/mla_attention.py

@@ -243,6 +243,7 @@
 )
 from vllm.model_executor.utils import replace_parameter
 from vllm.platforms import current_platform
+from vllm.utils.cublas import ensure_cublas_tail_padding
 from vllm.utils.flashinfer import has_flashinfer
 from vllm.utils.math_utils import cdiv, round_down
 from vllm.utils.torch_utils import (
@@ -604,6 +605,11 @@ def __init__(
             **extra_impl_args,
         )
         self.q_pad_num_heads = getattr(self.impl, "q_pad_num_heads", None)
+        self.mla_bmm_tail_padding_bytes = int(
+            getattr(self.impl, "mla_bmm_tail_padding_bytes", 0)
+        )
+        if self.mla_bmm_tail_padding_bytes < 0:
+            raise ValueError("mla_bmm_tail_padding_bytes must be non-negative")
         self.use_direct_call = not current_platform.opaque_attention_op()
 
         vllm_config = get_current_vllm_config()
@@ -966,6 +972,7 @@ def forward_impl(
             self.impl.dcp_world_size = get_dcp_group().world_size
 
         fp8_attention = is_quantized_kv_cache(self.kv_cache_dtype)
+        bmm_tail_padding_bytes = getattr(self, "mla_bmm_tail_padding_bytes", 0)
 
         num_actual_toks = attn_metadata.num_actual_tokens
 
@@ -1089,6 +1096,12 @@ def forward_impl(
                 N, B, P = mqa_q_nope.shape
                 _, _, L = self.W_UK_T.shape
 
+                if bmm_tail_padding_bytes:
+                    mqa_q_nope = ensure_cublas_tail_padding(
+                        mqa_q_nope,
+                        bmm_tail_padding_bytes,
+                    )
+
                 if self.q_pad_num_heads is not None:
                     mqa_ql_nope = mqa_q_nope.new_empty((self.q_pad_num_heads, B, L))
                     mqa_ql_nope.resize_((N, B, L))
@@ -1300,6 +1313,9 @@ def forward_impl(
                         use_b12x=dcp_use_b12x,
                         b12x_max_batch_size=self.dcp_max_batch_size,
                         b12x_query_head_dim=(self.kv_lora_rank + self.qk_rope_head_dim),
+                        output_tail_padding_bytes=(
+                            bmm_tail_padding_bytes if not project_before_merge else 0
+                        ),
                     )
                 else:
                     if project_before_merge:
@@ -1332,6 +1348,11 @@ def forward_impl(
                             lse,
                             get_dcp_group(),
                             is_lse_base_on_e=self.impl.lse_base_on_e,
+                            output_tail_padding_bytes=(
+                                bmm_tail_padding_bytes
+                                if not project_before_merge
+                                else 0
+                            ),
                         )
 
             if project_before_merge:
@@ -1599,6 +1620,12 @@ def _v_up_proj(self, x: torch.Tensor, out: torch.Tensor):
             )
         else:
             # Multiply + Transpose (N, B, L) x (N, L, V)->(N, B, V)->(B, N, V)
+            bmm_tail_padding_bytes = getattr(self, "mla_bmm_tail_padding_bytes", 0)
+            if bmm_tail_padding_bytes:
+                x = ensure_cublas_tail_padding(
+                    x,
+                    bmm_tail_padding_bytes,
+                )
             torch.bmm(x, self.W_UV, out=out.transpose(0, 1))
 
     def _v_up_proj_bmm(

vllm/model_executor/layers/linear.py

@@ -227,6 +227,13 @@ def apply(
         x: torch.Tensor,
         bias: torch.Tensor | None = None,
     ) -> torch.Tensor:
+        tail_padding_bytes = int(getattr(layer, "output_tail_padding_bytes", 0))
+        if tail_padding_bytes > 0 and bias is None:
+            return torch.ops.vllm.tail_padded_unquantized_gemm(
+                x,
+                layer.weight,
+                tail_padding_bytes,
+            )
         if envs.VLLM_BATCH_INVARIANT and current_platform.is_cuda_alike():
             return linear_batch_invariant(x, layer.weight, bias)
         return dispatch_unquantized_gemm()(layer, x, layer.weight, bias)