[Perf][DSv4.1] Pin the low-latency MXFP8 GEMM tactics at TP8 and without autotuning

tests/models/deepseek_v41/nvidia/specialized_decode/test_kernels.py

@@ -61,6 +61,7 @@
     HIDDEN,
     PADDED,
     SLICE,
+    SLICE_BY_PADDED,
     TOP_K,
     FrontConfig,
     RoutedExpertGemm,
@@ -112,8 +113,7 @@ def _routing(rows: int, generator: torch.Generator) -> torch.Tensor:
     return torch.stack(ids).to(torch.int32).to(DEVICE)
 
 
-@pytest.fixture(scope="module")
-def experts():
+def _make_experts(slice_: int, padded: int):
     """Random MXFP4 experts converted by vLLM for the TRT-LLM-gen backend, the
     tensors both the FlashInfer path and the low-batch kernels read."""
     from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (
@@ -122,7 +122,7 @@ def experts():
     )
 
     generator = torch.Generator().manual_seed(0)
-    e, p, h = EXPERTS, PADDED, HIDDEN
+    e, p, h = EXPERTS, padded, HIDDEN
 
     def randint(low, high, *shape):
         return torch.randint(low, high, shape, dtype=torch.uint8, generator=generator)
@@ -132,9 +132,9 @@ def randint(low, high, *shape):
         randint(118, 126, e, 2 * p, h // 32),
     )
     w2, w2_scale = randint(0, 256, e, h, p // 2), randint(118, 126, e, h, p // 32)
-    # The intermediate is 576 wide, padded with zeros to 640.
-    w13.view(e, 2, p, -1)[:, :, SLICE:] = 0
-    w2[:, :, SLICE // 2 :] = 0
+    # The per-rank intermediate (576 at TP4, 288 at TP8) is zero-padded.
+    w13.view(e, 2, p, -1)[:, :, slice_:] = 0
+    w2[:, :, slice_ // 2 :] = 0
     w13, w2, w13_scale, w2_scale, _, _ = convert_weight_to_mxfp4_moe_kernel_format(
         Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_MXFP8,
         nn.Module(),
@@ -150,6 +150,24 @@ def randint(low, high, *shape):
     return w13, w13_scale, w2, w2_scale, clamp
 
 
+@pytest.fixture(scope="module")
+def experts():
+    """TP4 experts (576-column intermediate slice padded to 640)."""
+    return _make_experts(SLICE, PADDED)
+
+
+@pytest.fixture(
+    scope="module",
+    params=sorted(SLICE_BY_PADDED.items(), reverse=True),
+    ids=lambda shape: f"padded{shape[0]}",
+)
+def experts_by_tp(request):
+    """Experts at every tensor-parallel slice the GEMMs take: TP4's 576 columns
+    padded to 640 and TP8's 288 padded to 384."""
+    padded, slice_ = request.param
+    return _make_experts(slice_, padded)
+
+
 def _gemm(experts) -> RoutedExpertGemm:
     from vllm.model_executor.layers.fused_moe.config import (
         mxfp4_mxfp8_moe_quant_config,
@@ -207,7 +225,7 @@ def _flashinfer_experts(experts, routing, output=None):
         top_k=TOP_K,
         n_group=None,
         topk_group=None,
-        intermediate_size=PADDED,
+        intermediate_size=w13.size(1) // 2,
         local_expert_offset=0,
         local_num_experts=EXPERTS,
         routed_scaling_factor=None,
@@ -242,12 +260,12 @@ def _fma_finalize(rows: torch.Tensor, permuted: torch.Tensor, weights: torch.Ten
 
 @pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"])
 @pytest.mark.parametrize("rows", [1, 5, 6, 8])
-def test_routed_experts_match_flashinfer(experts, rows, graph):
+def test_routed_experts_match_flashinfer(experts_by_tp, rows, graph):
     """Route + GEMM pair: the routed rows bitwise those of FlashInfer's
     TRT-LLM-gen MoE, and their top-k reduction within one BF16 ulp of its
     finalized output."""
     generator = torch.Generator().manual_seed(rows)
-    gemm = _gemm(experts)
+    gemm = _gemm(experts_by_tp)
 
     def inputs():
         x = torch.randn(rows, HIDDEN, generator=generator).to(torch.bfloat16)
@@ -268,10 +286,10 @@ def run():
     else:
         routing, out = run()
     routed = out.gemm2_permuted[out.expanded_idx_to_permuted_idx.long()]
-    assert torch.equal(routed, _flashinfer_experts(experts, routing))
+    assert torch.equal(routed, _flashinfer_experts(experts_by_tp, routing))
 
     finalized = torch.empty(rows, HIDDEN, dtype=torch.bfloat16, device=DEVICE)
-    _flashinfer_experts(experts, routing, finalized)
+    _flashinfer_experts(experts_by_tp, routing, finalized)
     ours = _fma_finalize(
         out.gemm2_permuted, out.expanded_idx_to_permuted_idx, out.expert_weights
     )
@@ -519,9 +537,13 @@ def _quantized(rows: int, k: int) -> QuantizedActivation:
 
 
 @pytest.mark.parametrize("rows", [1, 5, 6, 8])
-@pytest.mark.parametrize("shape", list(mxfp8_linear.PINNED_TACTICS), ids=str)
+@pytest.mark.parametrize(
+    "shape",
+    list(dict.fromkeys((n, k) for n, k, _ in mxfp8_linear.PINNED_TACTICS)),
+    ids=str,
+)
 def test_low_latency_mxfp8_mm(shape, rows):
-    """``fused_wqa_wkv``/``wq_b``/``wo_b`` at their TP4 shapes on
+    """``fused_wqa_wkv``/``wq_b``/``wo_b`` at their TP4 and TP8 shapes on
     ``cutedsl_low_latency``, eagerly and replayed, within BF16 tolerance of the
     stock ``cute-dsl`` layer; a BF16 input is left to the stock layer."""
     if not mxfp8_linear._backend_supported():
@@ -544,6 +566,17 @@ def test_low_latency_mxfp8_mm(shape, rows):
     assert mxfp8_linear.low_latency_mm(linear, x) is None
 
 
+def test_pinned_tactics_by_tensor_parallel_size():
+    """A runner problem (N, M, K, L) gets its shape's pinned tactic at the
+    tensor parallel size it belongs to, and only for 5-8 rows."""
+    pinned = mxfp8_linear._pinned
+    assert pinned((1792, 8, 5120, 1), 4) == (128, 10, 1, 4)
+    assert pinned((1792, 8, 5120, 1), 8) == (128, 10, 1, 2)
+    assert pinned((4096, 5, 1280, 1), 8) == (128, 10, 1, 2)
+    assert pinned((4096, 5, 1280, 1), 4) is None
+    assert pinned((8192, 4, 1280, 1), 4) is None
+
+
 @pytest.mark.parametrize("rows", [1, 5, 6, 8, 16, 17])
 def test_low_batch_lm_head(dist_init, default_vllm_config, rows):
     """The logits processor's LM head on CuTe DSL ``mm_bf16`` for 1-16 rows,

tests/models/deepseek_v41/nvidia/specialized_decode/test_numerics.py

@@ -17,8 +17,10 @@
   sets up to exactly tied scores);
 * one CUDA graph of all of these, replayed ``REPLAYS`` times bitwise.
 
-Needs four SM100-family GPUs with NVLink multicast and flashinfer-python
-0.7.2; about 5 minutes. For a race stress run, run this file as a module
+The fused all-reduce publishes over NVLink multicast where the mailbox gets a
+mapping and with unicast stores otherwise; the ``unicast`` case forces the
+latter. Needs four SM100-family GPUs and flashinfer-python 0.7.2; about 5
+minutes per case. For a race stress run, run this file as a module
 (``python -m tests.models.deepseek_v41.nvidia.specialized_decode.test_numerics``)
 with ``--replays 100000``.
 """
@@ -806,7 +808,9 @@ def run(replays: int) -> str | None:
 @pytest.mark.skipif(
     not current_platform.is_device_capability_family(100), reason="needs SM100"
 )
-def test_specialized_decode_numerics():
+@pytest.mark.parametrize("unicast", [False, True], ids=["multicast", "unicast"])
+def test_specialized_decode_numerics(unicast, monkeypatch):
+    monkeypatch.setenv("VLLM_DSV41_MHC_UNICAST", "1" if unicast else "0")
     if torch.accelerator.device_count() < TP:
         pytest.skip("needs four GPUs")
     from vllm.models.deepseek_v41.nvidia.specialized_decode import FLASHINFER_VERSION

vllm/envs.py

@@ -91,6 +91,7 @@
     VLLM_MAIN_CUDA_VERSION: str = "13.0"
     VLLM_FLOAT32_MATMUL_PRECISION: Literal["highest", "high", "medium"] = "highest"
     VLLM_BATCH_INVARIANT: bool = False
+    VLLM_DSV41_MHC_UNICAST: bool = False
     VLLM_TRITON_USE_TD: bool | None = None
     VLLM_GPU_SYNC_CHECK: Literal["warn", "error"] | None = None
     MAX_JOBS: str | None = None
@@ -640,6 +641,13 @@ def _resolve_rust_cli_path() -> str | None:
     # Enable batch-invariant mode: deterministic results regardless of
     # batch composition. Requires NVIDIA GPU with compute capability >= 9.0.
     "VLLM_BATCH_INVARIANT": lambda: bool(int(os.getenv("VLLM_BATCH_INVARIANT", "0"))),
+    # DeepSeek-V4.1's fused Lamport all-reduce + mHC publishes each rank's
+    # contribution with NVLink multicast when its mailbox gets a multicast
+    # mapping, and with unicast stores into every rank's mailbox otherwise.
+    # Set to 1 to force the unicast path (for testing).
+    "VLLM_DSV41_MHC_UNICAST": lambda: bool(
+        int(os.getenv("VLLM_DSV41_MHC_UNICAST", "0"))
+    ),
     "VLLM_REPLICATE_EMBED": lambda: (
         os.getenv("VLLM_REPLICATE_EMBED", "0").strip().lower() in ("1", "true")
     ),

vllm/models/deepseek_v41/nvidia/ops/cute_dsl/all_reduce_mhc.py

@@ -33,7 +33,9 @@
 from cutlass import BFloat16, Float32, Int32, Int64, Uint8, Uint32
 from cutlass.cute.runtime import make_fake_compact_tensor
 
+import vllm.envs as envs
 from vllm.distributed import get_tp_group
+from vllm.logger import init_logger
 from vllm.models.deepseek_v41.common.ops.query_quant import mxfp8_scale_bytes
 
 from .primitives import (
@@ -64,10 +66,15 @@
     store_global_u32x4,
     store_lamport_sentinel_u32x4,
     store_shared_cluster_f32,
+    stuc_bf16x2,
+    stuc_bf16x4,
+    stuc_bf16x8,
     to_cute,
     to_cute_dynamic,
 )
 
+logger = init_logger(__name__)
+
 LAMPORT_GENERATIONS = 3
 NEXT_STAGE = 0
 ACTIVE_STAGE = 1
@@ -126,6 +133,7 @@ def __init__(
         threads: int,
         release_before_store: bool,
         enable_pdl: bool,
+        unicast: bool = False,
     ) -> None:
         if elements_per_thread not in (1, QUAD_BF16, VEC_BF16):
             raise ValueError("elements_per_thread must be 1, 4, or 8")
@@ -141,6 +149,7 @@ def __init__(
         self.ctas_per_token = math.ceil(self.fragments / threads)
         self.release_before_store = release_before_store
         self.enable_pdl = enable_pdl
+        self.unicast = unicast
 
     @cute.jit
     def __call__(
@@ -242,30 +251,54 @@ def kernel(
                     (Int64(stage) * self.tp + self.rank) * self.capacity_m
                     + Int64(token)
                 ) * self.hidden + Int64(fragment)
-                stmc_bf16x2(
-                    contribution_mailbox_multicast_address + mailbox_element * 2,
-                    packed_word,
-                )
+                if cutlass.const_expr(self.unicast):
+                    stuc_bf16x2(
+                        contribution_mailbox_multicast_address,
+                        mailbox_element * 2,
+                        packed_word,
+                        self.tp,
+                    )
+                else:
+                    stmc_bf16x2(
+                        contribution_mailbox_multicast_address + mailbox_element * 2,
+                        packed_word,
+                    )
         elif cutlass.const_expr(self.elements_per_thread == QUAD_BF16):
             if fragment < self.fragments:
                 mailbox_element = (
                     (Int64(stage) * self.tp + self.rank) * self.capacity_m
                     + Int64(token)
                 ) * self.hidden + Int64(fragment) * QUAD_BF16
-                stmc_bf16x4(
-                    contribution_mailbox_multicast_address + mailbox_element * 2,
-                    packed_words,
-                )
+                if cutlass.const_expr(self.unicast):
+                    stuc_bf16x4(
+                        contribution_mailbox_multicast_address,
+                        mailbox_element * 2,
+                        packed_words,
+                        self.tp,
+                    )
+                else:
+                    stmc_bf16x4(
+                        contribution_mailbox_multicast_address + mailbox_element * 2,
+                        packed_words,
+                    )
         else:
             if fragment < self.fragments:
                 mailbox_element = (
                     (Int64(stage) * self.tp + self.rank) * self.capacity_m
                     + Int64(token)
                 ) * self.hidden + Int64(fragment) * VEC_BF16
-                stmc_bf16x8(
-                    contribution_mailbox_multicast_address + mailbox_element * 2,
-                    packed_words,
-                )
+                if cutlass.const_expr(self.unicast):
+                    stuc_bf16x8(
+                        contribution_mailbox_multicast_address,
+                        mailbox_element * 2,
+                        packed_words,
+                        self.tp,
+                    )
+                else:
+                    stmc_bf16x8(
+                        contribution_mailbox_multicast_address + mailbox_element * 2,
+                        packed_words,
+                    )
 
         if cutlass.const_expr(not self.release_before_store):
             if block == 0 and tidx == 0:
@@ -294,6 +327,7 @@ def __init__(
         enable_pdl: bool,
         prefetch_group: int,
         fp32_weights: bool = False,
+        unicast: bool = False,
     ) -> None:
         if hidden <= 0 or hidden % QUAD_BF16:
             raise ValueError("hidden must be a positive multiple of 4")
@@ -309,6 +343,7 @@ def __init__(
         self.include_shared_expert = include_shared_expert
         self.load_shared_expert_before_pdl = load_shared_expert_before_pdl
         self.enable_pdl = enable_pdl
+        self.unicast = unicast
         self.prefetch_group = prefetch_group
         self.prefetch_groups = (top_k + prefetch_group - 1) // prefetch_group
         self.fp32_weights = fp32_weights
@@ -516,10 +551,18 @@ def kernel(
             mailbox_element = (
                 (Int64(stage) * self.tp + self.rank) * self.capacity_m + Int64(token)
             ) * self.hidden + Int64(fragment) * QUAD_BF16
-            stmc_bf16x4(
-                contribution_mailbox_multicast_address + mailbox_element * 2,
-                packed,
-            )
+            if cutlass.const_expr(self.unicast):
+                stuc_bf16x4(
+                    contribution_mailbox_multicast_address,
+                    mailbox_element * 2,
+                    packed,
+                    self.tp,
+                )
+            else:
+                stmc_bf16x4(
+                    contribution_mailbox_multicast_address + mailbox_element * 2,
+                    packed,
+                )
 
         if block == 0 and tidx == 0:
             store_global_u32(
@@ -1004,6 +1047,28 @@ def __init__(
         # a 640-fragment token in one fully used trip, and the finalize stages
         # all top_k rows at once.
         with torch.accelerator.device_index(device.index):
+            self._mailbox = symm_mem.empty(
+                (LAMPORT_GENERATIONS, tp, capacity, hidden),
+                dtype=torch.bfloat16,
+                device=device,
+            )
+            self._mailbox_handle = symm_mem.rendezvous(self._mailbox, group)
+            multicast = int(self._mailbox_handle.multicast_ptr or 0)
+            # Without NVLink multicast (or when forced), the publish kernels
+            # store into every rank's mailbox through a table of their bases.
+            self.unicast = envs.VLLM_DSV41_MHC_UNICAST or not multicast
+            if self.unicast:
+                self._peer_table = torch.tensor(
+                    self._mailbox_handle.buffer_ptrs, dtype=torch.int64, device=device
+                )
+                multicast = self._peer_table.data_ptr()
+            elif multicast % 16:
+                raise RuntimeError("misaligned NVLink multicast mapping")
+            self._multicast = multicast
+            logger.info_once(
+                "DSV4.1 mHC: Lamport all-reduce publishes with %s.",
+                "unicast NVLink stores" if self.unicast else "NVLink multicast",
+            )
             self._publish = cute.compile(
                 _SharedOnlyPublishDeviceKernel(
                     hidden=hidden,
@@ -1014,6 +1079,7 @@ def __init__(
                     threads=128,
                     release_before_store=False,
                     enable_pdl=True,
+                    unicast=self.unicast,
                 ),
                 make_fake_dynamic_compact_tensor(
                     BFloat16, alignment=16, divisibility=hidden
@@ -1038,6 +1104,7 @@ def __init__(
                     enable_pdl=True,
                     prefetch_group=top_k,
                     fp32_weights=True,
+                    unicast=self.unicast,
                 ),
                 make_fake_dynamic_compact_tensor(
                     BFloat16, alignment=16, divisibility=hidden
@@ -1137,16 +1204,6 @@ def compile_collective(quantize: bool, residual_only: bool):
                     hidden, dtype=torch.bfloat16, device=device
                 )
 
-            self._mailbox = symm_mem.empty(
-                (LAMPORT_GENERATIONS, tp, capacity, hidden),
-                dtype=torch.bfloat16,
-                device=device,
-            )
-            self._mailbox_handle = symm_mem.rendezvous(self._mailbox, group)
-            multicast = int(self._mailbox_handle.multicast_ptr or 0)
-            if not multicast or multicast % 16:
-                raise RuntimeError("NVLink multicast mapping is unavailable")
-            self._multicast = multicast
             # Every slot starts as the Lamport sentinel, BF16 negative zero.
             self._mailbox.view(torch.int16).fill_(-32768)
             self._stage_state = torch.zeros(2, dtype=torch.int32, device=device)