Revert "sonic-moe: add support and testing for concatenated layout (#…

sonic-moe/build.toml

@@ -6,7 +6,6 @@ version = 1
 
 [general.hub]
 repo-id = "kernels-community/sonic-moe"
-branch = "concatenated-gate-up"
 
 [general.cuda]
 minver = "12.8"

sonic-moe/tests/test_moe.py

@@ -13,7 +13,7 @@
     pytest.skip("SonicMoE requires Hopper (SM90) or newer GPU", allow_module_level=True)
 
 try:
-    from sonic_moe import KernelBackendMoE, MoE, enable_quack_gemm, moe_general_routing_inputs
+    from sonic_moe import KernelBackendMoE, MoE, enable_quack_gemm
     from sonic_moe.enums import ActivationType
 except ImportError as e:
     pytest.skip(f"sonicmoe dependencies not available: {e}", allow_module_level=True)
@@ -118,176 +118,3 @@ def test_moe_quack_gemm(problem_shape):
         assert_close(y_kernel.float(), y_torch.float(), atol=1.4e-2, rtol=2e-2)
 
     torch.cuda.empty_cache()
-
-
-# ────────────────── is_concatenated_gate_up tests ──────────────────
-
-
-def _make_interleaved_and_concatenated_weights(E, I, H, device, dtype):
-    """Create matched interleaved and concatenated weight pairs."""
-    w1_inter = torch.randn(E, 2 * I, H, device=device, dtype=dtype).permute(1, 2, 0)
-    gate = w1_inter[0::2].permute(2, 0, 1)  # (E, I, H)
-    up = w1_inter[1::2].permute(2, 0, 1)    # (E, I, H)
-    w1_concat = torch.cat([gate, up], dim=1).contiguous().permute(1, 2, 0)
-    w2 = torch.randn(E, H, I, device=device, dtype=dtype).permute(1, 2, 0)
-    return w1_inter, w1_concat, w2
-
-
-def _make_routing(T, E, K, device):
-    """Create routing inputs for moe_general_routing_inputs."""
-    topk_indices = torch.randint(0, E, (T, K), device=device)
-    topk_weights = torch.randn(T, K, device=device, dtype=torch.bfloat16).softmax(dim=-1)
-    token_idx = torch.arange(T, device=device, dtype=torch.int32).unsqueeze(1).expand(-1, K).reshape(-1)
-    expert_ids = topk_indices.reshape(-1).to(torch.int32)
-    router_scores = topk_weights.reshape(-1).to(torch.bfloat16)
-    return token_idx, expert_ids, router_scores
-
-
-CONCAT_SHAPES = [
-    # (T, H, I, E, K)
-    (128, 768, 256, 128, 8),
-    (1024, 768, 512, 64, 4),
-    (256, 4096, 512, 128, 8),
-    (4096, 4096, 1024, 64, 4),
-]
-
-CONCAT_ACTIVATIONS = [ActivationType.SWIGLU, ActivationType.GEGLU, ActivationType.REGLU]
-
-
-@pytest.mark.parametrize("problem_shape", CONCAT_SHAPES)
-@pytest.mark.parametrize("activation_type", CONCAT_ACTIVATIONS)
-def test_concatenated_gate_up(problem_shape, activation_type):
-    """Verify concatenated layout produces bit-exact results vs interleaved."""
-    device = torch.device("cuda")
-    dtype = torch.bfloat16
-    T, H, I, E, K = problem_shape
-
-    set_seed(_SEED)
-
-    w1_inter, w1_concat, w2 = _make_interleaved_and_concatenated_weights(E, I, H, device, dtype)
-    x = 0.02 * torch.randn(T, H, device=device, dtype=dtype)
-    token_idx, expert_ids, router_scores = _make_routing(T, E, K, device)
-    stream_id = torch.cuda.current_stream(device).cuda_stream
-
-    out_ref, _ = moe_general_routing_inputs(
-        x, router_scores, token_idx, expert_ids,
-        w1_inter, None, w2, None,
-        E, stream_id, activation_type,
-        is_inference_mode_enabled=True,
-        is_concatenated_gate_up=False,
-    )
-    out_test, _ = moe_general_routing_inputs(
-        x, router_scores, token_idx, expert_ids,
-        w1_concat, None, w2, None,
-        E, stream_id, activation_type,
-        is_inference_mode_enabled=True,
-        is_concatenated_gate_up=True,
-    )
-
-    assert torch.equal(out_ref, out_test), (
-        f"Mismatch: max_diff={(out_ref.float() - out_test.float()).abs().max().item()}"
-    )
-
-    torch.cuda.empty_cache()
-
-
-CONCAT_BACKWARD_SHAPES = [
-    (256, 768, 256, 128, 8),
-    (1024, 4096, 512, 64, 4),
-]
-
-
-@pytest.mark.parametrize("problem_shape", CONCAT_BACKWARD_SHAPES)
-def test_concatenated_gate_up_backward(problem_shape):
-    """Verify gradients match between interleaved and concatenated layouts."""
-    device = torch.device("cuda")
-    dtype = torch.bfloat16
-    T, H, I, E, K = problem_shape
-
-    set_seed(_SEED)
-
-    w1_inter, w1_concat, w2 = _make_interleaved_and_concatenated_weights(E, I, H, device, dtype)
-    x = 0.02 * torch.randn(T, H, device=device, dtype=dtype, requires_grad=True)
-    x_clone = x.clone().detach().requires_grad_()
-    token_idx, expert_ids, router_scores = _make_routing(T, E, K, device)
-    stream_id = torch.cuda.current_stream(device).cuda_stream
-
-    # Need w1 to require grad for dw1
-    w1_inter_param = w1_inter.clone().detach().requires_grad_()
-    w1_concat_param = w1_concat.clone().detach().requires_grad_()
-
-    out_ref, _ = moe_general_routing_inputs(
-        x, router_scores, token_idx, expert_ids,
-        w1_inter_param, None, w2, None,
-        E, stream_id, ActivationType.SWIGLU,
-        is_inference_mode_enabled=False,
-        is_concatenated_gate_up=False,
-    )
-    out_test, _ = moe_general_routing_inputs(
-        x_clone, router_scores, token_idx, expert_ids,
-        w1_concat_param, None, w2, None,
-        E, stream_id, ActivationType.SWIGLU,
-        is_inference_mode_enabled=False,
-        is_concatenated_gate_up=True,
-    )
-
-    dy = 0.02 * torch.randn_like(out_ref)
-
-    grads_ref = torch.autograd.grad(out_ref, [x, w1_inter_param], grad_outputs=dy)
-    grads_test = torch.autograd.grad(out_test, [x_clone, w1_concat_param], grad_outputs=dy)
-
-    # dx should match exactly
-    assert_close(grads_ref[0].float(), grads_test[0].float(), atol=1e-2, rtol=1e-2), "dx mismatch"
-
-    # dw1: interleaved grad vs concatenated grad — compare after mapping to same layout
-    dw1_inter = grads_ref[1]  # interleaved layout
-    dw1_concat = grads_test[1]  # concatenated layout
-    # Convert interleaved dw1 to concatenated for comparison
-    dw1_inter_as_concat_gate = dw1_inter[0::2].permute(2, 0, 1)
-    dw1_inter_as_concat_up = dw1_inter[1::2].permute(2, 0, 1)
-    dw1_inter_as_concat = torch.cat([dw1_inter_as_concat_gate, dw1_inter_as_concat_up], dim=1).contiguous().permute(1, 2, 0)
-
-    assert_close(dw1_inter_as_concat.float(), dw1_concat.float(), atol=1e-2, rtol=1e-2), "dw1 mismatch"
-
-    torch.cuda.empty_cache()
-
-
-@pytest.mark.parametrize(
-    "problem_shape",
-    [(256, 768, 256, 128, 8), (1024, 4096, 512, 64, 4)],
-)
-def test_concatenated_gate_up_with_bias(problem_shape):
-    """Verify concatenated layout with bias produces bit-exact results vs interleaved."""
-    device = torch.device("cuda")
-    dtype = torch.bfloat16
-    T, H, I, E, K = problem_shape
-
-    set_seed(_SEED)
-
-    w1_inter, w1_concat, w2 = _make_interleaved_and_concatenated_weights(E, I, H, device, dtype)
-    b1 = torch.randn(E, 2 * I, device=device, dtype=dtype)
-    b2 = torch.randn(E, H, device=device, dtype=dtype)
-    x = 0.02 * torch.randn(T, H, device=device, dtype=dtype)
-    token_idx, expert_ids, router_scores = _make_routing(T, E, K, device)
-    stream_id = torch.cuda.current_stream(device).cuda_stream
-
-    out_ref, _ = moe_general_routing_inputs(
-        x, router_scores, token_idx, expert_ids,
-        w1_inter, b1, w2, b2,
-        E, stream_id, ActivationType.SWIGLU,
-        is_inference_mode_enabled=True,
-        is_concatenated_gate_up=False,
-    )
-    out_test, _ = moe_general_routing_inputs(
-        x, router_scores, token_idx, expert_ids,
-        w1_concat, b1, w2, b2,
-        E, stream_id, ActivationType.SWIGLU,
-        is_inference_mode_enabled=True,
-        is_concatenated_gate_up=True,
-    )
-
-    assert torch.equal(out_ref, out_test), (
-        f"Mismatch: max_diff={(out_ref.float() - out_test.float()).abs().max().item()}"
-    )
-
-    torch.cuda.empty_cache()

sonic-moe/torch-ext/sonic_moe/functional/__init__.py

@@ -70,7 +70,6 @@ def forward(
         is_varlen_K: bool,
         activation_type: ActivationType,
         is_inference_mode_enabled: bool,
-        is_concatenated_gate_up: bool = False,
     ) -> torch.Tensor:
         T, H = x.shape
         I, H, E = w1.shape
@@ -106,7 +105,6 @@ def forward(
                 activation_type=activation_type.value,
                 is_glu_activation=is_glu_activation,
                 is_inference_mode_enabled=is_inference_mode_enabled,
-                is_concatenated_gate_up=is_concatenated_gate_up,
             )
 
         ctx.T = T
@@ -117,7 +115,6 @@ def forward(
         ctx.I = I
         ctx.is_varlen_K = is_varlen_K
         ctx.is_glu_activation = is_glu_activation
-        ctx.is_concatenated_gate_up = is_concatenated_gate_up
         ctx.stream_id = stream_id
 
         ctx.save_for_backward(
@@ -149,7 +146,6 @@ def backward(ctx, _: None, dz: torch.Tensor):
         K = ctx.K
         H = ctx.H
         is_glu_activation = ctx.is_glu_activation
-        is_concatenated_gate_up = ctx.is_concatenated_gate_up
         is_varlen_K = ctx.is_varlen_K
         stream_id = ctx.stream_id
 
@@ -194,7 +190,6 @@ def backward(ctx, _: None, dz: torch.Tensor):
                 s_scatter_idx=s_scatter_idx,
                 is_glu_activation=is_glu_activation,
                 stream_id=stream_id,
-                is_concatenated_gate_up=is_concatenated_gate_up,
             )
 
             _up_projection_backward_weight(
@@ -206,7 +201,6 @@ def backward(ctx, _: None, dz: torch.Tensor):
                 x_gather_idx=x_gather_idx,
                 is_glu_activation=is_glu_activation,
                 stream_id=stream_id,
-                is_concatenated_gate_up=is_concatenated_gate_up,
             )
 
         dx_reduced = torch.empty(T, H, dtype=dz.dtype, device=dz.device)
@@ -221,7 +215,7 @@ def backward(ctx, _: None, dz: torch.Tensor):
             is_varlen_K=is_varlen_K,
         )
 
-        return dx_reduced, dw1, db1, *[None] * 13
+        return dx_reduced, dw1, db1, *[None] * 12
 
 
 class _DownProjection(torch.autograd.Function):
@@ -492,7 +486,6 @@ def moe_general_routing_inputs(
     stream_id: int,
     activation_type: ActivationType,
     is_inference_mode_enabled: bool = False,
-    is_concatenated_gate_up: bool = False,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     assert ((b1 is None) and (b2 is None)) or (
         (b1 is not None) and (b2 is not None)
@@ -538,7 +531,6 @@ def moe_general_routing_inputs(
         True,  # is_varlen_K
         activation_type,
         is_inference_mode_enabled,
-        is_concatenated_gate_up,
     )
 
     o = _DownProjection.apply(

sonic-moe/torch-ext/sonic_moe/functional/backward.py

@@ -206,7 +206,6 @@ def _up_projection_backward_act(
     s_scatter_idx: torch.Tensor,
     is_glu_activation: bool,
     stream_id: int,
-    is_concatenated_gate_up: bool = False,
 ) -> None:
     I, H, E = w1.size()
     if is_glu_activation:
@@ -229,9 +228,9 @@ def _up_projection_backward_act(
         mE_permute_order = convert_torch_tensor_to_cute_tensor(expert_schedule_order, (0,), 0, 4, 1, stream=stream_id)
     current_stream = cuda.CUstream(stream_id)
 
-    compile_dx_key = ("dx", E, H, I, is_glu_activation, dx_expanded.dtype, is_concatenated_gate_up)
+    compile_dx_key = ("dx", E, H, I, is_glu_activation, dx_expanded.dtype)
     if compile_dx_key not in _up_projection_backward_act.compile_cache:
-        dx_module = HopperWgmma_MoE_Up_proj_ActGrad_Bwd(E, H, I, is_glu_activation, is_concatenated_gate_up=is_concatenated_gate_up)
+        dx_module = HopperWgmma_MoE_Up_proj_ActGrad_Bwd(E, H, I, is_glu_activation)
         tensormaps = [dx_module.module.generate_tensormap(None, None, None) for _ in range(2)]
         _up_projection_backward_act.compile_cache[compile_dx_key] = cute.compile(
             dx_module,
@@ -245,9 +244,9 @@ def _up_projection_backward_act(
             mE_permute_order,
             current_stream,
         )
-        _up_projection_backward_act.compile_cache[(TENSORMAP, compile_dx_key)] = tensormaps
+        _up_projection_backward_act.compile_cache[f"dx-{TENSORMAP}"] = tensormaps
 
-    dx_tensormaps = _up_projection_backward_act.compile_cache[(TENSORMAP, compile_dx_key)]
+    dx_tensormaps = _up_projection_backward_act.compile_cache[f"dx-{TENSORMAP}"]
     _up_projection_backward_act.compile_cache[compile_dx_key](
         mDz,
         mW1_trans,
@@ -274,7 +273,6 @@ def _up_projection_backward_weight(
     x_gather_idx: torch.Tensor,
     is_glu_activation: bool,
     stream_id: int,
-    is_concatenated_gate_up: bool = False,
 ) -> None:
     I, H, E = dw1.size()
     if is_glu_activation:
@@ -295,9 +293,9 @@ def _up_projection_backward_weight(
         mE_permute_order = convert_torch_tensor_to_cute_tensor(expert_schedule_order, (0,), 0, 4, 1, stream=stream_id)
     current_stream = cuda.CUstream(stream_id)
 
-    compile_dw1_key = ("dw1", E, H, I, is_glu_activation, x.dtype, is_concatenated_gate_up)
+    compile_dw1_key = ("dw1", E, H, I, is_glu_activation, x.dtype)
     if compile_dw1_key not in _up_projection_backward_weight.compile_cache:
-        dw1_module = HopperWgmma_MoE_Up_proj_WeightGrad_Bwd(E, H, I, is_glu_activation, is_concatenated_gate_up=is_concatenated_gate_up)
+        dw1_module = HopperWgmma_MoE_Up_proj_WeightGrad_Bwd(E, H, I, is_glu_activation)
         tensormaps = [dw1_module.module.generate_tensormap(None, None, None) for _ in range(1)]
         _up_projection_backward_weight.compile_cache[compile_dw1_key] = cute.compile(
             dw1_module,
@@ -310,9 +308,9 @@ def _up_projection_backward_weight(
             mE_permute_order,
             current_stream,
         )
-        _up_projection_backward_weight.compile_cache[(TENSORMAP, compile_dw1_key)] = tensormaps
+        _up_projection_backward_weight.compile_cache[f"dw1-{TENSORMAP}"] = tensormaps
 
-    dw1_tensormaps = _up_projection_backward_weight.compile_cache[(TENSORMAP, compile_dw1_key)]
+    dw1_tensormaps = _up_projection_backward_weight.compile_cache[f"dw1-{TENSORMAP}"]
     _up_projection_backward_weight.compile_cache[compile_dw1_key](
         mX_trans,
         mDz_trans,
@@ -408,14 +406,14 @@ def _down_projection_backward_act(
             mE_permute_order,
             current_stream,
         )
-        _down_projection_backward_act.compile_cache[(TENSORMAP, compile_dz_key)] = tensormaps
+        _down_projection_backward_act.compile_cache[f"dz-{TENSORMAP}"] = tensormaps
 
     if ds_partial is None:
         ds_partial_N = _down_projection_backward_act.compile_cache["ds_partial_N"]
         ds_partial = torch.empty(TK, ds_partial_N, dtype=torch.float32, device=topk_scores.device)
         mDS_partial = convert_torch_tensor_to_cute_tensor(ds_partial, (0, 1), 1, 4, 1, stream=stream_id)
 
-    dz_tensormaps = _down_projection_backward_act.compile_cache[(TENSORMAP, compile_dz_key)]
+    dz_tensormaps = _down_projection_backward_act.compile_cache[f"dz-{TENSORMAP}"]
     _down_projection_backward_act.compile_cache[compile_dz_key](
         mDout,
         mW2_trans,
@@ -522,9 +520,9 @@ def _down_projection_backward_weight(
             mE_permute_order,
             current_stream,
         )
-        _down_projection_backward_weight.compile_cache[(TENSORMAP, compile_dw2_key)] = tensormaps
+        _down_projection_backward_weight.compile_cache[f"dw2-{TENSORMAP}"] = tensormaps
 
-    dw2_tensormaps = _down_projection_backward_weight.compile_cache[(TENSORMAP, compile_dw2_key)]
+    dw2_tensormaps = _down_projection_backward_weight.compile_cache[f"dw2-{TENSORMAP}"]
     _down_projection_backward_weight.compile_cache[compile_dw2_key](
         mDout_trans, mY1S_trans, mDw2, mE_offset, mX_gather, dw2_tensormaps, mE_permute_order, current_stream
     )