diff --git a/python/tokenspeed/runtime/models/kimi_k3.py b/python/tokenspeed/runtime/models/kimi_k3.py index 4d0806ad1f..f8075f1d58 100644 --- a/python/tokenspeed/runtime/models/kimi_k3.py +++ b/python/tokenspeed/runtime/models/kimi_k3.py @@ -1853,7 +1853,7 @@ def forward( max_num_tokens_per_gpu = num_global_tokens if self.native_latent_moe is not None: - if self._use_fused_decode_pipeline and hidden_states.shape[0] == 1: + if self._use_fused_decode_pipeline and 0 < hidden_states.shape[0] <= 4: output = self._forward_fused_decode_pipeline(hidden_states, prefix_sum) else: output = self.native_latent_moe( diff --git a/test/runtime/test_kimi_k3_config.py b/test/runtime/test_kimi_k3_config.py index 81d5ecb3d5..e577a0dbfb 100644 --- a/test/runtime/test_kimi_k3_config.py +++ b/test/runtime/test_kimi_k3_config.py @@ -592,6 +592,40 @@ def __init__(self, **kwargs): routed_input + 1, ) + def test_native_kimi_moe_zero_tokens_bypass_fused_pipeline(self): + from tokenspeed.runtime.models.kimi_k3 import KimiLinearMoE + + hidden_states = torch.empty(0, 64) + prefix_sum = torch.empty_like(hidden_states) + native_latent_moe = mock.Mock(return_value=prefix_sum) + fused_pipeline = mock.Mock( + side_effect=AssertionError("zero tokens must bypass the fused pipeline") + ) + layer = SimpleNamespace( + _gather_dp_tokens_for_moe=False, + native_latent_moe=native_latent_moe, + _use_fused_decode_pipeline=True, + _forward_fused_decode_pipeline=fused_pipeline, + ) + + output = KimiLinearMoE.forward( + layer, + hidden_states, + prefix_sum, + num_global_tokens=0, + max_num_tokens_per_gpu=0, + ) + + torch.testing.assert_close(output, prefix_sum) + self.assertEqual(tuple(output.shape), (0, 64)) + fused_pipeline.assert_not_called() + native_latent_moe.assert_called_once_with( + hidden_states, + num_global_tokens=0, + max_num_tokens_per_gpu=0, + prefix_sum=prefix_sum, + ) + def test_cross_dp_ep_gather_uses_dp_group_and_returns_local_offset(self): from tokenspeed.runtime.models.kimi_k3 import KimiLinearMoE diff --git a/tokenspeed-kernel-amd/python/tokenspeed_kernel_amd/ops/gfx950/gemm/fp16/rmsnorm_linear_add.py b/tokenspeed-kernel-amd/python/tokenspeed_kernel_amd/ops/gfx950/gemm/fp16/rmsnorm_linear_add.py index ce07aefc38..ed421b2e2e 100644 --- a/tokenspeed-kernel-amd/python/tokenspeed_kernel_amd/ops/gfx950/gemm/fp16/rmsnorm_linear_add.py +++ b/tokenspeed-kernel-amd/python/tokenspeed_kernel_amd/ops/gfx950/gemm/fp16/rmsnorm_linear_add.py @@ -25,7 +25,6 @@ import torch from tokenspeed_kernel_amd._triton import gl, gluon -_HIDDEN = gl.constexpr(7168) _LATENT = gl.constexpr(3584) _BLOCK_N = gl.constexpr(32) _BLOCK_K = gl.constexpr(512) @@ -41,9 +40,14 @@ def _rmsnorm_linear_add_kernel( residual_ptr, shared_ptr, output_ptr, + latent_stride_m, + residual_stride_m, + shared_stride_m, + output_stride_m, eps, ): - pid_n = gl.program_id(0) + pid_m = gl.program_id(0) + pid_n = gl.program_id(1) layout: gl.constexpr = gl.BlockedLayout( [1, _BLOCK_K // _LANES], [1, _LANES], @@ -60,7 +64,7 @@ def _rmsnorm_linear_add_kernel( k_mask = offs_k < _LATENT latent = gl.amd.cdna4.buffer_load( latent_ptr, - offs_k.to(gl.int32), + (pid_m * latent_stride_m + offs_k).to(gl.int32), mask=k_mask, other=0.0, ).to(gl.float32) @@ -73,7 +77,7 @@ def _rmsnorm_linear_add_kernel( k_mask = offs_k < _LATENT latent = gl.amd.cdna4.buffer_load( latent_ptr, - offs_k.to(gl.int32), + (pid_m * latent_stride_m + offs_k).to(gl.int32), mask=k_mask, other=0.0, ).to(gl.float32) @@ -101,9 +105,16 @@ def _rmsnorm_linear_add_kernel( # Match the materialized BF16 projection before the residual additions. acc = acc.to(gl.bfloat16).to(gl.float32) - acc += gl.amd.cdna4.buffer_load(residual_ptr, offs_n.to(gl.int32)).to(gl.float32) - acc += gl.amd.cdna4.buffer_load(shared_ptr, offs_n.to(gl.int32)).to(gl.float32) - gl.store(output_ptr + offs_n, acc) + residual_offset = pid_m * residual_stride_m + offs_n + shared_offset = pid_m * shared_stride_m + offs_n + output_offset = pid_m * output_stride_m + offs_n + acc += gl.amd.cdna4.buffer_load(residual_ptr, residual_offset.to(gl.int32)).to( + gl.float32 + ) + acc += gl.amd.cdna4.buffer_load(shared_ptr, shared_offset.to(gl.int32)).to( + gl.float32 + ) + gl.store(output_ptr + output_offset, acc) def gluon_rmsnorm_linear_add_gfx950( @@ -118,13 +129,18 @@ def gluon_rmsnorm_linear_add_gfx950( ) -> torch.Tensor: if out is None: out = torch.empty_like(residual) - _rmsnorm_linear_add_kernel[(_HIDDEN // _BLOCK_N,)]( + m = latent.shape[0] + _rmsnorm_linear_add_kernel[(m, projection_weight.shape[0] // _BLOCK_N)]( latent, norm_weight, projection_weight, residual, shared, out, + latent.stride(0), + residual.stride(0), + shared.stride(0), + out.stride(0), float(eps), num_warps=8, num_stages=1, diff --git a/tokenspeed-kernel-amd/python/tokenspeed_kernel_amd/ops/gfx950/moe/mxfp4/routing.py b/tokenspeed-kernel-amd/python/tokenspeed_kernel_amd/ops/gfx950/moe/mxfp4/routing.py index 3629f2d4e0..69dfaf443c 100644 --- a/tokenspeed-kernel-amd/python/tokenspeed_kernel_amd/ops/gfx950/moe/mxfp4/routing.py +++ b/tokenspeed-kernel-amd/python/tokenspeed_kernel_amd/ops/gfx950/moe/mxfp4/routing.py @@ -317,7 +317,7 @@ def invoke_softmax_topk_route_gluon( (M, topk), dtype=torch.float32, device=router_logits.device ) bias = correction_bias if correction_bias is not None else topk_weights - nw = 1 if M <= 2 else 4 + nw = min(_next_pow2(M), 4) _softmax_topk_route_gluon_kernel[(1,)]( router_logits, bias, @@ -358,7 +358,7 @@ def _launch_sigmoid_bias_topk_route_gluon( topk_weights = torch.empty( (M, topk), dtype=torch.float32, device=router_logits.device ) - nw = 1 if M <= 2 else 4 + nw = 1 if M == 1 else min(max(_next_pow2(M), 4), 8) _sigmoid_bias_topk_route_gluon_kernel[(1,)]( router_logits, correction_bias, diff --git a/tokenspeed-kernel-amd/python/tokenspeed_kernel_amd/ops/gfx950/moe/mxfp4/situ_decode.py b/tokenspeed-kernel-amd/python/tokenspeed_kernel_amd/ops/gfx950/moe/mxfp4/situ_decode.py index 309468afd9..d3beea9e56 100644 --- a/tokenspeed-kernel-amd/python/tokenspeed_kernel_amd/ops/gfx950/moe/mxfp4/situ_decode.py +++ b/tokenspeed-kernel-amd/python/tokenspeed_kernel_amd/ops/gfx950/moe/mxfp4/situ_decode.py @@ -51,9 +51,16 @@ WARP_DECODE_STAGE1_BLOCK_KB = 1024 WARP_DECODE_STAGE1_NUM_WARPS = 4 WARP_DECODE_STAGE2_BLOCK_N = 8 +WARP_DECODE_STAGE2_BATCHED_BLOCK_N = 4 WARP_DECODE_STAGE2_BLOCK_KB = 512 WARP_DECODE_STAGE2_NUM_WARPS = 8 -WARP_DECODE_STAGE2_M4_NUM_WARPS = 4 +WARP_DECODE_STAGE2_BATCHED_NUM_WARPS = 4 +WARP_DECODE_TP_STAGE1_BLOCK_N = 2 +WARP_DECODE_TP_STAGE1_NUM_WARPS = 1 +WARP_DECODE_TP_STAGE2_BLOCK_N = 8 +WARP_DECODE_TP_STAGE2_BLOCK_KB = 256 +WARP_DECODE_TP_STAGE2_NUM_WARPS = 1 +WARP_DECODE_TP_TOPK_PER_CTA = 1 _MIN_WARP_DECODE_BLOCK_KB = 128 _KIMI3_SHARED_K = gl.constexpr(768) _KIMI3_SHARED_BLOCK_K = gl.constexpr(512) @@ -290,6 +297,10 @@ def _stage2_a16w4_warp_gemv_combine( stride_ids, stride_twm, stride_tws, + stride_sim, + stride_sik, + stride_som, + stride_son, FUSE_SHARED_DOWN: gl.constexpr, NUM_ROUTED_PROGRAMS: gl.constexpr, TOP_K: gl.constexpr, @@ -297,13 +308,18 @@ def _stage2_a16w4_warp_gemv_combine( NUM_LOCAL_EXPERTS: gl.constexpr, LINEAR_WEIGHTS: gl.constexpr, NUM_PID_N: gl.constexpr, + NUM_TOPK_GROUPS: gl.constexpr, + NUM_SHARED_PID_N: gl.constexpr, BLOCK_N: gl.constexpr, BLOCK_KB: gl.constexpr, NUM_WARPS: gl.constexpr, + MASK_K_TAIL: gl.constexpr, ): pid = gl.program_id(0) if FUSE_SHARED_DOWN and pid >= NUM_ROUTED_PROGRAMS: - shared_pid_n = pid - NUM_ROUTED_PROGRAMS + shared_pid = pid - NUM_ROUTED_PROGRAMS + shared_token = shared_pid // NUM_SHARED_PID_N + shared_pid_n = shared_pid % NUM_SHARED_PID_N shared_layout: gl.constexpr = gl.BlockedLayout( [1, _KIMI3_SHARED_BLOCK_K // _LANES], [1, _LANES], @@ -323,7 +339,9 @@ def _stage2_a16w4_warp_gemv_combine( shared_k_valid = shared_offs_k < _KIMI3_SHARED_K shared_input = gl.amd.cdna4.buffer_load( ptr=shared_input_ptr, - offsets=shared_offs_k.to(gl.int32), + offsets=(shared_token * stride_sim + shared_offs_k * stride_sik).to( + gl.int32 + ), mask=shared_k_valid, other=0.0, ).to(gl.float32) @@ -345,13 +363,15 @@ def _stage2_a16w4_warp_gemv_combine( axis=1, ) gl.store( - shared_out_ptr + shared_offs_n, + shared_out_ptr + shared_token * stride_som + shared_offs_n * stride_son, shared_acc.to(shared_out_ptr.dtype.element_ty), ) return - token = pid // NUM_PID_N - pid_n = pid % NUM_PID_N + token = pid // (NUM_TOPK_GROUPS * NUM_PID_N) + token_pid = pid % (NUM_TOPK_GROUPS * NUM_PID_N) + topk_group = token_pid // NUM_PID_N + pid_n = token_pid % NUM_PID_N layout: gl.constexpr = gl.BlockedLayout( [(BLOCK_N + NUM_WARPS - 1) // NUM_WARPS, BLOCK_KB // _LANES], [1, _LANES], @@ -373,7 +393,12 @@ def _stage2_a16w4_warp_gemv_combine( packed_k = intermediate_dim // 2 acc = gl.zeros([BLOCK_N], gl.float32, expanded_n_layout) - for slot in gl.static_range(0, TOP_K): + for group_slot in gl.static_range( + 0, (TOP_K + NUM_TOPK_GROUPS - 1) // NUM_TOPK_GROUPS + ): + slot = ( + topk_group * ((TOP_K + NUM_TOPK_GROUPS - 1) // NUM_TOPK_GROUPS) + group_slot + ) expert = ( gl.load(local_ids_ptr + token * stride_idm + slot * stride_ids) - EXPERT_START @@ -391,9 +416,19 @@ def _stage2_a16w4_warp_gemv_combine( expanded_k = 2 * kb0 + gl.arange( 0, 2 * BLOCK_KB, layout=expanded_k_layout ) + if MASK_K_TAIL: + packed_k_valid = offs_kb < packed_k + expanded_k_valid = expanded_k < intermediate_dim + else: + packed_k_valid = gl.full([BLOCK_KB], True, gl.int1, layout=k_layout) + expanded_k_valid = gl.full( + [2 * BLOCK_KB], True, gl.int1, layout=expanded_k_layout + ) inter = gl.amd.cdna4.buffer_load( ptr=inter_ptr, offsets=(inter_row + expanded_k * stride_ipk).to(gl.int32), + mask=expanded_k_valid, + other=0.0, ).to(gl.float32) w_offsets = ( w_expert @@ -417,12 +452,16 @@ def _stage2_a16w4_warp_gemv_combine( packed = gl.amd.cdna4.buffer_load( ptr=w2_ptr, offsets=w_offsets.to(gl.int32), + mask=packed_k_valid[None, :], + other=0, ) weight = gl.amd.cdna4.scaled_upcast( packed, gl.amd.cdna4.buffer_load( ptr=w2_scale_ptr, offsets=scale_offsets.to(gl.int32), + mask=expanded_k_valid[None, :], + other=0, ), gl.bfloat16, axis=1, @@ -432,12 +471,45 @@ def _stage2_a16w4_warp_gemv_combine( # Match the reference's BF16 W2 result before route weighting. acc += route_weight * route_acc.to(gl.bfloat16).to(gl.float32) + out_row = token * NUM_TOPK_GROUPS + topk_group gl.store( - out_ptr + token * stride_om + expanded_offs_n * stride_on, + out_ptr + out_row * stride_om + expanded_offs_n * stride_on, acc.to(out_ptr.dtype.element_ty), ) +@gluon.jit +def _reduce_topk_groups( + partial_ptr, + out_ptr, + hidden_dim, + stride_pm, + stride_pn, + stride_om, + stride_on, + NUM_GROUPS: gl.constexpr, + BLOCK_N: gl.constexpr, +): + pid = gl.program_id(0) + num_pid_n = gl.cdiv(hidden_dim, BLOCK_N) + token = pid // num_pid_n + pid_n = pid % num_pid_n + layout: gl.constexpr = gl.BlockedLayout([4], [64], [1], [0]) + offs_n = pid_n * BLOCK_N + gl.arange(0, BLOCK_N, layout=layout) + acc = gl.zeros([BLOCK_N], gl.float32, layout=layout) + for group in gl.static_range(0, NUM_GROUPS): + acc += gl.load( + partial_ptr + (token * NUM_GROUPS + group) * stride_pm + offs_n * stride_pn, + mask=offs_n < hidden_dim, + other=0.0, + ).to(gl.float32) + gl.store( + out_ptr + token * stride_om + offs_n * stride_on, + acc.to(out_ptr.dtype.element_ty), + mask=offs_n < hidden_dim, + ) + + def gluon_a16w4_situ_warp_decode_ep_gfx950( hidden_states: torch.Tensor, w13_weight: torch.Tensor, @@ -486,12 +558,11 @@ def gluon_a16w4_situ_warp_decode_ep_gfx950( raise ValueError( "Kimi K3 shared input and weight must be provided together" ) - if tuple(shared_input.shape) != (1, 768) or tuple(shared_weight.shape) != ( - 7168, - 768, - ): + if tuple(shared_input.shape) != (hidden_states.shape[0], 768) or tuple( + shared_weight.shape + ) != (7168, 768): raise ValueError( - "Kimi K3 shared down fusion requires [1, 768] input and " + "Kimi K3 shared down fusion requires [M, 768] input and " "[7168, 768] weight" ) if ( @@ -571,10 +642,10 @@ def gluon_a16w4_situ_warp_decode_ep_gfx950( w13_stride_n = w13_weight.stride(2) w2_stride_k = w2_weight.stride(1) w2_stride_n = w2_weight.stride(2) - if hidden_dim % 256 or intermediate_dim % 256: + if hidden_dim % 256 or intermediate_dim % 128: raise ValueError( - "gfx950 warp decode requires hidden and intermediate dimensions " - "divisible by 256" + "gfx950 warp decode requires hidden dimensions divisible by 256 " + "and intermediate dimensions divisible by 128" ) local_topk_ids = local_topk_ids.to(torch.int32) @@ -583,9 +654,14 @@ def gluon_a16w4_situ_warp_decode_ep_gfx950( dtype=torch.bfloat16, device=hidden_states.device, ) - stage1_block_n = WARP_DECODE_STAGE1_BLOCK_N + tp_local = num_experts == 896 and intermediate_dim == 384 and top_k == 16 + stage1_block_n = ( + WARP_DECODE_TP_STAGE1_BLOCK_N if tp_local else WARP_DECODE_STAGE1_BLOCK_N + ) stage1_block_kb = WARP_DECODE_STAGE1_BLOCK_KB - stage1_warps = WARP_DECODE_STAGE1_NUM_WARPS + stage1_warps = ( + WARP_DECODE_TP_STAGE1_NUM_WARPS if tp_local else WARP_DECODE_STAGE1_NUM_WARPS + ) if ( intermediate_dim % stage1_block_n or stage1_block_kb < _MIN_WARP_DECODE_BLOCK_KB @@ -636,45 +712,72 @@ def gluon_a16w4_situ_warp_decode_ep_gfx950( raise ValueError("routed output must match the hidden-state shape and dtype") if not out.is_contiguous() or out.device != hidden_states.device: raise ValueError("routed output must be contiguous and colocated") - stage2_block_n = WARP_DECODE_STAGE2_BLOCK_N + # Finer output tiles keep work balanced when local route counts differ. + stage2_block_n = ( + WARP_DECODE_TP_STAGE2_BLOCK_N + if tp_local + else ( + WARP_DECODE_STAGE2_BATCHED_BLOCK_N + if num_tokens in (2, 4) + else WARP_DECODE_STAGE2_BLOCK_N + ) + ) packed_intermediate = intermediate_dim // 2 - stage2_block_kb = _largest_exact_block_kb( - packed_intermediate, - WARP_DECODE_STAGE2_BLOCK_KB, + stage2_block_kb = ( + WARP_DECODE_TP_STAGE2_BLOCK_KB + if tp_local + else _largest_exact_block_kb( + packed_intermediate, + WARP_DECODE_STAGE2_BLOCK_KB, + ) ) - # Smaller shapes and the M=1 joint path remain faster with eight waves. + # The M=1 joint path remains faster with eight waves. stage2_warps = ( - WARP_DECODE_STAGE2_M4_NUM_WARPS - if num_tokens == 4 and not fuse_shared_down - else WARP_DECODE_STAGE2_NUM_WARPS + WARP_DECODE_TP_STAGE2_NUM_WARPS + if tp_local + else ( + WARP_DECODE_STAGE2_BATCHED_NUM_WARPS + if num_tokens in (2, 4) + else WARP_DECODE_STAGE2_NUM_WARPS + ) ) - if hidden_dim % stage2_block_n or packed_intermediate % stage2_block_kb: - raise ValueError("unmasked stage2 requires exact N and packed-K tiles") - stage2_grid = num_tokens * triton.cdiv(hidden_dim, stage2_block_n) + if hidden_dim % stage2_block_n: + raise ValueError("stage2 requires exact output-column tiles") + num_topk_groups = triton.cdiv(top_k, WARP_DECODE_TP_TOPK_PER_CTA) if tp_local else 1 + stage2_grid = num_tokens * num_topk_groups * triton.cdiv(hidden_dim, stage2_block_n) + stage2_out = out + if num_topk_groups > 1: + stage2_out = torch.empty( + (num_tokens * num_topk_groups, hidden_dim), + dtype=torch.float32, + device=hidden_states.device, + ) if fuse_shared_down: if shared_out is None: shared_out = torch.empty( - (1, 7168), dtype=torch.bfloat16, device=hidden_states.device + (num_tokens, 7168), + dtype=torch.bfloat16, + device=hidden_states.device, ) if ( - shared_out.shape != (1, 7168) + shared_out.shape != (num_tokens, 7168) or shared_out.dtype != torch.bfloat16 or not shared_out.is_contiguous() or shared_out.device != hidden_states.device ): - raise ValueError("shared output must be contiguous BF16 [1, 7168]") + raise ValueError("shared output must be contiguous BF16 [M, 7168]") else: if shared_out is not None: raise ValueError("shared output requires fused shared down") shared_out = out total_stage2_grid = stage2_grid + ( - triton.cdiv(7168, stage2_block_n) if fuse_shared_down else 0 + num_tokens * triton.cdiv(7168, stage2_block_n) if fuse_shared_down else 0 ) _stage2_a16w4_warp_gemv_combine[(total_stage2_grid,)]( inter, w2_weight, w2_scale, - out, + stage2_out, local_topk_ids, topk_weights, out if shared_input is None else shared_input, @@ -690,12 +793,16 @@ def gluon_a16w4_situ_warp_decode_ep_gfx950( w2_scale.stride(0), w2_scale.stride(1), w2_scale.stride(2), - out.stride(0), - out.stride(1), + stage2_out.stride(0), + stage2_out.stride(1), local_topk_ids.stride(0), local_topk_ids.stride(1), topk_weights.stride(0), topk_weights.stride(1), + shared_input.stride(0) if shared_input is not None else out.stride(0), + shared_input.stride(1) if shared_input is not None else out.stride(1), + shared_out.stride(0), + shared_out.stride(1), FUSE_SHARED_DOWN=fuse_shared_down, NUM_ROUTED_PROGRAMS=stage2_grid, TOP_K=top_k, @@ -703,11 +810,27 @@ def gluon_a16w4_situ_warp_decode_ep_gfx950( NUM_LOCAL_EXPERTS=num_experts, LINEAR_WEIGHTS=linear_weights, NUM_PID_N=hidden_dim // stage2_block_n, + NUM_TOPK_GROUPS=num_topk_groups, + NUM_SHARED_PID_N=7168 // stage2_block_n, BLOCK_N=stage2_block_n, BLOCK_KB=stage2_block_kb, NUM_WARPS=stage2_warps, + MASK_K_TAIL=packed_intermediate % stage2_block_kb != 0, num_warps=stage2_warps, ) + if num_topk_groups > 1: + _reduce_topk_groups[(num_tokens * triton.cdiv(hidden_dim, 256),)]( + stage2_out, + out, + hidden_dim, + stage2_out.stride(0), + stage2_out.stride(1), + out.stride(0), + out.stride(1), + NUM_GROUPS=num_topk_groups, + BLOCK_N=256, + num_warps=1, + ) if fuse_shared_down: return out, shared_out return out diff --git a/tokenspeed-kernel/python/tokenspeed_kernel/ops/moe/gluon/latent_decode.py b/tokenspeed-kernel/python/tokenspeed_kernel/ops/moe/gluon/latent_decode.py index 60c35f2c01..25ba457279 100644 --- a/tokenspeed-kernel/python/tokenspeed_kernel/ops/moe/gluon/latent_decode.py +++ b/tokenspeed-kernel/python/tokenspeed_kernel/ops/moe/gluon/latent_decode.py @@ -46,7 +46,7 @@ ), priority=Priority.SPECIALIZED, traits={ - "tokens": frozenset({1}), + "tokens": frozenset({1, 2, 3, 4}), "latent_size": frozenset({3584}), "topk": frozenset({16}), "num_local_experts": frozenset({112}), diff --git a/tokenspeed-kernel/python/tokenspeed_kernel/ops/moe/gluon/sigmoid_topk.py b/tokenspeed-kernel/python/tokenspeed_kernel/ops/moe/gluon/sigmoid_topk.py index 924bbc11cf..8a805bea0d 100644 --- a/tokenspeed-kernel/python/tokenspeed_kernel/ops/moe/gluon/sigmoid_topk.py +++ b/tokenspeed-kernel/python/tokenspeed_kernel/ops/moe/gluon/sigmoid_topk.py @@ -49,9 +49,15 @@ def gluon_sigmoid_bias_topk_gfx950( normalize_topk_weights: bool, weights_dtype: torch.dtype = torch.float32, ) -> tuple[torch.Tensor, torch.Tensor]: + decode_supported = router_logits.shape[0] * topk <= 128 route = ( invoke_sigmoid_bias_topk_route_gluon - if router_logits.shape[0] * topk <= 128 + if decode_supported + and ( + router_logits.shape[0] == 1 + or router_logits.dtype != torch.float32 + or topk > 16 + ) else invoke_sigmoid_bias_topk_route_prefill_gluon ) topk_ids, topk_weights = route( diff --git a/tokenspeed-kernel/python/tokenspeed_kernel/ops/moe/latent_decode.py b/tokenspeed-kernel/python/tokenspeed_kernel/ops/moe/latent_decode.py index ac0789773d..93c36bffba 100644 --- a/tokenspeed-kernel/python/tokenspeed_kernel/ops/moe/latent_decode.py +++ b/tokenspeed-kernel/python/tokenspeed_kernel/ops/moe/latent_decode.py @@ -80,7 +80,7 @@ def latent_moe_decode_pipeline_available( *, topk: int, ) -> bool: - """Return whether joint routed/shared one-token decode is available. + """Return whether joint routed/shared small-batch decode is available. This construction-time probe owns backend, tensor-format, expert-plan, and exact-shape constraints needed by the orchestration-level optimization. @@ -119,26 +119,30 @@ def latent_moe_decode_pipeline_available( "linear_weights": True, "inputs_contiguous": True, } - for topk_weight_dtype in (torch.bfloat16, torch.float32): - signature = format_signature( - hidden_states=dense_tensor_format(routed_weight.dtype), - w13_weight=dense_tensor_format(w13_weight.dtype), - w13_scale=dense_tensor_format(w13_scale.dtype), - w2_weight=dense_tensor_format(w2_weight.dtype), - w2_scale=dense_tensor_format(w2_scale.dtype), - topk_weights=dense_tensor_format(topk_weight_dtype), - topk_ids=dense_tensor_format(torch.int32), - shared_input=dense_tensor_format(shared_gate_up_weight.dtype), - shared_weight=dense_tensor_format(shared_down_weight.dtype), - routed_out=dense_tensor_format(routed_weight.dtype), - shared_out=dense_tensor_format(shared_down_weight.dtype), - ) - try: - select_kernel("moe", "latent_expert_shared", signature, traits=traits) - except NoKernelFoundError: - continue - return True - return False + for tokens in range(1, 5): + traits["tokens"] = tokens + for topk_weight_dtype in (torch.bfloat16, torch.float32): + signature = format_signature( + hidden_states=dense_tensor_format(routed_weight.dtype), + w13_weight=dense_tensor_format(w13_weight.dtype), + w13_scale=dense_tensor_format(w13_scale.dtype), + w2_weight=dense_tensor_format(w2_weight.dtype), + w2_scale=dense_tensor_format(w2_scale.dtype), + topk_weights=dense_tensor_format(topk_weight_dtype), + topk_ids=dense_tensor_format(torch.int32), + shared_input=dense_tensor_format(shared_gate_up_weight.dtype), + shared_weight=dense_tensor_format(shared_down_weight.dtype), + routed_out=dense_tensor_format(routed_weight.dtype), + shared_out=dense_tensor_format(shared_down_weight.dtype), + ) + try: + select_kernel("moe", "latent_expert_shared", signature, traits=traits) + except NoKernelFoundError: + continue + break + else: + return False + return True def latent_moe_expert_shared( diff --git a/tokenspeed-kernel/test/ops/moe/test_gluon_mxfp4_amd.py b/tokenspeed-kernel/test/ops/moe/test_gluon_mxfp4_amd.py index f820d5c1e1..7c60b12f02 100644 --- a/tokenspeed-kernel/test/ops/moe/test_gluon_mxfp4_amd.py +++ b/tokenspeed-kernel/test/ops/moe/test_gluon_mxfp4_amd.py @@ -265,11 +265,11 @@ def record_stage2(*args, **kwargs): torch.testing.assert_close(actual.float(), expected, atol=2e-2, rtol=2e-2) -def test_bf16_activation_situ_moe() -> None: +@pytest.mark.parametrize("num_tokens", [1, 2, 3, 4]) +def test_bf16_activation_situ_moe(num_tokens: int) -> None: if not is_cdna4(): pytest.skip("BF16 SiTU activation is unavailable on this GPU") - num_tokens = 1 num_experts = 2 hidden_size = 3584 intermediate_size = 3072 @@ -306,7 +306,9 @@ def test_bf16_activation_situ_moe() -> None: topk_weights = torch.full( (num_tokens, top_k), 1.0 / top_k, dtype=torch.float32, device="cuda" ) - topk_ids = torch.tensor([[0, 1]], dtype=torch.int32, device="cuda") + topk_ids = torch.tensor([[0, 1]] * num_tokens, dtype=torch.int32, device="cuda") + shared_input = torch.randn(num_tokens, 768, dtype=torch.bfloat16, device="cuda") + shared_weight = torch.randn(7168, 768, dtype=torch.bfloat16, device="cuda") actual = gluon_a16w4_situ_warp_decode_ep_gfx950( hidden_states, @@ -320,11 +322,21 @@ def test_bf16_activation_situ_moe() -> None: situ_linear_beta=25.0, linear_weights=True, w13_interleaved=True, + shared_input=shared_input, + shared_weight=shared_weight, ) torch.cuda.synchronize() - assert actual.shape == hidden_states.shape - torch.testing.assert_close(actual, torch.zeros_like(actual), atol=0, rtol=0) + assert isinstance(actual, tuple) + routed, shared = actual + assert routed.shape == hidden_states.shape + torch.testing.assert_close(routed, torch.zeros_like(routed), atol=0, rtol=0) + torch.testing.assert_close( + shared, + torch.nn.functional.linear(shared_input, shared_weight), + atol=2e-2, + rtol=2e-2, + ) def test_static_fp8_activation_moe_gfx950_smoke() -> None: diff --git a/tokenspeed-kernel/test/ops/moe/test_gluon_mxfp4_routing_gfx950.py b/tokenspeed-kernel/test/ops/moe/test_gluon_mxfp4_routing_gfx950.py index 8b43421f07..8dac4371ca 100644 --- a/tokenspeed-kernel/test/ops/moe/test_gluon_mxfp4_routing_gfx950.py +++ b/tokenspeed-kernel/test/ops/moe/test_gluon_mxfp4_routing_gfx950.py @@ -20,6 +20,8 @@ from __future__ import annotations +from unittest import mock + import pytest import tokenspeed_kernel.ops.moe.gluon.sigmoid_topk as gluon_sigmoid_topk import torch @@ -101,6 +103,74 @@ def launch(route_input, bias, topk, **kwargs): assert actual_ids is sentinel_ids +def test_public_sigmoid_bias_topk_uses_per_token_float32_route( + monkeypatch: pytest.MonkeyPatch, +) -> None: + logits = torch.zeros((2, 896), device="cuda", dtype=torch.float32) + correction_bias = torch.zeros(896, device="cuda", dtype=torch.float32) + sentinel_ids = torch.empty((2, 16), device="cuda", dtype=torch.int32) + sentinel_weights = torch.empty((2, 16), device="cuda", dtype=torch.float32) + + def launch(route_input, bias, topk, **kwargs): + assert route_input is logits + assert bias is correction_bias + assert topk == 16 + return sentinel_ids, sentinel_weights + + monkeypatch.setattr( + gluon_sigmoid_topk, + "invoke_sigmoid_bias_topk_route_prefill_gluon", + launch, + ) + actual_weights, actual_ids = moe_sigmoid_bias_topk( + logits, + correction_bias, + 16, + routed_scaling_factor=2.827, + normalize_topk_weights=True, + ) + + assert actual_weights is sentinel_weights + assert actual_ids is sentinel_ids + + +def test_public_sigmoid_bias_topk_retains_decode_route_above_prefill_topk( + monkeypatch: pytest.MonkeyPatch, +) -> None: + logits = torch.zeros((2, 64), device="cuda", dtype=torch.float32) + correction_bias = torch.zeros(64, device="cuda", dtype=torch.float32) + sentinel_ids = torch.empty((2, 32), device="cuda", dtype=torch.int32) + sentinel_weights = torch.empty((2, 32), device="cuda", dtype=torch.float32) + + def launch(route_input, bias, topk, **kwargs): + assert route_input is logits + assert bias is correction_bias + assert topk == 32 + return sentinel_ids, sentinel_weights + + monkeypatch.setattr( + gluon_sigmoid_topk, + "invoke_sigmoid_bias_topk_route_gluon", + launch, + ) + monkeypatch.setattr( + gluon_sigmoid_topk, + "invoke_sigmoid_bias_topk_route_prefill_gluon", + mock.Mock(side_effect=AssertionError("top-k 32 must use the decode route")), + ) + actual_weights, actual_ids = moe_sigmoid_bias_topk( + logits, + correction_bias, + 32, + routed_scaling_factor=2.827, + normalize_topk_weights=True, + solution="gluon", + ) + + assert actual_weights is sentinel_weights + assert actual_ids is sentinel_ids + + @pytest.mark.parametrize("dtype", _ROUTE_DTYPES) def test_sigmoid_bias_topk_route_gluon_fuses_sigmoid( monkeypatch: pytest.MonkeyPatch, diff --git a/tokenspeed-kernel/test/ops/test_kimi3_projection_gfx950.py b/tokenspeed-kernel/test/ops/test_kimi3_projection_gfx950.py index 177edb6ab3..7305286dfb 100644 --- a/tokenspeed-kernel/test/ops/test_kimi3_projection_gfx950.py +++ b/tokenspeed-kernel/test/ops/test_kimi3_projection_gfx950.py @@ -152,13 +152,17 @@ def test_kimi3_latent_projection_add3_matches_torch_and_captures( torch.testing.assert_close(actual, expected, rtol=2e-2, atol=2e-2) -def test_kimi3_rmsnorm_linear_add_matches_composed_and_captures() -> None: +@pytest.mark.parametrize("num_tokens", [1, 2, 3, 4]) +def test_kimi3_rmsnorm_linear_add_matches_composed_and_captures( + num_tokens: int, +) -> None: torch.manual_seed(37) - latent = torch.randn(1, 3584, device="cuda", dtype=torch.bfloat16) + latent = torch.randn(num_tokens, 3584, device="cuda", dtype=torch.bfloat16) norm_weight = torch.randn(3584, device="cuda", dtype=torch.bfloat16) projection_weight = torch.randn(7168, 3584, device="cuda", dtype=torch.bfloat16) - prefix = torch.randn(1, 7168, device="cuda", dtype=torch.bfloat16) - shared = torch.randn(1, 7168, device="cuda", dtype=torch.bfloat16) + prefix = torch.randn(num_tokens, 7168, device="cuda", dtype=torch.bfloat16) + shared_lane = torch.randn(num_tokens, 10752, device="cuda", dtype=torch.bfloat16) + shared = shared_lane[:, 3584:] source = latent.float() normalized = ( source