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..0a40027a04a7 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 ) @@ -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, 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..a6c79f14ec11 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,15 @@ 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. TP4 is one GB300 + # tray; TP8 spans two trays of an NVL72 NVLink domain. + return ( + parallel.tensor_parallel_size in (4, 8) + 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..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. @@ -23,16 +25,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..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 @@ -121,7 +124,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 +145,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/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/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, 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 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():