Skip to content

Merge upstream→gfx11 #1413

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
eble-amd merged 3 commits into from
Oct 29, 2026
Merged

Merge upstream→gfx11 #1413

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
c01368f
[Quant] Use canonical N-first weight format for CT WNA16 MoE (#52798)
HDCharles Sep 28, 2026
55a5667
Merge commit 'c01368fd64eaecc0c7242e4d933ec80fe8e047f5'
eble-amd Oct 28, 2026
c273a29
Hybrid W4A16 MoE: convert N-first CT weights
eble-amd Oct 28, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
51 changes: 43 additions & 8 deletions tests/kernels/quantization/test_int4_emulation_moe.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
):
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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),
Expand Down Expand Up @@ -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),
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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),
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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)
Expand Down
53 changes: 44 additions & 9 deletions tests/quantization/test_moe_wna16.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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,
Expand All @@ -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):
Expand Down
Loading
Loading