Merge branch 'main' into patches

Signed-off-by: MOVZX <movzx@yahoo.com>

tests/model_executor/layers/test_mla_short_prefill_indexer.py

@@ -830,15 +830,19 @@ def test_sparse_mqa_logits_prefill_matches_dense_masked(capture: bool):
     _assert_same_selection(dense, out, ref)
 
 
-def _build_fp4_paged_cache(req_lens: torch.Tensor, page_kv: int, head_dim: int):
+def _build_fp4_paged_cache(
+    req_lens: torch.Tensor, page_kv: int, head_dim: int, page_pad: int = 0
+):
     """MXFP4 paged cache in DeepGEMM's fused layout: per page, all tokens'
-    fp4 data first, then their packed UE8M0 scales."""
+    fp4 data first, then their packed UE8M0 scales; ``page_pad`` bytes pad
+    each page like the other layers' pages in a packed block."""
     device = req_lens.device
     pages_per = (req_lens + page_kv - 1) // page_kv
     num_pages, max_pages = int(pages_per.sum()), int(pages_per.max())
+    head_bytes = head_dim // 2 + 4
     kv_cache = torch.zeros(
-        num_pages, page_kv, 1, head_dim // 2 + 4, dtype=torch.uint8, device=device
-    )
+        num_pages, page_kv * head_bytes + page_pad, dtype=torch.uint8, device=device
+    )[:, : page_kv * head_bytes].view(num_pages, page_kv, 1, head_bytes)
     block_table = torch.zeros(
         len(req_lens), max_pages, dtype=torch.int32, device=device
     )
@@ -864,7 +868,10 @@ def _build_fp4_paged_cache(req_lens: torch.Tensor, page_kv: int, head_dim: int):
 
 @pytest.mark.parametrize("next_n", [1, 2])
 @pytest.mark.parametrize("capture", [False, True])
-def test_sparse_mqa_logits_paged_matches_dense_masked(next_n: int, capture: bool):
+@pytest.mark.parametrize("page_pad", [0, 256])
+def test_sparse_mqa_logits_paged_matches_dense_masked(
+    next_n: int, capture: bool, page_pad: int
+):
     """Paged sparse-MQA decode selects exactly the values ``torch.topk`` picks
     from the dense bf16 masked logits, incl. MTP rows (next_n > 1 flattened to
     one query row each), CUDA graph capture/replay and schedule reuse by a
@@ -886,8 +893,10 @@ def test_sparse_mqa_logits_paged_matches_dense_masked(next_n: int, capture: bool
         .clamp(min=0)
         .to(torch.int32)
     )
-    kv_cache, block_table = _build_fp4_paged_cache(req_lens, page_kv, head_dim)
-    assert kv_cache.stride(0) % 512 == 0  # the paged sparse kernel requires it
+    kv_cache, block_table = _build_fp4_paged_cache(
+        req_lens, page_kv, head_dim, page_pad
+    )
+    assert kv_cache.stride(0) % 512 == page_pad  # padded pages are not 512B
 
     q_fp, q_sf = _quant_fp4(
         torch.randn(rows * num_heads, head_dim, device="cuda", dtype=torch.bfloat16)

tests/model_executor/layers/test_rocm_unquantized_gemm.py

@@ -17,6 +17,15 @@ if current_platform.is_cuda():
 from vllm.model_executor.layers import utils
 
 
+def _on_gfx11() -> bool:
+    """gfx11 check that is safe to evaluate at collection time on any platform."""
+    if not current_platform.is_rocm():
+        return False
+    from vllm.platforms.rocm import on_gfx11
+
+    return on_gfx11()
+
+
 @pytest.mark.parametrize("m", [1, 128])
 def test_rocm_unquantized_gemm_gfx1x_wvsplitk_path(monkeypatch, m):
     x = torch.randn(1, 64, dtype=torch.float16)
@@ -145,6 +154,101 @@ def test_rocm_unquantized_gemm_rejects_unsupported_skinny_layouts(
     assert torch.allclose(out, ref, atol=1e-3, rtol=1e-3)
 
 
+def _row_stride_padded(weight: torch.Tensor, pad: int) -> torch.Tensor:
+    """A weight whose row stride is offset by ``pad`` elements, as produced by
+    ``maybe_pad_weight_avoid_cache_aliasing``."""
+    padded = torch.nn.functional.pad(weight, (0, pad))[..., :-pad]
+    assert padded.shape == weight.shape
+    assert padded.stride() == (weight.shape[1] + pad, 1)
+    return padded
+
+
+def test_rocm_unquantized_gemm_wvsplitk_accepts_row_stride_padded_weight(monkeypatch):
+    # wvSplitK reads the weight row stride, so a row-stride-padded weight must
+    # still reach it instead of falling back to torch.nn.functional.linear.
+    x = torch.randn(4, 1024, dtype=torch.float16)
+    weight = _row_stride_padded(torch.randn(128, 1024, dtype=torch.float16), 64)
+
+    monkeypatch.setattr(utils, "use_aiter_triton_gemm", lambda *args: False)
+    monkeypatch.setattr(utils.rocm_aiter_ops, "is_tgemm_enabled", lambda: False)
+    monkeypatch.setattr(utils.envs, "VLLM_ROCM_USE_SKINNY_GEMM", True)
+    monkeypatch.setattr("vllm.platforms.rocm.on_gfx11", lambda: True)
+    monkeypatch.setattr("vllm.platforms.rocm.on_gfx1x", lambda: True)
+    monkeypatch.setattr("vllm.platforms.rocm.on_gfx9", lambda: False)
+    monkeypatch.setattr("vllm.platforms.rocm.on_gfx950", lambda: False)
+    monkeypatch.setattr("vllm.platforms.rocm.on_gfx1250", lambda: False)
+    monkeypatch.setattr(utils, "num_compute_units", lambda: 120)
+
+    wvsplitk_mock = MagicMock(side_effect=lambda w, x_view, _, __: x_view @ w.t())
+    monkeypatch.setattr(utils.ops, "wvSplitK", wvsplitk_mock)
+
+    out = utils.rocm_unquantized_gemm_impl(x, weight, None)
+    ref = torch.nn.functional.linear(x, weight, None)
+
+    wvsplitk_mock.assert_called_once()
+    assert wvsplitk_mock.call_args.args[0].stride() == weight.stride()
+    assert torch.allclose(out, ref, atol=1e-3, rtol=1e-3)
+
+
+@pytest.mark.parametrize(
+    "k,element_size,expected",
+    [
+        (4096, 2, 64),  # 8192 B row stride, aliases
+        (12288, 2, 64),  # 24576 B row stride, aliases
+        (1024, 2, 64),  # smallest aliasing bf16/fp16 K
+        (2560, 2, 0),  # 5120 B, no aliasing
+        (11008, 2, 0),  # 22016 B, no aliasing
+        (2048, 1, 128),  # fp8: 2048 B, aliases, twice the pad
+        (512, 4, 32),  # fp32: 2048 B, aliases, half the pad
+    ],
+)
+def test_cache_aliasing_pad_elems(k, element_size, expected):
+    assert utils.cache_aliasing_pad_elems(k, element_size) == expected
+
+
+def test_maybe_pad_weight_avoid_cache_aliasing_skips_non_gfx11(monkeypatch):
+    # Padding is gfx11-only, even for a K whose row stride would alias there.
+    monkeypatch.setattr(utils.current_platform, "is_rocm", lambda: True)
+    monkeypatch.setattr("vllm.platforms.rocm.on_gfx11", lambda: False)
+
+    weight = torch.randn(128, 4096, dtype=torch.bfloat16)
+    assert utils.maybe_pad_weight_avoid_cache_aliasing(weight) is weight
+
+
+def test_maybe_pad_weight_avoid_cache_aliasing(monkeypatch):
+    monkeypatch.setattr(utils.current_platform, "is_rocm", lambda: True)
+    monkeypatch.setattr("vllm.platforms.rocm.on_gfx11", lambda: True)
+
+    weight = torch.randn(128, 4096, dtype=torch.bfloat16)
+    padded = utils.maybe_pad_weight_avoid_cache_aliasing(weight)
+
+    assert padded.shape == weight.shape
+    assert padded.stride() == (4096 + 64, 1)
+    assert torch.equal(padded, weight)
+    # Idempotent: the padded stride no longer aliases, so a second call is a no-op
+    # and preserves the pointer captured by any graph.
+    repadded = utils.maybe_pad_weight_avoid_cache_aliasing(padded)
+    assert repadded.data_ptr() == padded.data_ptr()
+
+
+@pytest.mark.parametrize(
+    "weight",
+    [
+        torch.randn(128, 2560, dtype=torch.bfloat16),  # no aliasing
+        torch.randn(2560, 128, dtype=torch.bfloat16).t(),  # not contiguous
+        torch.randn(2, 128, 4096, dtype=torch.bfloat16),  # not 2D
+    ],
+)
+def test_maybe_pad_weight_avoid_cache_aliasing_skips(monkeypatch, weight):
+    monkeypatch.setattr(utils.current_platform, "is_rocm", lambda: True)
+    monkeypatch.setattr("vllm.platforms.rocm.on_gfx11", lambda: True)
+
+    assert (
+        utils.maybe_pad_weight_avoid_cache_aliasing(weight).data_ptr()
+        == weight.data_ptr()
+    )
+
+
 @pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm-only kernel test")
 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
 def test_rocm_unquantized_gemm_noncontiguous_activation_real_kernel(monkeypatch, dtype):
@@ -164,6 +268,26 @@ def test_rocm_unquantized_gemm_noncontiguous_activation_real_kernel(monkeypatch,
     torch.testing.assert_close(out, ref, atol=1e-2, rtol=1e-2)
 
 
+@pytest.mark.skipif(
+    not _on_gfx11(), reason="cache-aliasing padding only applies on gfx11"
+)
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
+def test_rocm_unquantized_gemm_padded_weight_real_kernel(monkeypatch, dtype):
+    x = torch.randn(4, 4096, device="cuda", dtype=dtype)
+    weight = _row_stride_padded(torch.randn(128, 4096, device="cuda", dtype=dtype), 64)
+
+    monkeypatch.setattr(utils.envs, "VLLM_ROCM_USE_SKINNY_GEMM", True)
+    original_wvsplitk = utils.ops.wvSplitK
+    wvsplitk_mock = MagicMock(side_effect=original_wvsplitk)
+    monkeypatch.setattr(utils.ops, "wvSplitK", wvsplitk_mock)
+
+    out = utils.rocm_unquantized_gemm_impl(x, weight, None)
+    ref = torch.nn.functional.linear(x, weight, None)
+
+    wvsplitk_mock.assert_called_once()
+    torch.testing.assert_close(out, ref, atol=1e-2, rtol=1e-2)
+
+
 @pytest.mark.parametrize("m", [1, 128])
 def test_rocm_unquantized_gemm_gfx1x_n_gt_5_falls_back(monkeypatch, m):
     # wvSplitK skinny GEMM handles n in [1, 5] (see PR #40687); n > 5 must

tests/model_executor/test_deep_gemm_warmup.py

@@ -123,19 +123,12 @@ def test_mega_moe_m_values_reach_mega_gate(monkeypatch) -> None:
     threshold (1 for 128 experts), so the forward actually reaches the kernel."""
     dg = SimpleNamespace(
         get_bf16_mega_gate_config=lambda m, *_: {"num_sms": -(-m // 16)},
-        get_block_m_for_mega_moe=lambda *_: 16,
     )
     monkeypatch.setattr(deep_gemm_warmup, "_import_deep_gemm", lambda: dg)
     monkeypatch.setattr(deep_gemm_warmup, "_mega_moe_uses_mega_gate", lambda _: True)
-    monkeypatch.setattr(
-        deep_gemm_warmup, "get_ep_group", lambda: SimpleNamespace(world_size=8)
-    )
-    buffer = SimpleNamespace(num_max_tokens_per_rank=64)
     module = SimpleNamespace(
         gate=SimpleNamespace(input_size=4096, output_size=128),
-        experts=SimpleNamespace(
-            top_k=6, num_experts=128, get_symm_buffer=lambda: buffer
-        ),
+        experts=SimpleNamespace(top_k=6, get_block_m=lambda num_tokens: 16),
         use_sequence_parallel=False,
     )
     assert deep_gemm_warmup._get_fp8_fp4_mega_moe_m_values(module, 40) == [16, 32, 40]

tests/model_executor/test_enabled_custom_ops.py

@@ -139,6 +139,122 @@ def test_topk_softmax_dispatch(use_rocm_aiter: bool):
         assert topk_func == vllm_topk_softmax
 
 
+def _topk_gating_launch(
+    num_tokens: int = 4,
+    num_experts: int = 512,
+    topk: int = 10,
+    num_shared_experts: int = 0,
+    weights_width: int | None = None,
+    ids_width: int | None = None,
+    contiguous: bool = True,
+) -> dict:
+    """Tensors for one softmax top-k launch, as dispatch_topk_softmax_func sees them."""
+    width = topk + num_shared_experts
+    weights = torch.empty(num_tokens, weights_width or width)
+    ids = torch.empty(num_tokens, ids_width or width, dtype=torch.int32)
+    ids = ids[:, :topk]
+    if contiguous:
+        gating = torch.empty(num_tokens, num_experts + num_shared_experts)
+    else:
+        gating = torch.empty(num_experts + num_shared_experts + 3, num_tokens).t()
+    return dict(topk_weights=weights, topk_indices=ids, gating_output=gating)
+
+
+@pytest.mark.parametrize(
+    "launch, num_shared_experts, scoring_func, expect_gating",
+    [
+        pytest.param({}, 0, "", True, id="plain-softmax"),
+        pytest.param({"num_tokens": 4096}, 0, "", True, id="max-tokens"),
+        pytest.param({"num_tokens": 4097}, 0, "", False, id="too-many-tokens"),
+        pytest.param({"contiguous": False}, 0, "", False, id="non-contiguous"),
+        pytest.param({}, 0, "sigmoid", False, id="scoring-without-shared"),
+        pytest.param({"num_shared_experts": 1}, 1, "sigmoid", True, id="one-shared"),
+        pytest.param({"num_shared_experts": 8}, 8, "sigmoid", True, id="eight-shared"),
+        pytest.param(
+            {"num_shared_experts": 3}, 3, "sigmoid", False, id="unsupported-count"
+        ),
+        pytest.param({"num_shared_experts": 1}, 1, "", False, id="no-shared-scoring"),
+        pytest.param(
+            {"num_shared_experts": 1}, 1, "softmax", False, id="wrong-shared-scoring"
+        ),
+        pytest.param(
+            {"num_shared_experts": 1, "weights_width": 10},
+            1,
+            "sigmoid",
+            False,
+            id="weights-too-narrow",
+        ),
+        pytest.param(
+            {"num_shared_experts": 1, "ids_width": 12},
+            1,
+            "sigmoid",
+            False,
+            id="row-stride-mismatch",
+        ),
+        pytest.param(
+            {"num_tokens": 4097, "num_shared_experts": 1},
+            1,
+            "sigmoid",
+            False,
+            id="shared-too-many-tokens",
+        ),
+    ],
+)
+@pytest.mark.parametrize("gating_enabled", [True, False])
+def test_topk_softmax_dispatch_aiter_topk_gating(
+    monkeypatch: pytest.MonkeyPatch,
+    launch: dict,
+    num_shared_experts: int,
+    scoring_func: str,
+    expect_gating: bool,
+    gating_enabled: bool,
+):
+    """AITER launches use topk_gating only where it is enabled, and only when it
+    supports them.
+
+    Every other AITER launch must reach the legacy topk_softmax, so this is the
+    one place the choice is made.
+    """
+    monkeypatch.setattr(
+        rocm_aiter_ops, "is_topk_gating_enabled", lambda: gating_enabled
+    )
+    topk_func = dispatch_topk_softmax_func(
+        True,
+        **_topk_gating_launch(**launch),
+        num_shared_experts=num_shared_experts,
+        shared_expert_scoring_func=scoring_func,
+    )
+    if expect_gating and gating_enabled:
+        assert topk_func == rocm_aiter_ops.topk_gating
+    else:
+        assert topk_func == rocm_aiter_ops.topk_softmax
+
+
+@pytest.mark.parametrize(
+    "cdna, fused_moe_enabled, expected",
+    [
+        pytest.param(3, True, True, id="cdna3"),
+        pytest.param(4, True, True, id="cdna4"),
+        pytest.param(2, True, False, id="cdna2"),
+        pytest.param(3, False, False, id="cdna3-aiter-moe-off"),
+        pytest.param(4, False, False, id="cdna4-aiter-moe-off"),
+    ],
+)
+def test_topk_gating_enabled_by_cdna(
+    monkeypatch: pytest.MonkeyPatch,
+    cdna: int,
+    fused_moe_enabled: bool,
+    expected: bool,
+):
+    """topk_gating is on for CDNA3 and CDNA4 with AITER MoE enabled, nowhere else."""
+    monkeypatch.setattr("vllm._aiter_ops.is_aiter_found_and_supported", lambda: True)
+    monkeypatch.setattr(
+        rocm_aiter_ops, "is_fused_moe_enabled", lambda: fused_moe_enabled
+    )
+    monkeypatch.setattr("vllm.platforms.rocm.get_cdna_version", lambda: cdna)
+    assert bool(rocm_aiter_ops.is_topk_gating_enabled()) is expected
+
+
 @pytest.mark.parametrize(
     "use_rocm_aiter", [True, False] if current_platform.is_rocm() else [False]
 )

tests/model_executor/test_routed_experts_capture.py

@@ -7,9 +7,16 @@ from unittest.mock import Mock, patch
 import pytest
 import torch
 
-from vllm.config import VllmConfig
+from vllm.config import CUDAGraphMode, VllmConfig
 from vllm.config.compilation import CompilationMode
 from vllm.distributed.eplb.eplb_state import EplbLayerState
+from vllm.forward_context import (
+    BatchDescriptor,
+    DPMetadata,
+    ForwardContext,
+    get_forward_context,
+    set_forward_context,
+)
 from vllm.model_executor.layers.fused_moe.config import RoutingMethodType
 from vllm.model_executor.layers.fused_moe.experts.cpu_int4_moe import CPUExpertsInt4
 from vllm.model_executor.layers.fused_moe.modular_kernel import (
@@ -36,6 +43,12 @@ pytestmark = pytest.mark.cpu_test
 _REC_MODULE = "vllm.model_executor.layers.fused_moe.routed_experts_capturer"
 
 
+def _forward_context(**kwargs):
+    return ForwardContext(
+        no_compile_layers={}, attn_metadata={}, slot_mapping={}, **kwargs
+    )
+
+
 def _capturer_with_buffer(
     *,
     max_tokens: int = 8,
@@ -319,13 +332,120 @@ def test_routed_experts_capturer_single_dp_no_metadata():
     """dp_metadata is None: capture writes the full topk_ids rows."""
     capturer = _capturer_with_buffer(dp_rank=0)
     topk = torch.tensor([[1, 2], [3, 4], [5, 6]], dtype=torch.int32)
-    ctx = SimpleNamespace(dp_metadata=None)
+    ctx = _forward_context()
     with patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx):
         capturer.capture(layer_id=0, topk_ids=topk)
     assert torch.equal(capturer.device_buffer[:3, 0, :], topk)
     assert capturer.device_buffer[3, 0, 0].item() == -1
 
 
+@pytest.mark.parametrize("tp_size, num_tokens", [(2, 6), (2, 3), (2, 1), (4, 2)])
+@pytest.mark.parametrize(
+    "cg_mode", [CUDAGraphMode.NONE, CUDAGraphMode.FULL, CUDAGraphMode.PIECEWISE]
+)
+def test_single_dp_ep_capture_restores_tp_shards(tp_size, num_tokens, cg_mode):
+    """Capture all current rows after model-level SP, including buffer reuse."""
+    capturer = _capturer_with_buffer(tp_size=tp_size)
+    config = SimpleNamespace(
+        parallel_config=SimpleNamespace(
+            data_parallel_size=1,
+            tensor_parallel_size=tp_size,
+            enable_expert_parallel=True,
+            use_sequence_parallel_moe=False,
+            is_moe_model=True,
+        ),
+        compilation_config=SimpleNamespace(
+            fast_moe_cold_start=False, static_forward_context={}
+        ),
+    )
+    tp_group = Mock()
+    shard_n = (num_tokens + tp_size - 1) // tp_size
+    with (
+        set_forward_context(
+            None,
+            config,
+            num_tokens=num_tokens,
+            cudagraph_runtime_mode=cg_mode,
+        ),
+        patch(f"{_REC_MODULE}.get_tp_group", return_value=tp_group),
+    ):
+        assert get_forward_context().dp_metadata is None
+        for offset in (20, 1):
+            expected = torch.arange(num_tokens * 2, dtype=torch.int32).reshape(-1, 2)
+            expected += offset
+            padded = torch.full((shard_n * tp_size, 2), 99, dtype=torch.int32)
+            padded[:num_tokens] = expected
+            tp_group.all_gather.return_value = padded
+            capturer.capture(0, padded[:shard_n])
+            assert torch.equal(
+                capturer.snapshot_routing_data(num_tokens)[:, 0].int(), expected
+            )
+            assert (capturer.device_buffer[num_tokens:, 0] == -1).all()
+    if num_tokens == 1:
+        tp_group.all_gather.assert_not_called()
+    else:
+        assert tp_group.all_gather.call_count == 2
+
+
+@pytest.mark.parametrize("num_tokens, num_rows", [(6, 6), (3, 4), (0, 0)])
+def test_single_dp_tp_capture_does_not_regather_full_batch(num_tokens, num_rows):
+    capturer = _capturer_with_buffer(tp_size=2)
+    topk = torch.arange(num_rows * 2, dtype=torch.int32).reshape(-1, 2)
+    ctx = _forward_context(batch_descriptor=BatchDescriptor(num_tokens=num_tokens))
+    with (
+        patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx),
+        patch(f"{_REC_MODULE}.get_tp_group") as get_tp_group,
+    ):
+        capturer.capture(0, topk)
+    get_tp_group.assert_not_called()
+    assert torch.equal(capturer.device_buffer[:num_tokens, 0], topk[:num_tokens])
+    assert (capturer.device_buffer[num_tokens:, 0] == -1).all()
+
+
+@pytest.mark.parametrize("num_tokens, num_rows", [(6, 2), (0, 1)])
+def test_single_dp_tp_capture_rejects_unexpected_batch_without_writing(
+    num_tokens, num_rows
+):
+    capturer = _capturer_with_buffer(tp_size=2)
+    topk = torch.ones(num_rows, 2, dtype=torch.int32)
+    ctx = _forward_context(batch_descriptor=BatchDescriptor(num_tokens=num_tokens))
+    with (
+        patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx),
+        pytest.raises(AssertionError, match="unexpected topk_ids batch dim"),
+    ):
+        capturer.capture(0, topk)
+    assert (capturer.device_buffer == -1).all()
+
+
+@pytest.mark.parametrize("tp_rank", [0, 1])
+def test_capture_uses_pre_sp_size_after_model_shards_padding_mask(tp_rank):
+    """Model-level SP shortens the mask; capture still needs all three rows."""
+    from vllm.models.common.ops import sequence_parallel as sp
+
+    capturer = _capturer_with_buffer(tp_size=2)
+    expected = torch.arange(8, dtype=torch.int32).reshape(4, 2)
+    ctx = _forward_context(
+        is_padding=torch.zeros(3, dtype=torch.bool),
+        batch_descriptor=BatchDescriptor(num_tokens=3),
+    )
+    with (
+        patch.object(sp, "get_tensor_model_parallel_world_size", return_value=2),
+        patch.object(sp, "get_tensor_model_parallel_rank", return_value=tp_rank),
+    ):
+        ctx.is_padding = sp.sp_padding_mask(ctx.is_padding, torch.empty(3, 1))
+    assert len(ctx.is_padding) == 2
+    tp_group = Mock()
+    tp_group.all_gather.return_value = expected
+    with (
+        patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx),
+        patch(f"{_REC_MODULE}.get_tp_group", return_value=tp_group),
+    ):
+        capturer.capture(0, expected[2 * tp_rank : 2 * (tp_rank + 1)])
+    tp_group.all_gather.assert_called_once()
+    torch.testing.assert_close(capturer.device_buffer[:3, 0], expected[:3])
+    assert (capturer.device_buffer[3:, 0] == -1).all()
+
+
 @pytest.mark.parametrize(
     ("num_experts", "dtype_name", "torch_dtype"),
     [(256, "uint8", torch.uint8), (257, "uint16", torch.uint16)],
@@ -359,7 +479,7 @@ def test_routed_experts_capturer_narrows_snapshot(output_dtype):
     capturer = _capturer_with_buffer(dtype=torch.int32)
     capturer.output_dtype = output_dtype
     topk = torch.tensor([[1, 2], [254, 255]], dtype=torch.int64)
-    ctx = SimpleNamespace(dp_metadata=None)
+    ctx = _forward_context()
     with patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx):
         capturer.capture(layer_id=0, topk_ids=topk)
 
@@ -375,8 +495,8 @@ def test_routed_experts_capturer_dp_naive_concatenated_all_ranks():
     """N == sum(num_tokens_dp): slice this rank's segment from concatenated topk."""
     capturer = _capturer_with_buffer(dp_rank=1)
     num_tokens_dp = torch.tensor([2, 3], dtype=torch.int32)
-    ctx = SimpleNamespace(
-        dp_metadata=SimpleNamespace(num_tokens_across_dp_cpu=num_tokens_dp)
+    ctx = _forward_context(
+        dp_metadata=DPMetadata(num_tokens_across_dp_cpu=num_tokens_dp)
     )
     # Concatenated order: rank0 rows then rank1 rows.
     topk = torch.tensor(
@@ -392,8 +512,8 @@ def test_routed_experts_capturer_dp_modular_local_tokens():
     """N == token_num_per_dp: topk is already local to this DP rank."""
     capturer = _capturer_with_buffer(dp_rank=1)
     num_tokens_dp = torch.tensor([2, 3], dtype=torch.int32)
-    ctx = SimpleNamespace(
-        dp_metadata=SimpleNamespace(num_tokens_across_dp_cpu=num_tokens_dp)
+    ctx = _forward_context(
+        dp_metadata=DPMetadata(num_tokens_across_dp_cpu=num_tokens_dp)
     )
     topk = torch.tensor([[10, 11], [12, 13], [14, 15]], dtype=torch.int32)
     with patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx):
@@ -404,8 +524,8 @@ def test_routed_experts_capturer_dp_modular_local_tokens():
 def test_routed_experts_capturer_dp_ep_gathered_shards():
     capturer = _capturer_with_buffer(dp_rank=1, tp_size=2)
     num_tokens_dp = torch.tensor([3, 2], dtype=torch.int32)
-    ctx = SimpleNamespace(
-        dp_metadata=SimpleNamespace(
+    ctx = _forward_context(
+        dp_metadata=DPMetadata(
             num_tokens_across_dp_cpu=num_tokens_dp,
             local_sizes=[2, 2, 1, 1],
         )
@@ -421,13 +541,17 @@ def test_routed_experts_capturer_dp_ep_gathered_shards():
     assert torch.equal(capturer.device_buffer[:2, 0], topk[4:])
 
 
-def test_routed_experts_capturer_sp_modular_gathers_tp_shards():
-    capturer = _capturer_with_buffer(dp_rank=1, tp_size=2)
+@pytest.mark.parametrize("single_dp", [False, True])
+def test_routed_experts_capturer_sp_modular_gathers_tp_shards(single_dp):
+    capturer = _capturer_with_buffer(dp_rank=0 if single_dp else 1, tp_size=2)
     num_tokens_dp = torch.tensor([2, 3], dtype=torch.int32)
-    ctx = SimpleNamespace(
-        dp_metadata=SimpleNamespace(num_tokens_across_dp_cpu=num_tokens_dp)
+    ctx = _forward_context(
+        dp_metadata=(
+            None if single_dp else DPMetadata(num_tokens_across_dp_cpu=num_tokens_dp)
+        ),
+        batch_descriptor=BatchDescriptor(num_tokens=3),
     )
-    local_shard = torch.tensor([[10, 11], [12, 13]], dtype=torch.int32)
+    local_shard = torch.tensor([[10, 12], [11, 13]], dtype=torch.int32).t()
     gathered = torch.tensor([[10, 11], [12, 13], [14, 15], [-1, -1]], dtype=torch.int32)
     tp_group = Mock()
     tp_group.all_gather.return_value = gathered
@@ -438,16 +562,83 @@ def test_routed_experts_capturer_sp_modular_gathers_tp_shards():
     ):
         capturer.capture(layer_id=0, topk_ids=local_shard)
 
-    tp_group.all_gather.assert_called_once_with(local_shard, dim=0)
+    tp_group.all_gather.assert_called_once()
+    gather_input = tp_group.all_gather.call_args.args[0]
+    assert gather_input.is_contiguous()
+    torch.testing.assert_close(gather_input, local_shard)
+    assert tp_group.all_gather.call_args.kwargs == {"dim": 0}
     assert torch.equal(capturer.device_buffer[:3, 0], gathered[:3])
 
 
+def test_routed_experts_capturer_rejects_incomplete_dp_shard_groups():
+    """Invalid gather metadata must not silently select another rank's rows."""
+    capturer = _capturer_with_buffer(dp_rank=1, tp_size=2)
+    ctx = _forward_context(
+        dp_metadata=DPMetadata(
+            num_tokens_across_dp_cpu=torch.tensor([3, 2]),
+            local_sizes=[2, 2, 2],
+        )
+    )
+    with (
+        patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx),
+        pytest.raises(ValueError, match="not divisible by DP rank count"),
+    ):
+        capturer.capture(0, torch.ones(6, 2, dtype=torch.int32))
+    assert (capturer.device_buffer == -1).all()
+
+
+@pytest.mark.parametrize("layer_id", [-1, 4])
+def test_routed_experts_capturer_rejects_layer_before_restoring_rows(layer_id):
+    capturer = _capturer_with_buffer(tp_size=2)
+    ctx = _forward_context(batch_descriptor=BatchDescriptor(num_tokens=3))
+    with (
+        patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx),
+        patch(f"{_REC_MODULE}.get_tp_group") as get_tp_group,
+        pytest.raises(IndexError, match="exceeds capture buffer"),
+    ):
+        capturer.capture(layer_id, torch.ones(2, 2, dtype=torch.int32))
+    get_tp_group.assert_not_called()
+    assert (capturer.device_buffer == -1).all()
+
+
+@pytest.mark.parametrize(
+    ("counts", "local_sizes", "start", "stop"),
+    [
+        ([3, 0], [2, 2, 0, 0], 3, 3),
+        ([0, 3], None, 0, 3),
+        ([0, 0], [0, 0, 0, 0], 0, 0),
+        ([2, 2], [1, 1, 1, 1], 2, 4),
+    ],
+    ids=["idle-rank", "only-active-rank", "all-idle", "unpadded-gather-tie"],
+)
+def test_routed_experts_capturer_dp_preserves_empty_and_unpadded_spans(
+    counts, local_sizes, start, stop
+):
+    capturer = _capturer_with_buffer(dp_rank=1, tp_size=2)
+    topk = torch.arange(sum(counts) * 2, dtype=torch.int32).reshape(-1, 2)
+    ctx = _forward_context(
+        dp_metadata=DPMetadata(
+            num_tokens_across_dp_cpu=torch.tensor(counts), local_sizes=local_sizes
+        )
+    )
+    with (
+        patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx),
+        patch(f"{_REC_MODULE}.get_tp_group") as get_tp_group,
+    ):
+        capturer.capture(0, topk)
+    get_tp_group.assert_not_called()
+    torch.testing.assert_close(
+        capturer.device_buffer[: stop - start, 0], topk[start:stop]
+    )
+    assert (capturer.device_buffer[stop - start :, 0] == -1).all()
+
+
 def test_routed_experts_capturer_dp_unexpected_batch_raises():
     """Mismatch between topk batch dim and DP layout: fail fast."""
     capturer = _capturer_with_buffer(dp_rank=0)
     num_tokens_dp = torch.tensor([2, 3], dtype=torch.int32)
-    ctx = SimpleNamespace(
-        dp_metadata=SimpleNamespace(num_tokens_across_dp_cpu=num_tokens_dp)
+    ctx = _forward_context(
+        dp_metadata=DPMetadata(num_tokens_across_dp_cpu=num_tokens_dp)
     )
     # total=5, local=2: n=1 matches neither naive (5) nor modular (2).
     topk = torch.tensor([[1, 2]], dtype=torch.int32)