hybrid unordered kernels: guard token_map_at_dispatch bit packing at JIT time

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;

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<int>(blockIdx.x);
     const auto thread_idx = static_cast<int>(threadIdx.x);

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,