[Quant] Use canonical N-first weight format for CT WNA16 MoE

tests/kernels/quantization/test_int4_emulation_moe.py

@@ -240,6 +240,29 @@ def _make_awq_moe_weights(E, K, N, group_size, asym=False):
     )
 
 
+def _canonicalize_gptq_moe_inputs(
+    w13: torch.Tensor,
+    w2: torch.Tensor,
+    w13_scale: torch.Tensor,
+    w2_scale: torch.Tensor,
+    w13_qzeros: torch.Tensor | None,
+    w2_qzeros: torch.Tensor | None,
+):
+    """Convert GPTQ K-first tensors to the oracle's canonical N-first layout."""
+
+    def transpose(value: torch.Tensor | None) -> torch.Tensor | None:
+        return value.transpose(1, 2).contiguous() if value is not None else None
+
+    return (
+        w13.transpose(1, 2).contiguous(),
+        w2.transpose(1, 2).contiguous(),
+        w13_scale.transpose(1, 2).contiguous(),
+        w2_scale.transpose(1, 2).contiguous(),
+        transpose(w13_qzeros),
+        transpose(w2_qzeros),
+    )
+
+
 def _run_emulation_forward(
     experts, w13_bf16, w2_bf16, hidden_states, topk_weights, topk_ids, E, K, N
 ):
@@ -478,7 +501,8 @@ def test_gptq_process_weights_shapes_and_values(E, K, N, group_size, asym):
     w13, w13s, w13z, w13_ref, w2, w2s, w2z, w2_ref = _make_gptq_moe_weights(
         E, K, N, group_size, asym
     )
-    result = _process_weights_emulation_gptq(w13, w2, w13s, w2s, w13z, w2z)
+    gptq_inputs = _canonicalize_gptq_moe_inputs(w13, w2, w13s, w2s, w13z, w2z)
+    result = _process_weights_emulation_gptq(*gptq_inputs)
     w13_out, w2_out = result[0], result[1]
 
     assert w13_out.shape == (E, 2 * N, K)
@@ -561,14 +585,15 @@ def test_gptq_awq_process_weights_agree(E, K, N, group_size, asym):
         g2s_list.append(s2)
         a2s_list.append(s2)
 
-    gptq_res = _process_weights_emulation_gptq(
+    gptq_inputs = _canonicalize_gptq_moe_inputs(
         torch.stack(g13_list),
         torch.stack(g2_list),
         torch.stack(g13s_list),
         torch.stack(g2s_list),
         torch.stack(g13z_list) if asym else None,
         torch.stack(g2z_list) if asym else None,
     )
+    gptq_res = _process_weights_emulation_gptq(*gptq_inputs)
     awq_res = _process_weights_emulation_awq(
         torch.stack(a13_list),
         torch.stack(a2_list),
@@ -615,14 +640,15 @@ def test_gptq_vs_awq_forward_agree(E, K, N, top_k, group_size, num_tokens):
         g2s_list.append(s2)
         a2s_list.append(s2.clone())
 
-    gptq_res = _process_weights_emulation_gptq(
+    gptq_inputs = _canonicalize_gptq_moe_inputs(
         torch.stack(g13_list),
         torch.stack(g2_list),
         torch.stack(g13s_list),
         torch.stack(g2s_list),
         None,
         None,
     )
+    gptq_res = _process_weights_emulation_gptq(*gptq_inputs)
     awq_res = _process_weights_emulation_awq(
         torch.stack(a13_list),
         torch.stack(a2_list),
@@ -750,14 +776,17 @@ def test_ep_output_matches_no_ep(E, K, N, top_k, group_size, num_tokens, ep_size
         if fmt == "gptq"
         else _process_weights_emulation_awq
     )
-    res = process_fn(
+    process_inputs = (
         torch.stack(packed13_list),
         torch.stack(packed2_list),
         torch.stack(scales13_list),
         torch.stack(scales2_list),
         None,
         None,
     )
+    if fmt == "gptq":
+        process_inputs = _canonicalize_gptq_moe_inputs(*process_inputs)
+    res = process_fn(*process_inputs)
     w13_all, w2_all = res[0], res[1]  # [E, 2N, K], [E, K, N]
 
     hidden_states = torch.randn(num_tokens, K, dtype=torch.bfloat16, device=device)
@@ -855,14 +884,15 @@ def test_ep_gptq_awq_agree(E, K, N, top_k, group_size, num_tokens, ep_size):
         g2s_list.append(s2)
         a2s_list.append(s2.clone())
 
-    gptq_res = _process_weights_emulation_gptq(
+    gptq_inputs = _canonicalize_gptq_moe_inputs(
         torch.stack(g13_list),
         torch.stack(g2_list),
         torch.stack(g13s_list),
         torch.stack(g2s_list),
         None,
         None,
     )
+    gptq_res = _process_weights_emulation_gptq(*gptq_inputs)
     awq_res = _process_weights_emulation_awq(
         torch.stack(a13_list),
         torch.stack(a2_list),
@@ -954,14 +984,15 @@ def test_ep_partial_rank_no_active_experts(
         scales13_list.append(s13)
         scales2_list.append(s2)
 
-    res = _process_weights_emulation_gptq(
+    gptq_inputs = _canonicalize_gptq_moe_inputs(
         torch.stack(packed13_list),
         torch.stack(packed2_list),
         torch.stack(scales13_list),
         torch.stack(scales2_list),
         None,
         None,
     )
+    res = _process_weights_emulation_gptq(*gptq_inputs)
     w13_all, w2_all = res[0], res[1]
 
     # Force topk_ids to only use experts in [0, num_local) — rank 0's slice
@@ -1029,14 +1060,15 @@ def test_ep_sum_equals_full_forward(E, K, N, top_k, group_size, num_tokens, ep_s
         scales13_list.append(s13)
         scales2_list.append(s2)
 
-    res = _process_weights_emulation_gptq(
+    gptq_inputs = _canonicalize_gptq_moe_inputs(
         torch.stack(packed13_list),
         torch.stack(packed2_list),
         torch.stack(scales13_list),
         torch.stack(scales2_list),
         None,
         None,
     )
+    res = _process_weights_emulation_gptq(*gptq_inputs)
     w13_all, w2_all = res[0], res[1]
 
     # Fix routing so every token uses exactly 2 consecutive experts (round-robin)
@@ -1143,14 +1175,17 @@ def test_emulation_output_close_to_bf16_reference(
         if fmt == "gptq"
         else _process_weights_emulation_awq
     )
-    res = process_fn(
+    process_inputs = (
         torch.stack(packed13_list),
         torch.stack(packed2_list),
         torch.stack(scales13_list),
         torch.stack(scales2_list),
         None,
         None,
     )
+    if fmt == "gptq":
+        process_inputs = _canonicalize_gptq_moe_inputs(*process_inputs)
+    res = process_fn(*process_inputs)
     w13_bf16, w2_bf16 = res[0], res[1]
 
     dummy_scale = torch.ones(1, dtype=torch.float16, device=device)

tests/quantization/test_moe_wna16.py

@@ -149,7 +149,11 @@ def test_wna16_oracle_rejects_incompatible_quant_structures(
     assert expected in reason
 
 
-def test_compressed_tensors_weights_are_transposed_for_triton():
+def test_compressed_tensors_nfirst_weights_pass_through_for_triton():
+    """CT provides N-first [E, N, K_packed] int32; Triton expects N-first uint8.
+
+    No transpose needed — just a view from int32 to uint8.
+    """
     quant_config = QuantizationArgs(
         num_bits=4,
         type=QuantizationType.INT,
@@ -158,10 +162,41 @@ def test_compressed_tensors_weights_are_transposed_for_triton():
         dynamic=False,
         group_size=32,
     )
-    w13 = torch.arange(16, dtype=torch.int32).reshape(1, 2, 8)
-    w2 = torch.arange(12, dtype=torch.int32).reshape(1, 2, 6)
-    w13_scale = torch.arange(32, dtype=torch.float16).reshape(1, 4, 8)
-    w2_scale = torch.arange(18, dtype=torch.float16).reshape(1, 3, 6)
+    # N-first layout: [E, N, K_packed]
+    w13 = torch.arange(16, dtype=torch.int32).reshape(1, 8, 2)
+    w2 = torch.arange(12, dtype=torch.int32).reshape(1, 6, 2)
+    w13_scale = torch.arange(32, dtype=torch.float16).reshape(1, 8, 4)
+    w2_scale = torch.arange(18, dtype=torch.float16).reshape(1, 6, 3)
+
+    converted = convert_to_wna16_moe_kernel_format(
+        backend=WNA16MoEBackend.TRITON,
+        layer=torch.nn.Module(),
+        quant_config=quant_config,
+        input_dtype=None,
+        w13=w13,
+        w2=w2,
+        w13_scale=w13_scale,
+        w2_scale=w2_scale,
+    )
+
+    assert converted is not None
+    assert torch.equal(converted[0], w13.contiguous().view(torch.uint8))
+    assert torch.equal(converted[1], w2.contiguous().view(torch.uint8))
+    assert torch.equal(converted[2], w13_scale)
+    assert torch.equal(converted[3], w2_scale)
+
+
+def test_gptq_nfirst_weights_pass_through_for_triton():
+    """GPTQ frontend normalizes to N-first before calling oracle.
+
+    Verifies the oracle handles N-first GPTQ inputs the same way as CT.
+    """
+    quant_config = AutoGPTQConfig(4, 128, False, True, False, {}, {})
+    # N-first layout (after GPTQ frontend transpose): [E, N, K_packed]
+    w13 = torch.arange(16, dtype=torch.int32).reshape(1, 8, 2)
+    w2 = torch.arange(12, dtype=torch.int32).reshape(1, 6, 2)
+    w13_scale = torch.arange(32, dtype=torch.float16).reshape(1, 8, 4)
+    w2_scale = torch.arange(18, dtype=torch.float16).reshape(1, 6, 3)
 
     converted = convert_to_wna16_moe_kernel_format(
         backend=WNA16MoEBackend.TRITON,
@@ -175,10 +210,10 @@ def test_compressed_tensors_weights_are_transposed_for_triton():
     )
 
     assert converted is not None
-    assert torch.equal(converted[0], w13.transpose(1, 2).contiguous().view(torch.uint8))
-    assert torch.equal(converted[1], w2.transpose(1, 2).contiguous().view(torch.uint8))
-    assert torch.equal(converted[2], w13_scale.transpose(1, 2).contiguous())
-    assert torch.equal(converted[3], w2_scale.transpose(1, 2).contiguous())
+    assert torch.equal(converted[0], w13.contiguous().view(torch.uint8))
+    assert torch.equal(converted[1], w2.contiguous().view(torch.uint8))
+    assert torch.equal(converted[2], w13_scale)
+    assert torch.equal(converted[3], w2_scale)
 
 
 def test_moe_wna16_setup_forwards_selected_backend(monkeypatch):