From 9babff311a5b1de5067f197aba49f5b20f16db5f Mon Sep 17 00:00:00 2001 From: jiagaoxiang Date: Thu, 3 Sep 2026 16:43:52 +0000 Subject: [PATCH 1/2] test: pin AttentionResidualMixer against an independent mix Closes PRPUNDIT-22. Independent RMS-score / unnormalised-value oracle, plus num_blocks=0 identity, hidden-size guard, and gradcheck. --- .../backends/megatron/kimi_k3/__init__.py | 2 + .../kimi_k3/test_attention_residual.py | 230 ++++++++++++++++++ 2 files changed, 232 insertions(+) create mode 100644 tests/unit_tests/backends/megatron/kimi_k3/__init__.py create mode 100644 tests/unit_tests/backends/megatron/kimi_k3/test_attention_residual.py diff --git a/tests/unit_tests/backends/megatron/kimi_k3/__init__.py b/tests/unit_tests/backends/megatron/kimi_k3/__init__.py new file mode 100644 index 000000000..89778402a --- /dev/null +++ b/tests/unit_tests/backends/megatron/kimi_k3/__init__.py @@ -0,0 +1,2 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. diff --git a/tests/unit_tests/backends/megatron/kimi_k3/test_attention_residual.py b/tests/unit_tests/backends/megatron/kimi_k3/test_attention_residual.py new file mode 100644 index 000000000..2b696a929 --- /dev/null +++ b/tests/unit_tests/backends/megatron/kimi_k3/test_attention_residual.py @@ -0,0 +1,230 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Tests for ``AttentionResidualMixer.forward``. + +The mixer replaces ``x = x + sublayer(x)`` with a learned softmax mixture over +the running residual stream (``prefix_sum``) and a set of cross-layer +checkpoints (``block_residual``); see the module docstring at +``primus/backends/megatron/core/transformer/kimi_k3/attention_residual.py`` +for the three easy-to-get-wrong details this file pins: + +1. the output mixes the *un-normalised* candidates -- RMS normalisation feeds + the softmax scores only; +2. the scorer is rank-1, ``norm_weight (elementwise*) proj_weight``; +3. everything runs in fp32 (or wider) internally and casts back to + ``prefix_sum``'s dtype at the end. + +``_reference_attn_res_mix`` below is an independent transcription of +``_apply_attn_res`` (``modeling_kimi_linear.py``) -- written without looking at +``attn_res_kernels._eager.reference.eager_attn_res_mix`` -- so a regression in +either implementation has a real chance of being caught rather than the test +merely checking the module against its own backend. +""" + +from __future__ import annotations + +import pytest + +torch = pytest.importorskip("torch") + +mcore_module = pytest.importorskip( + "megatron.core.transformer.module", + reason="megatron.core is not importable in this environment", +) + +from megatron.core.transformer.transformer_config import TransformerConfig # noqa: E402 + +from primus.backends.megatron.core.transformer.kimi_k3.attention_residual import ( # noqa: E402 + AttentionResidualMixer, +) + + +def _make_config(hidden_size=8, eps=1e-6, params_dtype=torch.float32, attn_res_backend="eager"): + config = TransformerConfig( + hidden_size=hidden_size, + num_attention_heads=1, + num_layers=1, + layernorm_epsilon=eps, + params_dtype=params_dtype, + use_cpu_initialization=True, + ) + config.attn_res_backend = attn_res_backend + return config + + +def _make_mixer(hidden_size=8, eps=1e-6, params_dtype=torch.float32, seed=0): + mixer = AttentionResidualMixer(_make_config(hidden_size, eps, params_dtype)) + # Reseed *after* construction (reset_parameters may or may not draw from + # the RNG depending on config.init_method) so the overwritten values below + # -- and everything a caller subsequently draws with torch.randn -- are + # deterministic regardless of that detail. + torch.manual_seed(seed) + with torch.no_grad(): + mixer.norm_weight.normal_(mean=1.0, std=0.1) + mixer.proj_weight.normal_(mean=0.0, std=1.0) + return mixer + + +def _reference_attn_res_mix(prefix_sum, block_residual, norm_weight, proj_weight, eps): + """Independent transcription of ``_apply_attn_res``. + + Deliberately written candidate-by-candidate with an explicit RMSNorm and + an explicit weighted sum, rather than the batched ``cat`` / + ``matmul`` formulation the production eager kernel uses, so this does not + just re-run the same code under a different name. + """ + compute_dtype = torch.promote_types(prefix_sum.dtype, torch.float32) + prefix_sum_f = prefix_sum.to(compute_dtype) + block_residual_f = block_residual.to(compute_dtype) + score_vec = (norm_weight.to(compute_dtype) * proj_weight.to(compute_dtype).squeeze(0)).to(compute_dtype) + + num_blocks = block_residual_f.shape[-2] + candidates = [block_residual_f[..., i, :] for i in range(num_blocks)] + [prefix_sum_f] + + raw_scores = [] + for cand in candidates: + variance = cand.pow(2).mean(dim=-1, keepdim=True) + normed = cand * torch.rsqrt(variance + eps) + raw_scores.append((normed * score_vec).sum(dim=-1)) + + scores = torch.stack(raw_scores, dim=-1) # [*, num_candidates] + probs = scores.softmax(dim=-1) + + mixed = torch.zeros_like(prefix_sum_f) + for i, cand in enumerate(candidates): + mixed = mixed + probs[..., i : i + 1] * cand + + return mixed.to(prefix_sum.dtype) + + +# --------------------------------------------------------------------------- +# Correctness against the independent reference +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("num_blocks", [1, 3]) +@pytest.mark.parametrize("leading_shape", [(5,), (2, 3)]) +def test_forward_matches_independent_reference(num_blocks, leading_shape): + """Bit-close parity against a from-scratch transcription of the reference. + + Covers both a flat token axis and Megatron's ``[seq, batch]`` layout: the + module docstring notes ``torch.matmul`` broadcasts over leading dims + identically either way, so both shapes should hit the same arithmetic. + """ + hidden_size = 8 + mixer = _make_mixer(hidden_size=hidden_size) + + prefix_sum = torch.randn(*leading_shape, hidden_size) + block_residual = torch.randn(*leading_shape, num_blocks, hidden_size) + + actual = mixer(prefix_sum, block_residual) + expected = _reference_attn_res_mix( + prefix_sum, block_residual, mixer.norm_weight, mixer.proj_weight, mixer.eps + ) + + assert actual.shape == prefix_sum.shape + torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-6) + + +def test_forward_output_dtype_matches_prefix_sum(): + """Detail 3: fp32 internally, cast back once to ``prefix_sum``'s dtype.""" + mixer = _make_mixer(hidden_size=8, params_dtype=torch.float32) + + prefix_sum = torch.randn(4, 8, dtype=torch.float64) + block_residual = torch.randn(4, 2, 8, dtype=torch.float64) + + out = mixer(prefix_sum, block_residual) + + assert out.dtype == torch.float64 + expected = _reference_attn_res_mix( + prefix_sum, block_residual, mixer.norm_weight, mixer.proj_weight, mixer.eps + ) + torch.testing.assert_close(out, expected, rtol=1e-6, atol=1e-7) + + +# --------------------------------------------------------------------------- +# Detail 1: mixing the un-normalised candidates, not the RMS-normalised ones +# --------------------------------------------------------------------------- + + +def test_forward_mixes_unnormalised_candidates(): + """A candidate with huge norm must not be rescaled away by the mixer. + + If the module mixed the RMS-normalised ``k`` instead of the raw ``v`` (the + bug the module docstring calls out), every candidate would end up at unit + RMS and a large-magnitude block checkpoint would look identical to a small + one at the output. Constructing one candidate 100x the others and checking + the output scales with it (rather than collapsing to O(1) RMS) catches + that swap. + """ + hidden_size = 8 + mixer = _make_mixer(hidden_size=hidden_size) + + prefix_sum = torch.randn(1, hidden_size) + big_block = torch.randn(1, 1, hidden_size) * 100.0 + + out = mixer(prefix_sum, big_block) + expected = _reference_attn_res_mix(prefix_sum, big_block, mixer.norm_weight, mixer.proj_weight, mixer.eps) + + torch.testing.assert_close(out, expected, rtol=1e-5, atol=1e-6) + # The un-normalised mix can only reach this magnitude if the raw (not + # RMS-normalised) big_block leaked into the weighted sum. + assert out.norm() > 5.0 + + +# --------------------------------------------------------------------------- +# num_blocks == 0: a legal no-op softmax over the stream alone +# --------------------------------------------------------------------------- + + +def test_forward_with_zero_blocks_is_identity(): + mixer = _make_mixer(hidden_size=8) + prefix_sum = torch.randn(3, 8) + block_residual = torch.empty(3, 0, 8) + + out = mixer(prefix_sum, block_residual) + + torch.testing.assert_close(out, prefix_sum, rtol=1e-5, atol=1e-6) + + +# --------------------------------------------------------------------------- +# Shape validation +# --------------------------------------------------------------------------- + + +def test_forward_raises_on_hidden_size_mismatch(): + mixer = _make_mixer(hidden_size=8) + prefix_sum = torch.randn(2, 8) + block_residual = torch.randn(2, 1, 4) + + with pytest.raises(ValueError, match="hidden"): + mixer(prefix_sum, block_residual) + + +# --------------------------------------------------------------------------- +# Differentiability (module docstring: gradcheck-safe end to end) +# --------------------------------------------------------------------------- + + +def test_forward_gradcheck(): + hidden_size = 4 + mixer = _make_mixer(hidden_size=hidden_size, params_dtype=torch.float64) + mixer = mixer.to(torch.float64) + + prefix_sum = torch.randn(2, hidden_size, dtype=torch.float64, requires_grad=True) + block_residual = torch.randn(2, 2, hidden_size, dtype=torch.float64, requires_grad=True) + mixer.norm_weight.requires_grad_(True) + mixer.proj_weight.requires_grad_(True) + + # norm_weight / proj_weight are passed through so gradcheck perturbs the + # exact tensors the mixer reads (they are the same objects, not copies). + assert torch.autograd.gradcheck( + lambda ps, br, nw, pw: mixer(ps, br), + (prefix_sum, block_residual, mixer.norm_weight, mixer.proj_weight), + eps=1e-6, + atol=1e-4, + ) From 529ccbf2b1730b55bb61b82a11370acc5dce88e9 Mon Sep 17 00:00:00 2001 From: Doug J Date: Thu, 3 Sep 2026 11:04:34 -0700 Subject: [PATCH 2/2] test: drop unused importorskip binding (Copilot review) pytest.importorskip's return value was never used; call it for its side effect (skip if megatron.core is unimportable) without binding it. --- .../backends/megatron/kimi_k3/test_attention_residual.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit_tests/backends/megatron/kimi_k3/test_attention_residual.py b/tests/unit_tests/backends/megatron/kimi_k3/test_attention_residual.py index 2b696a929..6b3701de1 100644 --- a/tests/unit_tests/backends/megatron/kimi_k3/test_attention_residual.py +++ b/tests/unit_tests/backends/megatron/kimi_k3/test_attention_residual.py @@ -31,7 +31,7 @@ torch = pytest.importorskip("torch") -mcore_module = pytest.importorskip( +pytest.importorskip( "megatron.core.transformer.module", reason="megatron.core is not importable in this environment", )