From a33850a4c3311c616d20a2cf74b5b887d951ad1f Mon Sep 17 00:00:00 2001 From: Michael Schmid Date: Mon, 26 Oct 2026 03:29:22 +0000 Subject: [PATCH 1/3] [Perf][DSv4.1] Fused all-reduce + mHC without NVLink multicast The Lamport all-reduce + mHC kernel published each rank's contribution with multimem stores, so it needed an NVLink multicast mapping of its mailbox and the specialized decode refused to start without one. An NVLink partition has a fixed number of multicast groups shared by every engine in it, so engines that start later run out of them. When the symmetric-memory rendezvous maps no multicast object (on any rank), the publish kernels and the publishing wo_b GEMM now store into every rank's mailbox through a table of their base addresses. Multicast stays the default where it is available; VLLM_DSV41_MHC_UNICAST=1 forces the unicast path. Co-authored-by: Claude Opus 5.5 Signed-off-by: Michael Schmid --- .../specialized_decode/test_numerics.py | 10 +- vllm/envs.py | 8 ++ .../nvidia/ops/cute_dsl/all_reduce_mhc.py | 109 +++++++++++++----- .../nvidia/ops/cute_dsl/primitives.py | 49 ++++++++ vllm/models/deepseek_v41/nvidia/ops/mhc.py | 23 ++-- .../nvidia/specialized_decode/AGENTS.md | 19 ++- .../nvidia/specialized_decode/__init__.py | 5 +- .../specialized_decode/boundary/all_reduce.py | 8 +- .../boundary/publishing_gemm.py | 14 ++- 9 files changed, 185 insertions(+), 60 deletions(-) diff --git a/tests/models/deepseek_v41/nvidia/specialized_decode/test_numerics.py b/tests/models/deepseek_v41/nvidia/specialized_decode/test_numerics.py index 3442d862e491..a49060d80b80 100644 --- a/tests/models/deepseek_v41/nvidia/specialized_decode/test_numerics.py +++ b/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 diff --git a/vllm/envs.py b/vllm/envs.py index a7bde6b336fa..6119d6af6653 100755 --- a/vllm/envs.py +++ b/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") ), diff --git a/vllm/models/deepseek_v41/nvidia/ops/cute_dsl/all_reduce_mhc.py b/vllm/models/deepseek_v41/nvidia/ops/cute_dsl/all_reduce_mhc.py index 191a928fe2e9..46b9bb11abe7 100644 --- a/vllm/models/deepseek_v41/nvidia/ops/cute_dsl/all_reduce_mhc.py +++ b/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) diff --git a/vllm/models/deepseek_v41/nvidia/ops/cute_dsl/primitives.py b/vllm/models/deepseek_v41/nvidia/ops/cute_dsl/primitives.py index 16d5f4721a33..2364021ebe6c 100644 --- a/vllm/models/deepseek_v41/nvidia/ops/cute_dsl/primitives.py +++ b/vllm/models/deepseek_v41/nvidia/ops/cute_dsl/primitives.py @@ -318,6 +318,55 @@ def shuffle_sync_idx_u32( ) +def _stuc(peer_table, offset, words, peers, loc, ip) -> None: + """Store 1, 2 or 4 32-bit words at ``offset`` into each of ``peers`` + buffers whose base addresses are the u64 entries of ``peer_table``: the + unicast twin of ``stmc_bf16x*`` for deployments without NVLink multicast.""" + vec = {1: "b32", 2: "v2.b32", 4: "v4.b32"}[len(words)] + regs = ", ".join(f"${i + 2}" for i in range(len(words))) + regs = regs if len(words) == 1 else "{" + regs + "}" + body = ["{", ".reg .u64 base, addr;"] + for peer in range(peers): + body += [ + f"ld.global.nc.u64 base, [$0+{8 * peer}];", + "add.u64 addr, base, $1;", + f"st.relaxed.sys.global.{vec} [addr], {regs};", + ] + body.append("}") + _asm( + None, + [peer_table.ir_value(loc=loc, ip=ip), offset.ir_value(loc=loc, ip=ip), *words], + "\n\t".join(body), + "l,l" + ",r" * len(words), + side_effects=True, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def stuc_bf16x2( + peer_table: Int64, offset: Int64, packed: Uint32, peers: int, *, loc=None, ip=None +) -> None: + _stuc(peer_table, offset, [packed.ir_value(loc=loc, ip=ip)], peers, loc, ip) + + +@dsl_user_op +def stuc_bf16x4( + peer_table: Int64, offset: Int64, values, peers: int, *, loc=None, ip=None +) -> None: + words = [values[i].ir_value(loc=loc, ip=ip) for i in range(2)] + _stuc(peer_table, offset, words, peers, loc, ip) + + +@dsl_user_op +def stuc_bf16x8( + peer_table: Int64, offset: Int64, values, peers: int, *, loc=None, ip=None +) -> None: + words = [values[i].ir_value(loc=loc, ip=ip) for i in range(4)] + _stuc(peer_table, offset, words, peers, loc, ip) + + @dsl_user_op def stmc_bf16x2(address: Int64, packed: Uint32, *, loc=None, ip=None) -> None: _asm( diff --git a/vllm/models/deepseek_v41/nvidia/ops/mhc.py b/vllm/models/deepseek_v41/nvidia/ops/mhc.py index edfdbd1ee40d..36e6160479b1 100644 --- a/vllm/models/deepseek_v41/nvidia/ops/mhc.py +++ b/vllm/models/deepseek_v41/nvidia/ops/mhc.py @@ -3,7 +3,7 @@ """Dispatch DSV4.1 mHC operations and overlap coefficient generation.""" from functools import partial -from typing import TYPE_CHECKING, cast +from typing import TYPE_CHECKING import torch @@ -27,8 +27,6 @@ ) if TYPE_CHECKING: - from vllm.distributed.device_communicators.cuda_communicator import CudaCommunicator - from .cute_dsl import AllReduceMHC logger = init_logger(__name__) @@ -59,17 +57,14 @@ def supports_mhc_overlap(vllm_config: VllmConfig) -> bool: def supports_mhc_all_reduce(vllm_config: VllmConfig) -> bool: parallel = vllm_config.parallel_config config = vllm_config.model_config.hf_config - if ( - parallel.tensor_parallel_size != 4 - or parallel.enable_expert_parallel - or config.hidden_size != 5120 - or config.hc_mult != 4 - ): - return False - # The custom all-reduce's MNNVL buffers succeed exactly when NVLink - # multicast does, which the fused kernel's own mailbox needs too. - comm = cast("CudaCommunicator", get_tp_group().device_communicator).ca_comm - return comm is not None and bool(comm.mnnvl_lamport_ag_multicast_ptr) + # AllReduceMHC publishes over NVLink multicast when its mailbox gets a + # multicast mapping and with unicast stores otherwise. + return ( + parallel.tensor_parallel_size == 4 + and not parallel.enable_expert_parallel + and config.hidden_size == 5120 + and config.hc_mult == 4 + ) def init_mhc_all_reduce(vllm_config: VllmConfig) -> None: diff --git a/vllm/models/deepseek_v41/nvidia/specialized_decode/AGENTS.md b/vllm/models/deepseek_v41/nvidia/specialized_decode/AGENTS.md index 0986fd485687..229f3e25ffdb 100644 --- a/vllm/models/deepseek_v41/nvidia/specialized_decode/AGENTS.md +++ b/vllm/models/deepseek_v41/nvidia/specialized_decode/AGENTS.md @@ -23,16 +23,15 @@ startup on every TP rank with the missing requirements: `FULL_DECODE_ONLY`). - No `torch.compile` of the model, no sleep mode. - No LoRA, KV transfer or batch invariance. -- DeepGEMM and NVLink multicast for the fused all-reduce + mHC. - -This check runs once, at startup. The configuration is the same on every -rank, but the last requirement is probed per rank: the custom all-reduce sets -up its NVLink multicast buffers on each rank and catches a failure there, so -one rank can be without them while the others have them. The ranks' results -are therefore combined with a MIN all-reduce over the TP group: if any rank -misses a requirement, every rank raises, rather than some ranks running the -specialized decode and others the stock one, which issue different -collectives and would hang. Which forward steps use the executor depends +- DeepGEMM for the fused all-reduce + mHC. Its Lamport mailbox publishes + over NVLink multicast when the symmetric-memory rendezvous maps one on every + rank, and otherwise with unicast stores into each rank's mailbox + (`VLLM_DSV41_MHC_UNICAST=1` forces the unicast path). + +This check runs once, at startup. The ranks' results are combined with a MIN +all-reduce over the TP group: if any rank misses a requirement, every rank +raises, rather than some ranks running the specialized decode and others the +stock one, which issue different collectives and would hang. Which forward steps use the executor depends only on the batch, which every TP rank receives identically, so that needs no collective. diff --git a/vllm/models/deepseek_v41/nvidia/specialized_decode/__init__.py b/vllm/models/deepseek_v41/nvidia/specialized_decode/__init__.py index 1d93d5e64f7d..310e66c25f12 100644 --- a/vllm/models/deepseek_v41/nvidia/specialized_decode/__init__.py +++ b/vllm/models/deepseek_v41/nvidia/specialized_decode/__init__.py @@ -121,7 +121,7 @@ def unsupported(vllm_config: VllmConfig) -> list[str]: and vllm_config.kv_transfer_config is None and not envs.VLLM_BATCH_INVARIANT ), - "DeepGEMM and NVLink multicast for the fused all-reduce + mHC": ( + "DeepGEMM for the fused all-reduce + mHC": ( supports_mhc_overlap(vllm_config) and supports_mhc_all_reduce(vllm_config) ), } @@ -142,8 +142,7 @@ def specialized_decode_config(vllm_config: VllmConfig) -> DecodeConfig | None: raise ValueError("Supported decode_backend values are auto/latency_specialized") missing = unsupported(vllm_config) supported = torch.tensor([0 if missing else 1], dtype=torch.int32) - # The NVLink multicast check is rank-local (the custom all-reduce catches a - # failed setup on each rank), and ranks that disagree would hang. + # Ranks that disagree would issue different collectives and hang. dist.all_reduce(supported, op=dist.ReduceOp.MIN, group=get_tp_group().cpu_group) if not supported.item(): reason = "; ".join(missing) or "another tensor parallel rank to support it" diff --git a/vllm/models/deepseek_v41/nvidia/specialized_decode/boundary/all_reduce.py b/vllm/models/deepseek_v41/nvidia/specialized_decode/boundary/all_reduce.py index 0bb4cf72b243..20bda5dd08ac 100644 --- a/vllm/models/deepseek_v41/nvidia/specialized_decode/boundary/all_reduce.py +++ b/vllm/models/deepseek_v41/nvidia/specialized_decode/boundary/all_reduce.py @@ -38,11 +38,13 @@ def __init__(self, config: PreTrainedConfig, quantize_layer_input: bool) -> None self._published: torch.Tensor | None = None def publish_target(self) -> tuple[int, int]: - """(multicast mailbox address, stage-state address) for a producer + """(mailbox publish address, stage-state address) for a producer that stores its contribution itself, as ``_SharedOnlyPublishDeviceKernel`` does: generation ``stage_state[0]``, slot ``[rank][token]``, BF16 with - negative zero cleared, then ``stage_state[1] = stage_state[0]``. Both - stay valid for this instance's lifetime.""" + negative zero cleared, then ``stage_state[1] = stage_state[0]``. The + publish address is the mailbox's multicast address, or with ``unicast`` + the address of the table of every rank's mailbox base. Both stay valid + for this instance's lifetime.""" return self._multicast, self._stage_state.data_ptr() def mark_published(self, x: torch.Tensor) -> None: diff --git a/vllm/models/deepseek_v41/nvidia/specialized_decode/boundary/publishing_gemm.py b/vllm/models/deepseek_v41/nvidia/specialized_decode/boundary/publishing_gemm.py index 87e7e45fa099..ed8c7c21a892 100644 --- a/vllm/models/deepseek_v41/nvidia/specialized_decode/boundary/publishing_gemm.py +++ b/vllm/models/deepseek_v41/nvidia/specialized_decode/boundary/publishing_gemm.py @@ -112,6 +112,7 @@ sanitize_negative_zero_u32x4, shuffle_sync_idx_u32, stmc_bf16x8, + stuc_bf16x8, ) from .. import LOW_BATCH_MAX_ROWS from ..attention.mxfp8 import Tactic, mxfp8_eligible @@ -162,6 +163,7 @@ def __init__( rank: int, capacity_m: int, hidden: int, + unicast: bool = False, **kwargs, ) -> None: super().__init__(**kwargs) @@ -173,6 +175,7 @@ def __init__( self.mailbox_multicast = mailbox_multicast self.stage_state = stage_state self.tp = tp + self.unicast = unicast self.rank = rank self.capacity_m = capacity_m self.hidden = hidden @@ -213,7 +216,15 @@ def publish_value( (Int64(stage) * self.tp + self.rank) * self.capacity_m + Int64(token) ) * self.hidden + Int64(column) - stmc_bf16x8(Int64(self.mailbox_multicast) + element * 2, clean) + if cutlass.const_expr(self.unicast): + stuc_bf16x8( + Int64(self.mailbox_multicast), + Int64(element) * 2, + clean, + self.tp, + ) + else: + stmc_bf16x8(Int64(self.mailbox_multicast) + element * 2, clean) @cute.jit def epilog_warp( @@ -536,6 +547,7 @@ def _compile(self, tactic: Tactic, device: torch.device): mailbox_multicast=multicast, stage_state=stage_state, tp=all_reduce.tp, + unicast=all_reduce.unicast, rank=all_reduce.rank, capacity_m=all_reduce.capacity, hidden=all_reduce.hidden_size, From 83fbb9fc96e7543bf845278c24fecc8476f912e4 Mon Sep 17 00:00:00 2001 From: Michael Schmid Date: Tue, 27 Oct 2026 02:57:08 +0000 Subject: [PATCH 2/3] [Perf][DSv4.1] Specialized decode at TP8 on GB300 Allow tensor parallel size 8 for the DeepSeek-V4.1 specialized decode and its fused all-reduce + mHC: two GB300 trays of an NVL72 NVLink domain. At TP8 each rank holds a 288-column slice of the routed experts' 2304-wide intermediate, padded to 384 by the TRT-LLM-gen weight conversion. The routed expert GEMMs now take the slice from the weights' padded width (576/640 at TP4, 288/384 at TP8): GEMM1 uses ceil(2 * slice / 128) m tiles (5 at TP8, whose last 32 H columns come from zero weights) and zeroes the remaining 64 padding columns from where the tiles end. The routed-expert test compares against FlashInfer's TRT-LLM-gen MoE at both slices. Signed-off-by: Michael Schmid --- .../nvidia/specialized_decode/test_kernels.py | 40 +++++++++++---- vllm/models/deepseek_v41/nvidia/ops/mhc.py | 5 +- .../nvidia/specialized_decode/AGENTS.md | 6 ++- .../nvidia/specialized_decode/__init__.py | 9 ++-- .../moe/routed_expert_gemm.py | 51 ++++++++++++------- .../moe/routed_expert_gemm_cute.py | 19 +++++-- 6 files changed, 90 insertions(+), 40 deletions(-) diff --git a/tests/models/deepseek_v41/nvidia/specialized_decode/test_kernels.py b/tests/models/deepseek_v41/nvidia/specialized_decode/test_kernels.py index 12ecc14a1303..8964e3bb1828 100644 --- a/tests/models/deepseek_v41/nvidia/specialized_decode/test_kernels.py +++ b/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 ) diff --git a/vllm/models/deepseek_v41/nvidia/ops/mhc.py b/vllm/models/deepseek_v41/nvidia/ops/mhc.py index 36e6160479b1..a6c79f14ec11 100644 --- a/vllm/models/deepseek_v41/nvidia/ops/mhc.py +++ b/vllm/models/deepseek_v41/nvidia/ops/mhc.py @@ -58,9 +58,10 @@ def supports_mhc_all_reduce(vllm_config: VllmConfig) -> bool: parallel = vllm_config.parallel_config config = vllm_config.model_config.hf_config # AllReduceMHC publishes over NVLink multicast when its mailbox gets a - # multicast mapping and with unicast stores otherwise. + # multicast mapping and with unicast stores otherwise. TP4 is one GB300 + # tray; TP8 spans two trays of an NVL72 NVLink domain. return ( - parallel.tensor_parallel_size == 4 + parallel.tensor_parallel_size in (4, 8) and not parallel.enable_expert_parallel and config.hidden_size == 5120 and config.hc_mult == 4 diff --git a/vllm/models/deepseek_v41/nvidia/specialized_decode/AGENTS.md b/vllm/models/deepseek_v41/nvidia/specialized_decode/AGENTS.md index 229f3e25ffdb..5a7cd3425d3a 100644 --- a/vllm/models/deepseek_v41/nvidia/specialized_decode/AGENTS.md +++ b/vllm/models/deepseek_v41/nvidia/specialized_decode/AGENTS.md @@ -14,8 +14,10 @@ startup on every TP rank with the missing requirements: - GPU: GB300 (SM103). - Runner: Model Runner V2. -- Parallelism: TP=4; no pipeline, data, expert or context parallelism, no - micro-batching. +- Parallelism: TP=4 (one GB300 tray) or TP=8 (two trays of an NVL72 NVLink + domain; each rank holds a 288-column slice of the routed experts' + intermediate, padded to 384); no pipeline, data, expert or context + parallelism, no micro-batching. - Speculative decoding: DSpark, without adaptive verification. - Attention: `FLASHINFER_MLA_SPARSE_DSV41`, fp8 KV cache, MXFP4 indexer with sparse logits. diff --git a/vllm/models/deepseek_v41/nvidia/specialized_decode/__init__.py b/vllm/models/deepseek_v41/nvidia/specialized_decode/__init__.py index 310e66c25f12..fe21f110d43b 100644 --- a/vllm/models/deepseek_v41/nvidia/specialized_decode/__init__.py +++ b/vllm/models/deepseek_v41/nvidia/specialized_decode/__init__.py @@ -36,6 +36,9 @@ MHC_OVERLAP_MAX_ROWS = 48 # Rows of the low-batch GEMMs: one DSpark verification step of up to 8 rows # (6 at one request). +# TP4 is one GB300 tray; TP8 spans two trays of an NVL72 NVLink domain, with +# each rank holding a 288-column slice of the routed experts' intermediate. +SUPPORTED_TP_SIZES = (4, 8) LOW_BATCH_MAX_ROWS = 8 # The FlashInfer (requirements/cuda.txt) whose internals the publishing wo_b # GEMM, the MXFP8 layer input and the pinned MXFP8 tactics were checked @@ -82,9 +85,9 @@ def unsupported(vllm_config: VllmConfig) -> list[str]: requirements = { "a GB300 (SM103) GPU": current_platform.is_device_capability(103), "Model Runner V2": vllm_config.use_v2_model_runner, - "tensor parallel size 4 without pipeline, data, expert or context " - "parallelism or micro-batching": ( - parallel.tensor_parallel_size == 4 + "tensor parallel size 4 or 8 without pipeline, data, expert or " + "context parallelism or micro-batching": ( + parallel.tensor_parallel_size in SUPPORTED_TP_SIZES and parallel.pipeline_parallel_size == 1 and parallel.data_parallel_size == 1 and not parallel.enable_expert_parallel diff --git a/vllm/models/deepseek_v41/nvidia/specialized_decode/moe/routed_expert_gemm.py b/vllm/models/deepseek_v41/nvidia/specialized_decode/moe/routed_expert_gemm.py index f7120d8948dc..0db920bf972c 100644 --- a/vllm/models/deepseek_v41/nvidia/specialized_decode/moe/routed_expert_gemm.py +++ b/vllm/models/deepseek_v41/nvidia/specialized_decode/moe/routed_expert_gemm.py @@ -76,6 +76,9 @@ HIDDEN = 5120 SLICE = 576 PADDED = 640 +# The routed experts' per-rank intermediate slice by its padded width: 576 +# (padded to 640) at TP4 and 288 (padded to 384) at TP8. +SLICE_BY_PADDED = {640: 576, 384: 288} # The fused front is built for exactly the target layers' routing. EXPERTS = 384 TOP_K = 6 @@ -142,7 +145,11 @@ def __init__( prepare_finalize: FusedMoEPrepareAndFinalizeModular, quant_config: FusedMoEQuantConfig, ): - check_weights(w13, w13_scale, w2, w2_scale, clamp, top_k) + self.padded = int(w13.size(1)) // 2 + self.slice = SLICE_BY_PADDED.get(self.padded, SLICE) + check_weights( + w13, w13_scale, w2, w2_scale, clamp, top_k, self.slice, self.padded + ) self.experts = int(w13.size(0)) self.top_k = top_k index = w13.device.index @@ -151,8 +158,8 @@ def __init__( ) # Compile the kernels now, not inside a CUDA graph capture. _route_kernel(self.device_index) - _gemm1_kernel(self.device_index) - _gemm2_kernel(self.device_index) + _gemm1_kernel(self.device_index, self.slice, self.padded) + _gemm2_kernel(self.device_index, self.padded) # The modular kernel's input quantization, reused so the activations # match its path bit for bit. self.prepare_finalize = prepare_finalize @@ -266,11 +273,15 @@ def gemm1(self, routing: Routing) -> tuple[torch.Tensor, torch.Tensor]: 128-row scale atom per token tile, rows 0-7 used).""" device = routing.x_quant.device padded_rows = TILE * routing.ctas - h_values = torch.empty((padded_rows, PADDED), dtype=torch.uint8, device=device) + h_values = torch.empty( + (padded_rows, self.padded), dtype=torch.uint8, device=device + ) h_scales = torch.empty( - routing.ctas * tile_scale_bytes(PADDED), dtype=torch.uint8, device=device + routing.ctas * tile_scale_bytes(self.padded), + dtype=torch.uint8, + device=device, ) - _gemm1_kernel(self.device_index)( + _gemm1_kernel(self.device_index, self.slice, self.padded)( _pointer(cutlass.Float4E2M1FN, self.w13, 32), _pointer(cutlass.Float8E8M0FNU, self.w13_scale, 16), _pointer(cutlass.Float8E4M3FN, routing.x_quant, 16), @@ -298,7 +309,7 @@ def gemm2( dtype=torch.bfloat16, device=routing.x_quant.device, ) - _gemm2_kernel(self.device_index)( + _gemm2_kernel(self.device_index, self.padded)( _pointer(cutlass.Float4E2M1FN, self.w2, 32), _pointer(cutlass.Float8E8M0FNU, self.w2_scale, 16), _pointer(cutlass.Float8E4M3FN, h_values, 16), @@ -373,11 +384,11 @@ def _front_kernel(device_index: int, image: bool): @cache -def _gemm1_kernel(device_index: int): - """``RoutedExpertGemm1`` compiled for a device.""" +def _gemm1_kernel(device_index: int, slice_: int = SLICE, padded: int = PADDED): + """``RoutedExpertGemm1`` compiled for a device and slice.""" with torch.accelerator.device_index(device_index): return cute.compile( - RoutedExpertGemm1(HIDDEN, SLICE, PADDED), + RoutedExpertGemm1(HIDDEN, slice_, padded), _fake(cutlass.Float4E2M1FN, 32), _fake(cutlass.Float8E8M0FNU, 16), _fake(cutlass.Float8E4M3FN, 16), @@ -396,11 +407,11 @@ def _gemm1_kernel(device_index: int): @cache -def _gemm2_kernel(device_index: int): - """``RoutedExpertGemm2`` compiled for a device.""" +def _gemm2_kernel(device_index: int, padded: int = PADDED): + """``RoutedExpertGemm2`` compiled for a device and padded intermediate.""" with torch.accelerator.device_index(device_index): return cute.compile( - RoutedExpertGemm2(HIDDEN, PADDED), + RoutedExpertGemm2(HIDDEN, padded), _fake(cutlass.Float4E2M1FN, 32), _fake(cutlass.Float8E8M0FNU, 16), _fake(cutlass.Float8E4M3FN, 16), @@ -428,15 +439,17 @@ def check_weights( w2_scale: torch.Tensor, clamp: torch.Tensor | None, top_k: int, + slice_: int = SLICE, + padded: int = PADDED, ) -> None: """Raise ``ValueError`` unless the weights have the layout the kernels read and the routing fits the route kernel.""" experts = w13.size(0) for tensor, shape in ( - (w13, (experts, 2 * PADDED, HIDDEN // 2)), - (w13_scale, (experts, 2 * PADDED, HIDDEN // 32)), - (w2, (experts, HIDDEN, PADDED // 2)), - (w2_scale, (experts, HIDDEN, PADDED // 32)), + (w13, (experts, 2 * padded, HIDDEN // 2)), + (w13_scale, (experts, 2 * padded, HIDDEN // 32)), + (w2, (experts, HIDDEN, padded // 2)), + (w2_scale, (experts, HIDDEN, padded // 32)), ): if ( tuple(tensor.shape) != shape @@ -464,8 +477,8 @@ def check_weights( # GEMM1 never reads w13's padding tile (intermediate 576..639), so the # padding must be zero, as must the matching K columns of w2. if ( - torch.count_nonzero(w13.view(torch.uint8)[:, 2 * SLICE :]).item() - or torch.count_nonzero(w2.view(torch.uint8)[:, :, SLICE // 2 :]).item() + torch.count_nonzero(w13.view(torch.uint8)[:, 2 * slice_ :]).item() + or torch.count_nonzero(w2.view(torch.uint8)[:, :, slice_ // 2 :]).item() ): raise ValueError("the intermediate padding of w13/w2 is not zero") diff --git a/vllm/models/deepseek_v41/nvidia/specialized_decode/moe/routed_expert_gemm_cute.py b/vllm/models/deepseek_v41/nvidia/specialized_decode/moe/routed_expert_gemm_cute.py index a3689a089f23..999a2eff5cac 100644 --- a/vllm/models/deepseek_v41/nvidia/specialized_decode/moe/routed_expert_gemm_cute.py +++ b/vllm/models/deepseek_v41/nvidia/specialized_decode/moe/routed_expert_gemm_cute.py @@ -1096,8 +1096,19 @@ def __init__( cta_k: int = 256, num_ab_stage: int = 6, ) -> None: - super().__init__(2 * padded, 2 * slice_ // 128, hidden, cta_k, num_ab_stage) + # 128 w13 rows (64 H columns) per m tile. TP4: 9 tiles cover the + # 576-column slice exactly. TP8: ceil(576 / 128) = 5 tiles cover the + # 288-column slice plus 32 columns of zero weights (H = 0 there); the + # epilogue zeroes the remaining 64 padding columns. + m_tiles = (2 * slice_ + 127) // 128 + super().__init__(2 * padded, m_tiles, hidden, cta_k, num_ab_stage) self.slice, self.padded = slice_, padded + self.computed = m_tiles * 64 # H columns the m tiles write + if padded - self.computed != 64: + raise ValueError( + f"slice {slice_} padded to {padded} leaves " + f"{padded - self.computed} padding columns, not 64" + ) @cute.jit def __call__( @@ -1269,12 +1280,14 @@ def store( thread = warp * 32 + lane pad_row = thread // 16 if pad_row < valid: - pad = (row0 + pad_row) * h_row + (self.slice + (thread % 16) * 4) // 2 + pad = (row0 + pad_row) * h_row + ( + self.computed + (thread % 16) * 4 + ) // 2 mH[pad] = Uint16(0) mH[pad + 1] = Uint16(0) if thread < 2 * TILE: pad_row = thread // 2 - pad_block = self.slice // 32 + thread % 2 + pad_block = self.computed // 32 + thread % 2 if pad_row < valid: mHS[ n_tile * tile_bytes From 0aec6150e1948d8b23121718d7ca1cdbcaea3ce0 Mon Sep 17 00:00:00 2001 From: Michael Schmid Date: Tue, 27 Oct 2026 03:12:14 +0000 Subject: [PATCH 3/3] [Perf][DSv4.1] Pin the low-latency MXFP8 GEMM tactics at TP8 and without autotuning The specialized decode pins one tactic per attention projection for 5-8 rows, because split-K tactics tuned in isolation are slower inside the decode step's PDL chain. The pins covered the TP4 shapes only and applied only while FlashInfer autotuned. Without autotuning (for instance at TP8 across two trays, where the cross-rank tuning sync is slow) the runner falls back to FlashInfer's generic first tactic, one CTA per 128 output columns: 14 CTAs for fused_wqa_wkv (N=1792) and 32 for wq_b at TP8 (N=4096). - Key the pins by (N, K, TP) and add TP8 (fused_wqa_wkv is replicated, so it shares its shape with TP4 but pins split-K 2 there; wq_b and wo_b halve). - Put the pinned tactic first in FlashInfer's tactic list for good, so the untuned runner and the publishing wo_b GEMM use it too. On DeepSeek-V4.1-Flash, TP8 over two GB300 trays with DSpark and autotuning off, one request's decode step went from 3,926 to 3,834 us (+2.4 % tok/s; same trays, 5-8 requests per variant). Signed-off-by: Michael Schmid --- .../nvidia/specialized_decode/test_kernels.py | 19 +++- .../specialized_decode/attention/mxfp8.py | 94 ++++++++++++++----- .../specialized_decode/runtime/warmup.py | 8 +- 3 files changed, 92 insertions(+), 29 deletions(-) diff --git a/tests/models/deepseek_v41/nvidia/specialized_decode/test_kernels.py b/tests/models/deepseek_v41/nvidia/specialized_decode/test_kernels.py index 8964e3bb1828..0a40027a04a7 100644 --- a/tests/models/deepseek_v41/nvidia/specialized_decode/test_kernels.py +++ b/tests/models/deepseek_v41/nvidia/specialized_decode/test_kernels.py @@ -537,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(): @@ -562,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, diff --git a/vllm/models/deepseek_v41/nvidia/specialized_decode/attention/mxfp8.py b/vllm/models/deepseek_v41/nvidia/specialized_decode/attention/mxfp8.py index 4014147ba0be..e2a8792206c6 100644 --- a/vllm/models/deepseek_v41/nvidia/specialized_decode/attention/mxfp8.py +++ b/vllm/models/deepseek_v41/nvidia/specialized_decode/attention/mxfp8.py @@ -10,7 +10,10 @@ The backend's tactics differ by under 1 us when tuned in isolation, but split-K tactics are slower inside the back-to-back PDL chain of a decode step, so the -5-8-row bucket of the GB300 TP4 shapes is tuned over one pinned tactic. +5-8-row bucket of the GB300 shapes is tuned over one pinned tactic per tensor +parallel size. Without FlashInfer autotuning the pinned tactic is still the +runner's first, rather than FlashInfer's generic default, whose one CTA per +128 output columns leaves most SMs idle for these narrow GEMMs. """ import importlib @@ -21,6 +24,7 @@ import torch from torch import nn +from vllm.distributed import get_tensor_model_parallel_world_size from vllm.model_executor.kernels.linear.mxfp8.flashinfer import ( FlashInferCutedslMxfp8LinearKernel, ) @@ -37,13 +41,17 @@ BACKEND = "cutedsl_low_latency" -# (cta_k, ab_stages, sfb_tmem_stages, split_k) per (N, K) on GB300 for 5-8 -# rows, chosen jointly for the chain of the three projections. +# (cta_k, ab_stages, sfb_tmem_stages, split_k) per (N, K, TP) on GB300 for 5-8 +# rows, chosen jointly for the chain of the three projections. fused_wqa_wkv is +# replicated, so it has the same shape at TP4 and TP8; wq_b and wo_b halve. Tactic = tuple[int, int, int, int] -PINNED_TACTICS: dict[tuple[int, int], Tactic] = { - (1792, 5120): (128, 10, 1, 4), - (8192, 1280): (256, 6, 1, 1), - (5120, 2048): (256, 4, 1, 2), +PINNED_TACTICS: dict[tuple[int, int, int], Tactic] = { + (1792, 5120, 4): (128, 10, 1, 4), + (8192, 1280, 4): (256, 6, 1, 1), + (5120, 2048, 4): (256, 4, 1, 2), + (1792, 5120, 8): (128, 10, 1, 2), + (4096, 1280, 8): (128, 10, 1, 2), + (5120, 1024, 8): (128, 8, 1, 1), } _PINNED_MIN_ROWS = LOW_BATCH_MAX_ROWS // 2 + 1 @@ -131,38 +139,72 @@ def run_zeros(linear: nn.Module, rows: int) -> torch.Tensor: return mxfp8_mm(linear, qa) -@contextmanager -def pinned_tactics() -> Iterator[None]: - """Offer FlashInfer's tuner only the pinned tactic for pinned problems. The - tuner's tactic list is a FlashInfer internal: off the checked version, - every tactic is tuned.""" +def _pinned(problem_sizes_mnkl: tuple[int, int, int, int], tp: int) -> Tactic | None: + """The pinned tactic of a runner problem at tensor parallel size ``tp``, if + any. The runner computes out.T: its problem is (N, M, K, L).""" + n, rows, k, _ = problem_sizes_mnkl + if not _PINNED_MIN_ROWS <= rows <= LOW_BATCH_MAX_ROWS: + return None + return PINNED_TACTICS.get((n, k, tp)) + + +def _tactics_module(): + """FlashInfer's low-latency GEMM module, whose ``autotune_tactics`` lists + a problem's tactics (a FlashInfer internal), or None off GB300 and the + checked FlashInfer version.""" if not ( current_platform.is_device_capability(103) and uses_checked_flashinfer("pinning the MXFP8 GEMM tactics") ): - yield - return + return None module = importlib.import_module( "flashinfer.gemm.kernels.cute_dsl.low_latency_blockscaled_gemm" ) - original: Callable[..., list[Tactic]] | None = getattr( - module, "autotune_tactics", None - ) - if original is None: - yield + return module if hasattr(module, "autotune_tactics") else None + + +def prefer_pinned_tactics() -> None: + """Put the pinned tactic first in FlashInfer's tactic list for pinned + problems, for good: the runner falls back to the first tactic when no tuned + one is cached, which is every call when FlashInfer autotuning is off, and + the publishing wo_b GEMM resolves its tactic the same way.""" + module = _tactics_module() + if module is None or getattr(module, "_dsv41_prefers_pinned", False): return - enumerate_tactics = original + enumerate_tactics: Callable[..., list[Tactic]] = module.autotune_tactics + tp = get_tensor_model_parallel_world_size() def autotune_tactics( problem_sizes_mnkl: tuple[int, int, int, int], *args: object, **kwargs: object ) -> list[Tactic]: tactics = enumerate_tactics(problem_sizes_mnkl, *args, **kwargs) - # The runner computes out.T: its problem is (N, M, K, L). - n, rows, k, _ = problem_sizes_mnkl - pinned = PINNED_TACTICS.get((n, k)) - if _PINNED_MIN_ROWS <= rows <= LOW_BATCH_MAX_ROWS and pinned in tactics: - return [pinned] - return tactics + pinned = _pinned(problem_sizes_mnkl, tp) + if pinned not in tactics: + return tactics + return [pinned] + [tactic for tactic in tactics if tactic != pinned] + + module.autotune_tactics = autotune_tactics # type: ignore[attr-defined] + module._dsv41_prefers_pinned = True # type: ignore[attr-defined] + + +@contextmanager +def pinned_tactics() -> Iterator[None]: + """Offer FlashInfer's tuner only the pinned tactic for pinned problems. The + tuner's tactic list is a FlashInfer internal: off the checked version, + every tactic is tuned.""" + module = _tactics_module() + if module is None: + yield + return + original: Callable[..., list[Tactic]] = module.autotune_tactics + tp = get_tensor_model_parallel_world_size() + + def autotune_tactics( + problem_sizes_mnkl: tuple[int, int, int, int], *args: object, **kwargs: object + ) -> list[Tactic]: + tactics = original(problem_sizes_mnkl, *args, **kwargs) + pinned = _pinned(problem_sizes_mnkl, tp) + return [pinned] if pinned in tactics else tactics module.autotune_tactics = autotune_tactics # type: ignore[attr-defined] try: diff --git a/vllm/models/deepseek_v41/nvidia/specialized_decode/runtime/warmup.py b/vllm/models/deepseek_v41/nvidia/specialized_decode/runtime/warmup.py index 4d7ae13e3158..bdab95f81183 100644 --- a/vllm/models/deepseek_v41/nvidia/specialized_decode/runtime/warmup.py +++ b/vllm/models/deepseek_v41/nvidia/specialized_decode/runtime/warmup.py @@ -21,7 +21,12 @@ from .. import FUSED_MHC_MAX_ROWS, LOW_BATCH_MAX_ROWS from ..attention.aux_projections import warmup_aux_projection_gemms -from ..attention.mxfp8 import pinned_tactics, run_zeros, shape_representatives +from ..attention.mxfp8 import ( + pinned_tactics, + prefer_pinned_tactics, + run_zeros, + shape_representatives, +) from ..attention.step import AttentionStep, key_weight, side_weights from ..logits.lm_head import autotune_low_batch_lm_head, warmup_low_batch_lm_head from ..moe.step import deferred_finalize @@ -38,6 +43,7 @@ def warm_up(decode: SpecializedDecode, vllm_config: VllmConfig) -> None: for linear in (step.attn.fused_wqa_wkv, step.attn.wq_b, step.attn.wo_b) ) rows = min(vllm_config.scheduler_config.max_num_batched_tokens, FUSED_MHC_MAX_ROWS) + prefer_pinned_tactics() with torch.inference_mode(): if decode.config.autotune: with _tuning():