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,