diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index 4210569993b5..4a0280c518be 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.buildkite/test_areas/models_basic.yaml @@ -613,6 +613,9 @@ steps: - tests/models/deepseek_v41/nvidia/specialized_decode/test_kernels.py - tests/models/deepseek_v41/nvidia/specialized_decode/test_numerics.py - tests/models/deepseek_v41/nvidia/specialized_decode/test_layers.py + - tests/models/deepseek_v41/nvidia/specialized_decode/test_decode_topk.py + - tests/models/deepseek_v41/nvidia/specialized_decode/test_sparse_mla_cute.py + - tests/models/deepseek_v41/nvidia/specialized_decode/test_all_gather.py - tests/distributed/test_custom_all_reduce.py - tests/kernels/test_engram.py commands: @@ -624,6 +627,11 @@ steps: - pytest -v -s models/deepseek_v41/nvidia/specialized_decode/test_kernels.py - pytest -v -s models/deepseek_v41/nvidia/specialized_decode/test_numerics.py - pytest -v -s models/deepseek_v41/nvidia/specialized_decode/test_layers.py + - pytest -v -s models/deepseek_v41/nvidia/specialized_decode/test_decode_topk.py + - pytest -v -s models/deepseek_v41/nvidia/specialized_decode/test_sparse_mla_cute.py + - pytest -v -s models/deepseek_v41/nvidia/specialized_decode/test_all_gather.py + - pytest -v -s distributed/test_custom_all_reduce.py -k all_reduce_mhc + - pytest -v -s kernels/test_engram.py - label: ':amd: (MI355) Kimi-K3 Specialized Decode' key: kimi-k3-specialized-decode-amd diff --git a/tests/distributed/test_custom_all_reduce.py b/tests/distributed/test_custom_all_reduce.py index b943f457d2dd..7fa7f38e7693 100644 --- a/tests/distributed/test_custom_all_reduce.py +++ b/tests/distributed/test_custom_all_reduce.py @@ -70,7 +70,7 @@ def _all_reduce_mhc(monkeypatch, tp_size, pp_size, rank, distributed_init_port): init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port) ensure_model_parallel_initialized(tp_size, pp_size) op = AllReduceMHC( - hidden_size=5120, hc_mult=4, max_num_tokens=16, top_k=6, device=device + hidden_size=5120, hc_mult=4, max_num_tokens=256, top_k=6, device=device ) def check_eager_and_replayed(fused, check, halved, doubled): @@ -98,12 +98,16 @@ def mhc_inputs(n): weight = torch.randn(5120, device=device, dtype=torch.bfloat16) return residual, post, comb, pre, weight - def run(n): + def contribution(n): torch.manual_seed(42 + rank) x = torch.randn(n, 5120, device=device, dtype=torch.bfloat16) # Packed +0/-0 pairs collide with the Lamport sentinel. x[:, :16] = 0 x[:, 9:16:2] = -0.0 + return x + + def run(n): + x = contribution(n) residual, post, comb, pre, weight = mhc_inputs(n) def fused(): @@ -167,10 +171,26 @@ def check(outputs): check_eager_and_replayed(fused, check, gemm2, shared) - # Changing shapes cover shrinking and growing batches in one mailbox. - for n in (1, 6, 12, 8, 16, 3, 5, 2, 4, 1): - run(n) - run_finalize(n) + # Changing shapes cover shrinking and growing batches in one mailbox: + # LL only, BT only (200 rows take the wider owner reduce), then both + # protocols in turn, switching at 8 rows. + shapes = (1, 6, 12, 8, 16, 3, 5, 2, 4, 1) + for bt_min_rows, ns in ((257, shapes), (1, (*shapes, 200)), (8, shapes)): + op.bt_min_rows = op.bt_min_finalize_rows = bt_min_rows + for n in ns: + run(n) + run_finalize(n) + + # Both protocols give the same bits. + for n in (6, 16, 200): + x = contribution(n) + inputs = mhc_inputs(n) + op.bt_min_rows = n + 1 + ll = op(x, *inputs, 1e-6) + op.bt_min_rows = n + bt = op(x, *inputs, 1e-6) + assert torch.equal(ll[0], bt[0]) + assert torch.equal(ll[1], bt[1]) @pytest.mark.skipif( diff --git a/tests/models/deepseek_v41/nvidia/specialized_decode/test_all_gather.py b/tests/models/deepseek_v41/nvidia/specialized_decode/test_all_gather.py new file mode 100644 index 000000000000..991a91ff2735 --- /dev/null +++ b/tests/models/deepseek_v41/nvidia/specialized_decode/test_all_gather.py @@ -0,0 +1,264 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The specialized decode's TP all-gather (``LamportAllGather``) at TP4. + +* ``accepts`` by size: 16-byte multiples up to the capacity. +* Eager gathers of int32, BF16 and FP32 tensors whose sizes grow and shrink + from call to call (1 pack to the full capacity), one rank late each call, + bitwise against NCCL. +* The Lamport sentinel word (FP32 negative zero, INT32_MIN) arrives as 0 on + every rank, the sending rank included. +* CUDA graphs of three gathers of different sizes each, some with the ranks + skewed, replayed in a seeded random order with new contents every replay; + a device-side count of words that differ from the expected concatenation + must stay 0. A word left in the mailbox by an earlier call of another size + would show up here as a mismatch. + +Needs four SM100-family GPUs with NVLink multicast; about a minute. For a +stress run, run this file as a module +(``python -m tests.models.deepseek_v41.nvidia.specialized_decode.test_all_gather``) +with ``--replays 2000000``. With ``--custom-all-gather`` it replays the same +graphs on vLLM's custom all-gather (``CustomAllreduce.custom_all_gather``, the +MNNVL Lamport kernel) instead and only counts the mismatched words: that +kernel never waits for a rank's own multicast copy, and words of an earlier +call can reach a later one. +""" + +import argparse +import os +import random +import tempfile + +import pytest +import torch + +from vllm.platforms import current_platform + +TP = 4 +CAPACITY = 4 << 20 # bytes a rank, as the executor builds it +REPLAYS = 2000 +LATE_CYCLES = 20_000 +SENTINEL = -(2**31) +# Pack counts (16 bytes each) of the eager sequence: growing, shrinking, a +# partial last CTA, one pack and the full capacity. +EAGER_PACKS = (1, 37, 389, 128, 5000, 3, CAPACITY // 16, 2, 777, 1) +# Words a rank of each graph's three gathers. +GRAPH_WORDS = ( + (4, 2800, 148), + (20000, 12, 2048), + (512, 8192, 4), + (CAPACITY // 4, 4, 64), +) + + +def _late(rank: int, step: int) -> None: + """Delay one rank per step, a different one each step.""" + if rank == step % TP: + torch.cuda._sleep(LATE_CYCLES) + + +def _random_words(generator, words: int) -> torch.Tensor: + """Random int32 words, none of them the sentinel.""" + x = torch.randint( + -(2**31), + 2**31 - 1, + (words,), + generator=generator, + device="cuda", + dtype=torch.int64, + ).to(torch.int32) + return torch.where(x == SENTINEL, 1, x) + + +def _check_accepts(all_gather): + assert not all_gather.accepts(torch.empty(0, device="cuda")) + assert not all_gather.accepts(torch.empty(2, device="cuda")) # 8 bytes + assert all_gather.accepts(torch.empty(4, device="cuda")) + assert all_gather.accepts(torch.empty(CAPACITY // 4, device="cuda")) + assert not all_gather.accepts(torch.empty(CAPACITY // 4 + 4, device="cuda")) + with pytest.raises(ValueError): + all_gather(torch.empty(2, device="cuda")) + + +def _check_eager(all_gather, rank): + from vllm.distributed import get_tp_group + + group = get_tp_group() + generator = torch.Generator(device="cuda").manual_seed(rank) + for step, packs in enumerate(EAGER_PACKS): + words = _random_words(generator, packs * 4) + for dtype in (torch.int32, torch.bfloat16, torch.float32): + # One 16-byte pack a row. + x = words.view(dtype).view(packs, -1) + _late(rank, step) + got = all_gather(x) + want = group.all_gather(x, 0) + assert got.dtype == dtype and got.shape == want.shape, (packs, dtype) + assert torch.equal(got.view(torch.int32), want.view(torch.int32)), ( + packs, + dtype, + ) + + +def _check_sentinel(all_gather, rank): + from vllm.distributed import get_tp_group + + generator = torch.Generator(device="cuda").manual_seed(100 + rank) + x = _random_words(generator, 64) + x[rank::7] = SENTINEL + got = all_gather(x) + want = get_tp_group().all_gather(x, 0) + want = torch.where(want == SENTINEL, 0, want) + assert torch.equal(got, want) + negative_zero = torch.full((8,), -0.0, device="cuda") + gathered = all_gather(negative_zero) + assert torch.equal(gathered.view(torch.int32), torch.zeros_like(gathered).int()) + + +def _pattern(rank: int, kind: int, words: int) -> torch.Tensor: + """(rank, kind, index) as int32 bits, never the sentinel; the replay + counter goes into bits 20-25.""" + index = torch.arange(words, dtype=torch.int32, device="cuda") & 0xFFFFF + return (rank << 28) | (kind << 26) | index + + +def _check_replays(all_gather, rank, replays, stop=True): + """Gather graphs replayed in a random order; returns the mismatch count, + after the first 1000 replays with a mismatch if ``stop``.""" + from vllm.distributed.parallel_state import graph_capture + + counter = torch.zeros(1, dtype=torch.int32, device="cuda") + errors = torch.zeros(1, dtype=torch.int64, device="cuda") + graphs = [] + with graph_capture(device=torch.device("cuda")) as context: + for sizes in GRAPH_WORDS: + mine = [_pattern(rank, k, n) for k, n in enumerate(sizes)] + want = [ + torch.cat([_pattern(r, k, n) for r in range(TP)]) + for k, n in enumerate(sizes) + ] + inputs = [torch.empty_like(m) for m in mine] + for skew in (False, True): + + def body(mine=mine, want=want, inputs=inputs, skew=skew): + counter.add_(1) + bits = (counter & 0x3F) << 20 + outs = [] + for kind in range(3): + torch.add(mine[kind], bits, out=inputs[kind]) + if skew: + torch.cuda._sleep(LATE_CYCLES * ((rank + kind) % TP)) + outs.append(all_gather(inputs[kind])) + for kind in range(3): + errors.add_((outs[kind] != want[kind] + bits).sum()) + + body() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=context.stream): + body() + # The graph keeps raw pointers: keep its tensors alive too. + graphs.append((graph, (mine, want, inputs))) + torch.accelerator.synchronize() + assert not stop or errors.item() == 0, "eager warm-up" + order = random.Random(1234) + for start in range(0, replays, 1000): + for _ in range(min(1000, replays - start)): + graphs[order.randrange(len(graphs))][0].replay() + torch.accelerator.synchronize() + if stop and errors.item(): + break + return int(errors.item()) + + +def _worker(rank: int, port: int, replays: int, skip_path: str, custom: bool) -> None: + from tests.utils import init_test_distributed_environment + from vllm.config import VllmConfig, set_current_vllm_config + from vllm.distributed import cleanup_dist_env_and_memory, get_tp_group + from vllm.models.deepseek_v41.nvidia.ops.cute_dsl import LamportAllGather + + torch.accelerator.set_device_index(rank) + with set_current_vllm_config(VllmConfig()), torch.inference_mode(): + init_test_distributed_environment(TP, 1, rank, str(port), local_rank=rank) + ca_comm = getattr(get_tp_group().device_communicator, "ca_comm", None) + supported = torch.tensor( + [int(ca_comm is not None and bool(ca_comm.mnnvl_lamport_ag_multicast_ptr))] + ) + torch.distributed.all_reduce( + supported, torch.distributed.ReduceOp.MIN, get_tp_group().cpu_group + ) + if not supported.item(): + if rank == 0: + with open(skip_path, "w") as f: + f.write("no NVLink multicast") + cleanup_dist_env_and_memory() + return + if custom: + assert ca_comm is not None + + def all_gather(x): + # It takes floating-point inputs only; the int32 words pass + # bitwise as FP32. + out = ca_comm.custom_all_gather(x.view(torch.float32)) + assert out is not None, x.nbytes + return out.view(x.dtype) + + else: + all_gather = LamportAllGather( + rank_capacity_bytes=CAPACITY, device=torch.device("cuda", rank) + ) + _check_accepts(all_gather) + _check_eager(all_gather, rank) + _check_sentinel(all_gather, rank) + errors = _check_replays(all_gather, rank, replays, stop=not custom) + total = torch.tensor([errors]) + torch.distributed.all_reduce(total, group=get_tp_group().cpu_group) + print( + f"rank {rank}: {replays} replays ({3 * replays} gathers) on the " + f"{'custom' if custom else 'Lamport'} all-gather, {errors} mismatched " + f"words here, {int(total.item())} on all ranks", + flush=True, + ) + assert custom or total.item() == 0 + torch.accelerator.synchronize() + cleanup_dist_env_and_memory() + + +def run(replays: int, custom: bool = False) -> str | None: + """Run the checks on four GPUs (only the replays on the custom all-gather + if ``custom``); returns why they were skipped, if so.""" + from vllm.utils.network_utils import get_open_port + + with tempfile.TemporaryDirectory() as tmp: + skip_path = os.path.join(tmp, "skipped") + torch.multiprocessing.spawn( + _worker, args=(get_open_port(), replays, skip_path, custom), nprocs=TP + ) + if os.path.exists(skip_path): + with open(skip_path) as f: + return f.read() + return None + + +@pytest.mark.distributed(num_gpus=4) +@pytest.mark.skipif( + not current_platform.is_device_capability_family(100), reason="needs SM100" +) +def test_lamport_all_gather(): + if torch.accelerator.device_count() < TP: + pytest.skip("needs four GPUs") + skipped = run(REPLAYS) + if skipped: + pytest.skip(skipped) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) + parser.add_argument("--replays", type=int, default=REPLAYS) + parser.add_argument( + "--custom-all-gather", + action="store_true", + help="replay the graphs on vLLM's custom all-gather and count mismatches", + ) + args = parser.parse_args() + skipped = run(args.replays, args.custom_all_gather) + print(f"skipped: {skipped}" if skipped else "done") diff --git a/tests/models/deepseek_v41/nvidia/specialized_decode/test_decode_topk.py b/tests/models/deepseek_v41/nvidia/specialized_decode/test_decode_topk.py new file mode 100644 index 000000000000..765b5b119ade --- /dev/null +++ b/tests/models/deepseek_v41/nvidia/specialized_decode/test_decode_topk.py @@ -0,0 +1,270 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""``decode_topk`` (``decode_topk_cute.py``) against ``torch.topk``. + +Token mode: each row's output must be an exact top-k of its BF16 logits in +the kernel's key order (every index whose key exceeds the k-th largest key, +then indices whose key equals it), or ``0 .. n-1`` and -1 for rows of at most +k values. Rows cover the sampled path, short rows (the direct path), crowded +and constant values, values rising toward the row's end (survivors in the last +CTA), NaN and infinities, the forced direct path, and a sample that misses the +row's top values. Block mode must +select an exact top-k of the 8-token block maxima, as the stock +``select_candidate_blocks`` on the upcast logits does, with the newest block +pinned and -inf blocks written as -1. +""" + +import pytest +import torch + +from vllm.models.deepseek_v41.nvidia.specialized_decode.attention.decode_topk_cute import ( # noqa: E501 + SAMPLES, + _resident_clusters, + cluster_for, + decode_topk, + supports_decode_topk, +) +from vllm.platforms import current_platform + +pytestmark = pytest.mark.skipif( + not ( + current_platform.is_cuda() and current_platform.is_device_capability_family(100) + ), + reason="needs an SM100-family GPU", +) + +DEVICE = "cuda" +WIDTH = 262144 + + +def order_keys(values: torch.Tensor) -> torch.Tensor: + """The kernel's order-preserving key of each BF16 value, as int32.""" + bits = values.view(torch.int16).int() & 0xFFFF + return torch.where(bits >= 0x8000, bits ^ 0xFFFF, bits ^ 0x8000) + + +def make_logits(rows: int, kind: str, seed: int) -> torch.Tensor: + """[rows, WIDTH] BF16 logits of one distribution.""" + g = torch.Generator(device=DEVICE).manual_seed(seed) + x = torch.randn(rows, WIDTH, device=DEVICE, generator=g) + if kind == "crowded": # few distinct values: long exact ties + x = torch.randint(0, 6, (rows, WIDTH), device=DEVICE, generator=g).float() + elif kind == "constant": + x = torch.full((rows, WIDTH), 3.0, device=DEVICE) + elif kind == "rising": # the largest values at the row's end + x = x * 0.1 + torch.linspace(0, 8, WIDTH, device=DEVICE) + elif kind == "special": + x[:, 7] = float("nan") + x[:, 1000::977] = float("inf") + x[:, 2000::991] = float("-inf") + return x.bfloat16() + + +def check_tokens(logits: torch.Tensor, lengths: list[int], out: torch.Tensor, k: int): + """Assert each row's output is an exact top-k (module docstring).""" + for r, n in enumerate(lengths): + got = out[r] + if n <= k: + want = torch.arange(k, device=DEVICE, dtype=torch.int32) + want[n:] = -1 + assert torch.equal(got, want), f"row {r} (n {n})" + continue + assert int((got < 0).sum()) == 0 and int(got.max()) < n, ( + f"row {r}: ids out of range" + ) + assert got.unique().numel() == k, f"row {r}: repeated ids" + keys = order_keys(logits[r, :n]) + kth = keys.topk(k).values[-1] + chosen = keys[got.long()] + assert int((chosen < kth).sum()) == 0, f"row {r}: a key below the k-th" + assert int((chosen > kth).sum()) == int((keys > kth).sum()), ( + f"row {r}: missed a key above" + ) + + +@pytest.mark.parametrize("kind", ["normal", "crowded", "constant", "rising", "special"]) +# One row, 8 CTAs a row (12), 4 CTAs a row (20), the most rows (64). +@pytest.mark.parametrize("rows", [1, 12, 20, 64]) +@pytest.mark.parametrize(("k", "fallback"), [(512, False), (2048, False), (2048, True)]) +def test_token_topk(kind, rows, k, fallback): + lengths_all = [ + WIDTH, + 0, + k - 1, + k, + k + 1, + 5000, + 8192, + 8193, + 30000, + 162000, + 99991, + 250000, + ] + lengths = [lengths_all[(r * 5) % len(lengths_all)] for r in range(rows)] + logits = make_logits(rows, kind, rows * 31 + k) + assert supports_decode_topk(logits, rows, k, block_mode=False) + ends = torch.tensor(lengths, dtype=torch.int32, device=DEVICE) + out = torch.full((rows, k), 7, dtype=torch.int32, device=DEVICE) + decode_topk(logits, ends, out, block_mode=False, fallback=fallback) + torch.accelerator.synchronize() + check_tokens(logits, lengths, out, k) + + +def test_token_topk_strided_output_and_graph(): + """A wider output buffer (row stride > k) and a CUDA graph replay on new + logits.""" + rows, k = 12, 512 + lengths = [162000, 140000, 600, 9000] * 3 + logits = make_logits(rows, "normal", 1) + ends = torch.tensor(lengths, dtype=torch.int32, device=DEVICE) + buffer = torch.full((rows, k + 64), 7, dtype=torch.int32, device=DEVICE) + out = buffer[:, :k] + decode_topk(logits, ends, out, block_mode=False) + torch.accelerator.synchronize() + check_tokens(logits, lengths, out, k) + assert torch.all(buffer[:, k:] == 7) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + decode_topk(logits, ends, out, block_mode=False) + logits.copy_(make_logits(rows, "rising", 2)) + graph.replay() + torch.accelerator.synchronize() + check_tokens(logits, lengths, out, k) + + +@pytest.mark.parametrize( + "rows,wide,cluster", + [(4, 2 * WIDTH, 8), (20, 2 * WIDTH, 8), (4, 4 * WIDTH, 16), (12, 4 * WIDTH, 16)], +) +@pytest.mark.parametrize("fallback", [False, True]) +def test_wide_rows(rows, wide, cluster, fallback): + """Rows of up to 524,288 logits (the ratio-2 layers' width at a 1M + context), which only 8 CTAs a row hold, and of up to 1,048,576 (the + candidate source's), which take 16; at 20 and 12 rows more clusters than + a GB300 runs at once, so they run in waves.""" + lengths = [ + [wide, wide * 3 // 4 + 1, wide // 2 + 1, 9000][r % 4] for r in range(rows) + ] + g = torch.Generator(device=DEVICE).manual_seed(11) + logits = (torch.randn(rows, wide, device=DEVICE, generator=g) - 1.5).bfloat16() + assert supports_decode_topk(logits, rows, 512, block_mode=False) + assert supports_decode_topk(logits, rows, 2048, block_mode=True) + assert cluster_for(rows, wide, torch.accelerator.current_device_index()) == cluster + ends = torch.tensor(lengths, dtype=torch.int32, device=DEVICE) + out = torch.full((rows, 512), 7, dtype=torch.int32, device=DEVICE) + decode_topk(logits, ends, out, block_mode=False, fallback=fallback) + blocks = torch.full((rows, 2048), 7, dtype=torch.int32, device=DEVICE) + decode_topk(logits, ends, blocks, block_mode=True) + torch.accelerator.synchronize() + check_tokens(logits, lengths, out, 512) + for r, end in enumerate(lengths): + nblocks = (end + 7) // 8 + if nblocks <= 2048: # every block, in order (no -inf blocks here) + want = torch.arange(2048, device=DEVICE, dtype=torch.int32) + want[nblocks:] = -1 + assert torch.equal(blocks[r], want), f"row {r}" + continue + keys = block_keys(logits[r], end) + best = keys.topk(2048).values.sort().values + assert torch.equal(selected_keys(blocks[r], keys), best), f"row {r}" + + +def test_candidate_source_waves(): + """Rows that need 16 CTAs are taken for at most two waves of clusters.""" + waves = 2 * _resident_clusters(torch.accelerator.current_device_index(), 16) + logits = torch.empty(waves + 1, 4 * WIDTH, dtype=torch.bfloat16, device=DEVICE) + for block_mode, k in ((False, 512), (True, 2048)): + assert supports_decode_topk(logits, waves, k, block_mode) + assert not supports_decode_topk(logits, waves + 1, k, block_mode) + narrow = logits[:, : 2 * WIDTH] # 8 CTAs a row: up to 64 rows + assert supports_decode_topk(narrow, waves + 1, 512, block_mode=False) + + +def test_token_topk_sampler_miss(): + """Large values only at the sampled positions: the sampled cutoff keeps + fewer than k keys, so the kernel must recover through the direct path.""" + rows, k, n = 2, 512, WIDTH + cluster = cluster_for(rows, WIDTH, torch.accelerator.current_device_index()) + per_cta = SAMPLES // cluster + chunk = ((n + cluster - 1) // cluster + 7) // 8 * 8 + positions = [] + for c in range(cluster): + length = min(n - c * chunk, chunk) + valid = per_cta * length // chunk + positions += [ + c * chunk + (2 * t + 1) * length // (2 * valid) for t in range(valid) + ] + x = torch.full((rows, WIDTH), -1.0, device=DEVICE) + g = torch.Generator(device=DEVICE).manual_seed(7) + x[:, positions] = torch.randn(rows, len(positions), device=DEVICE, generator=g) + 10 + logits = x.bfloat16() + ends = torch.tensor([n] * rows, dtype=torch.int32, device=DEVICE) + out = torch.full((rows, k), 7, dtype=torch.int32, device=DEVICE) + decode_topk(logits, ends, out, block_mode=False) + torch.accelerator.synchronize() + check_tokens(logits, [n] * rows, out, k) + + +NEG_INF_KEY = 0x007F + + +def block_keys(logits: torch.Tensor, end: int) -> torch.Tensor: + """Keys of a row's valid 8-token block maxima, the newest pinned at +inf.""" + nblocks = (end + 7) // 8 + padded = torch.full((nblocks * 8,), float("-inf"), device=DEVICE) + padded[:end] = logits[:end].float() + scores = padded.view(nblocks, 8).amax(dim=1) + scores[nblocks - 1] = float("inf") + return order_keys(scores.bfloat16()) + + +def selected_keys(ids: torch.Tensor, keys: torch.Tensor) -> torch.Tensor: + """Sorted keys of a candidate row's selected blocks, -1 slots as -inf.""" + assert int(ids.max()) < keys.numel() + valid = ids[ids >= 0] + assert valid.unique().numel() == valid.numel(), "repeated ids" + chosen = torch.full_like(ids, NEG_INF_KEY) + chosen[ids >= 0] = keys[valid.long()] + return chosen.sort().values + + +@pytest.mark.parametrize("rows", [3, 12, 20]) +@pytest.mark.parametrize("fallback", [False, True]) +def test_candidate_blocks_match_stock(workspace_init, rows, fallback): + """Block mode selects an exact top-k of the block maxima, as the stock + ``select_candidate_blocks`` on the upcast logits does: the same multiset + of block keys (BF16 maxima tie often, so the block sets may differ only + among equal maxima), the pinned newest block, -inf blocks as -1.""" + from vllm.model_executor.kernels.attention.dsa.candidate_blocks import ( + select_candidate_blocks, + ) + + k, block = 2048, 8 + ends_all = [WIDTH, 162000, 3000 * 8 + 5, 2048 * 8, 100 * 8 + 3, 0, 70001, 131072] + ends_list = [ends_all[r % len(ends_all)] for r in range(rows)] + g = torch.Generator(device=DEVICE).manual_seed(rows) + logits = torch.randn(rows, WIDTH, device=DEVICE, generator=g).bfloat16() + logits[0, : 50 * block] = float("-inf") + last = rows - 1 + end = ends_list[last] + if end: + logits[last, (end - 1) // block * block : end] = -100.0 # lowest, but pinned + ends = torch.tensor(ends_list, dtype=torch.int32, device=DEVICE) + assert supports_decode_topk(logits, rows, k, block_mode=True) + got = torch.full((rows, k), 7, dtype=torch.int32, device=DEVICE) + decode_topk(logits, ends, got, block_mode=True, fallback=fallback) + want = torch.full((rows, k), 7, dtype=torch.int32, device=DEVICE) + select_candidate_blocks(logits.float(), None, ends, k, block, want) + torch.accelerator.synchronize() + for r, end in enumerate(ends_list): + nblocks = (end + 7) // 8 + if nblocks <= k: + # Every block in order, -inf ones as -1; the stock path pads alike. + assert torch.equal(got[r].sort().values, want[r].sort().values), f"row {r}" + continue + keys = block_keys(logits[r], end) + best = keys.topk(k).values.sort().values + assert torch.equal(selected_keys(got[r], keys), best), f"row {r}" + assert torch.equal(selected_keys(want[r], keys), best), f"row {r} (stock)" + assert int(got[r].max()) == nblocks - 1, f"row {r}: newest block not pinned" 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 c3d1deb741b4..f1b32dfbd919 100644 --- a/tests/models/deepseek_v41/nvidia/specialized_decode/test_kernels.py +++ b/tests/models/deepseek_v41/nvidia/specialized_decode/test_kernels.py @@ -61,9 +61,11 @@ HIDDEN, PADDED, SLICE, + TILE, TOP_K, FrontConfig, RoutedExpertGemm, + cta_words, ) from vllm.models.deepseek_v41.nvidia.specialized_decode.runtime.metadata_graphs import ( AttnMetadataCudaGraphs, @@ -91,24 +93,26 @@ def _capture(fn): return graph, out -def _bf16_ulps(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: - """Distance in BF16 ulps between two BF16 tensors.""" - - def ordered(x): - bits = x.contiguous().view(torch.int16).to(torch.int32) - return torch.where(bits < 0, -32768 - bits, bits) - - return (ordered(a) - ordered(b)).abs() +def _bf16_ulp(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + """One BF16 ulp at the larger magnitude of two BF16 tensors, in FP64.""" + _, exponent = torch.frexp(torch.maximum(a.abs(), b.abs()).double()) + return torch.ldexp(torch.ones_like(exponent, dtype=torch.float64), exponent - 8) # --------------------------------------------------------- routed experts -def _routing(rows: int, generator: torch.Generator) -> torch.Tensor: - """Top-k ids of ``rows`` rows drawn from about 22 experts, as a decode - step of six rows routes.""" - pool = torch.randperm(EXPERTS, generator=generator)[:22] - ids = [pool[torch.randperm(22, generator=generator)[:TOP_K]] for _ in range(rows)] +def _routing( + rows: int, generator: torch.Generator, pool_size: int = 22 +) -> torch.Tensor: + """Top-k ids of ``rows`` rows drawn from ``pool_size`` experts (22, as a + decode step of six rows routes); above eight rows experts own several + tiles, and with a pool of six every expert takes every row.""" + pool = torch.randperm(EXPERTS, generator=generator)[:pool_size] + ids = [ + pool[torch.randperm(pool_size, generator=generator)[:TOP_K]] + for _ in range(rows) + ] return torch.stack(ids).to(torch.int32).to(DEVICE) @@ -216,7 +220,7 @@ def _flashinfer_experts(experts, routing, output=None): enable_pdl=True, activation_type=activation_to_flashinfer_int(MoEActivation.SILU), output=output, - tune_max_num_tokens=8, + tune_max_num_tokens=max(routing.topk_ids.size(0), 8), ) if output is not None: return output @@ -241,18 +245,25 @@ 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): +@pytest.mark.parametrize( + ("rows", "pool_size"), + # Both sides of every GEMM width (8, 16, 32, 64, 128, 192 rows), one DSpark + # request (6 rows), and crowded routing (6 experts) at the larger widths. + [(r, 22) for r in (1, 6, 8, 9, 16, 17, 32, 33, 64, 65, 128, 129, 171, 192)] + + [(64, 6), (128, 6), (192, 6)], +) +def test_routed_experts_match_flashinfer(experts, rows, pool_size, 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.""" + TRT-LLM-gen MoE, and their top-k reduction within one BF16 ulp (plus FP32 + summation-order error) of its finalized output.""" generator = torch.Generator().manual_seed(rows) gemm = _gemm(experts) def inputs(): x = torch.randn(rows, HIDDEN, generator=generator).to(torch.bfloat16) weights = torch.softmax(torch.randn(rows, TOP_K, generator=generator), -1) - return x.to(DEVICE), _routing(rows, generator), weights.to(DEVICE) + ids = _routing(rows, generator, pool_size) + return x.to(DEVICE), ids, weights.to(DEVICE) x, ids, weights = inputs() @@ -275,7 +286,59 @@ def run(): ours = _fma_finalize( out.gemm2_permuted, out.expanded_idx_to_permuted_idx, out.expert_weights ) - assert _bf16_ulps(ours, finalized.cpu()).max() <= 1 + # FlashInfer may add the k products in another order (newer releases split + # them into groups), so allow FP32's reordering error on top of the ulp. + magnitude = (out.expert_weights.cpu().double()[..., None] * routed.cpu()).abs() + slack = 2 * (TOP_K + 1) * 2**-24 * magnitude.sum(1) + distance = (ours.double() - finalized.cpu().double()).abs() + assert (distance <= _bf16_ulp(ours, finalized.cpu()) + slack).all() + + +@pytest.mark.parametrize("rows", [6, 12, 48, 192]) +def test_routed_experts_replays_agree(experts, rows): + """GEMM2 waits for each token tile's GEMM1 items (the record's ready + words), not for the GEMM1 grid. One captured route + GEMM pair is + replayed 500 times on two alternating input sets (different rows and + experts), and each replay's routed rows must be bitwise + FlashInfer's for its set: a tile read before its H is complete would + hold the other set's H. The used tiles' ready words must count GEMM1's + m tiles. Six and twelve rows leave GEMM1 a partial last wave, as a + decode step does.""" + generator = torch.Generator().manual_seed(1000 + rows) + gemm = _gemm(experts) + + def inputs(): + x = torch.randn(rows, HIDDEN, generator=generator).to(torch.bfloat16) + weights = torch.softmax(torch.randn(rows, TOP_K, generator=generator), -1) + ids = _routing(rows, generator, min(EXPERTS, TOP_K * rows)) + return x.to(DEVICE), ids, weights.to(DEVICE) + + sets = [inputs(), inputs()] + x, ids, weights = (t.clone() for t in sets[0]) + + def load(index): + for tensor, value in zip((x, ids, weights), sets[index]): + tensor.copy_(value) + + def run(): + routing = gemm.route(x, ids, weights) + return routing, gemm.gemms(routing) + + graph, (routing, out) = _capture(run) + wants = [] + for index in range(2): + load(index) + graph.replay() + wants.append(_flashinfer_experts(experts, routing)) + for replay in range(500): + load(replay % 2) + graph.replay() + routed = out.gemm2_permuted[out.expanded_idx_to_permuted_idx.long()] + assert torch.equal(routed, wants[replay % 2]), replay + ready = 4 + 2 * cta_words(routing.ctas) + TILE * routing.ctas + tiles = int(routing.record[1]) + counts = routing.record[ready : ready + tiles] + assert torch.all(counts == 2 * SLICE // 128), counts SENTINEL = 128815 # First of five image sentinel token ids. @@ -299,8 +362,9 @@ def _front_inputs(rows: int, generator: torch.Generator, image_rows: int): return logits.to(DEVICE), x.to(torch.bfloat16).to(DEVICE), ids.to(DEVICE) -def _chain(gemm: RoutedExpertGemm, x, logits, ids): - """What the front replaces: the DSv4 top-k, then quantization and route.""" +def _chain(gemm: RoutedExpertGemm, x, logits, ids, is_padding=None): + """What the front replaces: the DSv4 top-k (with the router's padding + mask), then quantization and route.""" from vllm.model_executor.layers.fused_moe.router.dsv4_topk import dsv4_topk config = gemm.front_config @@ -309,6 +373,7 @@ def _chain(gemm: RoutedExpertGemm, x, logits, ids): config.bias, torch.int32, config.scale, + is_padding=is_padding, input_ids=ids, bias_vl=config.bias_vl, image_sentinel_lo=config.image_sentinel_lo, @@ -333,12 +398,18 @@ def _assert_same_routing(got, want): assert got.ctas == want.ctas -@pytest.mark.parametrize("scale", [1.0, 1.5]) -@pytest.mark.parametrize("image", [None, torch.int32, torch.int64]) -@pytest.mark.parametrize("rows", [1, 5, 6, 8]) +@pytest.mark.parametrize( + ("image", "scale"), [(None, 1.0), (torch.int32, 1.5), (torch.int64, 1.0)] +) +# Both sides of each build width (8, 16, 32, 64, 128, 192 rows), which also +# changes the rows a selecting CTA holds, and one DSpark request (6 rows). +@pytest.mark.parametrize( + "rows", [1, 6, 8, 9, 16, 17, 32, 33, 64, 65, 128, 129, 171, 192] +) def test_front_matches_the_chain(experts, rows, image, scale): """The fused front: ids (in tie order), FP32 weights, MXFP8 bytes and - scales, permutation and routing record bitwise those of the chain.""" + scales, permutation and routing record bitwise those of the chain, with + and without CUDA graph padding rows.""" generator = torch.Generator().manual_seed(rows) gemm = _gemm(experts) bias = 0.05 * torch.randn(EXPERTS, generator=generator) @@ -351,19 +422,29 @@ def test_front_matches_the_chain(experts, rows, image, scale): image_sentinel_lo=SENTINEL if image else 0, scale=scale, ) - for trial in range(4): + + def padding(count): + return torch.arange(rows, device=DEVICE) >= rows - count + + for trial in range(6): logits, x, ids = _front_inputs(rows, generator, image_rows=trial % (rows + 1)) if image is not None: ids = ids.to(image) - _assert_same_routing(gemm.front(x, logits, ids), _chain(gemm, x, logits, ids)) + # Trials 4 and 5: the last row, then the last half of the rows, padding. + mask = None if trial < 4 else padding(1 if trial == 4 else rows // 2) + _assert_same_routing( + gemm.front(x, logits, ids, mask), _chain(gemm, x, logits, ids, mask) + ) - # Replayed from a CUDA graph on fresh logits and rows. - graph, routing = _capture(lambda: gemm.front(x, logits, ids)) + # Replayed from a CUDA graph on fresh logits, rows and padding. + mask = padding(0) + graph, routing = _capture(lambda: gemm.front(x, logits, ids, mask)) fresh_logits, fresh_x, _ = _front_inputs(rows, generator, image_rows=0) logits.copy_(fresh_logits) x.copy_(fresh_x) + mask.copy_(padding(rows // 3)) graph.replay() - _assert_same_routing(routing, _chain(gemm, x, logits, ids)) + _assert_same_routing(routing, _chain(gemm, x, logits, ids, mask)) # ------------------------------------------------- candidate selection @@ -501,6 +582,8 @@ def __init__(self, n: int, k: int): ) self.weight = nn.Parameter(w_q, requires_grad=False) self.weight_scale = nn.Parameter(w_s.view(n, k // 32), requires_grad=False) + # The kernel takes a BF16 input's output dtype from the layer. + self.params_dtype = torch.bfloat16 kernel = FlashInferCutedslMxfp8LinearKernel(Mxfp8LinearLayerConfig((n, k))) kernel.process_weights_after_loading(self) self.quant_method = SimpleNamespace(kernel=kernel) @@ -544,6 +627,69 @@ def test_low_latency_mxfp8_mm(shape, rows): assert mxfp8_linear.low_latency_mm(linear, x) is None +@pytest.mark.parametrize("rows", [9, 12, 18, 30, 48, 72, 192, 288, 384]) +@pytest.mark.parametrize("shape", list(mxfp8_linear.CUTE_DSL_TACTICS), ids=str) +def test_pinned_mxfp8_mm(shape, rows): + """``fused_wqa_wkv``/``wq_b``/``wo_b`` at their TP4 shapes on ``cute-dsl`` + with the pinned tactic of the row count, eagerly and replayed, within BF16 + tolerance of the stock layer, for MXFP8 and BF16 inputs.""" + if not current_platform.is_device_capability(103): + pytest.skip("the tactics are pinned on GB300") + n, k = shape + linear = _Mxfp8Linear(n, k) + qa = _quantized(rows, k) + out = mxfp8_linear.pinned_mm(linear, qa) + assert out is not None + torch.testing.assert_close(out, linear(qa), atol=2e-2, rtol=1.6e-2) + + graph, out = _capture(lambda: mxfp8_linear.pinned_mm(linear, qa)) + qa.data.copy_(_quantized(rows, k).data) + graph.replay() + torch.testing.assert_close(out, linear(qa), atol=2e-2, rtol=1.6e-2) + + x = torch.randn(rows, k, device=DEVICE).to(torch.bfloat16) + out = mxfp8_linear.pinned_mm(linear, x) + assert out is not None + torch.testing.assert_close(out, linear(x), atol=2e-2, rtol=1.6e-2) + + +@pytest.mark.parametrize("rows", [9, 12, 16, 17, 24, 32, 33, 48]) +@pytest.mark.parametrize("shape", list(mxfp8_linear.WIDE_TACTICS), ids=str) +def test_wide_low_latency_mxfp8_mm(shape, rows): + """``fused_wqa_wkv``/``wq_b``/``wo_b`` at their TP4 shapes on the + low-latency kernel with a 16-, 32- or 48-row activation tile, eagerly and + replayed, within BF16 tolerance of the stock layer, for MXFP8 and BF16 + inputs; above ``WIDE_LOW_LATENCY_MAX_ROWS`` rows it does not run.""" + if not mxfp8_linear._wide_supported(): + pytest.skip("the wide tactics are pinned on GB300, FlashInfer >= 0.7.2") + n, k = shape + linear = _Mxfp8Linear(n, k) + qa = _quantized(rows, k) + out = mxfp8_linear.wide_mm(linear, qa) + assert out is not None + torch.testing.assert_close(out, linear(qa), atol=2e-2, rtol=1.6e-2) + + graph, out = _capture(lambda: mxfp8_linear.wide_mm(linear, qa)) + qa.data.copy_(_quantized(rows, k).data) + graph.replay() + assert torch.equal(out, mxfp8_linear.wide_mm(linear, qa)) + torch.testing.assert_close(out, linear(qa), atol=2e-2, rtol=1.6e-2) + + x = torch.randn(rows, k, device=DEVICE).to(torch.bfloat16) + out = mxfp8_linear.wide_mm(linear, x) + assert out is not None + torch.testing.assert_close(out, linear(x), atol=2e-2, rtol=1.6e-2) + assert mxfp8_linear.wide_mm(linear, _quantized(49, k)) is None + + +def test_pinned_mxfp8_mm_rows(): + """Up to ``LOW_BATCH_MAX_ROWS`` rows, and for an unpinned shape, the + stock layer runs.""" + n, k = next(iter(mxfp8_linear.CUTE_DSL_TACTICS)) + assert mxfp8_linear.pinned_mm(_Mxfp8Linear(n, k), _quantized(8, k)) is None + assert mxfp8_linear.pinned_mm(_Mxfp8Linear(1024, k), _quantized(48, k)) 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, @@ -1033,14 +1179,19 @@ class _ChainBuilder: supports_update_block_table = False - def __init__(self, offset: int, max_tokens: int): + def __init__(self, offset: int, max_tokens: int, delay: int = 0): self.offset = offset + # GPU cycles to wait before the clamp, so that a build that does not + # wait for it reads the table first. + self.delay = delay self.token_to_req = torch.zeros(max_tokens, dtype=torch.int32, device=DEVICE) self.out = torch.zeros(max_tokens, dtype=torch.int64, device=DEVICE) def build(self, common_prefix_len, common_attn_metadata): cm = common_attn_metadata token_to_req = cm.token_to_req_indices(self.token_to_req) + if self.delay: + torch.cuda._sleep(self.delay) block_table = cm.block_table_tensor.clamp_(min=0) req = token_to_req.long() value = block_table[:, 0].index_select(0, req) * 1000 @@ -1052,11 +1203,33 @@ def build(self, common_prefix_len, common_attn_metadata): return {"out": out, "token_to_req": token_to_req, "block_table": block_table} -def test_metadata_graph_replay_matches_eager_build(): +class _ReaderBuilder(_ChainBuilder): + """A builder that, like the indexer's, only reads its block table: it + must see the table its group's writers clamped.""" + + def build(self, common_prefix_len, common_attn_metadata): + cm = common_attn_metadata + token_to_req = cm.token_to_req_indices(self.token_to_req) + req = token_to_req.long() + value = cm.block_table_tensor[:, 0].index_select(0, req).long() + for _ in range(4): + value = value * 3 + cm.positions + self.offset + out = self.out[: cm.num_actual_tokens] + torch.add(value, cm.slot_mapping, out=out) + return {"out": out, "token_to_req": token_to_req} + + +@pytest.mark.parametrize("readers", [False, True]) +def test_metadata_graph_replay_matches_eager_build(monkeypatch, readers): """A recorded metadata build of five KV-cache groups (two sharing a block table, so the build forks into four branches) refreshes, on replay, every buffer the FULL forward graphs read to the eager build's values, also - after an eager build of another batch left other values there.""" + after an eager build of another batch left other values there. With + readers, the two groups sharing a table each add a builder that only reads + it, which builds on a branch of its own after the table's clamps.""" + from vllm.models.deepseek_v41.nvidia.specialized_decode.runtime import ( + metadata_graphs, + ) from vllm.v1.worker.gpu.attn_utils import build_attn_metadata from vllm.v1.worker.utils import AttentionGroup @@ -1072,7 +1245,9 @@ def test_metadata_graph_replay_matches_eager_build(): torch.zeros(num_reqs, 8, dtype=torch.int32, device=DEVICE) for _ in range(4) ] block_tables = [tables[0], tables[1], tables[2], tables[3], tables[1]] - builders = [_ChainBuilder(i, 64) for i in range(groups)] + builders = [ + _ChainBuilder(i, 64, delay=1_000_000 if readers else 0) for i in range(groups) + ] attn_groups = [ [ AttentionGroup( @@ -1085,6 +1260,20 @@ def test_metadata_graph_replay_matches_eager_build(): ] for i, builder in enumerate(builders) ] + if readers: + monkeypatch.setattr( + metadata_graphs, "_BLOCK_TABLE_READERS", frozenset({_ReaderBuilder}) + ) + for i in (1, 4): + attn_groups[i].append( + AttentionGroup( + backend=None, + layer_names=[f"g{i}.r{j}" for j in range(2)], + kv_cache_spec=None, + kv_cache_group_id=i, + metadata_builders=[_ReaderBuilder(10 + i, 64)], + ) + ) # The runner's persistent query-start buffer, which the recorded build reads. persistent_query_start = query_start.to(DEVICE) @@ -1123,6 +1312,7 @@ def tensors(metadata): slot_mappings.copy_(torch.randint(-1, 100, slot_mappings.shape)) for table in tables: table.copy_(torch.randint(-1, 50, table.shape)) + table[0, 0] = -1 # The clamp changes what the readers see. inputs = [t.clone() for t in tables] # An eager step of another batch leaves its values in the buffers. build(attn_groups, other_query_start, other_query_start.to(DEVICE)) @@ -1135,3 +1325,35 @@ def tensors(metadata): for k, want in expected.items(): assert torch.equal(seen[k], want), (step, k) assert torch.equal(replayed[k], want), (step, k) + + +@pytest.mark.parametrize("rows", [1, 12, 48, 192]) +@torch.inference_mode() +def test_stream_embedding_matches_the_lookup(rows): + """The replicated table's lookup: each row of the embedding and of its + four hc streams is bitwise the table's row of its token, and zeros for a + token outside the table (as the sharded lookup's sum gives it); the + outputs are written in full and nothing past them.""" + from vllm.models.deepseek_v41.nvidia.specialized_decode.embedding.replicated import ( # noqa: E501 + embed_streams, + ) + + torch.manual_seed(0) + vocab, hidden, hc_mult = 1000, 5120, 4 + table = torch.randn(vocab, hidden, device="cuda").bfloat16() + ids = torch.randint(0, vocab, (rows,), dtype=torch.int32, device="cuda") + outside = torch.tensor([-1, vocab, vocab + 4096, vocab - 1], dtype=torch.int32) + ids[: min(rows, 4)] = outside[: min(rows, 4)].cuda() + want = torch.zeros(rows, hidden, device="cuda").bfloat16() + inside = (ids >= 0) & (ids < vocab) + want[inside] = table[ids[inside].long()] + + streams = torch.full((rows + 1, hc_mult, hidden), float("nan"), device="cuda") + streams = streams.bfloat16() + embedding = torch.full((rows + 1, hidden), float("nan"), device="cuda").bfloat16() + embed_streams(ids, table, embedding[:rows], streams[:rows]) + for c in range(hc_mult): + assert torch.equal(streams[:rows, c], want) + assert streams[rows:].isnan().all() + assert torch.equal(embedding[:rows], want) + assert embedding[rows:].isnan().all() diff --git a/tests/models/deepseek_v41/nvidia/specialized_decode/test_layers.py b/tests/models/deepseek_v41/nvidia/specialized_decode/test_layers.py index 8c86d6936d1f..5c70c2c196e0 100644 --- a/tests/models/deepseek_v41/nvidia/specialized_decode/test_layers.py +++ b/tests/models/deepseek_v41/nvidia/specialized_decode/test_layers.py @@ -13,8 +13,8 @@ candidate source, 4 a ratio-1 index source over 3's candidates. Twelve requests decode together once the last prefill is done, finish one by one and leave the last on its own: FULL decode graphs of 12 down to 1 requests -(6 rows each), above 48 rows (inline mHC), padded, row-sharded (12 -requests) and low-batch (one request). +(6 rows each), above 48 rows (the pinned projections), padded, +row-sharded (from two requests) and low-batch (one request). Every decoder layer's residual streams and FFN input and the model's output are kept for every decode token and every 16th prefill position, and the two @@ -153,9 +153,9 @@ def zeros(*shape, dtype=torch.bfloat16): # proves that the CUDA graphs hold the executor, not the stock layers. self.executor_sublayers = 0 - def sublayer_and_keep(boundaries, layer, ffn, state, step, **kwargs): + def sublayer_and_keep(boundaries, layer, ffn, state, **kwargs): self.executor_sublayers += 1 - out = sublayer(boundaries, layer, ffn, state, step, **kwargs) + out = sublayer(boundaries, layer, ffn, state, **kwargs) i = index.get(layer) if ffn and i is not None: residual = out[0].residual @@ -296,6 +296,7 @@ def layers(model): ), "stock_fused_mhc_max_rows": getattr(mhc._all_reduce_mhc, "capacity", None), "draft_attention": [type(m.attn).__name__ for m in layers(draft)], + "embedding": type(getattr(decode, "embedding", None)).__name__, } def record_steps(self, prompt_lens: list[int]) -> None: @@ -501,6 +502,8 @@ def test_specialized_decode_matches_the_model(vllm_runner, tmp_path): assert rank["logits"] == ["SpecializedDecodeLogitsProcessor"], rank assert rank["metadata_graphs"] > 0, rank assert rank["fused_mhc_max_rows"] == FUSED_MHC_MAX_ROWS, rank + # The executor embeds from the replicated table. + assert rank["embedding"] == "ReplicatedEmbedding", rank for rank in paths_off: assert rank["decode"] == "NoneType", rank assert rank["executor_sublayers"] == 0, rank 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 4b6cf95795fc..0f304e647546 100644 --- a/tests/models/deepseek_v41/nvidia/specialized_decode/test_numerics.py +++ b/tests/models/deepseek_v41/nvidia/specialized_decode/test_numerics.py @@ -8,13 +8,13 @@ plain, with the MXFP8 layer input, with a deferred MoE finalize and residual-only, bitwise against the separate all-reduce and mHC kernels; * the executor's sublayer boundary on it, with the next mHC coefficients on - the side stream and inline, against the separate kernels and the stock mHC - pre; + the side stream, against the separate kernels and the stock mHC pre; * the ``wo_b`` GEMM that publishes into the fused kernel's mailbox, per split-K tactic, bitwise against FlashInfer's GEMM plus the separate publish; -* logits gathered with the custom all-gather, bitwise against NCCL; -* the row-sharded decode indexer against the replicated one (equal top-k - sets up to exactly tied scores); +* logits gathered with the executor's all-gather (``LamportAllGather``), + bitwise against NCCL, and NCCL above its capacity; +* the row-sharded decode indexer, its slices gathered with it, against the + replicated one (equal top-k 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 @@ -187,7 +187,7 @@ def _check_all_reduce_mhc(arm, rank): def _check_boundaries(arm, rank): """``Boundaries.sublayer`` on a partial sum and on a deferred MoE output, - with the next coefficients on the side stream and inline: the streams + with the next coefficients on the side stream: the streams bitwise against the separate all-reduce and mHC post, the mixes bitwise against the stock mHC pre on those streams, and the layer input within the fused RMSNorm's ulp of it.""" @@ -262,28 +262,26 @@ def _check_boundaries(arm, rank): stream=side, ) torch.cuda.current_stream().wait_stream(side) - for side_stream in (True, False): - what = f"{name}, {rows} rows, side stream {side_stream}" - state = BoundaryState( - mix["residual"], mix["post"], mix["comb"], mix["pre"], pending - ) - decode_step = SimpleNamespace(side_stream=side_stream) - _late(rank, step := step + 1) - state, layer_input, _, aux = boundaries.sublayer( - layer, False, state, decode_step, capture_aux=True - ) - boundaries.join(layer, decode_step) - _assert_equal( - (state.residual, aux), - (streams, streams.mean(dim=1)), - f"streams, {what}", - ) - _assert_equal( - (state.post, state.comb, state.pre), - (want[0], want[1], want[3]), - f"mixes, {what}", - ) - assert _bf16_ulps(layer_input, want[2]) <= 1, what + what = f"{name}, {rows} rows" + state = BoundaryState( + mix["residual"], mix["post"], mix["comb"], mix["pre"], pending + ) + _late(rank, step := step + 1) + state, layer_input, _, aux = boundaries.sublayer( + layer, False, state, capture_aux=True + ) + boundaries.join(layer) + _assert_equal( + (state.residual, aux), + (streams, streams.mean(dim=1)), + f"streams, {what}", + ) + _assert_equal( + (state.post, state.comb, state.pre), + (want[0], want[1], want[3]), + f"mixes, {what}", + ) + assert _bf16_ulps(layer_input, want[2]) <= 1, what # ------------------------------------------------------- publishing wo_b @@ -410,22 +408,22 @@ def _check_publishing_wo_b(arm, publisher, rank): # -------------------------------------------------------------- logits -def _check_logits_gather(rank): - from vllm.distributed import get_tp_group +VOCAB = 129280 + + +def _check_logits_gather(rank, all_gather): from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.models.deepseek_v41.nvidia.specialized_decode.logits.logits import ( SpecializedDecodeLogitsProcessor, ) - vocab = 129280 - processor = SpecializedDecodeLogitsProcessor(vocab) + processor = SpecializedDecodeLogitsProcessor(VOCAB, all_gather) generator = torch.Generator(device="cuda").manual_seed(rank) - for rows in (1, 6, 16, 48): - logits = _randn(generator, rows, vocab // TP) - if rows <= 16: - # Otherwise both sides would take NCCL. - custom = get_tp_group().device_communicator.custom_all_gather(logits) - assert custom is not None, f"no custom all-gather for {rows} rows" + # Up to 64 BF16 rows fit the all-gather; 80 take NCCL on both sides. + for rows in (1, 6, 16, 48, 64, 80): + logits = _randn(generator, rows, VOCAB // TP) + assert all_gather.accepts(logits) == (rows <= 64), rows + _late(rank, rows) got = processor._gather_logits(logits) want = LogitsProcessor._gather_logits(processor, logits) _assert_equal((got,), (want,), f"logits, {rows} rows") @@ -440,8 +438,10 @@ def _check_logits_gather(rank): MAX_MODEL_LEN = 655_360 PAGE_TOKENS = 128 ROPE_POSITIONS = 2048 -# (requests, of which trailing padding requests). -INDEX_BATCHES = ((8, 0), (32, 10)) +# (requests, of which trailing padding requests): one request (replicated), +# requests split across ranks (the last rank idles), a short last slice, whole +# requests. +INDEX_BATCHES = ((1, 0), (2, 0), (5, 0), (8, 0), (32, 10)) # A dense index source (ratio 2) and the candidate source (ratio 1). INDEX_LAYERS = ((2, False), (1, True)) @@ -571,9 +571,10 @@ def _indexer(keys, ratio, candidates): ) -def _run_indexer(indexer, step, stock=False): - """One decode forward, stock or the executor's; returns the top-k and - candidate rows it wrote.""" +def _run_indexer(indexer, step, stock=False, all_gather=None): + """One decode forward, stock or the executor's (gathering its row slices + with ``all_gather``, else NCCL); returns the top-k and candidate rows it + wrote.""" from vllm.forward_context import ForwardContext, override_forward_context from vllm.models.deepseek_v41.nvidia.specialized_decode.attention.indexer import ( decode_rows, @@ -584,7 +585,7 @@ def _run_indexer(indexer, step, stock=False): if stock: indexer.forward_cuda(hidden, step.q, None, step.weights) else: - decode_rows(indexer, hidden, step.q, step.weights) + decode_rows(indexer, hidden, step.q, step.weights, all_gather) candidates = indexer.candidate_blocks return ( indexer.topk_indices_buffer[: step.rows], @@ -651,7 +652,7 @@ def _assert_same_sets(got, want, scores, limits, k, what): assert (row_scores[swapped] == boundary).all(), (what, row) -def _check_row_sharded_indexer(rank): +def _check_row_sharded_indexer(rank, all_gather): generator = torch.Generator(device="cuda").manual_seed(0) # equal on all ranks keys = {ratio: _index_keys(ratio, generator) for ratio, _ in INDEX_LAYERS} for requests, padding in INDEX_BATCHES: @@ -665,7 +666,7 @@ def _check_row_sharded_indexer(rank): for t in _run_indexer(indexer, step, stock=True) ] _late(rank, requests) - got = _run_indexer(indexer, step) + got = _run_indexer(indexer, step, all_gather=all_gather) logits, blocks = _index_scores(indexer, step) ends = step.metadata.decode.seq_lens.view(-1) _assert_same_sets(got[0], want[0], logits, ends, INDEX_TOPK, what) @@ -680,11 +681,12 @@ def _check_row_sharded_indexer(rank): # ------------------------------------------------------- graph replays -def _check_replays(arm, publisher, keys, rank, replays): +def _check_replays(arm, publisher, keys, rank, replays, all_gather): """One CUDA graph of a decode step's collectives: a published ``wo_b`` boundary, a deferred finalize with the MXFP8 layer input, a residual-only - boundary, a full-capacity boundary and the row-sharded indexer, with the - ranks arriving in a fixed skew. Every replay must equal the first; the + boundary, a full-capacity boundary, the row-sharded indexer and two logits + gathers of other sizes on the indexer's all-gather, with the ranks + arriving in a fixed skew. Every replay must equal the first; the indexer's top-k rows as sets up to boundary ties, since top-k writes them in any order and picks any of the indices that tie at the boundary.""" from vllm.distributed.parallel_state import graph_capture @@ -700,6 +702,8 @@ def _check_replays(arm, publisher, keys, rank, replays): partials = {rows: _partial(rows, rank, seed=rows) for rows in (48, CAPACITY)} step = _decode_step(8, 0, 2, keys[2], torch.Generator(device="cuda").manual_seed(1)) indexer = _indexer(keys[2], 2, candidates=False) + generator = torch.Generator(device="cuda").manual_seed(rank) + vocab_shards = [_randn(generator, rows, VOCAB // TP) for rows in (1, 6)] def decode_step(): torch.cuda._sleep(rank * LATE_CYCLES) @@ -714,7 +718,8 @@ def decode_step(): ) out.append(streams) out += arm(partials[CAPACITY], **mixes[CAPACITY]) - out.append(_run_indexer(indexer, step)[0]) + out += [all_gather(vocab_shards[0]), all_gather(vocab_shards[1])] + out.append(_run_indexer(indexer, step, all_gather=all_gather)[0]) return [t for t in out if t is not None] with graph_capture(device=torch.device("cuda")) as context: @@ -744,12 +749,16 @@ def _worker(rank: int, port: int, replays: int, skip_path: str) -> None: from vllm.config import VllmConfig, set_current_vllm_config from vllm.distributed import cleanup_dist_env_and_memory, get_tp_group from vllm.models.deepseek_v41.nvidia.ops import mhc + from vllm.models.deepseek_v41.nvidia.ops.cute_dsl import LamportAllGather from vllm.models.deepseek_v41.nvidia.specialized_decode.boundary.all_reduce import ( BoundaryAllReduce, ) from vllm.models.deepseek_v41.nvidia.specialized_decode.boundary.publishing_gemm import ( # noqa: E501 PublishingMxfp8Linear, ) + from vllm.models.deepseek_v41.nvidia.specialized_decode.runtime.decode import ( + ALL_GATHER_RANK_BYTES, + ) from vllm.v1.worker.workspace import init_workspace_manager torch.accelerator.set_device_index(rank) @@ -779,12 +788,15 @@ def _worker(rank: int, port: int, replays: int, skip_path: str) -> None: arm = BoundaryAllReduce(config.model_config.hf_config, True) assert arm.capacity == CAPACITY and arm.quantize_layer_input publisher = PublishingMxfp8Linear(arm) + all_gather = LamportAllGather( + rank_capacity_bytes=ALL_GATHER_RANK_BYTES, device=torch.device("cuda", rank) + ) _check_all_reduce_mhc(arm, rank) _check_boundaries(arm, rank) _check_publishing_wo_b(arm, publisher, rank) - _check_logits_gather(rank) - keys = _check_row_sharded_indexer(rank) - _check_replays(arm, publisher, keys, rank, replays) + _check_logits_gather(rank, all_gather) + keys = _check_row_sharded_indexer(rank, all_gather) + _check_replays(arm, publisher, keys, rank, replays, all_gather) torch.accelerator.synchronize() cleanup_dist_env_and_memory() diff --git a/tests/models/deepseek_v41/nvidia/specialized_decode/test_sparse_mla_cute.py b/tests/models/deepseek_v41/nvidia/specialized_decode/test_sparse_mla_cute.py new file mode 100644 index 000000000000..04489b9abd72 --- /dev/null +++ b/tests/models/deepseek_v41/nvidia/specialized_decode/test_sparse_mla_cute.py @@ -0,0 +1,186 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""``sparse_mla_decode`` (``sparse_mla_cute.py``) against an FP32 reference. + +Each token attends to its own columns: the first 128 index the sliding-window +pool and the rest the compressed pool, as storage offsets. The cases cover the +widths of 1 to 5 chunks (128, a partial second chunk, 384, 512, 640), masked +columns (-1 and columns past a token's length, including a token with no active +column, whose output is zero), batches from one token to the most the kernel +takes, sinks and none, and scales other than one. The tolerance allows for the +probabilities' FP8 rounding (TRTLLM-gen's own error is of the same size). +""" + +import pytest +import torch + +from vllm.models.deepseek_v41.nvidia.specialized_decode.attention.sparse_mla_cute import ( # noqa: E501 + _max_tokens, + sparse_mla_decode, + supports_sparse_mla_cute, +) +from vllm.platforms import current_platform + +pytestmark = pytest.mark.skipif( + not ( + current_platform.is_cuda() and current_platform.is_device_capability_family(100) + ), + reason="needs an SM100-family GPU", +) + +DEVICE = "cuda" +H, D, SWA = 16, 512, 128 +SWA_ROWS, COMP_ROWS, PAGE = 1024, 8192, 64 + + +def fp8(shape: tuple[int, ...], std: float, g: torch.Generator) -> torch.Tensor: + """Normal values of standard deviation ``std`` as FP8 E4M3.""" + return (torch.randn(shape, device=DEVICE, generator=g) * std).to( + torch.float8_e4m3fn + ) + + +def make_case(tokens: int, width: int, seed: int) -> dict: + """Pools, indices with masked columns, and lengths for ``tokens`` tokens.""" + g = torch.Generator(device=DEVICE).manual_seed(seed) + swa_pool = fp8((SWA_ROWS // PAGE, PAGE, D), 1.0, g) + comp_pool = fp8((COMP_ROWS // PAGE, PAGE, D), 1.0, g) + swa = torch.randint( + 0, SWA_ROWS, (tokens, min(width, SWA)), device=DEVICE, generator=g + ) + columns = [swa] + if width > SWA: + columns.append( + torch.randint( + 0, COMP_ROWS, (tokens, width - SWA), device=DEVICE, generator=g + ) + ) + indices = torch.cat(columns, dim=1).to(torch.int32) + drop = torch.rand(indices.shape, device=DEVICE, generator=g) < 0.05 + indices[drop] = -1 + lens = torch.full((tokens,), width, dtype=torch.int32, device=DEVICE) + if tokens > 1: + lens[1] = max(width - 77, 1) # a partial last chunk + if tokens > 2: + lens[2] = 0 # no active column + return dict( + q=fp8((tokens, H, D), 0.5, g), + swa_pool=swa_pool, + comp_pool=comp_pool, + indices=indices, + lens=lens, + sinks=torch.randn(H, device=DEVICE, generator=g), + ) + + +def reference( + c: dict, bmm1_scale: float, bmm2_scale: float, sinks: torch.Tensor | None +) -> torch.Tensor: + """FP32 sparse attention over each token's active columns (K = V).""" + swa = c["swa_pool"].reshape(-1, D).float() + comp = c["comp_pool"].reshape(-1, D).float() + q = c["q"].float() + out = torch.zeros(q.shape, device=DEVICE) + for t in range(q.shape[0]): + n = int(c["lens"][t]) + idx = c["indices"][t, :n].long() + cols = torch.arange(n, device=DEVICE) + keep = idx >= 0 + if not bool(keep.any()): + continue + rows = torch.where( + (cols < SWA)[:, None], + swa[idx.clamp(0, swa.shape[0] - 1)], + comp[idx.clamp(0, comp.shape[0] - 1)], + )[keep] + s = (q[t] @ rows.T) * bmm1_scale + m = s.max(dim=1).values + if sinks is not None: + m = torch.maximum(m, sinks) + p = torch.exp(s - m[:, None]) + denom = p.sum(dim=1) + if sinks is not None: + denom = denom + torch.exp(sinks - m) + out[t] = (p @ rows) / denom[:, None] * bmm2_scale + return out + + +@pytest.mark.parametrize("width", [128, 200, 384, 512, 640]) +@pytest.mark.parametrize("tokens", [1, 6, 48]) +@pytest.mark.parametrize("with_sinks", [True, False]) +def test_matches_reference(width: int, tokens: int, with_sinks: bool) -> None: + c = make_case(tokens, width, seed=width * 1000 + tokens) + # The kernel takes only tokens whose clusters are all resident at once, + # which depends on the device's SM count (read once the case's tensors + # have made the device's context current). + tokens = min(tokens, _max_tokens(c["q"].device.index, -(-width // 128))) + c.update({key: c[key][:tokens] for key in ("q", "indices", "lens")}) + sinks = c["sinks"] if with_sinks else None + bmm1_scale, bmm2_scale = D**-0.5 * 1.25, 0.75 + out = torch.full((tokens, H, D), float("nan"), dtype=torch.bfloat16, device=DEVICE) + args = (c["q"], c["swa_pool"], c["comp_pool"], c["indices"], c["lens"], out) + assert supports_sparse_mla_cute(*args, bmm1_scale, bmm2_scale, sinks) + sparse_mla_decode(*args, bmm1_scale, bmm2_scale, sinks) + torch.accelerator.synchronize() + ref = reference(c, bmm1_scale, bmm2_scale, sinks) + got = out.float() + assert not torch.isnan(got).any() + torch.testing.assert_close(got, ref, atol=0.03, rtol=0.05) + if tokens > 2: + assert bool((got[2] == 0).all()) + + +def test_cuda_graph_replay_matches_eager() -> None: + c = make_case(12, 640, seed=7) + args_in = (c["q"], c["swa_pool"], c["comp_pool"], c["indices"], c["lens"]) + eager = torch.empty((12, H, D), dtype=torch.bfloat16, device=DEVICE) + sparse_mla_decode(*args_in, eager, D**-0.5, 1.0, c["sinks"]) + graphed = torch.empty_like(eager) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + sparse_mla_decode(*args_in, graphed, D**-0.5, 1.0, c["sinks"]) + graph.replay() + torch.accelerator.synchronize() + assert torch.equal(graphed, eager) + + +def test_declines_unsupported() -> None: + c = make_case(4, 640, seed=3) + out = torch.empty((4, H, D), dtype=torch.bfloat16, device=DEVICE) + base = [c["q"], c["swa_pool"], c["comp_pool"], c["indices"], c["lens"], out] + assert supports_sparse_mla_cute(*base, 0.1, 1.0, c["sinks"]) + # Tensor scales, BF16 sinks, a 768-wide set, a BF16 query. + assert not supports_sparse_mla_cute(*base, torch.tensor(0.1), 1.0, c["sinks"]) + assert not supports_sparse_mla_cute(*base, 0.1, 1.0, c["sinks"].bfloat16()) + wide = torch.zeros((4, 768), dtype=torch.int32, device=DEVICE) + assert not supports_sparse_mla_cute(*base[:3], wide, *base[4:], 0.1, 1.0, None) + bf16_q = c["q"].bfloat16() + assert not supports_sparse_mla_cute(bf16_q, *base[1:], 0.1, 1.0, None) + # More tokens than one wave of clusters. + many = _max_tokens(torch.accelerator.current_device_index(), 5) + 1 + big = make_case(many, 640, seed=5) + big_out = torch.empty((many, H, D), dtype=torch.bfloat16, device=DEVICE) + big_args = ( + big["q"], + big["swa_pool"], + big["comp_pool"], + big["indices"], + big["lens"], + ) + assert not supports_sparse_mla_cute(*big_args, big_out, 0.1, 1.0, None) + + +def test_index_rows_wider_than_the_table() -> None: + # A [tokens, 640] view of a wider buffer, as when the builder's table is + # sliced: the columns past 640 hold indices the kernel must not read. + c = make_case(6, 640, seed=11) + wide = torch.full((6, 768), 5, dtype=torch.int32, device=DEVICE) + wide[:, :640] = c["indices"] + view = wide[:, :640] + args_in = (c["q"], c["swa_pool"], c["comp_pool"]) + dense = torch.empty((6, H, D), dtype=torch.bfloat16, device=DEVICE) + strided = torch.empty_like(dense) + sparse_mla_decode(*args_in, c["indices"], c["lens"], dense, 0.05, 1.0, c["sinks"]) + sparse_mla_decode(*args_in, view, c["lens"], strided, 0.05, 1.0, c["sinks"]) + torch.accelerator.synchronize() + assert torch.equal(strided, dense) diff --git a/vllm/models/common/ops/moe_tail_fusion/primitives.py b/vllm/models/common/ops/moe_tail_fusion/primitives.py index 53a2384bf265..f3ee78e3d3e1 100644 --- a/vllm/models/common/ops/moe_tail_fusion/primitives.py +++ b/vllm/models/common/ops/moe_tail_fusion/primitives.py @@ -112,6 +112,22 @@ def load_global_u32x2(pointer: cute.Pointer, *, loc=None, ip=None): return cute.TensorSSA(packed, 2, Uint32) +@dsl_user_op +def load_global_u64(address: Int64, *, loc=None, ip=None) -> Int64: + """Load one 64-bit word, such as a device address stored in global memory.""" + return Int64( + _asm( + T.i64(), + [address.ir_value(loc=loc, ip=ip)], + "ld.global.u64 $0, [$1];", + "=l,l", + side_effects=False, + loc=loc, + ip=ip, + ) + ) + + @dsl_user_op def load_global_bf16_as_f32(address: Int64, *, loc=None, ip=None) -> Float32: return Float32( diff --git a/vllm/models/deepseek_v41/common/ops/query_quant.py b/vllm/models/deepseek_v41/common/ops/query_quant.py index 729f55114e93..0d2246052adb 100644 --- a/vllm/models/deepseek_v41/common/ops/query_quant.py +++ b/vllm/models/deepseek_v41/common/ops/query_quant.py @@ -50,12 +50,20 @@ def _q_kv_norm_quant_kernel( KV_SIZE: tl.constexpr, BLOCK: tl.constexpr, LAUNCH_PDL: tl.constexpr, + TRIGGER_FIRST: tl.constexpr, ): row = tl.program_id(0).to(tl.int64) cols = tl.arange(0, BLOCK) if LAUNCH_PDL: + # Dependents read the outputs only after their own grid-dependency + # wait, so triggering first only lets them launch sooner: a decode + # step's query projection then loads its weights while this kernel + # waits for its input. + if TRIGGER_FIRST: + tl.extra.cuda.gdc_launch_dependents() tl.extra.cuda.gdc_wait() - tl.extra.cuda.gdc_launch_dependents() + if not TRIGGER_FIRST: + tl.extra.cuda.gdc_launch_dependents() if tl.program_id(1) == 0: groups = tl.arange(0, BLOCK // 32) if row < num_tokens: @@ -93,8 +101,12 @@ def fused_q_kv_rmsnorm_quant( q_weight: torch.Tensor, kv_weight: torch.Tensor, eps: float, + early_trigger: bool = False, ) -> tuple[QuantizedActivation, torch.Tensor]: - """Normalize Q/KV and quantize Q with FlashInfer's swizzled MXFP8 scales.""" + """Normalize Q/KV and quantize Q with FlashInfer's swizzled MXFP8 scales. + With ``early_trigger``, the kernel triggers its PDL dependents before its + own grid-dependency wait, for a caller whose query projection loads its + weights before its wait.""" assert qr.ndim == kv.ndim == 2 and qr.shape[0] == kv.shape[0] assert qr.stride(-1) == kv.stride(-1) == 1 assert q_weight.is_contiguous() and kv_weight.is_contiguous() @@ -126,6 +138,7 @@ def fused_q_kv_rmsnorm_quant( kv_size, block, launch_pdl, + launch_pdl and early_trigger, num_warps=8 if block >= 2048 else 4, launch_pdl=launch_pdl, ) diff --git a/vllm/models/deepseek_v41/nvidia/model.py b/vllm/models/deepseek_v41/nvidia/model.py index a532ac25c392..96b186a87c67 100644 --- a/vllm/models/deepseek_v41/nvidia/model.py +++ b/vllm/models/deepseek_v41/nvidia/model.py @@ -1000,9 +1000,19 @@ def forward( lookback_token_ids: torch.Tensor | None = None, decode_step: "DecodeStep | None" = None, ) -> torch.Tensor | IntermediateTensors: + # A specialized decode step embeds from the executor's copy of the + # table, which writes the first layer's hc streams too. + embedding = ( + decode_step.executor.embedding + if decode_step is not None and not self.use_sequence_parallel + else None + ) + streams: torch.Tensor | None = None if get_pp_group().is_first_rank: if inputs_embeds is not None: hidden_states = inputs_embeds + elif embedding is not None and embedding.accepts(input_ids): + hidden_states, streams = embedding(input_ids) else: hidden_states = self.embed_input_ids(input_ids) else: @@ -1097,7 +1107,12 @@ def forward( if decode_step is not None: return decode_step.run( - hidden_states, positions, input_ids, engram_hashes, engram_mask + hidden_states, + positions, + input_ids, + engram_hashes, + engram_mask, + streams=streams, ) residual, post_mix, res_mix = None, None, None @@ -1695,7 +1710,9 @@ def build_specialized_decode(self) -> "SpecializedDecode": self.specialized_decode = SpecializedDecode.for_target( self.model, self.specialized_decode_config, self.lm_head ) - self.logits_processor = SpecializedDecodeLogitsProcessor(self.config.vocab_size) + self.logits_processor = SpecializedDecodeLogitsProcessor( + self.config.vocab_size, self.specialized_decode.all_gather + ) return self.specialized_decode def get_model_state_cls(self): diff --git a/vllm/models/deepseek_v41/nvidia/model_state.py b/vllm/models/deepseek_v41/nvidia/model_state.py index afa08c2278a6..f524938ae80d 100644 --- a/vllm/models/deepseek_v41/nvidia/model_state.py +++ b/vllm/models/deepseek_v41/nvidia/model_state.py @@ -298,9 +298,16 @@ def prepare_inputs( self, input_batch: InputBatch, req_states: RequestState ) -> dict[str, torch.Tensor | None]: model_inputs = super().prepare_inputs(input_batch, req_states) - window = self.lookback_token_ids - if window is None: + if self.lookback_token_ids is None: return model_inputs + self._fill_lookback(input_batch, req_states) + model_inputs["lookback_token_ids"] = self.lookback_token_ids + return model_inputs + + def _fill_lookback(self, input_batch: InputBatch, req_states: RequestState) -> None: + """Gather the batch's lookback windows into ``lookback_token_ids``.""" + window = self.lookback_token_ids + assert window is not None all_token_ids = req_states.all_token_ids.gpu depth = window.shape[1] _gather_lookback_kernel[(window.shape[0],)]( @@ -314,8 +321,6 @@ def prepare_inputs( DEPTH=depth, BLOCK_DEPTH=triton.next_power_of_2(depth), ) - model_inputs["lookback_token_ids"] = window - return model_inputs def prepare_dummy_inputs(self, num_reqs: int, num_tokens: int) -> dict[str, Any]: model_inputs = super().prepare_dummy_inputs(num_reqs, num_tokens) diff --git a/vllm/models/deepseek_v41/nvidia/ops/cute_dsl/__init__.py b/vllm/models/deepseek_v41/nvidia/ops/cute_dsl/__init__.py index 5e126a42d607..03d8a11e0bc2 100644 --- a/vllm/models/deepseek_v41/nvidia/ops/cute_dsl/__init__.py +++ b/vllm/models/deepseek_v41/nvidia/ops/cute_dsl/__init__.py @@ -1,8 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""CuTe DSL kernels for DSV4.1 mHC.""" +"""CuTe DSL kernels for DSV4.1 mHC and the TP collectives of its decode.""" +from .all_gather import LamportAllGather from .all_reduce_mhc import AllReduceMHC -__all__ = ["AllReduceMHC"] +__all__ = ["AllReduceMHC", "LamportAllGather"] diff --git a/vllm/models/deepseek_v41/nvidia/ops/cute_dsl/all_gather.py b/vllm/models/deepseek_v41/nvidia/ops/cute_dsl/all_gather.py new file mode 100644 index 000000000000..4a4e927f0773 --- /dev/null +++ b/vllm/models/deepseek_v41/nvidia/ops/cute_dsl/all_gather.py @@ -0,0 +1,405 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# +# The inline-asm PTX helpers below follow the wrappers of FlashInfer's +# flashinfer/comm/mnnvl_cutedsl/cute_dsl_primitives.py +# (flashinfer-ai/flashinfer@139af6f6), which carries the notice below. +# Modified by the vLLM contributors. +# +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# FlashInfer's NOTICE: FlashInfer / Copyright 2025-2026 NVIDIA / +# Copyright 2023-2026 FlashInfer community (https://flashinfer.ai/) +"""A TP all-gather over NVLink multicast with a Lamport mailbox, for the +specialized decode. + +The specialized decode gathers small tensors across the TP group on its +critical path: each dense indexer layer of a row-sharded decode step gathers +its top-k rows (and the candidate source its candidate blocks), and every step +gathers the vocab-parallel logits. A few hundred kilobytes at most, so the +latency of the collective, not its bandwidth, is what counts. + +vLLM's custom all-gather has a Lamport kernel for NVLink multicast +(``mnnvl_lamport_all_gather``), but it never waits for the rank's own +multicast copy: it takes its own slot from registers, and a copy that lands +after the next call has cleaned its stage stays there. A later gather of +another size lays the slots out differently and can read that word as a +peer's payload. A stress test of back-to-back gathers of alternating sizes in +CUDA graphs hits this hundreds of times in a few million gathers. This module +is a gather of the same shape that cannot leave a word behind. + +The mailbox is symmetric memory with a multicast mapping, laid out as +16-byte packs:: + + mailbox[generation][rank][pack] 3 generations x TP ranks x capacity + +A call runs one kernel, one thread a 16-byte pack of the rank's input: + +1. Publish: the thread stores its pack through the multicast address into + every rank's ``mailbox[g][rank][pack]``. A word equal to the Lamport + sentinel (``0x80000000``: FP32 negative zero, INT32_MIN) goes out as 0. +2. Collect: it polls ``mailbox[g][r][pack]`` of every rank ``r``, its own + included, until no word of it is the sentinel; writes each rank's pack to + the output, rank-major as NCCL lays out a dim-0 gather; and stores the + sentinel back into the slots it read. Every slot a call writes is read + after it has landed and reset by the same call, so no call leaves a word + for a later one, whatever the sizes. +3. Generation: ``state[0]`` is the call's generation. Each CTA reads it once + and counts itself in ``state[1]``; the last CTA to arrive advances the + generation and clears the count. A rank can be at most one call ahead of + a peer (finishing a call needs every rank's payload), so a rank writing + the next generation never touches the slots a slower peer still reads; + three generations, as vLLM's and FlashInfer's Lamport kernels use, leave + one more in between. + +The kernel triggers its PDL dependents once it has published. The output +holds the sanitized words on every rank, its own slot included, so all ranks +see the same bits. The caller must issue calls in the same order on every +rank and on one stream: calls on one instance must not overlap. +""" + +from typing import Any + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import cutlass.utils as utils +import torch +import torch.distributed as dist +import torch.distributed._symmetric_memory as symm_mem +from cutlass import Int32, Int64, Uint32 +from cutlass._mlir.dialects import llvm +from cutlass.cute.runtime import make_fake_stream +from cutlass.cutlass_dsl import T, dsl_user_op + +from vllm.distributed import get_tp_group +from vllm.models.common.ops.moe_tail_fusion.primitives import current_cu_stream + +THREADS = 128 +PACK_BYTES = 16 +GENERATIONS = 3 +SENTINEL = 0x80000000 +# state words +GENERATION = 0 +ARRIVALS = 1 + + +# The PTX helpers below follow the inline-asm wrappers of FlashInfer's +# cute_dsl_primitives.py (see the notice at the top of this file). +def _asm(result, operands, text, constraints, *, loc, ip): + return llvm.inline_asm( + result, + operands, + text, + constraints, + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _load_u32x4(address: Int64, volatile: bool, *, loc=None, ip=None): + """One 16-byte pack as four words; volatile for the Lamport polls.""" + opcode = "ld.volatile.global.v4.u32" if volatile else "ld.global.v4.u32" + loaded = _asm( + llvm.StructType.get_literal([T.i32()] * 4), + [address.ir_value(loc=loc, ip=ip)], + f"{opcode} {{$0, $1, $2, $3}}, [$4];", + "=r,=r,=r,=r,l", + loc=loc, + ip=ip, + ) + return tuple( + Uint32(llvm.extractvalue(T.i32(), loaded, [index], loc=loc, ip=ip)) + for index in range(4) + ) + + +@dsl_user_op +def _store_u32x4( + address: Int64, + w0: Uint32, + w1: Uint32, + w2: Uint32, + w3: Uint32, + multicast: bool, + *, + loc=None, + ip=None, +) -> None: + """One 16-byte pack, through a multicast address (``multimem.st``, which + moves the bits unchanged) or to a local one.""" + opcode = ( + "multimem.st.relaxed.sys.global.v4.f32" if multicast else "st.global.v4.u32" + ) + _asm( + None, + [address.ir_value(loc=loc, ip=ip)] + + [word.ir_value(loc=loc, ip=ip) for word in (w0, w1, w2, w3)], + f"{opcode} [$0], {{$1, $2, $3, $4}};", + "l,r,r,r,r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _load_volatile_u32(address: Int64, *, loc=None, ip=None) -> Uint32: + return Uint32( + _asm( + T.i32(), + [address.ir_value(loc=loc, ip=ip)], + "ld.volatile.global.u32 $0, [$1];", + "=r,l", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _store_u32(address: Int64, value: Uint32, *, loc=None, ip=None) -> None: + _asm( + None, + [address.ir_value(loc=loc, ip=ip), value.ir_value(loc=loc, ip=ip)], + "st.global.u32 [$0], $1;", + "l,r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _arrive(address: Int64, *, loc=None, ip=None) -> Uint32: + """Add 1 to the arrival count (acquire-release, GPU scope); returns the + count before.""" + return Uint32( + _asm( + T.i32(), + [address.ir_value(loc=loc, ip=ip), Uint32(1).ir_value(loc=loc, ip=ip)], + "atom.add.acq_rel.gpu.global.u32 $0, [$1], $2;", + "=r,l,r", + loc=loc, + ip=ip, + ) + ) + + +class _AllGatherKernel: + """One thread a 16-byte pack of the rank's input, ``THREADS`` a CTA.""" + + def __init__(self, rank: int, world_size: int, rank_packs: int) -> None: + self.rank = rank + self.world_size = world_size + self.rank_packs = rank_packs + + @cute.jit + def __call__( + self, + inp: Int64, + out: Int64, + mailbox: Int64, + multicast: Int64, + state: Int64, + packs: Int32, + stream: cuda.CUstream, + ) -> None: + """Launch on ``stream``; addresses are byte addresses, ``packs`` the + rank's input size in 16-byte packs.""" + self.kernel(inp, out, mailbox, multicast, state, packs).launch( + grid=((packs + THREADS - 1) // THREADS, 1, 1), + block=(THREADS, 1, 1), + stream=stream, + use_pdl=True, + ) + + @cute.kernel + def kernel( + self, + inp: Int64, + out: Int64, + mailbox: Int64, + multicast: Int64, + state: Int64, + packs: Int32, + ) -> None: + block, _, _ = cute.arch.block_idx() + tidx, _, _ = cute.arch.thread_idx() + grid, _, _ = cute.arch.grid_dim() + pack = block * THREADS + tidx + active = pack < packs + smem = utils.SmemAllocator() + shared_generation = smem.allocate_tensor(Uint32, cute.make_layout(1), 4) + + cute.arch.griddepcontrol_wait() + + # One read of the generation a CTA, before the CTA counts itself: the + # last CTA to arrive can then advance it while the others still run. + if tidx == 0: + current = _load_volatile_u32(state + GENERATION * 4) + shared_generation[0] = current + if _arrive(state + ARRIVALS * 4) == Uint32(grid) - Uint32(1): + _store_u32( + state + GENERATION * 4, + (current + Uint32(1)) % Uint32(GENERATIONS), + ) + _store_u32(state + ARRIVALS * 4, Uint32(0)) + cute.arch.barrier() + generation = shared_generation[0] + # Pack offset of rank 0's slot of this pack in this generation. + base = Int64(generation) * Int64(self.world_size * self.rank_packs) + Int64( + pack + ) + + if active: + w0, w1, w2, w3 = _load_u32x4(inp + Int64(pack) * PACK_BYTES, False) + if w0 == Uint32(SENTINEL): + w0 = Uint32(0) + if w1 == Uint32(SENTINEL): + w1 = Uint32(0) + if w2 == Uint32(SENTINEL): + w2 = Uint32(0) + if w3 == Uint32(SENTINEL): + w3 = Uint32(0) + slot = base + Int64(self.rank * self.rank_packs) + _store_u32x4(multicast + slot * PACK_BYTES, w0, w1, w2, w3, True) + cute.arch.griddepcontrol_launch_dependents() + + # Every rank's slot, this rank's own included: its multicast copy must + # land before the slot is reset, or it would outlive the reset. + values = cute.make_rmem_tensor(cute.make_layout((self.world_size, 4)), Uint32) + values.fill(Uint32(0)) + dirty = active + while dirty: + dirty = False + for source in cutlass.range_constexpr(self.world_size): + slot = base + Int64(source * self.rank_packs) + p0, p1, p2, p3 = _load_u32x4(mailbox + slot * PACK_BYTES, True) + dirty = ( + dirty + | (p0 == Uint32(SENTINEL)) + | (p1 == Uint32(SENTINEL)) + | (p2 == Uint32(SENTINEL)) + | (p3 == Uint32(SENTINEL)) + ) + values[source, 0] = p0 + values[source, 1] = p1 + values[source, 2] = p2 + values[source, 3] = p3 + if active: + sentinel = Uint32(SENTINEL) + for source in cutlass.range_constexpr(self.world_size): + position = Int64(source) * Int64(packs) + Int64(pack) + _store_u32x4( + out + position * PACK_BYTES, + values[source, 0], + values[source, 1], + values[source, 2], + values[source, 3], + False, + ) + slot = base + Int64(source * self.rank_packs) + _store_u32x4( + mailbox + slot * PACK_BYTES, + sentinel, + sentinel, + sentinel, + sentinel, + False, + ) + + +class LamportAllGather: + """TP all-gather along dim 0 over NVLink multicast, for inputs of up to + ``rank_capacity_bytes`` a rank (a multiple of 16 bytes). + + Construction is collective over the TP group: every rank compiles the + kernel and allocates its mailbox, and it raises when NVLink multicast is + unavailable. ``accepts`` tells whether a tensor fits, from its size + alone, so every rank decides alike. + """ + + def __init__(self, *, rank_capacity_bytes: int, device: torch.device) -> None: + if rank_capacity_bytes <= 0 or rank_capacity_bytes % PACK_BYTES: + raise ValueError("the capacity must be a positive multiple of 16 bytes") + group = get_tp_group().device_group + self.world_size = world_size = dist.get_world_size(group) + rank = dist.get_rank(group) + self.rank_capacity_bytes = rank_capacity_bytes + rank_packs = rank_capacity_bytes // PACK_BYTES + device = torch.device(device) + with torch.accelerator.device_index(device.index): + # cute.compile returns an untyped JIT executor. + self._kernel: Any = cute.compile( + _AllGatherKernel(rank, world_size, rank_packs), + Int64(0), + Int64(0), + Int64(0), + Int64(0), + Int64(0), + Int32(1), + make_fake_stream(), + ) + self._mailbox = symm_mem.empty( + (GENERATIONS * world_size * rank_packs * PACK_BYTES // 4,), + dtype=torch.int32, + 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 % PACK_BYTES: + raise RuntimeError("NVLink multicast mapping is unavailable") + self._multicast = multicast + # Every slot starts as the sentinel. + self._mailbox.fill_(-(2**31)) + self._state = torch.zeros(2, dtype=torch.int32, device=device) + torch.accelerator.synchronize() + dist.barrier(group=group) + + def accepts(self, inp: torch.Tensor) -> bool: + """Whether ``inp`` fits: 1 to ``rank_capacity_bytes`` bytes, a + multiple of 16.""" + return 0 < inp.nbytes <= self.rank_capacity_bytes and not inp.nbytes % 16 + + def __call__(self, inp: torch.Tensor) -> torch.Tensor: + """Every rank's ``inp`` concatenated along dim 0 in rank order, on the + current stream. Words equal to ``0x80000000`` arrive as 0, on every + rank (FP32 negative zero becomes positive zero).""" + if not self.accepts(inp): + raise ValueError( + f"the all-gather takes 16-byte multiples of up to " + f"{self.rank_capacity_bytes} bytes a rank, not {inp.nbytes}" + ) + if not inp.is_contiguous() or inp.data_ptr() % PACK_BYTES: + # The kernel loads 16-byte packs. + inp = inp.clone(memory_format=torch.contiguous_format) + out = torch.empty( + (inp.shape[0] * self.world_size, *inp.shape[1:]), + dtype=inp.dtype, + device=inp.device, + ) + self._kernel( + Int64(inp.data_ptr()), + Int64(out.data_ptr()), + Int64(self._mailbox.data_ptr()), + Int64(self._multicast), + Int64(self._state.data_ptr()), + Int32(inp.nbytes // PACK_BYTES), + current_cu_stream(), + ) + return out 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 3bbd13cf2f49..166efc56ff76 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 @@ -1,7 +1,29 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -# Adapted from FlashInfer's low-latency MNNVL CuTe DSL all-reduce -# (flashinfer/comm/mnnvl_cutedsl/kernel_ll, flashinfer-ai/flashinfer@139af6f6). +# +# Adapted from FlashInfer's low-latency and balanced MNNVL CuTe DSL all-reduces +# (flashinfer/comm/mnnvl_cutedsl/kernel_ll and kernel_bt, +# flashinfer-ai/flashinfer@139af6f6), which carry the notice below. Modified +# by the vLLM contributors: the consumer applies DeepSeek-V4.1's mHC post-mix +# and collapse before its RMSNorm, and the kernels are fitted to its MoE +# finalize and decode shapes. +# +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# FlashInfer's NOTICE: FlashInfer / Copyright 2025-2026 NVIDIA / +# Copyright 2023-2026 FlashInfer community (https://flashinfer.ai/) """Lamport all-reduce fused with DSV4.1's mHC post, collapse and RMSNorm. @@ -13,6 +35,22 @@ the RMSNorm then normalizes. The finalize variant also reads FP32 routing weights. +Every rank receives every rank's whole contribution in LL, so its NVLink +traffic grows with ``tp`` times the rows. From ``BT_MIN_ROWS`` rows on +(``BT_MIN_FINALIZE_ROWS`` for the MoE finalize), the +two-shot protocol of FlashInfer's BT kernels takes over: the publish kernels +store each token's contribution only into its owner rank's mailbox (token +``t`` belongs to rank ``t % tp``), an owner kernel reduces its tokens and +multicasts the BF16 sums into a second mailbox, and the same mHC consumer reads +that one reduced value instead of ``tp`` contributions:: + + LL: publish ──multicast──► mailbox[gen][rank][token] ─► consumer (sum, mHC) + BT: publish ──unicast to owner──► mailbox[gen][rank][token // tp] + ─► owner reduce ──multicast──► sums[gen][token] ─► consumer (mHC) + +The owner sums the ranks in the same order and precision as the LL consumer +(FP32 from rank 0, rounded to BF16), so both protocols give the same bits. + Optionally the consumer also writes the normalized output as MXFP8 with UE8M0 scales in the swizzled 128x4 layout FlashInfer's MXFP8 GEMMs read. It quantizes the BF16-rounded output with FlashInfer's own quantize helpers, so the result is @@ -22,7 +60,7 @@ from __future__ import annotations import math -from typing import Any +from typing import Any, NamedTuple import cuda.bindings.driver as cuda import cutlass @@ -51,6 +89,7 @@ load_global_bf16_as_f32, load_global_u32x2, load_global_u32x4, + load_global_u64, load_volatile_u32, make_fake_dynamic_compact_tensor, map_shared_to_peer, @@ -63,6 +102,7 @@ stmc_bf16x4, stmc_bf16x8, store_global_u32, + store_global_u32x2, store_global_u32x4, store_lamport_sentinel_u32x4, store_shared_cluster_f32, @@ -71,6 +111,14 @@ ) from vllm.models.deepseek_v41.common.ops.query_quant import mxfp8_scale_bytes +# The fewest rows reduced with the two-shot (BT) protocol, for the plain +# all-reduce and for the MoE finalize (whose publish also folds the top-k +# rows and the shared expert); fewer use LL. BT reads each contribution once, +# at its owner, but waits twice (owner, then consumer), which costs more where +# ranks arrive skewed, as in a decode step. An instance whose capacity is +# below BT_MIN_ROWS builds no BT kernels or buffers. +BT_MIN_ROWS = 64 +BT_MIN_FINALIZE_ROWS = 64 # One UE8M0 scale per 32 values: four adjacent threads' BF16 fragments. MXFP8_BLOCK = 32 FRAGMENTS_PER_MXFP8_BLOCK = MXFP8_BLOCK // VEC_BF16 @@ -80,6 +128,11 @@ class _SharedOnlyPublishDeviceKernel: + """Publishes a rank's contribution. LL (``owner_capacity`` None): the + Int64 argument is the mailbox's multicast address. BT: it is the address + of an int64 table of every rank's unicast mailbox address, and each token + goes only to its owner's mailbox, ``owner_capacity`` tokens a rank.""" + def __init__( self, *, @@ -91,9 +144,13 @@ def __init__( threads: int, release_before_store: bool, enable_pdl: bool, + owner_capacity: int | None = None, ) -> None: if elements_per_thread not in (1, QUAD_BF16, VEC_BF16): raise ValueError("elements_per_thread must be 1, 4, or 8") + if owner_capacity is not None and elements_per_thread != VEC_BF16: + raise ValueError("the owner (BT) publish stores 8-element fragments") + self.owner_capacity = owner_capacity if hidden <= 0 or hidden % elements_per_thread: raise ValueError("hidden must divide evenly across thread fragments") self.hidden = hidden @@ -221,7 +278,7 @@ def kernel( contribution_mailbox_multicast_address + mailbox_element * 2, packed_words, ) - else: + elif cutlass.const_expr(self.owner_capacity is None): if fragment < self.fragments: mailbox_element = ( (Int64(stage) * self.tp + self.rank) * self.capacity_m @@ -231,6 +288,16 @@ def kernel( contribution_mailbox_multicast_address + mailbox_element * 2, packed_words, ) + else: + if fragment < self.fragments: + mailbox_element = ( + (Int64(stage) * self.tp + self.rank) * self.owner_capacity + + Int64(token // self.tp) + ) * self.hidden + Int64(fragment) * VEC_BF16 + owner_mailbox = load_global_u64( + contribution_mailbox_multicast_address + Int64(token % self.tp) * 8 + ) + store_global_u32x4(owner_mailbox + mailbox_element * 2, packed_words) if cutlass.const_expr(not self.release_before_store): if block == 0 and tidx == 0: @@ -244,6 +311,9 @@ def kernel( class _QuadFinalizePublishDeviceKernel: + """The MoE finalize (top-k weighted sum plus the shared expert) fused with + the publish, ``elements_per_thread`` (4 or 8) values a thread.""" + def __init__( self, *, @@ -259,9 +329,27 @@ def __init__( enable_pdl: bool, prefetch_group: int, fp32_weights: bool = False, + owner_capacity: int | None = None, + elements_per_thread: int = QUAD_BF16, ) -> None: - if hidden <= 0 or hidden % QUAD_BF16: - raise ValueError("hidden must be a positive multiple of 4") + if elements_per_thread not in (QUAD_BF16, VEC_BF16): + raise ValueError("elements_per_thread must be 4 or 8") + if hidden <= 0 or hidden % elements_per_thread: + raise ValueError("hidden must be a multiple of elements_per_thread") + # Fragment width: 8- or 16-byte loads and stores. + self.elements = elements_per_thread + self.words = elements_per_thread // 2 + quad = elements_per_thread == QUAD_BF16 + self.load_words = load_global_u32x2 if quad else load_global_u32x4 + self.unpack = packed_u32x2_to_bf16x4 if quad else packed_u32x4_to_bf16x8 + self.pack = bf16x4_to_packed_u32x2 if quad else bf16x8_to_packed_u32x4 + self.sanitize = ( + sanitize_negative_zero_u32x2 if quad else sanitize_negative_zero_u32x4 + ) + self.store_multicast = stmc_bf16x4 if quad else stmc_bf16x8 + self.store_unicast = store_global_u32x2 if quad else store_global_u32x4 + # As in _SharedOnlyPublishDeviceKernel: None for LL, else BT. + self.owner_capacity = owner_capacity self.hidden = hidden self.top_k = top_k self.tp = tp @@ -269,7 +357,7 @@ def __init__( self.capacity_m = capacity_m self.threads = threads self.routed_scaling_factor = routed_scaling_factor - self.fragments = hidden // QUAD_BF16 + self.fragments = hidden // elements_per_thread self.ctas_per_token = math.ceil(self.fragments / threads) self.include_shared_expert = include_shared_expert self.load_shared_expert_before_pdl = load_shared_expert_before_pdl @@ -360,23 +448,21 @@ def kernel( if cutlass.const_expr(self.include_shared_expert): shared_values = cute.make_rmem_tensor( - cute.make_layout((QUAD_BF16,)), BFloat16 + cute.make_layout((self.elements,)), BFloat16 ) shared_values.fill(BFloat16(0.0)) if cutlass.const_expr(self.load_shared_expert_before_pdl): # noqa: SIM102 if fragment < self.fragments: shared_element = ( - Int64(token) * self.hidden + Int64(fragment) * QUAD_BF16 + Int64(token) * self.hidden + Int64(fragment) * self.elements ) shared_pointer = cute.make_ptr( BFloat16, (shared_output.iterator + shared_element).llvm_ptr, cute.AddressSpace.gmem, - assumed_align=8, - ) - shared_values.store( - packed_u32x2_to_bf16x4(load_global_u32x2(shared_pointer)) + assumed_align=2 * self.elements, ) + shared_values.store(self.unpack(self.load_words(shared_pointer))) if cutlass.const_expr(self.enable_pdl): cute.arch.griddepcontrol_wait() @@ -386,21 +472,21 @@ def kernel( ): if fragment < self.fragments: shared_element = ( - Int64(token) * self.hidden + Int64(fragment) * QUAD_BF16 + Int64(token) * self.hidden + Int64(fragment) * self.elements ) shared_pointer = cute.make_ptr( BFloat16, (shared_output.iterator + shared_element).llvm_ptr, cute.AddressSpace.gmem, - assumed_align=8, - ) - shared_values.store( - packed_u32x2_to_bf16x4(load_global_u32x2(shared_pointer)) + assumed_align=2 * self.elements, ) + shared_values.store(self.unpack(self.load_words(shared_pointer))) stage = load_volatile_u32(stage_state.iterator + NEXT_STAGE) if fragment < self.fragments: - accumulator = cute.make_rmem_tensor(cute.make_layout((QUAD_BF16,)), Float32) + accumulator = cute.make_rmem_tensor( + cute.make_layout((self.elements,)), Float32 + ) accumulator.fill(Float32(0.0)) if cutlass.const_expr(self.prefetch_group == 1): for k in cutlass.range_constexpr(self.top_k): @@ -411,24 +497,22 @@ def kernel( ) if row != Int32(-1): source_element = ( - Int64(row) * self.hidden + Int64(fragment) * QUAD_BF16 + Int64(row) * self.hidden + Int64(fragment) * self.elements ) source_pointer = cute.make_ptr( BFloat16, (routed_output.iterator + source_element).llvm_ptr, cute.AddressSpace.gmem, - assumed_align=8, + assumed_align=2 * self.elements, ) accumulator.store( accumulator.load() - + packed_u32x2_to_bf16x4( - load_global_u32x2(source_pointer) - ).to(Float32) + + self.unpack(self.load_words(source_pointer)).to(Float32) * weight ) else: inputs = cute.make_rmem_tensor( - cute.make_layout((self.prefetch_group, 2)), + cute.make_layout((self.prefetch_group, self.words)), Uint32, ) inputs.fill(Uint32(0)) @@ -440,34 +524,34 @@ def kernel( (staged_indices + k).llvm_ptr, Int32, ) - for word in cutlass.range_constexpr(2): + for word in cutlass.range_constexpr(self.words): inputs[item, word] = Uint32(0) if row != Int32(-1): source_element = ( Int64(row) * self.hidden - + Int64(fragment) * QUAD_BF16 + + Int64(fragment) * self.elements ) source_pointer = cute.make_ptr( BFloat16, (routed_output.iterator + source_element).llvm_ptr, cute.AddressSpace.gmem, - assumed_align=8, + assumed_align=2 * self.elements, ) - source = load_global_u32x2(source_pointer) - for word in cutlass.range_constexpr(2): + source = self.load_words(source_pointer) + for word in cutlass.range_constexpr(self.words): inputs[item, word] = source[word] for item in cutlass.range_constexpr(self.prefetch_group): k = group * self.prefetch_group + item if cutlass.const_expr(k < self.top_k): source = cute.make_rmem_tensor( - cute.make_layout((2,)), + cute.make_layout((self.words,)), Uint32, ) - for word in cutlass.range_constexpr(2): + for word in cutlass.range_constexpr(self.words): source[word] = inputs[item, word] accumulator.store( accumulator.load() - + packed_u32x2_to_bf16x4(source.load()).to(Float32) + + self.unpack(source.load()).to(Float32) * cute.arch.load( (staged_weights + k).llvm_ptr, Float32, @@ -476,15 +560,26 @@ def kernel( result = accumulator.load() if cutlass.const_expr(self.include_shared_expert): result = result + shared_values.load().to(Float32) - result_packed = bf16x4_to_packed_u32x2(result.to(BFloat16)) - packed = sanitize_negative_zero_u32x2(result_packed) - 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, - ) + result_packed = self.pack(result.to(BFloat16)) + packed = self.sanitize(result_packed) + if cutlass.const_expr(self.owner_capacity is None): + mailbox_element = ( + (Int64(stage) * self.tp + self.rank) * self.capacity_m + + Int64(token) + ) * self.hidden + Int64(fragment) * self.elements + self.store_multicast( + contribution_mailbox_multicast_address + mailbox_element * 2, + packed, + ) + else: + mailbox_element = ( + (Int64(stage) * self.tp + self.rank) * self.owner_capacity + + Int64(token // self.tp) + ) * self.hidden + Int64(fragment) * self.elements + owner_mailbox = load_global_u64( + contribution_mailbox_multicast_address + Int64(token % self.tp) * 8 + ) + self.store_unicast(owner_mailbox + mailbox_element * 2, packed) if block == 0 and tidx == 0: store_global_u32( @@ -496,6 +591,135 @@ def kernel( cute.arch.griddepcontrol_launch_dependents() +class _OwnerReduceMulticastDeviceKernel: + """FlashInfer's BT ``_OwnerReduceMulticastDeviceKernel`` without its + residual add: each rank waits for every rank's contribution to the tokens + it owns, sums them in FP32 in rank order, and multicasts the BF16 sums into + every rank's reduced mailbox ``[generation][token]``. It then advances the + generation and resets its contribution slots.""" + + def __init__( + self, + *, + hidden: int, + tp: int, + rank: int, + capacity_m: int, + owner_capacity: int, + threads: int, + enable_pdl: bool, + ) -> None: + self.hidden = hidden + self.tp = tp + self.rank = rank + self.capacity_m = capacity_m + self.owner_capacity = owner_capacity + self.threads = threads + self.enable_pdl = enable_pdl + self.fragments = hidden // VEC_BF16 + self.ctas_per_token = math.ceil(self.fragments / threads) + + @cute.jit + def __call__( + self, + contribution_mailbox: cute.Tensor, + stage_state: cute.Tensor, + reduced_mailbox_multicast_address: Int64, + m: Int32, + stream: cuda.CUstream, + ) -> None: + owned_tokens = (m + Int32(self.tp - 1)) // Int32(self.tp) + self.kernel( + contribution_mailbox, + stage_state, + reduced_mailbox_multicast_address, + m, + ).launch( + grid=(owned_tokens * self.ctas_per_token, 1, 1), + block=(self.threads, 1, 1), + stream=stream, + use_pdl=self.enable_pdl, + ) + + @cute.kernel + def kernel( + self, + contribution_mailbox: cute.Tensor, + stage_state: cute.Tensor, + reduced_mailbox_multicast_address: Int64, + m: Int32, + ) -> None: + block, _, _ = cute.arch.block_idx() + tidx, _, _ = cute.arch.thread_idx() + local_token = block // self.ctas_per_token + fragment = (block % self.ctas_per_token) * self.threads + tidx + token = local_token * self.tp + self.rank + active = token < m and fragment < self.fragments + + if cutlass.const_expr(self.enable_pdl): + cute.arch.griddepcontrol_wait() + + stage = load_volatile_u32(stage_state.iterator + ACTIVE_STAGE) + rank_values = cute.make_rmem_tensor(cute.make_layout((self.tp, 4)), Uint32) + rank_values.fill(Uint32(0)) + dirty = active + while dirty: + dirty = False + for source_rank in cutlass.range_constexpr(self.tp): + source_element = ( + (Int64(stage) * self.tp + source_rank) * self.owner_capacity + + Int64(local_token) + ) * self.hidden + Int64(fragment) * VEC_BF16 + source_pointer = cute.make_ptr( + BFloat16, + (contribution_mailbox.iterator + source_element).llvm_ptr, + cute.AddressSpace.gmem, + assumed_align=16, + ) + packed = load_global_u32x4(source_pointer, volatile=True) + dirty = dirty | fragment_has_negative_zero(packed) + for word in cutlass.range_constexpr(4): + rank_values[source_rank, word] = packed[word] + if active: + reduced = cute.make_rmem_tensor(cute.make_layout((VEC_BF16,)), Float32) + reduced.fill(Float32(0.0)) + for source_rank in cutlass.range_constexpr(self.tp): + packed = cute.make_rmem_tensor(cute.make_layout((4,)), Uint32) + for word in cutlass.range_constexpr(4): + packed[word] = rank_values[source_rank, word] + reduced.store( + reduced.load() + packed_u32x4_to_bf16x8(packed.load()).to(Float32) + ) + reduced_element = ( + Int64(stage) * self.capacity_m + Int64(token) + ) * self.hidden + Int64(fragment) * VEC_BF16 + stmc_bf16x8( + reduced_mailbox_multicast_address + reduced_element * 2, + sanitize_negative_zero_u32x4( + bf16x8_to_packed_u32x4(reduced.load().to(BFloat16)) + ), + ) + + if block == 0 and tidx == 0: + store_global_u32( + (stage_state.iterator + NEXT_STAGE).toint(), + (stage + Uint32(1)) % Uint32(LAMPORT_GENERATIONS), + ) + cute.arch.barrier() + if cutlass.const_expr(self.enable_pdl): + cute.arch.griddepcontrol_launch_dependents() + + if active: + for source_rank in cutlass.range_constexpr(self.tp): + source_element = ( + (Int64(stage) * self.tp + source_rank) * self.owner_capacity + + Int64(local_token) + ) * self.hidden + Int64(fragment) * VEC_BF16 + store_lamport_sentinel_u32x4( + Int64((contribution_mailbox.iterator + source_element).toint()) + ) + + class _LamportMHCDeviceKernel: """FlashInfer's ``_LamportResidualRMSNormDeviceKernel`` with mHC. @@ -505,7 +729,9 @@ class _LamportMHCDeviceKernel: is also written as MXFP8 with swizzled 128x4 UE8M0 scales, including the zero scales of the padding rows up to a multiple of 128. With ``residual_only``, only the residual output is written; the pre-mix, gamma - and normalized output are unused. + and normalized output are unused. With ``reduced_input`` (BT), the mailbox + holds each token's reduced sum, ``[generation][token]``, from + ``_OwnerReduceMulticastDeviceKernel``, which also advanced the generation. """ def __init__( @@ -521,9 +747,13 @@ def __init__( enable_pdl: bool, quantize: bool = False, residual_only: bool = False, + reduced_input: bool = False, ) -> None: if rank_lanes not in (1, 2, 4, 8): raise ValueError("rank_lanes must be 1, 2, 4, or 8") + if reduced_input and rank_lanes != 1: + raise ValueError("a reduced mailbox has one source: rank_lanes must be 1") + self.reduced_input = reduced_input if tp % rank_lanes: raise ValueError("tp must be divisible by rank_lanes") if threads <= 0 or threads % WARP_SIZE or threads % rank_lanes: @@ -562,7 +792,18 @@ def __init__( self.fragment_stride = cluster_size * self.groups_per_cta self.trips = math.ceil(self.fragments / self.fragment_stride) self.warps = threads // WARP_SIZE - self.rank_waves = tp // rank_lanes + self.rank_waves = 1 if reduced_input else tp // rank_lanes + + def source_element(self, stage, source_rank, token, fragment) -> Int64: + """Mailbox element of a fragment: of ``source_rank``'s contribution, + or of the reduced sum (``reduced_input``).""" + if self.reduced_input: + slot = Int64(stage) * self.capacity_m + Int64(token) + else: + slot = ( + Int64(stage) * self.tp + Int64(source_rank) + ) * self.capacity_m + Int64(token) + return slot * self.hidden + Int64(fragment) * VEC_BF16 def smem_size_in_bytes(self) -> int: return (self.warps + self.cluster_size) * 4 @@ -720,11 +961,9 @@ def kernel( for wave in cutlass.range_constexpr(self.rank_waves): source_rank = wave * self.rank_lanes + rank_lane if fragment < self.fragments: - source_element = ( - (Int64(active_stage) * self.tp + Int64(source_rank)) - * self.capacity_m - + Int64(token) - ) * self.hidden + Int64(fragment) * VEC_BF16 + source_element = self.source_element( + active_stage, source_rank, token, fragment + ) source_pointer = cute.make_ptr( BFloat16, (contribution_mailbox.iterator + source_element).llvm_ptr, @@ -789,11 +1028,12 @@ def kernel( if cutlass.const_expr(self.write_layer_input): prenorm_fragments[trip, None].store(collapse.load().to(BFloat16)) - if token == 0 and cluster_rank == 0 and tidx == 0: - store_global_u32( - (stage_state.iterator + NEXT_STAGE).toint(), - (active_stage + Uint32(1)) % Uint32(LAMPORT_GENERATIONS), - ) + if cutlass.const_expr(not self.reduced_input): # noqa: SIM102 + if token == 0 and cluster_rank == 0 and tidx == 0: + store_global_u32( + (stage_state.iterator + NEXT_STAGE).toint(), + (active_stage + Uint32(1)) % Uint32(LAMPORT_GENERATIONS), + ) if cutlass.const_expr(self.enable_pdl): cute.arch.griddepcontrol_launch_dependents() @@ -802,11 +1042,9 @@ def kernel( for wave in cutlass.range_constexpr(self.rank_waves): source_rank = wave * self.rank_lanes + rank_lane if fragment < self.fragments: - source_element = ( - (Int64(active_stage) * self.tp + Int64(source_rank)) - * self.capacity_m - + Int64(token) - ) * self.hidden + Int64(fragment) * VEC_BF16 + source_element = self.source_element( + active_stage, source_rank, token, fragment + ) store_lamport_sentinel_u32x4( Int64((contribution_mailbox.iterator + source_element).toint()) ) @@ -923,10 +1161,27 @@ def kernel( word % words_per_row * SCALE_TILE_COLS, Int32(self.scale_cols), ) - store_global_u32(quant_scales.iterator + offset, Uint32(0)) + store_global_u32( + (quant_scales.iterator + offset).toint(), Uint32(0) + ) word = word + m * self.cluster_size * self.threads +class _Consumer(NamedTuple): + """A call's validated consumer: the compiled variant and its outputs.""" + + # The JIT executor ``cute.compile`` returns; CuTe DSL exposes no type for it. + kernel: Any + # None for the residual-only variant, which reads no pre-mix: the launch + # passes the post-mix (same shape) in its place. + pre: torch.Tensor | None + norm_weight: torch.Tensor + residual_output: torch.Tensor + norm_output: torch.Tensor + layer_input: torch.Tensor | None + quant_args: tuple[cute.Tensor, cute.Tensor] + + class AllReduceMHC: """Lamport TP all-reduce fused with mHC post, collapse and RMSNorm. @@ -939,7 +1194,9 @@ class AllReduceMHC: buffers that receive the normalized collapse as MXFP8. With ``residual_only``, both also accept ``pre`` and ``norm_weight`` of None and then write only the post-mixed streams, for a consumer that collapses them - itself. + itself. Calls of ``bt_min_rows`` rows or more (finalize: + ``bt_min_finalize_rows``) use the two-shot BT protocol, with its own + mailboxes and generation state. """ def __init__( @@ -1021,13 +1278,27 @@ def __init__( Int32(capacity), current_cu_stream(), ) + # BT, for instances that can take BT_MIN_ROWS rows; a smaller one + # (thresholds above its capacity) builds none of its kernels or + # buffers. Tokens are owned round-robin, so a rank owns at most + # owner_capacity of a call's rows. + bt_enabled = capacity >= BT_MIN_ROWS + self.bt_min_rows = BT_MIN_ROWS if bt_enabled else capacity + 1 + self.bt_min_finalize_rows = ( + BT_MIN_FINALIZE_ROWS if bt_enabled else capacity + 1 + ) + self._owner_capacity = owner_capacity = math.ceil(capacity / tp) + if bt_enabled: + self._compile_bt(hidden, top_k, tp, rank, capacity, owner_capacity) # Stand-ins for the MXFP8 outputs of the unquantized variants, # which never touch them. self._no_quant_output = torch.empty(16, dtype=torch.uint8, device=device) no_quant_cute = to_cute(self._no_quant_output, 16) self._no_quant_args = (no_quant_cute, no_quant_cute) - def compile_collective(quantize: bool, residual_only: bool): + def compile_collective( + quantize: bool, residual_only: bool, reduced_input: bool + ): if quantize: quant_fakes = ( make_fake_dynamic_compact_tensor( @@ -1053,10 +1324,16 @@ def compile_collective(quantize: bool, residual_only: bool): enable_pdl=True, quantize=quantize, residual_only=residual_only, + reduced_input=reduced_input, ), make_fake_compact_tensor( BFloat16, - (LAMPORT_GENERATIONS * tp * capacity * hidden,), + ( + LAMPORT_GENERATIONS + * (1 if reduced_input else tp) + * capacity + * hidden, + ), assumed_align=16, ), make_fake_dynamic_compact_tensor( @@ -1085,17 +1362,17 @@ def compile_collective(quantize: bool, residual_only: bool): current_cu_stream(), ) - self._collective = compile_collective(quantize=False, residual_only=False) - self._collective_quant = ( - compile_collective(quantize=True, residual_only=False) - if quantize_layer_input - else None - ) - self._collective_residual = ( - compile_collective(quantize=False, residual_only=True) - if residual_only - else None - ) + # Consumers by (BT, variant): "plain", "quant" (also the MXFP8 + # layer input) and "residual" (residual streams only). + self._consumers: dict[tuple[bool, str], Any] = {} + for bt in (False, True) if bt_enabled else (False,): + self._consumers[bt, "plain"] = compile_collective(False, False, bt) + if quantize_layer_input: + self._consumers[bt, "quant"] = compile_collective(True, False, bt) + if residual_only: + self._consumers[bt, "residual"] = compile_collective( + False, True, bt + ) if residual_only: # Stand-in for the norm weight the variant never reads. self._no_norm_weight = torch.empty( @@ -1115,9 +1392,141 @@ def compile_collective(quantize: bool, residual_only: bool): # 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) + if bt_enabled: + self._allocate_bt(group, tp, capacity, owner_capacity, device) torch.accelerator.synchronize() dist.barrier(group=group) + def _compile_bt( + self, + hidden: int, + top_k: int, + tp: int, + rank: int, + capacity: int, + owner_capacity: int, + ) -> None: + """Compile the BT publish, finalize publish and owner-reduce kernels.""" + self._bt_publish = cute.compile( + _SharedOnlyPublishDeviceKernel( + hidden=hidden, + tp=tp, + rank=rank, + capacity_m=capacity, + elements_per_thread=VEC_BF16, + threads=128, + release_before_store=False, + enable_pdl=True, + owner_capacity=owner_capacity, + ), + make_fake_dynamic_compact_tensor( + BFloat16, alignment=16, divisibility=hidden + ), + make_fake_compact_tensor(Int32, (2,), assumed_align=4), + Int64(0), + Int32(capacity), + current_cu_stream(), + ) + self._bt_finalize_publish = cute.compile( + _QuadFinalizePublishDeviceKernel( + hidden=hidden, + top_k=top_k, + tp=tp, + rank=rank, + capacity_m=capacity, + threads=128, + routed_scaling_factor=1.0, + include_shared_expert=True, + load_shared_expert_before_pdl=False, + enable_pdl=True, + prefetch_group=top_k, + fp32_weights=True, + owner_capacity=owner_capacity, + elements_per_thread=VEC_BF16, + ), + make_fake_dynamic_compact_tensor( + BFloat16, alignment=16, divisibility=hidden + ), + make_fake_dynamic_compact_tensor(Float32, alignment=4, divisibility=top_k), + make_fake_dynamic_compact_tensor(Int32, alignment=4, divisibility=top_k), + make_fake_dynamic_compact_tensor( + BFloat16, alignment=16, divisibility=hidden + ), + make_fake_compact_tensor(Int32, (2,), assumed_align=4), + Int64(0), + Int32(capacity), + current_cu_stream(), + ) + # FlashInfer's BT presets for TP4/H5120 reduce with 32 threads a + # CTA up to 192 rows and 320 above. + self._owner_reduce = { + threads: cute.compile( + _OwnerReduceMulticastDeviceKernel( + hidden=hidden, + tp=tp, + rank=rank, + capacity_m=capacity, + owner_capacity=owner_capacity, + threads=threads, + enable_pdl=True, + ), + make_fake_compact_tensor( + BFloat16, + (LAMPORT_GENERATIONS * tp * owner_capacity * hidden,), + assumed_align=16, + ), + make_fake_compact_tensor(Int32, (2,), assumed_align=4), + Int64(0), + Int32(capacity), + current_cu_stream(), + ) + for threads in (32, 320) + } + + def _allocate_bt( + self, + group: dist.ProcessGroup, + tp: int, + capacity: int, + owner_capacity: int, + device: torch.device, + ) -> None: + """Allocate the BT mailboxes (collective over ``group``).""" + # BT: contributions [generation][source rank][owned token] on each + # owner, stored by unicast through a table of the owners' addresses, + # and the owners' sums [generation][token] on every rank. + self._bt_mailbox = symm_mem.empty( + (LAMPORT_GENERATIONS, tp, owner_capacity, self.hidden_size), + dtype=torch.bfloat16, + device=device, + ) + self._bt_mailbox_handle = symm_mem.rendezvous(self._bt_mailbox, group) + # get_remote_tensor keeps the rendezvous offset of a pooled buffer. + owner_mailboxes = [ + self._bt_mailbox_handle.get_remote_tensor( + peer, self._bt_mailbox.shape, self._bt_mailbox.dtype + ).data_ptr() + for peer in range(tp) + ] + if not all(owner_mailboxes): + raise RuntimeError("NVLink peer mapping is unavailable") + self._bt_owner_mailboxes = torch.tensor( + owner_mailboxes, dtype=torch.int64, device=device + ) + self._bt_reduced = symm_mem.empty( + (LAMPORT_GENERATIONS, capacity, self.hidden_size), + dtype=torch.bfloat16, + device=device, + ) + self._bt_reduced_handle = symm_mem.rendezvous(self._bt_reduced, group) + bt_multicast = int(self._bt_reduced_handle.multicast_ptr or 0) + if not bt_multicast or bt_multicast % 16: + raise RuntimeError("NVLink multicast mapping is unavailable") + self._bt_multicast = bt_multicast + self._bt_mailbox.view(torch.int16).fill_(-32768) + self._bt_reduced.view(torch.int16).fill_(-32768) + self._bt_stage_state = torch.zeros(2, dtype=torch.int32, device=device) + def __call__( self, x: torch.Tensor, @@ -1150,16 +1559,15 @@ def __call__( self._check(m, residual, post, comb, pre) if x.shape != (m, self.hidden_size): raise ValueError("unexpected x shape") - self._publish( + bt = m >= self.bt_min_rows + consumer = self._consumer(residual, pre, norm_weight, m, mxfp8_out, bt) + (self._bt_publish if bt else self._publish)( to_cute_dynamic(x.flatten(), 16, divisibility=self.hidden_size), - to_cute(self._stage_state, 4), - Int64(self._multicast), + *self._publish_target(bt), Int32(m), current_cu_stream(), ) - return self._collect( - residual, post, comb, pre, norm_weight, norm_eps, m, mxfp8_out - ) + return self._collect(consumer, residual, post, comb, norm_eps, m, bt) def finalize( self, @@ -1208,21 +1616,29 @@ def finalize( or expanded_idx_to_permuted_idx.numel() != m * top_k ): raise ValueError("expanded_idx_to_permuted_idx must be int32 [M, top_k]") - self._finalize_publish( + bt = m >= self.bt_min_finalize_rows + consumer = self._consumer(residual, pre, norm_weight, m, mxfp8_out, bt) + (self._bt_finalize_publish if bt else self._finalize_publish)( to_cute_dynamic(gemm2_permuted.flatten(), 16, divisibility=hidden), to_cute_dynamic(expert_weights.flatten(), 4, divisibility=top_k), to_cute_dynamic( expanded_idx_to_permuted_idx.flatten(), 4, divisibility=top_k ), to_cute_dynamic(shared_output.flatten(), 16, divisibility=hidden), - to_cute(self._stage_state, 4), - Int64(self._multicast), + *self._publish_target(bt), Int32(m), current_cu_stream(), ) - return self._collect( - residual, post, comb, pre, norm_weight, norm_eps, m, mxfp8_out - ) + return self._collect(consumer, residual, post, comb, norm_eps, m, bt) + + def _publish_target(self, bt: bool) -> tuple[cute.Tensor, Int64]: + """The publish kernels' generation state and mailbox address: LL's + multicast mailbox, or BT's table of the owners' mailboxes.""" + if bt: + return to_cute(self._bt_stage_state, 4), Int64( + self._bt_owner_mailboxes.data_ptr() + ) + return to_cute(self._stage_state, 4), Int64(self._multicast) def _check( self, @@ -1244,65 +1660,107 @@ def _check( ): raise ValueError("unexpected mix shapes") - def _collect( + def _consumer( self, residual: torch.Tensor, - post: torch.Tensor, - comb: torch.Tensor, pre: torch.Tensor | None, norm_weight: torch.Tensor | None, - norm_eps: float, m: int, mxfp8_out: tuple[torch.Tensor, torch.Tensor] | None, - ) -> tuple[torch.Tensor, torch.Tensor | None]: - hidden, hc = self.hidden_size, self.hc_mult + bt: bool, + ) -> _Consumer: + """Select and validate the consumer variant of a call and allocate its + outputs. Callers do this before the call's first launch, so a bad + argument raises with the mailboxes and generation state untouched.""" + hidden = self.hidden_size residual_output = torch.empty_like(residual) if (pre is None) != (norm_weight is None): raise ValueError("pre and norm_weight are both given or both None") - layer_input: torch.Tensor | None = None if pre is None or norm_weight is None: - if self._collective_residual is None: + if (bt, "residual") not in self._consumers: raise ValueError("built without residual_only") if mxfp8_out is not None: raise ValueError("a residual-only boundary has no layer input") # The variant reads neither the pre-mix nor the norm weight and - # writes no layer input: stand-ins fill their arguments. - collective, quant_args = self._collective_residual, self._no_quant_args - pre, norm_weight, norm_output = post, self._no_norm_weight, residual_output + # writes no layer input: stand-ins fill their arguments (the + # post-mix has the pre-mix's shape). + return _Consumer( + kernel=self._consumers[bt, "residual"], + pre=None, + norm_weight=self._no_norm_weight, + residual_output=residual_output, + norm_output=residual_output, + layer_input=None, + quant_args=self._no_quant_args, + ) + layer_input = torch.empty( + (m, hidden), dtype=torch.bfloat16, device=residual.device + ) + if mxfp8_out is None: + kernel, quant_args = self._consumers[bt, "plain"], self._no_quant_args else: - layer_input = norm_output = torch.empty( - (m, hidden), dtype=torch.bfloat16, device=residual.device + if (bt, "quant") not in self._consumers: + raise ValueError("built without quantize_layer_input") + values, scales = mxfp8_out + if values.shape != (m, hidden) or scales.numel() != ( + mxfp8_scale_bytes(m, hidden) + ): + raise ValueError("mxfp8_out does not match the layer input") + kernel = self._consumers[bt, "quant"] + quant_args = ( + to_cute_dynamic( + values.view(torch.uint8).flatten(), 16, divisibility=hidden + ), + to_cute_dynamic(scales.flatten(), 16, divisibility=16), ) - if mxfp8_out is None: - collective, quant_args = self._collective, self._no_quant_args - else: - if self._collective_quant is None: - raise ValueError("built without quantize_layer_input") - values, scales = mxfp8_out - if values.shape != (m, hidden) or scales.numel() != ( - mxfp8_scale_bytes(m, hidden) - ): - raise ValueError("mxfp8_out does not match the layer input") - collective = self._collective_quant - quant_args = ( - to_cute_dynamic( - values.view(torch.uint8).flatten(), 16, divisibility=hidden - ), - to_cute_dynamic(scales.flatten(), 16, divisibility=16), - ) - collective( - to_cute(self._mailbox.flatten(), 16), + return _Consumer( + kernel=kernel, + pre=pre, + norm_weight=norm_weight, + residual_output=residual_output, + norm_output=layer_input, + layer_input=layer_input, + quant_args=quant_args, + ) + + def _collect( + self, + consumer: _Consumer, + residual: torch.Tensor, + post: torch.Tensor, + comb: torch.Tensor, + norm_eps: float, + m: int, + bt: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + """Launch the call's reduction after its publish: for BT the owner + reduce, then the consumer that ``_consumer`` selected.""" + hidden, hc = self.hidden_size, self.hc_mult + if bt: + # The owners reduce first; the consumer then reads their sums. + self._owner_reduce[32 if m <= 192 else 320]( + to_cute(self._bt_mailbox.flatten(), 16), + to_cute(self._bt_stage_state, 4), + Int64(self._bt_multicast), + Int32(m), + current_cu_stream(), + ) + pre = post if consumer.pre is None else consumer.pre + consumer.kernel( + to_cute((self._bt_reduced if bt else self._mailbox).flatten(), 16), to_cute_dynamic(residual.flatten(), 16, divisibility=hc * hidden), to_cute_dynamic(post.flatten(), 4, divisibility=hc), to_cute_dynamic(comb.flatten(), 4, divisibility=hc * hc), to_cute_dynamic(pre.flatten(), 4, divisibility=hc), - to_cute(norm_weight, 16), - to_cute_dynamic(residual_output.flatten(), 16, divisibility=hc * hidden), - to_cute_dynamic(norm_output.flatten(), 16, divisibility=hidden), - *quant_args, - to_cute(self._stage_state, 4), + to_cute(consumer.norm_weight, 16), + to_cute_dynamic( + consumer.residual_output.flatten(), 16, divisibility=hc * hidden + ), + to_cute_dynamic(consumer.norm_output.flatten(), 16, divisibility=hidden), + *consumer.quant_args, + to_cute(self._bt_stage_state if bt else self._stage_state, 4), Float32(norm_eps), Int32(m), current_cu_stream(), ) - return residual_output, layer_input + return consumer.residual_output, consumer.layer_input diff --git a/vllm/models/deepseek_v41/nvidia/ops/fused_wo_a.py b/vllm/models/deepseek_v41/nvidia/ops/fused_wo_a.py index b195d4b84937..b144f5b15cfd 100644 --- a/vllm/models/deepseek_v41/nvidia/ops/fused_wo_a.py +++ b/vllm/models/deepseek_v41/nvidia/ops/fused_wo_a.py @@ -92,12 +92,16 @@ class CompileKey: n_groups: int heads_per_group: int o_lora_rank: int + # Trigger the PDL dependent right after the TMEM allocation rather than + # once the last stage has loaded. + early_trigger: bool = False @staticmethod def kernel(compile_key: CompileKey) -> Any: tokens = compile_key.tokens_per_tile token_tiles = compile_key.token_tiles full_tiles = compile_key.full_tiles + early_trigger = compile_key.early_trigger acc_cols = 1 << (tokens - 1).bit_length() tile_m = max(8, acc_cols) tile_n = 128 @@ -233,6 +237,16 @@ def device_kernel( if warp == 0: cute.arch.alloc_tmem(tmem_cols, taddr) cute.arch.relinquish_tmem_alloc_permit() + # With early_trigger, for a dependent wo_b GEMM that loads its + # weights before its own grid wait: launching it here overlaps + # that weight read with this kernel's quantization, MMA and + # reduction, and it reads this kernel's outputs only after its + # wait. It launches only once every CTA has issued this, after + # its own TMEM allocation, so it cannot take TMEM this kernel + # still needs. + if cutlass.const_expr(early_trigger): + with cute.arch.elect_one(): + cute.arch.griddepcontrol_launch_dependents() copy_x = cute.make_copy_atom( cute.nvgpu.CopyUniversalOp(), BFloat16, num_bits_per_copy=64 ) @@ -343,7 +357,9 @@ def device_kernel( for stage in cutlass.range_constexpr(num_stages): cute.arch.mbarrier_wait(mbar_loaded + stage, 0) _tcgen05.fence_after_thread_sync() - if cutlass.const_expr(stage == num_stages - 1): + if cutlass.const_expr( + not early_trigger and stage == num_stages - 1 + ): with cute.arch.elect_one(): cute.arch.griddepcontrol_launch_dependents() adesc = sdesc | (sW[None, None, stage].iterator.toint() >> 4) @@ -501,7 +517,13 @@ def host_entrypoint( return host_entrypoint def dispatch( # type: ignore[override] - self, *, tokens: int, n_groups: int, heads_per_group: int, o_lora_rank: int + self, + *, + tokens: int, + n_groups: int, + heads_per_group: int, + o_lora_rank: int, + early_trigger: bool = False, ) -> CompileKey: token_tiles = (tokens + 31) // 32 # Host integer arithmetic: balance tiles in multiples of four tokens. @@ -517,10 +539,17 @@ def dispatch( # type: ignore[override] n_groups=n_groups, heads_per_group=heads_per_group, o_lora_rank=o_lora_rank, + early_trigger=early_trigger, ) def get_warmup_keys( - self, *, max_tokens: int, n_groups: int, heads_per_group: int, o_lora_rank: int + self, + *, + max_tokens: int, + n_groups: int, + heads_per_group: int, + o_lora_rank: int, + early_trigger: bool = False, ) -> list[CompileKey]: # Trace all accepted token counts and deduplicate shared tile shapes. return self._trace_dispatch(self.dispatch)( @@ -528,6 +557,7 @@ def get_warmup_keys( n_groups=n_groups, heads_per_group=heads_per_group, o_lora_rank=o_lora_rank, + early_trigger=early_trigger, ) def warmup_inputs(self, compile_key: CompileKey) -> tuple[Any, ...]: @@ -570,8 +600,11 @@ def __call__( rope: torch.Tensor, weight: torch.Tensor, weight_scale: torch.Tensor, + early_trigger: bool = False, ) -> CuTeDSLLaunchSpec["FusedWoAKernel.CompileKey"]: - """Return FP8 WO-A output and FlashInfer F8_128x4 scale bytes.""" + """Return FP8 WO-A output and FlashInfer F8_128x4 scale bytes. + ``early_trigger`` launches the PDL dependent right after the TMEM + allocation, for a dependent that loads its weights before its wait.""" tokens = x.shape[0] groups, rank = weight_scale.shape[:2] n = groups * rank @@ -592,6 +625,7 @@ def __call__( n_groups=groups, heads_per_group=x.shape[1] // groups, o_lora_rank=rank, + early_trigger=early_trigger, ) return compile_key, launch_args, (q, scales) diff --git a/vllm/models/deepseek_v41/nvidia/ops/o_proj.py b/vllm/models/deepseek_v41/nvidia/ops/o_proj.py index 88a0a4f8be4f..4de3ad74c455 100644 --- a/vllm/models/deepseek_v41/nvidia/ops/o_proj.py +++ b/vllm/models/deepseek_v41/nvidia/ops/o_proj.py @@ -52,6 +52,27 @@ def _can_fuse_wo_a(layer: nn.Module) -> bool: ) +def _fused_wo_a_shape(layer: nn.Module) -> dict[str, int]: + """The fused WO-A kernel's warmup arguments for ``layer``.""" + return dict( + max_tokens=_fused_wo_a_max_tokens(layer), + n_groups=layer.n_local_groups, + heads_per_group=layer.n_local_heads // layer.n_local_groups, + o_lora_rank=layer.o_lora_rank, + ) + + +def precompile_dsv41_o_proj_early_trigger(layer: nn.Module, max_tokens: int) -> None: + """Compile the fused WO-A variant ``dsv41_o_proj`` dispatches with + ``early_trigger``, for its token counts up to ``max_tokens``.""" + if _can_fuse_wo_a(layer): + from .fused_wo_a import _FUSED_WO_A_KERNEL + + shape = _fused_wo_a_shape(layer) + shape["max_tokens"] = min(shape["max_tokens"], max_tokens) + _FUSED_WO_A_KERNEL.warmup(**shape, early_trigger=True) + + def register_dsv41_o_proj_warmup(layer: nn.Module) -> None: """Warm the fused WO-A token counts ``dsv41_o_proj`` can dispatch.""" # Draft models load outside the warmup registry; their layers share the @@ -59,12 +80,7 @@ def register_dsv41_o_proj_warmup(layer: nn.Module) -> None: if _can_fuse_wo_a(layer): from .fused_wo_a import _FUSED_WO_A_KERNEL - _FUSED_WO_A_KERNEL.register_warmup( - max_tokens=_fused_wo_a_max_tokens(layer), - n_groups=layer.n_local_groups, - heads_per_group=layer.n_local_heads // layer.n_local_groups, - o_lora_rank=layer.o_lora_rank, - ) + _FUSED_WO_A_KERNEL.register_warmup(**_fused_wo_a_shape(layer)) def dsv41_o_proj( @@ -72,9 +88,12 @@ def dsv41_o_proj( attn_out: torch.Tensor, positions: torch.Tensor, wo_b_proj: Callable[..., torch.Tensor] | None = None, + early_trigger: bool = False, ) -> torch.Tensor: """``deep_gemm_fp8_o_proj`` with WO-A fused for small SM100/SM103 batches; - ``wo_b_proj`` replaces ``layer._wo_b_proj``.""" + ``wo_b_proj`` replaces ``layer._wo_b_proj``. With ``early_trigger``, the + fused WO-A launches ``wo_b_proj`` (PDL) right after its TMEM allocation, + for a ``wo_b_proj`` that loads its weights before its grid wait.""" wo_b_proj = wo_b_proj or layer._wo_b_proj o = attn_out[:, : layer.n_local_heads, :] cos_sin_cache = layer.rotary_emb.cos_sin_cache @@ -87,6 +106,7 @@ def dsv41_o_proj( rope=cos_sin_cache, weight=layer.wo_a.weight, weight_scale=layer.wo_a.weight_scale, + early_trigger=early_trigger, ) return wo_b_proj( QuantizedActivation(q, scales, o.dtype, q.shape, kMxfp8Dynamic) diff --git a/vllm/models/deepseek_v41/nvidia/specialized_decode/AGENTS.md b/vllm/models/deepseek_v41/nvidia/specialized_decode/AGENTS.md index 724b3aee63c1..b8ded8d05680 100644 --- a/vllm/models/deepseek_v41/nvidia/specialized_decode/AGENTS.md +++ b/vllm/models/deepseek_v41/nvidia/specialized_decode/AGENTS.md @@ -21,7 +21,7 @@ startup on every TP rank with the missing requirements: sparse logits. - CUDA graphs: FULL graphs for decode batches only (`FULL_AND_PIECEWISE` or `FULL_DECODE_ONLY`). -- No `torch.compile` of the model, no sleep mode. +- No `torch.compile` of the model, no sleep mode, no runtime weight updates. - No LoRA, KV transfer or batch invariance. - DeepGEMM and NVLink multicast for the fused all-reduce + mHC. @@ -36,7 +36,7 @@ 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. -Two features use FlashInfer internals and run on flashinfer-python 0.7.2 or +Four features use FlashInfer internals and run on flashinfer-python 0.7.2 or newer (`MIN_FLASHINFER_VERSION`): - the publishing `wo_b` GEMM: the CuTe DSL low-latency block-scaled GEMM @@ -44,16 +44,23 @@ newer (`MIN_FLASHINFER_VERSION`): (`flashinfer.gemm.gemm_base`) to look up its tuned tactic; - the fused kernel's MXFP8 layer input: the CuTe DSL MXFP8 quantize helpers (`flashinfer.quantization.quantization_cute_dsl_utils`, - `flashinfer.cute_dsl.fp4_common`). + `flashinfer.cute_dsl.fp4_common`); +- the pinned `cute-dsl` MXFP8 GEMM tactics above 8 rows: `mm_mxfp8`'s private + `cute-dsl` runner and workspace (`flashinfer.gemm.gemm_base`), called with + the pinned tactic; +- the low-latency MXFP8 GEMM above 8 rows: the kernel class + (`flashinfer.gemm.kernels.cute_dsl.low_latency_blockscaled_gemm`) built + with a 16-, 32- or 48-row activation tile. On an older version each logs a warning once and the executor goes without -it. `test_numerics.py` compares both with FlashInfer's own GEMM and quantize -bitwise; run it when `requirements/cuda.txt` moves to a new FlashInfer. +it. `test_numerics.py` compares the first two with FlashInfer's own GEMM and +quantize bitwise, and `test_kernels.py` the last two with the stock layer; run +them when `requirements/cuda.txt` moves to a new FlashInfer. Two more depend on FlashInfer's interface, not its version: -- the pinned MXFP8 GEMM tactics (GB300) wrap the low-latency GEMM's - `autotune_tactics`; without it, or when it does not list a pinned tactic, +- the pinned low-latency MXFP8 GEMM tactics up to 8 rows (GB300) wrap that + GEMM's `autotune_tactics`; without it, or when it does not list a pinned tactic, the GEMMs tune over every tactic; - the sparse-MLA decode's persistent KV counters need the `multi_ctas_kv_counter_buffer` argument (0.7.1 added it); without it each @@ -75,6 +82,7 @@ forward runs the stock layers, which the executor does not modify. | `boundary/` | The mHC between sublayers on the fused all-reduce + mHC kernel, and the publishing `wo_b` GEMM | | `attention/` | The attention step: low-batch MXFP8 and BF16 projections, the row-sharded indexer, candidate overlap, shared sparse-index tables and the sparse-MLA decode | | `moe/` | The MoE step with a deferred finalize, and the low-batch routed-expert GEMMs | +| `embedding/` | The embedding and first-layer hc streams from a replicated table | | `engram/` | The Engram WKV ahead of the decoder layers | | `logits/` | The low-batch LM head and the logits gather | @@ -82,24 +90,30 @@ forward runs the stock layers, which the executor does not modify. | Path | Steps | | --- | --- | -| Fused Lamport all-reduce + mHC at every sublayer boundary; the next mHC coefficients on the side stream up to 48 rows, inline beyond | all executor steps (up to 384 rows) | +| Fused Lamport all-reduce + mHC at every sublayer boundary; the next mHC coefficients on the side stream | all executor steps (up to 384 rows) | | Every MoE finalize deferred to the next boundary's fused kernel, in a residual-only variant before an Engram layer and after the last layer | all executor steps | | The fused kernel also writes the attention input as MXFP8, which `fused_wqa_wkv` reads instead of quantizing again | all executor steps | -| Aux hidden-state means on the mHC side stream | up to 48 rows | -| Routed experts on CuTe DSL routing and `tcgen05` GEMM kernels (`moe/routed_expert_gemm.py`) | 1-8 rows | +| Aux hidden-state means on the mHC side stream | all executor steps | +| Routed experts on CuTe DSL routing and `tcgen05` GEMM kernels (`moe/routed_expert_gemm.py`) | 1-192 rows | | `fused_wqa_wkv`, `wq_b` and `wo_b` on FlashInfer's low-latency CuTe DSL MXFP8 GEMM with pinned tactics | 1-8 rows | +| `fused_wqa_wkv`, `wq_b` and `wo_b` on the low-latency GEMM with a 16-, 32- or 48-row activation tile where its shape has a wide tactic, else on FlashInfer's `cute-dsl` MXFP8 GEMM with pinned tactics | 9-48 rows | +| `fused_wqa_wkv`, `wq_b` and `wo_b` on FlashInfer's `cute-dsl` MXFP8 GEMM with pinned tactics instead of per-start autotuning | 49-384 rows | | Compressor, indexer-weight and indexer `wk` projections on a low-batch BF16 GEMM; the index key store reads `wk`'s FP32 output | 1-8 rows | | `wo_b` stores its output into the fused kernel's mailbox, which then only collects | 1-8 rows | +| The decode Q/KV norm and the fused `wo_a` trigger their PDL dependents (`wq_b`, `wo_b`) early, as those load their weights before their own wait | 1-48 rows | +| The embedding and the first layer's hc streams looked up from a full copy of the table on every rank (1.32 GB a GPU), without the all-reduce | all executor steps | | Indexer query, compressor and index key on side streams, each joined at its first reader | all executor steps; the side projections up to `VLLM_MULTI_STREAM_GEMM_TOKEN_THRESHOLD` rows | -| Candidate consumers' candidate expansion and MQA-logits schedule on a side stream after the candidate source's indexer | up to 16 rows | -| Dense indexer rows sharded over the TP ranks and all-gathered | steps of at least 8 requests that TP divides | -| Candidate blocks selected with `cooperative_topk` | up to 64 rows per rank | +| Candidate consumers' candidate expansion and MQA-logits schedule on a side stream after the candidate source's indexer | all executor steps | +| Dense indexer rows sharded over the TP ranks, a contiguous even slice per rank (the last ones may be short or empty), and all-gathered with the executor's Lamport all-gather (`ops/cute_dsl/all_gather.py`), over NCCL above its 4 MiB a rank | uniform decode steps of at least 12 rows | +| Token top-k, and the candidate source's blocks, selected from the BF16 logits with `decode_topk` (`attention/decode_topk_cute.py`) | up to 64 rows per rank; the candidate source while its 16-CTA rows fit two waves (14 rows on a GB300) | +| Otherwise the candidate blocks selected with `cooperative_topk` on the upcast logits | up to 64 rows per rank | | Compressed layers sharing an index source reuse one sparse-index table; the first of them builds the other KV-cache groups' tables on a side stream | all executor steps | -| Sparse-MLA decode launched on a KV-counter buffer zeroed once | all executor steps | -| Engram WKV chain on a side stream right after each layer's lookup | up to 16 rows, with host-offloaded TP-only tables and a CuTe DSL MXFP8 `wkv` | +| Sparse-MLA decode on a CuTe DSL kernel, a thread block cluster per token with one CTA per 128 columns (`attention/sparse_mla_cute.py`) | the target's causal decode while every token's cluster is resident at once (at width 640, up to 8 requests with DSpark on a GB300) | +| Otherwise the TRTLLM-gen sparse-MLA decode, launched on a KV-counter buffer zeroed once | all other executor steps | +| Engram WKV chain on a side stream right after each layer's lookup | all executor steps, with host-offloaded TP-only tables and a CuTe DSL MXFP8 `wkv` | | Attention metadata builds recorded at CUDA graph capture, one per batch shape the FULL decode graphs serve, and replayed | FULL-graph uniform decode | | LM head on FlashInfer's CuTe DSL BF16 GEMM | 1-16 rows | -| Logits gathered with the custom all-gather | whenever it accepts the input | +| Logits gathered with the executor's Lamport all-gather | up to 4 MiB a rank (at TP4, 32 FP32 or 64 BF16 rows) | ## Dispatch @@ -108,4 +122,4 @@ rather than fall back. Every per-step gate reads only host state that all TP ranks share (row and request counts, graph mode, the attention metadata's host fields); nothing is read back from the GPU. The LM head and logits gather run outside the CUDA graph (Model Runner V2 computes logits eagerly), gated by row -count and by the custom all-gather. +count and by the all-gather's size limit. diff --git a/vllm/models/deepseek_v41/nvidia/specialized_decode/__init__.py b/vllm/models/deepseek_v41/nvidia/specialized_decode/__init__.py index a08a284cdac5..5f0f4df389f5 100644 --- a/vllm/models/deepseek_v41/nvidia/specialized_decode/__init__.py +++ b/vllm/models/deepseek_v41/nvidia/specialized_decode/__init__.py @@ -32,12 +32,12 @@ # Rows of the executor's fused all-reduce + mHC mailbox: the largest FULL # decode graph at 64 requests x 6 DSpark tokens. FUSED_MHC_MAX_ROWS = 384 -# Rows whose next mHC coefficients run on the side stream; larger steps compute -# them inline after the fused kernel. -MHC_OVERLAP_MAX_ROWS = 48 # Rows of the low-batch GEMMs: one DSpark verification step of up to 8 rows # (6 at one request). LOW_BATCH_MAX_ROWS = 8 +# Rows up to which the attention projections above LOW_BATCH_MAX_ROWS keep the +# low-latency MXFP8 GEMM, with a 16-, 32- or 48-row activation tile. +WIDE_LOW_LATENCY_MAX_ROWS = 48 # The oldest FlashInfer whose internals the publishing wo_b GEMM and the MXFP8 # layer input were checked against. Newer releases run them too: test_numerics.py # compares both with FlashInfer's own GEMM and quantize bitwise on the release @@ -118,6 +118,8 @@ def unsupported(vllm_config: VllmConfig) -> list[str]: ), # The routed-expert GEMMs and the publishing wo_b bind weight addresses. "no sleep mode": not vllm_config.model_config.enable_sleep_mode, + # The replicated embedding table is copied from the shards once, at load. + "no runtime weight updates": vllm_config.weight_transfer_config is None, "no LoRA, KV transfer or batch invariance": ( vllm_config.lora_config is None and vllm_config.kv_transfer_config is None diff --git a/vllm/models/deepseek_v41/nvidia/specialized_decode/attention/decode_topk_cute.py b/vllm/models/deepseek_v41/nvidia/specialized_decode/attention/decode_topk_cute.py new file mode 100644 index 000000000000..eee32bcb0552 --- /dev/null +++ b/vllm/models/deepseek_v41/nvidia/specialized_decode/attention/decode_topk_cute.py @@ -0,0 +1,1331 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Decode top-k of the dense indexers straight from their BF16 logits, in CuTe DSL. + +The problem. In a decode step each dense indexer layer scores every compressed +token of a row (SM100's paged MQA logits are BF16) and keeps the ``k`` largest: +the token top-512, and in the candidate source also the top-2048 of the rows' +8-token block maxima. ``cooperative_topk`` reads FP32, so the logits used to be +upcast first, and it builds a radix histogram of every value with shared-memory +atomics. Real indexer logits crowd a few histogram bins (a row of 150K logits +holds about 200 distinct BF16 values), where those atomics serialize. + +What this module computes. ``DecodeTopK`` writes, for each row, the indices of +its ``k`` largest elements in any order (exact: every index whose value exceeds +the k-th largest, then enough of those equal to it), or ``0 .. n-1`` and -1 +padding when the row has at most ``k`` elements. An element is one BF16 logit +(token mode), or the NaN-propagating maximum of an aligned group of 8 logits +(block mode, the candidate blocks), with the row's newest block pinned at +inf +and blocks scoring -inf written as -1. Values are compared through an +order-preserving 16-bit key of their BF16 bits, so ties are exact. + +How. One thread block cluster of ``C`` CTAs per row: 8 CTAs of 512 threads +while every row's cluster can be resident at once, else 4 of 1024 (a second +wave of clusters would nearly double the time). Each CTA copies one +contiguous chunk of the row into shared memory, as 16-byte vectors (8 logits, +or one block); the chunks split the row's valid length, not its width. A +thread stages 8 vectors and keeps them in registers too, so a cluster holds +262,144 logits. Wider rows need more: at a 1M context the ratio-2 layers' +logits are 524,288 wide, and a second build stages 16 vectors a thread at 8 +CTAs, the other 8 in shared memory only. The candidate source layer scores +every token (ratio 1), 1,048,576 at a 1M context, so a third build runs 16 +CTAs a row (a non-portable cluster size, which CuTe DSL allows on every +kernel) with 16 vectors a thread. Rows of the wide builds run in waves when +they outnumber the resident clusters; the 16-CTA build takes at most two +waves (`supports_decode_topk`). Two ways to the answer: + +- Sampled (token rows longer than one vector a thread):: + + sample the row ──► every CTA: the same cutoff (the q-th largest of the + row's 2048 samples, q aiming at 3k survivors) + └► keep keys > cutoff and up to k keys == cutoff, push both to CTA 0 + └► CTA 0: exact k-th largest among <= 8192 survivors (two 8-bit + radix passes), write k indices + +- Direct (block mode, short token rows, and a sample that misjudged the row: + more than 8192 survivors, or fewer than k at or above the cutoff): the exact + k-th largest key from two cluster-wide 8-bit radix histograms, then every + CTA writes its own selected indices straight to the output. + +Both filter the same way: a thread tests each of its vectors' maximum against +the cutoff, the vectors that can hold selected keys go to a CTA queue, and +the CTA's threads unpack the queued vectors evenly; a CTA scan gives each +thread its output slots and one cluster atomic reserves the CTA's range, so no +key costs an atomic. Whole-row histograms spread each warp's atomics over 8 +copies (lane l counts into copy l mod 8), since a crowded bin serializes the +atomics on its word; the radix selects over the few thousand samples or +survivors count into one histogram, where the copies' zeroing and summing +cost more than they save. + +Synchronization. Cluster barriers separate the phases: every CTA has started +and zeroed its counters before any distributed shared memory access; samples +are complete before the cutoff; survivors are complete before CTA 0 reads +them; histogram contributions are complete before they are read; and CTA 0 +waits until every peer is done with its shared memory before it exits. The +kernel waits on its predecessor grid before reading inputs and lets its +dependents launch once its chunk is in shared memory (programmatic dependent +launch). +""" + +from functools import cache + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import torch +from cutlass import Int32, Int64, Uint16, Uint32 +from cutlass._mlir.dialects import llvm +from cutlass.cute.runtime import make_fake_stream, make_ptr +from cutlass.cutlass_dsl import T, dsl_user_op + +from vllm.models.common.ops.moe_tail_fusion.primitives import current_cu_stream + +SAMPLES = 2048 # keys sampled a row, over the cluster +CAPACITY = 8192 # survivors a row +SURVIVOR_FACTOR = 3 # the cutoff aims at this many times k survivors +VECTORS = 8 # 16-byte vectors a thread keeps in registers (of 8, or 16 for wide rows) +NEG_INF_KEY = 0x007F # key of BF16 -inf +POS_INF_KEY = 0xFF80 # key of BF16 +inf +# ``ctl`` slots (shared Uint32 words). +BIN, REMAIN = 0, 1 # select_bin's result +BASE_GT, BASE_EQ = 2, 3 # this CTA's reserved first slots +TOTALS = 8 # at CTA 0: (gt, eq) totals, one 8-byte aligned pair a pass +CTL_WORDS = 16 +SCAN_WORDS = 160 +COPIES = 8 # copies of a whole-row histogram, [copy][bin]; lane l uses copy l % 8 +TOP_KS = (512, 1024, 2048) # the k values a kernel is built for + + +def _asm(result, operands, text, constraints, side_effects, loc, ip): + """One PTX inline assembly operation.""" + return llvm.inline_asm( + result, + operands, + text, + constraints, + has_side_effects=side_effects, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +def _v(value, loc, ip): + """IR value of a DSL scalar.""" + return value.ir_value(loc=loc, ip=ip) + + +def _u32x4(loaded, loc, ip): + """The four words of an inline-assembly struct result.""" + return tuple( + Uint32(llvm.extractvalue(T.i32(), loaded, [index], loc=loc, ip=ip)) + for index in range(4) + ) + + +@dsl_user_op +def load_u32x4(address: Int64, *, loc=None, ip=None): + """Four 32-bit words at a 16-byte aligned global address.""" + loaded = _asm( + llvm.StructType.get_literal([T.i32()] * 4), + [_v(address, loc, ip)], + "ld.global.v4.u32 {$0, $1, $2, $3}, [$4];", + "=r,=r,=r,=r,l", + False, + loc, + ip, + ) + return _u32x4(loaded, loc, ip) + + +@dsl_user_op +def shared_load_u32x4(address: Int32, *, loc=None, ip=None): + """Four 32-bit words at a 16-byte aligned shared address of this CTA.""" + loaded = _asm( + llvm.StructType.get_literal([T.i32()] * 4), + [_v(Int32(address), loc, ip)], + "ld.shared.v4.u32 {$0, $1, $2, $3}, [$4];", + "=r,=r,=r,=r,r", + True, + loc, + ip, + ) + return _u32x4(loaded, loc, ip) + + +@dsl_user_op +def shared_store_u32x4(address: Int32, w0, w1, w2, w3, *, loc=None, ip=None) -> None: + """Store four 32-bit words at a 16-byte aligned shared address of this CTA.""" + _asm( + None, + [_v(Int32(address), loc, ip)] + + [_v(Uint32(w), loc, ip) for w in (w0, w1, w2, w3)], + "st.shared.v4.u32 [$0], {$1, $2, $3, $4};", + "r,r,r,r,r", + True, + loc, + ip, + ) + + +@dsl_user_op +def store_s32(address: Int64, value: Int32, *, loc=None, ip=None) -> None: + """Store a 32-bit word to global memory.""" + _asm( + None, + [_v(address, loc, ip), _v(Int32(value), loc, ip)], + "st.global.s32 [$0], $1;", + "l,r", + True, + loc, + ip, + ) + + +@dsl_user_op +def max_nan_bf16x2(a: Uint32, b: Uint32, *, loc=None, ip=None) -> Uint32: + """Lane-wise maximum of two BF16 pairs, NaN if either is NaN.""" + return Uint32( + _asm( + T.i32(), + [_v(Uint32(a), loc, ip), _v(Uint32(b), loc, ip)], + "max.NaN.bf16x2 $0, $1, $2;", + "=r,r,r", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def shfl_up(value: Uint32, delta: int, *, loc=None, ip=None) -> Uint32: + """``value`` of the lane ``delta`` below (own value in the lowest lanes).""" + return Uint32( + _asm( + T.i32(), + [_v(Uint32(value), loc, ip)], + f"shfl.sync.up.b32 $0, $1, {delta}, 0, 0xffffffff;", + "=r,r", + True, + loc, + ip, + ) + ) + + +@dsl_user_op +def shared_red(address: Int32, value: Uint32, *, loc=None, ip=None) -> None: + """Atomic add to a shared word of this CTA.""" + _asm( + None, + [_v(Int32(address), loc, ip), _v(Uint32(value), loc, ip)], + "red.shared.add.u32 [$0], $1;", + "r,r", + True, + loc, + ip, + ) + + +@dsl_user_op +def peer_address(address: Int32, rank, *, loc=None, ip=None) -> Int32: + """The cluster address of this CTA's shared ``address`` in CTA ``rank``.""" + return Int32( + _asm( + T.i32(), + [_v(Int32(address), loc, ip), _v(Int32(rank), loc, ip)], + "mapa.shared::cluster.u32 $0, $1, $2;", + "=r,r,r", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def cluster_add_u64(address: Int32, value: Int64, *, loc=None, ip=None) -> Int64: + """Atomic add to an 8-byte aligned cluster shared word pair; returns the + old value.""" + return Int64( + _asm( + T.i64(), + [_v(Int32(address), loc, ip), _v(Int64(value), loc, ip)], + "atom.shared::cluster.add.u64 $0, [$1], $2;", + "=l,r,l", + True, + loc, + ip, + ) + ) + + +@dsl_user_op +def cluster_red(address: Int32, value: Uint32, *, loc=None, ip=None) -> None: + """Atomic add to a cluster shared word.""" + _asm( + None, + [_v(Int32(address), loc, ip), _v(Uint32(value), loc, ip)], + "red.shared::cluster.add.u32 [$0], $1;", + "r,r", + True, + loc, + ip, + ) + + +@dsl_user_op +def cluster_load(address: Int32, *, loc=None, ip=None) -> Uint32: + """Load a cluster shared word.""" + return Uint32( + _asm( + T.i32(), + [_v(Int32(address), loc, ip)], + "ld.shared::cluster.u32 $0, [$1];", + "=r,r", + True, + loc, + ip, + ) + ) + + +@dsl_user_op +def cluster_store_u32(address: Int32, value, *, loc=None, ip=None) -> None: + """Store a 32-bit word to cluster shared memory.""" + _asm( + None, + [_v(Int32(address), loc, ip), _v(Uint32(value), loc, ip)], + "st.shared::cluster.u32 [$0], $1;", + "r,r", + True, + loc, + ip, + ) + + +@dsl_user_op +def cluster_store_u16(address: Int32, value, *, loc=None, ip=None) -> None: + """Store the low half of a word to cluster shared memory.""" + _asm( + None, + [_v(Int32(address), loc, ip), _v(Uint32(value), loc, ip)], + "{\n.reg .b16 h;\ncvt.u16.u32 h, $1;\nst.shared::cluster.u16 [$0], h;\n}", + "r,r", + True, + loc, + ip, + ) + + +@cute.jit +def order_key(bits: Uint32) -> Uint32: + """The order-preserving key of the BF16 value in the low half of ``bits``: + a larger value has a larger key, NaNs of either sign outside the + numbers.""" + b = bits & Uint32(0xFFFF) + return b ^ (Uint32(0x8000) | ((Uint32(0) - (b >> Uint32(15))) & Uint32(0x7FFF))) + + +@cute.jit +def block_key(w0: Uint32, w1: Uint32, w2: Uint32, w3: Uint32) -> Uint32: + """The key of the NaN-propagating maximum of eight BF16 values.""" + m = max_nan_bf16x2(max_nan_bf16x2(w0, w1), max_nan_bf16x2(w2, w3)) + m = max_nan_bf16x2(m, (m >> Uint32(16)) | (m << Uint32(16))) + return order_key(m) + + +@cute.jit +def hist_add( + hist_address: Int32, bin_: Uint32, valid, copies: bool, lane: Int32 +) -> None: + """Count each valid lane's ``bin_`` with one shared atomic. ``copies``: + lane ``l`` counts into copy ``l % COPIES`` of a ``[COPIES][256]`` + histogram (whole-row histograms: crowded bins would serialize the + atomics of one copy); else into a plain 256-word one.""" + if valid: + slot = Int32(bin_) + if cutlass.const_expr(copies): + slot = slot + (lane & (COPIES - 1)) * 256 + shared_red(hist_address + slot * 4, Uint32(1)) + + +@cute.jit +def select_bin(hist: cute.Tensor, q: Uint32, ctl: cute.Tensor, lane: Int32) -> None: + """Warp 0: the bin ``b`` of a 256-bin histogram holding the ``q``-th + largest entry (q >= 1), into ``ctl[BIN]``, and ``q`` minus the entries + in the bins above ``b`` into ``ctl[REMAIN]``. Lane ``l`` scans bins + ``255 - 8 l`` down to ``248 - 8 l``.""" + low = shared_load_u32x4(Int32(hist.iterator.toint()) + (248 - 8 * lane) * 4) + high = shared_load_u32x4(Int32(hist.iterator.toint()) + (252 - 8 * lane) * 4) + counts = (high[3], high[2], high[1], high[0], low[3], low[2], low[1], low[0]) + local = Uint32(0) + for i in cutlass.range_constexpr(8): + local = local + counts[i] + inclusive = local + for shift in cutlass.range_constexpr(5): + below = shfl_up(inclusive, 1 << shift) + if lane >= (1 << shift): + inclusive = inclusive + below + above = inclusive - local + if (above < q) & (inclusive >= q): + found = Int32(-1) + remain = Uint32(0) + for i in cutlass.range_constexpr(8): + if (found < 0) & (above + counts[i] >= q): + found = Int32(255 - 8 * lane - i) + remain = q - above + above = above + counts[i] + ctl[BIN] = Uint32(found) + ctl[REMAIN] = remain + + +class DecodeTopK: + """The decode top-k (module docstring) of rows of BF16 logits. + + ``k``: 512, 1024 or 2048; ``cluster``: 4, 8 or 16 CTAs a row (4: 1024 + threads, else 512); ``vector_count``: 16-byte vectors a thread stages (8, + or 16 at 8 or 16 CTAs: the first 8 also stay in registers), so a row + holds at most ``cluster * threads * vector_count`` vectors; + ``block_mode``: elements are 8-logit block maxima. + """ + + def __init__( + self, k: int, cluster: int, vector_count: int, block_mode: bool + ) -> None: + assert k in TOP_KS and cluster in (4, 8, 16) + assert vector_count in vector_choices(cluster) + self.k = k + self.cluster = cluster + self.threads = threads_for(cluster) + self.warps = self.threads // 32 + self.block_mode = block_mode + self.per_vector = 1 if block_mode else 8 # elements a 16-byte vector holds + self.vector_count = vector_count + self.vectors = self.threads * vector_count # a CTA's vectors + self.samples = SAMPLES // cluster # a CTA's samples, one a thread + assert self.samples <= self.threads + self.max_elements = cluster * self.vectors * self.per_vector + # Token rows of at most one vector a thread take the direct path. + self.direct_elements = cluster * self.threads * self.per_vector + + @cute.jit + def __call__( + self, + logits_ptr: cute.Pointer, + lengths_ptr: cute.Pointer, + out_ptr: cute.Pointer, + rows: Int32, + row_bytes: Int64, + width: Int32, + out_stride: Int32, + fallback: Int32, + stream: cuda.CUstream, + ): + """``logits``: BF16 rows ``row_bytes`` apart (16-byte aligned), each + valid below ``lengths[r]`` (int32, clamped to [0, width]); ``out``: + int32 rows ``out_stride`` apart. ``fallback`` nonzero forces the + direct path (tests).""" + self.kernel( + cute.make_tensor(lengths_ptr, cute.make_layout(rows)), + cute.make_tensor(out_ptr, cute.make_layout(rows * out_stride)), + logits_ptr.toint(), + row_bytes, + width, + out_stride, + fallback, + ).launch( + grid=(rows * self.cluster, 1, 1), + block=[self.threads, 1, 1], + cluster=[self.cluster, 1, 1], + stream=stream, + use_pdl=True, + ) + + @cute.jit + def vector_words(self, data_address: Int32, vector: Int32): + """The four words of the chunk's vector ``vector`` in shared memory.""" + return shared_load_u32x4(data_address + vector * 16) + + @cute.jit + def vector_top(self, words, element: Int32, pinned: Int32) -> Uint32: + """The largest key of a staged vector whose first element is + ``element`` (block mode: the pinned newest block ranks first).""" + top = block_key(words[0], words[1], words[2], words[3]) + if cutlass.const_expr(self.block_mode): # noqa: SIM102 + if element == pinned: + top = Uint32(POS_INF_KEY) + return top + + @cute.jit + def word_key(self, words, e: int, element: Int32, pinned: Int32) -> Uint32: + """The key of element ``e`` of a vector's ``words`` (``element``: its + row index, for the pinned block).""" + key = Uint32(0) + if cutlass.const_expr(self.block_mode): + key = block_key(words[0], words[1], words[2], words[3]) + if element == pinned: + key = Uint32(POS_INF_KEY) + else: + key = order_key(Uint32(words[e // 2]) >> Uint32(16 * (e % 2))) + return key + + @cute.jit + def scan2(self, a: Uint32, b: Uint32, scan: cute.Tensor, tid: Int32): + """Exclusive CTA prefix sums of ``a`` and ``b`` over the threads, and + their totals (every thread calls it; three barriers).""" + lane = tid & 31 + warp = tid >> 5 + in_a = a + in_b = b + for shift in cutlass.range_constexpr(5): + below_a = shfl_up(in_a, 1 << shift) + below_b = shfl_up(in_b, 1 << shift) + if lane >= (1 << shift): + in_a = in_a + below_a + in_b = in_b + below_b + cute.arch.sync_threads() # the previous scan's readers are done + if lane == 31: + scan[warp] = in_a + scan[32 + warp] = in_b + cute.arch.sync_threads() + if warp == 0: + sum_a = Uint32(0) + sum_b = Uint32(0) + if lane < self.warps: + sum_a = Uint32(scan[lane]) + sum_b = Uint32(scan[32 + lane]) + run_a = sum_a + run_b = sum_b + for shift in cutlass.range_constexpr(5): + below_a = shfl_up(run_a, 1 << shift) + below_b = shfl_up(run_b, 1 << shift) + if lane >= (1 << shift): + run_a = run_a + below_a + run_b = run_b + below_b + scan[64 + lane] = run_a - sum_a + scan[96 + lane] = run_b - sum_b + if lane == 31: + scan[128] = run_a + scan[129] = run_b + cute.arch.sync_threads() + first_a = Uint32(scan[64 + warp]) + in_a - a + first_b = Uint32(scan[96 + warp]) + in_b - b + return first_a, first_b, Uint32(scan[128]), Uint32(scan[129]) + + @cute.jit + def radix_select( + self, + keys: cute.Tensor, + count: Int32, + q: Uint32, + hist: cute.Tensor, + hist2: cute.Tensor, + ctl: cute.Tensor, + tid: Int32, + ): + """The ``q``-th largest of ``keys[0:count]`` (shared Uint16, q <= count) + and how many of the keys equal to it rank within the top ``q``. Every + thread calls it.""" + lane = tid & 31 + h1 = Int32(hist.iterator.toint()) + h2 = Int32(hist2.iterator.toint()) + if tid < 256: + hist[tid] = Uint32(0) + hist2[tid] = Uint32(0) + cute.arch.sync_threads() + iterations = (count + self.threads - 1) // self.threads + for it in cutlass.range(iterations, unroll=4): + i = it * self.threads + tid + valid = i < count + key = Uint32(0) + if valid: + key = Uint32(keys[i]) + hist_add(h1, key >> Uint32(8), valid, False, lane) + cute.arch.sync_threads() + if tid < 32: + select_bin(hist, q, ctl, lane) + cute.arch.sync_threads() + high = Uint32(ctl[BIN]) + remain = Uint32(ctl[REMAIN]) + for it in cutlass.range(iterations, unroll=4): + i = it * self.threads + tid + key = Uint32(0) + if i < count: + key = Uint32(keys[i]) + hist_add( + h2, + key & Uint32(0xFF), + (i < count) & ((key >> Uint32(8)) == high), + False, + lane, + ) + cute.arch.sync_threads() + if tid < 32: + select_bin(hist2, remain, ctl, lane) + cute.arch.sync_threads() + tau = (high << Uint32(8)) | Uint32(ctl[BIN]) + need = Uint32(ctl[REMAIN]) + cute.arch.sync_threads() + return tau, need + + @cute.jit + def queue_vectors( + self, + regs, + data_address: Int32, + cutoff: Uint32, + start: Int32, + length: Int32, + pinned: Int32, + tid: Int32, + scan, + queue, + ) -> Int32: + """Queue this CTA's vectors that can hold a key >= ``cutoff`` (by + their maximum) and return how many (every thread).""" + threads, per = self.threads, self.per_vector + hits = Uint32(0) + for j in cutlass.range_constexpr(VECTORS): + first = (j * threads + tid) * per + if first < length: + words = (regs[j, 0], regs[j, 1], regs[j, 2], regs[j, 3]) + if self.vector_top(words, start + first, pinned) >= cutoff: + hits = hits | Uint32(1 << j) + if cutlass.const_expr(self.vector_count > VECTORS): # noqa: SIM102 + # Only a chunk reaching the shared-only vectors visits them. + if VECTORS * threads * per < length: + for j in cutlass.range_constexpr(VECTORS, self.vector_count): + first = (j * threads + tid) * per + if first < length: + words = self.vector_words(data_address, j * threads + tid) + if self.vector_top(words, start + first, pinned) >= cutoff: + hits = hits | Uint32(1 << j) + slot, _, total, _ = self.scan2( + Uint32(cute.arch.popc(hits)), Uint32(0), scan, tid + ) + position = Int32(slot) + if hits != Uint32(0): + for j in cutlass.range_constexpr(self.vector_count): + if (hits >> Uint32(j)) & Uint32(1) != Uint32(0): + queue[position] = Uint16(j * threads + tid) + position = position + 1 + cute.arch.sync_threads() + return Int32(total) + + @cute.jit + def count_queued( + self, + data_address: Int32, + queued: Int32, + cutoff: Uint32, + start: Int32, + length: Int32, + pinned: Int32, + tid: Int32, + queue, + ): + """This thread's keys above and equal to ``cutoff`` in its share of + the queued vectors (entries ``tid``, ``tid + threads``, ...).""" + per = self.per_vector + gt = Uint32(0) + eq = Uint32(0) + for q in cutlass.range(tid, queued, self.threads, unroll=1): + vector = Int32(queue[q]) + words = self.vector_words(data_address, vector) + for e in cutlass.range_constexpr(per): + element = vector * per + e + key = self.word_key(words, e, start + element, pinned) + if element < length: + if key > cutoff: + gt = gt + Uint32(1) + if key == cutoff: + eq = eq + Uint32(1) + return gt, eq + + @cute.jit + def reserve(self, gt: Uint32, eq: Uint32, totals: int, tid: Int32, ctl, scan): + """This thread's first slots for its ``gt`` / ``eq`` keys in the + cluster-wide lists counted at CTA 0's ``ctl[totals]`` pair: a CTA + scan, then one 64-bit cluster atomic for the CTA.""" + first_gt, first_eq, cta_gt, cta_eq = self.scan2(gt, eq, scan, tid) + if tid == 0: + old = cluster_add_u64( + peer_address(Int32(ctl.iterator.toint()) + totals * 4, 0), + Int64(cta_gt) | (Int64(cta_eq) << Int64(32)), + ) + ctl[BASE_GT] = Uint32(old & Int64(0xFFFFFFFF)) + ctl[BASE_EQ] = Uint32(old >> Int64(32)) + cute.arch.sync_threads() + return Int32(Uint32(ctl[BASE_GT]) + first_gt), Int32( + Uint32(ctl[BASE_EQ]) + first_eq + ) + + @cute.jit + def gather( + self, + regs, + data_address: Int32, + cutoff: Uint32, + start: Int32, + length: Int32, + pinned: Int32, + tid: Int32, + ctl, + scan, + queue, + g_key, + g_idx, + g_eq, + ): + """Sampled path: push this CTA's keys above ``cutoff`` and up to k + equal to it to CTA 0 and return the cluster's totals (every CTA).""" + k, per = self.k, self.per_vector + queued = self.queue_vectors( + regs, data_address, cutoff, start, length, pinned, tid, scan, queue + ) + gt, eq = self.count_queued( + data_address, queued, cutoff, start, length, pinned, tid, queue + ) + slot_gt, slot_eq = self.reserve(gt, eq, TOTALS, tid, ctl, scan) + key_remote = peer_address(Int32(g_key.iterator.toint()), 0) + idx_remote = peer_address(Int32(g_idx.iterator.toint()), 0) + eq_remote = peer_address(Int32(g_eq.iterator.toint()), 0) + for q in cutlass.range(tid, queued, self.threads, unroll=1): + vector = Int32(queue[q]) + words = self.vector_words(data_address, vector) + for e in cutlass.range_constexpr(per): + element = vector * per + e + key = self.word_key(words, e, start + element, pinned) + if element < length: + if key > cutoff: + if slot_gt < CAPACITY: + cluster_store_u16(key_remote + slot_gt * 2, key) + cluster_store_u32( + idx_remote + slot_gt * 4, Uint32(start + element) + ) + slot_gt = slot_gt + 1 + if key == cutoff: + if slot_eq < k: + cluster_store_u32( + eq_remote + slot_eq * 4, Uint32(start + element) + ) + slot_eq = slot_eq + 1 + cute.arch.cluster_arrive() + cute.arch.cluster_wait() + totals = peer_address(Int32(ctl.iterator.toint()) + TOTALS * 4, 0) + return cluster_load(totals), cluster_load(totals + 4) + + @cute.jit + def count_vector( + self, + level: int, + words, + first: Int32, + start: Int32, + length: Int32, + pinned: Int32, + high: Uint32, + hist_address: Int32, + lane: Int32, + ) -> None: + """Count a staged vector (``words``, its first element ``first``) into + a whole-row histogram: level 0 the high bytes of its keys, level 1 + the low bytes of those whose high byte is ``high``.""" + per = self.per_vector + if cutlass.const_expr(level == 0): + if cute.arch.vote_any_sync(first < length): + for e in cutlass.range_constexpr(per): + key = self.word_key(words, e, start + first + e, pinned) + hist_add( + hist_address, key >> Uint32(8), first + e < length, True, lane + ) + else: + # Only vectors reaching the selected high byte count. + top = self.vector_top(words, start + first, pinned) + reach = (first < length) & ((top >> Uint32(8)) >= high) + if cute.arch.vote_any_sync(reach): + for e in cutlass.range_constexpr(per): + key = self.word_key(words, e, start + first + e, pinned) + hist_add( + hist_address, + key & Uint32(0xFF), + reach & (first + e < length) & ((key >> Uint32(8)) == high), + True, + lane, + ) + + @cute.jit + def exact_cutoff( + self, + regs, + data_address: Int32, + start: Int32, + length: Int32, + pinned: Int32, + tid: Int32, + ctl, + hist, + hist_a, + hist_b, + ): + """The row's exact k-th largest key and how many keys equal to it + rank within the top k, from every CTA's vectors: two cluster-wide + radix histograms gathered at CTA 0 (every CTA).""" + threads, per = self.threads, self.per_vector + lane = tid & 31 + hist_address = Int32(hist.iterator.toint()) + high = Uint32(0) + remain = Uint32(self.k) + for level in cutlass.range_constexpr(2): + target = hist_a if level == 0 else hist_b + for i in cutlass.range(tid, 256 * COPIES, threads, unroll=1): + hist[i] = Uint32(0) + cute.arch.sync_threads() + for j in cutlass.range_constexpr(VECTORS): + self.count_vector( + level, + (regs[j, 0], regs[j, 1], regs[j, 2], regs[j, 3]), + (j * threads + tid) * per, + start, + length, + pinned, + high, + hist_address, + lane, + ) + if cutlass.const_expr(self.vector_count > VECTORS): # noqa: SIM102 + # Only a chunk reaching the shared-only vectors visits them, + # and only warps reaching a vector load it. + if VECTORS * threads * per < length: + for j in cutlass.range_constexpr(VECTORS, self.vector_count): + first = (j * threads + tid) * per + if cute.arch.vote_any_sync(first < length): + words = self.vector_words(data_address, j * threads + tid) + self.count_vector( + level, + words, + first, + start, + length, + pinned, + high, + hist_address, + lane, + ) + cute.arch.sync_threads() + target_remote = peer_address(Int32(target.iterator.toint()), 0) + if tid < 256: + count = Uint32(hist[tid]) + for c in cutlass.range_constexpr(1, COPIES): + count = count + Uint32(hist[c * 256 + tid]) + if count != Uint32(0): + cluster_red(target_remote + tid * 4, count) + cute.arch.cluster_arrive() + cute.arch.cluster_wait() + if tid < 256: + hist[tid] = cluster_load(target_remote + tid * 4) + cute.arch.sync_threads() + if tid < 32: + select_bin(hist, remain, ctl, lane) + cute.arch.sync_threads() + if cutlass.const_expr(level == 0): + high = Uint32(ctl[BIN]) + else: + high = (high << Uint32(8)) | Uint32(ctl[BIN]) + remain = Uint32(ctl[REMAIN]) + cute.arch.sync_threads() + return high, remain + + @cute.jit + def direct( + self, + regs, + data_address: Int32, + start: Int32, + length: Int32, + pinned: Int32, + tid: Int32, + out_row: Int64, + ctl, + hist, + hist_a, + hist_b, + scan, + queue, + ) -> None: + """Direct path: the exact cutoff, then every CTA writes its keys + above it and its share of the ``need`` equal ones to the output.""" + k, per = self.k, self.per_vector + tau, need = self.exact_cutoff( + regs, data_address, start, length, pinned, tid, ctl, hist, hist_a, hist_b + ) + queued = self.queue_vectors( + regs, data_address, tau, start, length, pinned, tid, scan, queue + ) + gt, eq = self.count_queued( + data_address, queued, tau, start, length, pinned, tid, queue + ) + slot_gt, slot_eq = self.reserve(gt, eq, TOTALS + 2, tid, ctl, scan) + above = Int32(k) - Int32(need) + for q in cutlass.range(tid, queued, self.threads, unroll=1): + vector = Int32(queue[q]) + words = self.vector_words(data_address, vector) + for e in cutlass.range_constexpr(per): + element = vector * per + e + key = self.word_key(words, e, start + element, pinned) + index = start + element + if cutlass.const_expr(self.block_mode): # noqa: SIM102 + if key == Uint32(NEG_INF_KEY): + index = Int32(-1) + if element < length: + if key > tau: + if slot_gt < above: + store_s32(out_row + Int64(slot_gt) * 4, index) + slot_gt = slot_gt + 1 + if key == tau: + if slot_eq < Int32(need): + store_s32(out_row + Int64(above + slot_eq) * 4, index) + slot_eq = slot_eq + 1 + + @cute.jit + def emit( + self, + total_gt: Uint32, + cutoff: Uint32, + out_row: Int64, + tid: Int32, + ctl, + hist, + hist2, + scan, + g_key, + g_idx, + g_eq, + ) -> None: + """CTA 0, sampled path: write the row's k indices from the gathered + survivors.""" + k, threads = self.k, self.threads + count = Int32(total_gt) + if total_gt >= Uint32(k): + tau, need = self.radix_select( + g_key, count, Uint32(k), hist, hist2, ctl, tid + ) + above_n = Uint32(0) + tie_n = Uint32(0) + for i in cutlass.range(tid, count, threads, unroll=1): + key = Uint32(g_key[i]) + if key > tau: + above_n = above_n + Uint32(1) + if key == tau: + tie_n = tie_n + Uint32(1) + slot_above, slot_tie, _, _ = self.scan2(above_n, tie_n, scan, tid) + above = Int32(k) - Int32(need) + for i in cutlass.range(tid, count, threads, unroll=1): + key = Uint32(g_key[i]) + index = g_idx[i] + if key > tau: + store_s32(out_row + Int64(Int32(slot_above)) * 4, index) + slot_above = slot_above + Uint32(1) + if key == tau: + if Int32(slot_tie) < Int32(need): + store_s32(out_row + Int64(above + Int32(slot_tie)) * 4, index) + slot_tie = slot_tie + Uint32(1) + else: + for i in cutlass.range(tid, count, threads, unroll=1): + store_s32(out_row + Int64(i) * 4, g_idx[i]) + for i in cutlass.range(tid, Int32(k) - count, threads, unroll=1): + store_s32(out_row + Int64(count + i) * 4, g_eq[i]) + + @cute.kernel + def kernel( + self, + lengths: cute.Tensor, + out: cute.Tensor, + logits_address: Int64, + row_bytes: Int64, + width: Int32, + out_stride: Int32, + fallback: Int32, + ): + """One cluster a row (module docstring).""" + k, cluster, threads, per = self.k, self.cluster, self.threads, self.per_vector + samples = self.samples + + @cute.struct + class Shared: + data: cute.struct.Align[cute.struct.MemRange[Uint32, self.vectors * 4], 16] + ctl: cute.struct.Align[cute.struct.MemRange[Uint32, CTL_WORDS], 16] + scan: cute.struct.MemRange[Uint32, SCAN_WORDS] + hist: cute.struct.Align[cute.struct.MemRange[Uint32, 256 * COPIES], 16] + hist2: cute.struct.Align[cute.struct.MemRange[Uint32, 256], 16] + hist_a: cute.struct.MemRange[Uint32, 256] + hist_b: cute.struct.MemRange[Uint32, 256] + g_idx: cute.struct.MemRange[Int32, CAPACITY] + g_eq: cute.struct.MemRange[Int32, k] + queue: cute.struct.MemRange[Uint16, self.vectors] + sample: cute.struct.MemRange[Uint16, SAMPLES] + g_key: cute.struct.MemRange[Uint16, CAPACITY] + + shared = cutlass.utils.SmemAllocator().allocate(Shared) + + def view(member, size): + return cute.make_tensor(member.data_ptr(), cute.make_layout(size)) + + data = view(shared.data, self.vectors * 4) + data_address = Int32(data.iterator.toint()) + ctl = view(shared.ctl, CTL_WORDS) + scan = view(shared.scan, SCAN_WORDS) + hist = view(shared.hist, 256 * COPIES) + hist2 = view(shared.hist2, 256) + hist_a = view(shared.hist_a, 256) + hist_b = view(shared.hist_b, 256) + g_idx = view(shared.g_idx, CAPACITY) + g_eq = view(shared.g_eq, k) + queue = view(shared.queue, self.vectors) + sample = view(shared.sample, SAMPLES) + g_key = view(shared.g_key, CAPACITY) + + tid, _, _ = cute.arch.thread_idx() + cta, _, _ = cute.arch.block_idx() + rank = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + row = cta // cluster + if tid < CTL_WORDS: + ctl[tid] = Uint32(0) + if tid < 256: + hist_a[tid] = Uint32(0) + hist_b[tid] = Uint32(0) + # Every CTA has started, and CTA 0's counters are zero, before any + # distributed shared memory access (the wait comes after the loads). + cute.arch.cluster_arrive() + cute.arch.griddepcontrol_wait() + tokens = cutlass.max(cutlass.min(lengths[row], width), Int32(0)) + n = tokens + if cutlass.const_expr(self.block_mode): + n = (tokens + 7) // 8 + pinned = n - 1 # block mode: the newest block + base = logits_address + Int64(row) * row_bytes + out_row = Int64(out.iterator.toint()) + Int64(row) * Int64(out_stride) * 4 + if n <= k: + cute.arch.cluster_wait() + if rank == 0: + for slot in cutlass.range(tid, k, threads, unroll=1): + value = Int32(-1) + if slot < n: + value = slot + if cutlass.const_expr(self.block_mode): + if slot == pinned: + value = slot + else: + w = load_u32x4(base + Int64(slot) * 16) + if block_key(w[0], w[1], w[2], w[3]) == Uint32( + NEG_INF_KEY + ): + value = Int32(-1) + store_s32(out_row + Int64(slot) * 4, value) + else: + # Chunks start on a 16-byte vector. + align = 1 if self.block_mode else 8 + chunk = ((n + cluster - 1) // cluster + align - 1) // align * align + start = rank * chunk + length = cutlass.max(cutlass.min(n - start, chunk), Int32(0)) + element_bytes = 16 if self.block_mode else 2 + regs = cute.make_rmem_tensor((VECTORS, 4), Uint32) + for j in cutlass.range_constexpr(VECTORS): + first = (j * threads + tid) * per + for w in cutlass.range_constexpr(4): + regs[j, w] = Uint32(0) + if first < length: + words = load_u32x4(base + Int64(start + first) * element_bytes) + for w in cutlass.range_constexpr(4): + regs[j, w] = words[w] + if cutlass.const_expr(self.vector_count > VECTORS): + # A wide row's other vectors go through registers to shared + # memory only. + more = cute.make_rmem_tensor((self.vector_count - VECTORS, 4), Uint32) + if VECTORS * threads * per < length: + for b in cutlass.range_constexpr(self.vector_count - VECTORS): + first = ((VECTORS + b) * threads + tid) * per + if first < length: + words = load_u32x4( + base + Int64(start + first) * element_bytes + ) + for w in cutlass.range_constexpr(4): + more[b, w] = words[w] + for b in cutlass.range_constexpr(self.vector_count - VECTORS): + vector = (VECTORS + b) * threads + tid + if vector * per < length: + shared_store_u32x4( + data_address + vector * 16, + more[b, 0], + more[b, 1], + more[b, 2], + more[b, 3], + ) + for j in cutlass.range_constexpr(VECTORS): + if (j * threads + tid) * per < length: + shared_store_u32x4( + data_address + (j * threads + tid) * 16, + regs[j, 0], + regs[j, 1], + regs[j, 2], + regs[j, 3], + ) + cute.arch.sync_threads() + cute.arch.griddepcontrol_launch_dependents() + use_direct = Int32(1) + cutoff = Uint32(0) + total_gt = Uint32(0) + if cutlass.const_expr(self.block_mode): + cute.arch.cluster_wait() + else: + if (n > self.direct_elements) & (fallback == 0): + use_direct = Int32(0) + # Samples: CTA c's valid ones spread over its chunk, at the + # same density in every CTA; the rest are key 0. + key = Uint32(0) + valid_samples = (samples * length) // chunk + if (use_direct == 0) & (tid < valid_samples): + offset = ((2 * tid + 1) * length) // (2 * valid_samples) + key = order_key( + Uint32(data[offset >> 1]) >> (Uint32(offset & 1) * Uint32(16)) + ) + cute.arch.cluster_wait() + if use_direct == 0: + if tid < samples: + address = ( + Int32(sample.iterator.toint()) + (rank * samples + tid) * 2 + ) + for peer in cutlass.range_constexpr(cluster): + cluster_store_u16(peer_address(address, peer), key) + cute.arch.cluster_arrive() + cute.arch.cluster_wait() + total_samples = Int32(0) + for c in cutlass.range_constexpr(cluster): + c_length = cutlass.max( + cutlass.min(n - c * chunk, chunk), Int32(0) + ) + total_samples = total_samples + (samples * c_length) // chunk + q = (SURVIVOR_FACTOR * k * samples + chunk - 1) // chunk + q = cutlass.max(cutlass.min(q, total_samples), Int32(1)) + cutoff, _ = self.radix_select( + sample, SAMPLES, Uint32(q), hist, hist2, ctl, tid + ) + total_gt, total_eq = self.gather( + regs, + data_address, + cutoff, + start, + length, + pinned, + tid, + ctl, + scan, + queue, + g_key, + g_idx, + g_eq, + ) + if (total_gt > Uint32(CAPACITY)) | ( + total_gt + total_eq < Uint32(k) + ): + use_direct = Int32(1) + if use_direct != 0: + self.direct( + regs, + data_address, + start, + length, + pinned, + tid, + out_row, + ctl, + hist, + hist_a, + hist_b, + scan, + queue, + ) + # CTA 0 exits only after every peer is done with its shared memory. + cute.arch.cluster_arrive() + if rank == 0: + if cutlass.const_expr(not self.block_mode): # noqa: SIM102 + if use_direct == 0: + self.emit( + total_gt, + cutoff, + out_row, + tid, + ctl, + hist, + hist2, + scan, + g_key, + g_idx, + g_eq, + ) + cute.arch.cluster_wait() + + +def _pointer(dtype: type[cutlass.Numeric], tensor: torch.Tensor, align: int): + """A CuTe DSL global-memory pointer to ``tensor``'s data.""" + return make_ptr( + dtype, tensor.data_ptr(), cute.AddressSpace.gmem, assumed_align=align + ) + + +def _fake(dtype: type[cutlass.Numeric], align: int): + """A placeholder pointer for compilation.""" + return make_ptr(dtype, align, cute.AddressSpace.gmem, assumed_align=align) + + +@cache +def _kernel( + device_index: int, k: int, cluster: int, vector_count: int, block_mode: bool +): + """``DecodeTopK`` compiled for a device, k, cluster size, staged vectors + a thread and mode.""" + with torch.accelerator.device_index(device_index): + return cute.compile( + DecodeTopK(k, cluster, vector_count, block_mode), + _fake(Uint32, 16), + _fake(Int32, 4), + _fake(Int32, 4), + Int32(1), + Int64(16), + Int32(8), + Int32(k), + Int32(0), + make_fake_stream(), + ) + + +def threads_for(cluster: int) -> int: + """Threads a CTA: 1024 at 4 CTAs a row, else 512.""" + return 1024 if cluster == 4 else 512 + + +def vector_choices(cluster: int) -> tuple[int, ...]: + """The 16-byte vectors a thread may stage: 8, or 16 at 512 threads (1024 + threads have the shared memory for 8).""" + return (8,) if cluster == 4 else (8, 16) + + +def vectors_for(width: int, cluster: int) -> int | None: + """The fewest staged vectors a thread whose cluster holds a row of + ``width`` logits (a 16-byte vector: 8 logits, or one block), or None.""" + for count in vector_choices(cluster): + if width // 8 <= cluster * threads_for(cluster) * count: + return count + return None + + +@cache +def _resident_clusters(device_index: int, cluster: int) -> int: + """Clusters of ``cluster`` CTAs the device runs at once, one CTA an SM + (the kernel's shared memory allows no more): 36 of 4, 15 of 8 and 7 of 16 + on a 152-SM GB300.""" + with torch.accelerator.device_index(device_index): + return cutlass.utils.HardwareInfo(device_index).get_max_active_clusters(cluster) + + +def cluster_for(rows: int, width: int, device_index: int) -> int: + """CTAs a row: 8 while every row's cluster is resident at once, else 4. + Rows wider than 4 CTAs hold (the ratio-2 layers' at a 1M context) always + take 8, and rows wider than 8 CTAs hold (the candidate source's) 16, in + waves when they outnumber the resident clusters.""" + if vectors_for(width, 8) is None: + return 16 + if rows <= _resident_clusters(device_index, 8) or vectors_for(width, 4) is None: + return 8 + return 4 + + +def precompile(device_index: int, width: int, k: int, block_mode: bool) -> None: + """Compile the ``DecodeTopK`` variants `decode_topk` launches on rows of + ``width`` logits for ``k`` and ``block_mode``, at any row count, and read + the cluster occupancy it plans with, before graph capture.""" + if k not in TOP_KS: + return + if vectors_for(width, 8) is not None: + clusters = (8,) if vectors_for(width, 4) is None else (8, 4) + elif vectors_for(width, 16) is not None: + clusters = (16,) + else: + return + _resident_clusters(device_index, clusters[0]) + for cluster in clusters: + vector_count = vectors_for(width, cluster) + assert vector_count is not None + _kernel(device_index, k, cluster, vector_count, block_mode) + + +def supports_decode_topk( + logits: torch.Tensor, rows: int, k: int, block_mode: bool +) -> bool: + """Whether `decode_topk` takes these BF16 logits (16-byte aligned rows + whose width, a multiple of 8, fits one cluster of at most 16 CTAs) for + ``rows`` rows. Rows that need 16 CTAs are taken for at most two waves of + clusters: beyond, the waves cost more than the upcast and the stock + top-k.""" + if logits.dtype != torch.bfloat16 or logits.dim() != 2 or logits.stride(1) != 1: + return False + width = logits.shape[1] + if vectors_for(width, 8) is None and ( + vectors_for(width, 16) is None + or rows > 2 * _resident_clusters(logits.device.index, 16) + ): + return False + return ( + 0 < rows <= 64 + and k in TOP_KS + and width % 8 == 0 + and logits.stride(0) % 8 == 0 + and logits.data_ptr() % 16 == 0 + ) + + +def decode_topk( + logits: torch.Tensor, + lengths: torch.Tensor, + out: torch.Tensor, + block_mode: bool, + fallback: bool = False, +) -> None: + """Launch `DecodeTopK` on the current stream. + + Args: + logits: [rows, width] BF16, row ``r`` valid below ``lengths[r]`` + (see `supports_decode_topk`). + lengths: Contiguous [rows] int32 valid lengths in logits columns. + out: [rows, k] int32 (unit column stride) for the indices: tokens, or + 8-token blocks in block mode (-1 for blocks scoring -inf). + block_mode: Select 8-token blocks (the candidate source) instead of + tokens. + fallback: Force the direct path (tests). + + """ + rows, width = logits.shape + k = out.shape[1] + assert supports_decode_topk(logits, rows, k, block_mode) + assert out.dtype == torch.int32 and out.stride(1) == 1 and out.shape[0] == rows + assert ( + lengths.dtype == torch.int32 + and lengths.is_contiguous() + and lengths.numel() == rows + ) + cluster = cluster_for(rows, width, logits.device.index) + vector_count = vectors_for(width, cluster) + _kernel(logits.device.index, k, cluster, vector_count, block_mode)( + _pointer(Uint32, logits, 16), + _pointer(Int32, lengths, 4), + _pointer(Int32, out, 4), + Int32(rows), + Int64(logits.stride(0) * 2), + Int32(width), + Int32(out.stride(0)), + Int32(1 if fallback else 0), + current_cu_stream(), + ) diff --git a/vllm/models/deepseek_v41/nvidia/specialized_decode/attention/indexer.py b/vllm/models/deepseek_v41/nvidia/specialized_decode/attention/indexer.py index afb5b337904e..691970438432 100644 --- a/vllm/models/deepseek_v41/nvidia/specialized_decode/attention/indexer.py +++ b/vllm/models/deepseek_v41/nvidia/specialized_decode/attention/indexer.py @@ -5,16 +5,30 @@ ``decode_rows`` runs a dense indexer (``SparseAttnIndexer``) on a copy of the varlen decode branch of ``sparse_attn_indexer``, with two additions: -- Row sharding. The indexer is replicated across the TP ranks. In a step of at - least ``ROW_SHARD_MIN_REQUESTS`` requests whose count the ranks divide, each - rank scores, selects candidates for and top-k's only its contiguous slice of - the rows; TP all-gathers of the int32 rows rebuild ``topk_indices_buffer`` - and, in the candidate source, the candidate blocks. -- The candidate source selects its candidate blocks with ``cooperative_topk`` - for up to 64 rows per rank (see ``candidate_topk``). +- Row sharding. The indexer is replicated across the TP ranks. In a uniform + decode step of at least ``ROW_SHARD_MIN_ROWS`` rows, each rank scores, + selects candidates for and top-k's only its contiguous slice of the rows + (an even number of them, so the last ranks' slices may be short or empty); + TP all-gathers of the int32 rows (see ``_all_gather_rows``; the executor's + ``LamportAllGather``) rebuild ``topk_indices_buffer`` and, in the candidate + source, the candidate blocks. + A slice may split a request: varlen rows are independent. +- Selection from the BF16 logits. For up to 64 rows per rank, ``decode_topk`` + (``decode_topk_cute``) selects the token top-k, and in the candidate source + the candidate blocks, from SM100's BF16 logits without upcasting them. The + candidate source's logits are twice as wide (compress ratio 1), so at a 1M + context they need 16 CTAs a row and take that path only up to two waves of + such clusters (14 rows on a GB300); otherwise the candidate source selects + with ``cooperative_topk`` (see ``candidate_topk``) on the upcast logits. Steps it does not take run the stock indexer. +``decode_rows`` runs in two parts. ``plan_decode_rows`` takes the gates, +resets the top-k rows and builds the slice's DeepGEMM schedule: none of it +reads the layer's index query or keys, so the attention step issues it while +the current stream would otherwise wait for them. ``finish_decode_rows`` +scores, selects and gathers once they are ready. + ``CandidateOverlap`` moves the candidate consumers' candidate expansion and sparse MQA-logits schedule onto a side stream. @@ -22,6 +36,9 @@ that every TP rank shares: the step's row and request counts. """ +from dataclasses import dataclass +from typing import TYPE_CHECKING + import torch from torch import nn @@ -41,7 +58,6 @@ kv_cache_as_quant_view, ) from vllm.model_executor.layers.sparse_mqa_indexer import SparseMQAIndexer -from vllm.models.common.ops.sequence_parallel import sp_all_gather from vllm.models.deepseek_v41.attention import DeepseekV4Indexer from vllm.models.deepseek_v41.common.ops import fused_indexer_q_rope_quant from vllm.triton_utils import triton @@ -50,6 +66,7 @@ get_paged_mqa_logits_metadata, get_paged_sparse_mqa_logits_metadata, mqa_logits_mx_inputs, + mqa_logits_use_mx_scales, ) from vllm.v1.attention.backends.mla.indexer import ( DeepSeekV32IndexerDecodeMetadata, @@ -64,35 +81,45 @@ select_candidate_blocks_cooperative, supports_cooperative_candidates, ) +from .decode_topk_cute import decode_topk, supports_decode_topk + +if TYPE_CHECKING: + from ...ops.cute_dsl import LamportAllGather + +# Fewest rows of a row-sharded decode step (two requests of six rows): a +# single request's slices save less than the all-gathers cost. +ROW_SHARD_MIN_ROWS = 12 + + +def row_shard_rows(num_rows: int, num_requests: int, world_size: int) -> int | None: + """Rows each rank's slice spans in a pure-decode step, or None when the + step stays replicated. Rank ``r`` scores rows ``[r n, (r + 1) n)`` + clipped to the step's rows. -# Fewest requests of a row-sharded decode step. -ROW_SHARD_MIN_REQUESTS = 8 -# Most rows of a step whose candidate expansion overlaps. -OVERLAP_MAX_ROWS = 16 + Slices start on an even row: DeepGEMM's paged MQA logits needs its int32 + ``indices`` 8-byte aligned. Varlen rows are independent, so a slice may + split a request. Steps whose rows are not a multiple of their requests + (some capture and profiling runs) stay replicated. + """ + if world_size < 2 or not num_requests or num_rows % num_requests: + return None + if num_rows < ROW_SHARD_MIN_ROWS: + return None + rows = -(-num_rows // world_size) + return rows + rows % 2 def row_shard_bounds( num_rows: int, num_requests: int, world_size: int, rank: int ) -> tuple[int, int] | None: - """Rows ``[start, end)`` that ``rank`` scores in a pure-decode step, or - None when the step stays replicated. - - Each slice holds whole requests of a uniform step; varlen rows are - independent, so the split is exact either way. Slices start on an even - row: DeepGEMM's paged MQA logits needs its int32 ``indices`` 8-byte - aligned. DSpark decode steps always split so; capture and profiling runs - of one row a request may not, and stay replicated. - """ - rows = num_rows // world_size - if ( - world_size < 2 - or num_requests < ROW_SHARD_MIN_REQUESTS - or num_requests % world_size - or num_rows % num_requests - or rows % 2 - ): + """Rows ``[start, end)`` that ``rank`` scores in a pure-decode step (an + empty range for a rank past the last row), or None when the step stays + replicated.""" + rows = row_shard_rows(num_rows, num_requests, world_size) + if rows is None: return None - return rank * rows, (rank + 1) * rows + start = min(rank * rows, num_rows) + return start, min(start + rows, num_rows) def rope_quant_q( @@ -114,15 +141,55 @@ def rope_quant_q( use_fp4=indexer.use_fp4_kv, weights_out_dtype=indexer.indexer_weights_dtype, ) + if mqa_logits_use_mx_scales(): + # The SM100 logits kernels take BF16 weights. Cast them here, on the + # query's side stream, rather than in the scoring path + # (``mqa_logits_mx_inputs``); a no-op for BF16 weights. + weights = weights.to(torch.bfloat16) if isinstance(q_quant, tuple): return q_quant[0], q_quant[1], weights return q_quant, None, weights -def _all_gather_rows(local: torch.Tensor, out: torch.Tensor) -> None: - # The custom all-gather takes floating-point rows and moves the int32 - # bits unchanged: its Lamport variant rewrites only the INT32_MIN pattern. - out.copy_(sp_all_gather(local.view(torch.float32)).view(torch.int32)) +def _all_gather_rows( + buffer: torch.Tensor, + rows: int, + num_rows: int, + all_gather: "LamportAllGather | None", +) -> None: + """Gather every rank's ``rows``-row slice of ``buffer`` into its first + ``num_rows`` rows; rows past ``num_rows`` (a short or empty last slice) + are dropped. The slices go through ``all_gather`` when it takes them, + else NCCL.""" + group = get_tp_group() + local = buffer[group.rank_in_group * rows : (group.rank_in_group + 1) * rows] + # Not the custom all-gather: its MNNVL Lamport kernel never waits for a + # rank's own multicast copy, so a late copy can be read as a peer's slot + # by a later gather of another size, and a step's gathers differ in size + # (the top-k rows, then the candidate blocks). The executor's + # LamportAllGather waits for every slot. The int32 rows hold no + # INT32_MIN, which it would turn into 0. + if all_gather is not None and all_gather.accepts(local): + gathered = all_gather(local) + else: + gathered = group.all_gather(local, 0) + buffer[:num_rows].copy_(gathered[:num_rows]) + + +@dataclass +class DecodeRowsPlan: + """The rows a rank scores in ``finish_decode_rows`` and their schedule.""" + + decode: DeepSeekV32IndexerDecodeMetadata + max_seq_len: int + num_rows: int + # Rows of each rank's slice, or None when the step stays replicated. + shard: int | None + start: int + end: int + fast_candidates: bool + # DeepGEMM's paged MQA-logits schedule of rows ``[start, end)``. + schedule: torch.Tensor def decode_rows( @@ -130,9 +197,19 @@ def decode_rows( hidden_states: torch.Tensor, q_quant: torch.Tensor | tuple[torch.Tensor, torch.Tensor], weights: torch.Tensor, + all_gather: "LamportAllGather | None" = None, ) -> torch.Tensor: """``op(hidden_states, q_quant, None, weights)`` for a pure-decode step, - row-sharded and with cooperative candidates where they apply.""" + row-sharded and with cooperative candidates where they apply; the row + slices are gathered with ``all_gather`` where it takes them.""" + plan = plan_decode_rows(op, hidden_states.shape[0]) + return finish_decode_rows(op, hidden_states, q_quant, weights, plan, all_gather) + + +def plan_decode_rows(op: SparseAttnIndexer, num_tokens: int) -> DecodeRowsPlan | None: + """The first part of ``decode_rows`` for a step of ``num_tokens`` tokens: + reset the top-k rows and build the slice's schedule. None when the stock + indexer runs instead (and nothing was issued).""" attn_metadata = get_forward_context().attn_metadata metadata = ( attn_metadata.get(op.k_cache.prefix) @@ -148,11 +225,22 @@ def decode_rows( or metadata.decode.global_seq_lens is not None or not op.skip_k_cache_insert ): - return op(hidden_states, q_quant, None, weights) + return None num_rows = metadata.num_decode_tokens tp_group = get_tp_group() - bounds = row_shard_bounds( - num_rows, metadata.num_decodes, tp_group.world_size, tp_group.rank_in_group + shard = row_shard_rows(num_rows, metadata.num_decodes, tp_group.world_size) + if shard is not None and any( + buffer is not None and buffer.shape[0] < tp_group.world_size * shard + for buffer in (op.topk_indices_buffer, op.candidate_blocks) + ): + # The rows past the step's that the slices span do not fit. + shard = None + bounds = ( + None + if shard is None + else row_shard_bounds( + num_rows, metadata.num_decodes, tp_group.world_size, tp_group.rank_in_group + ) ) start, end = bounds or (0, num_rows) fast_candidates = ( @@ -165,33 +253,62 @@ def decode_rows( ) ) if bounds is None and not fast_candidates: - return op(hidden_states, q_quant, None, weights) + return None - q_values, q_scale = q_quant if isinstance(q_quant, tuple) else (q_quant, None) - topk_indices_buffer = op.topk_indices_buffer - topk_indices_buffer[: hidden_states.shape[0]] = -1 - _decode_rows( - op, - q_values, - q_scale, - weights, - metadata.decode, + decode = metadata.decode + assert decode.indices is not None + op.topk_indices_buffer[:num_tokens] = -1 + schedule = decode.schedule_metadata + if end > start and end - start < decode.seq_lens.shape[0]: + # The slice's own schedule; a FULL graph records it. + kv_cache = kv_cache_as_quant_view( + op.k_cache.kv_cache, op.head_dim, op.use_fp4_cache + ) + schedule = get_paged_mqa_logits_metadata( + decode.seq_lens[start:end], + kv_cache.shape[1], + schedule.shape[0] - 1, + indices=decode.indices[start:end], + ) + return DecodeRowsPlan( + decode, metadata.max_seq_len, + num_rows, + shard, start, end, fast_candidates, + schedule, ) - if bounds is not None: + + +def finish_decode_rows( + op: SparseAttnIndexer, + hidden_states: torch.Tensor, + q_quant: torch.Tensor | tuple[torch.Tensor, torch.Tensor], + weights: torch.Tensor, + plan: DecodeRowsPlan | None, + all_gather: "LamportAllGather | None" = None, +) -> torch.Tensor: + """The rest of ``decode_rows`` once the index query and keys are ready: + score and select the planned rows and gather the slices (with + ``all_gather`` where it takes them).""" + if plan is None: + return op(hidden_states, q_quant, None, weights) + q_values, q_scale = q_quant if isinstance(q_quant, tuple) else (q_quant, None) + if plan.end > plan.start: + _decode_rows(op, q_values, q_scale, weights, plan) + if plan.shard is not None: _all_gather_rows( - topk_indices_buffer[start:end, : op.topk_tokens], - topk_indices_buffer[:num_rows, : op.topk_tokens], + op.topk_indices_buffer[:, : op.topk_tokens], + plan.shard, + plan.num_rows, + all_gather, ) if op.candidate_write: assert op.candidate_blocks is not None - _all_gather_rows( - op.candidate_blocks[start:end], op.candidate_blocks[:num_rows] - ) - return topk_indices_buffer + _all_gather_rows(op.candidate_blocks, plan.shard, plan.num_rows, all_gather) + return op.topk_indices_buffer def _decode_rows( @@ -199,27 +316,18 @@ def _decode_rows( q: torch.Tensor, q_scale: torch.Tensor | None, weights: torch.Tensor, - decode: DeepSeekV32IndexerDecodeMetadata, - max_seq_len: int, - start: int, - end: int, - fast_candidates: bool, + plan: DecodeRowsPlan, ) -> None: - """The varlen decode branch of `sparse_attn_indexer` on rows + """The varlen decode branch of `sparse_attn_indexer` on the planned rows ``[start, end)``: paged MQA logits, candidate selection or mask, and top-k into the same rows of the shared buffers.""" + decode, start, end = plan.decode, plan.start, plan.end assert decode.indices is not None kv_cache = kv_cache_as_quant_view( op.k_cache.kv_cache, op.head_dim, op.use_fp4_cache ) seq_lens = decode.seq_lens[start:end] indices = decode.indices[start:end] - schedule = decode.schedule_metadata - if end - start < decode.seq_lens.shape[0]: - # The slice's own schedule; a FULL graph records it. - schedule = get_paged_mqa_logits_metadata( - seq_lens, kv_cache.shape[1], schedule.shape[0] - 1, indices=indices - ) q = q[start:end].unsqueeze(1) if op.use_fp4_cache: q = q.view(torch.int8) @@ -229,23 +337,28 @@ def _decode_rows( q, q_scale, weights[start:end], op.logits_dtype ) row_ends = seq_lens.view(-1) - logits = as_fp32_logits( - fp8_fp4_paged_mqa_logits( - (q, q_scale), - kv_cache, - decode_weights, - seq_lens, - decode.block_table[start:end], - schedule, - max_model_len=op.max_model_len, - logits_dtype=op.logits_dtype, - indices=indices, - ), - row_ends, + raw = fp8_fp4_paged_mqa_logits( + (q, q_scale), + kv_cache, + decode_weights, + seq_lens, + decode.block_table[start:end], + plan.schedule, + max_model_len=op.max_model_len, + logits_dtype=op.logits_dtype, + indices=indices, ) + topk_rows = op.topk_indices_buffer[start:end, : op.topk_tokens] + if _selects_from_bf16(op, raw, plan.fast_candidates): + # Straight from the BF16 logits: no upcast. + if op.candidate_blocks is not None: + decode_topk(raw, row_ends, op.candidate_blocks[start:end], block_mode=True) + decode_topk(raw, row_ends, topk_rows, block_mode=False) + return + logits = as_fp32_logits(raw, row_ends) if op.candidate_blocks is not None: candidates = op.candidate_blocks[start:end] - if fast_candidates: + if plan.fast_candidates: select_candidate_blocks_cooperative( logits, row_ends, op.candidate_block_size, candidates ) @@ -266,9 +379,31 @@ def _decode_rows( logits, seq_lens, 1, - op.topk_indices_buffer[start:end, : op.topk_tokens], + topk_rows, op.topk_tokens, - max_seq_len, + plan.max_seq_len, + ) + + +def _selects_from_bf16( + op: SparseAttnIndexer, logits: torch.Tensor, fast_candidates: bool +) -> bool: + """Whether ``decode_topk`` selects the token top-k, and the candidate + blocks if the layer writes them, from these logits.""" + rows = logits.shape[0] + if not ( + op.topk_backend == "auto" + and supports_decode_topk(logits, rows, op.topk_tokens, block_mode=False) + ): + return False + if op.candidate_blocks is None: + return True + return ( + fast_candidates + and op.candidate_block_size == 8 + and supports_decode_topk( + logits, rows, op.candidate_blocks.shape[1], block_mode=True + ) ) 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 18a79fb02758..9d7c34516913 100644 --- a/vllm/models/deepseek_v41/nvidia/specialized_decode/attention/mxfp8.py +++ b/vllm/models/deepseek_v41/nvidia/specialized_decode/attention/mxfp8.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Low-batch MXFP8 attention projections. +"""MXFP8 attention projections. At up to ``LOW_BATCH_MAX_ROWS`` rows the executor runs ``fused_wqa_wkv``, ``wq_b`` and ``wo_b`` on FlashInfer ``mm_mxfp8`` with its @@ -8,15 +8,31 @@ stock layers keep ``cute-dsl``. Both backends read the same weights, weight scales and F8_128x4 activation scales. -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. +The backend's tactics are close 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. + +Up to ``WIDE_LOW_LATENCY_MAX_ROWS`` rows the GB300 TP4 shapes keep the +low-latency kernel, which the backend's runner offers only with an 8-row +activation tile: this module compiles FlashInfer's kernel class itself with a +16-, 32- or 48-row tile and a fixed tactic per tile. The kernel loads the +weights before its programmatic grid wait and lets its dependents launch once +its weight loads are issued, so in the PDL chain of a decode step the weight +read overlaps the previous kernel; ``cute-dsl`` waits before any load. + +Above that, the projections stay on ``cute-dsl``, whose tuner picks a tactic +per power-of-two row bucket at every server start. Timed in isolation the +candidates are close, so the pick varies from start to start, and inside a +decode step one start's pick can be much slower than another's. The GB300 +TP4 shapes therefore run a fixed tactic per row range, chosen in the chain of +the three projections, and skip the tuner. """ import importlib from collections.abc import Callable, Iterable, Iterator from contextlib import contextmanager from functools import cache +from typing import TYPE_CHECKING import torch from torch import nn @@ -28,12 +44,18 @@ QuantizedActivation, as_quantized_activation, ) +from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( + mxfp8_e4m3_quantize, +) from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp8Dynamic from vllm.models.deepseek_v41.common.ops.query_quant import mxfp8_scale_bytes from vllm.platforms import current_platform from vllm.utils import flashinfer as vllm_flashinfer -from .. import LOW_BATCH_MAX_ROWS +from .. import LOW_BATCH_MAX_ROWS, WIDE_LOW_LATENCY_MAX_ROWS, uses_checked_flashinfer + +if TYPE_CHECKING: + from flashinfer.autotuner import TunableRunner BACKEND = "cutedsl_low_latency" @@ -47,6 +69,40 @@ } _PINNED_MIN_ROWS = LOW_BATCH_MAX_ROWS // 2 + 1 +# The low-latency tactic per (N, K) and activation tile (rows rounded up to +# 16) on GB300 above LOW_BATCH_MAX_ROWS rows, chosen jointly for the chain of +# the three projections. +WIDE_TACTICS: dict[tuple[int, int], dict[int, Tactic]] = { + (1792, 5120): {16: (128, 9, 1, 4), 32: (128, 8, 1, 4), 48: (128, 7, 1, 4)}, + (8192, 1280): {16: (256, 5, 1, 1), 32: (256, 5, 1, 1), 48: (128, 9, 1, 1)}, + (5120, 2048): {16: (128, 8, 1, 2), 32: (128, 10, 1, 1), 48: (128, 9, 1, 1)}, +} + +# ((tile_m, tile_n), (cluster_m, cluster_n), swap_ab, prefetch, split_k) of +# the cute-dsl GEMM per (N, K) on GB300 above LOW_BATCH_MAX_ROWS rows, as +# (largest rows, tactic) in increasing rows. +CuteDslTactic = tuple[tuple[int, int], tuple[int, int], bool, bool, int] +CUTE_DSL_TACTICS: dict[tuple[int, int], tuple[tuple[int, CuteDslTactic], ...]] = { + (1792, 5120): ( + (12, ((128, 16), (1, 1), True, False, 2)), + (24, ((128, 8), (1, 1), True, False, 1)), + (48, ((256, 64), (2, 1), True, False, 1)), + (96, ((256, 64), (4, 1), True, False, 1)), + (384, ((256, 64), (2, 1), True, False, 1)), + ), + (8192, 1280): ( + (12, ((128, 8), (1, 1), True, False, 1)), + (96, ((128, 64), (1, 1), True, False, 1)), + (192, ((256, 128), (2, 1), False, False, 1)), + (384, ((256, 192), (2, 1), True, False, 1)), + ), + (5120, 2048): ( + (24, ((128, 16), (1, 1), True, False, 1)), + (192, ((256, 64), (2, 1), True, False, 1)), + (384, ((256, 128), (2, 1), True, False, 1)), + ), +} + def mxfp8_mm(layer: nn.Module, qa: QuantizedActivation) -> torch.Tensor: """``layer``'s GEMM on ``cutedsl_low_latency``.""" @@ -106,6 +162,236 @@ def low_latency_mm( return mxfp8_mm(layer, qa) +@cache +def _cute_dsl_runner(out_dtype: torch.dtype) -> "TunableRunner | None": + """``mm_mxfp8``'s ``cute-dsl`` runner (with PDL), when this device and + FlashInfer are the ones the pinned tactics were chosen on.""" + if not ( + current_platform.is_device_capability(103) + and uses_checked_flashinfer("pinning the MXFP8 GEMM tactics") + ): + return None + from flashinfer.gemm import gemm_base + + return gemm_base._cute_dsl_gemm_mxfp8_runner(10, 3, True, out_dtype) + + +# The pinned tactic per (N, K, rows); None where the runner cannot run it. +_pinned: dict[tuple[int, int, int], CuteDslTactic | None] = {} + + +def pinned_mm( + layer: nn.Module, x: torch.Tensor | QuantizedActivation +) -> torch.Tensor | None: + """``layer``'s ``cute-dsl`` GEMM of ``x`` with the pinned tactic of its + shape and row count, quantizing a BF16 ``x`` as the layer does; None when + no tactic is pinned for them.""" + kernel = _mxfp8_kernel(layer) + assert kernel is not None + k, n = layer.weight.shape # column-major [K, N] + ranges = CUTE_DSL_TACTICS.get((n, k)) + qa = as_quantized_activation(x, kernel.input_quant_key()) + if ranges is None or (qa is None and not isinstance(x, torch.Tensor)): + return None + if qa is not None: + a, out_dtype, shape = qa.data, qa.orig_dtype, qa.orig_shape + else: + assert isinstance(x, torch.Tensor) + a, out_dtype, shape = x.view(-1, k), x.dtype, x.shape + runner = _cute_dsl_runner(out_dtype) + if runner is None: + return None + + from flashinfer.gemm import gemm_base + + rows = a.shape[0] + key = (n, k, rows) + if key not in _pinned: + tactic = next((t for max_rows, t in ranges if rows <= max_rows), None) + # The runner reads only the operands' shapes and the output's layout. + out = torch.empty(rows, n, dtype=out_dtype, device=a.device) + operands = [a, layer.weight, None, None, out_dtype, out, None] + valid = runner.get_valid_tactics(operands, None) # type: ignore[arg-type] + runs = rows > LOW_BATCH_MAX_ROWS and tactic in valid + _pinned[key] = tactic if runs else None + tactic = _pinned[key] + if tactic is None: + return None + if qa is not None: + data, scale = qa.data, qa.scale + else: + data, scale = mxfp8_e4m3_quantize(a, is_sf_swizzled_layout=True) + out = torch.empty(rows, n, dtype=out_dtype, device=a.device) + workspace = gemm_base._get_cache_buf( + "mm_mxfp8_workspace", gemm_base.DEFAULT_WORKSPACE_SIZE, a.device + ) + inputs = [data, layer.weight, scale, layer.weight_scale, out_dtype, out, workspace] + runner(inputs=inputs, tactic=tactic) + return out.view(*shape[:-1], n) + + +@cache +def _wide_supported() -> bool: + """Whether this device and FlashInfer are the ones the wide tactics were + chosen on and the low-latency kernel class was checked against.""" + return current_platform.is_device_capability(103) and uses_checked_flashinfer( + "the low-latency MXFP8 GEMM above 8 rows" + ) + + +def wide_tactic(n: int, k: int, rows: int) -> tuple[int, Tactic] | None: + """(activation tile, tactic) of the wide low-latency GEMM of an (N, K) + weight at ``rows`` rows; None where it does not run.""" + if not LOW_BATCH_MAX_ROWS < rows <= WIDE_LOW_LATENCY_MAX_ROWS: + return None + tile = (rows + 15) // 16 * 16 + tactic = WIDE_TACTICS.get((n, k), {}).get(tile) + if tactic is None or not _wide_supported(): + return None + return tile, tactic + + +def _gmem_ptr(dtype, address: int, align: int): + import cutlass.cute as cute + from cutlass.cute.runtime import make_ptr + + return make_ptr(dtype, address, cute.AddressSpace.gmem, assumed_align=align) + + +def compile_low_latency(gemm, tile: int, cta_k: int, device: torch.device): + """``gemm``, a FlashInfer ``LowLatencyBlockscaledGemmKernel`` with a + ``tile``-row activation tile, compiled for MXFP8 operands and a BF16 + output without scale or bias, as FlashInfer's runner compiles it.""" + import cutlass + import cutlass.cute as cute + from cutlass.cute.runtime import make_fake_stream + + f8, sf, bf16 = cutlass.Float8E4M3FN, cutlass.Float8E8M0FNU, cutlass.BFloat16 + with torch.accelerator.device_index(device.index): + return cute.compile( + gemm, + _gmem_ptr(f8, 16, 16), + _gmem_ptr(sf, 32, 32), + _gmem_ptr(f8, 16, 16), + _gmem_ptr(sf, 16, 16), + _gmem_ptr(bf16, 16, 16), + None, + None, + tuple(cutlass.Int32(v) for v in (128, tile, cta_k, 1)), + make_fake_stream(), + ) + + +def launch_low_latency( + compiled, + weight: torch.Tensor, + weight_scale: torch.Tensor, + a: torch.Tensor, + a_scale: torch.Tensor, + out: torch.Tensor, +) -> None: + """Launch a compiled low-latency GEMM on the current stream: ``out`` = + ``a`` (rows x K, F8_128x4 scales) times the column-major [K, N] + ``weight``. The kernel computes ``out.T``: its M is the weight's N.""" + import cutlass + + from vllm.models.common.ops.moe_tail_fusion.primitives import current_cu_stream + + k, n = weight.shape + f8, sf, bf16 = cutlass.Float8E4M3FN, cutlass.Float8E8M0FNU, cutlass.BFloat16 + compiled( + _gmem_ptr(f8, weight.data_ptr(), 16), + _gmem_ptr(sf, weight_scale.data_ptr(), 32), + _gmem_ptr(f8, a.data_ptr(), 16), + _gmem_ptr(sf, a_scale.data_ptr(), 16), + _gmem_ptr(bf16, out.data_ptr(), 16), + None, + None, + tuple(cutlass.Int32(v) for v in (n, a.shape[0], k, 1)), + current_cu_stream(), + ) + + +# The compiled wide low-latency GEMM per (tactic, tile, device index). +_wide_kernels: dict[tuple[Tactic, int, int | None], object] = {} + + +def _wide_kernel(tactic: Tactic, tile: int, device: torch.device): + """The low-latency GEMM of ``tactic`` with a ``tile``-row activation + tile, compiled once per device.""" + key = (tactic, tile, device.index) + compiled = _wide_kernels.get(key) + if compiled is None: + import cutlass + from flashinfer.gemm.kernels.cute_dsl.low_latency_blockscaled_gemm import ( + LowLatencyBlockscaledGemmKernel, + ) + + cta_k, ab_stages, sfb_tmem_stages, split_k = tactic + gemm = LowLatencyBlockscaledGemmKernel( + acc_dtype=cutlass.Float32, + mma_tiler_mnk=(128, tile, cta_k), + num_ab_stage=ab_stages, + num_sfb_tmem_stage=sfb_tmem_stages, + sf_vec_size=32, + use_pdl=True, + split_k=split_k, + ) + compiled = compile_low_latency(gemm, tile, cta_k, device) + _wide_kernels[key] = compiled + return compiled + + +def wide_mm( + layer: nn.Module, x: torch.Tensor | QuantizedActivation +) -> torch.Tensor | None: + """``layer``'s low-latency GEMM of ``x`` with a wide activation tile, + quantizing a BF16 ``x`` as the layer does; None where it does not run.""" + kernel = _mxfp8_kernel(layer) + assert kernel is not None + k, n = layer.weight.shape # column-major [K, N] + qa = as_quantized_activation(x, kernel.input_quant_key()) + if qa is not None: + if qa.scale.ndim != 1 or not qa.data.is_contiguous(): + return None + a, out_dtype, shape = qa.data, qa.orig_dtype, qa.orig_shape + elif isinstance(x, torch.Tensor): + a, out_dtype, shape = x.reshape(-1, k), x.dtype, x.shape + else: + return None + choice = wide_tactic(n, k, a.shape[0]) + if choice is None or out_dtype != torch.bfloat16: + return None + tile, tactic = choice + if qa is not None: + data, scale = qa.data, qa.scale + else: + data, scale = mxfp8_e4m3_quantize(a, is_sf_swizzled_layout=True) + out = torch.empty(a.shape[0], n, dtype=out_dtype, device=a.device) + compiled = _wide_kernel(tactic, tile, a.device) + launch_low_latency(compiled, layer.weight, layer.weight_scale, data, scale, out) + return out.view(*shape[:-1], n) + + +def wide_or_pinned_mm( + layer: nn.Module, x: torch.Tensor | QuantizedActivation +) -> torch.Tensor | None: + """An eligible ``layer``'s GEMM of ``x`` above ``LOW_BATCH_MAX_ROWS`` + rows: on the wide low-latency GEMM, else with the pinned ``cute-dsl`` + tactic; None where neither applies.""" + out = wide_mm(layer, x) + return out if out is not None else pinned_mm(layer, x) + + +def projection_mm( + layer: nn.Module, x: torch.Tensor | QuantizedActivation, low_batch: bool +) -> torch.Tensor | None: + """An eligible ``layer``'s GEMM of ``x``: on ``cutedsl_low_latency`` in a + step of up to ``LOW_BATCH_MAX_ROWS`` rows (``low_batch``), else + ``wide_or_pinned_mm``; None where none applies.""" + return low_latency_mm(layer, x) if low_batch else wide_or_pinned_mm(layer, x) + + def shape_representatives(linears: Iterable[nn.Module]) -> list[nn.Module]: """One eligible linear per distinct (N, K), in a rank-independent order.""" by_shape: dict[tuple[int, int], nn.Module] = {} @@ -116,19 +402,35 @@ def shape_representatives(linears: Iterable[nn.Module]) -> list[nn.Module]: return [by_shape[shape] for shape in sorted(by_shape)] -def run_zeros(linear: nn.Module, rows: int) -> torch.Tensor: - """Call ``linear`` on an input laid out like the fused quantize producers'.""" +def _zeros(linear: nn.Module, rows: int) -> QuantizedActivation: + """An input of ``linear`` laid out like the fused quantize producers'.""" k = linear.weight.shape[0] # column-major [K, N] device = linear.weight.device shape = torch.Size((rows, k)) - qa = QuantizedActivation( + return QuantizedActivation( torch.zeros(shape, dtype=torch.float8_e4m3fn, device=device), torch.zeros(mxfp8_scale_bytes(rows, k), dtype=torch.uint8, device=device), torch.bfloat16, shape, kMxfp8Dynamic, ) - return mxfp8_mm(linear, qa) + + +def run_zeros(linear: nn.Module, rows: int) -> torch.Tensor: + """Call ``linear`` on ``cutedsl_low_latency`` on a zero input.""" + return mxfp8_mm(linear, _zeros(linear, rows)) + + +def compile_pinned(linear: nn.Module) -> None: + """Compile ``linear``'s pinned ``cute-dsl`` tactics, each at the largest + row count of its range, and its wide low-latency GEMMs.""" + k, n = linear.weight.shape + for rows, _ in CUTE_DSL_TACTICS.get((n, k), ()): + pinned_mm(linear, _zeros(linear, rows)) + for tile in WIDE_TACTICS.get((n, k), {}): + choice = wide_tactic(n, k, tile) + if choice is not None: + _wide_kernel(choice[1], tile, linear.weight.device) @contextmanager diff --git a/vllm/models/deepseek_v41/nvidia/specialized_decode/attention/sparse_mla.py b/vllm/models/deepseek_v41/nvidia/specialized_decode/attention/sparse_mla.py index a4ab3db24888..e5180f0a76f7 100644 --- a/vllm/models/deepseek_v41/nvidia/specialized_decode/attention/sparse_mla.py +++ b/vllm/models/deepseek_v41/nvidia/specialized_decode/attention/sparse_mla.py @@ -16,6 +16,11 @@ kernel returns every counter to zero when it finishes, so a buffer zeroed once stays valid for launches that do not overlap, as in the FlashInfer MLA backend. + +Batches whose clusters fit the GPU at once (at width 640, up to 8 requests with +DSpark on a GB300) run the CuTe DSL kernel of ``sparse_mla_cute`` instead of +TRTLLM-gen: it spreads a token's columns over one CTA per 128 instead of at most +3. """ import inspect @@ -33,6 +38,11 @@ DeepseekV4FlashInferMLAAttention, get_flashinfer_dsv4_workspace, ) +from .sparse_mla_cute import ( + SWA_COLUMNS, + sparse_mla_decode, + supports_sparse_mla_cute, +) logger = init_logger(__name__) @@ -396,12 +406,41 @@ def run( output = output.unsqueeze(1) decode_cu = None max_decode_query_len = 1 + decode_indices = sparse_indices[:num_tokens] + # The CuTe kernel reads the first SWA_COLUMNS columns from the SWA pool. + if ( + attn.window_size == SWA_COLUMNS + and swa_metadata.decode_swa_width == SWA_COLUMNS + and supports_sparse_mla_cute( + query, + swa_k_cache, + compressed_kv_cache, + decode_indices, + decode_topk_lens, + output, + bmm1_scale, + bmm2_scale, + attn.attn_sink, + ) + ): + sparse_mla_decode( + query, + swa_k_cache, + compressed_kv_cache, + decode_indices, + decode_topk_lens, + output, + cast(float, bmm1_scale), + cast(float, bmm2_scale), + attn.attn_sink, + ) + return trtllm_decode_sparse_mla_dsv4( query=query, swa_kv_cache=swa_k_cache, compressed_kv_cache=compressed_kv_cache, workspace_buffer=get_flashinfer_dsv4_workspace(q.device), - sparse_indices=sparse_indices[:num_tokens], + sparse_indices=decode_indices, sparse_topk_lens=decode_topk_lens, seq_lens=decode_seq_lens, out=output, diff --git a/vllm/models/deepseek_v41/nvidia/specialized_decode/attention/sparse_mla_cute.py b/vllm/models/deepseek_v41/nvidia/specialized_decode/attention/sparse_mla_cute.py new file mode 100644 index 000000000000..95c7282dea05 --- /dev/null +++ b/vllm/models/deepseek_v41/nvidia/specialized_decode/attention/sparse_mla_cute.py @@ -0,0 +1,986 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The sparse-MLA decode of small batches, in CuTe DSL. + +The problem. In a decode step every query token (DeepSeek-V4's MLA after the +query projection: 16 heads a rank at TP4, FP8 E4M3, head dim 512) attends to +its own sparse set of KV rows. A row is 512 FP8 bytes and serves as both K and +V. The set holds up to 640 columns: the 128 sliding-window rows of one KV pool, +then the indexer's top-k rows of the compressed pool, all addressed as storage +offsets into their pool, with -1 for an unused column and a per-token count of +active columns. TRTLLM-gen splits a token's 640 rows over at most 3 CTAs, so at +one or two requests (6 or 12 tokens with DSpark) its kernel runs about 10 µs +on a GPU of 152 SMs; most of that is the latency of one CTA's serial work. + +What this module computes. ``sparse_mla_decode`` writes, for each token ``t`` +and head ``h``, the BF16 attention output over the token's active columns ``c`` +(``c < lens[t]`` and ``indices[t, c] >= 0``):: + + s_c = bmm1_scale * q[t, h] . kv[c] + p_c = exp(s_c - m) + out[t, h] = bmm2_scale * sum_c p_c kv[c] / (sum_c p_c + exp(sink_h - m)) + +where ``kv[c]`` is row ``indices[t, c]`` of the sliding-window pool for the +first 128 columns and of the compressed pool for the others, ``m`` is the +maximum of the scores and the sink, and the sink term is omitted without sinks. +A token with no active column gets zeros. This is what +``trtllm_batch_decode_sparse_mla_dsv4`` computes for these shapes. + +How. One CTA of 128 threads takes one token and one 128-column chunk of its +sparse set (chunk 0 = the sliding-window columns); the ``n = ceil(width / 128)`` +chunks of a token form a thread block cluster, so a 640-wide layer runs 5 CTAs +a token in parallel:: + + gather the chunk's KV rows and the query (cp.async, one group per 128-byte K + block) into 128B-swizzled tiles + └► S^T = KV x Q^T on tcgen05 (M 128 rows, N 16 heads, K 512, FP8 -> FP32 TMEM) + └► chunk softmax per head in base 2: max m_c, sum l_c, and + P^T = 256 * exp2(S - m_c) as FP8 (scaled: small weights stay normal) + └► O^T = V^T x P^T on tcgen05, V the same tile read MN-major + (four MMAs of 128 head dims) ... + └► ... while the cluster exchanges (m_c, l_c) per head (st.async) + └► each CTA scales its partial O by exp2(m_c - M) / L + (M, L over the cluster's chunks and the sink) + └► block j of the 512 head dims goes to CTA j % n, which + sums the n partials and stores BF16 + +The partial blocks travel by bulk copies between the CTAs' shared memory into +the receiving CTA's KV tile, whose MMAs are done by then: a cluster barrier +after each CTA's second MMA separates the two uses. A last cluster barrier keeps +every CTA resident until the bulk copies reading its shared memory are complete. + +Synchronization. The kernel reads the sinks (weights) before its grid +dependency wait and everything else after it, and lets its dependents launch +once its S MMAs are issued (programmatic dependent launch). Masked columns +gather zeros and score -inf. + +Scope. Two CTAs fit an SM (about 78 KB of shared memory each), so a launch is +one wave only while every token's cluster can be resident: ``_max_tokens`` allows +2 x the device's resident clusters of ``n`` CTAs at one CTA an SM (on a 152-SM +GB300, 8 requests with DSpark at width 640). +Larger batches keep TRTLLM-gen, which is faster once the clusters run in waves. +""" + +import math +from functools import cache + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import cutlass.utils as utils +import torch +from cutlass import Float32, Int32, Int64, Uint32, Uint64 +from cutlass._mlir.dialects import llvm +from cutlass.cute.runtime import make_fake_stream +from cutlass.cutlass_dsl import T, dsl_user_op + +from vllm.cute_utils import _tcgen05, recast_val +from vllm.cute_utils import mbarrier as mbar_ops +from vllm.models.common.ops.moe_tail_fusion.primitives import current_cu_stream + +HEADS = 16 +HEAD_DIM = 512 +ROWS = 128 # KV rows (columns of the sparse set) a CTA +SWA_COLUMNS = ROWS # chunk 0: the sliding-window columns, from the SWA pool +THREADS = 128 +MAX_CHUNKS = 5 # widths up to 640 +BLOCK_BYTES = ROWS * 128 # one 128-byte K block of the KV tile +Q_BLOCK_BYTES = HEADS * 128 # one 128-byte K block of the query tile +PART_BYTES = ROWS * HEADS * 4 # one 128-dim block of a partial O, FP32 +TMEM_COLS = 128 # S: 16 columns, O: 4 x 16 +P_SCALE = 256.0 +LOG2E = 1.4426950408889634 +CTAS_PER_SM = 2 +# tcgen05 instruction descriptors, kind::f8f6f4: D FP32, A and B E4M3, M = 128, +# N = 16; bit 15 makes A MN-major (V^T read from the K-major KV tile). +IDESC_QK = (1 << 4) | ((HEADS >> 3) << 17) | ((ROWS >> 4) << 24) +IDESC_PV = IDESC_QK | (1 << 15) + + +def _sdesc(lbo: int, sbo: int) -> int: + """A 128B-swizzled shared-memory matrix descriptor without its start + address: leading and stride byte offsets, version 1.""" + return ((lbo >> 4) << 16) | ((sbo >> 4) << 32) | (1 << 46) | (2 << 61) + + +SDESC_K = _sdesc(0, 1024) # K-major operands: 8-row groups 1024 bytes apart +SDESC_MN = _sdesc(BLOCK_BYTES, 1024) # V^T: 128-row K blocks a tile block apart + + +def _asm(result, operands, text, constraints, side_effects, loc, ip): + """One PTX inline assembly operation.""" + return llvm.inline_asm( + result, + operands, + text, + constraints, + has_side_effects=side_effects, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +def _v(value, loc, ip): + """IR value of a DSL scalar.""" + return value.ir_value(loc=loc, ip=ip) + + +@dsl_user_op +def cp_async_16(smem: Int32, gmem: Int64, *, loc=None, ip=None) -> None: + """Copy 16 bytes from global to this CTA's shared memory (cp.async).""" + _asm( + None, + [_v(Int32(smem), loc, ip), _v(Int64(gmem), loc, ip)], + "cp.async.cg.shared.global [$0], [$1], 16;", + "r,l", + True, + loc, + ip, + ) + + +@dsl_user_op +def st_shared_zero16(smem: Int32, *, loc=None, ip=None) -> None: + """Store 16 zero bytes at a 16-byte aligned shared address.""" + _asm( + None, + [_v(Int32(smem), loc, ip), _v(Int32(0), loc, ip)], + "st.shared.v4.b32 [$0], {$1, $1, $1, $1};", + "r,r", + True, + loc, + ip, + ) + + +@dsl_user_op +def st_shared_v4( + smem: Int32, w0: Uint32, w1: Uint32, w2: Uint32, w3: Uint32, *, loc=None, ip=None +) -> None: + """Store four 32-bit words at a 16-byte aligned shared address.""" + _asm( + None, + [_v(Int32(smem), loc, ip)] + [_v(Uint32(w), loc, ip) for w in (w0, w1, w2, w3)], + "st.shared.v4.b32 [$0], {$1, $2, $3, $4};", + "r,r,r,r,r", + True, + loc, + ip, + ) + + +@dsl_user_op +def st_shared_bf16(smem: Int32, value: Float32, *, loc=None, ip=None) -> None: + """Store ``value`` rounded to BF16 at a shared address.""" + _asm( + None, + [_v(Int32(smem), loc, ip), _v(Float32(value), loc, ip)], + "{\n\t.reg .b16 b;\n\tcvt.rn.bf16.f32 b, $1;\n\tst.shared.b16 [$0], b;\n\t}", + "r,f", + True, + loc, + ip, + ) + + +@dsl_user_op +def ld_global_s32(address: Int64, *, loc=None, ip=None) -> Int32: + """A 32-bit integer from global memory an earlier grid wrote (side-effecting, + so it stays after the grid dependency wait).""" + return Int32( + _asm( + T.i32(), + [_v(Int64(address), loc, ip)], + "ld.global.b32 $0, [$1];", + "=r,l", + True, + loc, + ip, + ) + ) + + +@dsl_user_op +def ld_global_f32(address: Int64, *, loc=None, ip=None) -> Float32: + """A float from read-only global memory.""" + return Float32( + _asm( + T.f32(), + [_v(Int64(address), loc, ip)], + "ld.global.nc.f32 $0, [$1];", + "=f,l", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def copy16_shared_to_global(smem: Int32, gmem: Int64, *, loc=None, ip=None) -> None: + """Copy 16 bytes from this CTA's shared memory to global memory.""" + _asm( + None, + [_v(Int32(smem), loc, ip), _v(Int64(gmem), loc, ip)], + "{\n\t.reg .b32 a, b, c, d;\n\tld.shared.v4.b32 {a, b, c, d}, [$0];\n\t" + "st.global.v4.b32 [$1], {a, b, c, d};\n\t}", + "r,l", + True, + loc, + ip, + ) + + +@dsl_user_op +def pack_e4m3x4( + a: Float32, b: Float32, c: Float32, d: Float32, *, loc=None, ip=None +) -> Uint32: + """Four floats as E4M3 bytes, ``a`` in the lowest byte (round to nearest, + saturating).""" + return Uint32( + _asm( + T.i32(), + [_v(Float32(v), loc, ip) for v in (a, b, c, d)], + "{\n\t.reg .b16 lo, hi;\n\tcvt.rn.satfinite.e4m3x2.f32 lo, $2, $1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 hi, $4, $3;\n\tmov.b32 $0, {lo, hi};\n\t}", + "=r,f,f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def mma_f8( + d_tmem: Int32, + adesc: Uint64, + bdesc: Uint64, + idesc: Uint32, + accumulate: Int32, + *, + loc=None, + ip=None, +) -> None: + """One tcgen05 FP8 MMA of this CTA: D (TMEM) = A x B (+ D if + ``accumulate``), A and B from shared memory descriptors.""" + _asm( + None, + [ + _v(Int32(d_tmem), loc, ip), + _v(Uint64(adesc), loc, ip), + _v(Uint64(bdesc), loc, ip), + _v(Uint32(idesc), loc, ip), + _v(Int32(accumulate), loc, ip), + ], + "{\n\t.reg .pred p;\n\tsetp.ne.b32 p, $4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f8f6f4 [$0], $1, $2, $3, p;\n\t}", + "r,l,l,r,r", + True, + loc, + ip, + ) + + +@dsl_user_op +def st_async_peer_b32( + smem: Int32, mbar: Int32, peer: Int32, value: Uint32, *, loc=None, ip=None +) -> None: + """Store a word at ``smem`` of cluster CTA ``peer`` (the same offset as a + local address), completing 4 bytes on the peer's ``mbar``.""" + _asm( + None, + [_v(Int32(v), loc, ip) for v in (smem, mbar, peer)] + + [_v(Uint32(value), loc, ip)], + "{\n\t.reg .b32 ra, rb;\n\tmapa.shared::cluster.u32 ra, $0, $2;\n\t" + "mapa.shared::cluster.u32 rb, $1, $2;\n\t" + "st.async.shared::cluster.mbarrier::complete_tx::bytes.b32 " + "[ra], $3, [rb];\n\t}", + "r,r,r,r", + True, + loc, + ip, + ) + + +@dsl_user_op +def bulk_copy_to_peer( + dest: Int32, + source: Int32, + size: Int32, + mbar: Int32, + peer: Int32, + *, + loc=None, + ip=None, +) -> None: + """Bulk-copy ``size`` bytes of this CTA's shared memory at ``source`` to + ``dest`` of cluster CTA ``peer`` (the same offset as a local address), + completing on the peer's ``mbar``.""" + _asm( + None, + [_v(Int32(v), loc, ip) for v in (dest, source, size, mbar, peer)], + "{\n\t.reg .b32 rd, rb;\n\tmapa.shared::cluster.u32 rd, $0, $4;\n\t" + "mapa.shared::cluster.u32 rb, $3, $4;\n\t" + "cp.async.bulk.shared::cluster.shared::cta.mbarrier::complete_tx::bytes " + "[rd], [$1], $2, [rb];\n\t}", + "r,r,r,r,r", + True, + loc, + ip, + ) + + +class SparseMlaDecodeCute: + """The sparse-MLA decode (module docstring) of tokens whose sparse sets + are at most ``chunks`` x 128 columns wide; ``has_sinks``: add each head's + attention sink to the softmax denominator.""" + + def __init__(self, chunks: int, has_sinks: bool) -> None: + assert 1 <= chunks <= MAX_CHUNKS + self.n = chunks + self.has_sinks = has_sinks + # Received partial blocks a CTA owns at most (block j -> CTA j % n). + self.slots = (4 + chunks - 1) // chunks + # Partials go out from the KV tile's upper half, part 4 + j for block j, + # and arrive in its lower half (at most 4 parts: slots x (n - 1) <= 4). + self.stage_part = 4 + assert self.slots * (chunks - 1) <= self.stage_part + + @cute.jit + def __call__( + self, + q_ptr: Int64, + q_stride: Int32, + swa_ptr: Int64, + swa_stride: Int32, + comp_ptr: Int64, + comp_stride: Int32, + idx_ptr: Int64, + idx_stride: Int32, + idx_cols: Int32, + lens_ptr: Int64, + sinks_ptr: Int64, + out_ptr: Int64, + out_stride: Int32, + scale_log2: Float32, + out_scale: Float32, + tokens: Int32, + stream: cuda.CUstream, + ): + """Launch on ``stream``: one cluster of ``n`` CTAs a token. Pointers + are addresses; strides are in elements (bytes for the FP8 query and + pools, BF16 elements for ``out``); ``scale_log2`` is + ``bmm1_scale * log2(e)``.""" + self.kernel( + q_ptr, + q_stride, + swa_ptr, + swa_stride, + comp_ptr, + comp_stride, + idx_ptr, + idx_stride, + idx_cols, + lens_ptr, + sinks_ptr, + out_ptr, + out_stride, + scale_log2, + out_scale, + ).launch( + grid=(self.n, tokens, 1), + block=(THREADS, 1, 1), + cluster=(self.n, 1, 1), + stream=stream, + use_pdl=True, + ) + + @cute.kernel + def kernel( + self, + q_ptr: Int64, + q_stride: Int32, + swa_ptr: Int64, + swa_stride: Int32, + comp_ptr: Int64, + comp_stride: Int32, + idx_ptr: Int64, + idx_stride: Int32, + idx_cols: Int32, + lens_ptr: Int64, + sinks_ptr: Int64, + out_ptr: Int64, + out_stride: Int32, + scale_log2: Float32, + out_scale: Float32, + ): + n = self.n + tid, _, _ = cute.arch.thread_idx() + chunk, token, _ = cute.arch.block_idx() + warp = cute.arch.make_warp_uniform(tid // 32) + lane = tid % 32 + # Softmax groups: 8 threads a head, 16 rows each. + head = tid // 8 + part = tid % 8 + neg_inf = Float32(-math.inf) + + smem = utils.SmemAllocator() + kv = smem.allocate_array(cutlass.Uint8, 4 * BLOCK_BYTES, byte_alignment=1024) + sq = smem.allocate_array(cutlass.Uint8, 4 * Q_BLOCK_BYTES, byte_alignment=1024) + sp = smem.allocate_array(cutlass.Uint8, HEADS * 128, byte_alignment=1024) + s_addr = smem.allocate_tensor(Int64, cute.make_layout(ROWS), 16) + # Per head: chunk max, chunk sum, merge factor, sink (base 2). + s_stat = smem.allocate_tensor(Float32, cute.make_layout((HEADS, 4)), 16) + # Per head and chunk of the cluster: max (rows 0-15), sum (rows 16-31). + s_ml = smem.allocate_tensor(Float32, cute.make_layout((2 * HEADS, n)), 16) + mbars = smem.allocate_array(Int64, 3, byte_alignment=8) + taddr = smem.allocate_array(Int32, 4, byte_alignment=16) + mbar_mma = mbars + mbar_ml = mbars + 1 + mbar_o = mbars + 2 + if tid == 0: + cute.arch.mbarrier_init(mbar_mma, 1) + cute.arch.mbarrier_init(mbar_ml, 1) + cute.arch.mbarrier_init(mbar_o, 1) + cute.arch.mbarrier_init_fence() + # The sinks are weights: read them before the grid wait. + if tid < HEADS: + sink = neg_inf + if cutlass.const_expr(self.has_sinks): + sink = ld_global_f32(sinks_ptr + Int64(tid * 4)) * Float32(LOG2E) + s_stat[tid, 3] = sink + cute.arch.sync_threads() + if cutlass.const_expr(n > 1): + # Peers must see these barriers initialized before their st.async. + cute.arch.cluster_arrive_relaxed() + # Indices, query and KV all come from earlier kernels. + cute.arch.griddepcontrol_wait() + if warp == 0: + cute.arch.alloc_tmem(TMEM_COLS, taddr) + cute.arch.relinquish_tmem_alloc_permit() + + # 1. This thread's KV row: its address, or 0 when masked. The length + # and the index load independently. + col = chunk * ROWS + tid + length = ld_global_s32(lens_ptr + Int64(token) * 4) + index = ld_global_s32( + idx_ptr + + (Int64(token) * Int64(idx_stride) + Int64(cute.min(col, idx_cols - 1))) + * 4 + ) + address = Int64(0) + if (col < length) & (col < idx_cols): # noqa: SIM102 + if index >= 0: + if chunk == 0: + address = swa_ptr + Int64(index) * Int64(swa_stride) + else: + address = comp_ptr + Int64(index) * Int64(comp_stride) + masked = address == Int64(0) + s_addr[tid] = address + cute.arch.sync_threads() + + # Gather in four cp.async groups, one per 128-byte K block of the rows (with + # the query's K block): in a group a warp copies 4 rows a step, 8 lanes a + # row, each 16-byte piece to its 128B-swizzled slot. The MMAs of a K block + # start as soon as its group has landed, while the later groups load. + kv_base = kv.toint() + sq_base = sq.toint() + sp_base = sp.toint() + q_row = q_ptr + Int64(token) * Int64(q_stride) + q_head = tid // 8 + q_piece = tid % 8 + piece = lane % 8 + for kb in cutlass.range_constexpr(4): + cp_async_16( + sq_base + + kb * Q_BLOCK_BYTES + + q_head * 128 + + ((q_piece ^ (q_head % 8)) * 16), + q_row + Int64(q_head * HEAD_DIM + kb * 128 + q_piece * 16), + ) + for i in cutlass.range_constexpr(ROWS // 16): + row = (i * 4 + warp) * 4 + lane // 8 + source = s_addr[row] + dest = ( + kv_base + kb * BLOCK_BYTES + row * 128 + ((piece ^ (row % 8)) * 16) + ) + if source != 0: + cp_async_16(dest, source + Int64(kb * 128 + piece * 16)) + else: + st_shared_zero16(dest) + cute.arch.cp_async_commit_group() + + # 2. S^T = KV x Q^T, 4 MMAs of K = 32 a K block. + tmem = cute.make_tensor(taddr, cute.make_layout(1))[0] + for kb in cutlass.range_constexpr(4): + cute.arch.cp_async_wait_group(3 - kb) + cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.sync_threads() + if warp == 0: + _tcgen05.fence_after_thread_sync() + for k4 in cutlass.range_constexpr(4): + adesc = Uint64(SDESC_K) | Uint64( + (kv_base + kb * BLOCK_BYTES + k4 * 32) >> 4 + ) + bdesc = Uint64(SDESC_K) | Uint64( + (sq_base + kb * Q_BLOCK_BYTES + k4 * 32) >> 4 + ) + with cute.arch.elect_one(): + mma_f8( + tmem, + adesc, + bdesc, + Uint32(IDESC_QK), + Int32(1 if kb + k4 > 0 else 0), + ) + if cutlass.const_expr(kb == 3): + _tcgen05.commit(mbar_mma) + with cute.arch.elect_one(): + cute.arch.griddepcontrol_launch_dependents() + cute.arch.mbarrier_wait(mbar_mma, 0) + _tcgen05.fence_after_thread_sync() + + # 3. Chunk softmax. Each thread's TMEM lane is one KV row; the scores are + # transposed through the (now free) query tile so that 8 threads reduce + # one head's 128 rows. + s = cute.make_rmem_tensor((HEADS,), Float32) + s.store(_tcgen05.ld(warp * 32, tmem, "32x32b", HEADS)) + _tcgen05.wait_ld() + s_t = cute.make_tensor( + cute.recast_ptr(sq, dtype=Float32), + cute.make_layout((HEADS, ROWS), stride=(ROWS, 1)), + ) + for h in cutlass.range_constexpr(HEADS): + value = s[h] * scale_log2 + if masked: + value = neg_inf + s_t[h, tid] = value + cute.arch.sync_threads() + vals = cute.make_rmem_tensor((16,), Float32) + for i in cutlass.range_constexpr(16): + vals[i] = s_t[head, part * 16 + i] + peak = vals[0] + for i in cutlass.range_constexpr(1, 16): + peak = cute.arch.fmax(peak, vals[i]) + peak = cute.arch.warp_reduction_max(peak, threads_in_group=8) + base = peak + if peak == neg_inf: + base = Float32(0.0) + total = Float32(0.0) + for i in cutlass.range_constexpr(16): + vals[i] = cute.math.exp2(vals[i] - base, fastmath=True) + total = total + vals[i] + total = cute.arch.warp_reduction_sum(total, threads_in_group=8) + # P^T as the K-major B operand: head h's 128 rows are one 128-byte row, + # this thread's 16 at its swizzled 16-byte slot. + scale_p = Float32(P_SCALE) + st_shared_v4( + sp_base + head * 128 + ((part ^ (head % 8)) * 16), + pack_e4m3x4( + vals[0] * scale_p, + vals[1] * scale_p, + vals[2] * scale_p, + vals[3] * scale_p, + ), + pack_e4m3x4( + vals[4] * scale_p, + vals[5] * scale_p, + vals[6] * scale_p, + vals[7] * scale_p, + ), + pack_e4m3x4( + vals[8] * scale_p, + vals[9] * scale_p, + vals[10] * scale_p, + vals[11] * scale_p, + ), + pack_e4m3x4( + vals[12] * scale_p, + vals[13] * scale_p, + vals[14] * scale_p, + vals[15] * scale_p, + ), + ) + if part == 0: + s_stat[head, 0] = peak + s_stat[head, 1] = total + cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.sync_threads() + + # 4. O^T = V^T x P^T, four MMAs of 128 head dims (K = 128 rows each); + # the statistics exchange and the merge factors overlap them. + if warp == 0: + _tcgen05.fence_after_thread_sync() + for j in cutlass.range_constexpr(4): + for kk in cutlass.range_constexpr(ROWS // 32): + adesc = Uint64(SDESC_MN) | Uint64( + (kv_base + j * BLOCK_BYTES + kk * 32 * 128) >> 4 + ) + bdesc = Uint64(SDESC_K) | Uint64((sp_base + kk * 32) >> 4) + with cute.arch.elect_one(): + mma_f8( + tmem + HEADS + HEADS * j, + adesc, + bdesc, + Uint32(IDESC_PV), + Int32(1 if kk > 0 else 0), + ) + _tcgen05.commit(mbar_mma) + if cutlass.const_expr(n > 1): + mbar_ml_addr = mbar_ml.toint() + s_ml_base = s_ml.iterator.toint() + cute.arch.cluster_wait() + if tid == 0: + mbar_ops.arrive_expect_tx(mbar_ml, n * 2 * HEADS * 4) + if tid < 2 * HEADS: + stat = s_stat[tid % HEADS, tid // HEADS] + for peer in cutlass.range_constexpr(n): + st_async_peer_b32( + s_ml_base + (chunk * 2 * HEADS + tid) * 4, + mbar_ml_addr, + Int32(peer), + recast_val(stat, Uint32), + ) + cute.arch.mbarrier_wait(mbar_ml, 0) + # 5. Merge factors, one thread a head: exp2(m_c - M) / L, L with the sink. + if tid < HEADS: + own = s_stat[tid, 0] + sink = s_stat[tid, 3] + top = sink + denom = Float32(0.0) + if cutlass.const_expr(n > 1): + for c in cutlass.range_constexpr(n): + top = cute.arch.fmax(top, s_ml[tid, c]) + for c in cutlass.range_constexpr(n): + peak_c = s_ml[tid, c] + if peak_c != neg_inf: + denom = ( + denom + + cute.math.exp2(peak_c - top, fastmath=True) + * s_ml[HEADS + tid, c] + ) + else: + top = cute.arch.fmax(top, own) + if own != neg_inf: + denom = cute.math.exp2(own - top, fastmath=True) * s_stat[tid, 1] + if cutlass.const_expr(self.has_sinks): + denom = denom + cute.math.exp2(sink - top, fastmath=True) + f = Float32(0.0) + if own != neg_inf: # noqa: SIM102 + if denom > Float32(0.0): + f = ( + cute.math.exp2(own - top, fastmath=True) + / denom + * out_scale + * Float32(1.0 / P_SCALE) + ) + s_stat[tid, 2] = f + cute.arch.mbarrier_wait(mbar_mma, 1) + _tcgen05.fence_after_thread_sync() + cute.arch.sync_threads() + if cutlass.const_expr(n > 1): + # This CTA's MMAs are done: peers may now write into its tile. + cute.arch.cluster_arrive() + o = cute.make_rmem_tensor((4 * HEADS,), Float32) + o.store(_tcgen05.ld(warp * 32, tmem + HEADS, "32x32b", 4 * HEADS)) + _tcgen05.wait_ld() + factor = cute.make_rmem_tensor((HEADS,), Float32) + for h in cutlass.range_constexpr(HEADS): + factor[h] = s_stat[h, 2] + for j in cutlass.range_constexpr(4): + for h in cutlass.range_constexpr(HEADS): + o[j * HEADS + h] = o[j * HEADS + h] * factor[h] + # Thread tid holds head dim (128 j + tid) of every head. + out_row = out_ptr + Int64(token) * Int64(out_stride) * 2 + if cutlass.const_expr(n == 1): + # Stage [block][head][dim] BF16 in the tile, then 16-byte stores. + for j in cutlass.range_constexpr(4): + for h in cutlass.range_constexpr(HEADS): + st_shared_bf16( + kv_base + ((j * HEADS + h) * 128 + tid) * 2, o[j * HEADS + h] + ) + cute.arch.sync_threads() + for k in cutlass.range_constexpr(4 * HEADS * 16 // THREADS): + piece = k * THREADS + tid + rowid = piece // 16 + d16 = piece % 16 + blk = rowid // HEADS + h = rowid % HEADS + copy16_shared_to_global( + kv_base + piece * 16, + out_row + Int64((h * HEAD_DIM + blk * 128 + d16 * 8) * 2), + ) + else: + mbar_o_addr = mbar_o.toint() + owned = Int32(0) + for j in cutlass.range_constexpr(4): + if chunk == j % n: + owned = owned + 1 + # The tile as 8 FP32 parts of [head][row]: head-major, so a warp's + # 32 rows of one head are 32 consecutive words (no bank conflicts + # on either side). + parts = cute.make_tensor( + cute.recast_ptr(kv, dtype=Float32), + cute.make_layout((HEADS, ROWS, 8), stride=(ROWS, 1, ROWS * HEADS)), + ) + for j in cutlass.range_constexpr(4): + if chunk != j % n: + for h in cutlass.range_constexpr(HEADS): + parts[h, tid, self.stage_part + j] = o[j * HEADS + h] + cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.sync_threads() + if tid == 0: # noqa: SIM102 + if owned > 0: + mbar_ops.arrive_expect_tx(mbar_o, owned * (n - 1) * PART_BYTES) + # Every peer's MMAs are done, so its tile's lower half can take + # partials: sender s of block j lands in part (j // n) (n - 1) + s, + # the senders numbered without the owner. + cute.arch.cluster_wait() + if tid == 0: + for j in cutlass.range_constexpr(4): + owner = j % n + if chunk != owner: + sender = chunk + if chunk > owner: + sender = chunk - 1 + bulk_copy_to_peer( + kv_base + ((j // n) * (n - 1) + sender) * PART_BYTES, + kv_base + (self.stage_part + j) * PART_BYTES, + Int32(PART_BYTES), + mbar_o_addr, + Int32(owner), + ) + if owned > 0: + cute.arch.mbarrier_wait(mbar_o, 0) + # Owned block j (slot j // n): the sum, BF16 [slot][head][dim] + # in the query tile, then 16-byte stores. + for j in cutlass.range_constexpr(4): + if chunk == j % n: + for h in cutlass.range_constexpr(HEADS): + acc = o[j * HEADS + h] + for c in cutlass.range_constexpr(n - 1): + acc = acc + parts[h, tid, (j // n) * (n - 1) + c] + st_shared_bf16( + sq_base + (((j // n) * HEADS + h) * 128 + tid) * 2, acc + ) + cute.arch.sync_threads() + for slot in cutlass.range_constexpr(self.slots): + if chunk + slot * n < 4: + for k in cutlass.range_constexpr(HEADS * 16 // THREADS): + piece = k * THREADS + tid + h = piece // 16 + d16 = piece % 16 + copy16_shared_to_global( + sq_base + (slot * HEADS * 128) * 2 + piece * 16, + out_row + + Int64( + (h * HEAD_DIM + (chunk + slot * n) * 128 + d16 * 8) + * 2 + ), + ) + # The bulk copies read this CTA's staged partials asynchronously: + # stay resident until every owner has received (its arrive follows + # its mbarrier wait). + cute.arch.cluster_arrive_relaxed() + cute.arch.cluster_wait() + _tcgen05.fence_before_thread_sync() + cute.arch.sync_threads() + if warp == 0: + _tcgen05.fence_after_thread_sync() + cute.arch.dealloc_tmem( + cute.make_ptr(Float32, tmem, cute.AddressSpace.tmem), TMEM_COLS + ) + + +@cache +def _kernel(device_index: int, chunks: int, has_sinks: bool): + """``SparseMlaDecodeCute`` compiled for a device, chunk count and sinks.""" + with torch.accelerator.device_index(device_index): + return cute.compile( + SparseMlaDecodeCute(chunks, has_sinks), + Int64(0), + Int32(0), + Int64(0), + Int32(0), + Int64(0), + Int32(0), + Int64(0), + Int32(0), + Int32(1), + Int64(0), + Int64(0), + Int64(0), + Int32(0), + Float32(0), + Float32(0), + Int32(1), + make_fake_stream(), + ) + + +@cache +def _max_tokens(device_index: int, chunks: int) -> int: + """Tokens whose clusters are all resident at once: ``CTAS_PER_SM`` times + the resident clusters of ``chunks`` CTAs at one CTA an SM.""" + with torch.accelerator.device_index(device_index): + clusters = cutlass.utils.HardwareInfo(device_index).get_max_active_clusters( + chunks + ) + return CTAS_PER_SM * clusters + + +def precompile(device_index: int, has_sinks: bool) -> None: + """Compile the kernels of DeepSeek-V4's two widths (128 sliding-window + columns alone, and 640) and count the resident clusters, outside any CUDA + graph capture; other widths compile at their first launch. + + Args: + device_index: The CUDA device to compile for. + has_sinks: Whether the layers add attention sinks. + + """ + for chunks in (1, MAX_CHUNKS): + _kernel(device_index, chunks, has_sinks) + _max_tokens(device_index, chunks) + + +def _is_pool(cache: torch.Tensor) -> bool: + """An FP8 KV pool of 512-byte rows ``stride(-2)`` apart, 16-byte aligned.""" + return ( + cache.dtype == torch.float8_e4m3fn + and cache.dim() >= 2 + and cache.shape[-1] == HEAD_DIM + and cache.stride(-1) == 1 + and cache.stride(-2) % 16 == 0 + and cache.data_ptr() % 16 == 0 + ) + + +def supports_sparse_mla_cute( + query: torch.Tensor, + swa_kv_cache: torch.Tensor, + compressed_kv_cache: torch.Tensor, + sparse_indices: torch.Tensor, + sparse_topk_lens: torch.Tensor, + out: torch.Tensor, + bmm1_scale: float | torch.Tensor, + bmm2_scale: float | torch.Tensor, + sinks: torch.Tensor | None, +) -> bool: + """Whether ``sparse_mla_decode`` takes these arguments: 16 FP8 query heads + of 512 (BF16 output), FP8 pools, int32 index rows of at most 640 columns, + scalar scales, FP32 sinks, and a batch whose clusters are all resident.""" + if ( + query.dtype != torch.float8_e4m3fn + or query.dim() != 3 + or query.shape[1:] != (HEADS, HEAD_DIM) + or query.stride(2) != 1 + or query.stride(1) != HEAD_DIM + or query.stride(0) % 16 != 0 + or query.data_ptr() % 16 != 0 + ): + return False + tokens = query.shape[0] + if ( + out.dtype != torch.bfloat16 + or out.shape != query.shape + or out.stride(2) != 1 + or out.stride(1) != HEAD_DIM + or out.stride(0) % 8 != 0 + or out.data_ptr() % 16 != 0 + ): + return False + if ( + sparse_indices.dtype != torch.int32 + or sparse_indices.dim() != 2 + or sparse_indices.shape[0] != tokens + or sparse_indices.stride(1) != 1 + or not 0 < sparse_indices.shape[1] <= MAX_CHUNKS * ROWS + ): + return False + if ( + sparse_topk_lens.dtype != torch.int32 + or not sparse_topk_lens.is_contiguous() + or sparse_topk_lens.numel() != tokens + ): + return False + if sinks is not None and ( + sinks.dtype != torch.float32 + or not sinks.is_contiguous() + or sinks.numel() < HEADS + ): + return False + if not ( + isinstance(bmm1_scale, float | int) and isinstance(bmm2_scale, float | int) + ): + return False + if not (_is_pool(swa_kv_cache) and _is_pool(compressed_kv_cache)): + return False + chunks = (sparse_indices.shape[1] + ROWS - 1) // ROWS + return 0 < tokens <= _max_tokens(query.device.index, chunks) + + +def sparse_mla_decode( + query: torch.Tensor, + swa_kv_cache: torch.Tensor, + compressed_kv_cache: torch.Tensor, + sparse_indices: torch.Tensor, + sparse_topk_lens: torch.Tensor, + out: torch.Tensor, + bmm1_scale: float, + bmm2_scale: float, + sinks: torch.Tensor | None, +) -> None: + """Launch `SparseMlaDecodeCute` on the current stream (see + `supports_sparse_mla_cute` for the accepted layouts). + + Args: + query: [tokens, 16, 512] FP8 E4M3. + swa_kv_cache: FP8 pool of 512-byte rows ``stride(-2)`` apart; the first + 128 columns of ``sparse_indices`` are its row offsets. + compressed_kv_cache: The pool of the other columns, likewise. + sparse_indices: [tokens, width] int32 row offsets, -1 for none + (``width`` <= 640). + sparse_topk_lens: [tokens] int32 active columns of each token. + out: [tokens, 16, 512] BF16 attention output. + bmm1_scale: Score scale (the softmax scale times the FP8 scales). + bmm2_scale: Output scale. + sinks: [16] FP32 attention sinks (natural log), or None. + + """ + assert supports_sparse_mla_cute( + query, + swa_kv_cache, + compressed_kv_cache, + sparse_indices, + sparse_topk_lens, + out, + bmm1_scale, + bmm2_scale, + sinks, + ) + tokens = query.shape[0] + width = sparse_indices.shape[1] + chunks = (width + ROWS - 1) // ROWS + _kernel(query.device.index, chunks, sinks is not None)( + Int64(query.data_ptr()), + Int32(query.stride(0)), + Int64(swa_kv_cache.data_ptr()), + Int32(swa_kv_cache.stride(-2)), + Int64(compressed_kv_cache.data_ptr()), + Int32(compressed_kv_cache.stride(-2)), + Int64(sparse_indices.data_ptr()), + Int32(sparse_indices.stride(0)), + Int32(width), + Int64(sparse_topk_lens.data_ptr()), + Int64(sinks.data_ptr() if sinks is not None else 0), + Int64(out.data_ptr()), + Int32(out.stride(0)), + Float32(bmm1_scale * LOG2E), + Float32(bmm2_scale), + Int32(tokens), + current_cu_stream(), + ) diff --git a/vllm/models/deepseek_v41/nvidia/specialized_decode/attention/step.py b/vllm/models/deepseek_v41/nvidia/specialized_decode/attention/step.py index 9eebbfc34934..9c64ba8237df 100644 --- a/vllm/models/deepseek_v41/nvidia/specialized_decode/attention/step.py +++ b/vllm/models/deepseek_v41/nvidia/specialized_decode/attention/step.py @@ -7,8 +7,9 @@ low-batch kernels where they apply: - ``fused_wqa_wkv`` reads the MXFP8 layer input the fused all-reduce + mHC - wrote instead of quantizing it again; up to ``LOW_BATCH_MAX_ROWS`` rows it, - ``wq_b`` and ``wo_b`` run on the low-latency MXFP8 GEMM (``mxfp8``). + wrote instead of quantizing it again; up to ``WIDE_LOW_LATENCY_MAX_ROWS`` + rows it, ``wq_b`` and ``wo_b`` run on the low-latency MXFP8 GEMM, and above + on the stock ``cute-dsl`` GEMM with pinned tactics (``mxfp8``). - The side projections (compressor ``kv_score``, indexer weights) join their side streams at their first reader rather than before the query path, and the indexer inputs are built on side streams: the indexer query once ``qr`` @@ -17,10 +18,12 @@ ``wk`` run on ``ll_bf16`` kernels (``aux_projections``), the key store reads ``wk``'s FP32 output, and ``wo_b`` publishes its output into the mailbox of the fused all-reduce that reduces it next (``boundary.publishing_gemm``). -- Dense indexers row-shard their decode rows (``indexer.decode_rows``); the - candidate source starts the candidate consumers' expansion on a side stream - right after its indexer, and index sources first join the side-stream - sparse-index builds, which read the top-k indices they overwrite. +- Dense indexers row-shard their decode rows (``indexer.decode_rows``) and + plan them (top-k reset, slice schedule) while the current stream waits for + the index key; the candidate source starts the candidate consumers' + expansion on a side stream right after its indexer, and index sources first + join the side-stream sparse-index builds, which read the top-k indices they + overwrite. - The sparse MLA reuses the step's sparse-index tables and runs on persistent KV counters (``sparse_mla``). """ @@ -34,16 +37,30 @@ from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation from vllm.model_executor.layers.sparse_attn_indexer import SparseAttnIndexer from vllm.models.deepseek_v41.attention import DeepseekV4Indexer +from vllm.models.deepseek_v41.common.ops.query_quant import fused_q_kv_rmsnorm_quant from ...flashinfer_sparse import DeepseekV4FlashInferMLAAttention from ...ops.o_proj import dsv41_o_proj +from .. import WIDE_LOW_LATENCY_MAX_ROWS from .aux_projections import aux_projection_gemm, aux_projection_gemm_accepts -from .indexer import CandidateOverlap, decode_rows, rope_quant_q +from .indexer import ( + CandidateOverlap, + DecodeRowsPlan, + finish_decode_rows, + plan_decode_rows, + rope_quant_q, +) from .indexer_k_store import indexer_k_norm_rope_store_fp32 -from .mxfp8 import low_latency_mm, mxfp8_eligible +from .mxfp8 import ( + low_latency_mm, + mxfp8_eligible, + projection_mm, + wide_or_pinned_mm, +) from .sparse_mla import SparseIndexTables, SparseMlaDecode if TYPE_CHECKING: + from ...ops.cute_dsl import LamportAllGather from ..boundary.publishing_gemm import PublishingMxfp8Linear from ..runtime.decode import DecodeStep @@ -78,10 +95,13 @@ def __init__( mxfp8_input: bool, source: CandidateOverlap | None, consumer: CandidateOverlap | None, + all_gather: "LamportAllGather | None", ) -> None: self.attn = attn self.tables = tables self.mla = mla + # Gathers the row-sharded indexer's slices. + self.all_gather = all_gather self.publisher = ( publisher if publisher is not None and publisher.supports(attn.wo_b, 1) @@ -115,17 +135,34 @@ def run( out = attn._alloc_attn_out(x.shape[0], x) assert isinstance(out, torch.Tensor) qr_kv, kv_score, indexer_weights = self._input_projections(x, quant, step) - qr, qr_scale, kv = attn._split_qkv_and_norm(qr_kv) - q, index_q, index_q_scale, index_weights = self._prepare( - qr, kv, qr_scale, kv_score, indexer_weights, positions, step + qr, qr_scale, kv = self._split_qkv_and_norm(qr_kv) + q, index_q, index_q_scale, index_weights, plan = self._prepare( + qr, kv, qr_scale, kv_score, indexer_weights, positions, x.shape[0], step ) indexer = attn.indexer if indexer is not None and index_q is not None: assert index_weights is not None - self._indexer(indexer, x, index_q, index_q_scale, index_weights, step) + self._indexer(indexer, x, index_q, index_q_scale, index_weights, plan, step) self.mla.run(attn, q, out) + # The fused wo_a triggers wo_b early up to WIDE_LOW_LATENCY_MAX_ROWS, + # where wo_b runs on the low-latency GEMM (above 8 rows, where its shape + # has a wide tactic), which loads its weights before its own wait. if step.low_batch and (self.publisher is not None or self.wo_b_mxfp8): - return dsv41_o_proj(attn, out, positions, wo_b_proj=self._wo_b_low_batch) + return dsv41_o_proj( + attn, + out, + positions, + wo_b_proj=self._wo_b_low_batch, + early_trigger=True, + ) + if self.wo_b_mxfp8: + return dsv41_o_proj( + attn, + out, + positions, + wo_b_proj=self._wo_b_pinned, + early_trigger=step.rows <= WIDE_LOW_LATENCY_MAX_ROWS, + ) return attn._o_proj(out, positions) def _input_projections( @@ -175,12 +212,34 @@ def _input_projections( qr_kv = attn._fused_wqa_wkv_gemm(x) else: qr_kv = None - if step.low_batch and self.wqa_mxfp8: - qr_kv = low_latency_mm(attn.fused_wqa_wkv, quant) + if self.wqa_mxfp8: + qr_kv = projection_mm(attn.fused_wqa_wkv, quant, step.low_batch) if qr_kv is None: qr_kv = attn.fused_wqa_wkv(quant)[0] return qr_kv, results[0], results[1] + def _split_qkv_and_norm( + self, qr_kv: torch.Tensor + ) -> tuple[torch.Tensor | QuantizedActivation, torch.Tensor | None, torch.Tensor]: + """``attn._split_qkv_and_norm``, with the fused norm triggering its + PDL dependents early up to ``WIDE_LOW_LATENCY_MAX_ROWS`` rows: there + this step's ``wq_b`` runs on the low-latency GEMM (above 8 rows, where + its shape has a wide tactic), which loads its weights before its own + wait, overlapping them with the norm.""" + attn = self.attn + if attn.q_lora_rank % 32 or not attn._can_fuse_query_quant: + return attn._split_qkv_and_norm(qr_kv) + qr, kv = qr_kv.split([attn.q_lora_rank, attn.head_dim], dim=-1) + qr_quant, kv = fused_q_kv_rmsnorm_quant( + qr, + kv, + attn.q_norm.weight.data, + attn.kv_norm.weight.data, + attn.eps, + early_trigger=qr.shape[0] <= WIDE_LOW_LATENCY_MAX_ROWS, + ) + return qr_quant, None, kv + def _prepare( self, qr: torch.Tensor | QuantizedActivation, @@ -189,9 +248,14 @@ def _prepare( kv_score: torch.Tensor | None, indexer_weights: torch.Tensor | None, positions: torch.Tensor, + num_tokens: int, step: "DecodeStep", ) -> tuple[ - torch.Tensor, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None + torch.Tensor, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + DecodeRowsPlan | None, ]: """``_prepare_and_attn`` up to the indexer, with the indexer inputs on side streams. @@ -199,9 +263,10 @@ def _prepare( From the moment ``qr`` is ready, aux stream 2 runs the indexer query (``wq_b``, then RoPE/quantize once the weights GEMM is done) and aux stream 0, after its ``kv_score`` GEMM, the compressor and the index - key; the current stream runs the main query and KV insertion, then the - compressed-cache insertion once the latent is ready. Returns the query - and the indexer query, its scale and weights. + key; the current stream runs the main query and KV insertion, a dense + indexer's plan (``plan_decode_rows``), then the compressed-cache + insertion once the latent is ready. Returns the query, the indexer + query, its scale and weights, and the dense indexer's plan. """ attn = self.attn indexer, compressor = attn.indexer, attn.compressor @@ -239,21 +304,28 @@ def _prepare( index_k_done.record() q = None - if step.low_batch and self.wq_b_mxfp8: - q = low_latency_mm(attn.wq_b, qr) + if self.wq_b_mxfp8: + q = projection_mm(attn.wq_b, qr, step.low_batch) if q is None: q = attn._wq_b_proj(qr, qr_scale) q = q.view(-1, attn.n_local_heads, attn.head_dim) q = attn._fused_qnorm_rope_kv_insert( q, kv, positions, get_forward_context().attn_metadata ) + # The plan reads neither the index query nor the keys: this stream + # would otherwise wait for them here. + plan = None + if indexer is not None and isinstance(indexer.indexer_op, SparseAttnIndexer): + # The side-stream table builds read the top-k rows it resets. + self.tables.join() + plan = plan_decode_rows(indexer.indexer_op, num_tokens) if compressor is not None: latent_ready.wait() compressor.insert_cache(latent, positions, attn.rotary_emb) index_k_done.wait() if indexer is not None: index_q_done.wait() - return q, index_q, index_q_scale, index_weights + return q, index_q, index_q_scale, index_weights, plan def _produce_k( self, @@ -295,21 +367,24 @@ def _indexer( index_q: torch.Tensor, index_q_scale: torch.Tensor | None, index_weights: torch.Tensor, + plan: DecodeRowsPlan | None, step: "DecodeStep", ) -> None: """The layer's indexer: the top-k indices (and in the candidate - source, the candidate blocks) of the step.""" - # The side-stream table builds read the top-k indices this overwrites. - self.tables.join() + source, the candidate blocks) of the step; a dense indexer finishes + its ``plan``.""" q_quant = (index_q, index_q_scale) if index_q_scale is not None else index_q op = indexer.indexer_op if isinstance(op, SparseAttnIndexer): - decode_rows(op, x, q_quant, index_weights) + finish_decode_rows(op, x, q_quant, index_weights, plan, self.all_gather) else: + # The side-stream table builds read the top-k indices this + # overwrites. + self.tables.join() if self.consumer is not None: self.consumer.join() op(x, q_quant, None, index_weights) - if self.source is not None and step.overlap_candidates and self.source.warmed: + if self.source is not None and self.source.warmed: self.source.fork() def _wo_b_low_batch(self, z: torch.Tensor | QuantizedActivation) -> torch.Tensor: @@ -320,3 +395,9 @@ def _wo_b_low_batch(self, z: torch.Tensor | QuantizedActivation) -> torch.Tensor return self.publisher(wo_b, z) out = low_latency_mm(wo_b, z) if self.wo_b_mxfp8 else None return out if out is not None else self.attn._wo_b_proj(z) + + def _wo_b_pinned(self, z: torch.Tensor | QuantizedActivation) -> torch.Tensor: + """``wo_b`` above ``LOW_BATCH_MAX_ROWS`` rows: on the wide low-latency + or the pinned ``cute-dsl`` GEMM where one runs.""" + out = wide_or_pinned_mm(self.attn.wo_b, z) + return out if out is not None else self.attn._wo_b_proj(z) 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..6f466453557a 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 @@ -42,7 +42,8 @@ def publish_target(self) -> tuple[int, int]: 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.""" + stay valid for this instance's lifetime. This is the LL mailbox, for + steps of fewer than ``bt_min_rows`` rows.""" return self._multicast, self._stage_state.data_ptr() def mark_published(self, x: torch.Tensor) -> None: @@ -75,10 +76,13 @@ def __call__( if published is not x: raise RuntimeError("reducing a tensor other than the published one") m = x.shape[0] + if m >= self.bt_min_rows: + # The producer stored into the LL mailbox (``publish_target``); + # ``PublishingMxfp8Linear.supports`` admits only smaller calls. + raise RuntimeError(f"a published contribution of {m} rows takes BT") self._check(m, residual, post, comb, pre) - return self._collect( - residual, post, comb, pre, norm_weight, norm_eps, m, mxfp8_out - ) + consumer = self._consumer(residual, pre, norm_weight, m, mxfp8_out, False) + return self._collect(consumer, residual, post, comb, norm_eps, m) def finalize( self, diff --git a/vllm/models/deepseek_v41/nvidia/specialized_decode/boundary/boundary.py b/vllm/models/deepseek_v41/nvidia/specialized_decode/boundary/boundary.py index e7404327903f..f39e8e4aced0 100644 --- a/vllm/models/deepseek_v41/nvidia/specialized_decode/boundary/boundary.py +++ b/vllm/models/deepseek_v41/nvidia/specialized_decode/boundary/boundary.py @@ -8,8 +8,7 @@ normalizes the streams into the next sublayer's input (mHC pre). Here the fused all-reduce + mHC kernel (``BoundaryAllReduce``) does all of that for every step the executor runs; the next sublayer's mHC coefficients then come from the -stock ``mhc_pre_delayed_overlap``, on the shared mHC side stream for up to -``MHC_OVERLAP_MAX_ROWS`` rows and inline after the fused kernel above. +stock ``mhc_pre_delayed_overlap`` on the shared mHC side stream. The fused kernel also finalizes a MoE output left unfinalized, collects a ``wo_b`` output the publishing GEMM already stored into its mailbox, writes the @@ -19,16 +18,12 @@ """ from dataclasses import dataclass -from functools import partial from typing import TYPE_CHECKING import torch from vllm.distributed import tensor_model_parallel_all_reduce -from vllm.model_executor.kernels.mhc.tilelang import ( - mhc_post_tilelang, - mhc_pre_delayed_tilelang, -) +from vllm.model_executor.kernels.mhc.tilelang import mhc_post_tilelang from vllm.model_executor.layers.fused_moe.moe_output import ( MoEOutput, UnfinalizedMoEOutput, @@ -42,7 +37,6 @@ if TYPE_CHECKING: from ...model import DeepseekV4DecoderLayer - from ..runtime.decode import DecodeStep from .all_reduce import BoundaryAllReduce @@ -84,30 +78,36 @@ def __init__(self, all_reduce: "BoundaryAllReduce") -> None: self.all_reduce = all_reduce def first( - self, layer: "DeepseekV4DecoderLayer", x: torch.Tensor, step: "DecodeStep" + self, + layer: "DeepseekV4DecoderLayer", + x: torch.Tensor, + streams: torch.Tensor | None = None, ) -> tuple[BoundaryState, torch.Tensor]: """The first layer's attention entry: the embedding broadcast to the - hc streams, whose identity pre-mix selects stream zero.""" + hc streams (or written with them, ``streams``), whose identity pre-mix + selects stream zero.""" assert layer.hc_attn_fn_broadcast is not None - residual = x.unsqueeze(1).expand(-1, layer.hc_mult, -1).contiguous() - return self._pre(layer, residual, layer.hc_attn_fn_broadcast, None, x, step) + residual = ( + x.unsqueeze(1).expand(-1, layer.hc_mult, -1).contiguous() + if streams is None + else streams + ) + return self._pre(layer, residual, layer.hc_attn_fn_broadcast, None, x) def pre( self, layer: "DeepseekV4DecoderLayer", residual: torch.Tensor, state: BoundaryState, - step: "DecodeStep", ) -> tuple[BoundaryState, torch.Tensor]: """An Engram layer's attention entry, on the injected streams.""" - return self._pre(layer, residual, layer.hc_attn_fn, state.pre, None, step) + return self._pre(layer, residual, layer.hc_attn_fn, state.pre, None) def sublayer( self, layer: "DeepseekV4DecoderLayer", ffn: bool, state: BoundaryState, - step: "DecodeStep", *, quantize: bool = False, capture_aux: bool = False, @@ -119,7 +119,7 @@ def sublayer( Returns the new state, the layer input, the same input as MXFP8 with ``quantize``, and with ``capture_aux`` the mean of the streams (on the - side stream when the coefficients run there: join it before reading). + side stream, after the coefficients: join it before reading). """ if ffn: fn, scale, base = layer.hc_ffn_fn, layer.hc_ffn_scale, layer.hc_ffn_base @@ -140,14 +140,12 @@ def sublayer( None if quant is None else (quant.data, quant.scale), ) assert layer_input is not None - stream = layer.mhc_stream if step.side_stream else torch.cuda.current_stream() + stream = layer.mhc_stream assert stream is not None aux = None - if capture_aux and step.side_stream: + if capture_aux: streams_ready = torch.cuda.Event() streams_ready.record() - elif capture_aux: - aux = residual.mean(dim=1) post, comb, layer_input, next_pre = mhc_pre_delayed_overlap( residual, fn, @@ -164,7 +162,7 @@ def sublayer( stream=stream, layer_input=layer_input, ) - if capture_aux and step.side_stream: + if capture_aux: # residual.mean(dim=1), after the coefficients on the side stream. aux = residual.new_empty(rows, hidden) stream.wait_event(streams_ready) @@ -184,11 +182,10 @@ def post_only(self, state: BoundaryState) -> torch.Tensor: x = tensor_model_parallel_all_reduce(state.pending) return mhc_post_tilelang(x, state.residual, state.post, state.comb) - def join(self, layer: "DeepseekV4DecoderLayer", step: "DecodeStep") -> None: + def join(self, layer: "DeepseekV4DecoderLayer") -> None: """Make the current stream wait for the side-stream coefficients.""" - if step.side_stream: - assert layer.mhc_stream is not None - torch.cuda.current_stream().wait_stream(layer.mhc_stream) + assert layer.mhc_stream is not None + torch.cuda.current_stream().wait_stream(layer.mhc_stream) def _pre( self, @@ -197,16 +194,10 @@ def _pre( fn: torch.Tensor, pre_mix: torch.Tensor | None, x: torch.Tensor | None, - step: "DecodeStep", ) -> tuple[BoundaryState, torch.Tensor]: """The stock attention entry without a reduction: mHC pre of - ``residual`` on the side stream, or inline with TileLang.""" - mhc_pre = ( - partial(mhc_pre_delayed_overlap, stream=layer.mhc_stream) - if step.side_stream - else mhc_pre_delayed_tilelang - ) - post, comb, layer_input, next_pre = mhc_pre( + ``residual`` on the side stream.""" + post, comb, layer_input, next_pre = mhc_pre_delayed_overlap( residual, fn, layer.hc_attn_scale, @@ -220,6 +211,7 @@ def _pre( x=x, norm_weight=layer.attn_norm.weight, norm_eps=layer.attn_norm.variance_epsilon, + stream=layer.mhc_stream, ) return BoundaryState(residual, post, comb, next_pre), layer_input 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 3677a3432b3c..3a7e25d517be 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 @@ -460,11 +460,17 @@ def __init__(self, all_reduce: BoundaryAllReduce) -> None: self._compiled: dict[tuple[Tactic, int], object] = {} def supports(self, linear: torch.nn.Module, num_tokens: int) -> bool: - """Whether ``linear`` on ``num_tokens`` rows runs the publishing GEMM.""" + """Whether ``linear`` on ``num_tokens`` rows runs the publishing GEMM. + It publishes into the LL mailbox, so it takes no call of the all-reduce's + BT sizes.""" + all_reduce = self.all_reduce + max_rows = min( + LOW_BATCH_MAX_ROWS, all_reduce.capacity, all_reduce.bt_min_rows - 1 + ) return ( mxfp8_eligible(linear) - and 0 < num_tokens <= min(LOW_BATCH_MAX_ROWS, self.all_reduce.capacity) - and linear.weight.shape[1] == self.all_reduce.hidden_size + and 0 < num_tokens <= max_rows + and linear.weight.shape[1] == all_reduce.hidden_size and getattr(linear, "bias", None) is None ) diff --git a/vllm/models/deepseek_v41/nvidia/specialized_decode/embedding/__init__.py b/vllm/models/deepseek_v41/nvidia/specialized_decode/embedding/__init__.py new file mode 100644 index 000000000000..208f01a7cb5e --- /dev/null +++ b/vllm/models/deepseek_v41/nvidia/specialized_decode/embedding/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/models/deepseek_v41/nvidia/specialized_decode/embedding/replicated.py b/vllm/models/deepseek_v41/nvidia/specialized_decode/embedding/replicated.py new file mode 100644 index 000000000000..6eb22506a035 --- /dev/null +++ b/vllm/models/deepseek_v41/nvidia/specialized_decode/embedding/replicated.py @@ -0,0 +1,247 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""A decode step's token embedding from a full copy of the table on every rank. + +Under tensor parallelism vLLM shards the token embedding table by vocabulary +(``VocabParallelEmbedding``): each rank looks up the tokens its shard holds, +writes zeros for the others, and an all-reduce sums the ranks' rows. The +DeepSeek-V4.1 backbone then copies the embedding into its ``hc_mult`` +hyper-connection streams ([n, H] -> [n, hc, H]) at the first layer. At a +few rows both are fixed costs on the critical path of every step. + +``ReplicatedEmbedding`` keeps a full copy of the table on every rank: the +ranks broadcast their shards into it once, when the decode executor is built +(1.32 GB a GPU for DeepSeek-V4.1-Flash), which is before the KV cache is +sized, so the memory profile counts it. A captured decode step then looks its +tokens up in one launch of ``StreamEmbedding`` that writes the embedding (for +the target's first layer, whose pre-mix reads it) and its streams, and needs +no collective. A token outside the vocabulary gets zeros, as the sharded +lookup's sum gives it, so the result equals the all-reduced embedding (up to +the sign of a zero, which the sum makes positive). + +The copy serves the table as it was when the executor was built. Loading +weights again rewrites the shards in place but not the copy, so a deployment +that updates weights at runtime must not use the specialized decode. +""" + +from functools import cache + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import torch +from cutlass import Int32 +from cutlass.cute.runtime import make_fake_stream, make_ptr +from torch import nn + +from vllm.distributed import get_tp_group +from vllm.model_executor.layers.vocab_parallel_embedding import ( + UnquantizedEmbeddingMethod, + VocabParallelEmbedding, +) +from vllm.models.common.ops.moe_tail_fusion.primitives import current_cu_stream + +THREADS = 256 +VECTOR_BYTES = 16 +MAX_CTAS = 1024 + + +class StreamEmbedding: + """Rows of a ``[vocab, H]`` BF16 table into an ``[n, H]`` embedding and + its ``[n, hc_mult, H]`` streams, by 16-byte vectors; a token outside + ``[0, vocab)`` gets zeros.""" + + def __init__(self, hidden: int, hc_mult: int) -> None: + assert hidden * 2 % VECTOR_BYTES == 0 + self.row_vectors = hidden * 2 // VECTOR_BYTES + self.hc_mult = hc_mult + + @cute.jit + def __call__( + self, + ids_ptr: cute.Pointer, + table_ptr: cute.Pointer, + embedding_ptr: cute.Pointer, + streams_ptr: cute.Pointer, + rows: Int32, + vocab: Int32, + num_ctas: Int32, + stream: cuda.CUstream, + ): + """Every tensor as int32 words; a row's vector ``v`` is the column + ``v`` of four words.""" + row_vectors = self.row_vectors + self.kernel( + cute.make_tensor(ids_ptr, cute.make_layout((rows,))), + cute.make_tensor(table_ptr, cute.make_layout((4, row_vectors, vocab))), + cute.make_tensor(embedding_ptr, cute.make_layout((4, row_vectors, rows))), + cute.make_tensor( + streams_ptr, cute.make_layout((4, row_vectors, self.hc_mult, rows)) + ), + rows, + vocab, + num_ctas, + ).launch(grid=(num_ctas, 1, 1), block=[THREADS, 1, 1], stream=stream) + + @cute.kernel + def kernel( + self, + ids: cute.Tensor, + table: cute.Tensor, + embedding: cute.Tensor, + streams: cute.Tensor, + rows: Int32, + vocab: Int32, + num_ctas: Int32, + ): + """Thread t of CTA c takes the (row, vector) items c * THREADS + t, + then every num_ctas * THREADS after it.""" + thread, _, _ = cute.arch.thread_idx() + cta, _, _ = cute.arch.block_idx() + vector = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), Int32, num_bits_per_copy=128 + ) + words = cute.make_rmem_tensor((4,), Int32) + row_vectors = self.row_vectors + for item in cutlass.range( + cta * THREADS + thread, rows * row_vectors, num_ctas * THREADS, unroll=1 + ): + row = item // row_vectors + v = item - row * row_vectors + token = ids[row] + words.fill(0) + if token >= 0: # noqa: SIM102 + if token < vocab: + cute.copy(vector, table[None, v, token], words) + cute.copy(vector, words, embedding[None, v, row]) + for c in cutlass.range_constexpr(self.hc_mult): + cute.copy(vector, words, streams[None, v, c, row]) + + +def _pointer(address: int, align: int = VECTOR_BYTES): + """A CuTe DSL global-memory pointer to int32 words.""" + return make_ptr(Int32, address, cute.AddressSpace.gmem, assumed_align=align) + + +@cache +def _kernel(device_index: int, hidden: int, hc_mult: int): + """``StreamEmbedding`` compiled for a device.""" + with torch.accelerator.device_index(device_index): + return cute.compile( + StreamEmbedding(hidden, hc_mult), + _pointer(16, 4), + _pointer(16), + _pointer(16), + _pointer(16), + Int32(1), + Int32(1), + Int32(1), + make_fake_stream(), + ) + + +def embed_streams( + ids: torch.Tensor, + table: torch.Tensor, + embedding: torch.Tensor, + streams: torch.Tensor, +) -> None: + """Write ``table``'s rows of ``ids`` (int32 ``[n]``) into ``embedding`` + (``[n, H]``) and ``streams`` (``[n, hc, H]``), on the current stream; + zeros for an id outside the table.""" + rows, hc_mult, hidden = streams.shape + assert ids.dtype == torch.int32 and ids.shape == (rows,) and ids.is_contiguous() + assert table.dtype == torch.bfloat16 and table.shape[1] == hidden + assert table.is_contiguous() and streams.is_contiguous() + assert streams.dtype == torch.bfloat16 + assert embedding.shape == (rows, hidden) and embedding.is_contiguous() + assert embedding.dtype == torch.bfloat16 + if rows == 0: + return + items = rows * (hidden * 2 // VECTOR_BYTES) + num_ctas = min((items + THREADS - 1) // THREADS, MAX_CTAS) + _kernel(table.device.index, hidden, hc_mult)( + _pointer(ids.data_ptr(), 4), + _pointer(table.data_ptr()), + _pointer(embedding.data_ptr()), + _pointer(streams.data_ptr()), + Int32(rows), + Int32(table.shape[0]), + Int32(num_ctas), + current_cu_stream(), + ) + + +class ReplicatedEmbedding: + """A full copy of a vocabulary-parallel embedding table on every TP rank, + and a decode step's lookup from it.""" + + @classmethod + def create(cls, embed: nn.Module, hc_mult: int) -> "ReplicatedEmbedding | None": + """The copy of ``embed``'s table, or None when ``embed`` is not the + plain sharded BF16 lookup this replaces. The decision reads only the + configuration, which every TP rank shares; building the copy is + collective.""" + if not isinstance(embed, VocabParallelEmbedding): + return None + weight = embed.weight + indices = embed.shard_indices + if ( + embed.tp_size == 1 + or not embed.reduce_results + or embed.parallel_group is not None + or not isinstance(embed.quant_method, UnquantizedEmbeddingMethod) + or weight.dtype != torch.bfloat16 + or weight.dim() != 2 + or not weight.is_contiguous() + or weight.shape[1] * 2 % VECTOR_BYTES + or indices.added_vocab_end_index != indices.added_vocab_start_index + ): + return None + return cls(embed, hc_mult) + + def __init__(self, embed: VocabParallelEmbedding, hc_mult: int) -> None: + group = get_tp_group() + weight = embed.weight.data + indices = embed.shard_indices + bounds = torch.tensor( + [indices.org_vocab_start_index, indices.org_vocab_end_index], + dtype=torch.int64, + device=weight.device, + ) + ranges = group.all_gather(bounds, dim=0).view(group.world_size, 2).tolist() + vocab = embed.org_vocab_size + covered = sorted((start, end) for start, end in ranges if end > start) + assert ( + covered[0][0] == 0 + and covered[-1][1] == vocab + and all(a[1] == b[0] for a, b in zip(covered, covered[1:])) + ), f"vocabulary shards {ranges} do not tile [0, {vocab})" + self.table = weight.new_empty(vocab, weight.shape[1]) + for source, (start, end) in enumerate(ranges): + rows = self.table[start:end] + if source == group.rank_in_group: + rows.copy_(weight[: end - start]) + group.broadcast(rows, src=source) + self.hc_mult = hc_mult + # Compile and load the kernel now: steps use it only in captures. + self(torch.zeros(1, dtype=torch.int32, device=weight.device)) + torch.accelerator.synchronize(weight.device) + + def accepts(self, ids: torch.Tensor) -> bool: + """Whether ``__call__`` takes ``ids``: a non-empty int32 vector.""" + return ( + ids.dtype == torch.int32 + and ids.dim() == 1 + and ids.shape[0] > 0 + and ids.is_contiguous() + ) + + def __call__(self, ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """The embedding ``[n, H]`` and its streams ``[n, hc, H]`` of the + tokens ``ids``.""" + rows, hidden = ids.shape[0], self.table.shape[1] + streams = self.table.new_empty(rows, self.hc_mult, hidden) + embedding = self.table.new_empty(rows, hidden) + embed_streams(ids, self.table, embedding, streams) + return embedding, streams diff --git a/vllm/models/deepseek_v41/nvidia/specialized_decode/engram/prefetch.py b/vllm/models/deepseek_v41/nvidia/specialized_decode/engram/prefetch.py index 5d77a70b74c2..e202f7ca1af2 100644 --- a/vllm/models/deepseek_v41/nvidia/specialized_decode/engram/prefetch.py +++ b/vllm/models/deepseek_v41/nvidia/specialized_decode/engram/prefetch.py @@ -5,8 +5,8 @@ With host-offloaded tables, ``Engram`` looks up its rows on a prefetch stream before the decoder layers. The rest of the WKV chain (head all-gather, MXFP8 quantize, column-parallel GEMM, output all-gather) reads only those rows, so -for small decode steps ``EngramPrefetch`` runs it on a side stream right after -each layer's lookup, and the layer only waits for it before the fused post +for decode steps ``EngramPrefetch`` runs it on a side stream right after each +layer's lookup, and the layer only waits for it before the fused post kernel. The kernels, inputs and collectives are those of ``Engram.forward``. The chain issues TP collectives on the side stream, so whether a step starts it diff --git a/vllm/models/deepseek_v41/nvidia/specialized_decode/logits/logits.py b/vllm/models/deepseek_v41/nvidia/specialized_decode/logits/logits.py index 811e04bb6f07..e62b3198e139 100644 --- a/vllm/models/deepseek_v41/nvidia/specialized_decode/logits/logits.py +++ b/vllm/models/deepseek_v41/nvidia/specialized_decode/logits/logits.py @@ -1,30 +1,39 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Logits of the specialized decode: the low-batch LM head GEMM (``lm_head``) -and vocab-parallel logits gathered with the TP group's custom all-gather. - -For decode batches the one-shot custom all-gather is several times faster than -the NCCL ring. The kernel falls back to NCCL per call by dtype and size alone, -so every TP rank takes the same path. It shares its staging buffers and -signals with the group's other custom collectives, so the logits must be -gathered on the forward stream after any side stream has rejoined it. +and vocab-parallel logits gathered with the executor's all-gather. + +For decode batches a one-shot all-gather over NVLink multicast is several +times faster than the NCCL ring. The executor's ``LamportAllGather`` takes the +logits when they fit, by size alone, so every TP rank takes the same path; +larger ones go to NCCL. vLLM's custom all-gather is not used: its MNNVL +Lamport kernel never waits for a rank's own multicast copy (see +``ops.cute_dsl.all_gather``). The all-gather's mailbox also serves the +row-sharded indexers, so the logits must be gathered on the forward stream +after any side stream has rejoined it. """ +from typing import TYPE_CHECKING + import torch -from vllm.distributed import get_tp_group -from vllm.distributed.device_communicators.cuda_communicator import CudaCommunicator from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding from vllm.platforms import current_platform from .lm_head import low_batch_lm_head, low_batch_lm_head_applies +if TYPE_CHECKING: + from ...ops.cute_dsl import LamportAllGather + class SpecializedDecodeLogitsProcessor(LogitsProcessor): - def __init__(self, vocab_size: int) -> None: + def __init__( + self, vocab_size: int, all_gather: "LamportAllGather | None" = None + ) -> None: super().__init__(vocab_size) self.pdl = current_platform.is_arch_support_pdl() + self.all_gather = all_gather def _apply_head( self, @@ -41,10 +50,12 @@ def _apply_head( return super()._apply_head(lm_head, hidden_states, embedding_bias) def _gather_logits(self, logits: torch.Tensor) -> torch.Tensor: - communicator = get_tp_group().device_communicator - if self.use_all_gather and isinstance(communicator, CudaCommunicator): - gathered = communicator.custom_all_gather(logits) - if gathered is not None: - # Rank-major rows to rank-major vocab shards, as NCCL lays out. - return gathered.view(-1, *logits.shape).movedim(0, -2).flatten(-2) + if ( + self.use_all_gather + and self.all_gather is not None + and self.all_gather.accepts(logits) + ): + gathered = self.all_gather(logits) + # Rank-major rows to rank-major vocab shards, as NCCL lays out. + return gathered.view(-1, *logits.shape).movedim(0, -2).flatten(-2) return super()._gather_logits(logits) diff --git a/vllm/models/deepseek_v41/nvidia/specialized_decode/moe/routed_expert_front_cute.py b/vllm/models/deepseek_v41/nvidia/specialized_decode/moe/routed_expert_front_cute.py index 871b3032ddef..cdb1285de994 100644 --- a/vllm/models/deepseek_v41/nvidia/specialized_decode/moe/routed_expert_front_cute.py +++ b/vllm/models/deepseek_v41/nvidia/specialized_decode/moe/routed_expert_front_cute.py @@ -34,7 +34,9 @@ log(1 + exp(x))), ``biased = score + bias`` (``bias_vl`` for image sentinel tokens), the six largest biased scores (lowest expert id on exact ties, NaN ranked as -1e30), and ``weights`` = the unbiased scores times scale / sum - (sum replaced by 1 unless positive); + (sum replaced by 1 unless positive); a CUDA graph padding row + (``VLLM_MOE_SKIP_PADDING``'s mask) gets ids -1 and weights 0, so no expert + computes it; 2. the routed input's MXFP8 quantization (FlashInfer's linear-layout ``mxfp8_quantize``): per 32-value block, a UE8M0 scale rounded up from amax / 448 and E4M3 values rounded to nearest with saturation; @@ -51,15 +53,25 @@ floating-point operation is left to the compiler, which could contract or flush it. -Work split (grid ``1 + ceil(rows / 3)``, 1024 threads): - -* CTA 0: four warps per row each take the top-6 of 96 experts (three per - lane); one warp per row merges the 24 candidates and normalizes; then the - whole CTA builds the routing record from the selected ids. The selection is - a lexicographic (value, lowest id) maximum, so splitting it does not change - the result. -* CTA ``1 + c``: quantizes rows ``3c .. 3c + 2`` of the BF16 input, two - threads per 32-value block. +Work split (1024 threads a CTA). The kernel is built for calls of up to 8, +16, 32, 64, 128 or 192 rows (``max_rows``); the first ``S = 8`` CTAs select +and form one thread block cluster, the rest quantize: + +* Selecting CTA ``s`` takes rows ``s``, ``s + S``, ``s + 2 S``, ... (at most + ``max_rows / S`` rows: 1, 2, 4 or 8 up to 64 rows, 16 at 128, 24 at 192). + Four warps per row (two at 16 rows a CTA, one at 24) each take the top-6 of + their share of the experts (3, 6 or 12 per lane); one warp per row merges + the candidates and normalizes. The top-k is bound by instruction issue, + so its time grows with the rows a CTA holds; the rows are therefore spread + over all eight CTAs, even at 8 rows or fewer. The + selection is a lexicographic (value, lowest id) maximum, so splitting it + does not change the result. The merging warps store the selected ids to + CTA 0's shared memory (through distributed shared memory) and release + them with the cluster barrier, then CTA 0 builds the routing record from + them. The ids and weights go to global memory only after that release, + which would otherwise wait for those stores to reach L2. +* Quantizing CTA ``S + c``: rows ``3c .. 3c + 2`` of the BF16 input, two + threads per 32-value block. The grid is rounded up to whole clusters. PDL: every CTA triggers its dependents on entry (GEMM1's set-up overlaps this kernel; GEMM1 reads the outputs only after its own grid-dependency wait). The @@ -78,9 +90,11 @@ from .routed_expert_route_cute import ( MAX_EXPERTS, - WORDS, build_record, clear_scratch, + levels_for, + mask_words_for, + present_words, record_words, zero_record, ) @@ -88,19 +102,16 @@ HIDDEN = 5120 EXPERTS = 384 TOP_K = 6 -MAX_ROWS = 8 -PARTS = 4 # warps per row -PER_LANE = EXPERTS // 32 // PARTS # three experts per lane -CANDIDATES = PARTS * TOP_K # 24 per row, merged by one warp BLOCKS = HIDDEN // 32 # quantization blocks per row FRONT_THREADS = 1024 ROW_THREADS = 2 * BLOCKS # 320 threads quantize one row QUANT_ROWS = FRONT_THREADS // ROW_THREADS # rows per quantizing CTA +SELECT_CTAS = 8 # the selecting CTAs: a portable cluster +MAX_PARTS = 4 # the most warps per row: 4 x 6 candidates fit one merging warp NAN_KEY_VALUE = 0xF149F2CA # -1e30, Triton's rank of a NaN score -assert FRONT_THREADS == MAX_ROWS * PARTS * 32 -assert CANDIDATES <= 32 assert EXPERTS <= MAX_EXPERTS +assert EXPERTS % 32 == 0 def _asm(result, operands: list, text: str, constraints: str, loc, ip): @@ -213,7 +224,7 @@ def lane_best(key: cute.Tensor): """The lane's best remaining entry: the largest key, then the lowest j.""" best_key = Uint32(key[0]) best_j = Int32(0) - for j in cutlass.range_constexpr(1, PER_LANE): + for j in cutlass.range_constexpr(1, cute.size(key)): if Uint32(key[j]) > best_key: best_key = Uint32(key[j]) best_j = Int32(j) @@ -222,15 +233,17 @@ def lane_best(key: cute.Tensor): @cute.jit def part_top_k(logit: cute.Tensor, bias: cute.Tensor, base: Int32, lane: Int32): - """One warp's top-6 of its 96 experts (expert = base + lane + 32 j), in + """One warp's top-6 of its share of the experts (expert = base + lane + + 32 j, ``j`` below the size of ``logit``), in selection order: each round takes the warp's largest key with one ``redux.sync`` max, then the lowest expert id holding it with one min (Triton's ``tl.max``, then the lowest index equal to it); only the winning lane rescans. Lane ``s < 6`` gets slot ``s``: its key, id and unbiased score.""" - score = cute.make_rmem_tensor((PER_LANE,), Float32) - key = cute.make_rmem_tensor((PER_LANE,), Uint32) - for j in cutlass.range_constexpr(PER_LANE): + per_lane = cute.size(logit) + score = cute.make_rmem_tensor((per_lane,), Float32) + key = cute.make_rmem_tensor((per_lane,), Uint32) + for j in cutlass.range_constexpr(per_lane): score[j], key[j] = score_and_key(logit[j], bias[j]) best_key, best_j = lane_best(key) out_key = Uint32(0) @@ -244,7 +257,7 @@ def part_top_k(logit: cute.Tensor, bias: cute.Tensor, base: Int32, lane: Int32): winner = Int32(cute.arch.warp_redux_sync(mine, "umin")) won = Float32(0.0) if winner - base == lane + 32 * best_j: - for j in cutlass.range_constexpr(PER_LANE): + for j in cutlass.range_constexpr(per_lane): if best_j == j: won = score[j] key[j] = Uint32(0) # below every value's key @@ -263,15 +276,13 @@ def merge_row( expert: Int32, score: Float32, scale: Float32, - row: Int32, lane: Int32, - topk_ids: cute.Tensor, - topk_weights: cute.Tensor, - ids_shared: cute.Tensor, + padding: Int32, ): - """Merge a row's 24 candidates (lane ``c < 24`` holds candidate ``c``) into - its top-6, normalize as Triton does and store; lane ``s < 6`` ends with - slot ``s``.""" + """Merge a row's candidates (lane ``c`` holds candidate ``c``) into + its top-6 and normalize as Triton does; lane ``s < 6`` returns slot + ``s``'s expert id and weight. A padding row (``padding`` nonzero) gets + id -1 and weight 0, as the Triton kernel's padding mask does.""" slot_weight = Float32(0.0) # lane s < 8 keeps slot s (Triton's 8-wide vector) slot_id = Int32(0) for slot in cutlass.range_constexpr(TOP_K): @@ -293,10 +304,34 @@ def merge_row( total = slot_weight for shift in cutlass.range_constexpr(2, -1, -1): total = add_f32(total, cute.arch.shuffle_sync_bfly(total, 1 << shift)) + weight = Float32(0.0) if lane < TOP_K: - topk_ids[row * TOP_K + lane] = slot_id - topk_weights[row * TOP_K + lane] = normalized_weight(scale, total, slot_weight) - ids_shared[row * TOP_K + lane] = slot_id + weight = normalized_weight(scale, total, slot_weight) + if padding != 0: + slot_id = Int32(-1) + weight = Float32(0.0) + return slot_id, weight + + +@dsl_user_op +def store_cta0_shared(pointer: cute.Pointer, value, *, loc=None, ip=None): + """Store the 32-bit ``value`` at ``pointer``'s shared-memory offset in CTA 0 + of the cluster (distributed shared memory).""" + llvm.inline_asm( + None, + [ + Int32(pointer.toint(loc=loc, ip=ip)).ir_value(loc=loc, ip=ip), + Int32(value).ir_value(loc=loc, ip=ip), + ], + "{\n.reg .u32 a;\nmapa.shared::cluster.u32 a, $0, 0;\n" + "st.shared::cluster.u32 [a], $1;\n}", + "r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) # ---- MXFP8 quantization (FlashInfer's linear layout, two threads per block) @@ -474,10 +509,21 @@ def quantize_row( class RoutedExpertFront: """The fused front (module docstring); ``image``: whether the router - routes image sentinel tokens with ``bias_vl``.""" + routes image sentinel tokens with ``bias_vl``; ``max_rows``: 8, 16, 32, + 64, 128 or 192, the most rows of a call.""" - def __init__(self, image: bool) -> None: + def __init__(self, image: bool, max_rows: int = 8) -> None: + assert max_rows in (8, 16, 32, 64, 128, 192) self.image = image + self.max_rows = max_rows + self.selecting = SELECT_CTAS # CTAs, one cluster + self.select_rows = max_rows // self.selecting # 1-8, 16 at 128 rows, 24 at 192 + # Warps per row: four, or as many as fit above 8 rows a CTA. + self.parts = min(FRONT_THREADS // 32 // self.select_rows, MAX_PARTS) + self.per_lane = EXPERTS // 32 // self.parts # experts per lane + self.candidates = self.parts * TOP_K # per row, merged by one warp + self.levels = levels_for(max_rows) + assert EXPERTS % (32 * self.parts) == 0 and self.candidates <= 32 @cute.jit def __call__( @@ -489,6 +535,8 @@ def __call__( ids_are_64bit: Int32, sentinel_lo: Int32, scale: Float32, + is_padding_ptr: cute.Pointer, + has_padding: Int32, x_ptr: cute.Pointer, topk_ids_ptr: cute.Pointer, topk_weights_ptr: cute.Pointer, @@ -502,11 +550,16 @@ def __call__( """``logits``: FP32 [rows, 384]; ``bias``, ``bias_vl``: FP32 [384]; ``input_ids``: [rows] int32, or int64 if ``ids_are_64bit``, as Int32 words (``bias_vl`` and ``input_ids`` are read only with image routing; - pass any pointers otherwise); ``x``: BF16 [rows, 5120] as Uint32 + pass any pointers otherwise); ``is_padding``: bool [rows], read only + if ``has_padding``; ``x``: BF16 [rows, 5120] as Uint32 words; outputs: ``topk_ids`` int32 / ``topk_weights`` FP32 [rows, 6], ``x_quant`` [rows, 5120] E4M3 bytes, ``x_scale`` [rows, 160] UE8M0, - ``permuted`` [rows, 6] and the routing record. rows in [1, 8].""" + ``permuted`` [rows, 6] and the routing record. rows in [1, max_rows].""" slots = cute.make_layout(rows * TOP_K) + selecting = self.selecting + quantizing = (rows + QUANT_ROWS - 1) // QUANT_ROWS + # Whole clusters: the last one's spare CTAs exit at once. + quantizing = (quantizing + selecting - 1) // selecting * selecting self.kernel( cute.make_tensor(logits_ptr, cute.make_layout(rows * EXPERTS)), cute.make_tensor(bias_ptr, cute.make_layout(EXPERTS)), @@ -515,6 +568,8 @@ def __call__( ids_are_64bit, sentinel_lo, scale, + cute.make_tensor(is_padding_ptr, cute.make_layout(rows)), + has_padding, cute.make_tensor(x_ptr, cute.make_layout(rows * HIDDEN // 2)), cute.make_tensor(topk_ids_ptr, slots), cute.make_tensor(topk_weights_ptr, slots), @@ -526,10 +581,11 @@ def __call__( ), rows, ).launch( - grid=(1 + (rows + QUANT_ROWS - 1) // QUANT_ROWS, 1, 1), + grid=(selecting + quantizing, 1, 1), block=[FRONT_THREADS, 1, 1], stream=stream, use_pdl=True, + cluster=[selecting, 1, 1], ) @cute.kernel @@ -542,6 +598,8 @@ def kernel( ids_are_64bit: Int32, sentinel_lo: Int32, scale: Float32, + is_padding: cute.Tensor, + has_padding: Int32, x: cute.Tensor, topk_ids: cute.Tensor, topk_weights: cute.Tensor, @@ -551,12 +609,13 @@ def kernel( record: cute.Tensor, rows: Int32, ): - """CTA 0 selects and routes; the others quantize (module docstring).""" + """The first ``selecting`` CTAs select (CTA 0 also routes); the others + quantize (module docstring).""" cute.arch.griddepcontrol_launch_dependents() thread, _, _ = cute.arch.thread_idx() cta, _, _ = cute.arch.block_idx() - if cta != 0: - row = (cta - 1) * QUANT_ROWS + thread // ROW_THREADS + if cta >= self.selecting: + row = (cta - self.selecting) * QUANT_ROWS + thread // ROW_THREADS if thread < QUANT_ROWS * ROW_THREADS and row < rows: cute.arch.griddepcontrol_wait() quantize_row(x, row, thread % ROW_THREADS, x_quant, x_scale) @@ -569,11 +628,14 @@ def kernel( ids_are_64bit, sentinel_lo, scale, + is_padding, + has_padding, topk_ids, topk_weights, permuted, record, rows, + cta, thread, ) @@ -587,53 +649,69 @@ def select_and_route( ids_are_64bit: Int32, sentinel_lo: Int32, scale: Float32, + is_padding: cute.Tensor, + has_padding: Int32, topk_ids: cute.Tensor, topk_weights: cute.Tensor, permuted: cute.Tensor, record: cute.Tensor, rows: Int32, + rank: Int32, thread: Int32, ): - """CTA 0: top-k of every row, then the routing record.""" + """Selecting CTA ``rank``: top-k of its rows; then CTA 0 builds the + routing record from every row's ids.""" + selecting, levels = self.selecting, self.levels + bitmap_words = present_words(levels) + mask_size = MAX_EXPERTS * mask_words_for(levels) + selected = self.max_rows * TOP_K + parts, per_lane, candidates = self.parts, self.per_lane, self.candidates + held = self.select_rows * candidates @cute.struct class Shared: - present: cute.struct.MemRange[Uint32, WORDS] - rows_of: cute.struct.MemRange[Uint32, MAX_EXPERTS] - ids: cute.struct.MemRange[Int32, MAX_ROWS * TOP_K] - candidate_key: cute.struct.MemRange[Uint32, MAX_ROWS * CANDIDATES] - candidate_id: cute.struct.MemRange[Int32, MAX_ROWS * CANDIDATES] - candidate_score: cute.struct.MemRange[Float32, MAX_ROWS * CANDIDATES] + present: cute.struct.MemRange[Uint32, bitmap_words] + rows_of: cute.struct.MemRange[Uint32, mask_size] + ids: cute.struct.MemRange[Int32, selected] + candidate_key: cute.struct.MemRange[Uint32, held] + candidate_id: cute.struct.MemRange[Int32, held] + candidate_score: cute.struct.MemRange[Float32, held] shared = cutlass.utils.SmemAllocator().allocate(Shared) def view(member, size): return cute.make_tensor(member.data_ptr(), cute.make_layout(size)) - present = view(shared.present, WORDS) - rows_of = view(shared.rows_of, MAX_EXPERTS) - ids_shared = view(shared.ids, MAX_ROWS * TOP_K) - candidate_key = view(shared.candidate_key, MAX_ROWS * CANDIDATES) - candidate_id = view(shared.candidate_id, MAX_ROWS * CANDIDATES) - candidate_score = view(shared.candidate_score, MAX_ROWS * CANDIDATES) - clear_scratch(present, rows_of, thread, FRONT_THREADS) + present = view(shared.present, bitmap_words) + rows_of = view(shared.rows_of, mask_size) + ids_shared = view(shared.ids, selected) + candidate_key = view(shared.candidate_key, held) + candidate_id = view(shared.candidate_id, held) + candidate_score = view(shared.candidate_score, held) + # Every CTA of the cluster must have started before the stores to CTA + # 0's shared memory; this phase completes during the top-k. + cute.arch.cluster_arrive_relaxed() + if rank == 0: + clear_scratch(present, rows_of, thread, FRONT_THREADS, levels) warp = thread >> 5 lane = thread & 31 - row = warp // PARTS - part = warp % PARTS - base = part * 32 * PER_LANE # this warp's experts: base + lane + 32 j - row_bias = cute.make_rmem_tensor((PER_LANE,), Float32) - row_bias_vl = cute.make_rmem_tensor((PER_LANE,), Float32) + local_row = warp // parts + part = warp % parts + row = rank + selecting * local_row + base = part * 32 * per_lane # this warp's experts: base + lane + 32 j + row_bias = cute.make_rmem_tensor((per_lane,), Float32) + row_bias_vl = cute.make_rmem_tensor((per_lane,), Float32) if row < rows: - for j in cutlass.range_constexpr(PER_LANE): + for j in cutlass.range_constexpr(per_lane): row_bias[j] = bias[base + lane + 32 * j] if cutlass.const_expr(self.image): row_bias_vl[j] = bias_vl[base + lane + 32 * j] cute.arch.griddepcontrol_wait() - zero_record(record, rows, TOP_K, EXPERTS, thread, FRONT_THREADS) + if rank == 0: + zero_record(record, rows, TOP_K, EXPERTS, thread, FRONT_THREADS) if row < rows: - logit = cute.make_rmem_tensor((PER_LANE,), Float32) - for j in cutlass.range_constexpr(PER_LANE): + logit = cute.make_rmem_tensor((per_lane,), Float32) + for j in cutlass.range_constexpr(per_lane): logit[j] = logits[row * EXPERTS + base + lane + 32 * j] if cutlass.const_expr(self.image): token = Int64(input_ids[row]) @@ -642,35 +720,51 @@ def view(member, size): token = (Int64(input_ids[2 * row + 1]) << 32) | low_word low = Int64(sentinel_lo) if token >= low and token < low + 5: - for j in cutlass.range_constexpr(PER_LANE): + for j in cutlass.range_constexpr(per_lane): row_bias[j] = row_bias_vl[j] key, expert, score = part_top_k(logit, row_bias, base, lane) if lane < TOP_K: - slot = row * CANDIDATES + part * TOP_K + lane + slot = local_row * candidates + part * TOP_K + lane candidate_key[slot] = key candidate_id[slot] = expert candidate_score[slot] = score cute.arch.sync_threads() + cute.arch.cluster_wait() + slot_id = Int32(-1) + slot_weight = Float32(0.0) if row < rows and part == 0: key = Uint32(0) expert = Int32(-1) score = Float32(0.0) - if lane < CANDIDATES: - key = Uint32(candidate_key[row * CANDIDATES + lane]) - expert = candidate_id[row * CANDIDATES + lane] - score = candidate_score[row * CANDIDATES + lane] - merge_row( - key, expert, score, scale, row, lane, topk_ids, topk_weights, ids_shared + if lane < candidates: + key = Uint32(candidate_key[local_row * candidates + lane]) + expert = candidate_id[local_row * candidates + lane] + score = candidate_score[local_row * candidates + lane] + padding = Int32(0) + if has_padding != 0: + padding = Int32(is_padding[row]) + slot_id, slot_weight = merge_row(key, expert, score, scale, lane, padding) + if lane < TOP_K: + store_cta0_shared(ids_shared.iterator + row * TOP_K + lane, slot_id) + # Release the stores to CTA 0; CTA 0 acquires them. The release orders + # every earlier store of the thread, so the global stores follow it + # rather than delay it by their round trip to L2 (about 0.2-0.4 us). + cute.arch.cluster_arrive() + if row < rows and part == 0 and lane < TOP_K: + topk_ids[row * TOP_K + lane] = slot_id + topk_weights[row * TOP_K + lane] = slot_weight + cute.arch.cluster_wait() + if rank == 0: + build_record( + present, + rows_of, + ids_shared, + permuted, + record, + rows, + Int32(TOP_K), + Int32(EXPERTS), + thread, + FRONT_THREADS, + levels, ) - cute.arch.sync_threads() - build_record( - present, - rows_of, - ids_shared, - permuted, - record, - rows, - Int32(TOP_K), - Int32(EXPERTS), - thread, - ) 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 590fe2675a2a..59b3aa616f1e 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 @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Low-batch routed-expert GEMMs for DeepSeek-V4.1 decode (1-8 rows). +"""Low-batch routed-expert GEMMs for DeepSeek-V4.1 decode (1-192 rows). At concurrency 1 with DSpark, every target MoE layer applies its routed experts to six rows (384 experts, top-6; at TP4 every rank holds all experts at @@ -8,7 +8,9 @@ TRT-LLM-gen routing kernel and two batched MXFP8 x MXFP4 GEMMs; at this size they stream the weights of about 22 experts per layer, so their cost is HBM time plus fixed per-kernel latency. This module computes the same values from the same -tensors with CuTe DSL kernels sized for at most eight rows:: +tensors with CuTe DSL kernels sized for at most 192 rows (8-row token tiles; +the route kernel and the front are built for calls of up to 8, 16, 32, 64, 128 +or 192 rows, and GEMM1 gathers each tile's rows):: router logits ─► front ───────────────────────┐ (top-k, MXFP8 input, routing) top-k ids, MXFP8 input ─► route ──────────────┤ @@ -53,6 +55,8 @@ from cutlass import Float32, Int32 from cutlass.cute.runtime import make_fake_stream, make_ptr +import vllm.envs as envs +from vllm.forward_context import get_forward_context, is_forward_context_available from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput from vllm.models.common.ops.moe_tail_fusion.primitives import current_cu_stream @@ -63,7 +67,7 @@ RoutedExpertGemm2, tile_scale_bytes, ) -from .routed_expert_route_cute import RoutedExpertRoute +from .routed_expert_route_cute import RoutedExpertRoute, levels_for if TYPE_CHECKING: from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig @@ -84,7 +88,25 @@ MAX_EXPERTS = 384 MAX_TOP_K = 8 TILE = 8 -MAX_ROWS = 8 +MAX_ROWS = 192 +# The row counts the route kernel and the front are built for. +WIDTHS = tuple(w for w in (8, 16, 32, 64, 128, 192) if w <= MAX_ROWS) + + +def width(rows: int) -> int: + """The kernels' build for calls of ``rows`` rows: the narrowest that + fits.""" + return next(w for w in WIDTHS if rows <= w) + + +def padding_mask(rows: int) -> torch.Tensor | None: + """The CUDA graph padding rows of this forward's first ``rows`` (bool), + the mask the stock router applies with ``VLLM_MOE_SKIP_PADDING``; None + without one.""" + if envs.VLLM_MOE_SKIP_PADDING and is_forward_context_available(): + is_padding = get_forward_context().is_padding + return is_padding[:rows] if is_padding is not None else None + return None def max_ctas(rows: int, top_k: int, experts: int) -> int: @@ -100,8 +122,9 @@ def cta_words(ctas: int) -> int: def route_words(ctas: int) -> int: - """int32 length of the routing record for ``ctas`` token tiles.""" - return 4 + 2 * cta_words(ctas) + TILE * ctas + """int32 length of the routing record for ``ctas`` token tiles (its + layout is in ``routed_expert_route_cute.py``).""" + return 4 + 3 * cta_words(ctas) + TILE * ctas @dataclass @@ -150,7 +173,8 @@ def __init__( torch.accelerator.current_device_index() if index is None else index ) # Compile the kernels now, not inside a CUDA graph capture. - _route_kernel(self.device_index) + for built in WIDTHS: + _route_kernel(self.device_index, built) _gemm1_kernel(self.device_index) _gemm2_kernel(self.device_index) # The modular kernel's input quantization, reused so the activations @@ -163,13 +187,14 @@ def __init__( def set_front(self, config: FrontConfig) -> None: """Use the fused front with the router's ``config``.""" - _front_kernel(self.device_index, config.bias_vl is not None) + for built in WIDTHS: + _front_kernel(self.device_index, config.bias_vl is not None, built) self.front_config = config def accepts(self, x: torch.Tensor, runner: torch.nn.Module) -> bool: - """Whether a call with router input ``x`` can take these kernels: 1-8 - BF16 rows, and the runner's expert weights still the tensors bound here - (a weight reload replaces them).""" + """Whether a call with router input ``x`` can take these kernels: 1 to + ``MAX_ROWS`` BF16 rows, and the runner's expert weights still the + tensors bound here (a weight reload replaces them).""" return ( 1 <= x.size(0) <= MAX_ROWS and x.dtype == torch.bfloat16 @@ -203,7 +228,7 @@ def route( ctas = max_ctas(rows, top_k, self.experts) permuted = torch.empty((rows, top_k), dtype=torch.int32, device=x.device) record = torch.empty(route_words(ctas), dtype=torch.int32, device=x.device) - _route_kernel(self.device_index)( + _route_kernel(self.device_index, width(rows))( *(_pointer(cutlass.Int32, t, 4) for t in (ids, permuted, record)), Int32(rows), Int32(top_k), @@ -213,13 +238,18 @@ def route( return Routing(x_quant, x_scale, ids, topk_weights, permuted, record, ctas) def front( - self, x: torch.Tensor, logits: torch.Tensor, input_ids: torch.Tensor | None + self, + x: torch.Tensor, + logits: torch.Tensor, + input_ids: torch.Tensor | None, + is_padding: torch.Tensor | None = None, ) -> Routing: """Top-k, MXFP8 quantization and routing record in one launch. ``x``: BF16 router input [rows, 5120]; ``logits``: FP32 router logits [rows, 384]; ``input_ids``: [rows] int32 or int64, read only with image - routing. + routing; ``is_padding``: bool [rows], rows no expert computes (ids -1, + weights 0), as the router's padding mask (:func:`padding_mask`). """ config = self.front_config assert config is not None @@ -237,7 +267,7 @@ def front( if image: assert input_ids is not None ids = input_ids - _front_kernel(self.device_index, image)( + _front_kernel(self.device_index, image, width(rows))( _pointer(cutlass.Float32, logits, 16), _pointer(cutlass.Float32, config.bias, 16), _pointer(cutlass.Float32, config.bias_vl if image else config.bias, 16), @@ -245,6 +275,8 @@ def front( Int32(int(ids.dtype == torch.int64)), Int32(config.image_sentinel_lo), Float32(config.scale), + _pointer(cutlass.Uint8, topk_ids if is_padding is None else is_padding, 1), + Int32(is_padding is not None), _pointer(cutlass.Uint32, x, 16), _pointer(cutlass.Int32, topk_ids, 16), _pointer(cutlass.Float32, topk_weights, 16), @@ -294,7 +326,7 @@ def gemm2( """Launch GEMM2 on GEMM1's outputs.""" rows, top_k = routing.topk_ids.shape gemm2_rows = torch.empty( - (rows * top_k * TILE, HIDDEN), + (TILE * routing.ctas, HIDDEN), dtype=torch.bfloat16, device=routing.x_quant.device, ) @@ -331,11 +363,12 @@ def _fake(dtype: type[cutlass.Numeric], align: int): @cache -def _route_kernel(device_index: int): - """``RoutedExpertRoute`` compiled for a device.""" +def _route_kernel(device_index: int, built: int): + """``RoutedExpertRoute`` compiled for a device and calls of up to + ``built`` rows.""" with torch.accelerator.device_index(device_index): return cute.compile( - RoutedExpertRoute(), + RoutedExpertRoute(levels_for(built)), _fake(cutlass.Int32, 4), _fake(cutlass.Int32, 4), _fake(cutlass.Int32, 4), @@ -347,12 +380,12 @@ def _route_kernel(device_index: int): @cache -def _front_kernel(device_index: int, image: bool): +def _front_kernel(device_index: int, image: bool, built: int): """``RoutedExpertFront`` compiled for a device, with or without image - routing.""" + routing, for calls of up to ``built`` rows.""" with torch.accelerator.device_index(device_index): return cute.compile( - RoutedExpertFront(image), + RoutedExpertFront(image, built), _fake(cutlass.Float32, 16), _fake(cutlass.Float32, 16), _fake(cutlass.Float32, 16), @@ -360,6 +393,8 @@ def _front_kernel(device_index: int, image: bool): Int32(0), Int32(0), Float32(1.0), + _fake(cutlass.Uint8, 1), + Int32(0), _fake(cutlass.Uint32, 16), _fake(cutlass.Int32, 16), _fake(cutlass.Float32, 16), @@ -400,7 +435,7 @@ def _gemm2_kernel(device_index: int): """``RoutedExpertGemm2`` compiled for a device.""" with torch.accelerator.device_index(device_index): return cute.compile( - RoutedExpertGemm2(HIDDEN, PADDED), + RoutedExpertGemm2(HIDDEN, PADDED, 2 * SLICE // 128), _fake(cutlass.Float4E2M1FN, 32), _fake(cutlass.Float8E8M0FNU, 16), _fake(cutlass.Float8E4M3FN, 16), 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..5ab2ff9cf09d 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 @@ -40,7 +40,7 @@ """Routed-expert GEMMs of the DeepSeek-V4.1 low-batch MoE in CuTe DSL. -The low-batch MoE (``routed_expert_gemm.py``) multiplies at most eight token +The low-batch MoE (``routed_expert_gemm.py``) multiplies at most 192 token rows by the weights of the experts they chose. Its routing record (``routed_expert_route_cute.py``) lists, for each 8-row token tile, the tile's expert and its row limit (the end of the expert's rows within the @@ -62,22 +62,36 @@ Warps: 0 loads A and its scales, 1 loads B and its scales, 2 issues the MMAs, 4-7 run the epilogue (warp 3 is idle). -Programmatic dependent launch. GEMM2's weights stream before GEMM1 finishes: -only its B load warp waits for GEMM1, and it reads the record without waiting -(the record comes from the kernel before GEMM1, and GEMM1 triggers its -dependents only after its own wait for that kernel). GEMM1 needs the record to -pick its experts, so all its warps wait first. +Programmatic dependent launch. GEMM1 needs the record to pick its experts, so +all its warps wait for the kernel before it first; it then triggers its +dependents at once. GEMM2 never waits for the GEMM1 grid. It reads the record +without waiting (the record comes from the kernel before GEMM1, and GEMM1 +triggers only after its own wait for that kernel). Its CTAs take the SMs that +GEMM1's CTAs leave (one CTA of either fills an SM's shared memory) and stream +their weights at once. Its B load warp waits for each item's token tile +alone: GEMM1 counts a tile's finished items in the record's ready words +(``record[ready + c]``, zeroed with the record). After an item's epilogue +stores, GEMM1's four epilogue warps meet at a barrier and one thread adds 1 +with release semantics at GPU scope. GEMM2's B load warp spins with acquire +loads until tile ``c`` counts all of GEMM1's m tiles, then fences the async +proxy before its TMA loads of H and HS. So GEMM2 works on the tiles of +GEMM1's first waves while GEMM1's last, partial wave runs. A record serves +one GEMM1 and GEMM2 pair: the ready words are zeroed only when the record is +built. GEMM2 completes only after every GEMM1 item's epilogue, after which +GEMM1 reads and writes no global memory (its CTAs read the record's tile +count before triggering, and a release fence precedes the trigger); so a +kernel that waits for GEMM2 also sees GEMM1's work done. The weights are those vLLM prepares for the TRT-LLM-gen backend, whose rows are shuffled for 128-row epilogue tiles: in each block of 32, shuffled row ``s0 + 8 * s1`` (``s0 < 8``, ``s1 < 4``) holds row ``4 * s0 + s1``. Their scales are in the 128x4 layout the kernel reads for A. -GEMM1 (:class:`RoutedExpertGemm1`) multiplies the whole routed input (at most -eight rows) by each tile's expert and keeps, for every row of the tile, the -column of the token routed to it: clamped SwiGLU of the interleaved up/gate -rows, quantized to MXFP8 (H and its block scales). The epilogue's arithmetic -follows the GEMM1 kernel from Cake (vllm-project/vllm#59725). +GEMM1 (:class:`RoutedExpertGemm1`) gathers each tile's tokens from the +routed input (``record[4 + 2C + 8c + r]``, the token of tile row ``r``) and +multiplies them by the tile's expert: clamped SwiGLU of the interleaved +up/gate rows, quantized to MXFP8 (H and its block scales). The epilogue's +arithmetic follows the GEMM1 kernel from Cake (vllm-project/vllm#59725). GEMM2 (:class:`RoutedExpertGemm2`) reads token tile ``c`` of the intermediate H, the contiguous rows ``8c .. 8c + 7``, and writes BF16 rows of the output in @@ -89,7 +103,7 @@ from __future__ import annotations -from typing import Optional, cast +from typing import Optional import cuda.bindings.driver as cuda import cutlass @@ -130,6 +144,15 @@ class RoutedExpertGemm: # GEMM1 reads the record only after its grid dependency wait. route_before_wait = True + # GEMM1's B load warp gathers its rows (``gather_b_warp``) instead of + # loading B tiles with TMA. + gather_b = False + # GEMM1 counts each token tile's finished items in the record's ready + # words and triggers its dependents right after its wait; GEMM2 waits, + # item by item, for its tile's count to reach ``ready_target`` (GEMM1's + # m tiles) instead of for the GEMM1 grid (module docstring). + signal_ready = False + ready_target = 0 # Operand types, which a subclass's ``__call__`` sets. a_dtype: type[cutlass.Numeric] b_dtype: type[cutlass.Numeric] @@ -207,11 +230,13 @@ def _launch( epi: tuple, b_side: tuple, cta_words: Int32, + ready: Int32, ctas: Int32, stream: cuda.CUstream, ): """Build the TMA atoms and launch ``ctas`` persistent CTAs. ``sfb`` - None: the B load warp writes B's scales itself (``b_side``).""" + None: the B load warp writes B's scales itself (``b_side``). + ``ready``: the record's first ready word.""" self.mxf8f6f4 = self.a_dtype.width != self.b_dtype.width self.smem_alloc_a_dtype = ( cutlass.Int8 if self.mxf8f6f4 and self.a_dtype.width < 8 else self.a_dtype @@ -307,6 +332,7 @@ def _launch( epi, b_side, cta_words, + ready, self.cluster_layout_vmnk, self.a_smem_layout_staged, self.b_smem_layout_staged, @@ -337,6 +363,7 @@ def kernel( epi: tuple, b_side: tuple, cta_words: Int32, + ready: Int32, cluster_layout_vmnk: cute.Layout, a_smem_layout_staged: cute.ComposedLayout, b_smem_layout_staged: cute.ComposedLayout, @@ -437,6 +464,10 @@ class SharedStorage: count = Int32(0) if bid < total: count = (total - 1 - bid) // grid + 1 + if cutlass.const_expr(self.signal_ready): + # Every warp has read the tile count before warp 0 triggers + # GEMM2, which may finish before a CTA without items exits. + cute.arch.barrier() thr_mma = tiled_mma.get_slice(mma_tile_coord_v) if warp_idx == 0: @@ -466,6 +497,11 @@ class SharedStorage: ) tAsSFA = cute.filter_zeros(tAsSFA) tAgSFA = cute.filter_zeros(tAgSFA) + if cutlass.const_expr(self.signal_ready): + # A release fence first, so that GEMM2's start synchronizes + # with the record (and its zeroed ready words) GEMM1 waited for. + cute.arch.fence_acq_rel_gpu() + cute.arch.griddepcontrol_launch_dependents() policy = Int64(EVICT_FIRST) state = make_pipeline_state(PipelineUserType.Producer, AB_Stage) for i in cutlass.range(count, unroll=1): @@ -496,80 +532,46 @@ class SharedStorage: cache_policy=policy, ) state.advance() - cute.arch.griddepcontrol_launch_dependents() + if cutlass.const_expr(not self.signal_ready): + cute.arch.griddepcontrol_launch_dependents() # Keep shared memory alive until the MMA warp releases every stage. for _ in cutlass.range(AB_Stage, unroll=1): cute.arch.mbarrier_wait(ab_empty_bar + state.index, state.phase) state.advance() elif warp_idx == 1: - # B and its scales, from the previous kernel. - if cutlass.const_expr(self.route_before_wait): + if cutlass.const_expr(self.route_before_wait and self.ready_target == 0): cute.arch.griddepcontrol_wait() - b_cta_layout = cute.make_layout( - cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape - ) - gB_nkl = cute.local_tile( - mB_nkl, (self.cta_n, self.cta_k), (None, None, None) - ) - tBsB, tBgB = cpasync.tma_partition( - tma_atom_b, - block_in_cluster_coord_vmnk[1], - b_cta_layout, - cute.group_modes(sB, 0, 3), - cute.group_modes(thr_mma.partition_B(gB_nkl), 0, 3), - ) - tBsSFB = None - tBgSFB = None - if cutlass.const_expr(tma_atom_sfb is not None): - gSFB_nkl = cute.local_tile( - mSFB_nkl, (128, self.cta_k), (None, None, None) + if cutlass.const_expr(self.gather_b): + # B and its scales, gathered from the previous kernel's rows. + self.gather_b_warp( + sB, + sSFB, + ab_full_bar, + ab_empty_bar, + mRoute, + b_side, + cta_words, + count, ) - tBsSFB, tBgSFB = cpasync.tma_partition( + else: + # B and its scales, from the previous kernel. + self.load_b_warp( + tma_atom_b, + mB_nkl, tma_atom_sfb, - block_in_cluster_coord_vmnk[1], - b_cta_layout, - cute.group_modes(sSFB, 0, 3), - cute.group_modes( - tiled_mma_sfb.get_slice(mma_tile_coord_v).partition_B(gSFB_nkl), - 0, - 3, - ), + mSFB_nkl, + tiled_mma_sfb, + thr_mma, + cluster_layout_vmnk, + block_in_cluster_coord_vmnk, + sB, + sSFB, + ab_full_bar, + ab_empty_bar, + mRoute, + ready, + count, ) - tBsSFB = cute.filter_zeros(tBsSFB) - tBgSFB = cute.filter_zeros(tBgSFB) - state = make_pipeline_state(PipelineUserType.Producer, AB_Stage) - for i in cutlass.range(count, unroll=1): - item = bid + i * grid - n_tile = item // self.m_tiles - b_tile = self.b_tile_index(n_tile) - tBgB_item = tBgB[(None, b_tile, None, 0)] - for k_tile in cutlass.range(self.k_tiles, unroll=1): - cute.arch.mbarrier_wait(ab_empty_bar + state.index, state.phase) - if cutlass.const_expr(tma_atom_sfb is None): - self.write_b_scales(sSFB, state.index, k_tile, b_side) - with cute.arch.elect_one(): - cute.arch.mbarrier_arrive_and_expect_tx( - ab_full_bar + state.index, self.tma_bytes_b - ) - cute.copy( - tma_atom_b, - tBgB_item[(None, k_tile)], - tBsB[(None, state.index)], - tma_bar_ptr=ab_full_bar + state.index, - ) - if cutlass.const_expr(tma_atom_sfb is not None): - cute.copy( - tma_atom_sfb, - cast(cute.Tensor, tBgSFB)[(None, b_tile, None, 0)][ - (None, k_tile) - ], - cast(cute.Tensor, tBsSFB)[(None, state.index)], - tma_bar_ptr=ab_full_bar + state.index, - ) - state.advance() - for _ in cutlass.range(AB_Stage, unroll=1): - cute.arch.mbarrier_wait(ab_empty_bar + state.index, state.phase) - state.advance() elif warp_idx == 2: self.mma_warp( ab_full_bar, @@ -647,6 +649,13 @@ class SharedStorage: self.store( tDrAcc, tDcC, epi, mRoute, cta_words, n_tile, m_tile, i, exchange ) + if cutlass.const_expr(self.signal_ready): + # Every epilogue warp's stores of the item, then one + # release of its tile's count (module docstring). + cute.arch.barrier(barrier_id=2, number_of_threads=128) + if epi_tid == 0: + cute.arch.fence_proxy("async.global") + red_add_release((mRoute.iterator + ready + n_tile).toint()) cute.arch.mbarrier_arrive(tmem_bar) cute.arch.mbarrier_wait(tmem_bar, 1) cute.arch.barrier() @@ -781,11 +790,100 @@ def store( raise NotImplementedError @cute.jit - def write_b_scales( - self, sSFB: cute.Tensor, stage: Int32, k_tile: Int32, b_side: tuple + def load_b_warp( + self, + tma_atom_b: cute.CopyAtom, + mB_nkl: cute.Tensor, + tma_atom_sfb: cute.CopyAtom, + mSFB_nkl: cute.Tensor, + tiled_mma_sfb: cute.TiledMma, + thr_mma: cute.ThrMma, + cluster_layout_vmnk: cute.Layout, + block_in_cluster_coord_vmnk: tuple, + sB: cute.Tensor, + sSFB: cute.Tensor, + ab_full_bar: cute.Pointer, + ab_empty_bar: cute.Pointer, + mRoute: cute.Tensor, + ready: Int32, + count: Int32, ): - """Write B's scales of one stage (a subclass whose B scales are not - loaded with TMA).""" + """The B load warp: TMA of each item's B tile and its scales, each + item's once its token tile's ready count reaches ``ready_target`` + (when set).""" + bid, _, _ = cute.arch.block_idx() + grid, _, _ = cute.arch.grid_dim() + b_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape + ) + gB_nkl = cute.local_tile(mB_nkl, (self.cta_n, self.cta_k), (None, None, None)) + tBsB, tBgB = cpasync.tma_partition( + tma_atom_b, + block_in_cluster_coord_vmnk[1], + b_cta_layout, + cute.group_modes(sB, 0, 3), + cute.group_modes(thr_mma.partition_B(gB_nkl), 0, 3), + ) + gSFB_nkl = cute.local_tile(mSFB_nkl, (128, self.cta_k), (None, None, None)) + tBsSFB, tBgSFB = cpasync.tma_partition( + tma_atom_sfb, + block_in_cluster_coord_vmnk[1], + b_cta_layout, + cute.group_modes(sSFB, 0, 3), + cute.group_modes( + tiled_mma_sfb.get_slice(Int32(0)).partition_B(gSFB_nkl), 0, 3 + ), + ) + tBsSFB = cute.filter_zeros(tBsSFB) + tBgSFB = cute.filter_zeros(tBgSFB) + state = make_pipeline_state(PipelineUserType.Producer, self.num_ab_stage) + for i in cutlass.range(count, unroll=1): + item = bid + i * grid + n_tile = item // self.m_tiles + if cutlass.const_expr(self.ready_target > 0): + wait_ready( + (mRoute.iterator + ready + n_tile).toint(), self.ready_target + ) + b_tile = self.b_tile_index(n_tile) + tBgB_item = tBgB[(None, b_tile, None, 0)] + tBgSFB_item = tBgSFB[(None, b_tile, None, 0)] + for k_tile in cutlass.range(self.k_tiles, unroll=1): + cute.arch.mbarrier_wait(ab_empty_bar + state.index, state.phase) + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx( + ab_full_bar + state.index, self.tma_bytes_b + ) + cute.copy( + tma_atom_b, + tBgB_item[(None, k_tile)], + tBsB[(None, state.index)], + tma_bar_ptr=ab_full_bar + state.index, + ) + cute.copy( + tma_atom_sfb, + tBgSFB_item[(None, k_tile)], + tBsSFB[(None, state.index)], + tma_bar_ptr=ab_full_bar + state.index, + ) + state.advance() + for _ in cutlass.range(self.num_ab_stage, unroll=1): + cute.arch.mbarrier_wait(ab_empty_bar + state.index, state.phase) + state.advance() + + @cute.jit + def gather_b_warp( + self, + sB: cute.Tensor, + sSFB: cute.Tensor, + ab_full_bar: cute.Pointer, + ab_empty_bar: cute.Pointer, + mRoute: cute.Tensor, + b_side: tuple, + cta_words: Int32, + count: Int32, + ): + """The B load warp of a subclass that gathers B's rows itself + (``gather_b``).""" raise NotImplementedError @@ -800,9 +898,17 @@ class RoutedExpertGemm2(RoutedExpertGemm): tile's permuted order.""" def __init__( - self, hidden: int, intermediate: int, cta_k: int = 128, num_ab_stage: int = 10 + self, + hidden: int, + intermediate: int, + gemm1_m_tiles: int, + cta_k: int = 128, + num_ab_stage: int = 10, ) -> None: + """``gemm1_m_tiles``: GEMM1's items per token tile, the ready count + at which a tile's H is complete.""" super().__init__(hidden, hidden // 128, intermediate, cta_k, num_ab_stage) + self.ready_target = gemm1_m_tiles @cute.jit def __call__( @@ -845,8 +951,10 @@ def __call__( c = cute.make_tensor( c_ptr, cute.make_ordered_layout((self.m, rows), order=(0, 1)) ) - route = cute.make_tensor(route_ptr, cute.make_layout(4 + 2 * cta_words)) - self._launch(a, sfa, b, sfb, route, (c,), (), cta_words, ctas, stream) + # The ready words follow the token table (one row per padded row). + ready = 4 + 2 * cta_words + rows + route = cute.make_tensor(route_ptr, cute.make_layout(ready + cta_words)) + self._launch(a, sfa, b, sfb, route, (c,), (), cta_words, ready, ctas, stream) @cute.jit def store( @@ -970,6 +1078,115 @@ def _asm(result, operands: list, text: str, constraints: str, loc, ip): ) +@dsl_user_op +def cp_async_16(dst: Int32, src: Int64, *, loc=None, ip=None) -> None: + """``cp.async`` of 16 bytes from global address ``src`` to shared + address ``dst`` (cached in L2 only).""" + llvm.inline_asm( + None, + [Int32(dst).ir_value(loc=loc, ip=ip), Int64(src).ir_value(loc=loc, ip=ip)], + "cp.async.cg.shared.global [$0], [$1], 16;", + "r,l", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def cp_async_4(dst: Int32, src: Int64, *, loc=None, ip=None) -> None: + """``cp.async`` of 4 bytes from global address ``src`` to shared address + ``dst``.""" + llvm.inline_asm( + None, + [Int32(dst).ir_value(loc=loc, ip=ip), Int64(src).ir_value(loc=loc, ip=ip)], + "cp.async.ca.shared.global [$0], [$1], 4;", + "r,l", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def red_add_release(addr: Int64, *, loc=None, ip=None) -> None: + """Add 1 to the 32-bit word at global address ``addr`` with release + semantics at GPU scope.""" + llvm.inline_asm( + None, + [Int64(addr).ir_value(loc=loc, ip=ip)], + "red.release.gpu.global.add.u32 [$0], 1;", + "l", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def load_acquire(addr: Int64, *, loc=None, ip=None) -> Int32: + """The 32-bit word at global address ``addr``, loaded with acquire + semantics at GPU scope (never hoisted out of a polling loop).""" + return Int32( + llvm.inline_asm( + T.i32(), + [Int64(addr).ir_value(loc=loc, ip=ip)], + "ld.acquire.gpu.global.b32 $0, [$1];", + "=r,l", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def nanosleep(ns: int, *, loc=None, ip=None) -> None: + """Suspend the calling thread for about ``ns`` nanoseconds.""" + llvm.inline_asm( + None, + [], + f"nanosleep.u32 {ns};", + "", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@cute.jit +def wait_ready(addr: Int64, target: cutlass.Constexpr[int]): + """Spin until the ready count at ``addr`` reaches ``target`` (acquire + loads), then fence the async proxy, so that the caller's later TMA loads + see what the count's releasing writers stored (module docstring).""" + seen = load_acquire(addr) + while seen < target: + nanosleep(100) + seen = load_acquire(addr) + cute.arch.fence_proxy("async.global") + + +@cute.jit +def publish_stage(full_bar: cute.Pointer): + """Publish the calling warp's completed shared-memory writes of a stage to + the MMA warp: fence them to the async proxy, through which the MMA reads + shared memory, then one lane arrives on the stage's full barrier.""" + cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.sync_warp() + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive(full_bar) + + @dsl_user_op def clamped_swiglu_x2(up0, up1, gate0, gate1, limit, neg_limit, *, loc=None, ip=None): """``h`` of two tokens: ``up`` clamped to [-limit, limit] and ``gate`` to @@ -1074,11 +1291,11 @@ class RoutedExpertGemm1(RoutedExpertGemm): """GEMM1: ``h = MXFP8(clamped SwiGLU(x @ w13[expert]^T))`` of each token tile, in the tile's permuted rows. - B is the whole routed input ``x`` (at most eight rows; rows past its end - read as zero) with its linear UE8M0 scales, which the B load warp writes - into the layout the MMA reads. Accumulator column ``t`` is token ``t``; - the epilogue writes it to the tile row that token was routed to - (``record[4 + 2 C + 8 c + r] == t``), if any. w13's rows interleave up and + B is tile ``c``'s rows of the routed input ``x``: tile row ``r`` below + the tile's limit is token ``record[4 + 2 C + 8 c + r]``. The B load warp + gathers those rows and their linear UE8M0 scales into the layout the MMA + reads, so accumulator column ``r`` is tile row ``r``, and the epilogue + writes the columns below the limit. w13's rows interleave up and gate, so each thread pairs them, clamps them, applies SwiGLU, scales each 32-column block by its amax / 448 rounded up to a power of two (UE8M0), and stores E4M3 values with the block scales in the per-tile layout GEMM2 @@ -1087,6 +1304,8 @@ class RoutedExpertGemm1(RoutedExpertGemm): """ route_before_wait = False + gather_b = True + signal_ready = True def __init__( self, @@ -1145,42 +1364,129 @@ def __call__( ), ) clamp = cute.make_tensor(clamp_ptr, cute.make_layout(experts)) - route = cute.make_tensor( - route_ptr, cute.make_layout(4 + 2 * cta_words + h_rows) - ) + # The ready words follow the token table (one row per padded row). + ready = 4 + 2 * cta_words + h_rows + route = cute.make_tensor(route_ptr, cute.make_layout(ready + cta_words)) self._launch( - a, sfa, b, None, route, (h, hs, clamp), (b_scale,), cta_words, ctas, stream + a, + sfa, + b, + None, + route, + (h, hs, clamp), + (b, b_scale), + cta_words, + ready, + ctas, + stream, ) @cute.jit def b_tile_index(self, n_tile: Int32) -> Int32: - """Every item multiplies all of ``x``.""" + """B is gathered per item; the epilogue's coordinates stay in tile 0.""" return Int32(0) @cute.jit - def write_b_scales( - self, sSFB: cute.Tensor, stage: Int32, k_tile: Int32, b_side: tuple + def gather_b_warp( + self, + sB: cute.Tensor, + sSFB: cute.Tensor, + ab_full_bar: cute.Pointer, + ab_empty_bar: cute.Pointer, + mRoute: cute.Tensor, + b_side: tuple, + cta_words: Int32, + count: Int32, ): - """Copy the stage's scales of ``x`` (four per 32-bit word) into rows - 0-7 of B's scale atoms; rows past ``x``'s end get scale 1.""" - (b_scale,) = b_side - atoms = self.cta_k // 128 - stage_bytes = 512 * atoms - words = cute.make_tensor( - cute.recast_ptr(sSFB.iterator, dtype=Uint32), - cute.make_layout(self.num_ab_stage * stage_bytes // 4), - ) + """Gather each item's tile rows of ``x`` and their scales into B. + + Per stage, lane ``l`` copies the 16-byte chunk ``l % 16`` of tile rows + ``l // 16 + 2 j`` (``j < 4``) with ``cp.async``, to its address under + the 128-byte swizzle that TMA and the MMA apply to shared addresses, + and lane ``l < 8 x atoms`` copies one 32-bit word (four K blocks) of + row ``l // atoms``'s scales. Rows at or past the tile's limit are not + loaded; their accumulator columns are never stored. The warp never + waits for a copy it has just issued: each stage is published two + stages late (after issuing a stage's copies, the warp waits for those + of two stages before, fences them to the async proxy and arrives on + that stage's full barrier).""" + x, b_scale = b_side + bid, _, _ = cute.arch.block_idx() + grid, _, _ = cute.arch.grid_dim() lane = cute.arch.lane_idx() - if lane < TILE * atoms: - row = lane // atoms - atom = lane - row * atoms - value = Uint32(0x7F7F7F7F) - if row < cute.size(b_scale, mode=[0]): - value = Uint32(b_scale[row, k_tile * atoms + atom]) - words[(stage * stage_bytes + atom * 512 + row * 16) // 4] = value - # The MMA warp copies them to tensor memory through the async proxy. - cute.arch.fence_proxy("async.shared", space="cta") - cute.arch.sync_warp() + atoms = self.cta_k // 128 # 128-column swizzle atoms and scale words per row + b_stage_bytes = TILE * self.cta_k + sf_stage_bytes = 512 * atoms + b_base = sB.iterator.toint() + sf_base = sSFB.iterator.toint() + chunk = lane % 16 + # The chunk's offset in a row of an 8-row stage: atoms of 8 x 128 bytes. + chunk_offset = (chunk // 8) * (TILE * 128) + (chunk % 8) * 16 + sf_row = lane // atoms + sf_atom = lane - sf_row * atoms + # Row r < 32 of a 128-row scale atom: byte 16 r, four K blocks. + sf_offset = sf_atom * 512 + sf_row * 16 + token = cute.make_rmem_tensor((4,), Int32) + state = make_pipeline_state(PipelineUserType.Producer, self.num_ab_stage) + older = Int32(-1) # stages issued but not yet published, oldest first + newer = Int32(-1) + for i in cutlass.range(count, unroll=1): + item = bid + i * grid + n_tile = item // self.m_tiles + row0 = n_tile * TILE + valid = mRoute[4 + cta_words + n_tile] - row0 + for j in cutlass.range_constexpr(4): + token[j] = Int32(-1) + if lane // 16 + 2 * j < valid: + token[j] = mRoute[4 + 2 * cta_words + row0 + lane // 16 + 2 * j] + sf_token = Int32(-1) + if lane < TILE * atoms and sf_row < valid: + sf_token = mRoute[4 + 2 * cta_words + row0 + sf_row] + for k_tile in cutlass.range(self.k_tiles, unroll=1): + cute.arch.mbarrier_wait(ab_empty_bar + state.index, state.phase) + if sf_token >= 0: + cp_async_4( + sf_base + state.index * sf_stage_bytes + sf_offset, + ( + b_scale.iterator + + sf_token * cute.size(b_scale, mode=[1]) + + k_tile * atoms + + sf_atom + ).toint(), + ) + for j in cutlass.range_constexpr(4): + if token[j] >= 0: + linear = ( + b_base + + state.index * b_stage_bytes + + (lane // 16 + 2 * j) * 128 + + chunk_offset + ) + cp_async_16( + linear ^ (((linear >> 7) & 7) << 4), + ( + x.iterator + + token[j] * self.k + + k_tile * self.cta_k + + chunk * 16 + ).toint(), + ) + cute.arch.cp_async_commit_group() + if older >= 0: + cute.arch.cp_async_wait_group(2) + publish_stage(ab_full_bar + older) + older = newer + newer = state.index + state.advance() + if older >= 0: + cute.arch.cp_async_wait_group(1) + publish_stage(ab_full_bar + older) + if newer >= 0: + cute.arch.cp_async_wait_group(0) + publish_stage(ab_full_bar + newer) + for _ in cutlass.range(self.num_ab_stage, unroll=1): + cute.arch.mbarrier_wait(ab_empty_bar + state.index, state.phase) + state.advance() @cute.jit def store( @@ -1195,8 +1501,8 @@ def store( item: Int32, exchange: cute.Tensor, ): - """H and HS of the item (class docstring). The thread holds, for - tokens ``t, t + 1`` (``t`` = 2 x lane % 4), the up and gate rows of H + """H and HS of the item (class docstring). The thread holds, for tile + rows ``t, t + 1`` (``t`` = 2 x lane % 4), the up and gate rows of H columns ``p`` (registers 0-3) and ``p + 1`` (4-7), ``p`` = 64 x m_tile + 16 x warp + 2 x lane / 4.""" mH, mHS, mClamp = epi @@ -1207,16 +1513,13 @@ def store( limit = mClamp[mRoute[4 + n_tile]] row0 = n_tile * TILE valid = mRoute[4 + cta_words + n_tile] - row0 - # The tile rows of the thread's two tokens (-1: not routed here). + # The thread's two tile rows (-1: past the tile's limit). row_a = Int32(-1) row_b = Int32(-1) - for r in cutlass.range_constexpr(TILE): - if r < valid: - routed = mRoute[4 + 2 * cta_words + row0 + r] - if routed == token: - row_a = Int32(r) - if routed == token + 1: - row_b = Int32(r) + if token < valid: + row_a = Int32(token) + if token + 1 < valid: + row_b = Int32(token + 1) lo_a, lo_b = clamped_swiglu_x2( tDrAcc[0], tDrAcc[1], tDrAcc[2], tDrAcc[3], limit, -limit diff --git a/vllm/models/deepseek_v41/nvidia/specialized_decode/moe/routed_expert_route_cute.py b/vllm/models/deepseek_v41/nvidia/specialized_decode/moe/routed_expert_route_cute.py index 14d1567b445d..3500af33bbd7 100644 --- a/vllm/models/deepseek_v41/nvidia/specialized_decode/moe/routed_expert_route_cute.py +++ b/vllm/models/deepseek_v41/nvidia/specialized_decode/moe/routed_expert_route_cute.py @@ -4,7 +4,7 @@ The routed-expert GEMMs (``routed_expert_gemm_cute.py``) multiply token tiles of at most eight rows by the weights of one expert each. The top-k expert ids of -up to eight rows become the expert-major, tile-padded row permutation and the +up to 192 rows become the expert-major, tile-padded row permutation and the per-tile tables both GEMMs read. The result is that of FlashInfer's TRT-LLM-gen routing kernel with precomputed ids and an 8-token tile: @@ -21,18 +21,30 @@ [4, 4 + C) tile -> local expert [4 + C, 4 + 2C) tile -> row limit min((c + 1) * 8, first + count) [4 + 2C, 4 + 2C + 8 M) padded row -> token (GEMM1 only) + [4 + 2C + 8 M, + C) tile -> GEMM1 items done (the ready words), zero where ``M`` is :func:`max_tiles` (the most tiles any routing of the call's - slots needs) and ``C`` is ``M`` rounded up to a multiple of four. + slots needs) and ``C`` is ``M`` rounded up to a multiple of four. GEMM1 + counts up the ready words and GEMM2 waits on them + (``routed_expert_gemm_cute.py``), so a record serves one GEMM1 and GEMM2 + pair. The ids come from a top-k selection, so the valid ids of one row are distinct: -an expert is chosen by at most ``rows <= 8`` slots and owns exactly one tile. -The expert's tile is then the number of distinct selected experts below it (a -presence bitmap), and a slot's rank the number of earlier rows that chose the -expert (a row mask per expert): a few population counts, no scans. +an expert is chosen by at most ``rows`` slots, and a slot's rank is the number +of earlier rows that chose the expert (a row mask per expert, one 32-bit word +per 32 rows). Up to eight rows every chosen expert owns exactly one tile, so +its first tile is the number of distinct selected experts below it (a +presence bitmap). Above eight rows an expert chosen by ``n`` rows owns +``ceil(n / 8)`` tiles: level ``l`` of the bitmap marks the experts chosen by +more than ``8 l`` rows, and the first tile is the sum over levels of the +marked experts below it. Either way the record takes a few population +counts, no scans. Above two levels, one warp first sums each bitmap word's +counts over the levels and prefixes them over the words, so a slot reads its +expert's word's prefix instead of summing every level's words. -:class:`RoutedExpertRoute` builds the record from precomputed ids in one -64-thread CTA; the fused front builds it with the same helpers. +:class:`RoutedExpertRoute` builds the record from precomputed ids in one CTA +(64 threads per level, at most 1,024; above 16 levels a thread takes two +slots); the fused front builds it with the same helpers. """ import cuda.bindings.driver as cuda @@ -43,7 +55,29 @@ TILE = 8 MAX_EXPERTS = 384 WORDS = MAX_EXPERTS // 32 -ROUTE_THREADS = 64 # >= rows * top_k (at most 8 * 8 slots) +MAX_THREADS = 1024 +# Above this many levels, build_record prefixes the bitmap words' counts. +PREFIX_LEVELS = 2 + + +def levels_for(max_rows: int) -> int: + """Bitmap levels of the record build for calls of up to ``max_rows`` rows: + the most tiles one expert can own.""" + return (max_rows + TILE - 1) // TILE + + +def mask_words_for(levels: int) -> int: + """32-bit words of an expert's row mask for ``8 * levels`` rows.""" + return (levels * TILE + 31) // 32 + + +def present_words(levels: int) -> int: + """32-bit words of the ``present`` scratch: the bitmap's ``levels`` levels, + then, above ``PREFIX_LEVELS`` levels, the word prefix of their counts + (``WORDS + 1`` words, the last the total).""" + if levels > PREFIX_LEVELS: + return levels * WORDS + WORDS + 1 + return levels * WORDS @cute.jit @@ -66,7 +100,7 @@ def tile_words(tiles: Int32) -> Int32: def record_words(rows: Int32, top_k: Int32, experts: Int32) -> Int32: """int32 length of the record.""" tiles = max_tiles(rows, top_k, experts) - return 4 + 2 * tile_words(tiles) + TILE * tiles + return 4 + 3 * tile_words(tiles) + TILE * tiles @cute.jit @@ -75,13 +109,17 @@ def clear_scratch( rows_of: cute.Tensor, thread: Int32, threads: cutlass.Constexpr[int], + levels: cutlass.Constexpr[int] = 1, ): - """Zero the expert bitmap and the per-expert row masks; a block barrier - must separate it from :func:`build_record`.""" - for expert in cutlass.range(thread, MAX_EXPERTS, threads, unroll=1): - rows_of[expert] = Uint32(0) - if expert < WORDS: - present[expert] = Uint32(0) + """Zero the expert bitmap's ``levels`` levels and the per-expert row + masks (not the word prefix, which :func:`build_record` writes whole); a + block barrier must separate it from :func:`build_record`.""" + for word in cutlass.range( + thread, MAX_EXPERTS * mask_words_for(levels), threads, unroll=1 + ): + rows_of[word] = Uint32(0) + if word < levels * WORDS: + present[word] = Uint32(0) @cute.jit @@ -102,29 +140,26 @@ def zero_record( @cute.jit -def build_record( - present: cute.Tensor, - rows_of: cute.Tensor, - ids: cute.Tensor, - permuted: cute.Tensor, - record: cute.Tensor, - rows: Int32, - top_k: Int32, - experts: Int32, - thread: Int32, -): - """Build the permutation and the record from ``ids`` (visible to every - thread of the block) into a zeroed record. Every thread of the block calls - it (it synchronizes the block once); thread ``s < rows * top_k`` handles - slot ``s``. ``experts``: a multiple of 32, at most 384; ``rows <= 8``.""" - slots = rows * top_k - words = tile_words(max_tiles(rows, top_k, experts)) - row = thread // top_k +def slot_expert(ids: cute.Tensor, slot: Int32, slots: Int32, experts: Int32) -> Int32: + """Slot ``slot``'s expert: its id, or -1 past the call's slots or for an + id outside ``[0, experts)``.""" expert = Int32(-1) - if thread < slots: - expert = Int32(ids[thread]) + if slot < slots: + expert = Int32(ids[slot]) if expert < 0 or expert >= experts: expert = Int32(-1) + return expert + + +@cute.jit +def mark_slot( + present: cute.Tensor, + rows_of: cute.Tensor, + expert: Int32, + row: Int32, + mask_words: cutlass.Constexpr[int], +): + """Mark ``expert`` in the bitmap's level 0 and ``row`` in its row mask.""" if expert >= 0: cute.arch.atomic_or( present.iterator + (expert >> 5), @@ -132,42 +167,183 @@ def build_record( scope="cta", ) cute.arch.atomic_or( - rows_of.iterator + expert, Uint32(1) << Uint32(row), scope="cta" + rows_of.iterator + expert * mask_words + (row >> 5), + Uint32(1) << Uint32(row & 31), + scope="cta", ) - cute.arch.sync_threads() - if thread < slots: + + +@cute.jit +def count_and_rank( + rows_of: cute.Tensor, + expert: Int32, + row: Int32, + mask_words: cutlass.Constexpr[int], +): + """The rows that chose ``expert``, and those of them before ``row`` (the + slot's rank).""" + count = Int32(0) + rank = Int32(0) + if expert >= 0: + earlier = (Uint32(1) << Uint32(row & 31)) - Uint32(1) + for word in cutlass.range_constexpr(mask_words): + chosen_by = rows_of[expert * mask_words + word] + count += Int32(cute.arch.popc(chosen_by)) + if word < (row >> 5): + rank += Int32(cute.arch.popc(chosen_by)) + if word == (row >> 5): + rank += Int32(cute.arch.popc(chosen_by & earlier)) + return count, rank + + +@cute.jit +def mark_level(present: cute.Tensor, expert: Int32, rank: Int32): + """Level ``l`` marks the experts chosen by more than ``8 l`` rows; the + slot of rank ``8 l`` sets it.""" + if expert >= 0 and rank >= TILE and rank % TILE == 0: + cute.arch.atomic_or( + present.iterator + (rank // TILE) * WORDS + (expert >> 5), + Uint32(1) << Uint32(expert & 31), + scope="cta", + ) + + +@cute.jit +def write_slot( + present: cute.Tensor, + permuted: cute.Tensor, + record: cute.Tensor, + slot: Int32, + slots: Int32, + words: Int32, + expert: Int32, + row: Int32, + count: Int32, + rank: Int32, + levels: cutlass.Constexpr[int], +): + """Store slot ``slot``'s permuted row, its token and, for the first slot + of each of its expert's tiles, the tile's expert and row limit; slot 0 + also stores the tile count.""" + if slot < slots: + prefix = levels * WORDS tiles = Int32(0) tile = Int32(0) - for word in cutlass.range_constexpr(WORDS): - bits = Int32(cute.arch.popc(present[word])) - tiles += bits - if word < (expert >> 5): # never for expert -1 - tile += bits - if thread == 0: + if cutlass.const_expr(levels > PREFIX_LEVELS): + tiles = Int32(present[prefix + WORDS]) + if expert >= 0: + tile = Int32(present[prefix + (expert >> 5)]) + else: + for level in cutlass.range_constexpr(levels): + for word in cutlass.range_constexpr(WORDS): + bits = Int32(cute.arch.popc(present[level * WORDS + word])) + tiles += bits + if word < (expert >> 5): # never for expert -1 + tile += bits + if slot == 0: record[0] = tiles * TILE record[1] = tiles if expert < 0: - permuted[thread] = Int32(-1) + permuted[slot] = Int32(-1) else: below = (Uint32(1) << Uint32(expert & 31)) - Uint32(1) - tile += Int32(cute.arch.popc(present[expert >> 5] & below)) - chosen_by = rows_of[expert] - count = Int32(cute.arch.popc(chosen_by)) - earlier = (Uint32(1) << Uint32(row)) - Uint32(1) - rank = Int32(cute.arch.popc(chosen_by & earlier)) + for level in cutlass.range_constexpr(levels): + tile += Int32( + cute.arch.popc(present[level * WORDS + (expert >> 5)] & below) + ) first = tile * TILE - permuted[thread] = first + rank + permuted[slot] = first + rank record[4 + 2 * words + first + rank] = row - if rank == 0: - record[4 + tile] = expert + if rank % TILE == 0: + # The expert's tile holding ranks rank .. rank + 7. + own = tile + rank // TILE + record[4 + own] = expert limit = first + count - if count > TILE: - limit = first + TILE - record[4 + words + tile] = limit + if count - rank > TILE: + limit = first + rank + TILE + record[4 + words + own] = limit + + +@cute.jit +def build_record( + present: cute.Tensor, + rows_of: cute.Tensor, + ids: cute.Tensor, + permuted: cute.Tensor, + record: cute.Tensor, + rows: Int32, + top_k: Int32, + experts: Int32, + thread: Int32, + threads: cutlass.Constexpr[int], + levels: cutlass.Constexpr[int] = 1, +): + """Build the permutation and the record from ``ids`` (visible to every + thread of the block) into a zeroed record. Every thread of the block's + ``threads`` calls it (it synchronizes the block up to three times); thread + ``t`` handles slots ``t``, ``t + threads``, ... below ``rows * top_k``. + ``experts``: a multiple of 32, at most 384; ``rows <= 8 * levels``, + ``top_k <= 8``; ``present`` holds ``present_words(levels)`` words and + ``rows_of`` ``mask_words_for(levels)`` words per expert.""" + mask_words = mask_words_for(levels) + # Slots a thread handles: up to 8 rows of 8 slots per level. + per_thread = (64 * levels + threads - 1) // threads + slots = rows * top_k + words = tile_words(max_tiles(rows, top_k, experts)) + owned = [] + for j in cutlass.range_constexpr(per_thread): + slot = thread + j * threads + expert = slot_expert(ids, slot, slots, experts) + row = slot // top_k + mark_slot(present, rows_of, expert, row, mask_words) + owned.append((slot, expert, row)) + cute.arch.sync_threads() + counted = [] + for j in cutlass.range_constexpr(per_thread): + slot, expert, row = owned[j] + count, rank = count_and_rank(rows_of, expert, row, mask_words) + if cutlass.const_expr(levels > 1): + mark_level(present, expert, rank) + counted.append((slot, expert, row, count, rank)) + if cutlass.const_expr(levels > 1): + cute.arch.sync_threads() + if cutlass.const_expr(levels > PREFIX_LEVELS): + prefix = levels * WORDS + # Warp 0: lane w < WORDS sums word w's counts over the levels; word + # w's prefix is the sum of the lanes below, the total at WORDS. + if thread < 32: + column = Int32(0) + if thread < WORDS: + for level in cutlass.range_constexpr(levels): + column += Int32(cute.arch.popc(present[level * WORDS + thread])) + words_below = Int32(0) + for word in cutlass.range_constexpr(WORDS): + column_of_word = cute.arch.shuffle_sync(column, word) + if word < thread: + words_below += column_of_word + if thread <= WORDS: + present[prefix + thread] = Uint32(words_below) + cute.arch.sync_threads() + for j in cutlass.range_constexpr(per_thread): + slot, expert, row, count, rank = counted[j] + write_slot( + present, + permuted, + record, + slot, + slots, + words, + expert, + row, + count, + rank, + levels, + ) class RoutedExpertRoute: - """The route kernel: the record from precomputed top-k ids (int32). + """The route kernel: the record from precomputed top-k ids (int32), for + calls of up to ``8 * levels`` rows. PDL: dependents may launch at once (the trigger is the first instruction), so GEMM1's set-up overlaps this kernel; GEMM1 reads the outputs only after @@ -175,6 +351,12 @@ class RoutedExpertRoute: kernel's own wait. """ + def __init__(self, levels: int = 1) -> None: + self.levels = levels + # 8 rows of up to 8 slots per level, at most 1,024 threads (two slots a + # thread above 16 levels). + self.threads = min(64 * levels, MAX_THREADS) + @cute.jit def __call__( self, @@ -186,8 +368,8 @@ def __call__( experts: Int32, stream: cuda.CUstream, ): - """``ids``/``permuted``: int32 [rows, top_k]; rows in [1, 8], - ``rows * top_k <= 64``.""" + """``ids``/``permuted``: int32 [rows, top_k]; rows in + [1, 8 * levels], ``rows * top_k <= 64 * levels``.""" slots = cute.make_layout(rows * top_k) self.kernel( cute.make_tensor(ids_ptr, slots), @@ -199,7 +381,7 @@ def __call__( top_k, experts, ).launch( - grid=(1, 1, 1), block=[ROUTE_THREADS, 1, 1], stream=stream, use_pdl=True + grid=(1, 1, 1), block=[self.threads, 1, 1], stream=stream, use_pdl=True ) @cute.kernel @@ -212,24 +394,38 @@ def kernel( top_k: Int32, experts: Int32, ): - """One CTA of :data:`ROUTE_THREADS` threads.""" + """One CTA of ``min(64 * levels, 1024)`` threads.""" cute.arch.griddepcontrol_launch_dependents() + bitmap_words = present_words(self.levels) + mask_size = MAX_EXPERTS * mask_words_for(self.levels) @cute.struct class Scratch: - present: cute.struct.MemRange[Uint32, WORDS] - rows_of: cute.struct.MemRange[Uint32, MAX_EXPERTS] + present: cute.struct.MemRange[Uint32, bitmap_words] + rows_of: cute.struct.MemRange[Uint32, mask_size] scratch = cutlass.utils.SmemAllocator().allocate(Scratch) - present = cute.make_tensor(scratch.present.data_ptr(), cute.make_layout(WORDS)) + present = cute.make_tensor( + scratch.present.data_ptr(), cute.make_layout(bitmap_words) + ) rows_of = cute.make_tensor( - scratch.rows_of.data_ptr(), cute.make_layout(MAX_EXPERTS) + scratch.rows_of.data_ptr(), cute.make_layout(mask_size) ) thread, _, _ = cute.arch.thread_idx() - clear_scratch(present, rows_of, thread, ROUTE_THREADS) + clear_scratch(present, rows_of, thread, self.threads, self.levels) cute.arch.sync_threads() cute.arch.griddepcontrol_wait() - zero_record(record, rows, top_k, experts, thread, ROUTE_THREADS) + zero_record(record, rows, top_k, experts, thread, self.threads) build_record( - present, rows_of, ids, permuted, record, rows, top_k, experts, thread + present, + rows_of, + ids, + permuted, + record, + rows, + top_k, + experts, + thread, + self.threads, + self.levels, ) diff --git a/vllm/models/deepseek_v41/nvidia/specialized_decode/moe/step.py b/vllm/models/deepseek_v41/nvidia/specialized_decode/moe/step.py index 2e7f0ba5ca5a..f94cbc50ffcc 100644 --- a/vllm/models/deepseek_v41/nvidia/specialized_decode/moe/step.py +++ b/vllm/models/deepseek_v41/nvidia/specialized_decode/moe/step.py @@ -7,8 +7,8 @@ fused kernel, including the layers before an Engram layer and the last one, whose boundaries take the kernel's residual-only variant: -- up to ``LOW_BATCH_MAX_ROWS`` rows, on the low-batch routed-expert GEMMs - (``routed_expert_gemm``); +- up to ``routed_expert_gemm.MAX_ROWS`` rows, on the low-batch routed-expert + GEMMs (``routed_expert_gemm``); - above that, on the runner's own ``_forward_impl`` with the deferral raised to the step's rows for the call (``deferred_finalize``). """ @@ -32,7 +32,7 @@ SharedExpertsOrder, ) -from .routed_expert_gemm import RoutedExpertGemm, front_applies +from .routed_expert_gemm import RoutedExpertGemm, front_applies, padding_mask if TYPE_CHECKING: from ...model import DeepseekV4MoE @@ -88,7 +88,7 @@ def run( if not self.deferrable: return self.ffn(x, input_ids, None) gemm = self.gemm - if step.low_batch and gemm is not None and gemm.accepts(x, self.runner): + if gemm is not None and gemm.accepts(x, self.runner): routed, shared_output = self._forward_low_batch(gemm, x, input_ids) return MoEOutput(routed=routed, shared_output=shared_output) # The runner's custom op returns tensors only, so run its body directly. @@ -119,7 +119,13 @@ def _forward_low_batch( router_logits, _ = runner.gate(hidden_states) runner._maybe_apply_shared_experts(hidden_states, SharedExpertsOrder.NO_OVERLAP) if front_applies(gemm, runner): - routing = gemm.front(hidden_states, router_logits, input_ids) + # The router's padding mask, as select_experts applies it. + routing = gemm.front( + hidden_states, + router_logits, + input_ids, + padding_mask(hidden_states.size(0)), + ) else: topk_weights, topk_ids = runner.router.select_experts( hidden_states=hidden_states, diff --git a/vllm/models/deepseek_v41/nvidia/specialized_decode/runtime/decode.py b/vllm/models/deepseek_v41/nvidia/specialized_decode/runtime/decode.py index 8a12d3e27785..dfeb4f0b6a49 100644 --- a/vllm/models/deepseek_v41/nvidia/specialized_decode/runtime/decode.py +++ b/vllm/models/deepseek_v41/nvidia/specialized_decode/runtime/decode.py @@ -44,11 +44,10 @@ from .. import ( FUSED_MHC_MAX_ROWS, LOW_BATCH_MAX_ROWS, - MHC_OVERLAP_MAX_ROWS, DecodeConfig, uses_checked_flashinfer, ) -from ..attention.indexer import OVERLAP_MAX_ROWS, CandidateOverlap +from ..attention.indexer import CandidateOverlap from ..attention.sparse_mla import SparseIndexTables, SparseMlaDecode from ..attention.step import AttentionStep from ..boundary.boundary import Boundaries, BoundaryState @@ -61,28 +60,23 @@ from ...model import DeepseekV4DecoderLayer, DeepseekV4Model -# Most rows of a step whose Engram WKV runs ahead of the layers. -ENGRAM_PREFETCH_MAX_ROWS = 16 - DecodeOutput: TypeAlias = torch.Tensor | tuple[torch.Tensor, list[torch.Tensor]] +# The executor's all-gather takes up to this much a rank, the custom +# all-gather's limit at TP4, so the logits gather keeps its sizes: up to 32 +# FP32 (64 BF16) rows of a 129,280-token vocabulary. Larger ones take NCCL. +ALL_GATHER_RANK_BYTES = 4 << 20 + class DecodeStep(NamedTuple): """The host-side gates of one captured decode step, by its rows.""" rows: int - # The next mHC coefficients and aux means run on the mHC side stream; - # otherwise inline after the fused all-reduce + mHC. - side_stream: bool # The low-batch GEMMs: routed experts, MXFP8 projections, ll_bf16 side # projections, publishing wo_b. low_batch: bool # The attention side projections run on aux streams. overlap_inputs: bool - # The candidate source forks the consumers' candidate expansion. - overlap_candidates: bool - # The Engram WKV runs ahead of the layers. - engram_prefetch: bool executor: "SpecializedDecode" def run( @@ -92,10 +86,18 @@ def run( input_ids: torch.Tensor | None, engram_hashes: torch.Tensor | None, engram_mask: torch.Tensor | None, + streams: torch.Tensor | None = None, ) -> DecodeOutput: - """The backbone's output after its embedding.""" + """The backbone's output after its embedding (and, from the + executor's ``embedding``, the embedding's hc ``streams``).""" return self.executor.run( - self, hidden_states, positions, input_ids, engram_hashes, engram_mask + self, + hidden_states, + positions, + input_ids, + engram_hashes, + engram_mask, + streams=streams, ) @@ -103,11 +105,8 @@ def step_gates(rows: int, executor: "SpecializedDecode") -> DecodeStep: """The step of ``rows`` rows (0 < rows <= ``FUSED_MHC_MAX_ROWS``).""" return DecodeStep( rows=rows, - side_stream=rows <= MHC_OVERLAP_MAX_ROWS, low_batch=rows <= LOW_BATCH_MAX_ROWS, overlap_inputs=rows <= envs.VLLM_MULTI_STREAM_GEMM_TOKEN_THRESHOLD, - overlap_candidates=rows <= OVERLAP_MAX_ROWS, - engram_prefetch=rows <= ENGRAM_PREFETCH_MAX_ROWS, executor=executor, ) @@ -167,6 +166,7 @@ def __init__( ) # The CuTe DSL kernels, imported only when the executor is built. + from ...ops.cute_dsl import LamportAllGather from ..boundary.all_reduce import BoundaryAllReduce from ..boundary.publishing_gemm import PublishingMxfp8Linear @@ -175,6 +175,13 @@ def __init__( backbone.config, uses_checked_flashinfer("the fused kernel's MXFP8 layer input"), ) + # The row-sharded indexers' gathers, and the logits gather of the + # model's specialized logits processor; NVLink multicast is there, + # since the fused all-reduce needs it too. + self.all_gather = LamportAllGather( + rank_capacity_bytes=ALL_GATHER_RANK_BYTES, + device=torch.device("cuda", torch.accelerator.current_device_index()), + ) self.boundaries = Boundaries(self.all_reduce) self.publisher = ( PublishingMxfp8Linear(self.all_reduce) @@ -212,7 +219,14 @@ def __init__( self.candidates = CandidateOverlap(indexer.indexer_op) self.engram = EngramPrefetch( [layer.engram for layer in layers if layer.engram is not None], - ENGRAM_PREFETCH_MAX_ROWS, + FUSED_MHC_MAX_ROWS, + ) + from ..embedding.replicated import ReplicatedEmbedding + + # The steps' embedding and hc streams from a full copy of the table, + # without the all-reduce (collective: every TP rank builds it). + self.embedding = ReplicatedEmbedding.create( + backbone.embed_tokens, backbone.hc_mult ) gemms = build_routed_expert_gemms([layer.ffn.experts for layer in layers]) quantize = self.all_reduce.quantize_layer_input @@ -228,6 +242,7 @@ def __init__( mxfp8_input=quantize and can_fuse_query_quant([attn.fused_wqa_wkv]), source=self.candidates if attn is source else None, consumer=self.candidates if attn in consumers else None, + all_gather=self.all_gather, ), moe=MoEStep(layer.ffn, gemm, FUSED_MHC_MAX_ROWS), engram=layer.engram, @@ -277,9 +292,11 @@ def run( input_ids: torch.Tensor | None, engram_hashes: torch.Tensor | None, engram_mask: torch.Tensor | None, + streams: torch.Tensor | None = None, ) -> DecodeOutput: """The decoder layers and the final collapse of the step, on the - embedding ``x``; returns what ``DeepseekV4Model.forward`` returns.""" + embedding ``x`` (with its hc ``streams`` when the step's embedding + wrote them); returns what ``DeepseekV4Model.forward`` returns.""" backbone = self.backbone aux_layers = backbone.aux_hidden_state_layers aux: dict[int, torch.Tensor] = {} @@ -287,14 +304,14 @@ def run( self.tables.reset() self.all_reduce.reset() try: - if engram_hashes is not None and step.engram_prefetch: + if engram_hashes is not None: self.engram.start(x.shape[0]) state: BoundaryState | None = None for plan in self.plans: layer = plan.layer quant = None if state is None: - state, layer_input = bound.first(layer, x, step) + state, layer_input = bound.first(layer, x, streams) elif plan.engram is not None and engram_hashes is not None: streams = bound.post_only(state) if plan.layer_id in aux_layers: @@ -305,23 +322,22 @@ def run( engram_hashes[:, plan.engram.layer_hash_index], engram_mask, ) - state, layer_input = bound.pre(layer, injected, state, step) + state, layer_input = bound.pre(layer, injected, state) else: state, layer_input, quant, layer_aux = bound.sublayer( layer, False, state, - step, quantize=plan.attention.mxfp8_input, capture_aux=plan.layer_id in aux_layers, ) if layer_aux is not None: aux[plan.layer_id] = layer_aux state.pending = plan.attention.run(layer_input, quant, positions, step) - bound.join(layer, step) - state, ffn_input, _, _ = bound.sublayer(layer, True, state, step) + bound.join(layer) + state, ffn_input, _, _ = bound.sublayer(layer, True, state) state.pending = plan.moe.run(ffn_input, input_ids, step) - bound.join(layer, step) + bound.join(layer) assert state is not None streams = bound.post_only(state) if backbone.end_layer in aux_layers: diff --git a/vllm/models/deepseek_v41/nvidia/specialized_decode/runtime/metadata_graphs.py b/vllm/models/deepseek_v41/nvidia/specialized_decode/runtime/metadata_graphs.py index 7af6b9ac5e4c..dc0f98c2c678 100644 --- a/vllm/models/deepseek_v41/nvidia/specialized_decode/runtime/metadata_graphs.py +++ b/vllm/models/deepseek_v41/nvidia/specialized_decode/runtime/metadata_graphs.py @@ -48,7 +48,13 @@ } ) +# Recordable builders that only read their block table: on a table that other +# builders clamp in place, they build on branches of their own once those have. +_BLOCK_TABLE_READERS = frozenset({DeepseekV41SparseIndexerMetadataBuilder}) + GraphKey = tuple[int, int, int, int] +# A metadata build branch: (KV-cache group index, attention group) pairs. +Branch = list[tuple[int, AttentionGroup]] # Builds the metadata of the attention groups it is given. BuildFn = Callable[[list[list[AttentionGroup]]], dict[str, Any]] @@ -177,21 +183,34 @@ def _forked_build( block_tables: Sequence[torch.Tensor], build_fn: BuildFn, ) -> dict[str, Any]: - """Build the KV-cache groups as concurrent branches on side streams. + """Build the attention groups as concurrent branches on side streams. The groups share only read-only inputs, with two exceptions. The sparse - MLA builder clamps its block table in place, so groups whose block - tables share storage stay on one branch, in serial order. And the first - group that needs the token-to-request map builds it into its own buffer, - which later groups reuse. Each branch builds its own map; the serial - build's producer is also the first of its branch to need one, so it - fills the buffer that the FULL forward graphs read. + MLA and compressor builders clamp their block table in place, so the + groups of one block table storage build on one branch, in serial order; + the groups that only read the table (``_BLOCK_TABLE_READERS``) then + build on branches of their own, each started after that branch. And + the first group that needs the token-to-request map builds it into its + own buffer, which later groups reuse. Each branch builds its own map; + the serial build's producer is also the first of its branch to need + one, so it fills the buffer that the FULL forward graphs read. """ - branches: dict[int, list[int]] = {} + tables: dict[int, tuple[Branch, Branch]] = {} for i, groups in enumerate(attn_groups): if groups: storage = block_tables[i].untyped_storage().data_ptr() - branches.setdefault(storage, []).append(i) + writers, readers = tables.setdefault(storage, ([], [])) + for group in groups: + reads = type(group.get_metadata_builder()) in _BLOCK_TABLE_READERS + (readers if reads else writers).append((i, group)) + # Each branch with the index of the branch it starts after, if any. + branches: list[tuple[Branch, int | None]] = [] + for writers, readers in tables.values(): + after = None + if writers: + after = len(branches) + branches.append((writers, None)) + branches += [([reader], after) for reader in readers] if len(branches) < 2: return build_fn(attn_groups) # Reused, so a warm-up creates per-stream state (e.g. DeepGEMM's) on the @@ -200,12 +219,17 @@ def _forked_build( self._side_streams.append(torch.cuda.Stream()) main_stream = current_stream() built: dict[str, Any] = {} - for stream, branch in zip(self._side_streams, branches.values()): - stream.wait_stream(main_stream) + for stream, (branch, after) in zip(self._side_streams, branches): + stream.wait_stream( + main_stream if after is None else self._side_streams[after] + ) with torch.cuda.stream(stream): built.update( build_fn( - [g if i in branch else [] for i, g in enumerate(attn_groups)] + [ + [group for j, group in branch if j == i] + for i in range(len(attn_groups)) + ] ) ) for stream in self._side_streams[: len(branches)]: diff --git a/vllm/models/deepseek_v41/nvidia/specialized_decode/runtime/model_state.py b/vllm/models/deepseek_v41/nvidia/specialized_decode/runtime/model_state.py index 251acde01edc..7ae01018482e 100644 --- a/vllm/models/deepseek_v41/nvidia/specialized_decode/runtime/model_state.py +++ b/vllm/models/deepseek_v41/nvidia/specialized_decode/runtime/model_state.py @@ -9,7 +9,9 @@ models, and before memory profiling); - CUDA graphs of the attention metadata builds of FULL decode steps (``metadata_graphs``), recorded with the inner CUDA graphs and replayed in - place of the eager build; + place of the eager build. A step that replays one gathers its engram + lookback windows on a side stream beside it, so the metadata graph hands + off to the forward's graph with no eager kernel between them; - no replay-start upload for decode steps whose starts are known to be zero. """ @@ -21,10 +23,12 @@ from torch import nn from vllm.config import CUDAGraphMode, VllmConfig +from vllm.utils.torch_utils import current_stream from vllm.v1.kv_cache_interface import KVCacheConfig from vllm.v1.worker.gpu.block_table import BlockTables from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache +from vllm.v1.worker.gpu.states import RequestState from vllm.v1.worker.utils import AttentionGroup from ...model_state import DeepseekV41ModelState @@ -83,6 +87,14 @@ def __init__( self._zero_replay_starts = 0 # The key whose metadata build capture_inner_cudagraphs records. self._recording: GraphKey | None = None + # The runner's request state, from the first prepare_inputs on; the + # lookback gather that runs beside a metadata graph reads it. + self._req_states: RequestState | None = None + # The gather's side stream and the event recorded after it. + self._lookback_side: tuple[torch.cuda.Stream, torch.cuda.Event] | None = None + # The batch whose lookback gather was issued beside its metadata graph, + # and the event that ends it. + self._staged_lookback: tuple[InputBatch, torch.cuda.Event] | None = None def initialize_kv_cache( self, kv_cache_config: KVCacheConfig, block_tables: BlockTables @@ -100,6 +112,43 @@ def _upload_replay_start(self, values: np.ndarray, num_reqs: int) -> torch.Tenso self._zero_replay_starts = 0 if has_replay else num_reqs return self._replay_start[:num_reqs] + def prepare_inputs( + self, input_batch: InputBatch, req_states: RequestState + ) -> dict[str, torch.Tensor | None]: + self._req_states = req_states + return super().prepare_inputs(input_batch, req_states) + + def _stage_lookback(self, input_batch: InputBatch) -> None: + """Issue the batch's lookback gather on a side stream, beside the + metadata graph about to replay. It reads the batch's request indices + and the token history, which the runner has written by now.""" + if self.lookback_token_ids is None or self._req_states is None: + return + if self._lookback_side is None: + self._lookback_side = (torch.cuda.Stream(), torch.cuda.Event()) + self._join_staged_lookback() + stream, event = self._lookback_side + stream.wait_stream(current_stream()) + with torch.cuda.stream(stream): + super()._fill_lookback(input_batch, self._req_states) + event.record(stream) + self._staged_lookback = (input_batch, event) + + def _join_staged_lookback(self) -> InputBatch | None: + """Make the current stream wait for a staged lookback gather, if any; + returns the batch it gathered.""" + staged, self._staged_lookback = self._staged_lookback, None + if staged is None: + return None + current_stream().wait_event(staged[1]) + return staged[0] + + def _fill_lookback(self, input_batch: InputBatch, req_states: RequestState) -> None: + # Even a stale gather (a batch whose inputs were never prepared) is + # joined, so that no two gathers write the window at once. + if self._join_staged_lookback() is not input_batch: + super()._fill_lookback(input_batch, req_states) + def _build_attn_metadata( self, input_batch: InputBatch, @@ -133,6 +182,7 @@ def _build_attn_metadata( if key is not None: # decode_backend=latency_specialized refuses micro-batching. assert ubatch_idx == 0 + self._stage_lookback(input_batch) attn_metadata = self.graphs.replay(key) if attn_metadata is not None: return attn_metadata 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..970c45870c4b 100644 --- a/vllm/models/deepseek_v41/nvidia/specialized_decode/runtime/warmup.py +++ b/vllm/models/deepseek_v41/nvidia/specialized_decode/runtime/warmup.py @@ -18,10 +18,18 @@ from vllm.config import VllmConfig from vllm.distributed import get_tp_group from vllm.forward_context import set_forward_context +from vllm.model_executor.layers.sparse_attn_indexer import SparseAttnIndexer -from .. import FUSED_MHC_MAX_ROWS, LOW_BATCH_MAX_ROWS +from ...ops.o_proj import precompile_dsv41_o_proj_early_trigger +from .. import FUSED_MHC_MAX_ROWS, LOW_BATCH_MAX_ROWS, WIDE_LOW_LATENCY_MAX_ROWS +from ..attention import decode_topk_cute, sparse_mla_cute from ..attention.aux_projections import warmup_aux_projection_gemms -from ..attention.mxfp8 import pinned_tactics, run_zeros, shape_representatives +from ..attention.mxfp8 import ( + compile_pinned, + 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 @@ -52,6 +60,7 @@ def warm_up(decode: SpecializedDecode, vllm_config: VllmConfig) -> None: for linear in linears: for m in range(1, LOW_BATCH_MAX_ROWS + 1): run_zeros(linear, m) + compile_pinned(linear) # The publishing wo_b compiles the tactics selected above; one layer # per shape compiles every layer's kernels. publishing = { @@ -65,10 +74,36 @@ def warm_up(decode: SpecializedDecode, vllm_config: VllmConfig) -> None: step.publisher.precompile(step.attn.wo_b, m) warmup_aux_projection_gemms(_aux_weights(steps)) warmup_low_batch_lm_head(decode.lm_head) + device = torch.accelerator.current_device_index() + for has_sinks in {step.attn.attn_sink is not None for step in steps}: + sparse_mla_cute.precompile(device, has_sinks) + for width, k, block_mode in _decode_topk_variants(steps): + decode_topk_cute.precompile(device, width, k, block_mode) + # The fused WO-A variant the steps' o_proj launches (early_trigger); + # one layer per shape compiles every layer's kernels. + for attn in { + (a.n_local_groups, a.n_local_heads, a.o_lora_rank): a + for a in (step.attn for step in steps) + }.values(): + precompile_dsv41_o_proj_early_trigger(attn, WIDE_LOW_LATENCY_MAX_ROWS) if decode.config.jit_warmup: _warm_up_mhc(decode, vllm_config, rows) +def _decode_topk_variants(steps: list[AttentionStep]) -> set[tuple[int, int, bool]]: + """(logits width, k, block mode) of every ``decode_topk`` call the dense + indexers can make: the paged MQA logits are ``max_model_len`` wide.""" + variants: set[tuple[int, int, bool]] = set() + for step in steps: + op = None if step.attn.indexer is None else step.attn.indexer.indexer_op + if not isinstance(op, SparseAttnIndexer) or op.topk_backend != "auto": + continue + variants.add((op.max_model_len, op.topk_tokens, False)) + if op.candidate_blocks is not None and op.candidate_block_size == 8: + variants.add((op.max_model_len, op.candidate_blocks.shape[1], True)) + return variants + + @contextmanager def _tuning() -> Iterator[None]: """FlashInfer's tuning context, its timings reduced over the TP group.""" @@ -97,7 +132,7 @@ def _deferred_moe( decode: SpecializedDecode, vllm_config: VllmConfig, rows: int, tune: bool ) -> None: """One deferring MoE layer on ``rows`` random rows with the finalize - deferred, as the executor runs it above ``LOW_BATCH_MAX_ROWS`` rows: + deferred, as the executor runs it above the routed-expert GEMMs' rows: tunes its experts over the serving buckets, and allocates what the deferred launch allocates on first use before memory profiling.""" moe = next((plan.moe for plan in decode.plans if plan.moe.deferrable), None)