From 18bba3b789858efe1be60fc5d61a13001c4ae39b Mon Sep 17 00:00:00 2001 From: Juan Muneton Date: Tue, 25 Aug 2026 14:56:45 +0000 Subject: [PATCH 1/4] [xpu] Fix the Inkling gate dispatch predicate so XPU stops entering CUDA JIT `torch.version.hip is None` means "not ROCm", which is also true on XPU, so the Inkling MoE gate entered the CUDA JIT path on XPU and died on an `is_cuda` assert one frame later: - sigmoid_gate_topk_renorm -> inkling_gate_topk_renorm_v2 -> `assert logits.is_cuda` (inkling_gate_topk_renorm.py:77), at every token count - InklingGate.forward_fused -> inkling_gate_gemv[_fused] -> `assert x.is_cuda` (inkling_gate_topk_renorm.py:203), at <= 4 tokens Every other condition guarding those JIT gates is satisfied by Inkling's real inputs on XPU (k == 6, n_shared == 2, G == 258, stride % 8 == 0, ptr % 32 == 0, and both SGLANG_OPT_* env vars default to true), so all 64 MoE layers raised at default env, at any batch size. Measured on Arc Pro B60: AssertionError on all 12 benchmark shapes and every token count from 1 to 8192. Gate on `.is_cuda and torch.version.hip is None` rather than swapping one for the other. Neither conjunct is sufficient alone: - `torch.version.hip is None` admits XPU, which is the bug above. - `.is_cuda` admits ROCm. PyTorch gives HIP tensors device type `cuda`, so `.is_cuda` is true there; the repo already works around the same aliasing in `is_cuda()` (utils/common.py:151, which needs `torch.version.cuda is not None`) and in `is_arch_support_pdl()` (jit/utils/arch.py:179). Admitting ROCm would send it into inkling_gate_topk_renorm_v2_kernel, which uses `__reduce_max_sync` (inkling_gate_topk_renorm.cuh:301) with no HIP equivalent and no arch guard, and assumes a 32-lane warp (:22) against a 64-wide wavefront. The conjunction also keeps `forward_fused` in agreement with `InklingGate.__init__`, whose `ensure_gate_gemv_fused_scratch` call is guarded by `torch.cuda.is_available() and torch.version.hip is None`. On ROCm a bare `.is_cuda` would take the FUSED GEMV path with that scratch never allocated. CUDA and ROCm behaviour is unchanged; only XPU moves, and it moves from raising to the Triton kernel in the same module. --- python/sglang/kernels/ops/moe/sigmoid_gate_topk_renorm.py | 5 ++++- python/sglang/srt/models/inkling_common/moe.py | 3 +++ 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/python/sglang/kernels/ops/moe/sigmoid_gate_topk_renorm.py b/python/sglang/kernels/ops/moe/sigmoid_gate_topk_renorm.py index 0f55f5047b75..6c8c0594ee74 100644 --- a/python/sglang/kernels/ops/moe/sigmoid_gate_topk_renorm.py +++ b/python/sglang/kernels/ops/moe/sigmoid_gate_topk_renorm.py @@ -179,13 +179,16 @@ def sigmoid_gate_topk_renorm( A = k + n_shared_experts assert bias.numel() == N and bias.stride(-1) == 1, f"{bias.shape=} expected [{N}]" - # The production shape uses the specialized CUDA JIT kernel. + # The production shape uses the specialized CUDA JIT kernel. Both conjuncts are + # load-bearing: XPU satisfies `torch.version.hip is None`, and ROCm reports + # `.is_cuda` -- while the kernel needs `__reduce_max_sync`, which HIP lacks. if ( k == 6 and n_shared_experts == 2 and G == 258 and logits.stride(0) % 8 == 0 and logits.data_ptr() % 32 == 0 + and logits.is_cuda and torch.version.hip is None and envs.SGLANG_OPT_USE_GATE_TOPK_JIT.get() ): diff --git a/python/sglang/srt/models/inkling_common/moe.py b/python/sglang/srt/models/inkling_common/moe.py index f5b6c52e51ee..95f22acf8869 100644 --- a/python/sglang/srt/models/inkling_common/moe.py +++ b/python/sglang/srt/models/inkling_common/moe.py @@ -364,6 +364,9 @@ def forward_fused( and x.dtype == torch.bfloat16 and x.is_contiguous() and x.shape[-1] == _INKLING_GATE_GEMV_HIDDEN + # Both GEMV branches below are nvcc JIT kernels. ROCm reports `.is_cuda` + # too, so match the hip exclusion `__init__` uses for their scratch. + and x.is_cuda and torch.version.hip is None ): if gemv_mode >= GateGemvMode.FUSED: From fe268caa4fb0333bd08f742a58dfd8e9e26dcc03 Mon Sep 17 00:00:00 2001 From: Juan Muneton Date: Tue, 25 Aug 2026 14:56:48 +0000 Subject: [PATCH 2/4] [xpu] Normalize the sigmoid gate renorm in log space to avoid a 0/0 NaN `_sigmoid_gate_topk_renorm_kernel` formed the weights as `sigmoid(x) / sum(sigmoid(x))`. fp32 sigmoid goes subnormal below x ~ -87 and flushes to zero below x ~ -104, so once all K + S active logits sit in that region the denominator is exactly 0 and every routed weight and both shared gammas come back NaN. Normalize as `exp(lp - logsumexp(lp))` with `lp = logsigmoid(x)` instead. This is algebraically identical -- exp(log s_i - log sum_j s_j) == s_i / sum_j s_j -- and is the same form the eager reference already uses (`_logsigmoid_normalize`, srt/models/inkling_common/moe.py:140), so the fast path was the one that disagreed with the reference, not the other way round. Measured on Arc Pro B60 against an fp64 oracle: the all-negative-logit cases at base -90 and -200 go from NaN to 3.6e-7 / 4.8e-7 max abs error, base -40 was already clean, and all 10 other benchmark shapes are unaffected. Cost is within noise -- 58.74 -> 58.96 us at T=1, 181.31 -> 181.23 at T=512, 1237.28 -> 1239.06 at T=8192 (+0.37% / -0.04% / +0.14%). The same algebra appears in the CUDA counterpart (kernels/jit/csrc/moe/inkling_gate_topk_renorm.cuh:121-126 and :336-341). That half is NOT touched here and has NOT been measured -- there is no NVIDIA GPU on the machine this was developed on. It needs its own change from someone with CUDA hardware. --- .../kernels/ops/moe/sigmoid_gate_topk_renorm.py | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/python/sglang/kernels/ops/moe/sigmoid_gate_topk_renorm.py b/python/sglang/kernels/ops/moe/sigmoid_gate_topk_renorm.py index 6c8c0594ee74..bb33a455bc5a 100644 --- a/python/sglang/kernels/ops/moe/sigmoid_gate_topk_renorm.py +++ b/python/sglang/kernels/ops/moe/sigmoid_gate_topk_renorm.py @@ -113,10 +113,18 @@ def _sigmoid_gate_topk_renorm_kernel( active = tl.where(mask_k[None, :], routed_vals, shared) A: tl.constexpr = K + S - probs = tl.sigmoid(active) mask_a = offs_a < A - probs = tl.where(mask_a[None, :], probs, 0.0) - weights = probs / tl.sum(probs, axis=1, keep_dims=True) + # Normalize in LOG space, not as sigmoid(x) / sum(sigmoid(x)). The two are + # algebraically identical -- exp(log s_i - log sum_j s_j) == s_i / sum_j s_j -- + # but fp32 sigmoid goes subnormal below x ~ -87 and flushes to 0 below x ~ -104, + # so once every one of the A active logits is in that region the explicit form + # divides 0/0 and every routed weight and shared gamma comes back NaN. This is + # the same form as the eager reference `_logsigmoid_normalize` (moe.py:140). + # log(sigmoid(x)) = min(x, 0) - log1p(exp(-|x|)), exact for large |x| either sign. + lp = tl.minimum(active, 0.0) - tl.log(1.0 + tl.exp(-tl.abs(active))) + lp = tl.where(mask_a[None, :], lp, float("-inf")) + e = tl.where(mask_a[None, :], tl.exp(lp - tl.max(lp, axis=1)[:, None]), 0.0) + weights = e / tl.sum(e, axis=1, keep_dims=True) weights *= (route_scale * tl.load(global_scale_ptr)).to(weights.dtype) mask_rk = mask_m[:, None] & mask_k[None, :] From 050ac9505a6a53823ed682294e8b6c5e17b4f9f3 Mon Sep 17 00:00:00 2001 From: Juan Muneton Date: Tue, 25 Aug 2026 14:56:50 +0000 Subject: [PATCH 3/4] [xpu] Add a LOGSIGMOID_SINK epilogue to the unified router; 22-64x on Xe `_sigmoid_gate_topk_renorm_kernel` builds its top-k on `tl.topk` / `tl.bitonic_merge` over a packed (value << 16 | ~index) key. That codegen is pathological on Intel Xe: 1237 us at T=8192 against 27 us for the unified router's iterative masked-max on the same shape. The unified router (`_router_triton_kernel`) is already structurally right for this op -- iterative masked-max instead of `tl.topk`, BLOCK_M=1, num_warps=1 at BLOCK_N=256. Only its epilogue is wrong for Inkling. So parameterize the epilogue rather than adding a kernel: - `EPILOGUE` selects SUM_NORM (existing, default) or LOGSIGMOID_SINK. - `SHARED_SINK` counts trailing sink columns of the score row. They take no bias (bias stays [N], not [N + SHARED_SINK]) and never enter the top-k, but they do join the normalizer, in slots K_ROUTED .. K. - The value carried out of the top-k loop becomes the winner's RAW logit rather than its activated score, while ranking still happens on `sigmoid(logits) + bias`. Two quantities per column is why a plain "top-k and keep the winning score" kernel cannot serve this gate. - `HAS_GLOBAL_SCALE` applies the model's scalar global_scale, and `RETURN_PACKED` emits the (id << 16 | bf16 weight) form that InklingGate.emit_packed_topk consumes. The wrapper gains `shared_sink`, `global_scale` and `return_packed` kwargs and returns a 4-tuple (routed_weights, topk_indices, shared_weights, packed_topk) when `shared_sink > 0`. `sigmoid_gate_topk_renorm` dispatches here only when `not logits.is_cuda`. CUDA is excluded on purpose: every number below is an Xe measurement, there is no NVIDIA A/B for this epilogue, and the new tests are XPU-only, so a CUDA caller that misses the JIT gate above (k != 6, G != 258, unaligned row, or SGLANG_OPT_USE_GATE_TOPK_JIT=0) keeps the sort-based kernel it has today rather than silently changing kernel and numerics. Lifting the exclusion is a separate PR that needs NVIDIA hardware and a CUDA-side case in test/registered/kernels/ops/moe/test_moe_fused_gate.py. `SGLANG_OPT_USE_ROUTER_GATE_EPILOGUE` (EnvBool(True), registered in srt/environ.py next to the other Inkling gate knobs) turns the epilogue off on the devices that do reach it, leaving the sort-based kernel as the fallback. Correctness, Arc Pro B60, fp64 oracle, rtol=atol=2e-3: - pass on all 12 shapes, max abs 3.6e-7, ZERO index flips at every token count, including the ties case (4 distinct selection values over 256 columns), shared_outranks_routed, and a non-multiple-of-BLOCK tail. - free invariant holds: routed + shared weights sum to route_scale * global_scale (10.399998 .. 10.400002 vs an exact 10.400000). - packed mode: ids bit-match the plain path, weights bit-match `w.bfloat16()`, shared gammas bit-match. Non-regression at the default `shared_sink == 0`: 44 configs (11 shapes x 4 token counts) captured before and after are bit-identical in both weights and ids -- DeepSeek-V3 grouped routing, grouped + fused shared experts, ungrouped 256/384, Kimi-K2-class 896 experts (the BLOCK_N=1024 / num_warps=4 branch), all three scoring functions, tanh softcapping, renorm on/off, scale on/off. Device time for those callers is unchanged (2.13 -> 2.34 us at T=1, 28.64 -> 26.76 at T=8192, i.e. noise in both directions). That evidence is from XPU; the argument that it carries to CUDA rests on LOGSIGMOID_SINK living entirely behind the `EPILOGUE` / `SHARED_SINK` constexprs, which are 0 for every existing caller. Device us (profiler self-time) through `sigmoid_gate_topk_renorm`, same process, interleaved, min of 3 reps: T=1 T=8 T=64 T=512 T=4096 T=8192 tl.topk 58.81 58.93 62.22 181.17 641.59 1237.80 router 2.43 2.45 2.75 3.43 15.63 28.09 speedup 24.2x 24.0x 22.6x 52.8x 41.1x 44.1x In packed mode -- what InklingGate actually requests -- the old kernel costs 37-41% more while the router costs 0.4% more, so the speedup is 34.1x / 32.6x / 31.5x / 72.6x / 41.2x / 44.3x on the same token counts. Two caveats. This is reachable on XPU only together with the `.is_cuda` dispatch fix earlier in this branch; without it the CUDA JIT branch is selected first and still asserts. And the op is submission-bound: 49.6 us of host submit against 2.45 us of device time at T=1, so this kernel win is NOT visible end-to-end until the outputs are preallocated per layer or the decode step is captured into a graph. Do not read these ratios as end-to-end numbers. --- .../sglang/kernels/ops/moe/moe_fused_gate.py | 269 ++++++++++++++---- .../ops/moe/sigmoid_gate_topk_renorm.py | 25 ++ python/sglang/srt/environ.py | 3 + 3 files changed, 246 insertions(+), 51 deletions(-) diff --git a/python/sglang/kernels/ops/moe/moe_fused_gate.py b/python/sglang/kernels/ops/moe/moe_fused_gate.py index b60ecce2ec54..b9f9d6eb29b1 100644 --- a/python/sglang/kernels/ops/moe/moe_fused_gate.py +++ b/python/sglang/kernels/ops/moe/moe_fused_gate.py @@ -88,16 +88,19 @@ def moe_fused_gate_jit( @triton.jit def _router_triton_kernel( - scores_ptr, # [M, N] fp32, GEMM output (raw logits) + scores_ptr, # [M, N (+ SHARED_SINK)] fp32, GEMM output (raw logits) bias_ptr, # [N] fp32/fp16/bf16 (upcast to fp32 on load) - out_weights_ptr, # [M, K] fp32 - out_indices_ptr, # [M, K] int32 + out_weights_ptr, # [M, K] fp32 (LOGSIGMOID_SINK: [M, K_ROUTED]) + out_indices_ptr, # [M, K] int32 (LOGSIGMOID_SINK: [M, K_ROUTED]) + out_shared_ptr, # [M, SHARED_SINK] fp32 (LOGSIGMOID_SINK only) + out_packed_ptr, # [M, K_ROUTED] int32 (LOGSIGMOID_SINK + RETURN_PACKED only) + global_scale_ptr, # [1] fp32 (HAS_GLOBAL_SCALE only) M, routed_scaling_factor, moe_softcapping, N: tl.constexpr, - K: tl.constexpr, # total topk (includes fused shared experts) - K_ROUTED: tl.constexpr, # K - num_fused_shared_experts + K: tl.constexpr, # total topk (includes fused shared experts / sink columns) + K_ROUTED: tl.constexpr, # K - num_fused_shared_experts - SHARED_SINK BLOCK_M: tl.constexpr, # rows processed per program (row tiling) BLOCK_N: tl.constexpr, # >= N, power of 2 BLOCK_K: tl.constexpr, # >= K, power of 2 @@ -117,6 +120,16 @@ def _router_triton_kernel( stride_wk, stride_im, stride_ik, + # Epilogue parameterization. EPILOGUE 0 (SUM_NORM) is the historical + # behaviour and the only one any existing call site reaches. EPILOGUE 1 + # (LOGSIGMOID_SINK) is Inkling's gate: the last SHARED_SINK columns of the + # row are shared-expert sink logits that take no bias and never enter the + # top-k, but do join the normalizer, and the quantity normalized is the + # winner's RAW logit rather than its activated score. + EPILOGUE: tl.constexpr = 0, + SHARED_SINK: tl.constexpr = 0, + HAS_GLOBAL_SCALE: tl.constexpr = False, + RETURN_PACKED: tl.constexpr = False, ) -> None: # Row-tiled: each program handles BLOCK_M rows; all reductions run along the # expert (N) axis. Tiling rows keeps CTAs large enough to stay occupancy-bound @@ -206,51 +219,130 @@ def _router_triton_kernel( selected_vals = tl.zeros([BLOCK_M, BLOCK_K], dtype=tl.float32) selected_idx = tl.zeros([BLOCK_M, BLOCK_K], dtype=tl.int32) + # The quantity carried out of the top-k loop is not the quantity it ranks on. + # SUM_NORM normalizes the activated score, so it carries `activated`. + # LOGSIGMOID_SINK normalizes the RAW logit, so it carries `scores` -- ranking + # still happens on `biased`. Two quantities per column is what makes a plain + # "top-k and keep the winning score" kernel unusable for Inkling. + if EPILOGUE == 1: + carried = scores + else: + carried = activated + cur = biased # [BLOCK_M, BLOCK_N] for k in tl.static_range(K_ROUTED): max_val = tl.max(cur, axis=1)[:, None] # [BLOCK_M, 1] is_max = cur == max_val lane_id = tl.where(is_max, offs_n[None, :], N + 1) # lowest expert id wins ties win_lane = tl.min(lane_id, axis=1)[:, None].to(tl.int32) # [BLOCK_M, 1] - win_activated = tl.sum( - tl.where(offs_n[None, :] == win_lane, activated, 0.0), axis=1 + win_carried = tl.sum( + tl.where(offs_n[None, :] == win_lane, carried, 0.0), axis=1 )[:, None] # [BLOCK_M, 1] slot = offs_k[None, :] == k # [1, BLOCK_K] - selected_vals = tl.where(slot, win_activated, selected_vals) + selected_vals = tl.where(slot, win_carried, selected_vals) selected_idx = tl.where(slot, win_lane, selected_idx) cur = tl.where(offs_n[None, :] == win_lane, -float("inf"), cur) - routed_sum = tl.sum(tl.where(mask_k_routed[None, :], selected_vals, 0.0), axis=1)[ - :, None - ] # [BLOCK_M, 1] - - # Fill fused-shared-expert slots: weight = routed_sum / routed_scaling_factor, - # id = num_experts + (slot - K_ROUTED). - if K_ROUTED < K: - is_shared = (offs_k[None, :] >= K_ROUTED) & mask_k_total[None, :] - shared_weight = routed_sum / routed_scaling_factor # [BLOCK_M, 1] - shared_idx = (N + (offs_k - K_ROUTED)).to(tl.int32)[None, :] # [1, BLOCK_K] - selected_vals = tl.where(is_shared, shared_weight, selected_vals) - selected_idx = tl.where(is_shared, shared_idx, selected_idx) - - if USE_PDL: - tl.extra.cuda.gdc_launch_dependents() - - if RENORMALIZE: - norm = tl.where(routed_sum > 0.0, routed_sum, 1.0) # [BLOCK_M, 1] - selected_vals = selected_vals / norm - if APPLY_SCALE: + if EPILOGUE == 1: + # Sink columns N .. N + SHARED_SINK of the same row occupy slots + # K_ROUTED .. K. They take no bias (bias is [N], not [N + SHARED_SINK]) + # and never enter the top-k, but they do join the normalizer. + offs_s = offs_k - K_ROUTED # [BLOCK_K] + mask_s = (offs_s >= 0) & (offs_s < SHARED_SINK) + sink_raw = tl.load( + scores_ptr + + offs_m[:, None] * stride_sm + + (N + offs_s[None, :]) * stride_sn, + mask=mask_m[:, None] & mask_s[None, :], + other=0.0, + ).to(tl.float32) + active = tl.where(mask_s[None, :], sink_raw, selected_vals) + + # exp(logsigmoid(x) - logsumexp(logsigmoid(x))) over the K active slots. + # Algebraically sigmoid(x) / sum(sigmoid(x)), but fp32 sigmoid underflows + # to 0 below x ~ -104, so the explicit form divides 0/0 and returns NaN + # once every active logit is that negative. + # log(sigmoid(x)) = min(x, 0) - log1p(exp(-|x|)), exact for large |x|. + lp = tl.minimum(active, 0.0) - tl.log(1.0 + tl.exp(-tl.abs(active))) + lp = tl.where(mask_k_total[None, :], lp, float("-inf")) + e = tl.where( + mask_k_total[None, :], tl.exp(lp - tl.max(lp, axis=1)[:, None]), 0.0 + ) + selected_vals = e / tl.sum(e, axis=1, keep_dims=True) + # route_scale, then the model's scalar global_scale. RENORMALIZE / + # APPLY_SCALE do not apply: this epilogue always renormalizes and always + # scales (the sum over routed + sink weights is route_scale*global_scale). selected_vals = selected_vals * routed_scaling_factor + if HAS_GLOBAL_SCALE: + selected_vals = selected_vals * tl.load(global_scale_ptr) + + if USE_PDL: + tl.extra.cuda.gdc_launch_dependents() + + # Routed weights/ids are [M, K_ROUTED]; the sink gammas get their own + # [M, SHARED_SINK] output rather than trailing slots of the routed one. + store_mask = mask_m[:, None] & mask_k_routed[None, :] + if RETURN_PACKED: + weight_bits = ( + selected_vals.to(tl.bfloat16).to(tl.int16, bitcast=True).to(tl.int32) + ) + tl.store( + out_packed_ptr + offs_m[:, None] * K_ROUTED + offs_k[None, :], + (selected_idx << 16) | weight_bits, + mask=store_mask, + ) + else: + tl.store( + out_weights_ptr + + offs_m[:, None] * stride_wm + + offs_k[None, :] * stride_wk, + selected_vals, + mask=store_mask, + ) + tl.store( + out_indices_ptr + + offs_m[:, None] * stride_im + + offs_k[None, :] * stride_ik, + selected_idx, + mask=store_mask, + ) + tl.store( + out_shared_ptr + offs_m[:, None] * SHARED_SINK + offs_s[None, :], + selected_vals, + mask=mask_m[:, None] & mask_s[None, :], + ) + else: + routed_sum = tl.sum( + tl.where(mask_k_routed[None, :], selected_vals, 0.0), axis=1 + )[:, None] # [BLOCK_M, 1] - out_w_ptr = ( - out_weights_ptr + offs_m[:, None] * stride_wm + offs_k[None, :] * stride_wk - ) - out_i_ptr = ( - out_indices_ptr + offs_m[:, None] * stride_im + offs_k[None, :] * stride_ik - ) - store_mask = mask_m[:, None] & mask_k_total[None, :] - tl.store(out_w_ptr, selected_vals, mask=store_mask) - tl.store(out_i_ptr, selected_idx, mask=store_mask) + # Fill fused-shared-expert slots: weight = routed_sum / routed_scaling_factor, + # id = num_experts + (slot - K_ROUTED). + if K_ROUTED < K: + is_shared = (offs_k[None, :] >= K_ROUTED) & mask_k_total[None, :] + shared_weight = routed_sum / routed_scaling_factor # [BLOCK_M, 1] + shared_idx = (N + (offs_k - K_ROUTED)).to(tl.int32)[None, :] # [1, BLOCK_K] + selected_vals = tl.where(is_shared, shared_weight, selected_vals) + selected_idx = tl.where(is_shared, shared_idx, selected_idx) + + if USE_PDL: + tl.extra.cuda.gdc_launch_dependents() + + if RENORMALIZE: + norm = tl.where(routed_sum > 0.0, routed_sum, 1.0) # [BLOCK_M, 1] + selected_vals = selected_vals / norm + if APPLY_SCALE: + selected_vals = selected_vals * routed_scaling_factor + + out_w_ptr = ( + out_weights_ptr + offs_m[:, None] * stride_wm + offs_k[None, :] * stride_wk + ) + out_i_ptr = ( + out_indices_ptr + offs_m[:, None] * stride_im + offs_k[None, :] * stride_ik + ) + store_mask = mask_m[:, None] & mask_k_total[None, :] + tl.store(out_w_ptr, selected_vals, mask=store_mask) + tl.store(out_i_ptr, selected_idx, mask=store_mask) @debug_kernel_api @@ -266,7 +358,13 @@ def moe_fused_gate( moe_softcapping: float = 0.0, num_expert_group: int = 1, topk_group: int = 1, -) -> Tuple[torch.Tensor, torch.Tensor]: + shared_sink: int = 0, + global_scale: torch.Tensor | None = None, + return_packed: bool = False, +) -> ( + Tuple[torch.Tensor, torch.Tensor] + | Tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor, torch.Tensor | None] +): """Triton fused router: scoring + bias + topk + (optional) renorm/scale. Mirrors the semantics of :func:`moe_fused_gate_jit` (the CUDA JIT kernel). @@ -274,6 +372,17 @@ def moe_fused_gate( (per-group top-2-sum group scores, keep ``topk_group`` groups, then top-k within). The first argument is named ``scores`` (raw GEMM logits) to match the existing call sites. + + ``shared_sink > 0`` selects the Inkling gate epilogue: ``scores`` is + ``[M, num_experts + shared_sink]``, the trailing ``shared_sink`` columns are + shared-expert sink logits that take no bias and never enter the top-k but do + join the normalizer, the normalized quantity is the winner's RAW logit rather + than its activated score, and the weights are renormalized in log space (see + :mod:`sglang.kernels.ops.moe.sigmoid_gate_topk_renorm`). It returns a 4-tuple + ``(routed_weights, topk_indices, shared_weights, packed_topk)``. + + At the default ``shared_sink == 0`` every code path, launch geometry and + output tensor is unchanged and the return value is the historical 2-tuple. """ scoring_func_int = _SCORING_FUNC_MAP.get(scoring_func.lower()) assert scoring_func_int is not None, ( @@ -285,10 +394,12 @@ def moe_fused_gate( torch.bfloat16, ), "scores must be float32/float16/bfloat16" assert scores.ndim == 2, "scores must be 2D" + assert shared_sink >= 0, f"{shared_sink=} must be non-negative" if bias is None: assert scoring_func.lower() == "softmax", ( "bias is required for non-softmax routing" ) + assert shared_sink == 0, "shared_sink requires a bias" else: # The kernel loads the bias and upcasts it to fp32 in-register (see # _router_triton_kernel), so a non-fp32 bias (DeepSeek-V4 stores the @@ -299,10 +410,23 @@ def moe_fused_gate( torch.bfloat16, ), "bias must be float32/float16/bfloat16" assert bias.ndim == 1, "bias must be 1D" - assert scores.size(1) == bias.size(0), ( - "scores and bias must have same num_experts" + assert scores.size(1) == bias.size(0) + shared_sink, ( + "scores must be [M, num_experts + shared_sink] and bias [num_experts]" ) assert topk > num_fused_shared_experts, "topk must be > num_fused_shared_experts" + if shared_sink > 0: + # The sink columns are read directly off the score row, so only + # column-stride-1 is required -- Inkling's gate logits are a [T, 258] + # slice of a padded [T, 264] fp32 GEMM output, i.e. not contiguous but + # column-contiguous, and need no copy. + assert scores.stride(1) == 1, f"{scores.stride()=} needs column stride 1" + assert num_fused_shared_experts == 0, "shared_sink and fused shared experts" + assert num_expert_group <= 1, "shared_sink does not support grouped routing" + assert scoring_func.lower() == "sigmoid", "shared_sink is sigmoid-gate only" + assert moe_softcapping == 0.0, "shared_sink does not support softcapping" + else: + assert global_scale is None, "global_scale requires shared_sink > 0" + assert not return_packed, "return_packed requires shared_sink > 0" if routed_scaling_factor is None: routed_scaling_factor = 1.0 @@ -318,6 +442,7 @@ def moe_fused_gate( and num_fused_shared_experts == 0 and num_expert_group <= 1 and moe_softcapping == 0.0 + and shared_sink == 0 ): radix_args = ( scores, @@ -330,17 +455,50 @@ def moe_fused_gate( if moe_route_radix.covered(scores, bias, topk): return moe_route_radix.route_radix(*radix_args, sorted=False) - M, N = scores.shape - K = topk - K_routed = topk - num_fused_shared_experts + M = scores.shape[0] + # bias is [N] while the score row is [N + shared_sink], and bias may be None + # for plain softmax routing, so derive N from the row rather than the bias. + N = scores.shape[1] - shared_sink + if shared_sink > 0: + # Sink columns take slots K_routed .. K of the normalizer, so they widen + # K without widening the routed outputs. + K = topk + shared_sink + K_routed = topk + else: + K = topk + K_routed = topk - num_fused_shared_experts if num_expert_group > 1: assert N % num_expert_group == 0, "num_experts must be divisible by group count" assert 1 <= topk_group <= num_expert_group, "invalid topk_group" experts_per_group = N // num_expert_group BLOCK_G = triton.next_power_of_2(num_expert_group) - weights = torch.empty((M, K), dtype=torch.float32, device=scores.device) - indices = torch.empty((M, K), dtype=torch.int32, device=scores.device) + shared_weights = packed_topk = None + if shared_sink > 0: + shared_weights = torch.empty( + (M, shared_sink), dtype=torch.float32, device=scores.device + ) + # In packed mode the kernel writes only the packed tensor; the unused + # pointer args still need a valid (never-stored) address. + if return_packed: + packed_topk = torch.empty( + (M, K_routed), dtype=torch.int32, device=scores.device + ) + weights = indices = None + w_arg = i_arg = p_arg = packed_topk + else: + weights = torch.empty( + (M, K_routed), dtype=torch.float32, device=scores.device + ) + indices = torch.empty( + (M, K_routed), dtype=torch.int32, device=scores.device + ) + w_arg, i_arg, p_arg = weights, indices, indices + s_arg = shared_weights + else: + weights = torch.empty((M, K), dtype=torch.float32, device=scores.device) + indices = torch.empty((M, K), dtype=torch.int32, device=scores.device) + w_arg, i_arg, p_arg, s_arg = weights, indices, indices, indices BLOCK_N = triton.next_power_of_2(N) # 256 -> 256, 384 -> 512 BLOCK_K = triton.next_power_of_2(K) # 6 -> 8, 8 -> 8 @@ -359,8 +517,11 @@ def moe_fused_gate( _router_triton_kernel[grid]( scores, bias if bias is not None else scores, - weights, - indices, + w_arg, + i_arg, + s_arg, + p_arg, + global_scale if global_scale is not None else scores, M, float(routed_scaling_factor), float(moe_softcapping), @@ -382,11 +543,17 @@ def moe_fused_gate( USE_PDL=use_pdl, stride_sm=scores.stride(0), stride_sn=scores.stride(1), - stride_wm=weights.stride(0), - stride_wk=weights.stride(1), - stride_im=indices.stride(0), - stride_ik=indices.stride(1), + stride_wm=w_arg.stride(0), + stride_wk=w_arg.stride(1), + stride_im=i_arg.stride(0), + stride_ik=i_arg.stride(1), + EPILOGUE=1 if shared_sink > 0 else 0, + SHARED_SINK=shared_sink, + HAS_GLOBAL_SCALE=global_scale is not None, + RETURN_PACKED=bool(return_packed), num_warps=num_warps, **extra, ) + if shared_sink > 0: + return weights, indices, shared_weights, packed_topk return weights, indices diff --git a/python/sglang/kernels/ops/moe/sigmoid_gate_topk_renorm.py b/python/sglang/kernels/ops/moe/sigmoid_gate_topk_renorm.py index bb33a455bc5a..36f546352a5b 100644 --- a/python/sglang/kernels/ops/moe/sigmoid_gate_topk_renorm.py +++ b/python/sglang/kernels/ops/moe/sigmoid_gate_topk_renorm.py @@ -21,6 +21,7 @@ from sglang.kernels.ops.moe.inkling_gate_topk_renorm import ( inkling_gate_topk_renorm_v2, ) +from sglang.kernels.ops.moe.moe_fused_gate import moe_fused_gate from sglang.srt.environ import envs @@ -209,6 +210,30 @@ def sigmoid_gate_topk_renorm( enable_pdl=is_arch_support_pdl(), ) + # On non-CUDA devices prefer the unified router's LOGSIGMOID_SINK epilogue: it + # is the same op with the `tl.topk` / `tl.bitonic_merge` sort below replaced by + # k masked-max passes in registers, worth 22-46x on Intel Xe. CUDA is excluded + # deliberately -- the win is an Xe measurement and there is no NVIDIA A/B or + # test for this epilogue, so CUDA keeps the sort-based kernel when it misses + # the JIT gate above. + if ( + n_shared_experts > 0 + and not logits.is_cuda + and envs.SGLANG_OPT_USE_ROUTER_GATE_EPILOGUE.get() + ): + return moe_fused_gate( + logits, + bias, + topk=k, + scoring_func="sigmoid", + renormalize=False, # LOGSIGMOID_SINK always renormalizes + routed_scaling_factor=route_scale, + apply_routed_scaling_factor_on_output=False, # ditto for the scale + shared_sink=n_shared_experts, + global_scale=global_scale, + return_packed=return_packed_topk, + ) + shared_w = torch.empty( (M, n_shared_experts), dtype=logits.dtype, device=logits.device ) diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index a2bb9274fc19..8f1b532f23b3 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -1414,6 +1414,9 @@ class Envs: # Inside the fused gate: replace the cublas gate linear with the # expert-per-block GEMV JIT kernel at small token counts (GateGemvMode). SGLANG_OPT_GATE_GEMV_MODE = EnvInt(GateGemvMode.PAIR) + # Inside the fused gate on non-CUDA devices: use the unified router's + # LOGSIGMOID_SINK epilogue instead of the tl.topk kernel. Inert on CUDA. + SGLANG_OPT_USE_ROUTER_GATE_EPILOGUE = EnvBool(True) # Capture all multi-layer EAGLE draft-extend steps and the in-graph chain # rotation into ONE CUDA graph instead of one captured graph per step. SGLANG_ENABLE_SINGLE_CG_DRAFT = EnvBool(True) From 26ba18dc67d2a0ee9bf320c3009d45c1543d06c0 Mon Sep 17 00:00:00 2001 From: Juan Muneton Date: Tue, 25 Aug 2026 15:34:03 +0000 Subject: [PATCH 4/4] [xpu] Add XPU tests for the Inkling gate epilogue `sigmoid_gate_topk_renorm` had no test references anywhere in the tree, and `moe_fused_gate` was covered only by test_moe_fused_gate.py, which hardcodes DEVICE = "cuda" and registers register_cuda_ci only. So all three changes earlier in this branch shipped with no automated coverage on XPU. Adds test/registered/xpu/test_inkling_gate_epilogue.py: 9 cases / 33 subtests, ~20 s on an Arc Pro B60. Both dispatch branches (the unified router's LOGSIGMOID_SINK epilogue and the tl.topk fallback) are checked against an fp64 recompute of the gate contract, plus the free normalization invariant, packed mode, ties, a dominating shared sink, a non-multiple-of-BLOCK tail, and the requirement that moe_fused_gate is unchanged at the default shared_sink == 0. The branch is forced per case with `envs.SGLANG_OPT_USE_ROUTER_GATE_EPILOGUE.override(...)`. Per the register_xpu_ci BKM this is a dedicated file under test/registered/xpu/ registering XPU CI only, so no per-class device skip is needed. Deliberately no register_cuda_ci: CI collects by file, so adding it would make CUDA runners pick up these XPU-only cases. A CUDA-side counterpart for the shared _router_triton_kernel epilogue belongs in test/registered/kernels/ops/moe/test_moe_fused_gate.py instead. Each case was verified to fail without the fix it guards: - Reverting all three changes: 21 failed / 4 passed. - With only the dispatch fix applied, so the NaN is reachable: all_negative_90 and all_negative_200 fail on "NaN in routed weights" while all_negative_40 passes -- the intended discrimination. - Mutating the router epilogue to carry the activated score instead of the winner's raw logit (the mistake that "silently changes routing"): every router subtest fails against the oracle, tl.topk subtests still pass. Note the sum invariant alone does NOT catch this -- a wrong carry still yields a normalized set -- which is why the oracle check is the gate. Not covered here, and not coverable from an XPU runner: the ROCm half of the `.is_cuda and torch.version.hip is None` dispatch predicate, and the CUDA exclusion on the router epilogue. Both are argued in their own commit messages. --- .../xpu/test_inkling_gate_epilogue.py | 392 ++++++++++++++++++ 1 file changed, 392 insertions(+) create mode 100644 test/registered/xpu/test_inkling_gate_epilogue.py diff --git a/test/registered/xpu/test_inkling_gate_epilogue.py b/test/registered/xpu/test_inkling_gate_epilogue.py new file mode 100644 index 000000000000..8b1f6b207e44 --- /dev/null +++ b/test/registered/xpu/test_inkling_gate_epilogue.py @@ -0,0 +1,392 @@ +"""Inkling MoE gate epilogue on Intel XPU. + +Covers the three things that were broken or untested on XPU: + + A. Dispatch. `sigmoid_gate_topk_renorm` gated its CUDA JIT fast path on + `torch.version.hip is None` ("not ROCm"), which is also true on XPU, so it + selected a CUDA kernel and raised `AssertionError` on every call. The gate + is now `.is_cuda and torch.version.hip is None` -- ROCm reports `.is_cuda`, + so neither conjunct alone is right. Guarded by + `test_dispatch_reaches_a_triton_kernel`. + + B. Underflow. The fallback `tl.topk` kernel normalized as + `sigmoid(x) / sum(sigmoid(x))`. fp32 sigmoid flushes to zero below + x ~ -104, so all-negative active logits made that a 0/0 and every weight + came back NaN. Guarded by `test_underflow_logits_are_nan_free`. + + C. The `LOGSIGMOID_SINK` epilogue in the unified router, which replaces the + `tl.topk` sort with iterative masked-max, plus the requirement that the + router's default `SHARED_SINK == 0` behaviour is untouched. Guarded by the + oracle / invariant / packed / non-regression cases below. + +This file lives under `test/registered/xpu/` and registers XPU CI only, so no +per-class device skip is needed -- see the `register_xpu_ci` BKM. Do not add +`register_cuda_ci` here: CI collects by file, so that would make CUDA runners +pick up these XPU-only cases. A CUDA-side counterpart for the shared +`_router_triton_kernel` epilogue belongs in +`test/registered/kernels/ops/moe/test_moe_fused_gate.py`. + +The oracle is an fp64 recompute of the gate contract: + + sel_j = sigmoid(logits[:, j]) + bias[j] for j < n_routed # selection only + idx = topk(sel, k) # lowest id wins ties + act = logits[idx] ++ logits[n_routed:] # RAW logits, k + s + w = exp(logsigmoid(act) - logsumexp(logsigmoid(act))) * route_scale * global_scale +""" + +import unittest + +import torch +import torch.nn.functional as F + +from sglang.kernels.ops.moe.moe_fused_gate import moe_fused_gate +from sglang.kernels.ops.moe.sigmoid_gate_topk_renorm import sigmoid_gate_topk_renorm +from sglang.srt.environ import envs +from sglang.test.ci.ci_register import register_xpu_ci +from sglang.test.test_utils import CustomTestCase + +register_xpu_ci(est_time=30, suite="stage-b-test-1-gpu-xpu") + +DEVICE = "xpu" +N_ROUTED = 256 +N_SHARED = 2 +G_COLS = N_ROUTED + N_SHARED # 258 +G_PAD = 264 # the production gate GEMM writes a [T, 264] fp32 row pitch +TOPK = 6 +ROUTE_SCALE = 8.0 +GLOBAL_SCALE = 1.3 + +# `SGLANG_OPT_USE_ROUTER_GATE_EPILOGUE` selects the dispatch branch inside +# `sigmoid_gate_topk_renorm`: True (default) is the unified router's +# LOGSIGMOID_SINK epilogue, False is the older tl.topk kernel, which stays +# reachable as a fallback. Both must satisfy the contract, so both are exercised. + + +def _make_logits(tokens: int, seed: int): + """Production-shaped inputs: a [T, 264] pad sliced to [:, :258], stride (264, 1). + + The real gate logits are a non-contiguous but column-contiguous slice of a + padded fp32 GEMM output, so the kernels must not require contiguity. + """ + g = torch.Generator(device=DEVICE).manual_seed(seed) + pad = torch.randn((tokens, G_PAD), generator=g, device=DEVICE, dtype=torch.float32) + logits = pad[:, :G_COLS] + bias = ( + torch.randn((N_ROUTED,), generator=g, device=DEVICE, dtype=torch.float32) * 0.1 + ) + global_scale = torch.tensor([GLOBAL_SCALE], device=DEVICE, dtype=torch.float32) + return logits, bias, global_scale + + +def _oracle(logits, bias, global_scale, k=TOPK, s=N_SHARED): + """fp64 recompute of the contract. Returns (routed_w, shared_w, indices).""" + l64 = logits.double() + sel = torch.sigmoid(l64[:, :-s]) + bias.double() + idx = sel.topk(k, dim=-1).indices + act = torch.cat([l64[:, :-s].gather(-1, idx), l64[:, -s:]], dim=-1) + lp = F.logsigmoid(act) + w = torch.exp(lp - torch.logsumexp(lp, -1, keepdim=True)) + w = w * ROUTE_SCALE * global_scale.double() + return w[:, :k].contiguous(), w[:, k:].contiguous(), idx.to(torch.int32) + + +def _gate(logits, bias, global_scale, *, use_router: bool, packed: bool = False): + """Call the production entry point with one dispatch branch forced.""" + with envs.SGLANG_OPT_USE_ROUTER_GATE_EPILOGUE.override(use_router): + return sigmoid_gate_topk_renorm( + logits, + TOPK, + N_SHARED, + ROUTE_SCALE, + global_scale, + bias, + return_packed_topk=packed, + ) + + +class TestInklingGateEpilogue(CustomTestCase): + # --- A: the dispatch predicate ------------------------------------------- + + def test_dispatch_reaches_a_triton_kernel(self): + """XPU must not select the CUDA JIT path (it asserts on `is_cuda`). + + This is the regression guard for `torch.version.hip is None` -> + `logits.is_cuda and torch.version.hip is None`. Every JIT gate other + than the device check is satisfied by these inputs (k == 6, + n_shared == 2, G == 258, stride % 8 == 0, ptr % 32 == 0), so before the + fix this raised at every token count. The ROCm half of that predicate + cannot be covered from here; it is argued in the commit message. + """ + for tokens in (1, 4, 8, 512): + logits, bias, global_scale = _make_logits(tokens, seed=tokens) + self.assertEqual(logits.stride(1), 1) + self.assertEqual(logits.stride(0) % 8, 0) + routed_w, indices, shared_w, packed = sigmoid_gate_topk_renorm( + logits, TOPK, N_SHARED, ROUTE_SCALE, global_scale, bias + ) + self.assertEqual(routed_w.shape, (tokens, TOPK)) + self.assertEqual(indices.shape, (tokens, TOPK)) + self.assertEqual(shared_w.shape, (tokens, N_SHARED)) + self.assertIsNone(packed) + self.assertTrue(torch.isfinite(routed_w).all(), f"{tokens=}") + self.assertTrue(torch.isfinite(shared_w).all(), f"{tokens=}") + + # --- B and C: agreement with the fp64 oracle ------------------------------ + + def _check_against_oracle(self, logits, bias, global_scale, *, use_router, msg): + routed_w, indices, shared_w, _ = _gate( + logits, bias, global_scale, use_router=use_router + ) + o_routed, o_shared, o_idx = _oracle(logits, bias, global_scale) + + self.assertFalse(torch.isnan(routed_w).any(), f"{msg}: NaN in routed weights") + self.assertFalse(torch.isnan(shared_w).any(), f"{msg}: NaN in shared weights") + + # fp32 top-k selection is a genuine knife edge: two selection scores can + # be bit-identical in fp32 and be ordered differently by two sigmoid + # implementations. Judge the weights only where selection agrees with + # the oracle, and require that no token disagrees on these inputs. + agree = (indices.int() == o_idx).all(dim=-1) + self.assertEqual( + int((~agree).sum()), 0, f"{msg}: top-k indices differ from the fp64 oracle" + ) + torch.testing.assert_close( + routed_w[agree].float(), + o_routed[agree].float(), + rtol=2e-3, + atol=2e-3, + msg=f"{msg}: routed weights", + ) + torch.testing.assert_close( + shared_w[agree].float(), + o_shared[agree].float(), + rtol=2e-3, + atol=2e-3, + msg=f"{msg}: shared weights", + ) + + def test_matches_fp64_oracle(self): + """Both dispatch branches must reproduce the contract. + + Token counts include 1 (decode), a non-multiple-of-BLOCK tail (37), and + a prefill-sized batch. + """ + for use_router in (True, False): + branch = "router" if use_router else "tl.topk" + for tokens in (1, 8, 37, 512): + logits, bias, global_scale = _make_logits(tokens, seed=tokens) + with self.subTest(branch=branch, tokens=tokens): + self._check_against_oracle( + logits, + bias, + global_scale, + use_router=use_router, + msg=f"{branch} T={tokens}", + ) + + def test_underflow_logits_are_nan_free(self): + """All-negative logits must not produce NaN (change B). + + fp32 sigmoid goes subnormal below x ~ -87 and flushes to zero below + x ~ -104. At base -90 and -200 every one of the k + s active logits is + in that region, so `sigmoid(x) / sum(sigmoid(x))` divides 0/0. Base -40 + is included as the control that was always fine. + """ + for use_router in (True, False): + branch = "router" if use_router else "tl.topk" + for base in (-40.0, -90.0, -200.0): + g = torch.Generator(device=DEVICE).manual_seed(7) + pad = torch.full((16, G_PAD), base, device=DEVICE, dtype=torch.float32) + pad[:, :G_COLS] += ( + torch.randn( + (16, G_COLS), generator=g, device=DEVICE, dtype=torch.float32 + ) + * 0.5 + ) + bias = ( + torch.randn((N_ROUTED,), generator=g, device=DEVICE) * 0.1 + ).float() + global_scale = torch.tensor( + [GLOBAL_SCALE], device=DEVICE, dtype=torch.float32 + ) + with self.subTest(branch=branch, base=base): + self._check_against_oracle( + pad[:, :G_COLS], + bias, + global_scale, + use_router=use_router, + msg=f"{branch} all_negative_{int(abs(base))}", + ) + + def test_shared_sink_outranking_every_routed_expert(self): + """A sink column can dominate the normalizer; it still never enters top-k.""" + g = torch.Generator(device=DEVICE).manual_seed(11) + pad = ( + torch.randn((16, G_PAD), generator=g, device=DEVICE, dtype=torch.float32) + - 3.0 + ) + pad[:, N_ROUTED:G_COLS] = 10.0 + bias = (torch.randn((N_ROUTED,), generator=g, device=DEVICE) * 0.1).float() + global_scale = torch.tensor([GLOBAL_SCALE], device=DEVICE, dtype=torch.float32) + for use_router in (True, False): + with self.subTest(branch="router" if use_router else "tl.topk"): + self._check_against_oracle( + pad[:, :G_COLS], + bias, + global_scale, + use_router=use_router, + msg="shared_outranks_routed", + ) + _, indices, _, _ = _gate( + pad[:, :G_COLS], bias, global_scale, use_router=use_router + ) + # Sink columns live at N_ROUTED.., outside the routed id range. + self.assertTrue((indices < N_ROUTED).all(), "sink id leaked into top-k") + + def test_ties_do_not_produce_nan_or_break_the_invariant(self): + """Only 4 distinct selection values over 256 columns, bias identically 0. + + Indices are deliberately not compared: with exact ties the winning set is + implementation-defined. What must hold is that the weights are finite and + still normalized. + """ + pad = torch.zeros((16, G_PAD), device=DEVICE, dtype=torch.float32) + vals = torch.tensor([2.0, 2.0, 1.0, 0.5], device=DEVICE) + pad[:, :N_ROUTED] = vals.repeat(N_ROUTED // 4) + pad[:, N_ROUTED:G_COLS] = 0.25 + bias = torch.zeros(N_ROUTED, device=DEVICE, dtype=torch.float32) + global_scale = torch.tensor([GLOBAL_SCALE], device=DEVICE, dtype=torch.float32) + for use_router in (True, False): + routed_w, _, shared_w, _ = _gate( + pad[:, :G_COLS], bias, global_scale, use_router=use_router + ) + with self.subTest(branch="router" if use_router else "tl.topk"): + self.assertTrue(torch.isfinite(routed_w).all()) + self.assertTrue(torch.isfinite(shared_w).all()) + total = routed_w.float().sum(-1) + shared_w.float().sum(-1) + torch.testing.assert_close( + total, + torch.full_like(total, ROUTE_SCALE * GLOBAL_SCALE), + rtol=2e-3, + atol=2e-3, + ) + + def test_weights_sum_to_route_scale_times_global_scale(self): + """Free invariant: the k routed plus s shared weights are a normalized set.""" + for use_router in (True, False): + for tokens in (1, 37, 512): + logits, bias, global_scale = _make_logits(tokens, seed=tokens + 100) + routed_w, _, shared_w, _ = _gate( + logits, bias, global_scale, use_router=use_router + ) + total = routed_w.float().sum(-1) + shared_w.float().sum(-1) + with self.subTest( + branch="router" if use_router else "tl.topk", tokens=tokens + ): + torch.testing.assert_close( + total, + torch.full_like(total, ROUTE_SCALE * GLOBAL_SCALE), + rtol=2e-3, + atol=2e-3, + ) + + def test_packed_matches_plain(self): + """Packed mode is what InklingGate.emit_packed_topk requests. + + The packed int32 is `(expert_id << 16) | bf16_bits(weight)`, so it must + carry the same ids and the bf16 rounding of the same weights. + """ + for use_router in (True, False): + for tokens in (1, 37, 512): + logits, bias, global_scale = _make_logits(tokens, seed=tokens + 200) + routed_w, indices, shared_w, packed_none = _gate( + logits, bias, global_scale, use_router=use_router + ) + p_w, p_idx, p_shared, packed = _gate( + logits, bias, global_scale, use_router=use_router, packed=True + ) + with self.subTest( + branch="router" if use_router else "tl.topk", tokens=tokens + ): + self.assertIsNone(packed_none) + self.assertIsNone(p_w) + self.assertIsNone(p_idx) + self.assertIsNotNone(packed) + self.assertEqual(packed.shape, (tokens, TOPK)) + self.assertEqual(packed.dtype, torch.int32) + unpacked_idx = (packed >> 16).to(torch.int32) + unpacked_w = ( + (packed << 16 >> 16).to(torch.int16).view(torch.bfloat16) + ) + self.assertTrue(torch.equal(unpacked_idx, indices.to(torch.int32))) + self.assertTrue( + torch.equal(unpacked_w.float(), routed_w.bfloat16().float()) + ) + self.assertTrue(torch.equal(p_shared, shared_w)) + + # --- C: the router's default epilogue must be untouched ------------------- + + def test_moe_fused_gate_default_epilogue_unchanged(self): + """`shared_sink == 0` must keep the historical SUM_NORM behaviour. + + The Inkling epilogue is behind the `SHARED_SINK` / `EPILOGUE` constexprs, + so every pre-existing caller must be unaffected: a 2-tuple return, and + weights that still match `activated / sum(activated) * scale`. + """ + for num_experts, topk, rsf, apply_scale in ( + (256, 8, 2.5, True), + (256, 6, 8.0, True), + (384, 8, 1.0, False), + ): + g = torch.Generator(device=DEVICE).manual_seed(num_experts + topk) + scores = torch.randn( + (64, num_experts), generator=g, device=DEVICE, dtype=torch.float32 + ) + bias = ( + torch.randn((num_experts,), generator=g, device=DEVICE) * 0.1 + ).float() + + out = moe_fused_gate( + scores, + bias, + topk=topk, + scoring_func="sigmoid", + renormalize=True, + routed_scaling_factor=rsf, + apply_routed_scaling_factor_on_output=apply_scale, + ) + with self.subTest(num_experts=num_experts, topk=topk): + # Still a 2-tuple: the third/fourth outputs appear only when + # shared_sink > 0. + self.assertEqual(len(out), 2) + weights, indices = out + + activated = torch.sigmoid(scores.double()) + ref_idx = (activated + bias.double()).topk(topk, dim=-1).indices + sel = activated.gather(-1, ref_idx) + ref_w = sel / sel.sum(-1, keepdim=True) + if apply_scale: + ref_w = ref_w * rsf + self.assertTrue(torch.equal(indices.int(), ref_idx.to(torch.int32))) + torch.testing.assert_close( + weights.float(), ref_w.float(), rtol=2e-3, atol=2e-3 + ) + + def test_shared_sink_kwargs_rejected_without_shared_sink(self): + """The new kwargs are meaningless at `shared_sink == 0` and must not be silent.""" + g = torch.Generator(device=DEVICE).manual_seed(3) + scores = torch.randn((8, 256), generator=g, device=DEVICE, dtype=torch.float32) + bias = (torch.randn((256,), generator=g, device=DEVICE) * 0.1).float() + global_scale = torch.tensor([GLOBAL_SCALE], device=DEVICE, dtype=torch.float32) + with self.assertRaises(AssertionError): + moe_fused_gate(scores, bias, topk=6, global_scale=global_scale) + with self.assertRaises(AssertionError): + moe_fused_gate(scores, bias, topk=6, return_packed=True) + # bias covers the routed experts only, so a [M, N] score row with + # shared_sink=2 is a shape error rather than a silent misread. + with self.assertRaises(AssertionError): + moe_fused_gate(scores, bias, topk=6, shared_sink=N_SHARED) + + +if __name__ == "__main__": + unittest.main()