diff --git a/tests/distributed/test_dcp_a2a.py b/tests/distributed/test_dcp_a2a.py index 84c4d075444b..0f071341b501 100644 --- a/tests/distributed/test_dcp_a2a.py +++ b/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: diff --git a/tests/model_executor/kernels/test_b12x_mxfp8_linear.py b/tests/model_executor/kernels/test_b12x_mxfp8_linear.py index 3f91a980e197..1b5251141aec 100644 --- a/tests/model_executor/kernels/test_b12x_mxfp8_linear.py +++ b/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)) diff --git a/tests/model_executor/layers/test_unquantized_linear_tail_padding.py b/tests/model_executor/layers/test_unquantized_linear_tail_padding.py new file mode 100644 index 000000000000..4c3d90afb73e --- /dev/null +++ b/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) diff --git a/vllm/model_executor/kernels/linear/mxfp8/b12x.py b/vllm/model_executor/kernels/linear/mxfp8/b12x.py index 3a8430082d34..be47cb32a9f0 100644 --- a/vllm/model_executor/kernels/linear/mxfp8/b12x.py +++ b/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) diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index d87241456b91..c2e49cfbbda4 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/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( diff --git a/vllm/model_executor/layers/linear.py b/vllm/model_executor/layers/linear.py index f1ec65603a71..ae83c75b02ae 100644 --- a/vllm/model_executor/layers/linear.py +++ b/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) diff --git a/vllm/model_executor/layers/mla.py b/vllm/model_executor/layers/mla.py index 1d008699d086..38a04d0d57ab 100644 --- a/vllm/model_executor/layers/mla.py +++ b/vllm/model_executor/layers/mla.py @@ -217,6 +217,11 @@ def __init__( indexer=self.indexer, topk_indices_buffer=mla_modules.topk_indices_buffer, ) + q_projection = self.q_b_proj if self.q_b_proj is not None else self.q_proj + if q_projection is not None and self.mla_attn.mla_bmm_tail_padding_bytes: + q_projection.output_tail_padding_bytes = ( + self.mla_attn.mla_bmm_tail_padding_bytes + ) # The deployed NVFP4 writer accepts a scale tensor but discards it and # quantizes with outer scale 1.0. Feeding x/s_l is exactly the missing diff --git a/vllm/model_executor/layers/utils.py b/vllm/model_executor/layers/utils.py index 4d2e50420f42..a53d1b2e64b2 100644 --- a/vllm/model_executor/layers/utils.py +++ b/vllm/model_executor/layers/utils.py @@ -11,6 +11,7 @@ from vllm._aiter_ops import rocm_aiter_ops from vllm.logger import init_logger from vllm.platforms import CpuArchEnum, current_platform +from vllm.utils.cublas import tail_padded_empty from vllm.utils.platform_utils import num_compute_units from vllm.utils.torch_utils import direct_register_custom_op @@ -98,6 +99,43 @@ def default_unquantized_gemm( return torch.nn.functional.linear(x, weight, bias) +def tail_padded_unquantized_gemm_impl( + x: torch.Tensor, + weight: torch.Tensor, + tail_padding_bytes: int, +) -> torch.Tensor: + """Run an unbiased linear directly into storage with a mapped tail.""" + output_shape = (*x.shape[:-1], weight.shape[0]) + output = tail_padded_empty( + output_shape, + device=x.device, + dtype=x.dtype, + tail_padding_bytes=tail_padding_bytes, + ) + torch.mm( + x.reshape(-1, x.shape[-1]), + weight.t(), + out=output.view(-1, weight.shape[0]), + ) + return output + + +def tail_padded_unquantized_gemm_fake( + x: torch.Tensor, + weight: torch.Tensor, + tail_padding_bytes: int, +) -> torch.Tensor: + del tail_padding_bytes + return x.new_empty((*x.shape[:-1], weight.shape[0])) + + +direct_register_custom_op( + op_name="tail_padded_unquantized_gemm", + op_func=tail_padded_unquantized_gemm_impl, + fake_impl=tail_padded_unquantized_gemm_fake, +) + + def use_aiter_triton_gemm(n, m, k, dtype): if ( not rocm_aiter_ops.is_triton_gemm_enabled() diff --git a/vllm/utils/cublas.py b/vllm/utils/cublas.py new file mode 100644 index 000000000000..adecc9e5f467 --- /dev/null +++ b/vllm/utils/cublas.py @@ -0,0 +1,59 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Storage helpers for CUDA library operations with bounded read-ahead.""" + +from __future__ import annotations + +import math + +import torch + +CUBLAS_BMM_TAIL_PADDING_BYTES = 64 * 1024 + + +def tail_padded_empty( + shape: tuple[int, ...], + *, + device: torch.device, + dtype: torch.dtype, + tail_padding_bytes: int = CUBLAS_BMM_TAIL_PADDING_BYTES, +) -> torch.Tensor: + """Allocate a contiguous tensor with live storage after its last item.""" + if tail_padding_bytes < 0: + raise ValueError("tail_padding_bytes must be non-negative") + numel = math.prod(shape) + pad_numel = math.ceil(tail_padding_bytes / dtype.itemsize) + storage = torch.empty(numel + pad_numel, device=device, dtype=dtype) + return storage[:numel].view(shape) + + +def storage_tail_bytes(tensor: torch.Tensor) -> int: + """Return storage bytes following the tensor's highest addressed item.""" + if tensor.numel() == 0: + end = tensor.storage_offset() * tensor.element_size() + else: + if any(stride < 0 for stride in tensor.stride()): + return 0 + last_item = tensor.storage_offset() + sum( + (size - 1) * stride + for size, stride in zip(tensor.shape, tensor.stride(), strict=True) + ) + end = (last_item + 1) * tensor.element_size() + return max(0, tensor.untyped_storage().nbytes() - end) + + +def ensure_cublas_tail_padding( + tensor: torch.Tensor, + tail_padding_bytes: int = CUBLAS_BMM_TAIL_PADDING_BYTES, +) -> torch.Tensor: + """Ensure a downstream cuBLAS BMM can perform bounded mapped read-ahead.""" + if storage_tail_bytes(tensor) >= tail_padding_bytes: + return tensor + output = tail_padded_empty( + tuple(tensor.shape), + device=tensor.device, + dtype=tensor.dtype, + tail_padding_bytes=tail_padding_bytes, + ) + output.copy_(tensor) + return output diff --git a/vllm/v1/attention/backends/mla/b12x_mla_sparse.py b/vllm/v1/attention/backends/mla/b12x_mla_sparse.py index 5fdffc083e48..4a82da1d50b5 100644 --- a/vllm/v1/attention/backends/mla/b12x_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/b12x_mla_sparse.py @@ -976,9 +976,12 @@ def __init__( self.v_head_dim: int = mla_args.get("v_head_dim", 512) # GLM_NSA contract: q_head_dim = kv_lora_rank (512) + qk_rope (64) = 576. self.q_head_dim = self.kv_lora_rank + self.qk_rope_head_dim - self.force_contiguous_mla_bmm_input = True - self.force_contiguous_mla_bmm_weight = True - self.force_contiguous_mla_bmm_output = True + # SM120 cuBLAS strided BMM may read one tile beyond a logical operand. + # Producers use this contract to keep that bounded access mapped without + # adding per-step layout copies. + from vllm.utils.cublas import CUBLAS_BMM_TAIL_PADDING_BYTES + + self.mla_bmm_tail_padding_bytes = CUBLAS_BMM_TAIL_PADDING_BYTES # The indexer carries the shared buffer for normal layers and tests; # the explicitly-passed buffer covers backbone skip layers, whose diff --git a/vllm/v1/attention/ops/common.py b/vllm/v1/attention/ops/common.py index a1688b92d4af..42a022010b53 100644 --- a/vllm/v1/attention/ops/common.py +++ b/vllm/v1/attention/ops/common.py @@ -231,6 +231,7 @@ def cp_lse_ag_out_rs( ctx: CPTritonContext | None = None, return_lse: bool = False, is_lse_base_on_e=True, + output_tail_padding_bytes: int = 0, ): """ cp_attn_out: [ B, H, D ] @@ -240,6 +241,10 @@ def cp_lse_ag_out_rs( cp_attn_out, cp_attn_lse, cp_group, ctx=ctx, is_lse_base_on_e=is_lse_base_on_e ) out = cp_group.reduce_scatter(out, dim=1) + if output_tail_padding_bytes: + from vllm.utils.cublas import ensure_cublas_tail_padding + + out = ensure_cublas_tail_padding(out, output_tail_padding_bytes) if return_lse: cp_num_heads = lse.shape[1] // cp_group.world_size diff --git a/vllm/v1/attention/ops/dcp_alltoall.py b/vllm/v1/attention/ops/dcp_alltoall.py index f647124a47e5..7bebf2f0c895 100644 --- a/vllm/v1/attention/ops/dcp_alltoall.py +++ b/vllm/v1/attention/ops/dcp_alltoall.py @@ -30,6 +30,7 @@ import vllm.envs as envs from vllm.logger import init_logger from vllm.triton_utils import tl, triton +from vllm.utils.cublas import tail_padded_empty from vllm.utils.multi_stream_utils import is_vllm_cudagraph_capture_active if TYPE_CHECKING: @@ -220,6 +221,7 @@ def _try_b12x_dcp_lse_reduce( is_lse_base_on_e: bool, max_batch_size: int | None, query_head_dim: int | None, + output_tail_padding_bytes: int = 0, ) -> torch.Tensor | None: """Use the low-latency B12X PCIe path when its contract is satisfied.""" world_size = cp_group.world_size @@ -286,9 +288,27 @@ def _try_b12x_dcp_lse_reduce( if not cp_attn_lse.is_contiguous(): cp_attn_lse = cp_attn_lse.contiguous() + output = None + if output_tail_padding_bytes: + # The GLM MLA V up-projection uses a strided cuBLAS BMM whose SM120 + # kernel reads ahead beyond the logical operand. Write the collective + # result directly into storage that keeps this bounded read mapped. + output = tail_padded_empty( + (batch, total_heads // world_size, head_dim), + device=cp_attn_out.device, + dtype=cp_attn_out.dtype, + tail_padding_bytes=output_tail_padding_bytes, + ) + if output is None: + return pool.lse_reduce_scatter( + cp_attn_out, + cp_attn_lse, + is_lse_base_on_e=is_lse_base_on_e, + ) return pool.lse_reduce_scatter( cp_attn_out, cp_attn_lse, + out=output, is_lse_base_on_e=is_lse_base_on_e, ) @@ -600,8 +620,7 @@ def _dcp_a2a_send_recv_buffers( # and free the captured address. Eager calls therefore use ordinary temporary # tensors, while captured calls retain fixed-size owners below. if device.type == "cuda" and ( - is_vllm_cudagraph_capture_active() - or torch.cuda.is_current_stream_capturing() + is_vllm_cudagraph_capture_active() or torch.cuda.is_current_stream_capturing() ): # FULL graphs share a global graph pool. Without a live Python owner, # a larger descriptor's staging allocation can be recycled while a @@ -851,13 +870,23 @@ def _dcp_a2a_unpack_combine( lse_pack_dim: int, return_lse: bool, is_lse_base_on_e: bool, + output_tail_padding_bytes: int = 0, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: world_size, num_tokens, h_per_rank, _ = recv_buffer.shape - out = torch.empty( - (num_tokens, h_per_rank, head_dim), - device=recv_buffer.device, - dtype=recv_buffer.dtype, - ) + output_shape = (num_tokens, h_per_rank, head_dim) + if output_tail_padding_bytes: + out = tail_padded_empty( + output_shape, + device=recv_buffer.device, + dtype=recv_buffer.dtype, + tail_padding_bytes=output_tail_padding_bytes, + ) + else: + out = torch.empty( + output_shape, + device=recv_buffer.device, + dtype=recv_buffer.dtype, + ) out_lse = torch.empty( (num_tokens, h_per_rank) if return_lse else (1, 1), device=recv_buffer.device, @@ -898,6 +927,7 @@ def dcp_a2a_lse_reduce( use_b12x: bool = False, b12x_max_batch_size: int | None = None, b12x_query_head_dim: int | None = None, + output_tail_padding_bytes: int = 0, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: """ Combine partial attention outputs across DCP ranks using All-to-All. @@ -915,6 +945,8 @@ def dcp_a2a_lse_reduce( use_b12x: Try the low-latency B12X PCIe path before NCCL A2A b12x_max_batch_size: Configured token capacity for B12X staging b12x_query_head_dim: Query width when it differs from output width + output_tail_padding_bytes: Mapped storage required after the result for + a following cuBLAS BMM Returns: Combined output [B, H/N, D] (head-scattered) @@ -936,6 +968,7 @@ def dcp_a2a_lse_reduce( is_lse_base_on_e=is_lse_base_on_e, max_batch_size=b12x_max_batch_size, query_head_dim=b12x_query_head_dim, + output_tail_padding_bytes=output_tail_padding_bytes, ) if b12x_result is not None: return b12x_result @@ -975,5 +1008,10 @@ def dcp_a2a_lse_reduce( work.wait() return _dcp_a2a_unpack_combine( - recv_buffer, D, lse_pack_dim, return_lse, is_lse_base_on_e + recv_buffer, + D, + lse_pack_dim, + return_lse, + is_lse_base_on_e, + output_tail_padding_bytes, )