[Codoge]

Motivation

RowParallelLinear resolves parallel_group="replicated" to no TP group, but its output allocation still enters the global TP symmetric-memory context whenever allocation symmetry is enabled. Kimi-K3's replicated shared experts produce complete local outputs and do not need that collective memory window.

Modifications

Disable the TP symmetric-memory allocation context when the linear has no TP group. Keep the existing attention-TP allocation path and the policy for sharded linears. Apply the fix in RowParallelLinear.forward so callers do not need model-specific allocation overrides.

Accuracy Tests

Upstream b146a8576de930cf3adc104e601feb2790cdf3f6 plus this PR only, CPU PyTorch 2.14.1:

python -m pytest test/registered/unit/layers/test_kimi_mlp_parallel_groups.py -q
# 2 passed

The existing test matrix runs real linear weight loading and forward operations across TP, attention-TP, shared-expert TP and replicated layouts. It now enables the surrounding allocation-symmetry policy and verifies that only replicated outputs disable the TP allocation context. Numerical comparisons still pass.

Additional distributed validation at 7861b6ce65 on 8× B300, BF16, with actual NCCL symmetric memory enabled: PASS. The 12 cases per rank cover replicated, TP and attention-TP outputs at nominal token counts 1 / 17 / 257 / 8192. is_tensor_in_symmetric_mempool confirms replicated outputs use ordinary allocations and sharded outputs remain in the NCCL pool. Replicated batch sizes vary by rank; sharded outputs complete the real TP AllReduce and match the reference exactly. This is a deterministic linear/allocation smoke test, not a full-model accuracy benchmark.

Speed Tests and Profiling

No standalone speedup is claimed. Earlier tuning combined local allocation with metadata and scheduling changes, so that aggregate improvement cannot be attributed to this patch. This PR corrects allocation scope without changing linear arithmetic. A full-model run with symmetric memory enabled has not been repeated on this branch.

Checklist

• Format changed files with the repository's isort / Ruff rules; AST, whitespace and registered-test checks pass.
• Add regression coverage to registered test suites.
• Follow the SGLang code style guidance.
• Full-model accuracy and serving benchmark on this branch (pending).

---

CI States

Latest PR Test (Base): ❌ Run #41299337224
Latest PR Test (Extra): ❌ Run #41299337141
Latest PR Test (AMD ROCm 10): ✅ Run #41299337140

Tip

Base CI only covers the core cases, and its workload may not exercise your change. Use /rerun-test to run the tests your change touches, or copy this prompt for an agent:

Find the registered tests this PR affects and run them with /rerun-test, following https://github.com/sgl-project/sglang/blob/main/.agents/skills/rerun-test/SKILL.md

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[Codoge]

@ch-wan Ready for review. Could an authorized maintainer enable CI with /tag-and-rerun-ci and review this allocation-scope fix?

The current red workflow results stopped at Block draft PR; the PR is now non-draft and satisfies the repository's MIN_BASE_SHA requirement.

At head 7861b6ce65, both existing linear-group tests pass with the added allocation assertions. Their matrix exercises real weight loading and forward operations across TP, attention-TP, shared-expert TP and replicated layouts. The fix retains the attention-TP path and disables global-TP symmetric allocation for replicated outputs. No standalone throughput gain is claimed; the distributed allocation smoke test below has now also passed.

Focused CI coverage:

/rerun-test test/registered/unit/layers/test_kimi_mlp_parallel_groups.py

This PR is independent of #47742 and #47743.

Additional distributed validation at 7861b6ce65 on 8× B300, BF16, with actual NCCL symmetric memory enabled: PASS. The 12 cases per rank cover replicated, TP and attention-TP outputs at nominal token counts 1 / 17 / 257 / 8192. is_tensor_in_symmetric_mempool confirms replicated outputs use ordinary allocations and sharded outputs remain in the NCCL pool. Replicated batch sizes vary by rank; sharded outputs complete the real TP AllReduce and match the reference exactly. This is a deterministic linear/allocation smoke test, not a full-model accuracy benchmark.

Reproduce the distributed allocation smoke test

Save this as validate_linear_symm.py and run against this PR:

NCCL_CUMEM_ENABLE=1 NCCL_NVLS_ENABLE=1 PYTHONPATH=python python -m torch.distributed.run --standalone --nproc_per_node=8 validate_linear_symm.py

"""Manual distributed validation of PR #47744 on its exact source tree."""

import json
import os

import torch
import torch.distributed as dist
import torch.nn.functional as F

from sglang.srt.distributed import parallel_state as ps
from sglang.srt.distributed.device_communicators.pynccl_allocator import (
    is_symmetric_memory_enabled,
    is_tensor_in_symmetric_mempool,
)
from sglang.srt.layers.linear import RowParallelLinear
from sglang.test.test_utils import publish_build_topology


@torch.inference_mode()
def main():
    rank = int(os.environ["RANK"])
    size = int(os.environ["WORLD_SIZE"])
    torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
    ps.set_custom_all_reduce(False)
    ps.init_distributed_environment(
        world_size=size, rank=rank, local_rank=rank, timeout=120
    )
    publish_build_topology(
        tp_size=size,
        world_rank=rank,
        enable_symm_mem=True,
        enable_nccl_nvls=True,
        disable_custom_all_reduce=True,
        device="cuda",
    )
    ps.initialize_model_parallel(enable_symm_mem=True)
    assert is_symmetric_memory_enabled()
    cases = []
    try:
        for layout in ("replicated", "tp", "attn_tp"):
            sharded = layout != "replicated"
            layer = RowParallelLinear(
                256,
                512,
                bias=False,
                input_is_parallel=False,
                params_dtype=torch.bfloat16,
                reduce_results=False,
                parallel_group=layout,
                use_dp_attention_reduce=layout == "attn_tp",
            ).cuda()
            layer.weight.fill_(1.0 / 256)
            for nominal_tokens in (1, 17, 257, 8192):
                # Replicated modules may see different local batch sizes.
                tokens = nominal_tokens if sharded else nominal_tokens + rank
                x = torch.ones(tokens, 256, device="cuda", dtype=torch.bfloat16)
                output, _ = layer(x)
                symmetric = is_tensor_in_symmetric_mempool(output)
                assert symmetric == sharded, (rank, layout, tokens, symmetric)
                expected_local = F.linear(x, torch.full_like(layer.weight, 1 / 256)) if not sharded else torch.full_like(output, 1 / size)
                torch.testing.assert_close(output, expected_local, rtol=0, atol=0)
                if sharded:
                    reduced = ps.get_tp_group().all_reduce(output)
                    torch.testing.assert_close(
                        reduced, torch.ones_like(reduced), rtol=0, atol=0
                    )
                cases.append({"layout": layout, "tokens": tokens, "symmetric": symmetric})
                dist.barrier()
        torch.cuda.synchronize()
        gathered = [None] * size
        dist.all_gather_object(gathered, cases)
        if rank == 0:
            print(json.dumps({"status": "PASS", "world_size": size, "cases_per_rank": len(cases), "rank_cases": gathered}), flush=True)
    finally:
        ps.destroy_model_parallel()
        ps.destroy_distributed_environment()


if __name__ == "__main__":
    main()

[BBuf]

[BBuf Review Agent] LGTM. Please run the weekly Kimi parallel-group test before merging.

[BBuf]

/tag-and-rerun-ci

[BBuf]

/rerun-test test/registered/unit/layers/test_kimi_mlp_parallel_groups.py

[github-actions]

🚀 ubuntu-latest (1 test): ✅ View workflow run

cd test/ && python3 registered/unit/layers/test_kimi_mlp_parallel_groups.py

[qybnb]

After #47744, an existing unit-test expectation is inconsistent with the updated symmetric-memory policy.

RowParallelLinear.forward() now disables symmetric-memory allocation when self.tp_group is None. For parallel_group="replicated", the group resolves to None, so disabled=True is expected.

However, test_replicated_model_mlps_keep_weights_and_tp_output_allocation still unconditionally asserts:
allocator.assert_called_once_with(tp, disabled=False)
This causes the ExaoneMoEMLP, group="replicated" case to fail. The numerical comparison immediately before this assertion passes; the failure is in the allocator-policy assertion.

Could you update the expectation to account for replicated layers, while preserving the numerical checks? For example:
allocator.assert_called_once_with(
tp, disabled=(group == "replicated")
)
@Codoge @BBuf