From 5b3608db636b5b2df05be5d07b0b47e1cbd0ddbc Mon Sep 17 00:00:00 2001 From: Xin Guan <294044352+xguannv@users.noreply.github.com> Date: Tue, 8 Sep 2026 10:42:47 +0000 Subject: [PATCH 1/3] [None][perf] Helix post-process: stop writing partial_o twice on both alltoall backends A CP rank that owns no KV blocks for a token attends to zero keys, so the attention kernel normalizes by a zero softmax sum and produces NaN. Before the exchange those rows are forced to a no-op contribution (zeroed partial_o, softmax stats of (-inf, 0)) by _helix_sanitize_empty_kv, which materializes a full extra copy of partial_o and softmax_stats on every layer. Fold that fill into work the two backends already do: * MNNVL fifo v2: the all-to-all sender already streams every byte of an entry through shared memory before packing it, so the masked entries are rewritten there, behind a new optional zero_kv_mask argument on trtllm::alltoall_helix_native. The entry is not packed yet, so the overwrite costs no extra global traffic. * NCCL: _helix_sanitize_empty_kv now runs inside the same torch.compile region as the transpose/split that feeds alltoall_helix, so inductor folds the fill into the transposed store instead of writing partial_o and reading it back. fifo v1 keeps calling _helix_sanitize_empty_kv eagerly, and passing zero_kv_mask=None preserves the previous behaviour everywhere, so the sender change is a no-op unless a caller opts in. The outputs are numerically unchanged: the same rows are forced to the same no-op values, just earlier. On GB300 with cp16 the post-process block drops from 10 kernels to 4 and from 28.90 to 16.84 us per layer, with the 16 per-rank ranges non-overlapping ([25.98, 29.81] vs [14.47, 17.47]). The all-to-all kernel itself is unchanged (9.76 -> 9.96 us), confirming the shared-memory rewrite is free. An independent in-process A/B under CUDA graph capture measures -9.71 us per call. Signed-off-by: Xin Guan <294044352+xguannv@users.noreply.github.com> --- cpp/tensorrt_llm/kernels/helixAllToAll.cu | 27 +++++++ cpp/tensorrt_llm/kernels/helixAllToAll.h | 8 +++ cpp/tensorrt_llm/thop/alltoallOp.cpp | 25 ++++++- tensorrt_llm/_torch/attention/attention.py | 70 +++++++++++++++---- .../_torch/custom_ops/cpp_custom_ops.py | 7 +- tensorrt_llm/_torch/distributed/ops.py | 10 ++- 6 files changed, 128 insertions(+), 19 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/helixAllToAll.cu b/cpp/tensorrt_llm/kernels/helixAllToAll.cu index 304b255a11f6..00a006bec899 100644 --- a/cpp/tensorrt_llm/kernels/helixAllToAll.cu +++ b/cpp/tensorrt_llm/kernels/helixAllToAll.cu @@ -398,6 +398,33 @@ __global__ void helixAllToAllKernel(HelixAllToAllParams params) // note: we don't need to pack anything, fields are already packed in // shared memory + // Zero-local-KV sanitization: the entry is in shared memory and + // not yet packed, so overwriting it costs no global traffic. + if (params.zeroKvMask != nullptr && params.zeroKvMask[entryIdx / params.zeroKvMaskDivisor] != 0) + { + // field0Size is a multiple of 16 (the op checks it), so the + // int4 store never runs off the end of the field. + int const field0Size = getFieldSize(params.sendFields[0]); + for (int off = laneId * static_cast(sizeof(int4)); off < field0Size; + off += WARP_SIZE * static_cast(sizeof(int4))) + { + *reinterpret_cast(shmem + off) = make_int4(0, 0, 0, 0); + } + if (laneId == 0) + { + // Field 1 sits at getFieldSize(field 0), not at + // align_up(..., 16) as computeTotalUnpackedSize computes; + // they agree only because field 0 is 16-byte aligned. + auto* stats = reinterpret_cast(shmem + field0Size); + int const statsCount = getFieldSize(params.sendFields[1]) / static_cast(sizeof(float2)); + for (int i = 0; i < statsCount; ++i) + { + stats[i] = make_float2(-INFINITY, 0.F); + } + } + __syncwarp(); + } + LL128Proto::protoPack(shmem, head, singlePacked128ByteCount, fifoEntry128ByteIndexBase, laneId); uint64_t* fifoEntry = fifoBase + fifoEntryIndex * (HELIX_FIFO_ENTRY_BYTES / sizeof(uint64_t)); diff --git a/cpp/tensorrt_llm/kernels/helixAllToAll.h b/cpp/tensorrt_llm/kernels/helixAllToAll.h index 95ab959ff3fc..ac51e012adbf 100644 --- a/cpp/tensorrt_llm/kernels/helixAllToAll.h +++ b/cpp/tensorrt_llm/kernels/helixAllToAll.h @@ -47,6 +47,14 @@ struct HelixAllToAllParams int cpSize; int channelCount; // use 0 to auto-compute int maxChannelCount; + + // Rows this rank owns no KV for. The sender replaces them in shared memory + // with a no-op contribution for the combine: field 0 zeros, field 1 + // (max, sum) = (-inf, 0). nullptr when the caller already sanitized. + uint8_t const* zeroKvMask; + // entryCount / zeroKvMask length: 1 when an entry is a token (fifo v2), + // num_heads when it is a (token, head) pair (fifo v1). + int zeroKvMaskDivisor; }; // ============================================================================ diff --git a/cpp/tensorrt_llm/thop/alltoallOp.cpp b/cpp/tensorrt_llm/thop/alltoallOp.cpp index 8a775b1cf6b8..9c076850d29f 100644 --- a/cpp/tensorrt_llm/thop/alltoallOp.cpp +++ b/cpp/tensorrt_llm/thop/alltoallOp.cpp @@ -21,6 +21,7 @@ #include "tensorrt_llm/runtime/utils/mpiUtils.h" #include "tensorrt_llm/thop/thUtils.h" +#include #include TRTLLM_NAMESPACE_BEGIN @@ -125,10 +126,13 @@ std::vector alltoall_helix( * @param workspace Workspace tensor (uint64, strided across ranks) * @param cp_rank Current context parallel rank * @param cp_size Total number of context parallel ranks + * @param zero_kv_mask Optional bool mask of rows this rank owns no KV for. The + * sender replaces them with a no-op contribution, so the caller must not + * sanitize itself. Its length must divide entry_count. * @return tuple of (partial_o_out, softmax_stats_out) with same shapes as inputs */ -std::tuple alltoall_helix_native( - torch::Tensor partial_o, torch::Tensor softmax_stats, torch::Tensor workspace, int64_t cp_rank, int64_t cp_size) +std::tuple alltoall_helix_native(torch::Tensor partial_o, torch::Tensor softmax_stats, + torch::Tensor workspace, int64_t cp_rank, int64_t cp_size, std::optional zero_kv_mask) { // Input validation @@ -224,6 +228,21 @@ std::tuple alltoall_helix_native( params.channelCount = 0; // auto-compute params.maxChannelCount = tensorrt_llm::kernels::computeHelixMaxChannelCount(cp_size); + // Optional zero-local-KV sanitization, applied by the sender in shared memory + params.zeroKvMask = nullptr; + params.zeroKvMaskDivisor = 1; + if (zero_kv_mask.has_value()) + { + auto const& mask = zero_kv_mask.value(); + CHECK_TH_CUDA(mask); + CHECK_CONTIGUOUS(mask); + CHECK_TYPE(mask, at::ScalarType::Bool); + TORCH_CHECK(mask.numel() > 0 && entry_count % mask.numel() == 0, "zero_kv_mask numel (", mask.numel(), + ") must divide the all-to-all entry count (", entry_count, ")"); + params.zeroKvMask = reinterpret_cast(mask.data_ptr()); + params.zeroKvMaskDivisor = entry_count / mask.numel(); + } + // Launch kernel auto stream = at::cuda::getCurrentCUDAStream(); tensorrt_llm::kernels::launchHelixAllToAll(params, allowVariableField1, stream); @@ -260,7 +279,7 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m) m.def("alltoall_helix(Tensor[] input_list, int[] group, int? num_lists) -> Tensor[]"); m.def( "alltoall_helix_native(Tensor partial_o, Tensor softmax_stats, Tensor(a!) workspace, int " - "cp_rank, int cp_size) -> (Tensor, Tensor)"); + "cp_rank, int cp_size, Tensor? zero_kv_mask=None) -> (Tensor, Tensor)"); m.def( "initialize_helix_workspace(Tensor(a!) workspace, int cp_rank, int cp_size) " "-> ()"); diff --git a/tensorrt_llm/_torch/attention/attention.py b/tensorrt_llm/_torch/attention/attention.py index 94344ff8247b..7565481f54e4 100644 --- a/tensorrt_llm/_torch/attention/attention.py +++ b/tensorrt_llm/_torch/attention/attention.py @@ -193,6 +193,46 @@ def _helix_sanitize_empty_kv( return partial_o, softmax_stats +@torch.compile(dynamic=False) +def _helix_nccl_pre_alltoall( + partial_o: torch.Tensor, + softmax_stats: torch.Tensor, + zero_kv_mask: Optional[torch.Tensor], + cp_size: int, +) -> List[torch.Tensor]: + """Sanitize zero-local-KV rows and reformat into the alltoall send layout. + + ``_helix_sanitize_empty_kv`` followed by the transpose and split + ``_helix_post_process`` used to do inline. Both live in one compiled region + so inductor folds the fill into the transposed store instead of writing + ``partial_o`` and reading it straight back. Dynamo inlines the call, so the + sanitize has exactly one definition and the fusion is unaffected. + + ``dynamic=False`` is deliberate. With dynamic shapes inductor emits a much + slower transposed store, and picks different grids on different ranks for + the same shape, which cancels the win at the collective. + + The cost is one specialization per CUDA-graph batch bucket. Exceeding + dynamo's ``cache_size_limit`` falls back to eager SILENTLY -- the symptom is + ``triton_poi_fused_*`` disappearing from the trace, not an error. + """ + partial_o, softmax_stats = _helix_sanitize_empty_kv(partial_o, + softmax_stats, + zero_kv_mask) + chunks = [] + for t in (partial_o, softmax_stats): + t = t.transpose(1, 0).contiguous() + chunks.extend(torch.split(t, t.shape[0] // cp_size)) + return chunks + + +@torch.compile(dynamic=False) +def _helix_nccl_post_alltoall( + gathered: List[torch.Tensor]) -> List[torch.Tensor]: + """Reformat the gathered partials into the helix_post_process layout.""" + return [t.transpose(1, 2).contiguous() for t in gathered] + + def _helix_post_process( partial_o: torch.Tensor, softmax_stats: torch.Tensor, @@ -210,25 +250,23 @@ def _helix_post_process( dimension that differs between the two callers is *value_dim* (``head_dim`` for MHA, ``kv_lora_rank`` for MLA). - zero_kv_mask marks tokens for which this CP rank owns no KV blocks; those rows - are forced to a no-op contribution before the exchange (see - _helix_sanitize_empty_kv). + zero_kv_mask marks tokens this CP rank owns no KV for; those rows are forced + to a no-op contribution before the exchange, exactly once per backend: + NCCL in _helix_nccl_pre_alltoall, fused with the reformat + fifo v2 in the all-to-all sender, while the entry is in shared memory + fifo v1 here, via _helix_sanitize_empty_kv When *aux_stream* and *ln_events* are provided the two ``.contiguous()`` calls in the FIFO-v1 path are overlapped on separate CUDA streams for better performance. """ - partial_o, softmax_stats = _helix_sanitize_empty_kv(partial_o, - softmax_stats, - zero_kv_mask) if mapping.cp_config.get("use_nccl_for_alltoall", True): - # NCCL-based implementation using alltoall_helix. - chunks = [] - for t in [partial_o, softmax_stats]: - t = t.transpose(1, 0).contiguous() - chunks.extend(torch.split(t, t.shape[0] // mapping.cp_size)) + # NCCL path. Sanitize is folded into _helix_nccl_pre_alltoall so + # inductor can fuse it into the reformat. + chunks = _helix_nccl_pre_alltoall(partial_o, softmax_stats, + zero_kv_mask, mapping.cp_size) gathered = alltoall_helix(chunks, mapping.cp_group) - gathered = [t.transpose(1, 2).contiguous() for t in gathered] + gathered = _helix_nccl_post_alltoall(gathered) return torch.ops.trtllm.helix_post_process(gathered[0], gathered[1], 1.0) else: @@ -239,6 +277,8 @@ def _helix_post_process( fifo_version = mapping.cp_config.get("fifo_version", 2) if fifo_version == 1: + partial_o, softmax_stats = _helix_sanitize_empty_kv( + partial_o, softmax_stats, zero_kv_mask) def reshape_o(): return partial_o.view(num_tokens, cp_size, num_heads_tp_cp, @@ -265,12 +305,16 @@ def reshape_s(): return torch.ops.trtllm.helix_post_process_native( partial_o_out, softmax_stats_out, 1.0, 2) else: + # fifo_v2: one entry is one token, and the sender already streams + # every byte through shared memory, so the sanitize rides along + # there instead of 6 separate elementwise kernels (~40% of the block). partial_o = partial_o.view(num_tokens, cp_size, num_heads_tp_cp * value_dim) softmax_stats = softmax_stats.view(num_tokens, cp_size, num_heads_tp_cp * 2) partial_o_out, softmax_stats_out = helix.alltoall_native( - partial_o, softmax_stats) + partial_o, softmax_stats, + None if zero_kv_mask is None else zero_kv_mask[:num_tokens]) gathered_o = partial_o_out.view(num_tokens, cp_size, num_heads_tp_cp, value_dim) gathered_stats = softmax_stats_out.view(num_tokens, cp_size, diff --git a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py index 0d6658167766..62e1cf12a6ec 100644 --- a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py @@ -1340,7 +1340,12 @@ def _(input_list, group, num_lists): ] @torch.library.register_fake("trtllm::alltoall_helix_native") - def _(partial_o, softmax_stats, workspace, cp_rank, cp_size): + def _(partial_o, + softmax_stats, + workspace, + cp_rank, + cp_size, + zero_kv_mask=None): # Returns outputs with same shapes as inputs return partial_o.new_empty(partial_o.shape), softmax_stats.new_empty( softmax_stats.shape) diff --git a/tensorrt_llm/_torch/distributed/ops.py b/tensorrt_llm/_torch/distributed/ops.py index b3ea554da186..506da1d389b3 100644 --- a/tensorrt_llm/_torch/distributed/ops.py +++ b/tensorrt_llm/_torch/distributed/ops.py @@ -543,14 +543,19 @@ def get(mapping: Mapping) -> "HelixAllToAllNative": return HelixAllToAllNative._cache[mapping] - def alltoall_native(self, partial_o: torch.Tensor, - softmax_stats: torch.Tensor): + def alltoall_native(self, + partial_o: torch.Tensor, + softmax_stats: torch.Tensor, + zero_kv_mask: Optional[torch.Tensor] = None): """ Perform all-to-all data exchange. Args: partial_o: Tensor with shape [..., cp_size, kv_lora_rank], dtype half. softmax_stats: Tensor with shape [..., cp_size, 2], dtype float32. + zero_kv_mask: Optional bool mask over the entry dimension, True + where this rank owns no KV. The sender rewrites those rows to a + no-op contribution, so the caller must not sanitize them itself. Returns: Tuple of (partial_o_out, softmax_stats_out) with same shapes as inputs. @@ -561,6 +566,7 @@ def alltoall_native(self, partial_o: torch.Tensor, self.workspace_tensor, self.mapping.cp_rank, self.mapping.cp_size, + zero_kv_mask, ) return partial_o_out, softmax_stats_out From bff7230b69b58de3a3e0eaafc1338f55990e7c05 Mon Sep 17 00:00:00 2001 From: Xin Guan <294044352+xguannv@users.noreply.github.com> Date: Tue, 8 Sep 2026 10:50:59 +0000 Subject: [PATCH 2/3] [None][test] Cover the zero-local-KV rows of Helix post-processing across ranks A CP rank that owns no KV blocks for a token attends to zero keys and hands the combine NaN with a finite sentinel in the softmax stats, so those rows have to be neutralized before the exchange. Nothing exercised that through the exchange: test_helix_postprocess.py calls _helix_sanitize_empty_kv and _helix_zero_kv_mask directly, and test_mla_helix.py gives every rank an equal, non-zero slice of the KV (ctx_len_per_gpu = ctx_len // world_size), so its zero_kv_mask is all False and the neutralization never runs -- on any of the three backends. Add a 2-rank test that drives _helix_post_process with a per-rank mask, poisons the masked rows with NaN, and compares against a float64 reference, for nccl, fifo v1 and fifo v2, plus one CUDA-graph case for fifo v2 since that is how production runs it. A negative control withholds the mask and requires NaN in the output, so a pass means the check can actually fail. Signed-off-by: Xin Guan <294044352+xguannv@users.noreply.github.com> --- .../attention/multi_gpu/test_helix_zero_kv.py | 240 ++++++++++++++++++ 1 file changed, 240 insertions(+) create mode 100644 tests/unittest/_torch/attention/multi_gpu/test_helix_zero_kv.py diff --git a/tests/unittest/_torch/attention/multi_gpu/test_helix_zero_kv.py b/tests/unittest/_torch/attention/multi_gpu/test_helix_zero_kv.py new file mode 100644 index 000000000000..7700debefd05 --- /dev/null +++ b/tests/unittest/_torch/attention/multi_gpu/test_helix_zero_kv.py @@ -0,0 +1,240 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Multi-rank coverage for the zero-local-KV rows of Helix post-processing. + +A CP rank that owns no KV blocks for a token attends to zero keys, so the +attention kernel normalizes by a zero softmax sum and hands back NaN together +with a finite sentinel in the softmax stats. Those rows have to reach the +combine as an exact no-op, and *where* they are neutralized differs per backend: +inside the compiled pre-alltoall region for NCCL, inside the all-to-all sender +for fifo v2, eagerly for fifo v1. Checking any one of those in isolation would +not say the contract holds, so this drives ``_helix_post_process`` itself. + +Nothing else covers this. ``test_helix_postprocess.py`` exercises +``_helix_sanitize_empty_kv`` and ``_helix_zero_kv_mask`` directly but never the +exchange, and ``test_mla_helix.py`` drives all three backends end to end while +handing every rank an equal, non-zero slice of the KV +(``ctx_len_per_gpu = ctx_len // world_size``), so its ``zero_kv_mask`` is all +False and the neutralization never runs. + +Masked rows are poisoned with NaN before the exchange, and +``test_zero_kv_negative_control`` withholds the mask to prove that poison does +reach the output. Without that control a PASS here would only mean the test ran. + +Scope: this runs at cp_size 2, so it does not cover the sender's entry stride at +larger cp (``entryIdx`` advances by a channel count derived from cp_size). That +was checked out of tree at cp16. +""" + +import pickle +import sys + +import _torch.attention.multi_gpu.helix_test_utils as helix_utils +import cloudpickle +import pytest +import torch +from _torch.attention.multi_gpu.helix_test_utils import parse_comms_medium, run_single_rank +from mpi4py import MPI +from mpi4py.futures import MPIPoolExecutor +from utils.util import skip_pre_blackwell + +import tensorrt_llm +from tensorrt_llm._torch.attention.attention import _helix_post_process +from tensorrt_llm.mapping import CpType, Mapping + +cloudpickle.register_pickle_by_value(sys.modules[__name__]) +cloudpickle.register_pickle_by_value(helix_utils) +MPI.pickle.__init__( + cloudpickle.dumps, + cloudpickle.loads, + pickle.HIGHEST_PROTOCOL, +) + +WORLD_SIZE = 2 +NUM_HEADS = 6 # heads this rank owns after TP/CP splitting +VALUE_DIM = 512 # kv_lora_rank +# bf16 partials through an fp32 combine. Three orders below the ~1.0 scale of +# the data, and loose on purpose: the failure this hunts is NaN, not rounding. +TOLERANCE = 2e-2 + + +def _zero_kv_mask(rank: int, world_size: int, num_tokens: int, device: torch.device): + """Per-rank mask marking the tokens this rank owns no KV for. + + Deliberately different on every rank, so an implementation that applies some + other rank's mask fails. Only token ``rank`` plus a strided tail are masked, + which leaves most tokens combining every rank. An earlier version masked + ``t % world_size == rank`` and left exactly one contributor per token, so + every combine weight was 1.0 and the output came back bit-identical to the + input -- a real number from a test that never combined anything. + """ + idx = torch.arange(num_tokens, device=device) + mask = idx == rank + mask |= (idx >= world_size) & (idx % (2 * world_size + 1) == rank) + return mask + + +def _reference(all_o, all_stats, all_mask, my_rank, cp_size): + """float64 combine over every rank's pre-sanitize tensors. + + ``all_o[r]`` is rank r's partial output laid out + ``[num_tokens, cp_size, num_heads, value_dim]`` where the cp dimension + indexes the *destination* rank, so this rank's share is ``[:, my_rank]``. + """ + num_tokens = all_o[0].shape[0] + o = torch.stack([t.view(num_tokens, cp_size, NUM_HEADS, VALUE_DIM)[:, my_rank] for t in all_o]) + s = torch.stack([t.view(num_tokens, cp_size, NUM_HEADS, 2)[:, my_rank] for t in all_stats]) + m = torch.stack(all_mask)[:, :, None] # [ranks, tokens, 1] + + o = torch.where(m[..., None], torch.zeros((), dtype=torch.float64), o.double()) + smax = torch.where(m, torch.full((), float("-inf"), dtype=torch.float64), s[..., 0].double()) + ssum = torch.where(m, torch.zeros((), dtype=torch.float64), s[..., 1].double()) + + weight = ssum * torch.exp(smax - smax.max(dim=0).values) + weight = weight / weight.sum(dim=0) + return (o * weight[..., None]).sum(dim=0).reshape(num_tokens, NUM_HEADS * VALUE_DIM) + + +def _zero_kv_rank(rank, world_size, num_tokens, comms_medium, use_cuda_graph, withhold_mask): + """Rank body: build poisoned inputs, post-process, compare to a reference. + + Returns ``(nan_count, max_abs_error)``. ``max_abs_error`` is NaN when the + output contains NaN, since the comparison is meaningless there. + """ + comm = tensorrt_llm.mpi_comm() + device = torch.device("cuda", torch.cuda.current_device()) + + torch.manual_seed(1234 + rank) + partial_o = torch.randn( + num_tokens, world_size * NUM_HEADS * VALUE_DIM, device=device, dtype=torch.bfloat16 + ) + stats = torch.empty(num_tokens, world_size * NUM_HEADS, 2, device=device, dtype=torch.float32) + stats[..., 0].normal_(0.0, 2.0) # max + stats[..., 1].uniform_(0.5, 2.0) # sum, strictly positive + + mask = _zero_kv_mask(rank, world_size, num_tokens, device) + # Poison exactly what the neutralization is supposed to overwrite. The large + # finite max makes an unsanitized row win the combine outright rather than + # being rounded away, and 0 * NaN = NaN means its partial_o then poisons the + # result on every rank. + partial_o[mask] = float("nan") + stats[mask, :, 0] = 1e4 + stats[mask, :, 1] = 7.0 + + all_o = [torch.from_numpy(x) for x in comm.allgather(partial_o.float().cpu().numpy())] + all_s = [torch.from_numpy(x) for x in comm.allgather(stats.cpu().numpy())] + all_m = [torch.from_numpy(x) for x in comm.allgather(mask.cpu().numpy())] + reference = _reference(all_o, all_s, all_m, rank, world_size).to(device) + + use_nccl_for_alltoall, fifo_version = parse_comms_medium(comms_medium) + mapping = Mapping( + world_size=world_size, + rank=rank, + cp_size=world_size, + cp_config={ + "cp_type": CpType.HELIX, + "use_nccl_for_alltoall": use_nccl_for_alltoall, + "fifo_version": fifo_version, + }, + ) + zero_kv_mask = None if withhold_mask else mask + + if use_cuda_graph: + # Production captures this region, and the masked branch has only ever + # been exercised eagerly. Warm up on a side stream first, then capture + # on every rank at once -- the all-to-all is collective. + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + for _ in range(3): + _helix_post_process( + partial_o, stats, mapping, NUM_HEADS, VALUE_DIM, zero_kv_mask=zero_kv_mask + ) + torch.cuda.current_stream().wait_stream(side) + torch.cuda.synchronize() + comm.barrier() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output = _helix_post_process( + partial_o, stats, mapping, NUM_HEADS, VALUE_DIM, zero_kv_mask=zero_kv_mask + ) + torch.cuda.synchronize() + comm.barrier() + graph.replay() + torch.cuda.synchronize() + else: + output = _helix_post_process( + partial_o.clone(), + stats.clone(), + mapping, + NUM_HEADS, + VALUE_DIM, + zero_kv_mask=zero_kv_mask, + ) + + output = output.double() + nan_count = int(torch.isnan(output).sum()) + max_err = float("nan") if nan_count else float((output - reference).abs().max()) + return nan_count, max_err + + +def _launch(num_tokens, comms_medium, use_cuda_graph=False, withhold_mask=False): + """Run ``_zero_kv_rank`` on WORLD_SIZE ranks and collect their results.""" + args = (_zero_kv_rank, WORLD_SIZE, num_tokens, comms_medium, use_cuda_graph, withhold_mask) + with MPIPoolExecutor(max_workers=WORLD_SIZE) as executor: + return list(executor.map(run_single_rank, *zip(*[args] * WORLD_SIZE))) + + +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="needs 2 GPUs to run this test") +@skip_pre_blackwell +@pytest.mark.parametrize("num_tokens", [17, 96]) +@pytest.mark.parametrize("comms_medium", ["nccl", "fifo_v1", "fifo_v2"]) +def test_zero_kv_rows_are_neutral(comms_medium: str, num_tokens: int): + """Masked rows must contribute nothing, on every backend.""" + for rank, (nan_count, max_err) in enumerate(_launch(num_tokens, comms_medium)): + assert nan_count == 0, ( + f"rank {rank}: {nan_count} NaN in the output, so a zero-local-KV row " + f"reached the combine unsanitized" + ) + assert max_err < TOLERANCE, f"rank {rank}: max|out - ref| = {max_err:.3e}" + + +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="needs 2 GPUs to run this test") +@skip_pre_blackwell +def test_zero_kv_rows_are_neutral_under_cuda_graph(): + """Same contract inside a CUDA graph, which is how production runs it. + + fifo v2 only: it is the backend that neutralizes inside the all-to-all + sender, and NCCL collectives under capture are a separate question. + """ + for rank, (nan_count, max_err) in enumerate(_launch(96, "fifo_v2", use_cuda_graph=True)): + assert nan_count == 0, f"rank {rank}: {nan_count} NaN in the replayed output" + assert max_err < TOLERANCE, f"rank {rank}: max|out - ref| = {max_err:.3e}" + + +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="needs 2 GPUs to run this test") +@skip_pre_blackwell +@pytest.mark.parametrize("comms_medium", ["nccl", "fifo_v1", "fifo_v2"]) +def test_zero_kv_negative_control(comms_medium: str): + """Withholding the mask must produce NaN. + + This is what makes the tests above mean something: it shows the poison in + the masked rows really does reach the output when nothing neutralizes it. + """ + results = _launch(96, comms_medium, withhold_mask=True) + assert any(nan_count > 0 for nan_count, _ in results), ( + "no NaN with the mask withheld, so these tests cannot detect a missing " + "sanitize and their PASS means nothing" + ) From 73d0f2cd11850138c23fdc2bf52c624965013ee8 Mon Sep 17 00:00:00 2001 From: Xin Guan <294044352+xguannv@users.noreply.github.com> Date: Thu, 10 Sep 2026 06:13:51 +0000 Subject: [PATCH 3/3] [None][chore] Address review feedback on the Helix zero-KV sanitize Three things from review, plus the docstrings the coverage check asked for. alltoall_helix_native only asserted that zero_kv_mask is a CUDA tensor. The kernel dereferences it on partial_o's device, so a mask allocated on a different device passed validation and became an invalid access instead of an error. Require the two to match. Cover every zero_kv_mask validation rule with a test. These run on one GPU: everything the op does before the mask check is host-side shape validation and pointer setup, so an invalid mask raises without a collective and without an initialized MNNVL workspace. They live next to the existing test_helix_postprocess_native_invalid_inputs rather than in the multi-GPU file, which would spend two GPUs on a host-side check. The cross-device case needs a second GPU and is separate. Replace the two comm.barrier() calls that bracket the collective CUDA graph capture with a deadline-bounded Ibarrier. A rank that dies during capture used to leave its peers blocked forever, so the run reported as a hang with no traceback; now it fails and names the rank and the phase. No MPI_Abort: these tests share an MPIPoolExecutor with the rest of the session. Signed-off-by: Xin Guan <294044352+xguannv@users.noreply.github.com> --- cpp/tensorrt_llm/kernels/helixAllToAll.cu | 16 +++++ cpp/tensorrt_llm/kernels/helixAllToAll.h | 7 ++ cpp/tensorrt_llm/thop/alltoallOp.cpp | 5 ++ .../_torch/custom_ops/cpp_custom_ops.py | 10 ++- .../test_helix_postprocess.py | 64 +++++++++++++++++++ .../attention/multi_gpu/test_helix_zero_kv.py | 34 +++++++++- 6 files changed, 133 insertions(+), 3 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/helixAllToAll.cu b/cpp/tensorrt_llm/kernels/helixAllToAll.cu index 00a006bec899..6a885e0d4b68 100644 --- a/cpp/tensorrt_llm/kernels/helixAllToAll.cu +++ b/cpp/tensorrt_llm/kernels/helixAllToAll.cu @@ -294,6 +294,22 @@ __host__ __device__ __forceinline__ int computeProtoTransferSize(HelixFieldInfo // Main All-to-All Kernel // ============================================================================ +//! Exchange one entry per peer rank over the MNNVL FIFOs, both directions in one +//! launch. +//! +//! blockIdx.z picks the role: a block is either a sender or a receiver for one +//! (peer, channel) pair. The sender stages an entry's fields into shared memory +//! with TMA, packs them with LL128Proto and pushes the result into the peer's +//! FIFO; the receiver pops, unpacks and writes out. Both walk their channel's +//! entries strided by the channel count, entryIdx = channel, channel + +//! runChannelCount, ... +//! +//! params.zeroKvMask, when non-null, marks entries this rank owns no KV for. +//! They are rewritten in shared memory before packing -- see the sanitization +//! block below -- which is why it costs no extra global traffic. +//! +//! \tparam ALLOW_VARIABLE_FIELD1 field 1 carries more than one (max, sum) pair +//! per entry, i.e. the fifo v1 layout where an entry is (token, head). template __global__ void helixAllToAllKernel(HelixAllToAllParams params) { diff --git a/cpp/tensorrt_llm/kernels/helixAllToAll.h b/cpp/tensorrt_llm/kernels/helixAllToAll.h index ac51e012adbf..d9c316fc9234 100644 --- a/cpp/tensorrt_llm/kernels/helixAllToAll.h +++ b/cpp/tensorrt_llm/kernels/helixAllToAll.h @@ -36,6 +36,13 @@ struct HelixFieldInfo int stride; // Stride between rows in bytes }; +//! Arguments for one Helix CP all-to-all: what to send, where to receive, and +//! which peers participate. +//! +//! Two fields per direction, matching the pair the Helix combine needs: field 0 +//! is the partial attention output, field 1 the per-row (max, sum) softmax +//! stats. An "entry" is one unit of exchange per peer rank, so a rank sends +//! entryCount entries to each of the cpSize peers. struct HelixAllToAllParams { HelixFieldInfo sendFields[2]; diff --git a/cpp/tensorrt_llm/thop/alltoallOp.cpp b/cpp/tensorrt_llm/thop/alltoallOp.cpp index 9c076850d29f..fb973b686a41 100644 --- a/cpp/tensorrt_llm/thop/alltoallOp.cpp +++ b/cpp/tensorrt_llm/thop/alltoallOp.cpp @@ -235,6 +235,11 @@ std::tuple alltoall_helix_native(torch::Tensor par { auto const& mask = zero_kv_mask.value(); CHECK_TH_CUDA(mask); + // CHECK_TH_CUDA only asserts "is a CUDA tensor". The kernel dereferences + // this pointer on partial_o's device, so a mask on a different device + // would be an invalid access rather than an error. + TORCH_CHECK(mask.device() == partial_o.device(), "zero_kv_mask must be on the same device as partial_o (got ", + mask.device(), " vs ", partial_o.device(), ")"); CHECK_CONTIGUOUS(mask); CHECK_TYPE(mask, at::ScalarType::Bool); TORCH_CHECK(mask.numel() > 0 && entry_count % mask.numel() == 0, "zero_kv_mask numel (", mask.numel(), diff --git a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py index 62e1cf12a6ec..786cb38572c1 100644 --- a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py @@ -15,6 +15,11 @@ def _register_fake(): + """Register the meta ("fake") kernels for the trtllm C++ custom ops. + + Each entry mirrors one op's output shapes and dtypes without touching data, + which is what torch.compile and export need to trace through the op. + """ @torch.library.register_fake("trtllm::allreduce") def allreduce( @@ -1346,7 +1351,10 @@ def _(partial_o, cp_rank, cp_size, zero_kv_mask=None): - # Returns outputs with same shapes as inputs + """Exchanged tensors keep the shape and dtype of their inputs. + + zero_kv_mask only changes the values the sender writes, never a shape. + """ return partial_o.new_empty(partial_o.shape), softmax_stats.new_empty( softmax_stats.shape) diff --git a/tests/unittest/_torch/attention/kernels/parallel_hw_agnostic/test_helix_postprocess.py b/tests/unittest/_torch/attention/kernels/parallel_hw_agnostic/test_helix_postprocess.py index 8b9b33b3afdf..7449f04a94a9 100644 --- a/tests/unittest/_torch/attention/kernels/parallel_hw_agnostic/test_helix_postprocess.py +++ b/tests/unittest/_torch/attention/kernels/parallel_hw_agnostic/test_helix_postprocess.py @@ -410,6 +410,70 @@ def test_helix_postprocess_native_invalid_inputs(self): with pytest.raises(RuntimeError): torch.ops.trtllm.helix_post_process_native(gathered_o, gathered_stats, 1.0, 2) + @parameterized.expand( + [ + ("cpu",), + ("non_contiguous",), + ("wrong_dtype",), + ("empty",), + ("not_a_divisor",), + ] + ) + def test_alltoall_helix_native_rejects_bad_zero_kv_mask(self, case): + """Every zero_kv_mask validation rule must reject before the kernel launches. + + Single GPU on purpose. Everything alltoall_helix_native does ahead of the + mask check is host-side shape validation and pointer setup -- nothing + dereferences the workspace or needs an initialized MNNVL region -- so an + invalid mask raises without a collective. Only invalid masks belong here: + a valid one would go on to launch the kernel against this uninitialized + workspace. + """ + device = torch.device("cuda") + num_tokens, cp_size, value_dim = 8, 2, 64 + partial_o = torch.randn(num_tokens, cp_size, value_dim, dtype=torch.float16, device=device) + softmax_stats = torch.randn(num_tokens, cp_size, 2, dtype=torch.float32, device=device) + workspace = torch.zeros(cp_size, 8, dtype=torch.uint64, device=device) + + if case == "cpu": + mask = torch.zeros(num_tokens, dtype=torch.bool) + elif case == "non_contiguous": + mask = torch.zeros(num_tokens, 2, dtype=torch.bool, device=device)[:, 0] + elif case == "wrong_dtype": + mask = torch.zeros(num_tokens, dtype=torch.uint8, device=device) + elif case == "empty": + mask = torch.zeros(0, dtype=torch.bool, device=device) + elif case == "not_a_divisor": + # 3 does not divide entry_count == 8. + mask = torch.zeros(3, dtype=torch.bool, device=device) + else: + raise AssertionError(f"unhandled case: {case}") + + with pytest.raises(RuntimeError): + torch.ops.trtllm.alltoall_helix_native( + partial_o, softmax_stats, workspace, 0, cp_size, mask + ) + + @unittest.skipIf(torch.cuda.device_count() < 2, "needs 2 GPUs") + def test_alltoall_helix_native_rejects_cross_device_zero_kv_mask(self): + """A mask on another CUDA device must be an error, not an invalid access. + + Separate from the parameterized cases above because it is the one rule + that cannot be checked with a single GPU. + """ + num_tokens, cp_size, value_dim = 8, 2, 64 + partial_o = torch.randn( + num_tokens, cp_size, value_dim, dtype=torch.float16, device="cuda:0" + ) + softmax_stats = torch.randn(num_tokens, cp_size, 2, dtype=torch.float32, device="cuda:0") + workspace = torch.zeros(cp_size, 8, dtype=torch.uint64, device="cuda:0") + mask = torch.zeros(num_tokens, dtype=torch.bool, device="cuda:1") + + with pytest.raises(RuntimeError): + torch.ops.trtllm.alltoall_helix_native( + partial_o, softmax_stats, workspace, 0, cp_size, mask + ) + @parameterized.expand( [ # (layout,) — "nccl", "fifo_v1", "fifo_v2". diff --git a/tests/unittest/_torch/attention/multi_gpu/test_helix_zero_kv.py b/tests/unittest/_torch/attention/multi_gpu/test_helix_zero_kv.py index 7700debefd05..f619c0226d92 100644 --- a/tests/unittest/_torch/attention/multi_gpu/test_helix_zero_kv.py +++ b/tests/unittest/_torch/attention/multi_gpu/test_helix_zero_kv.py @@ -40,6 +40,7 @@ import pickle import sys +import time import _torch.attention.multi_gpu.helix_test_utils as helix_utils import cloudpickle @@ -65,11 +66,40 @@ WORLD_SIZE = 2 NUM_HEADS = 6 # heads this rank owns after TP/CP splitting VALUE_DIM = 512 # kv_lora_rank +# Generous: the barrier only has to outlast a CUDA graph capture, and a false +# timeout would be a flaky test. It exists to bound a hang, not to be tight. +BARRIER_TIMEOUT_S = 300.0 # bf16 partials through an fp32 combine. Three orders below the ~1.0 scale of # the data, and loose on purpose: the failure this hunts is NaN, not rounding. TOLERANCE = 2e-2 +def _bounded_barrier(comm, label: str, timeout_s: float = BARRIER_TIMEOUT_S): + """Barrier that fails instead of hanging when a peer never arrives. + + The two barriers here bracket a collective CUDA graph capture. With a plain + ``comm.barrier()``, a rank that dies during capture leaves its peers blocked + forever and the run reports as a hang -- no traceback, no failing assertion, + just a job that has to be killed. Polling a non-blocking barrier turns that + into a normal test failure that names the rank and the phase. + + Deliberately no ``MPI_Abort``: these tests share an ``MPIPoolExecutor`` with + the rest of the session, and aborting the world would take unrelated tests + down with it. Raising is enough for the executor to surface the failure. + """ + request = comm.Ibarrier() + deadline = time.monotonic() + timeout_s + while not request.Test(): + if time.monotonic() > deadline: + request.Cancel() + raise TimeoutError( + f"rank {comm.Get_rank()} waited {timeout_s:.0f}s at the '{label}' " + f"barrier; a peer never arrived (most likely it died during CUDA " + f"graph capture)" + ) + time.sleep(0.01) + + def _zero_kv_mask(rank: int, world_size: int, num_tokens: int, device: torch.device): """Per-rank mask marking the tokens this rank owns no KV for. @@ -164,14 +194,14 @@ def _zero_kv_rank(rank, world_size, num_tokens, comms_medium, use_cuda_graph, wi ) torch.cuda.current_stream().wait_stream(side) torch.cuda.synchronize() - comm.barrier() + _bounded_barrier(comm, "before capture") graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): output = _helix_post_process( partial_o, stats, mapping, NUM_HEADS, VALUE_DIM, zero_kv_mask=zero_kv_mask ) torch.cuda.synchronize() - comm.barrier() + _bounded_barrier(comm, "after capture") graph.replay() torch.cuda.synchronize() else: