From eb5597e11ebc3da5b16f9de868e8bf4817b21da7 Mon Sep 17 00:00:00 2001 From: Xuan Jiang Date: Wed, 2 Sep 2026 20:48:29 +0000 Subject: [PATCH] hybrid unordered kernels: guard `token_map_at_dispatch` bit packing `pack_combine_recv_addr` in `combine_utils.cuh` packs rank/slot/channel into 5/14/12 bits by shifting and ORing without any mask. Callers pass unchecked ints: - `hybrid_dispatch_unordered.cuh` writes the packed word into `token_map_at_dispatch[token][k]` - `hybrid_combine_unordered.cuh` reads it back and asserts only that the channel bits match the channel it is running on If `slot` exceeds 16383 (its 14-bit field) it bleeds into the rank field at bit 26; the channel bits stay intact so the runtime `EP_DEVICE_ASSERT(channel == channel_idx)` still passes. The result is a misrouted combine partial with no visible failure, or a hang in the arrival-count wait when the corrupted rank never sends. Slot is `combine_slot`, bounded by `kNumMaxTokensPerChannel * kNumTopk`. For a large `kNumMaxTokensPerRank` with few channels and topk >= 8 this can exceed 16383. rank and channel have the same silent-overflow shape at 32 / 4096. All three bounds are compile-time constants on the two kernel templates, so `EP_STATIC_ASSERT` at the top of each kernel body catches the bad instantiation at NVRTC JIT time with a clear message that names the offending template parameter, instead of at runtime with a misrouted put. --- deep_ep/include/deep_ep/impls/combine_utils.cuh | 6 ++++++ .../deep_ep/impls/hybrid_combine_unordered.cuh | 14 ++++++++++++++ .../deep_ep/impls/hybrid_dispatch_unordered.cuh | 13 +++++++++++++ 3 files changed, 33 insertions(+) diff --git a/deep_ep/include/deep_ep/impls/combine_utils.cuh b/deep_ep/include/deep_ep/impls/combine_utils.cuh index 4c7d975c3..3bf37e55d 100644 --- a/deep_ep/include/deep_ep/impls/combine_utils.cuh +++ b/deep_ep/include/deep_ep/impls/combine_utils.cuh @@ -34,6 +34,12 @@ constexpr int kCombineRecvMapRankMask = (1 << kCombineRecvMapRankBits) - 1; constexpr int kCombineRecvMapSlotShift = kCombineRecvMapChannelBits; constexpr int kCombineRecvMapRankShift = kCombineRecvMapChannelBits + kCombineRecvMapSlotBits; +// Callers of `pack_combine_recv_addr` are expected to ensure rank/slot/channel +// fit in 5/14/12 bits respectively (kCombineRecvMap*Bits). No runtime mask is +// applied, so an out-of-range value bleeds into the neighbouring field silently +// — and `pack_combine_recv_addr`'s receiver checks only match on the channel +// bits. The two hybrid unordered kernels enforce these bounds at compile time +// on the template arguments they instantiate this with. __device__ __host__ __forceinline__ int pack_combine_recv_addr(const int& rank, const int& slot, const int& channel) { return (rank << kCombineRecvMapRankShift) | (slot << kCombineRecvMapSlotShift) | channel; diff --git a/deep_ep/include/deep_ep/impls/hybrid_combine_unordered.cuh b/deep_ep/include/deep_ep/impls/hybrid_combine_unordered.cuh index e67b36a99..4c7e305ce 100644 --- a/deep_ep/include/deep_ep/impls/hybrid_combine_unordered.cuh +++ b/deep_ep/include/deep_ep/impls/hybrid_combine_unordered.cuh @@ -57,6 +57,20 @@ hybrid_unordered_combine_impl(nv_bfloat16* x, void* buffer, void* workspace, const int scaleout_rank_idx, const int scaleup_rank_idx, int num_reduced_tokens, const int num_combined_tokens) { + // Guard `token_map_at_dispatch` bit packing: rank/slot/channel are ORed + // together with no masking in `pack_combine_recv_addr`, so a value that + // exceeds its field width bleeds into the neighbouring field silently. + // The `channel == channel_idx` assertion in the read-back stays intact + // when only slot overflows (channel is the lowest 12 bits), so nothing + // catches this at runtime; a compile-time gate keeps invalid template + // instantiations off the shelf. + EP_STATIC_ASSERT(kNumScaleoutRanks <= (1 << kCombineRecvMapRankBits), + "kNumScaleoutRanks exceeds the 5-bit dst_scaleout_rank field in combine_recv_addr"); + EP_STATIC_ASSERT(kNumMaxTokensPerChannel * kNumTopk <= (1 << kCombineRecvMapSlotBits), + "kNumMaxTokensPerChannel * kNumTopk exceeds the 14-bit slot field in combine_recv_addr"); + EP_STATIC_ASSERT(kNumChannels <= (1 << kCombineRecvMapChannelBits), + "kNumChannels exceeds the 12-bit channel field in combine_recv_addr"); + // Utils const auto sm_idx = static_cast(blockIdx.x); const auto thread_idx = static_cast(threadIdx.x); diff --git a/deep_ep/include/deep_ep/impls/hybrid_dispatch_unordered.cuh b/deep_ep/include/deep_ep/impls/hybrid_dispatch_unordered.cuh index 816f0571d..4aa4ca8c4 100644 --- a/deep_ep/include/deep_ep/impls/hybrid_dispatch_unordered.cuh +++ b/deep_ep/include/deep_ep/impls/hybrid_dispatch_unordered.cuh @@ -178,6 +178,19 @@ hybrid_unordered_dispatch_impl( EP_STATIC_ASSERT(kNumScaleoutWarps == kNumForwardWarps, "Invalid warp size"); EP_STATIC_ASSERT(kNumParts <= layout::WorkspaceLayout::kNumMaxParts, "kNumParts exceeds the per-part header workspace capacity"); + + // Guard the `token_map_at_dispatch` bit packing this kernel writes into + // (see `pack_combine_recv_addr` and its callsite below). Rank/slot/channel + // are ORed together with no masking, so an out-of-range value bleeds into + // the neighbouring field silently — and the receiver's channel == channel_idx + // assertion doesn't catch slot-overflow because the channel bits stay intact. + // Refuse to instantiate a template combination that could produce that. + EP_STATIC_ASSERT(kNumScaleoutRanks <= (1 << kCombineRecvMapRankBits), + "kNumScaleoutRanks exceeds the 5-bit dst_scaleout_rank field in combine_recv_addr"); + EP_STATIC_ASSERT(kNumMaxTokensPerChannel * kNumTopk <= (1 << kCombineRecvMapSlotBits), + "kNumMaxTokensPerChannel * kNumTopk exceeds the 14-bit slot field in combine_recv_addr"); + EP_STATIC_ASSERT(kNumChannels <= (1 << kCombineRecvMapChannelBits), + "kNumChannels exceeds the 12-bit channel field in combine_recv_addr"); EP_STATIC_ASSERT(kNumParts >= 1, "Invalid part count"); EP_STATIC_ASSERT(kNumSubParts >= 1, "Invalid sub-part count"); EP_STATIC_ASSERT(kNumSubParts <= kBatchSize,