From 193922d6ff5f7e213eacadd3cb7dd96612b30ed3 Mon Sep 17 00:00:00 2001 From: Fangzhou Ai <31551580+Fangzhou-Ai@users.noreply.github.com> Date: Fri, 9 Oct 2026 21:53:18 -0400 Subject: [PATCH] [ROCm][Perf][DSv4.1] Mono decode layer for DeepSeek-V4.1-Flash on MI355X (#60397) Signed-off-by: Fangzhou Ai Co-authored-by: Claude Opus 5.5 --- .buildkite/test_areas/models_basic.yaml | 20 +- tests/models/test_dsv41_mono_decode.py | 317 ++++ tests/models/test_dsv41_mono_layers.py | 335 ++++ tests/models/test_dsv41_mono_numerics.py | 529 ++++++ vllm/envs.py | 7 + vllm/models/deepseek_v41/amd/model.py | 17 + vllm/models/deepseek_v41/amd/mono/README.md | 94 + vllm/models/deepseek_v41/amd/mono/__init__.py | 5 + .../amd/mono/attention/__init__.py | 4 + .../deepseek_v41/amd/mono/attention/back.py | 884 +++++++++ .../deepseek_v41/amd/mono/attention/device.py | 504 +++++ .../deepseek_v41/amd/mono/attention/front.py | 662 +++++++ .../deepseek_v41/amd/mono/attention/gemv.py | 234 +++ .../deepseek_v41/amd/mono/attention/plan.py | 173 ++ .../deepseek_v41/amd/mono/common/__init__.py | 9 + .../models/deepseek_v41/amd/mono/common/mx.py | 78 + .../deepseek_v41/amd/mono/common/ops.py | 300 +++ .../amd/mono/common/peer_memory.py | 104 ++ .../deepseek_v41/amd/mono/common/plan.py | 135 ++ .../deepseek_v41/amd/mono/common/sync.py | 177 ++ vllm/models/deepseek_v41/amd/mono/layer.py | 940 ++++++++++ vllm/models/deepseek_v41/amd/mono/runner.py | 437 +++++ vllm/models/deepseek_v41/amd/mono/sources.py | 8 + .../deepseek_v41/amd/mono/stages/__init__.py | 4 + .../deepseek_v41/amd/mono/stages/dims.py | 60 + .../deepseek_v41/amd/mono/stages/gemv.py | 294 +++ .../deepseek_v41/amd/mono/stages/moe.py | 1635 +++++++++++++++++ .../deepseek_v41/amd/mono/stages/moe_plan.py | 54 + .../deepseek_v41/amd/mono/stages/moe_shape.py | 229 +++ .../deepseek_v41/amd/mono/stages/router.py | 86 + .../deepseek_v41/amd/mono/stages/seam.py | 406 ++++ vllm/models/deepseek_v41/amd/mono_decode.py | 376 ++++ vllm/models/deepseek_v41/amd/rocm.py | 4 +- 33 files changed, 9119 insertions(+), 2 deletions(-) create mode 100644 tests/models/test_dsv41_mono_decode.py create mode 100644 tests/models/test_dsv41_mono_layers.py create mode 100644 tests/models/test_dsv41_mono_numerics.py create mode 100644 vllm/models/deepseek_v41/amd/mono/README.md create mode 100644 vllm/models/deepseek_v41/amd/mono/__init__.py create mode 100644 vllm/models/deepseek_v41/amd/mono/attention/__init__.py create mode 100644 vllm/models/deepseek_v41/amd/mono/attention/back.py create mode 100644 vllm/models/deepseek_v41/amd/mono/attention/device.py create mode 100644 vllm/models/deepseek_v41/amd/mono/attention/front.py create mode 100644 vllm/models/deepseek_v41/amd/mono/attention/gemv.py create mode 100644 vllm/models/deepseek_v41/amd/mono/attention/plan.py create mode 100644 vllm/models/deepseek_v41/amd/mono/common/__init__.py create mode 100644 vllm/models/deepseek_v41/amd/mono/common/mx.py create mode 100644 vllm/models/deepseek_v41/amd/mono/common/ops.py create mode 100644 vllm/models/deepseek_v41/amd/mono/common/peer_memory.py create mode 100644 vllm/models/deepseek_v41/amd/mono/common/plan.py create mode 100644 vllm/models/deepseek_v41/amd/mono/common/sync.py create mode 100644 vllm/models/deepseek_v41/amd/mono/layer.py create mode 100644 vllm/models/deepseek_v41/amd/mono/runner.py create mode 100644 vllm/models/deepseek_v41/amd/mono/sources.py create mode 100644 vllm/models/deepseek_v41/amd/mono/stages/__init__.py create mode 100644 vllm/models/deepseek_v41/amd/mono/stages/dims.py create mode 100644 vllm/models/deepseek_v41/amd/mono/stages/gemv.py create mode 100644 vllm/models/deepseek_v41/amd/mono/stages/moe.py create mode 100644 vllm/models/deepseek_v41/amd/mono/stages/moe_plan.py create mode 100644 vllm/models/deepseek_v41/amd/mono/stages/moe_shape.py create mode 100644 vllm/models/deepseek_v41/amd/mono/stages/router.py create mode 100644 vllm/models/deepseek_v41/amd/mono/stages/seam.py create mode 100644 vllm/models/deepseek_v41/amd/mono_decode.py diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index 1d9a7fe2f3..483a9aceac 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.buildkite/test_areas/models_basic.yaml @@ -91,8 +91,9 @@ steps: - tests/models/test_deepseek_v4_dspark_rocm.py - tests/models/test_deepseek_v4_rocm_wo_a.py - tests/models/test_deepseek_v4_rocm_compressor_gemm_fusion.py + - tests/models/test_dsv41_mono_decode.py commands: - - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py models/test_deepseek_v4_vl_rocm.py models/transformers/test_layer_registry_e2e.py models/test_deepseek_v41_replay_start.py models/test_deepseek_v41_decoder_replay_layers.py models/test_deepseek_v41_replay_batch.py models/test_deepseek_v4_dspark_rocm.py models/test_deepseek_v4_rocm_wo_a.py models/test_deepseek_v4_rocm_compressor_gemm_fusion.py + - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py models/test_deepseek_v4_vl_rocm.py models/transformers/test_layer_registry_e2e.py models/test_deepseek_v41_replay_start.py models/test_deepseek_v41_decoder_replay_layers.py models/test_deepseek_v41_replay_batch.py models/test_deepseek_v4_dspark_rocm.py models/test_deepseek_v4_rocm_wo_a.py models/test_deepseek_v4_rocm_compressor_gemm_fusion.py models/test_dsv41_mono_decode.py - VLLM_ROCM_USE_AITER=1 pytest -v -s models/test_hyv4_rocm.py -m 'not distributed' - label: ":nvidia: (B200) Inkling" @@ -259,3 +260,20 @@ steps: device: cpu-small commands: - pytest -v -s models/test_utils.py models/test_vision.py models/test_adapters.py models/test_qwen3_5_mtp_config.py models/transformers/fusers/ model_executor/layers/test_activation.py models/test_language_model_cache_is_weak.py models/deepseek_v32/test_sequence_parallel.py models/test_deepseek_v4_fi_moe_ep.py models/test_dspark_mla.py models/transformers/test_layer_registry.py + +- label: ':amd: (MI355) DeepSeek-V4.1 Mono Decode Numerics' + key: dsv41-mono-decode-numerics-amd + dind: false + device: mi355_4 + num_devices: 4 + timeout_in_minutes: 40 + working_dir: /vllm-workspace/ + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/models/deepseek_v41/amd/ + - tests/models/test_dsv41_mono_numerics.py + - tests/models/test_dsv41_mono_layers.py + commands: + - pytest -v -s tests/models/test_dsv41_mono_numerics.py + - pytest -v -s tests/models/test_dsv41_mono_layers.py diff --git a/tests/models/test_dsv41_mono_decode.py b/tests/models/test_dsv41_mono_decode.py new file mode 100644 index 0000000000..684e72c004 --- /dev/null +++ b/tests/models/test_dsv41_mono_decode.py @@ -0,0 +1,317 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The DeepSeek-V4.1 mono decode layer runs only where its kernels are exact: +decode-only steps of at most MAX_ROWS rows with causal SWA windows, eager or in +a FULL graph. A PIECEWISE capture is replayed for mixed batches, so it must +never record the mono launches. Layers outside the whole-layer kernels run the +FFN launch on the same steps, from wo_b's unreduced output.""" + +import sys +from types import ModuleType, SimpleNamespace + +import pytest +import torch + +from vllm.config import CUDAGraphMode +from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import Mxfp4MoeBackend +from vllm.models.deepseek_v41.amd import mono_decode as md + +M = 6 + + +class _Runner: + def __init__(self): + self.calls = [] + + def forward(self, *args, **kwargs): + self.calls.append((args, kwargs)) + return ("mono outputs",) + + def ffn(self, *args): + self.calls.append((args, {})) + return ("ffn outputs",) + + +def _layer(comp="comp"): + cache = lambda: torch.zeros(4, 128 * md.RECORD, dtype=torch.uint8) # noqa: E731 + attn = SimpleNamespace( + swa_cache_layer=SimpleNamespace(prefix="swa", kv_cache=cache()), + compressed_cache_prefix=comp, + topk_indices_buffer=torch.zeros(64, 512, dtype=torch.int32), + compress_ratio=1, + _compressed_kv_cache=cache, + ) + return SimpleNamespace(attn=attn) + + +def _swa(**overrides): + fields = dict( + num_prefills=0, + num_decodes=1, + num_decode_tokens=M, + decode_swa_width=md.SWA_WIDTH, + decode_swa_indices=torch.zeros(64, 1, md.SWA_WIDTH, dtype=torch.int32), + decode_swa_lens=torch.zeros(64, dtype=torch.int32), + token_to_req_indices=torch.zeros(64, dtype=torch.int32), + slot_mapping=torch.zeros(64, dtype=torch.int64), + block_size=32, + ) + fields.update(overrides) + return SimpleNamespace(**fields) + + +def _inputs(rows=M): + return dict( + x=torch.zeros(rows, 5120, dtype=torch.bfloat16), + positions=torch.zeros(rows, dtype=torch.int64), + residual=torch.zeros(rows, 4, 5120, dtype=torch.bfloat16), + post_mix=torch.zeros(rows, 4, 1), + res_mix=torch.zeros(rows, 4, 4), + pre_mix=torch.zeros(rows, 4), + ) + + +@pytest.fixture +def run(monkeypatch): + """Call the mono path for one layer under a forward context built from + ``mode`` and ``metadata``; returns (result, runner calls).""" + + def _run( + mode=CUDAGraphMode.FULL, metadata=None, ffn_only=False, entry=None, **inputs + ): + if metadata is None: + metadata = { + "swa": _swa(), + "comp": SimpleNamespace(block_size=128, block_table=None), + } + fc = SimpleNamespace(cudagraph_runtime_mode=mode, attn_metadata=metadata) + runner = _Runner() + monkeypatch.setattr(md, "is_forward_context_available", lambda: True) + monkeypatch.setattr(md, "get_forward_context", lambda: fc) + monkeypatch.setattr(md, "_mono_runner", lambda device: runner) + monkeypatch.setattr(md.MonoDecodeLayer, "weights", lambda self, layer: None) + args = {**_inputs(), **inputs} + mono = md.MonoDecodeLayer(ffn_only) + if (entry or ("ffn" if ffn_only else "layer")) == "ffn": + # x stands for wo_b's unreduced output; the first layer's attention + # has no compressed cache + part, _ = args.pop("x"), args.pop("positions") + return mono.ffn(_layer(comp=None), part, **args), runner.calls + return mono(_layer(), **args), runner.calls + + return _run + + +@pytest.mark.parametrize("mode", [CUDAGraphMode.NONE, CUDAGraphMode.FULL]) +def test_decode_step_runs_mono(run, mode): + out, calls = run(mode) + assert out == ("mono outputs",) and len(calls) == 1 + args, kwargs = calls[0] + # the caches reach the kernels as [blocks, block, record] bytes + assert args[8].shape == (4, 32, md.RECORD) and args[8].stride() == ( + 128 * md.RECORD, + md.RECORD, + 1, + ) + assert kwargs["comp_cache"].shape == (4, 128, md.RECORD) + + +def test_padded_decode_step_runs_mono(run): + """A FULL graph pads the step to its capture size: the kernels take it and + skip the pad rows (their slot is -1).""" + swa = _swa(num_decode_tokens=M - 2) + comp = SimpleNamespace(block_size=128, block_table=None) + out, calls = run(metadata={"swa": swa, "comp": comp}) + assert out == ("mono outputs",) and len(calls) == 1 + + +@pytest.mark.parametrize( + "case", + [ + dict(mode=CUDAGraphMode.PIECEWISE), + dict( + metadata={ + "swa": _swa(num_prefills=1), + "comp": SimpleNamespace(block_size=128), + } + ), + dict( + metadata={ + "swa": _swa(decode_swa_width=256), + "comp": SimpleNamespace(block_size=128), + } + ), + dict( + metadata={ + "swa": _swa(num_decode_tokens=M + 1), + "comp": SimpleNamespace(block_size=128), + } + ), + dict(metadata={}), # a profile / dummy run: no attention metadata + dict(**_inputs(md.MAX_ROWS + 6)), + dict(residual=None), # the first layer's seam broadcasts the embedding + ], + ids=[ + "piecewise", + "prefill", + "noncausal-window", + "tokens-past-rows", + "no-metadata", + "rows", + "first-layer", + ], +) +def test_other_steps_take_the_original_path(run, case): + out, calls = run(**case) + assert out is None and not calls + + +def _decoder_layer(engram=None, **attn): + attn = ( + dict( + layer_id=21, + compressor=None, + indexer=None, + compress_ratio=1, + kv_mxfp8=False, + kv_cache_dtype="fp8_ds_mla", + ) + | attn + ) + backend = Mxfp4MoeBackend.AITER_MXFP4_BF16 + quant_method = SimpleNamespace(mxfp4_backend=backend) + ffn = SimpleNamespace( + experts=SimpleNamespace( + routed_experts=SimpleNamespace(quant_method=quant_method) + ), + shared_experts=object(), + gate=SimpleNamespace(tid2eid=None), + n_routed_experts=384, + n_activated_experts=6, + routed_scaling_factor=1.5, + swiglu_limit=10.0, + scoring_func="sqrtsoftplus", + renormalize=True, + ) + return SimpleNamespace( + attn=SimpleNamespace(**attn), + ffn=ffn, + engram=engram, + use_sequence_parallel=False, + fuse_seam_norm=True, + ) + + +@pytest.mark.parametrize( + "layer, path", + [ + (_decoder_layer(), "whole"), + (_decoder_layer(compress_ratio=0), "ffn"), + (_decoder_layer(engram=object()), "ffn"), + (_decoder_layer(compressor=object(), indexer=object()), "ffn"), + (_decoder_layer(indexer=object()), "ffn"), + (_decoder_layer(kv_cache_dtype="auto"), "ffn"), + (_decoder_layer(layer_id=40), None), + ], + ids=["standard", "first", "engram", "kv-source", "index-source", "kv", "draft"], +) +def test_layers_take_whole_or_ffn_launch(monkeypatch, layer, path): + """The whole-layer kernels serve the standard layers; every other backbone + layer keeps vLLM's attention and runs its FFN half as one launch.""" + _deployment(monkeypatch) + mono = md.MonoDecodeLayer.create(layer, _config()) + got = None if mono is None else "ffn" if mono.ffn_only else "whole" + assert got == path + + +@pytest.mark.parametrize("mode", [CUDAGraphMode.NONE, CUDAGraphMode.FULL]) +def test_decode_step_runs_ffn_launch(run, mode): + out, calls = run(mode, ffn_only=True, metadata={"swa": _swa()}) + assert out == ("ffn outputs",) and len(calls) == 1 + assert calls[0][0][1].shape == (M, 5120) # wo_b's unreduced output + + +def test_ffn_layer_never_runs_whole(run): + out, calls = run(ffn_only=True, entry="layer") + assert out is None and not calls + + +@pytest.mark.parametrize( + "case", + [ + dict(mode=CUDAGraphMode.PIECEWISE), + dict(metadata={"swa": _swa(num_prefills=1)}), + dict(metadata={"swa": _swa(num_decode_tokens=M + 1)}), + dict(metadata={}), + dict(**_inputs(md.MAX_ROWS + 6)), + ], + ids=["piecewise", "prefill", "tokens-past-rows", "no-metadata", "rows"], +) +def test_other_steps_reduce_wo_b_themselves(run, case): + """None: the layer all-reduces wo_b's partial and runs the FFN as before.""" + case.setdefault("metadata", {"swa": _swa()}) + out, calls = run(ffn_only=True, **case) + assert out is None and not calls + + +def _config(**parallel): + parallel = ( + dict(enable_expert_parallel=False, enable_eplb=False, data_parallel_size=1) + | parallel + ) + return SimpleNamespace( + model_config=SimpleNamespace(hf_config=SimpleNamespace(num_hidden_layers=40)), + parallel_config=SimpleNamespace(**parallel), + ) + + +def _deployment(monkeypatch, cdna=4, cus=256): + """A deployment the kernels take (the runner module stubbed: no FlyDSL).""" + import vllm.platforms.rocm as rocm + + monkeypatch.setenv("VLLM_ROCM_MONO_DECODE", "1") + monkeypatch.delenv("VLLM_ROCM_USE_AITER_MOE_A4W4_DSV4", raising=False) + monkeypatch.setattr(rocm, "get_cdna_version", lambda: cdna) + monkeypatch.setattr(md, "get_tensor_model_parallel_world_size", lambda: 2) + monkeypatch.setattr(md, "_compute_units", lambda: cus) + name = "vllm.models.deepseek_v41.amd.mono.runner" + runner = ModuleType(name) + runner.BLOCKS = 256 # type: ignore[attr-defined] + monkeypatch.setitem(sys.modules, name, runner) + + +@pytest.mark.parametrize("cdna", [3, 5]) +def test_opt_in_outside_the_kernels_cdna_raises(monkeypatch, cdna): + _deployment(monkeypatch, cdna=cdna) + with pytest.raises(ValueError, match="needs CDNA4"): + md.MonoDecodeLayer.create(_decoder_layer(), _config()) + + +def _other_backend(): + layer = _decoder_layer() + experts = layer.ffn.experts.routed_experts + experts.quant_method.mxfp4_backend = Mxfp4MoeBackend.TRITON + return layer + + +@pytest.mark.parametrize( + "case, match", + [ + (dict(parallel=dict(enable_expert_parallel=True)), "no expert"), + (dict(parallel=dict(data_parallel_size=2)), "no expert"), + (dict(parallel=dict(enable_eplb=True)), "no expert"), + (dict(a4w4=True), "A8W4"), + (dict(layer=_other_backend), "AITER_MXFP4_BF16"), + (dict(cus=128), "256 compute units"), + ], + ids=["expert-parallel", "data-parallel", "eplb", "a4w4", "moe-backend", "cus"], +) +def test_opt_in_outside_the_kernels_moe_raises(monkeypatch, case, match): + """The kernels read all 384 experts on each rank in AITER's A8W4 layout, with + all 256 CTAs resident: any other deployment is an error, not garbage.""" + _deployment(monkeypatch, cus=case.get("cus", 256)) + if case.get("a4w4"): + monkeypatch.setenv("VLLM_ROCM_USE_AITER_MOE_A4W4_DSV4", "1") + layer = case.get("layer", _decoder_layer)() + with pytest.raises(ValueError, match=match): + md.MonoDecodeLayer.create(layer, _config(**case.get("parallel", {}))) diff --git a/tests/models/test_dsv41_mono_layers.py b/tests/models/test_dsv41_mono_layers.py new file mode 100644 index 0000000000..580462a823 --- /dev/null +++ b/tests/models/test_dsv41_mono_layers.py @@ -0,0 +1,335 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The DeepSeek-V4.1 mono decode path through the real model: a five-layer +DeepSeek-V4.1-Flash at TP2 and TP4 on CDNA4, served with ``VLLM_ROCM_MONO_DECODE`` on +and off on the same random weights: every decoder layer's outputs at every +decode step must match. Unlike ``test_dsv41_mono_numerics.py``, which +checks the kernels against vLLM's ops one stage at a time, this runs +``DeepseekV4DecoderLayer.forward`` as the engine does -- the model's hooks, +the per-step path choice, the weights taken off the layer, the metadata and +the indexer's top-k -- so a change to the decoder layer that the kernels do +not follow fails here. + +The five layers hold both kinds the mono path takes: +- 0: sliding window only -> the FFN launch +- 1: ratio-2 KV and index source -> the FFN launch +- 2: ratio-2 consumer -> the whole layer (K1 + K2) +- 3: ratio-1 KV, index and candidate source -> the FFN launch +- 4: ratio-1 consumer -> the whole layer +""" + +import zlib + +import pytest +import torch + +from vllm.model_executor.model_loader import register_model_loader +from vllm.model_executor.model_loader.dummy_loader import DummyModelLoader +from vllm.platforms import current_platform + +MODEL = "deepseek-ai/DeepSeek-V4.1-Flash" +LOAD_FORMAT = "dsv41_mono_random" +TRUNCATED = { + "num_hidden_layers": 5, + # one entry a layer, the three MTP layers' included + "compress_ratios": [0, 2, 2, 1, 1, 0, 0, 0], + "kv_source_layer_ids": [1, 3], + "index_source_layer_ids": [1, 3], + "candidate_source_layer_id": 3, + "engram_layer_ids": [], + "engram_num_embeddings": [], + "dspark_target_layer_ids": [2, 3, 4], +} +WHOLE, FFN = (2, 4), (0, 1, 3) +STEPS = 5 # decode steps after the prefill's token + + +def _on_cdna4() -> bool: + if not current_platform.is_rocm(): + return False + from vllm.platforms.rocm import get_cdna_version + + return get_cdna_version() == 4 + + +@torch.no_grad() +def _fill(name: str, t: torch.Tensor) -> None: + """A random value for every weight, seeded by its name (alike on every + rank and in every run), in the ranges the checkpoint's formats hold.""" + if t.device.type == "meta": + return # materialized and filled by the dummy loader's own path + g = torch.Generator(device=t.device).manual_seed(zlib.crc32(name.encode())) + + def randn(scale=1.0, shift=0.0): + return torch.randn(t.shape, generator=g, device=t.device) * scale + shift + + if t.dtype == torch.uint8 and "scale" in name: + # E8M0 2^-12 .. 2^-8, alike over each 32 rows as in the checkpoint (the + # loader folds per-row scales into 32 x 32 blocks only then). Larger + # random layers make the stack chaotic: atomics' order alone moves + # vLLM's own logits run to run. + rows = t.shape[-2] if t.dim() > 1 else 1 + shape = (*t.shape[:-2], -(-rows // 32), t.shape[-1]) if t.dim() > 1 else t.shape + e8m0 = torch.randint(115, 120, shape, generator=g, device=t.device) + if t.dim() > 1: + e8m0 = e8m0.repeat_interleave(32, dim=-2)[..., :rows, :] + t.copy_(e8m0) + elif t.dtype in (torch.uint8, torch.float4_e2m1fn_x2): + u8 = t.view(torch.uint8) + u8.copy_(torch.randint(0, 256, u8.shape, generator=g, device=t.device)) + elif t.dtype == torch.float8_e4m3fn: + t.copy_(randn().clamp(-6, 6)) + elif not torch.is_floating_point(t): + return + elif "norm" in name: + t.copy_(randn(0.1, 1.0)) + elif name.endswith(("hc_attn_scale", "hc_ffn_scale")): + t.copy_(randn(0.5).abs() + 0.5) + elif name.endswith(("hc_attn_base", "hc_ffn_base", "attn_sink")): + t.copy_(randn(0.1)) + elif name.endswith("embed_tokens.weight"): + t.copy_(randn()) + else: + t.copy_(randn(0.02)) + + +@register_model_loader(LOAD_FORMAT) +class RandomWeightLoader(DummyModelLoader): + """Dummy loading with weights in the checkpoint's ranges, before post-load + processing, so every processed copy follows (plain dummy weights leave the + E8M0 scales zero, and both paths would agree on zeros).""" + + def load_weights(self, model, model_config) -> None: + from vllm.distributed import get_tensor_model_parallel_world_size + + super().load_weights(model, model_config) + for name, t in model.state_dict().items(): + _fill(name, t) + # the routed experts' intermediate past this rank's share is the + # loader's zero padding (576 -> 640 at TP4): the kernels skip it + inter = model_config.hf_text_config.moe_intermediate_size + inter //= get_tensor_model_parallel_world_size() + for m in model.modules(): + w13, w2 = getattr(m, "w13_weight", None), getattr(m, "w2_weight", None) + if w13 is None or w2 is None or w13.device.type == "meta": + continue + padded = w13.shape[1] // 2 + if padded > inter: + w13, w2 = w13.data.view(torch.uint8), w2.data.view(torch.uint8) + for half in (0, padded): + w13[:, half + inter : half + padded] = 0 + w2[..., inter // 2 :] = 0 + + +class MonoProbe: + """Worker extension (``worker_extension_cls``): it also brings this module, + and so the loader, into the workers.""" + + def _decoder_layers(self) -> dict: + from vllm.models.deepseek_v41.amd.model import DeepseekV4DecoderLayer + + model = self.model_runner.model # type: ignore[attr-defined] + return { + m.attn.layer_id: m + for m in model.modules() + if isinstance(m, DeepseekV4DecoderLayer) + } + + def capture_steps(self, lens: list[int]) -> None: + """Keep each decoder layer's outputs (x, residual) for every decode + token taken in a decode-only step (the steps the mono path takes), + keyed by (layer, prompt length, position): a request's step and row + can differ between engines (requests reach the engine core while it + runs), its prompt length and position cannot. ``lens``: the prompt + lengths, all distinct; a row's request is the longest prompt not past + its position.""" + from vllm.forward_context import get_forward_context + + self._captured: dict = {} + lens = sorted(lens) + + def hook(layer_id): + def keep(module, args, kwargs, out): + md = get_forward_context().attn_metadata + if not isinstance(md, dict): + return + swa = md.get(module.attn.swa_cache_layer.prefix) + if swa is None or swa.num_prefills != 0: + return + pos = kwargs.get("positions", args[1] if len(args) > 1 else None) + for row, p in enumerate(pos.tolist()): + n = max(n for n in lens if n <= p) + kept = [t[row].float().cpu() for t in out[:2]] + self._captured.setdefault((layer_id, n, p), kept) + + return keep + + def gap_hook(layer_id, layer): + def keep(ffn, args, kwargs): + # vLLM's path only (the mono kernels route inside): each row's + # top-6 margin, sqrt(softplus(logits)) + bias as vLLM routes + md = get_forward_context().attn_metadata + if not isinstance(md, dict): + return + swa = md.get(layer.attn.swa_cache_layer.prefix) + if swa is None or swa.num_prefills != 0: + return + h = args[0] if args else kwargs["hidden_states"] + logits = h.float() @ ffn.gate.weight.float().T + score = torch.nn.functional.softplus(logits).sqrt() + score = score + ffn.gate.e_score_correction_bias.float() + top = score.topk(7, dim=-1).values + gap = (top[:, 5] - top[:, 6]).tolist() + pos = self._positions + for row, p in enumerate(pos.tolist()): + n = max(n for n in lens if n <= p) + self._gaps.setdefault((layer_id, n, p), gap[row]) + + return keep + + def pos_hook(module, args, kwargs): + self._positions = kwargs.get( + "positions", args[1] if len(args) > 1 else None + ) + + self._gaps: dict = {} + for i, layer in self._decoder_layers().items(): + layer.register_forward_hook(hook(i), with_kwargs=True) + layer.register_forward_pre_hook(pos_hook, with_kwargs=True) + gap = gap_hook(i, layer) + layer.ffn.register_forward_pre_hook(gap, with_kwargs=True) + + def save_capture(self, path: str) -> None: + rank = self.rank # type: ignore[attr-defined] + torch.save((self._captured, self._gaps), f"{path}.rank{rank}") + + def mono_layers(self) -> dict: + from vllm.models.deepseek_v41.amd import mono_decode + + layers = { + i: None + if m.mono is None + else ["ffn" if m.mono.ffn_only else "whole", m.mono._weights is not None] + for i, m in self._decoder_layers().items() + } + runner = mono_decode._runner + return { + "layers": layers, + "epoch": -1 if runner is None else int(runner.epoch[0]), + } + + +def _decode(vllm_runner, monkeypatch, tp: int, mono: bool, prompts, path: str): + """Greedy decode of ``prompts`` with the mono path on or off: (each + prompt's output, the workers' mono layers); each decoder layer's outputs + for the decode tokens of decode-only steps, and vLLM's routing margins, + saved under ``path``.""" + from vllm import SamplingParams + from vllm.inputs import TokensPrompt + + monkeypatch.setenv("VLLM_ROCM_MONO_DECODE", "1" if mono else "0") + monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1") + monkeypatch.setenv("VLLM_ROCM_USE_AITER_MOE", "1") + with vllm_runner( + MODEL, + tensor_parallel_size=tp, + load_format=LOAD_FORMAT, + hf_overrides=TRUNCATED, + worker_extension_cls=f"{__name__}.MonoProbe", + moe_backend="aiter", + language_model_only=True, + tokenizer_mode="deepseek_v41", + attention_config={"indexer_kv_dtype": "mxfp4", "indexer_sparse_logits": True}, + block_size=128, + enforce_eager=True, + max_model_len=4096, + max_num_seqs=len(prompts), + # a prompt prefills in one step, so its decode tokens never share one + max_num_batched_tokens=16384, + enable_prefix_caching=False, + gpu_memory_utilization=0.5, + ) as runner: + runner.llm.collective_rpc("capture_steps", args=([len(p) for p in prompts],)) + params = SamplingParams(temperature=0.0, max_tokens=STEPS + 1) + outs = runner.llm.generate( + [TokensPrompt(prompt_token_ids=p) for p in prompts], params + ) + runner.llm.collective_rpc("save_capture", args=(path,)) + probe = runner.llm.collective_rpc("mono_layers") + return [o.outputs[0] for o in outs], probe + + +def _rel(a: torch.Tensor, b: torch.Tensor) -> float: + return ((a - b).norm() / b.norm()).item() + + +@pytest.mark.skipif(not _on_cdna4(), reason="the mono kernels target CDNA4") +@pytest.mark.parametrize("tp", [2, 4]) +def test_mono_decode_matches_the_model(vllm_runner, monkeypatch, tmp_path, tp): + if torch.accelerator.device_count() < tp: + pytest.skip(f"needs {tp} GPUs") + pytest.importorskip("flydsl") + g = torch.Generator().manual_seed(0) + # contexts past both ratios' 512 compressed entries, so top-k selects + lens = torch.randperm(1700, generator=g)[:8].add(300).tolist() # distinct + prompts = [torch.randint(1000, 60000, (n,), generator=g).tolist() for n in lens] + + on, probe_on = _decode( + vllm_runner, monkeypatch, tp, True, prompts, f"{tmp_path}/on" + ) + off, probe_off = _decode( + vllm_runner, monkeypatch, tp, False, prompts, f"{tmp_path}/off" + ) + + # both layer kinds took the mono path on every rank, at every decode step + launches = len(WHOLE + FFN) # one launch (pair) a layer a step + for rank in probe_on: + assert rank["layers"] == { + i: ["whole" if i in WHOLE else "ffn", True] for i in WHOLE + FFN + }, rank + assert rank["epoch"] >= launches * STEPS, rank + assert rank["epoch"] % launches == 0, rank + assert all(set(r["layers"].values()) == {None} for r in probe_off) + # each decode token's input matches between the runs while the greedy + # tokens before it do: (prompt length, position) pairs to compare + same: set[tuple[int, int]] = set() + for prompt, a, b in zip(prompts, on, off): + for step in range(STEPS): + if a.token_ids[: step + 1] != b.token_ids[: step + 1]: + break + same.add((len(prompt), len(prompt) + step)) + + # each layer's outputs for those tokens, mono against vLLM's path, token by + # token. Layer 0 sees the same inputs on both paths at a request's first + # decode token: there the bound is test_dsv41_mono_numerics' (about 3x the + # kernels' 0.4%); its later tokens attend over KV each path wrote, 3%. Each + # later layer takes the previous one's slightly different output, which + # random weights amplify (2-4% by layer 4): 8%. A token whose top-6 routing + # is a near-tie on vLLM's path (6th / 7th score within 5e-3: a 0.4% input + # difference moves scores near 0.8 by about 3e-3) may pick another expert; + # once it has, its residual stream differs and the later layers are exempt + # for it. The residual stream: 3% (0.3% seen). A layer the kernels no + # longer follow is off by far more, on every token. + for rank in range(tp): + got, _ = torch.load(f"{tmp_path}/on.rank{rank}") + want, gaps = torch.load(f"{tmp_path}/off.rank{rank}") + keys: list[tuple[int, int]] = sorted( + same & {k[1:] for k in got} & {k[1:] for k in want} + ) + assert len(keys) >= len(prompts), keys + # the bound that sees a 3% error: a request's first decode token + assert any(p == n for n, p in keys), keys + flipped: set = set() + for i in sorted(WHOLE + FFN): + for k in keys: + if k in flipped: + continue + (x, res), (x_ref, res_ref) = got[(i, *k)], want[(i, *k)] + err = _rel(res, res_ref) + assert err < 0.03, (rank, i, k, "residual", err) + bound = (0.015 if k[1] == k[0] else 0.03) if i == 0 else 0.08 + err = _rel(x, x_ref) + if err >= bound: + assert gaps[(i, *k)] < 5e-3, (rank, i, k, "x", err) + flipped.add(k) + assert len(flipped) <= len(keys) // 4, (rank, sorted(flipped)) diff --git a/tests/models/test_dsv41_mono_numerics.py b/tests/models/test_dsv41_mono_numerics.py new file mode 100644 index 0000000000..f7cd877caf --- /dev/null +++ b/tests/models/test_dsv41_mono_numerics.py @@ -0,0 +1,529 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Numerics of the DeepSeek-V4.1 mono decode kernels at TP2 on CDNA4 against the +ops vLLM runs for the same decode step: random weights in the checkpoint's +formats, the step's KV caches as vLLM's insert writes them. + +- The whole layer (K1 + K2) against the mHC seams, the ROCm attention op chain, + the TP all-reduce and ``DeepseekV4MoE``. +- The FFN launch against the same FFN half, from wo_b's unreduced output. + +Each stage is compared on its own input, so a failure names the stage; the MoE +is checked on the kernels' own FFN input, where routing cannot differ.""" + +import pytest +import ray +import torch + +from vllm.platforms import current_platform + +from ..utils import ( + ensure_model_parallel_initialized, + init_test_distributed_environment, + multi_process_parallel, +) + +MODEL = "deepseek-ai/DeepSeek-V4.1-Flash" +HIDDEN, Q_RANK, HEAD_DIM, NOPE, ROPE = 5120, 1280, 512, 448, 64 +N_HEADS, O_GROUPS, O_RANK = 64, 8, 1024 +WINDOW, TOPK, SWA_BLOCK, CACHE_BLOCK = 128, 512, 32, 128 +RECORD, DATA, ALIGN = 584, 576, 576 +REPLAYS = 300 +HC = 4 + + +def _on_cdna4() -> bool: + if not current_platform.is_rocm(): + return False + from vllm.platforms.rocm import get_cdna_version + + return get_cdna_version() == 4 + + +def _cos(a: torch.Tensor, b: torch.Tensor) -> float: + a, b = a.float().reshape(-1), b.float().reshape(-1) + return torch.nn.functional.cosine_similarity(a, b, dim=0).item() + + +def _rel(a: torch.Tensor, b: torch.Tensor) -> float: + """Relative L2 error of ``a`` against ``b``: unlike the cosine, it sees a + wrong scale.""" + a, b = a.float().reshape(-1), b.float().reshape(-1) + return ((a - b).norm() / b.norm()).item() + + +# ---------------------------------------------------------------- caches + + +def _paged_cache(num_blocks: int, block: int, device) -> torch.Tensor: + """[blocks, block, 584] u8 over pages of round_up(block * 584, 576) bytes, + the way vLLM binds a packed V4 cache.""" + page = -(-block * RECORD // ALIGN) * ALIGN + raw = torch.zeros(num_blocks * page, dtype=torch.uint8, device=device) + return torch.as_strided(raw, (num_blocks, block, RECORD), (page, RECORD, 1)) + + +def _rope(kv: torch.Tensor, positions: torch.Tensor, cos_sin: torch.Tensor): + """GPT-J RoPE of the last 64 dims (fp32 math, bf16 out).""" + cs = cos_sin[positions.long()] + cos, sin = cs[:, :32], cs[:, 32:] + x = kv.float() + e, o = x[:, NOPE::2].clone(), x[:, NOPE + 1 :: 2].clone() + x[:, NOPE::2] = e * cos - o * sin + x[:, NOPE + 1 :: 2] = e * sin + o * cos + return x.to(torch.bfloat16) + + +def _write_records(cache, block: int, slots, kv): + """fp8_ds_mla records of roped bf16 kv rows at ``slots``: the 448 nope dims + in e4m3 with a UE8M0 per 64, the rope dims in bf16, as vLLM's insert.""" + nope = kv.float()[:, :NOPE].view(-1, 7, 64) + exp = torch.ceil(torch.log2(nope.abs().amax(-1).clamp_min(1e-4) / 448.0)) + q = (nope * torch.exp2(-exp)[..., None]).clamp(-448, 448).reshape(-1, NOPE) + data = torch.cat( + [ + q.to(torch.float8_e4m3fn).view(torch.uint8), + kv[:, NOPE:].contiguous().view(torch.uint8), + ], + 1, + ) + scales = torch.zeros(kv.shape[0], 8, dtype=torch.uint8, device=kv.device) + scales[:, :7] = (exp + 127).clamp(0, 255).to(torch.uint8) + n = cache.untyped_storage().nbytes() - cache.storage_offset() + flat = torch.as_strided(cache, (n,), (1,)) + slots = slots.to(kv.device).long() + b, o = slots // block, slots % block + d0 = b * cache.stride(0) + o * DATA + s0 = b * cache.stride(0) + block * DATA + o * 8 + ar = torch.arange(DATA, device=kv.device) + flat[(d0[:, None] + ar).reshape(-1)] = data.reshape(-1) + ar = torch.arange(8, device=kv.device) + flat[(s0[:, None] + ar).reshape(-1)] = scales.reshape(-1) + + +def _normed_kv(n: int, kv_norm, g, device): + x = torch.randn(n, HEAD_DIM, generator=g).to(device) + x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True)) * kv_norm.float() + return x.to(torch.bfloat16) + + +def _decode_step(ratio: int, R: int, T: int, kv_norm, cos_sin, seed: int): + """R requests x T tokens past random contexts: the step's metadata as + vLLM's builders lay it out, its window rows and (ratio > 0) the compressed + cache filled, a random top-512 a token.""" + g = torch.Generator().manual_seed(seed) + dev = kv_norm.device + ctx = [int(c) for c in torch.randint(2000, 40000, (R,), generator=g)] + M = R * T + spans = [ + (max(0, c - WINDOW - T) // SWA_BLOCK, (c + T - 1) // SWA_BLOCK) for c in ctx + ] + n_swa = 1 + sum(hi - lo + 1 for lo, hi in spans) + perm = (torch.randperm(n_swa - 1, generator=g) + 1).tolist() + bt = torch.zeros(R, max(hi for _, hi in spans) + 1, dtype=torch.int64) + for r, (lo, hi) in enumerate(spans): + for lb in range(lo, hi + 1): + bt[r, lb] = perm.pop() + swa = _paged_cache(n_swa, SWA_BLOCK, dev) + + def slot(r, p): + return int(bt[r, p // SWA_BLOCK]) * SWA_BLOCK + p % SWA_BLOCK + + pos = [ctx[r] + t for r in range(R) for t in range(T)] + t2r = [r for r in range(R) for _ in range(T)] + swa_idx = torch.full((M, WINDOW), -1, dtype=torch.int32) + swa_len = torch.zeros(M, dtype=torch.int32) + for i in range(M): + ps = range(max(pos[i] - WINDOW + 1, 0), pos[i] + 1) + swa_idx[i, : len(ps)] = torch.tensor([slot(t2r[i], p) for p in ps]) + swa_len[i] = len(ps) + hist = [ + (r, p) for r in range(R) for p in range(max(0, ctx[r] - WINDOW - T), ctx[r]) + ] + hp = torch.tensor([p for _, p in hist], device=dev) + kv = _rope(_normed_kv(len(hist), kv_norm, g, dev), hp, cos_sin) + _write_records(swa, SWA_BLOCK, torch.tensor([slot(r, p) for r, p in hist]), kv) + + step = dict( + M=M, + positions=torch.tensor(pos, dtype=torch.int64, device=dev), + slot_mapping=torch.tensor( + [slot(t2r[i], pos[i]) for i in range(M)], dtype=torch.int64, device=dev + ), + token_to_req=torch.tensor(t2r, dtype=torch.int32, device=dev), + swa_indices=swa_idx.to(dev), + swa_lens=swa_len.to(dev), + swa_cache=swa, + topk=None, + comp_cache=None, + comp_block_table=None, + comp_entries=0, + ) + if ratio: + entries = CACHE_BLOCK // ratio + lens = [(c + T) // ratio for c in ctx] + nblk = [-(-n // entries) for n in lens] + perm = (torch.randperm(sum(nblk), generator=g) + 1).tolist() + cbt = torch.zeros(R, max(nblk), dtype=torch.int32) + for r in range(R): + for b in range(nblk[r]): + cbt[r, b] = perm.pop() + comp = _paged_cache(1 + sum(nblk), entries, dev) + for r in range(R): + idx = torch.arange(lens[r]) + slots = cbt[r, idx // entries].long() * entries + idx % entries + kv = _rope( + _normed_kv(lens[r], kv_norm, g, dev), (idx * ratio).to(dev), cos_sin + ) + _write_records(comp, entries, slots, kv) + topk = torch.full((M, TOPK), -1, dtype=torch.int32) + for i in range(M): + cand = (pos[i] + 1) // ratio + n = min(cand, TOPK) + topk[i, :n] = torch.randperm(cand, generator=g)[:n].to(torch.int32) + step.update( + topk=topk.to(dev), + comp_cache=comp, + comp_block_table=cbt.to(dev), + comp_entries=entries, + ) + return step + + +# ---------------------------------------------------------------- weights + + +def _mxfp8(n: int, k: int, g, device): + """A random MXFP8 weight [n, k] e4m3 and its E8M0 32 x 32 block scales.""" + w = torch.randn(n, k, generator=g).clamp(-6, 6).to(torch.float8_e4m3fn) + s = torch.randint(118, 123, (n // 32, k // 32), generator=g, dtype=torch.uint8) + return w.to(device), s.to(device) + + +def _random_moe(moe, gd, gr, inter: int) -> None: + """Random parameters in the loaded layout, before post-load processing: this + rank's shards (``gd``, a device generator) of the MXFP4 routed experts + (random nibbles, E8M0 scales; the intermediate past ``inter`` zero, as the + loader pads it) and the MXFP8 shared expert (one E8M0 a 32 x 32 block); the + replicated router from ``gr``, seeded alike on every rank.""" + u8 = torch.uint8 + for lin in (moe.shared_experts.gate_up_proj, moe.shared_experts.down_proj): + n = lin.weight.shape[0] + w = torch.randn(lin.weight.shape, generator=gd, device=gd.device) + lin.weight.data.copy_(w.clamp(-6, 6).to(lin.weight.dtype)) + sc = lin.weight_scale.data.view(u8) + rows = -(-n // 32) if sc.shape[0] == n else sc.shape[0] + blocks = torch.randint( + 118, 123, (rows, sc.shape[1]), generator=gd, device=gd.device, dtype=u8 + ) + sc.copy_(blocks.repeat_interleave(32, 0)[:n] if sc.shape[0] == n else blocks) + e = moe.experts.routed_experts + for name, lo, hi in (("weight", 0, 256), ("weight_scale", 118, 122)): + for w in ("w13", "w2"): + t = getattr(e, f"{w}_{name}").data.view(u8) + t.copy_( + torch.randint(lo, hi, t.shape, generator=gd, device=gd.device, dtype=u8) + ) + padded = e.w13_weight.shape[1] // 2 + if padded > inter: + w13, w2 = e.w13_weight.data.view(u8), e.w2_weight.data.view(u8) + for half in (0, padded): + w13[:, half + inter : half + padded] = 0 + w2[..., inter // 2 :] = 0 + for t in (moe.gate.weight, moe.gate.e_score_correction_bias): + t.data.copy_(torch.randn(t.shape, generator=gr, device=gr.device) * 0.02) + + +# ---------------------------------------------------------------- the worker + + +@ray.remote(num_gpus=1, max_calls=1) +def _mono_numerics(monkeypatch, tp_size, pp_size, rank, distributed_init_port): + from vllm._aiter_ops import rocm_aiter_ops + from vllm.v1.worker.workspace import init_workspace_manager + + with monkeypatch.context() as m: + m.delenv("CUDA_VISIBLE_DEVICES", raising=False) + # the checkpoint's MXFP4 experts run on AITER's MoE; its flags are read + # at import, so refresh them + m.setenv("VLLM_ROCM_USE_AITER", "1") + m.setenv("VLLM_ROCM_USE_AITER_MOE", "1") + rocm_aiter_ops.refresh_env_variables() + device = torch.device(f"cuda:{rank}") + torch.accelerator.set_device_index(device) + init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port) + ensure_model_parallel_initialized(tp_size, pp_size) + init_workspace_manager(device) + _check_layers(rank, tp_size, device) + + +def _check_layers(rank: int, tp: int, device) -> None: + from vllm._aiter_ops import rocm_aiter_ops + from vllm.config import VllmConfig, set_current_vllm_config + from vllm.distributed import get_tp_group, tensor_model_parallel_all_reduce + from vllm.engine.arg_utils import EngineArgs + from vllm.forward_context import set_forward_context + from vllm.model_executor.kernels.linear.mxfp8.rocm_block32_gemm import ( + rocm_mxfp8_block32_gemm, + ) + from vllm.model_executor.layers.mhc import MHCPreDelayedOp + from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( + mxfp8_e4m3_quantize, + ) + from vllm.model_executor.model_loader.utils import process_weights_after_loading + from vllm.models.common.ops import fused_q_kv_rmsnorm + from vllm.models.deepseek_v41.amd.model import DeepseekV4MoE + from vllm.models.deepseek_v41.amd.mono.layer import scratch_layout + from vllm.models.deepseek_v41.amd.mono.runner import ( + AttnWeights, + DSV41MonoLayer, + MonoLayerWeights, + ) + from vllm.models.deepseek_v41.amd.rocm import ( + compute_global_topk_ragged_indices_and_indptr, + ) + from vllm.models.deepseek_v41.common.rope import build_deepseek_v4_rope + from vllm.utils.torch_utils import set_default_torch_dtype + from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( + build_ragged_indices_from_dense, + rocm_inverse_rope_mxfp8_rows, + ) + + vc = EngineArgs( + model=MODEL, + tensor_parallel_size=tp, + moe_backend="aiter", + load_format="dummy", + max_model_len=65536, + language_model_only=True, + attention_config={"indexer_kv_dtype": "mxfp4", "indexer_sparse_logits": True}, + ).create_engine_config() + cfg = vc.model_config.hf_config + cfg = getattr(cfg, "text_config", cfg) + # replicated tensors alike on every rank; each rank's own TP shards + g = torch.Generator().manual_seed(1000) + gs = torch.Generator().manual_seed(1001 + rank) + H, G = N_HEADS // tp, O_GROUPS // tp + bf = torch.bfloat16 + + with set_current_vllm_config(vc), set_default_torch_dtype(bf), device: + moe = DeepseekV4MoE(vc, prefix="model.layers.21.ffn") + mhc = MHCPreDelayedOp() + _random_moe( + moe, + torch.Generator(device=device).manual_seed(2001 + rank), + torch.Generator(device=device).manual_seed(2000), + inter=2304 // tp, + ) + process_weights_after_loading(moe, vc.model_config, device) + + def f32(*shape, s=1.0, gen=g): + return (torch.randn(*shape, generator=gen) * s).to(device) + + mix = (2 + HC) * HC + seams = { + sub: dict( + fn=f32(mix, HC * HIDDEN, s=0.01), + scale=f32(3, s=0.5).abs() + 0.5, + base=f32(mix, s=0.1), + norm=(1 + f32(HIDDEN, s=0.1)).to(bf), + ) + for sub in ("attn", "ffn") + } + + def seam(sub, residual, x, post, comb, pre): + s = seams[sub] + return mhc( + residual, s["fn"], s["scale"], s["base"], cfg.rms_norm_eps, + cfg.hc_eps, cfg.hc_eps, 2.0, cfg.hc_sinkhorn_iters, + pre_mix=pre, sublayer_out=x, post_layer_mix=post, comb_res_mix=comb, + norm_weight=s["norm"], norm_eps=cfg.rms_norm_eps, + ) # fmt: skip + + def moe_out(x): + with set_forward_context(None, vc): + ids = torch.zeros(x.shape[0], dtype=torch.int64, device=device) + return moe(x.clone(), ids) + + wqkv, wqkv_s = _mxfp8(1792, HIDDEN, g, device) + wq_b, wq_b_s = _mxfp8(H * HEAD_DIM, Q_RANK, gs, device) + wo_a, wo_a_s = _mxfp8(G * O_RANK, H // G * HEAD_DIM, gs, device) + wo_b, wo_b_s = _mxfp8(HIDDEN, G * O_RANK, gs, device) + q_norm, kv_norm = (1 + f32(Q_RANK, s=0.1)).to(bf), (1 + f32(HEAD_DIM, s=0.1)).to(bf) + sink = f32(H, gen=gs) + e, sh = moe.experts.routed_experts, moe.shared_experts + u8 = lambda t: t.view(torch.uint8) # noqa: E731 + runner = DSV41MonoLayer(tp, rank, get_tp_group().cpu_group, device) + gw, bias = moe.gate.weight.float(), moe.gate.e_score_correction_bias.float() + + def top6(x): + """The top-6 experts as vLLM routes: sqrt(softplus(logits)) + bias.""" + sc = torch.nn.functional.softplus(x.float() @ gw.T).sqrt() + bias + return sc.topk(6, dim=1).indices.sort(1).values + + steps = [] # each case's launches, on fixed buffers, for the replay stress + for ratio, R, T in ((1, 1, 6), (2, 8, 6)): + with set_current_vllm_config(VllmConfig()): + rope = build_deepseek_v4_rope( + cfg, head_dim=HEAD_DIM, rope_head_dim=ROPE, + max_position_embeddings=cfg.max_position_embeddings, + compress_ratio=ratio, + ) # fmt: skip + cos_sin = rope.cos_sin_cache.to(device=device, dtype=torch.float32) + s = _decode_step(ratio, R, T, kv_norm, cos_sin, seed=7 + ratio) + M = s["M"] + mw = MonoLayerWeights( + attn=AttnWeights( + layer_id=21, wqkv=wqkv, wqkv_scale=wqkv_s, q_norm=q_norm, + kv_norm=kv_norm, wq_b=wq_b, wq_b_scale=wq_b_s, wo_a=wo_a, + wo_a_scale=wo_a_s, wo_b=wo_b, wo_b_scale=wo_b_s, attn_sink=sink, + cos_sin=cos_sin, ratio=ratio, + ), + hc_attn_fn=seams["attn"]["fn"], hc_attn_scale=seams["attn"]["scale"], + hc_attn_base=seams["attn"]["base"], attn_norm=seams["attn"]["norm"], + hc_ffn_fn=seams["ffn"]["fn"], hc_ffn_scale=seams["ffn"]["scale"], + hc_ffn_base=seams["ffn"]["base"], ffn_norm=seams["ffn"]["norm"], + gate_w=moe.gate.weight, bias=moe.gate.e_score_correction_bias, + w13=e.w13_weight, w13_s=e.w13_weight_scale, w2=e.w2_weight, + w2_s=e.w2_weight_scale, sgu=sh.gate_up_proj.weight, + sgu_s=u8(sh.gate_up_proj.weight_scale), sw2=sh.down_proj.weight, + sw2_s=u8(sh.down_proj.weight_scale), + ) # fmt: skip + + # the layer's inputs at the attention seam (the same on every rank) + gi = torch.Generator().manual_seed(11) + res0 = torch.randn(M, HC, HIDDEN, generator=gi) + res0 = ( + (res0 * torch.rsqrt(res0.pow(2).mean(-1, keepdim=True))).to(bf).to(device) + ) + x0 = (torch.randn(M, HIDDEN, generator=gi) * 0.5).to(bf).to(device) + comb0 = torch.rand(M, HC, HC, generator=gi) + 0.1 + for _ in range(20): + comb0 = comb0 / comb0.sum(-1, keepdim=True) + comb0 = comb0 / comb0.sum(-2, keepdim=True) + comb0 = comb0.to(device) + post0 = (torch.sigmoid(torch.randn(M, HC, 1, generator=gi)) * 2).to(device) + pre0 = (torch.sigmoid(torch.randn(M, HC, generator=gi)) + 1e-6).to(device) + + # vLLM's decode chain + res, post, comb, xa, attn_pre = seam("attn", res0, x0, post0, comb0, pre0) + swa_ind, swa_ptr = build_ragged_indices_from_dense( + s["swa_indices"].reshape(M, -1), s["swa_lens"] + ) + top_ind = top_ptr = None + if ratio: + valid = torch.ones(M, dtype=torch.bool, device=device) + top_ind, top_ptr, _ = compute_global_topk_ragged_indices_and_indptr( + s["topk"], s["token_to_req"], s["comp_block_table"], + s["comp_entries"], valid, + ) # fmt: skip + x8, xs = mxfp8_e4m3_quantize(xa) + qr_kv = rocm_mxfp8_block32_gemm(x8, xs, wqkv, wqkv_s, bf) + qr, kv = fused_q_kv_rmsnorm( + qr_kv[:, :Q_RANK], qr_kv[:, Q_RANK:], q_norm, kv_norm, cfg.rms_norm_eps + ) + q8, qs = mxfp8_e4m3_quantize(qr) + q = rocm_mxfp8_block32_gemm(q8, qs, wq_b, wq_b_s, bf).view(M, H, HEAD_DIM) + swa2d = s["swa_cache"].view(s["swa_cache"].shape[0], -1) + q_roped = torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( + q, kv, swa2d, s["slot_mapping"], s["positions"], cos_sin, H, + cfg.rms_norm_eps, SWA_BLOCK, False, False, True, False, + ) # fmt: skip + attn = torch.empty(M, H, HEAD_DIM, dtype=bf, device=device) + rocm_aiter_ops.triton_sparse_mla_fwd( + q_roped, s["swa_cache"], attn, HEAD_DIM**-0.5, swa_ptr, swa_ind, + attn_sink=sink, extra_kv_buffer=s["comp_cache"], + extra_kv_indptr=top_ptr, extra_kv_indices=top_ind, + ) # fmt: skip + a8 = torch.empty(M, H * HEAD_DIM, dtype=torch.float8_e4m3fn, device=device) + a_s = torch.empty(M, H * HEAD_DIM // 32, dtype=torch.uint8, device=device) + rocm_inverse_rope_mxfp8_rows(attn, s["positions"], cos_sin, ROPE, a8, a_s) + z = torch.ops.vllm.rocm_dsv41_mxfp8_wo_a_bmm(a8, a_s, wo_a, wo_a_s, G, O_RANK) + z8, zs = mxfp8_e4m3_quantize(z) + part = rocm_mxfp8_block32_gemm(z8, zs, wo_b, wo_b_s, bf).contiguous() + a = tensor_model_parallel_all_reduce(part.clone()) + res2, post2, comb2, xf, ffn_pre = seam("ffn", res, a, post, comb, attn_pre) + out = moe_out(xf) + + # ---- the whole layer: K1 + K2 + o = runner.forward( + mw, x0, res0, post0, comb0, pre0, s["positions"], s["slot_mapping"], + s["swa_cache"], s["swa_indices"], s["swa_lens"], s["token_to_req"], + topk_indices=s["topk"], comp_cache=s["comp_cache"], + comp_block_table=s["comp_block_table"], + ) # fmt: skip + off = scratch_layout(M, tp)["normed"][0] + normed = runner.scratch(M)[off : off + M * HIDDEN * 2].view(bf).view(M, HIDDEN) + tag = f"ratio {ratio}, {R} x {T}" + assert _cos(runner.res_mid[:M], res) > 0.99999, tag + for got, want in ( + (runner.post_a, post), + (runner.comb_a, comb), + (runner.pre_a, attn_pre), + ): + torch.testing.assert_close( + got[:M].view_as(want), want, rtol=1e-4, atol=1e-5 + ) + same_in = moe_out(normed) + # relative errors about 3x what the kernels show (0.3-0.5%) + assert _cos(o[1], res2) > 0.9999 and _rel(o[1], res2) < 0.01, tag + assert _cos(normed, xf) > 0.999 and _rel(normed, xf) < 0.012, tag + assert _cos(o[0], same_in) > 0.9999 and _rel(o[0], same_in) < 0.015, tag + # per token, where both FFN inputs pick the same experts + same = (top6(xf) == top6(normed)).all(1) + assert same.float().mean() >= 0.5, tag + tok = torch.nn.functional.cosine_similarity(o[0].float(), out.float(), dim=1) + assert tok[same].min() > 0.99, tag + + # ---- the FFN launch, from the same unreduced wo_b output + f = runner.ffn(mw, part, res, post, comb, attn_pre) + torch.testing.assert_close(f[1], res2, rtol=2**-8, atol=1e-6) + torch.testing.assert_close(normed, xf, rtol=2**-7, atol=1e-6) + assert _cos(f[0], out) > 0.9999 and _rel(f[0], out) < 0.015, tag + torch.accelerator.synchronize() + + whole, ffn = runner._outs(M, res0), runner._outs(M, res) + args = (s["positions"], s["slot_mapping"], s["swa_cache"], s["swa_indices"]) + meta = dict( + topk_indices=s["topk"], comp_cache=s["comp_cache"], + comp_block_table=s["comp_block_table"], + ) # fmt: skip + + def step(mw=mw, x0=x0, res0=res0, post0=post0, comb0=comb0, pre0=pre0, + args=args, s=s, meta=meta, whole=whole, part=part, res=res, + post=post, comb=comb, attn_pre=attn_pre, ffn=ffn): # fmt: skip + runner.forward( + mw, x0, res0, post0, comb0, pre0, *args, s["swa_lens"], + s["token_to_req"], **meta, outs=whole, + ) # fmt: skip + runner.ffn(mw, part, res, post, comb, attn_pre, outs=ffn) + return [*whole, *ffn] + + steps.append(step) + + # ---- the hand-off protocol under back-to-back launches: both step widths + # alternating in one graph, every replay bit for bit the first's + outs = [t for step in steps for t in step()] + torch.accelerator.synchronize() + ref = [t.clone() for t in outs] + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + for step in steps: + step() + for i in range(REPLAYS): + graph.replay() + if i % 50 == 49 or i == REPLAYS - 1: + torch.accelerator.synchronize() + bad = [ + j for j, (t, r) in enumerate(zip(outs, ref)) if not torch.equal(t, r) + ] + assert not bad, f"replay {i}: outputs {bad} differ from the first run" + + +@pytest.mark.skipif(not _on_cdna4(), reason="the mono kernels target CDNA4") +@pytest.mark.parametrize("tp", [2, 4]) +def test_mono_decode_layer_matches_vllm(monkeypatch: pytest.MonkeyPatch, tp: int): + if torch.accelerator.device_count() < tp: + pytest.skip(f"needs {tp} GPUs") + pytest.importorskip("flydsl") + multi_process_parallel(monkeypatch, tp, 1, _mono_numerics) diff --git a/vllm/envs.py b/vllm/envs.py index 851741d525..02c282fd8e 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -155,6 +155,7 @@ if TYPE_CHECKING: VLLM_ROCM_USE_AITER_FP4BMM: bool = True VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION: bool = False VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS: bool = False + VLLM_ROCM_MONO_DECODE: bool = False VLLM_ROCM_USE_AITER_TRITON_GEMM: bool = True VLLM_ROCM_USE_SKINNY_GEMM: bool = True VLLM_ROCM_FP8_PADDING: bool = True @@ -1387,6 +1388,12 @@ environment_variables: dict[str, Callable[[], Any]] = { os.getenv("VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS", "False").lower() in ("true", "1") ), + # Run the eligible layers of decode steps on the model's mono decode + # kernels: persistent launches a layer, the TP all-reduces in-kernel + # (DeepSeek-V4.1 on CDNA4). By default is disabled. + "VLLM_ROCM_MONO_DECODE": lambda: ( + os.getenv("VLLM_ROCM_MONO_DECODE", "False").lower() in ("true", "1") + ), # Whether to use aiter triton kernels for gemm ops. # By default is enabled. "VLLM_ROCM_USE_AITER_TRITON_GEMM": lambda: ( diff --git a/vllm/models/deepseek_v41/amd/model.py b/vllm/models/deepseek_v41/amd/model.py index cff72f12dc..5de173a870 100644 --- a/vllm/models/deepseek_v41/amd/model.py +++ b/vllm/models/deepseek_v41/amd/model.py @@ -14,6 +14,7 @@ from vllm.distributed import ( get_pp_group, get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, + tensor_model_parallel_all_reduce, ) from vllm.forward_context import get_forward_context, is_forward_context_available from vllm.logger import init_logger @@ -58,6 +59,7 @@ from vllm.models.common.ops.sequence_parallel import ( from vllm.models.deepseek_v4.amd.model import ( DeepseekV4MoE as DeepseekV4MoEBase, ) +from vllm.models.deepseek_v41.amd.mono_decode import MonoDecodeLayer from vllm.models.deepseek_v41.amd.rocm import DeepseekV41ROCMAiterMLAAttention from vllm.models.deepseek_v41.attention import DeepseekV4Attention from vllm.sequence import IntermediateTensors @@ -255,6 +257,12 @@ class DeepseekV4DecoderLayer(nn.Module): # attn_norm / ffn_norm into its collapse, so the separate norms are # skipped for the seams it takes. self.fuse_seam_norm = HAS_AITER_MHC_FUSED_POST_PRE_DELAYED_RMS_NORM + # Decode steps of an eligible layer may run as the mono decode layer. + self.mono = MonoDecodeLayer.create(self, vllm_config) + if self.mono is not None and self.mono.ffn_only: + # A decode step hands wo_b's TP partial to the FFN launch; forward + # reduces it on the others. + self.attn.wo_b.reduce_results = False @staticmethod def _hc_collapse(x: torch.Tensor, pre_mix: torch.Tensor) -> torch.Tensor: @@ -273,6 +281,10 @@ class DeepseekV4DecoderLayer(nn.Module): engram_hashes: torch.Tensor | None = None, engram_mask: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + if self.mono is not None: + out = self.mono(self, x, positions, residual, post_mix, res_mix, pre_mix) + if out is not None: + return out # Layer 0's attention seam projects the 2-D embedding with the folded # hc_attn_fn_broadcast instead of the 4-stream residual with hc_attn_fn. # The fused kernel only takes the latter, so that seam keeps the @@ -371,6 +383,11 @@ class DeepseekV4DecoderLayer(nn.Module): x = self.attn(positions, x, None) if self.use_sequence_parallel: x = sp_reduce_scatter(x) + elif self.mono is not None and self.mono.ffn_only: + out = self.mono.ffn(self, x, residual, post_mix, res_mix, attn_pre) + if out is not None: + return out + x = tensor_model_parallel_all_reduce(x) residual, post_mix, res_mix, x, ffn_pre = self.mhc_pre_delayed( residual, diff --git a/vllm/models/deepseek_v41/amd/mono/README.md b/vllm/models/deepseek_v41/amd/mono/README.md new file mode 100644 index 0000000000..88e40d73dd --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono/README.md @@ -0,0 +1,94 @@ +# DeepSeek-V4.1-Flash mono decode layer (gfx950, FlyDSL) + +One fused decoder layer for vLLM's decode steps on MI355X: the attention seam, +the attention, the attention all-reduce, the FFN seam, the MoE and the MoE +all-reduce, in two persistent launches a layer. The other backbone layers keep +vLLM's attention and run the rest -- the attention all-reduce, the FFN seam, +the MoE and its all-reduce -- as one launch (the FFN launch). ATOM's V4.1 mono decode +(ROCm/ATOM `atom/models/deepseek_v41/mono`) is the reference design: the seam +and MoE stages (`stages/`) and the framework under them (`common/`) are adapted +from it, each file naming its origin; the attention stages (`attention/`) are +new, reading vLLM's KV records and weights as loaded. + +## Scope + +- Layers: the whole layer on every backbone layer with the standard seam and no + compressor / indexer: 3-7, 9-13, 15-19 (compress ratio 2) and 21-23, 25-27, + 29-31, 33-35, 37-39 (ratio 1), 30 of 40. The FFN launch on the other 10: + layer 0 (embedding-broadcast seam), layer 1 (Engram at the seam) and the + compressor / indexer layers 2, 8, 14 (also Engram), 20, 24, 28, 32, 36, + whose wo_b leaves its TP reduce to it. The DSpark draft layers keep vLLM's + path. The paths interleave freely (same tensors at the layer boundary). +- Steps: decode only (FULL graph or eager), M <= 48 rows (8 requests x 6 + DSpark tokens), TP2 / TP4, `fp8_ds_mla` KV records (ROCm), text tokens (the + routing bias is `e_score_correction_bias`). +- Contract: a drop-in for `DeepseekV4DecoderLayer.forward`: + in `x` (the previous FFN output, reduced, bf16 [M, 5120]), `residual` + (bf16 [M, 4, 5120]), `post_mix` [M, 4, 1], `res_mix` [M, 4, 4], + `pre_mix` [M, 4] (f32, the previous seam's) + out the MoE output (reduced), the residual after the FFN seam, and that + seam's post / comb / pre mixes. + +## Launches + + K1 seam (stages/seam.py): slice (160 x 32 columns): fold x into the + residual (post x + sum comb r), the collapsed input (pre . R'), partial + mixes (fn . R') and sums of squares -> gate (S): 24 mixes -> pre / post + / comb (Sinkhorn 20) -> norm (S): attn_norm -> vLLM MXFP8 -> X8 + attention front (attention/front.py): wqkv, q / kv norms + KV insert, + wq_b + RoPE -> Q + K2 attention back (attention/back.py): split attention, combine + inverse + RoPE, wo_a, wo_b; the bf16 wo_b partial pushed to every rank + FFN seam: reduce (the ranks' partials summed in rank order, fp32 -> + bf16: the attention output), slice, gate, norm -> the MoE input + MoE (stages/moe.py): xq, router (bf16 GEMV, f32 logits), route + (sqrtsoftplus + bias, top-6, renormalized x 1.5), shared expert (MXFP8 + GEMVs, clamped SiLU), ug (MXFP4 x MXFP8 per expert of the step's + union), down (+ routing weight, top-k order; + shared) -> push, rank-order + sum -> out + FFN K2 without the attention back: push (160 x 32 columns) of vLLM's + unreduced wo_b output, as its bits, to every rank's attention region, + then K2's FFN seam and MoE (the same code, `layer.run_ffn`) + +## Execution model and hand-offs + +- Grid = 256 CTAs (one a CU, all resident), 512 threads a CTA, LDS a per-stage + union. A stage is a task table: CTA `b` runs tasks `b, b + 256, ...`, stages + in order; a consumer only waits on earlier stages, so no deadlock while the + grid is resident. A GEMV task issues its weight loads before it polls its + activations. +- In a launch: tagged pairs (value + tag in one 8 B store) or plain data + + flags, at device scope. The tag is the launch pair's epoch (K1 `2e + 1`, + K2 or the FFN launch `2e + 2`), a device word their CTA 0 moves on once every + CTA has marked the launch done: no clearing between steps, no host step hook, every argument + a fixed device pointer (graph-safe). +- Across ranks: one symmetric uncached buffer a rank (aiter `UncachedIpcHeap`), + every rank holding every peer's address; pushes are system-scope stores of + (value, tag) pairs. Every rank runs the same launches, so the epochs agree; + peer regions alternate by epoch parity, and a rank cannot run two steps ahead + of a peer (each step's reduce waits on every rank's push). + +## Numerics (the vLLM path each stage reproduces) + +- Seams: aiter `mhc_fused_post_pre_delayed_rmsnorm` (post fold order, fn as + bf16 hi + lo, products summed in order, Triton sigmoid / exp / rcp). +- Attention: vLLM's ROCm chain. `mxfp8(x)`: one E8M0 per 32 values, + `code = clamp(ceil(log2(amax / 448)) + 127, 0, 254)`; weights in vLLM's + row-major e4m3 [N, K] with [N/32, K/32] E8M0 scales (one shared copy); KV + records as vLLM's insert writes them (dims 0..447 e4m3 with a UE8M0 per 64, + RoPE dims bf16, 584 B); keys = the sliding window plus the top-k compressed + rows, softmax in base 2 with the attention sink folded in. +- MoE: vLLM's router (bf16 gate GEMV, f32 logits, `topk_hash_softplus_sqrt`) + and AITER a8w4 `fused_moe` (interleaved gate / up, swiglu limit 10, weights + in aiter's (16, 16) shuffle and `shuffle_scale` order); the shared expert as + vLLM's MXFP8 block-32 linears with `silu_and_mul_with_clamp`. +- All-reduces: exact, fp32 in rank order -> bf16. vLLM's decode all-reduces at + M <= 48 are below its INT4 quick-reduce threshold, so vLLM runs them exact too. + +## Validation + +- Per layer against vLLM's own modules in a TP process group (`DeepseekV4MoE`, + the mHC seams, the attention chain, the all-reduce), real checkpoint weights. +- Bench: layers of one type back to back in one HIP graph, cold weights, + us per layer, M = 1..48, TP2 / TP4; then e2e (AgentX, InferenceX#3690 + recipe). diff --git a/vllm/models/deepseek_v41/amd/mono/__init__.py b/vllm/models/deepseek_v41/amd/mono/__init__.py new file mode 100644 index 0000000000..72a6ec2ab1 --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono/__init__.py @@ -0,0 +1,5 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""DeepSeek-V4.1-Flash mono decode layer: one decoder layer of a decode step in +two persistent FlyDSL launches, its TP all-reduces in-kernel (README.md). The +host side is ``runner``; vLLM docks it in ``..mono_decode``.""" diff --git a/vllm/models/deepseek_v41/amd/mono/attention/__init__.py b/vllm/models/deepseek_v41/amd/mono/attention/__init__.py new file mode 100644 index 0000000000..da1bf8c269 --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono/attention/__init__.py @@ -0,0 +1,4 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The mono layer's attention stages: ``front`` (K1: wqkv, q / kv norms and the +KV insert, wq_b) and ``back`` (K2: split attention, combine, wo_a, wo_b).""" diff --git a/vllm/models/deepseek_v41/amd/mono/attention/back.py b/vllm/models/deepseek_v41/amd/mono/attention/back.py new file mode 100644 index 0000000000..4526f8eda7 --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono/attention/back.py @@ -0,0 +1,884 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""K2 ``back``: one DeepSeek-V4.1 attention layer from the rotated query to the +rank's wo_b partial. + +One launch a layer and rank, ``BLOCKS`` x ``THREADS``, S rows: + + split (S x chunks of CK keys): flash-decoding over one chunk of a token's + keys (K1's key table: top-k compressed rows, then window rows) for + every head of the rank. The token's q in LDS; each wave one + (head tile, 16-key tile): its fp8_ds_mla rows loaded once, + dequantized into its QK operands and kept raw in LDS; QK (bf16 + MFMA), the chunk's max / sum per head, bf16 p in LDS; PV over the + chunk's raw rows (wave w: dims 64 w ..) -> the chunk's partial + (bf16) and its (max, sum) per head, published + combine (S x H): a (token, head)'s chunks merged with the sink folded in, + bf16, inverse GPT-J RoPE of dims 448.., vLLM's MXFP8 of each 32 + -> XO (published) + wo_a (G x 1024 / 32, two 16-row tiles): the group's XO rows times wo_a + -> bf16 -> MXFP8 of the 32 rows -> X8B (published) + wo_b (320, 16 rows): X8B times wo_b -> bf16 -> ``out`` + +Hand-offs are plain device-scope data behind tagged flags (half the bytes of a +tagged pair per value). Adapted from ATOM's V4.1 ``attn_post`` (K2a), which +reads a bf16 pool and ends in a peer all-reduce. +""" + +from dataclasses import dataclass + +import flydsl.expr as fx +from aiter.ops.flydsl.kernels import buffer_ops as bo +from flydsl.expr import const_expr, gpu, range_constexpr, rocdl +from flydsl.expr.typing import T, as_ir_value + +from .device import ( + BLOCKS, + CM_DEV, + THREADS, + WAVES, + Mailbox, + bf16_round, + bf_hi, + bf_lo, + butterfly, + clamp_fp8, + first_task, + fp8_pack4, + fp8x8_bf16, + fresh, + gload, + gstore, + hw_exp2, + lane_gather, + ld_f32, + ld_i32, + mfma_bf16, + mx_code, + mx_mul, + pow2, + rsrc, + spin_until, + stamp, + traced, + xshfl, +) +from .gemv import RED, gemv_loads, gemv_mfmas, lds_row, row_sum, rows_loads, rows_store +from .plan import ( + DATA, + HALF, + HEAD_DIM, + HEAD_TILE, + KEYS, + NOPE, + ROPE, + ROWS, + WOA_ROWS, + Dims, + back_scratch, + cdiv, + front_end, + tile_rows, + token_tiles, +) + +_STREAM = fx.Stream(None) +NEG = -3.4028234663852886e38 +LOG2E = 1.4426950408889634 +REC = DATA + 8 # a raw record row in LDS: 576 data bytes + 8 scale bytes +Q_ROW = ( + HEAD_DIM // 2 + 4 +) # a q row in LDS (i32 words): 512 bf16 + 16 B, no bank aliasing +P_PAD = 2 # i32 words past a P row: 16 head rows read together hit distinct banks +# wo_a's K order: waves 0..3 / 4..7 take the two 16-row tiles, a contiguous K +# range a wave, summed in wave order +WOA_SPLIT_K = WAVES * ROWS // WOA_ROWS +BACK_POINTS = 16 # start, stage ends, split internals, GEMV internals + + +@dataclass(frozen=True) +class BackBuild: + tokens: int + tp: int + ratio: int # 0: window only; 1 / 2: top-k over a compressed cache + timeline: bool = False + + +def plus(t0, v): + return v if isinstance(t0, int) and t0 == 0 else v + t0 + + +def head_groups(d: Dims, s: int) -> int: + """Wave groups a split unit's head tiles are spread over: two (a tile each, + 64-key chunks, twice the units) while the units fit the grid, else one (a + wave's keys loaded and dequantized once for every head tile).""" + return 2 if d.head_tiles == 2 and s * cdiv(KEYS, 64) <= BLOCKS else 1 + + +def chunk_keys(d: Dims, s: int) -> int: + """A split unit's keys: a 16-key tile a wave of a head group.""" + return 16 * (WAVES // head_groups(d, s)) + + +def p_row(d: Dims, s: int) -> int: + """A P row (one head's chunk of bf16 p) in LDS, i32 words.""" + return chunk_keys(d, s) // 2 + P_PAD + + +@traced +def publish(c, region, i): + """This CTA's stores visible device-wide, then the flag pair (region, i).""" + rocdl.s_waitcnt(vmcnt=0) + gpu.barrier() + if c["tid"] == 0: + c["mb"].put(region, i, fx.Int32(1)) + + +@traced +def await_flags(c, region, n, index): + """Thread u < n polls flag ``index(u)``; then the CTA proceeds.""" + tid = c["tid"] + if tid < n: + c["mb"].poll([(region, index(fx.min(tid, n - 1)), 1)]) + gpu.barrier() + + +# ---------------------------------------------------------------- keys + + +def key_rec(c, t, k): + """(data address, scale address, valid) of token t's key k from K1's key + table: slot | 1 << 31 for a compressed row, a window slot otherwise, -1 for + none (it reads a zero record: its scores are masked, its values zeros).""" + a, ratio = c["args"], c["ratio"] + val = ld_i32(a["kt"], t * KEYS + fx.min(fx.max(k, 0), KEYS - 1)) + ok = val != fx.Int32(-1) + slot = ok.select(val & fx.Int32(0x7FFFFFFF), fx.Int32(0)) + if const_expr(ratio > 0): + is_c = ok & (val < 0) + blk = is_c.select(a["comp_block"], a["swa_block"]) + stride = is_c.select(a["comp_stride"], a["swa_stride"]) + base = is_c.select(a["comp"], a["swa"]) + else: + blk, stride, base = a["swa_block"], a["swa_stride"], a["swa"] + b = slot // blk + off = slot % blk + page = base + fx.Int64(b) * stride + z = a["zrec"] + data = ok.select(page + fx.Int64(off * DATA), z) + sc = ok.select(page + fx.Int64(blk * DATA + off * 8), z + fx.Int64(DATA)) + return data, sc, ok + + +def _scale_byte(scw, blk64): + return (scw[blk64 // 4] >> (blk64 % 4 * 8)) & 0xFF + + +def _fp8_pair_f32(word, scale, hi): + v = rocdl.cvt_scalef32_pk_f32_fp8( + fx.Vector.make_type(2, fx.Float32), + as_ir_value(fx.Int32(word)), + as_ir_value(fx.Float32(scale)), + hi, + ) + v = fx.Vector(v) + return [v[0], v[1]] + + +# ---------------------------------------------------------------- split + + +@traced +def stage_split(c, unit): + """Flash-decoding over chunk ``unit % nchunk`` of token ``unit // nchunk``, + every head of the rank -> the chunk's partial and (max, sum) per head. Wave + (head group hg, key tile kt): keys 16 kt .. of the chunk (loaded and + dequantized once) against the group's head tiles; then dims 64 w .. of the + PV over every key and head.""" + tid, lane, wave, d, s = c["tid"], c["lane"], c["wave"], c["d"], c["S"] + a, sl = c["args"], c["sl"] + nchunk = c["nchunk"] + t = unit // nchunk + ch = unit % nchunk + ck = chunk_keys(d, s) + ngrp = head_groups(d, s) + nkt = WAVES // ngrp # 16-key tiles a chunk + nht = d.head_tiles // ngrp # head tiles a wave + hgi = wave // nkt + kt = wave % nkt + g = lane // 16 + r16 = lane % 16 + k0 = ch * ck + # the key count, this lane's key and q: independent loads, out together + # (K1's key table holds -1 past a token's keys) + kv_len = ld_i32(a["klen"], t) + key = k0 + kt * 16 + r16 + data, sc, okv = key_rec(c, t, key) + qrow = d.heads * HEAD_DIM // 8 + qv = [] + for i in range_constexpr(cdiv(qrow, THREADS)): + u = fresh(tid) + THREADS * i + qv.append( + fx.Vector( + bo.buffer_load( + rsrc(a["q"]), + (t * d.q_rows) // 2 + fx.min(u, qrow - 1) * 4, + vec_width=4, + dtype=T.i32, + ) + ) + ) + if k0 < kv_len: + # a key's 448 e4m3 as 7 x 16 B (lane g: dims 64 c + 16 g ..), its 64 + # bf16 rope dims as 2 x 16 B (32 c + 8 g ..), its 7 scale bytes + scw = gload(sc, words=2) + nope = [ + gload(data + fx.Int64(64 * c4 + 16 * g), words=4) + for c4 in range(NOPE // 64) + ] + rope = [ + gload(data + fx.Int64(NOPE + (cc * 32 + 8 * g) * 2), words=4) + for cc in range(ROPE // 32) + ] + # q of every head into LDS (bf16 rows of 512 + pad), 16 B a thread + for i in range_constexpr(cdiv(qrow, THREADS)): + u = fresh(tid) + THREADS * i + if u < qrow: + fx.ptr_store( + qv[i], + sl["q"] + + (u // (HEAD_DIM // 8)) * Q_ROW + + (u % (HEAD_DIM // 8)) * 4, + ) + # the raw rows to LDS for the PV (head group 0's waves) + if hgi == 0: + rawb = sl["raw"] + (kt * 16 + r16) * (REC // 4) + for c4 in range_constexpr(NOPE // 64): + fx.ptr_store(nope[c4], rawb + (64 * c4 + 16 * g) // 4) + for cc in range_constexpr(ROPE // 32): + fx.ptr_store(rope[cc], rawb + (NOPE + (cc * 32 + 8 * g) * 2) // 4) + if g == 0: + fx.ptr_store(scw, rawb + DATA // 4) + gpu.barrier() + stamp(c, 5) + # QK: A = the keys (dequantized once), B = q of a head tile (LDS); the + # contraction's 32-dim chunk 2 c + h of the nope dims is {64 c + 16 g + + # 8 h ..} (lane g's 16 B load), q read in the same order + scale = fx.Int32(a["qk_scale"]).bitcast(fx.Float32) + st = sl["st"] + accs = [fx.Vector.filled(4, 0.0, fx.Float32) for _ in range(nht)] + for c4 in range_constexpr(NOPE // 64): + ksc = pow2(_scale_byte(scw, c4)) + for h2 in range_constexpr(2): + kw = fp8x8_bf16(nope[c4][2 * h2], nope[c4][2 * h2 + 1], ksc).bitcast( + fx.BFloat16 + ) + for ht in range_constexpr(nht): + qb = sl["q"] + ((hgi * nht + ht) * HEAD_TILE + r16) * Q_ROW + qw = fx.Vector( + fx.ptr_load( + qb + 32 * c4 + 8 * g + 4 * h2, + result_type=fx.Vector.make_type(4, fx.Int32), + ) + ) + accs[ht] = mfma_bf16(kw, qw.bitcast(fx.BFloat16), accs[ht]) + for cc in range_constexpr(ROPE // 32): + kw = rope[cc].bitcast(fx.BFloat16) + for ht in range_constexpr(nht): + qb = sl["q"] + ((hgi * nht + ht) * HEAD_TILE + r16) * Q_ROW + qw = fx.Vector( + fx.ptr_load( + qb + (NOPE + cc * 32 + 8 * g) // 2, + result_type=fx.Vector.make_type(4, fx.Int32), + ) + ) + accs[ht] = mfma_bf16(kw, qw.bitcast(fx.BFloat16), accs[ht]) + # scores of keys 4 g + i (C rows), head r16 of a tile (C column) + ok = (key < kv_len) & okv + okw = ok.select(fx.Int32(1), fx.Int32(0)) + valid = [lane_gather(okw, 4 * g + i) != 0 for i in range(4)] + svs = [] + for ht in range_constexpr(nht): + sv = [ + fx.Float32(valid[i].select(accs[ht][i] * scale, fx.Float32(NEG))) + for i in range(4) + ] + mw = butterfly( + fx.max(fx.max(sv[0], sv[1]), fx.max(sv[2], sv[3])), (32, 16), fx.max + ) + if g == 0: + fx.ptr_store(mw, st + ((hgi * nht + ht) * nkt + kt) * HEAD_TILE + r16) + svs.append(sv) + gpu.barrier() + lst = st + d.heads * nkt + for ht in range_constexpr(nht): + hti = hgi * nht + ht + m = fx.ptr_load(st + (hti * nkt) * HEAD_TILE + r16) + for j in range_constexpr(1, nkt): + m = fx.max(m, fx.ptr_load(st + (hti * nkt + j) * HEAD_TILE + r16)) + sv = svs[ht] + p = [ + fx.Float32(valid[i].select(hw_exp2(sv[i] - m), fx.Float32(0.0))) + for i in range(4) + ] + lw = (p[0] + p[2]) + (p[1] + p[3]) + lw = lw + xshfl(lw, 32) + lw = lw + xshfl(lw, 16) + pw = ( + fx.Vector.from_elements(p, fx.Float32).to(fx.BFloat16).bitcast(fx.Int32) + ) + pbase = ( + sl["p"] + (hti * HEAD_TILE + r16) * p_row(d, s) + (kt * 16 + 4 * g) // 2 + ) + fx.ptr_store(fx.Vector.from_elements([pw[0], pw[1]], fx.Int32), pbase) + if g == 0: + fx.ptr_store(lw, lst + (hti * nkt + kt) * HEAD_TILE + r16) + gpu.barrier() + # chunk stats: (max, sum) per head + if tid < d.heads: + hts = tid // HEAD_TILE + hr = tid % HEAD_TILE + mm = fx.ptr_load(st + (hts * nkt) * HEAD_TILE + hr) + ll = fx.ptr_load(lst + (hts * nkt) * HEAD_TILE + hr) + for j in range_constexpr(1, nkt): + mm = fx.max(mm, fx.ptr_load(st + (hts * nkt + j) * HEAD_TILE + hr)) + ll = ll + fx.ptr_load(lst + (hts * nkt + j) * HEAD_TILE + hr) + bo.buffer_store( + fx.Vector.from_elements([mm, ll], fx.Float32), + rsrc(c["pstat"]), + ((t * nchunk + ch) * d.heads + tid) * 2, + cache_modifier=CM_DEV, + ) + # PV: wave w dims 64 w + 4 r16 + dt (dt = 0..3: one LDS word holds the + # four of a key); the rope wave (dims 448..) reads bf16 pairs. Lane + # (r16, g) takes keys 4 g + j of a 16-key tile: A[dim][key] + is_rope = wave == WAVES - 1 + dbase = 64 * wave + sh = (wave % 4) * 8 + accs = [ + [fx.Vector.filled(4, 0.0, fx.Float32) for _ in range(4)] + for _ in range(d.head_tiles) + ] + for kt2 in range_constexpr(nkt): + rb0 = sl["raw"] + (kt2 * 16 + 4 * g) * (REC // 4) + vv = fx.Vector.filled(16, 0.0, fx.Float32) # [dt][j] + if is_rope: + rv = [] + for j in range_constexpr(4): + rw = fx.Vector( + fx.ptr_load( + rb0 + j * (REC // 4) + (NOPE + 8 * r16) // 4, + result_type=fx.Vector.make_type(2, fx.Int32), + ) + ) + rv.append([bf_lo(rw[0]), bf_hi(rw[0]), bf_lo(rw[1]), bf_hi(rw[1])]) + vv = fx.Vector.from_elements( + [rv[j][dt] for dt in range(4) for j in range(4)], fx.Float32 + ) + else: + nv = [] + for j in range_constexpr(4): + rb = rb0 + j * (REC // 4) + w = fx.Int32(fx.ptr_load(rb + (dbase + 4 * r16) // 4)) + scb = fx.Int32(fx.ptr_load(rb + DATA // 4 + wave // 4)) + scl = pow2((scb >> sh) & 0xFF) + nv.append( + _fp8_pair_f32(w, scl, False) + _fp8_pair_f32(w, scl, True) + ) + vv = fx.Vector.from_elements( + [nv[j][dt] for dt in range(4) for j in range(4)], fx.Float32 + ) + for ht2 in range_constexpr(d.head_tiles): + pb = fx.Vector( + fx.ptr_load( + sl["p"] + + (ht2 * HEAD_TILE + r16) * p_row(d, s) + + (kt2 * 16 + 4 * g) // 2, + result_type=fx.Vector.make_type(2, fx.Int32), + ) + ).bitcast(fx.Int16) + for dt in range_constexpr(4): + av = ( + fx.Vector.from_elements( + [vv[dt * 4 + j] for j in range(4)], fx.Float32 + ) + .to(fx.BFloat16) + .bitcast(fx.Int16) + ) + accs[ht2][dt] = fx.Vector( + rocdl.mfma_f32_16x16x16bf16_1k( + T.vec(4, T.f32), [av, pb, accs[ht2][dt]] + ) + ) + stamp(c, 6) + # partial[t][ch][head][dim] bf16: lane (r16, g) holds C rows 4 g + i -> + # dims 64 w + 16 g + 4 i + dt of head ht2 * 16 + r16: 32 contiguous + # bytes, two 16 B stores (4 lanes fill a 128 B line) + for ht2 in range_constexpr(d.head_tiles): + hh = ht2 * HEAD_TILE + r16 + words = [] + for i in range_constexpr(4): + v4 = fx.Vector.from_elements( + [accs[ht2][dt][i] for dt in range(4)], fx.Float32 + ) + w2 = v4.to(fx.BFloat16).bitcast(fx.Int32) + words += [w2[0], w2[1]] + off = (((t * nchunk + ch) * d.heads + hh) * HEAD_DIM + dbase + 16 * g) // 2 + for hlf in range_constexpr(2): + bo.buffer_store( + fx.Vector.from_elements(words[4 * hlf : 4 * hlf + 4], fx.Int32), + rsrc(c["part"]), + off + 4 * hlf, + cache_modifier=CM_DEV, + ) + publish(c, c["srdy"], t * nchunk + ch) + stamp(c, 7) + + +# ---------------------------------------------------------------- combine + +COMB_HEADS = WAVES # heads a combine unit: a wave each + + +@traced +def stage_combine(c, unit): + """(token, 8 heads) ``unit``, a wave a head, 8 dims a lane: the chunks + merged with the sink folded in, bf16, inverse RoPE, MXFP8 of each 32 -> XO, + published (a flag a (token, head group)).""" + lane, wave, d, s = c["lane"], c["wave"], c["d"], c["S"] + a = c["args"] + hg_n = d.heads // COMB_HEADS + t = unit // hg_n + hg = unit % hg_n + h = hg * COMB_HEADS + wave + nchunk = c["nchunk"] + ck = chunk_keys(d, s) + kv_len = ld_i32(a["klen"], t) + nlive = (kv_len + ck - 1) // ck + pos = ld_i32(a["pos"], 2 * t) + sink = ld_f32(a["sink"], h) * fx.Float32(LOG2E) + d0 = lane * 8 # this lane's 8 dims + rope_lane = d0 >= NOPE + jj = fx.max((d0 - NOPE) // 2, 0) # 4 pairs from here + cs = [ + rope_lane.select(ld_f32(a["cos_sin"], pos * ROPE + jj + i), fx.Float32(1.0)) + for i in range(4) + ] + sn = [ + rope_lane.select( + ld_f32(a["cos_sin"], pos * ROPE + HALF + jj + i), fx.Float32(0.0) + ) + for i in range(4) + ] + await_flags(c, c["srdy"], nlive, lambda u: t * nchunk + u) + m = fx.Float32(NEG) + ms, ls, vs = [], [], [] + for cc in range_constexpr(nchunk): + live = cc < nlive + stv = fx.Vector( + bo.buffer_load( + rsrc(c["pstat"]), + ((t * nchunk + cc) * d.heads + h) * 2, + vec_width=2, + dtype=T.f32, + ) + ) + pv = fx.Vector( + bo.buffer_load( + rsrc(c["part"]), + (((t * nchunk + cc) * d.heads + h) * HEAD_DIM + d0) // 2, + vec_width=4, + dtype=T.i32, + ) + ) + mc = fx.Float32(live.select(stv[0], fx.Float32(NEG))) + ms.append(mc) + ls.append(fx.Float32(live.select(stv[1], fx.Float32(0.0)))) + vals = [] + for e in range_constexpr(4): + vals += [bf_lo(pv[e]), bf_hi(pv[e])] + vs.append([fx.Float32(live.select(v, fx.Float32(0.0))) for v in vals]) + m = fx.max(m, mc) + m_final = fx.max(m, sink) + lsum = hw_exp2(sink - m_final) + tot = [fx.Float32(0.0)] * 8 + for cc in range_constexpr(nchunk): + wgt = hw_exp2(ms[cc] - m_final) + lsum = lsum + ls[cc] * wgt + tot = [tot[e] + vs[cc][e] * wgt for e in range(8)] + inv = fx.Float32(1.0) / lsum + o = [bf16_round(tot[e] * inv) for e in range(8)] + # inverse GPT-J RoPE of each pair: e' = e cos + o sin, o' = o cos - e sin + y = [] + for i in range_constexpr(4): + e, od = o[2 * i], o[2 * i + 1] + y += [e * cs[i] + od * sn[i], od * cs[i] - e * sn[i]] + amax = abs(y[0]) + for v in y[1:]: + amax = fx.max(amax, abs(v)) + amax = butterfly(amax, (1, 2), fx.max) # 32 dims = 4 lanes + code = mx_code(amax) + mul = mx_mul(code) + w0 = fp8_pack4(*[clamp_fp8(y[i] * mul) for i in range(4)]) + w1 = fp8_pack4(*[clamp_fp8(y[4 + i] * mul) for i in range(4)]) + bo.buffer_store( + fx.Vector.from_elements([w0, w1], fx.Int32), + rsrc(c["xo"]), + t * (d.q_rows // 4) + (h * HEAD_DIM + d0) // 4, + cache_modifier=CM_DEV, + ) + if lane % 4 == 0: + bo.buffer_store( + code, + rsrc(c["xos"]), + t * (d.q_rows // 32) + (h * HEAD_DIM + d0) // 32, + cache_modifier=CM_DEV, + ) + publish(c, c["crdy"], t * hg_n + hg) + + +# ---------------------------------------------------------------- wo_a / wo_b + + +@traced +def stage_woa(c, task, ops, t0, n): + """Rows 32 task .. of wo_a -> bf16 -> vLLM's MXFP8 of the 32 -> X8B.""" + tid, red, d, lane = c["tid"], c["red"], c["d"], c["lane"] + split = WOA_SPLIT_K + gemv_mfmas(c, d.group_k, c["xl"], c["xsl"], red, ops, split, tiled=True, rows=n) + gpu.barrier() + mine = tid < WOA_ROWS * n + tl = fx.min(tid // WOA_ROWS, n - 1) + t = plus(t0, tl) + row = tid % WOA_ROWS + tile = row // ROWS + y = bf16_round( + row_sum(red, row % ROWS, tl, [tile * split + w for w in range(split)]) + ) + amax = butterfly(abs(y), (1, 2, 4, 8, 16), fx.max) + code = mx_code(amax) + q = clamp_fp8(y * mx_mul(code)) + w = fp8_pack4( + q, + *[ + lane_gather(q.bitcast(fx.Int32), fx.min(lane + k, 63)).bitcast(fx.Float32) + for k in (1, 2, 3) + ], + ) + if mine & (row % 4 == 0): + bo.buffer_store( + w, + rsrc(c["x8b"]), + t * (d.o_rows // 4) + (task * WOA_ROWS + row) // 4, + cache_modifier=CM_DEV, + ) + if mine & (row == 0): + bo.buffer_store( + code, rsrc(c["x8bs"]), t * (d.o_rows // 32) + task, cache_modifier=CM_DEV + ) + gpu.barrier() + + +def woa_loads(c, task): + a = c["args"] + return gemv_loads( + c, + a["woa"], + a["woa_s"], + c["d"].group_k, + task * WOA_ROWS, + WOA_SPLIT_K, + tiled=True, + ) + + +def wob_loads(c, task, pred=None): + a = c["args"] + return gemv_loads(c, a["wob"], a["wob_s"], c["d"].o_rows, task * ROWS, pred=pred) + + +def _wob_store(a, t, col, v0, v1): + word = ( + fx.Vector.from_elements([v0, v1], fx.Float32) + .to(fx.BFloat16) + .bitcast(fx.Int32)[0] + ) + bo.buffer_store(word, rsrc(a["out"]), (t * a["out_stride"] + col) // 2) + + +@traced +def stage_wob(c, task, ops, task2, ops2, has2, t0, n, mark=False): + """Rows 16 task .. (and 16 task2 .., ``has2``) of wo_b -> bf16 -> out, or + (``c["wob_out"]``, the mono layer) the caller's hook with each bf16 pair.""" + tid, red = c["tid"], c["red"] + o_rows = c["d"].o_rows + a = c["args"] + gemv_mfmas(c, o_rows, c["xl"], c["xsl"], red, ops, rows=n) + gemv_mfmas(c, o_rows, c["xl"], c["xsl"], red + RED, ops2, rows=n) + gpu.barrier() + if const_expr(mark): + stamp(c, 15) + per = ROWS * n // 2 # bf16 pairs a task + k = tid // per + emit = c.get("wob_out") or (lambda t, col, v0, v1: _wob_store(a, t, col, v0, v1)) + if (tid < 2 * per) & ((k == 0) | has2): + u = tid % per + rp = u % (ROWS // 2) + tl = u // (ROWS // 2) + t = plus(t0, tl) + rk = red + k * RED + col = (k == 0).select(task, task2) * ROWS + 2 * rp + v0 = row_sum(rk, 2 * rp, tl) + v1 = row_sum(rk, 2 * rp + 1, tl) + emit(t, col, v0, v1) + gpu.barrier() + + +# ---------------------------------------------------------------- kernel + + +def back_plan(key): + """(nchunk, n_split, n_comb, COMB0, WOA0, WOB0, early_wob) of a build.""" + s, d = key.tokens, Dims(key.tp) + ck = chunk_keys(d, s) + nchunk = cdiv(KEYS, ck) + n_split = s * nchunk + n_comb = s * (d.heads // COMB_HEADS) + COMB0 = n_split % BLOCKS + assert d.woa_tasks <= BLOCKS + assert BLOCKS < d.wob_tasks <= 2 * BLOCKS + # the GEMVs past the attention units, so CTAs without attention work have + # their weights in flight from the start + WOA0 = (n_split + n_comb) % BLOCKS + WOB0 = (WOA0 + d.woa_tasks) % BLOCKS + # wo_b's weights before the wo_a stage only while the wo_a CTAs have no + # attention work (loads complete in issue order: else they would hold up + # wo_a's activation rows) + early_wob = n_split + n_comb + d.woa_tasks <= BLOCKS + return nchunk, n_split, n_comb, COMB0, WOA0, WOB0, early_wob + + +def back_lds_members(s: int, d: Dims) -> dict: + """K2's attention LDS structs: the split stage's and the GEMVs' (a union's + members: one at a time).""" + rows = tile_rows(s) + ck = chunk_keys(d, s) + + @fx.struct + class SplitLds: + q: fx.Array[fx.Int32, d.heads * Q_ROW, 16] + raw: fx.Array[fx.Int32, ck * REC // 4, 16] + p: fx.Array[fx.Int32, d.heads * p_row(d, s), 16] + st: fx.Array[fx.Float32, 2 * d.heads * (WAVES // head_groups(d, s)), 16] + + @fx.struct + class GemvLds: + red: fx.Array[fx.Float32, 2 * RED, 16] + xl: fx.Array[fx.Int32, rows * lds_row(d.group_k // 4), 16] + xsl: fx.Array[fx.Int32, rows * lds_row(d.group_k // 32), 16] + + return {"split": SplitLds, "gemv": GemvLds} + + +def back_lds(s: int, d: Dims): + """K2's attention LDS: the union of ``back_lds_members``.""" + m = back_lds_members(s, d) + SplitLds, GemvLds = m["split"], m["gemv"] + + @fx.union + class BackLds: + split: SplitLds # type: ignore[valid-type] + gemv: GemvLds # type: ignore[valid-type] + + return BackLds + + +def back_context(key, lds, args: dict, scratch, tag, layout=None) -> dict: + """The attention back's context: LDS views (``lds``: an allocated + ``back_lds``, or a dict of the members' views), mailbox (tag ``tag``), + scratch regions, arguments by name.""" + s, d = key.tokens, Dims(key.tp) + tid = fx.thread_idx.x + if isinstance(lds, dict): + spl, gml = lds["split"], lds["gemv"] + else: + spl, gml = lds.split.peek(), lds.gemv.peek() + c = { + "S": s, + "tid": tid, + "bid": fx.block_idx.x, + "lane": tid % 64, + "wave": tid // 64, + "red": gml.red.ptr, + "xl": gml.xl.ptr, + "xsl": gml.xsl.ptr, + "d": d, + "ratio": key.ratio, + "mb": Mailbox(tag), + "nchunk": back_plan(key)[0], + "sl": {"q": spl.q.ptr, "raw": spl.raw.ptr, "p": spl.p.ptr, "st": spl.st.ptr}, + "args": args, + } + for region, (off, _) in ( + layout or back_scratch(s, d, start=front_end(s, d)) + ).items(): + c[region] = scratch + fx.Int64(off) + return c + + +EPOCH_MARKS = 4 # word offset of the per-CTA marks in the epoch buffer + + +@traced +def epoch_begin(c, epoch, ep): + """This CTA has read the launch pair's epoch: its own mark word (no + contended atomic) := ep + 1, which ``epoch_end`` checks.""" + if c["tid"] == 0: + bo.buffer_store( + ep + 1, rsrc(epoch), EPOCH_MARKS + c["bid"], cache_modifier=CM_DEV + ) + + +@traced +def epoch_end(c, epoch, ep, reset=None): + """CTA 0, once every CTA marked this epoch (thread u checks CTA u's mark): + ``reset()`` (state later pairs must find cleared), then the epoch moved on. + The next pair's hand-offs carry a tag no earlier pair wrote; later launches' + kernels start after this one ends.""" + if c["bid"] == 0: + tid = c["tid"] + if tid < BLOCKS: + r = rsrc(epoch) + + def load_mark(): + return fx.Int32( + bo.buffer_load( + r, + EPOCH_MARKS + tid, + vec_width=1, + dtype=T.i32, + cache_modifier=CM_DEV, + ) + ) + + _ = spin_until(load_mark, lambda v: v != ep + 1) + gpu.barrier() + if tid == 0: + if const_expr(reset is not None): + reset() + gstore(epoch, ep + 1, words=1) + + +@traced +def run_back(c, key, bid): + """The attention back's stages on this CTA: split, combine, wo_a, wo_b.""" + s = key.tokens + dd = Dims(key.tp) + # plain ints: a Python object read inside traced control flow is carried + woa_tasks, groups, heads, group_k = dd.woa_tasks, dd.groups, dd.heads, dd.group_k + q_rows, o_rows, wob_tasks = dd.q_rows, dd.o_rows, dd.wob_tasks + a = c["args"] + tiles = token_tiles(s) + _, n_split, n_comb, COMB0, WOA0, WOB0, early_wob = back_plan(key) + stamp(c, 0) + for unit in range(first_task(bid, 0), n_split, BLOCKS): + stage_split(c, unit) + stamp(c, 1) + for unit in range(first_task(bid, COMB0), n_comb, BLOCKS): + stage_combine(c, unit) + gpu.barrier() + stamp(c, 2) + # the GEMV weights, every CTA right after its attention work: its wo_a + # task (if any), then its one or two wo_b tasks + woa_task = first_task(bid, WOA0) + has_woa = woa_task < woa_tasks + woa_t = fx.min(woa_task, woa_tasks - 1) + ops = gemv_loads( + c, + a["woa"], + a["woa_s"], + group_k, + woa_t * WOA_ROWS, + WOA_SPLIT_K, + tiled=True, + pred=has_woa, + ) + wob_task = first_task(bid, WOB0) + wob_second = wob_task + BLOCKS + has_second = wob_second < wob_tasks + + def wob_weights(): + return ( + wob_loads(c, wob_task), + wob_loads(c, fx.min(wob_second, wob_tasks - 1), pred=has_second), + ) + + if const_expr(early_wob): + ops1, ops2 = wob_weights() + if const_expr(c.get("tl") is not None): + rocdl.s_waitcnt(vmcnt=0) + gpu.barrier() + stamp(c, 8) + # wo_a: every token tile of its group's XO rows, the next tile's in flight + if has_woa: + grp = woa_t // (woa_tasks // groups) + await_flags(c, c["crdy"], s, lambda u: u * (heads // COMB_HEADS) + grp) + stamp(c, 9) + xo = ( + c["xo"], + c["xos"], + q_rows // 4, + grp * group_k // 4, + group_k // 4, + q_rows // 32, + grp * group_k // 32, + group_k // 32, + ) + regs = rows_loads(c, *xo, *tiles[0]) + for ti in range_constexpr(len(tiles)): + t0, n = tiles[ti] + rows_store(c, regs, group_k // 4, group_k // 32, c["xl"], c["xsl"], n) + if const_expr(ti + 1 < len(tiles)): + regs = rows_loads(c, *xo, *tiles[ti + 1]) + stage_woa(c, woa_t, ops, t0, n) + stamp(c, 10) + publish(c, c["ardy"], woa_t) + stamp(c, 3) + # wo_b + if const_expr(not early_wob): + ops1, ops2 = wob_weights() + stamp(c, 11) + await_flags(c, c["ardy"], woa_tasks, lambda u: u) + stamp(c, 12) + x8b = ( + c["x8b"], + c["x8bs"], + o_rows // 4, + 0, + o_rows // 4, + o_rows // 32, + 0, + o_rows // 32, + ) + regs = rows_loads(c, *x8b, *tiles[0]) + for ti in range_constexpr(len(tiles)): + t0, n = tiles[ti] + rows_store(c, regs, o_rows // 4, o_rows // 32, c["xl"], c["xsl"], n) + if const_expr(ti == 0): + stamp(c, 14) + if const_expr(ti + 1 < len(tiles)): + regs = rows_loads(c, *x8b, *tiles[ti + 1]) + stage_wob( + c, + wob_task, + ops1, + fx.min(wob_second, wob_tasks - 1), + ops2, + has_second, + t0, + n, + mark=ti == 0, + ) + if const_expr(ti == 0): + stamp(c, 13) + stamp(c, 4) diff --git a/vllm/models/deepseek_v41/amd/mono/attention/device.py b/vllm/models/deepseek_v41/amd/mono/attention/device.py new file mode 100644 index 0000000000..ec230cd707 --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono/attention/device.py @@ -0,0 +1,504 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Device primitives of the DeepSeek-V4.1 attention mega kernel. + +Wave reductions and shuffles, hardware math, bf16 / fp8 / MX packing, the +scaled MFMA, buffer and 64-bit global accesses, and the tagged-pair mailbox. +Adapted from ATOM's mono decode framework (``atom/mono/device``), itself ported +from the FlyDSL GLM-5 layer kernel. +""" + +import flydsl.expr as fx +from aiter.ops.flydsl.kernels import buffer_ops as bo +from aiter.ops.flydsl.kernels.dpp_utils import update_dpp_i32 +from aiter.ops.flydsl.kernels.kernels_common import kernel_signature +from flydsl._mlir import ir +from flydsl._mlir.dialects import llvm +from flydsl.compiler.ast_rewriter import ASTRewriter +from flydsl.expr import const_expr, gpu, range_constexpr, rocdl +from flydsl.expr import math as fmath +from flydsl.expr.typing import T, as_ir_value + +# The execution model: one CTA a CU of an MI355X, all resident. +BLOCKS = 256 +THREADS = 512 +WAVES = THREADS // 64 + +# gfx95x cache policy bits (LLVM CPol): SC0 = 1, SC1 = 16, NT = 2. SC1 alone is +# device scope (past the per-XCD L2s); NT streams data read once. +CM_DEV = 16 +CM_NT = 2 + +FP8_MAX = 448.0 +FP8, FP4 = 0, 4 # f8f6f4 MFMA operand formats +UNIT_SCALE = 127 # E8M0 code of 1.0 +FLT_MIN = 1.1754943508222875e-38 + +POLL_MAX = 12 # mailbox pairs polled a batch + + +def traced(fn): + """FlyDSL's AST rewrite for a helper, so its ``if`` / ``while`` / ``for`` on + traced values lower to scf ops (kernels rewrite only their own body).""" + return ASTRewriter.transform(fn) + + +def kernel_symbol(stem, **params): + """Kernel symbol that profilers demangle to ``aiter::_``.""" + name = f"{stem}_{kernel_signature(**params)}" + return f"_ZN5aiter{len(name)}{name}E" + + +def first_task(bid, base): + """CTA ``bid``'s first task of a stage placed from CTA ``base``, then every + BLOCKS-th.""" + return (bid - base % BLOCKS + BLOCKS) % BLOCKS + + +# ----------------------------------------------------------------------------- memory + + +def rsrc(addr, nbytes=None): + return bo.create_buffer_resource_from_addr(addr, num_records_bytes=nbytes) + + +def ld_i32(ptr, i, cm=0): + return fx.Int32( + bo.buffer_load(rsrc(ptr), i, vec_width=1, dtype=T.i32, cache_modifier=cm) + ) + + +def ld_f32(ptr, i): + return fx.Float32(bo.buffer_load(rsrc(ptr), i, vec_width=1, dtype=T.f32)) + + +def ld_bf(ptr, i): + return fx.Float32( + fx.BFloat16(bo.buffer_load(rsrc(ptr), i, vec_width=1, dtype=T.bf16)) + ) + + +def ld_i64(ptr, i): + """Element i of an i64 buffer (two dwords).""" + w = fx.Vector(bo.buffer_load(rsrc(ptr), i * 2, vec_width=2, dtype=T.i32)) + return (fx.Int64(w[1]) << 32) | fx.Int64(fx.Uint32(w[0])) + + +def _gptr(addr): + return llvm.IntToPtrOp( + ir.Type.parse("!llvm.ptr<1>"), as_ir_value(fx.Int64(addr)) + ).result + + +def gload(addr, words=4, nt=False): + """``words`` i32 (1, 2 or 4) at a 64-bit global address: a cache past 4 GiB is + reachable only this way (a buffer offset is 32 bits).""" + i32 = ir.IntegerType.get_signless(32) + ty = i32 if words == 1 else ir.VectorType.get([words], i32) + v = llvm.LoadOp(ty, _gptr(addr), alignment=4 * words, nontemporal=nt or None).result + return fx.Int32(v) if words == 1 else fx.Vector(v) + + +def atomic_add_dev(addr, value): + """Device-scope atomic add on the global i32 at ``addr``; the old value. + Relaxed: an acquire / release here would write back and invalidate L2.""" + return fx.Int32( + fx.atomic_add( + _gptr(addr), + fx.Int32(value), + syncscope="agent", + ordering=fx.AtomicOrdering.Monotonic, + ) + ) + + +def gstore(addr, value, words=4): + """Store ``words`` i32 (an Int32 or a Vector) at a 64-bit global address.""" + v = as_ir_value(fx.Int32(value) if words == 1 else value) + llvm.StoreOp(v, _gptr(addr), alignment=4 * words) + + +# ----------------------------------------------------------------------------- math + + +def fresh(v): + """``v`` through an opaque move, so the lane math built on it is computed at + its use instead of being hoisted out of a task loop and held in registers.""" + return fx.Int32( + llvm.inline_asm( + T.i32, + [fx.Int32(v).ir_value()], + "v_mov_b32 $0, $1", + "=v,v", + has_side_effects=True, + ) + ) + + +def tile_map(tid, width): + """(tid // width, tid % width) for 0 <= tid < THREADS without a division: + shifts for a power of two, else a multiply-shift exact on that range.""" + if width & (width - 1) == 0: + sh = width.bit_length() - 1 + return tid >> sh, tid & (width - 1) + m = -(-(1 << 20) // width) + assert all((t * m) >> 20 == t // width for t in range(THREADS)), width + row = (tid * m) >> 20 + return row, tid - row * width + + +def uniform(v): + return fx.Int32(rocdl.readfirstlane(T.i32, fx.Int32(v).ir_value())) + + +def _hw_f32(name, x): + return fx.Float32( + llvm.call_intrinsic(T.f32, name, [fx.Float32(x).ir_value()], [], []) + ) + + +def hw_rsq(x): + return _hw_f32("llvm.amdgcn.rsq.f32", x) + + +def hw_rcp(x): + return _hw_f32("llvm.amdgcn.rcp.f32", x) + + +def hw_exp2(x): + return _hw_f32("llvm.amdgcn.exp2.f32", x) + + +def fma(a, b, c): + return fx.Float32(fmath.fma(fx.Float32(a), fx.Float32(b), fx.Float32(c))) + + +def div_rn(x, d, r): + """``x / d`` rounded to nearest (IEEE), given ``r = RN(1 / d)``.""" + q = x * r + return fma(fma(-q, d, x), r, q) + + +def permlane_swap(off, x, y): + swap = rocdl.permlane32_swap if off == 32 else rocdl.permlane16_swap + pr = swap( + llvm.StructType.get_literal([T.i32, T.i32]), + as_ir_value(fx.Int32(x)), + as_ir_value(fx.Int32(y)), + False, + False, + ) + return tuple(fx.Int32(llvm.extractvalue(T.i32, pr, [j])) for j in range(2)) + + +def lane_gather(v, src_lane): + """Lane ``src_lane``'s value of v (ds_bpermute).""" + return fx.Int32( + rocdl.ds_bpermute(T.i32, fx.Int32(src_lane) * 4, fx.Int32(v).ir_value()) + ) + + +def lane_f32(v, src_lane): + return lane_gather(v.bitcast(fx.Int32), src_lane).bitcast(fx.Float32) + + +def readlane(v, src_lane): + bits = v.bitcast(fx.Int32) if isinstance(v, fx.Float32) else fx.Int32(v) + r = fx.Int32(rocdl.readlane(T.i32, bits, fx.Int32(src_lane))) + return r.bitcast(fx.Float32) if isinstance(v, fx.Float32) else r + + +def xshfl(v, off): + """Value of lane ``lane ^ off`` (DPP in a row, permlane swaps across).""" + assert off in (1, 2, 4, 8, 16, 32), off + if off >= 16: + return v.shuffle_xor(off, 64) + is_f = isinstance(v, fx.Float32) + x = v.bitcast(fx.Int32) if is_f else fx.Int32(v) + if off == 8: + y = fx.Int32(update_dpp_i32(x, x, 0x118, 0xF, 0xC, False)) + y = fx.Int32(update_dpp_i32(y, x, 0x108, 0xF, 0x3, False)) + elif off == 4: + y = fx.Int32(update_dpp_i32(x, x, 0x114, 0xF, 0xA, False)) + y = fx.Int32(update_dpp_i32(y, x, 0x104, 0xF, 0x5, False)) + elif off == 2: + y = fx.Int32(update_dpp_i32(x, x, 0x4E, 0xF, 0xF, False)) + else: + y = fx.Int32(update_dpp_i32(x, x, 0xB1, 0xF, 0xF, False)) + return y.bitcast(fx.Float32) if is_f else y + + +def xred(v, off, op): + if off < 16: + return op(v, xshfl(v, off)) + is_f = isinstance(v, fx.Float32) + x = v.bitcast(fx.Int32) if is_f else fx.Int32(v) + a, b = permlane_swap(off, x, x) + if is_f: + return op(a.bitcast(fx.Float32), b.bitcast(fx.Float32)) + return op(type(v)(a), type(v)(b)) + + +def _add(a, b): + return a + b + + +def butterfly(v, offsets, op=_add): + """Butterfly reduction over the xor ``offsets``, in that order.""" + for off in offsets: + v = xred(v, off, op) + return v + + +def wave_sum(v): + return butterfly(v, (32, 16, 8, 4, 2, 1)) + + +def wave_max(v): + return butterfly(v, (32, 16, 8, 4, 2, 1), fx.max) + + +def row_shr(v, lane, k): + """Lane - k of the same 16-lane row, else 0 (DPP row_shr).""" + moved = lane_gather(v.bitcast(fx.Int32), fx.max(lane - k, 0)).bitcast(fx.Float32) + return fx.Float32((lane % 16 >= k).select(moved, fx.Float32(0.0))) + + +# ----------------------------------------------------------------------------- packing + + +def bf16_round(a): + return fx.Float32(fx.Float32(a).to(fx.BFloat16)) + + +def bf16_pair(a, b): + """Two f32 -> one f32-typed word holding (bf16(a), bf16(b)).""" + return ( + fx.Vector.from_elements([a, b], fx.Float32) + .to(fx.BFloat16) + .bitcast(fx.Float32)[0] + ) + + +def bf_lo(w): + return (w << 16).bitcast(fx.Float32) + + +def bf_hi(w): + return (w & fx.Int32(-65536)).bitcast(fx.Float32) + + +def fp8_pack4(a, b, c, d): + """Four f32 (in range) -> one dword of four E4M3 bytes, RNE.""" + lo = fx.Int32(rocdl.cvt_pk_fp8_f32(T.i32, a, b, fx.Int32(0), False)) + return fx.Int32(rocdl.cvt_pk_fp8_f32(T.i32, c, d, lo, True)) + + +def fp8x8_bf16(w0, w1, scale): + """Eight E4M3 of two dwords times an f32 ``scale`` -> eight bf16 as a + 4-dword Vector (exact for a power-of-two scale in range).""" + out = [] + for w in (w0, w1): + for hi in (False, True): + p = rocdl.cvt_scalef32_pk_bf16_fp8( + T.bf16x2, as_ir_value(fx.Int32(w)), as_ir_value(fx.Float32(scale)), hi + ) + out.append(fx.Vector(p).bitcast(fx.Int32)[0]) + return fx.Vector.from_elements(out, fx.Int32) + + +def clamp_fp8(v): + return fx.max(fx.min(v, fx.Float32(FP8_MAX)), fx.Float32(-FP8_MAX)) + + +def pow2(code): + """The fp32 power of two a biased exponent (1..254) stands for.""" + return (code << 23).bitcast(fx.Float32) + + +def rcp_pow2(code): + """1 / ``pow2(code)``, exact for code in [1, 253].""" + return ((254 - code) << 23).bitcast(fx.Float32) + + +def ceil_exp(v): + """Biased exponent of the smallest power of two >= v (v > 0 finite, normal).""" + bits = v.bitcast(fx.Int32) + return ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0).select( + fx.Int32(1), fx.Int32(0) + ) + + +def mx_code(amax): + """VLLM's ROCm MXFP8 scale code of a 32-group: clamp(ceil(log2(amax / 448)) + + 127, 0, 254), amax floored at FLT_MIN.""" + a = fx.max(amax, fx.Float32(FLT_MIN)) + q = div_rn(a, fx.Float32(FP8_MAX), fx.Float32(1.0 / FP8_MAX)) + return fx.min(fx.max(ceil_exp(q), fx.Int32(0)), fx.Int32(254)) + + +def mx_mul(code): + """2^(127 - code) as fp32, for code in [0, 254] (the quant multiplier).""" + return ((254 - code) << 23).bitcast(fx.Float32) + + +# ----------------------------------------------------------------------------- MFMA + + +def mfma_scaled(a, b, c, sa, sb, a_fmt=FP8, b_fmt=FP8): + """16x16x128 scaled MFMA (gfx950 f8f6f4): lane l (row l % 16) holds K bytes + 16 (l // 16) .. +16 and 64 + 16 (l // 16) .. +16 of the 128-step, and the + E8M0 of K block l // 16.""" + return fx.Vector( + rocdl.mfma_scale_f32_16x16x128_f8f6f4( + T.vec(4, T.f32), [a, b, c, a_fmt, b_fmt, 0, sa, 0, sb] + ) + ) + + +def mfma_bf16(a, b, c): + """16x16x32 bf16 MFMA: lane l (row / col l % 16) holds K 8 (l // 16) .. +8.""" + return fx.Vector(rocdl.mfma_f32_16x16x32_bf16(T.vec(4, T.f32), [a, b, c])) + + +# ----------------------------------------------------------------------------- mailbox + + +@traced +def _spin(load_all, pending): + v = load_all() + while pending(v): + rocdl.s_sleep(1) + v = load_all() + return v + + +def spin_until(load, pending): + """``load()`` again (s_sleep between) while ``pending(value)``; the value.""" + return _spin(load, pending) + + +class Mailbox: + """Tagged-pair hand-off between the CTAs of one launch (device scope). + + Every 32-bit value is stored next to the launch's tag in one 8-byte store; a + consumer polls until every tag matches, so a hand-off is one round trip with + no separate flag. The tag is ``epoch * 64 + layer + 1`` (``epoch`` a device + counter bumped once a step): a pair from an earlier layer or step never + carries a live tag, and nothing is cleared between launches.""" + + def __init__(self, tag): + self.tag = tag + + def put(self, base, i, v): + bits = v.bitcast(fx.Int32) if isinstance(v, fx.Float32) else fx.Int32(v) + bo.buffer_store( + fx.Vector.from_elements([bits, self.tag], fx.Int32), + rsrc(base), + i * 2, + cache_modifier=CM_DEV, + ) + + def put_bf(self, base, i, vs): + """Elements i .. i + len(vs) (2 or 4, i aligned) as bf16 pairs, one store.""" + words = [] + for j in range_constexpr(len(vs) // 2): + words += [bf16_pair(vs[2 * j], vs[2 * j + 1]).bitcast(fx.Int32), self.tag] + bo.buffer_store( + fx.Vector.from_elements(words, fx.Int32), + rsrc(base), + i, + cache_modifier=CM_DEV, + ) + + def put_words(self, base, i, words): + """Raw 32-bit words (1 or 2) at pairs i.. (fp8 words must come this way).""" + vals = [] + for w in words: + vals += [fx.Int32(w), self.tag] + bo.buffer_store( + fx.Vector.from_elements(vals, fx.Int32), + rsrc(base), + i * 2, + cache_modifier=CM_DEV, + ) + + def poll(self, specs, batch=POLL_MAX): + """Batched poll of ``specs`` = [(base, pair index, npairs in {1, 2})]: all + re-loaded while any tag is stale. Returns the value words.""" + if const_expr(len(specs) == 0): + return [] + if const_expr(len(specs) > batch): + return self.poll(specs[:batch], batch) + self.poll(specs[batch:], batch) + tag = self.tag + + def load_all(): + words = [] + for b, i, n in specs: + w = fx.Vector( + bo.buffer_load( + rsrc(b), + fx.Int32(i) * 2, + vec_width=2 * n, + dtype=T.i32, + cache_modifier=CM_DEV, + ) + ) + words += [w[e] for e in range(2 * n)] + return fx.Vector.from_elements(words, fx.Int32) + + nw = sum(2 * n for _, _, n in specs) + + def pending(v): + bad = v[1] != tag + for e in range_constexpr(3, nw, 2): + bad = bad | (v[e] != tag) + return bad + + v = _spin(load_all, pending) + outs, e = [], 0 + for _, _, n in specs: + outs.append([v[e + 2 * q] for q in range(n)]) + e += 2 * n + return outs + + +LOAD_BATCH = 4 + + +@traced +def batched_rounds(tid, n, load, use, batch=LOAD_BATCH): + """A CTA's strided loop over v = tid + THREADS i < n: ``batch`` rounds' loads + issued before their uses; a round past n loads element n - 1, uses nothing.""" + rounds = -(-n // THREADS) + for b0 in range_constexpr(0, rounds, batch): + idx = range(b0, min(b0 + batch, rounds)) + got = [load(fx.min(tid + THREADS * i, n - 1)) for i in idx] + for i, g in zip(idx, got): + v = tid + THREADS * i + if v < n: + use(v, g) + + +def barrier(): + gpu.barrier() + + +# ----------------------------------------------------------------------------- timeline + + +def memrealtime(): + """s_memrealtime: the 100 MHz device clock.""" + return fx.Int64(llvm.call_intrinsic(T.i64, "llvm.amdgcn.s.memrealtime", [], [], [])) + + +@traced +def stamp(c, point): + """A timeline build's stamp: thread 0 of the CTA stores the clock at + ``point`` (``c["tl"]`` [BLOCKS, points] i64; None: not a timeline build).""" + if const_expr(c.get("tl") is None): + return + if c["tid"] == 0: + t = memrealtime() + w = fx.Vector.from_elements([fx.Int32(t), fx.Int32(t >> 32)], fx.Int32) + bo.buffer_store(w, rsrc(c["tl"]), (c["bid"] * c["tl_points"] + point) * 2) diff --git a/vllm/models/deepseek_v41/amd/mono/attention/front.py b/vllm/models/deepseek_v41/amd/mono/attention/front.py new file mode 100644 index 0000000000..3ac9e3fe9b --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono/attention/front.py @@ -0,0 +1,662 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""K1 ``front``: one DeepSeek-V4.1 attention layer from its normed input row to +the rotated query and the step's KV rows in the sliding-window cache. + +One launch a layer and rank, ``BLOCKS`` CTAs x ``THREADS`` threads, S rows. +Every CTA issues its weight loads (its wqkv part and its wq_b tasks) at the +kernel start, before it waits on anything; the stages: + + xq (S > 16: S x 2 K halves, on the 32 CTAs without a wqkv part): the + bf16 input's half -> vLLM's MXFP8 -> X8 (published, a flag each) + kt (S): the token's attention keys for K2 -- top-k compressed slots + through the kv source's block table, then window slots -> KT + wqkv_a (112 row tiles x 2 K halves, 16 rows): the input's MXFP8 half -- + quantized in the CTA (S <= 16) or X8's -- a token tile at a time + (the next tile's rows in flight) times wqkv's -> fp32 partials -> + QKVP + qkv (S): the halves summed -> bf16; q_norm -> bf16 -> MXFP8 -> QX8 + (published, a flag a token); kv_norm -> bf16 -> GPT-J RoPE -> bf16 + -> the fp8_ds_mla record at the token's slot + wq_b (H x 512 / 16, 16 rows): QX8 times wq_b -> bf16 -> GPT-J RoPE of + each head's dims 448.. -> q (global, bf16) + +Each elementwise step follows vLLM's ROCm path (``mxfp8_quantize``, +``FusedQKVRMSNorm``, ``fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert``); +the GEMV reductions differ from the Triton GEMMs' only in their order. +Adapted from ATOM's V4.1 ``attn_pre`` (K1), less its mHC seam. +""" + +from dataclasses import dataclass + +import flydsl.expr as fx +from aiter.ops.flydsl.kernels import buffer_ops as bo +from flydsl.expr import const_expr, gpu, range_constexpr, rocdl +from flydsl.expr.typing import T + +from .device import ( + BLOCKS, + CM_DEV, + THREADS, + Mailbox, + bf16_round, + bf_hi, + bf_lo, + butterfly, + ceil_exp, + clamp_fp8, + div_rn, + first_task, + fp8_pack4, + fresh, + gstore, + hw_rsq, + ld_f32, + ld_i32, + mx_code, + mx_mul, + rsrc, + stamp, + tile_map, + traced, + xshfl, +) +from .gemv import RED, gemv_loads, gemv_mfmas, lds_row, row_sum, rows_loads, rows_store +from .plan import ( + DATA, + EPS, + HALF, + HEAD_DIM, + HIDDEN, + KEYS, + KV_DIM, + NOPE, + Q_RANK, + QKV_ROWS, + ROPE, + ROWS, + TILE, + TOPK, + WINDOW, + WQKV_KPARTS, + WQKV_TASKS, + Dims, + front_scratch, + tile_rows, + token_tiles, +) + +_STREAM = fx.Stream(None) +X_WORDS = HIDDEN // 4 +X_GROUPS = HIDDEN // 32 +QX_WORDS = Q_RANK // 4 +KSTEPS = HIDDEN // 128 // WQKV_KPARTS # a wqkv part's 128-steps +KWORDS = KSTEPS * 32 +KGROUPS = KSTEPS * 4 +XH = HIDDEN // WQKV_KPARTS # a K half's values +XH_COLS = XH // 8 // 2 # 16 B items (8 values) a half of a K half +XH_RPR = THREADS // XH_COLS # half-rows a round +XQ_ITEMS = XH // 8 # an xq unit's 16 B items (one a thread) +FRONT_POINTS = 6 # start, kt done, wqkv done, qkv done, wq_b done, (spare) + + +@dataclass(frozen=True) +class FrontBuild: + tokens: int + tp: int + ratio: int # the layer's compress ratio (0: no top-k keys) + timeline: bool = False + + +def plus(t0, v): + return v if isinstance(t0, int) and t0 == 0 else v + t0 + + +@traced +def publish(c, region, i): + """This CTA's stores visible device-wide, then the flag pair (region, i).""" + rocdl.s_waitcnt(vmcnt=0) + gpu.barrier() + if c["tid"] == 0: + c["mb"].put(region, i, fx.Int32(1)) + + +@traced +def await_flags(c, region, n, index): + """Thread u < n polls flag ``index(u)``; then the CTA proceeds.""" + tid = c["tid"] + if tid < n: + c["mb"].poll([(region, index(fx.min(tid, n - 1)), 1)]) + gpu.barrier() + + +# ---------------------------------------------------------------- xq + + +def xq_loads(c, unit, pred): + """(token, K half) ``unit``'s bf16 input: its load only, 8 values a thread + (threads past the half mirror its last item).""" + a = c["args"] + xr = rsrc(a["x"], pred.select(fx.Int64(0xFFFFFFFF), fx.Int64(0))) + col = fx.min(c["tid"], XQ_ITEMS - 1) + e = (unit >> 1) * a["x_stride"] + (unit & 1) * XH + col * 8 + return fx.Vector(bo.buffer_load(xr, e >> 1, vec_width=4, dtype=T.i32)) + + +@traced +def stage_xq(c, unit, w, live): + """``xq_loads``' values -> vLLM's MXFP8 -> X8 words / X8S codes of the + unit (published by the caller).""" + tid, lane = c["tid"], c["lane"] + vals = [] + for d in range_constexpr(4): + vals += [bf_lo(w[d]), bf_hi(w[d])] + amax = abs(vals[0]) + for v in vals[1:]: + amax = fx.max(amax, abs(v)) + amax = butterfly(amax, (1, 2), fx.max) + code = mx_code(amax) + mul = mx_mul(code) + w0 = fp8_pack4(*[clamp_fp8(vals[q] * mul) for q in range(4)]) + w1 = fp8_pack4(*[clamp_fp8(vals[4 + q] * mul) for q in range(4)]) + t = unit >> 1 + part = unit & 1 + if live & (tid < XQ_ITEMS): + bo.buffer_store( + fx.Vector.from_elements([w0, w1], fx.Int32), + rsrc(c["x8"]), + t * X_WORDS + part * KWORDS + tid * 2, + cache_modifier=CM_DEV, + ) + if lane % 4 == 0: + bo.buffer_store( + code, + rsrc(c["x8s"]), + t * X_GROUPS + part * KGROUPS + (tid >> 2), + cache_modifier=CM_DEV, + ) + + +# ---------------------------------------------------------------- kt + + +@traced +def stage_kt(c, t): + """Token t's keys for K2: ``KT[t][k] = slot | 1 << 31`` for its top-k compressed + rows (k < ntopk), the window slot for k in [ntopk, kv_len), -1 for an + invalid or absent key; KLEN[t] = kv_len (0 for a pad row), KLEN[S + t] = + ntopk.""" + tid = c["tid"] + a, ratio = c["args"], c["ratio"] + slot = ld_i32(a["slot"], 2 * t) + live = slot >= 0 + nswa = live.select(ld_i32(a["swa_lens"], t), fx.Int32(0)) + if const_expr(ratio > 0): + pos = ld_i32(a["pos"], 2 * t) + ntopk = live.select(fx.min((pos + 1) // ratio, fx.Int32(TOPK)), fx.Int32(0)) + req = ld_i32(a["t2r"], t) + else: + ntopk = fx.Int32(0) + kv_len = ntopk + nswa + for i in range_constexpr((KEYS + THREADS - 1) // THREADS): + k = fresh(tid) + THREADS * i + sidx = fx.min(fx.max(k - ntopk, 0), WINDOW - 1) + s_slot = ld_i32(a["swa_idx"], t * WINDOW + sidx) + val = (k < kv_len).select(s_slot, fx.Int32(-1)) + if const_expr(ratio > 0): + local = ld_i32(a["topk"], t * TOPK + fx.min(k, TOPK - 1)) + lc = fx.max(local, 0) + e = a["comp_block"] + blk = ld_i32(a["comp_bt"], req * a["bt_stride"] + lc // e) + c_slot = (blk * e + lc % e) | fx.Int32(-2147483648) + c_val = (local >= 0).select(c_slot, fx.Int32(-1)) + val = (k < ntopk).select(c_val, val) + if k < KEYS: + bo.buffer_store(val, rsrc(a["kt"]), t * KEYS + k) + if tid == 0: + bo.buffer_store(kv_len, rsrc(a["klen"]), t) + bo.buffer_store(ntopk, rsrc(a["klen"]), c["S"] + t) + + +# ---------------------------------------------------------------- wqkv_a + + +def x_half_loads(c, part, pred, t0, n): + """Tokens t0 .. t0 + n's bf16 input, K half ``part``: its loads only (16 B a + lane-item, 8 values; a thread one column of [2 n half-rows, 160]), masked + off on CTAs without a wqkv part.""" + a = c["args"] + xr = rsrc(a["x"], pred.select(fx.Int64(0xFFFFFFFF), fx.Int64(0))) + hr0, col = tile_map(c["tid"], XH_COLS) + e0 = part * XH + col * 8 + out = [] + for i in range((2 * n + XH_RPR - 1) // XH_RPR): + hr = fx.min(hr0 + XH_RPR * i, 2 * n - 1) + e = (t0 + (hr >> 1)) * a["x_stride"] + (hr & 1) * (XH // 2) + e0 + out.append(fx.Vector(bo.buffer_load(xr, e >> 1, vec_width=4, dtype=T.i32))) + return out + + +@traced +def quant_x_half(c, regs, n): + """``x_half_loads``' values -> vLLM's MXFP8 rows of the K half in LDS (a + 32-group = 4 lanes).""" + tid, lane = c["tid"], c["lane"] + x_row, s_row = lds_row(KWORDS), lds_row(KGROUPS) + hr0, col = tile_map(tid, XH_COLS) + act = tid < XH_RPR * XH_COLS + for i in range_constexpr(len(regs)): + hr = hr0 + XH_RPR * i + w = regs[i] + vals = [] + for d in range_constexpr(4): + vals += [bf_lo(w[d]), bf_hi(w[d])] + amax = abs(vals[0]) + for v in vals[1:]: + amax = fx.max(amax, abs(v)) + amax = butterfly(amax, (1, 2), fx.max) + code = mx_code(amax) + mul = mx_mul(code) + w0 = fp8_pack4(*[clamp_fp8(vals[q] * mul) for q in range(4)]) + w1 = fp8_pack4(*[clamp_fp8(vals[4 + q] * mul) for q in range(4)]) + tok = hr >> 1 + half = hr & 1 + if act & (hr < 2 * n): + fx.ptr_store( + fx.Vector.from_elements([w0, w1], fx.Int32), + c["xl"] + (tok * x_row + half * (KWORDS // 2) + col * 2), + ) + if lane % 4 == 0: + fx.ptr_store( + code, c["xsl"] + (tok * s_row + half * (KGROUPS // 2) + (col >> 2)) + ) + gpu.barrier() + + +@traced +def stage_wqkv(c, task, part, ops, t0, n, live): + """Rows 16 task .. of wqkv over K half ``part`` against tokens t0 .. t0 + n + -> fp32 partials QKVP[part] (``live``: a CTA with a wqkv part).""" + tid, red = c["tid"], c["red"] + gemv_mfmas(c, HIDDEN, c["xl"], c["xsl"], red, ops, rows=n, nsteps=KSTEPS) + gpu.barrier() + if (tid < ROWS * n) & live: + r = tid % ROWS + tl = tid // ROWS + t = plus(t0, tl) + v = row_sum(red, r, tl) + c["mb"].put(c["qkvp"], (part * c["S"] + t) * QKV_ROWS + task * ROWS + r, v) + gpu.barrier() + + +# ---------------------------------------------------------------- q / kv norms + + +@traced +def stage_qkv(c, t): + """Token t: the K halves summed -> bf16; q_norm -> bf16 -> MXFP8 (QX8); + kv_norm -> bf16 -> RoPE -> bf16 -> the fp8_ds_mla record at + slot_mapping[t]. 256 threads (8 q columns and 2 kv columns each); threads + 256.. mirror thread 255 and store nothing.""" + tid, lane, wave, red = c["tid"], c["lane"], c["wave"], c["red"] + a, s = c["args"], c["S"] + tt = fx.min(tid, 255) + mine = tid < 256 + qc = [fx.min(tt * 8 + j, Q_RANK - 1) for j in range(8)] + # independent loads first: they overlap the partials' wait + qw = [c["ldbf"](a["qn_w"], qc[j]) for j in range(8)] + kw = [c["ldbf"](a["kvn_w"], tt * 2 + j) for j in range(2)] + pos = ld_i32(a["pos"], 2 * t) + slot = ld_i32(a["slot"], 2 * t) + jr = tt - NOPE // 2 + is_rope = jr >= 0 + jj = fx.max(jr, 0) + cs = is_rope.select(ld_f32(a["cos_sin"], pos * ROPE + jj), fx.Float32(1.0)) + sn = is_rope.select(ld_f32(a["cos_sin"], pos * ROPE + HALF + jj), fx.Float32(0.0)) + specs = [] + for p in range_constexpr(2): + specs += [(c["qkvp"], (p * s + t) * QKV_ROWS + qc[j], 1) for j in range(8)] + specs += [ + (c["qkvp"], (p * s + t) * QKV_ROWS + Q_RANK + tt * 2 + j, 1) + for j in range(2) + ] + got = c["mb"].poll(specs) + vals = [ + bf16_round(got[j][0].bitcast(fx.Float32) + got[10 + j][0].bitcast(fx.Float32)) + for j in range(10) + ] + xq = [((tt * 8 + j) < Q_RANK).select(vals[j], fx.Float32(0.0)) for j in range(8)] + xk = vals[8:10] + ssq = xq[0] * xq[0] + for j in range_constexpr(1, 8): + ssq = ssq + xq[j] * xq[j] + ssq = butterfly(ssq, (32, 16, 8, 4, 2, 1)) + ssk = butterfly(xk[0] * xk[0] + xk[1] * xk[1], (32, 16, 8, 4, 2, 1)) + if (lane == 0) & mine: + fx.ptr_store(ssq, red + wave) + fx.ptr_store(ssk, red + 4 + wave) + gpu.barrier() + tq = (fx.ptr_load(red + 0) + fx.ptr_load(red + 1)) + ( + fx.ptr_load(red + 2) + fx.ptr_load(red + 3) + ) + tk = (fx.ptr_load(red + 4) + fx.ptr_load(red + 5)) + ( + fx.ptr_load(red + 6) + fx.ptr_load(red + 7) + ) + rq = hw_rsq( + div_rn(tq, fx.Float32(float(Q_RANK)), fx.Float32(1.0 / Q_RANK)) + + fx.Float32(EPS) + ) + rk = hw_rsq( + div_rn(tk, fx.Float32(float(KV_DIM)), fx.Float32(1.0 / KV_DIM)) + + fx.Float32(EPS) + ) + # q: bf16 normed row, then vLLM's MXFP8 (a 32-group = 4 threads) + ys = [bf16_round((xq[j] * rq) * qw[j]) for j in range(8)] + amax = abs(ys[0]) + for y in ys[1:]: + amax = fx.max(amax, abs(y)) + code = mx_code(butterfly(amax, (1, 2), fx.max)) + mul = mx_mul(code) + if mine & (tt * 8 < Q_RANK): + w0 = fp8_pack4(*[clamp_fp8(ys[i] * mul) for i in range(4)]) + w1 = fp8_pack4(*[clamp_fp8(ys[4 + i] * mul) for i in range(4)]) + bo.buffer_store( + fx.Vector.from_elements([w0, w1], fx.Int32), + rsrc(c["qx8"]), + t * QX_WORDS + 2 * tt, + cache_modifier=CM_DEV, + ) + if tt % 4 == 0: + bo.buffer_store( + code, + rsrc(c["qx8s"]), + t * (Q_RANK // 32) + tt // 4, + cache_modifier=CM_DEV, + ) + # kv: this thread's pair (2 tt, 2 tt + 1) + kn = [bf16_round((xk[j] * rk) * kw[j]) for j in range(2)] + e0 = bf16_round(kn[0] * cs - kn[1] * sn) + e1 = bf16_round(kn[0] * sn + kn[1] * cs) + # record quant of the nope dims: a 64-block = 32 threads (half a wave) + amax = butterfly(fx.max(abs(e0), abs(e1)), (1, 2, 4, 8, 16), fx.max) + rcode = ceil_exp( + div_rn( + fx.max(amax, fx.Float32(1e-4)), fx.Float32(448.0), fx.Float32(1.0 / 448.0) + ) + ) + inv = mx_mul(rcode) + q0 = clamp_fp8(e0 * inv) + q1 = clamp_fp8(e1 * inv) + n0 = xshfl(q0, 1) + n1 = xshfl(q1, 1) + word = fp8_pack4(q0, q1, n0, n1) + rope_word = ( + fx.Vector.from_elements([e0, e1], fx.Float32) + .to(fx.BFloat16) + .bitcast(fx.Int32)[0] + ) + blk = a["swa_block"] + sblock = slot // blk + soff = slot % blk + base = a["swa"] + fx.Int64(sblock) * fx.Int64(a["swa_stride"]) + rec = base + fx.Int64(soff * DATA) + if mine & (tt % 32 == 0) & (tt < NOPE // 2): + fx.ptr_store(rcode.bitcast(fx.Float32), red + 8 + tt // 32) + gpu.barrier() + if mine & (slot >= 0): + if tt < NOPE // 2: + if tt % 2 == 0: + gstore(rec + fx.Int64(2 * tt), word, words=1) + else: + gstore(rec + fx.Int64(NOPE + 4 * (tt - NOPE // 2)), rope_word, words=1) + if tt == 0: + sb = [fx.ptr_load(red + 8 + b).bitcast(fx.Int32) for b in range(7)] + w0 = sb[0] | (sb[1] << 8) | (sb[2] << 16) | (sb[3] << 24) + w1 = sb[4] | (sb[5] << 8) | (sb[6] << 16) + sc = base + fx.Int64(blk * DATA + soff * 8) + gstore(sc, fx.Vector.from_elements([w0, w1], fx.Int32), words=2) + publish(c, c["qrdy"], t) + + +# ---------------------------------------------------------------- wq_b + + +def wqb_rope(c, task0, pos_t, n): + """(cos, sin) of this thread's pair in ``stage_wqb``'s epilogue over a tile + of n rows (``pos_t``: its token's position; every task of the CTA shares + the pair's dims: tasks BLOCKS apart are whole heads apart).""" + a = c["args"] + rp = c["tid"] % (ROWS // 2) + dim = (task0 * ROWS) % HEAD_DIM + 2 * rp + jj = fx.max((dim - NOPE) >> 1, 0) + return ld_f32(a["cos_sin"], pos_t * ROPE + jj), ld_f32( + a["cos_sin"], pos_t * ROPE + HALF + jj + ) + + +def wqb_pos(c, t0, n): + """This thread's token position for ``stage_wqb``'s epilogue of tile + t0 .. t0 + n.""" + per = ROWS * n // 2 + tl = fx.min((c["tid"] % per) // (ROWS // 2), n - 1) + return ld_i32(c["args"]["pos"], 2 * (t0 + tl)) + + +@traced +def stage_wqb(c, task0, bops, t0, n, rope): + """Rows 16 (task0 + BLOCKS k) .. of the rank's H x 512 query rows, every + task k of the CTA: GEMVs -> bf16 -> GPT-J RoPE of a head's dims 448.. + (fp32) -> bf16 -> q.""" + tid, red, d = c["tid"], c["red"], c["d"] + a = c["args"] + nt = len(bops) + for k in range_constexpr(nt): + gemv_mfmas(c, Q_RANK, c["qxl"], c["qxsl"], red + k * RED, bops[k], rows=n) + gpu.barrier() + per = ROWS * n // 2 # bf16 pairs a task + if tid < nt * per: + k = tid // per + u = tid % per + rp = u % (ROWS // 2) + tl = u // (ROWS // 2) + t = plus(t0, tl) + r0 = (task0 + BLOCKS * k) * ROWS + 2 * rp + rk = red + k * RED + v0 = bf16_round(row_sum(rk, 2 * rp, tl)) + v1 = bf16_round(row_sum(rk, 2 * rp + 1, tl)) + is_rope = (task0 * ROWS) % HEAD_DIM + 2 * rp >= NOPE + cs, sn = rope + e0 = is_rope.select(v0 * cs - v1 * sn, v0) + e1 = is_rope.select(v0 * sn + v1 * cs, v1) + word = ( + fx.Vector.from_elements([e0, e1], fx.Float32) + .to(fx.BFloat16) + .bitcast(fx.Int32)[0] + ) + bo.buffer_store(word, rsrc(a["q"]), (t * d.q_rows + r0) // 2) + gpu.barrier() + + +# ---------------------------------------------------------------- kernel + + +def front_lds(s: int, d: Dims): + """K1's LDS (the attention front's stages).""" + rows = tile_rows(s) + wqb_per_cta = d.wqb_tasks // BLOCKS + + @fx.struct + class FrontLds: + red: fx.Array[fx.Float32, wqb_per_cta * RED + 16, 16] + qxl: fx.Array[fx.Int32, rows * lds_row(QX_WORDS), 16] + qxsl: fx.Array[fx.Int32, rows * lds_row(Q_RANK // 32), 16] + xl: fx.Array[fx.Int32, rows * lds_row(KWORDS), 16] + xsl: fx.Array[fx.Int32, rows * lds_row(KGROUPS), 16] + + return FrontLds + + +def front_context(key, lds, args: dict, scratch, tag, layout=None) -> dict: + """The front stages' context: LDS views, mailbox (tag ``tag``), scratch + regions, arguments (``args``: the kernel's, by name).""" + s, d = key.tokens, Dims(key.tp) + tid = fx.thread_idx.x + c = { + "S": s, + "tid": tid, + "bid": fx.block_idx.x, + "lane": tid % 64, + "wave": tid // 64, + "red": lds.red.ptr, + "xl": lds.xl.ptr, + "xsl": lds.xsl.ptr, + "qxl": lds.qxl.ptr, + "qxsl": lds.qxsl.ptr, + "mb": Mailbox(tag), + "d": d, + "ratio": key.ratio, + "ldbf": lambda ptr, i: fx.Float32( + fx.BFloat16(bo.buffer_load(rsrc(ptr), i, vec_width=1, dtype=T.bf16)) + ), + "args": args, + } + for region, (off, _) in (layout or front_scratch(s, d)).items(): + c[region] = scratch + fx.Int64(off) + return c + + +@traced +def run_front(c, key, bid, x8_given=False, before_wqkv=None, after_wqkv=None): + """The front's stages on this CTA. ``x8_given``: a stage before this one + (the mono layer's seam norm) publishes X8 / XRDY for every token (no input + quantization here). ``before_wqkv()`` runs ahead of the wqkv weights' loads + (a latency-critical stage of the caller's), ``after_wqkv()`` right after.""" + s = key.tokens + wqb_tasks = Dims(key.tp).wqb_tasks + a = c["args"] + tid = c["tid"] + tiles = token_tiles(s) + wqb_per_cta = wqb_tasks // BLOCKS + assert wqb_tasks % BLOCKS == 0 + n_wqkv = WQKV_TASKS * WQKV_KPARTS # 224 + assert n_wqkv <= BLOCKS + n_free = BLOCKS - n_wqkv # CTAs without a wqkv part: xq, kt + QKV0 = 0 # on wqkv CTAs + WQB0 = 0 + in_cta_quant = s <= TILE and not x8_given + xq_stage = s > TILE and not x8_given + xq_units = WQKV_KPARTS * s + xq_per = (xq_units + n_free - 1) // n_free + stamp(c, 0) + + # ---- Loads complete in issue order, so every CTA issues its latency- + # critical loads before its later stages' weights: a wqkv CTA its first + # input tile and its wqkv part, the next tiles while it quantizes, then + # its wq_b tasks; a CTA with a key-table task those after that work. + # The q / kv norms run on wqkv CTAs, their partials' wait behind the + # wq_b weights that are due by then anyway. + wqkv_unit = first_task(bid, 0) + has_wqkv = wqkv_unit < n_wqkv + task = fx.min(wqkv_unit, n_wqkv - 1) % WQKV_TASKS + part = fx.min(wqkv_unit, n_wqkv - 1) // WQKV_TASKS + wqb_task = first_task(bid, WQB0) + free = bid - n_wqkv + is_free = free >= 0 + if const_expr(in_cta_quant): + xregs = [x_half_loads(c, part, has_wqkv, *tiles[0])] + if const_expr(xq_stage): + xunits = [free + n_free * i for i in range(xq_per)] + xlive = [is_free & (u < xq_units) for u in xunits] + xq_regs = [ + xq_loads(c, fx.min(fx.max(u, 0), xq_units - 1), lv) + for u, lv in zip(xunits, xlive) + ] + if const_expr(before_wqkv is not None): + before_wqkv() + kops = gemv_loads( + c, + a["wqkv"], + a["wqkv_s"], + HIDDEN, + task * ROWS, + kbeg=part * KSTEPS, + nsteps=KSTEPS, + pred=has_wqkv, + ) + if const_expr(after_wqkv is not None): + after_wqkv() + if const_expr(xq_stage): # noqa: SIM102 (a build-time and a run-time test) + if is_free: + for i in range_constexpr(xq_per): + stage_xq( + c, fx.min(fx.max(xunits[i], 0), xq_units - 1), xq_regs[i], xlive[i] + ) + rocdl.s_waitcnt(vmcnt=0) + gpu.barrier() + if tid < xq_per: + u = free + n_free * tid + if u < xq_units: + c["mb"].put(c["xrdy"], u, fx.Int32(1)) + for t in range(is_free.select(free, fx.Int32(s)), s, n_free): + stage_kt(c, t) + stamp(c, 1) + poss = [wqb_pos(c, t0, n) for t0, n in tiles] + + # ---- wqkv_a: one (row tile, K half) a CTA, every token tile; the wq_b + # weights issued behind the input rows' loads + def wqb_weights(): + return [ + gemv_loads(c, a["wqb"], a["wqb_s"], Q_RANK, (wqb_task + BLOCKS * i) * ROWS) + for i in range(wqb_per_cta) + ] + + if const_expr(in_cta_quant): + if has_wqkv: + quant_x_half(c, xregs[0], tiles[0][1]) + bops = wqb_weights() + if has_wqkv: + stage_wqkv(c, task, part, kops, *tiles[0], has_wqkv) + else: + await_flags(c, c["xrdy"], s, lambda u: WQKV_KPARTS * u + part) + xr = ( + c["x8"], + c["x8s"], + X_WORDS, + part * KWORDS, + KWORDS, + X_GROUPS, + part * KGROUPS, + KGROUPS, + ) + regs = rows_loads(c, *xr, tiles[0][0], tiles[0][1], pred=has_wqkv) + if const_expr(len(tiles) == 1): + bops = wqb_weights() + for ti in range_constexpr(len(tiles)): + t0, n = tiles[ti] + rows_store(c, regs, KWORDS, KGROUPS, c["xl"], c["xsl"], n) + if const_expr(ti + 1 < len(tiles)): + t1, n1 = tiles[ti + 1] + regs = rows_loads(c, *xr, t1, n1, pred=has_wqkv) + if const_expr(ti + 2 == len(tiles)): + bops = wqb_weights() + stage_wqkv(c, task, part, kops, t0, n, has_wqkv) + stamp(c, 2) + + # ---- q / kv norms, KV insert + for t in range(first_task(bid, QKV0), s, BLOCKS): + stage_qkv(c, t) + stamp(c, 3) + + # ---- wq_b: every task's rows of a token tile at a time, the next + # tile's rows in flight + ropes = [wqb_rope(c, wqb_task, poss[ti], tiles[ti][1]) for ti in range(len(tiles))] + await_flags(c, c["qrdy"], s, lambda u: u) + qx = (c["qx8"], c["qx8s"], QX_WORDS, 0, QX_WORDS, Q_RANK // 32, 0, Q_RANK // 32) + regs = rows_loads(c, *qx, *tiles[0]) + for ti in range_constexpr(len(tiles)): + t0, n = tiles[ti] + rows_store(c, regs, QX_WORDS, Q_RANK // 32, c["qxl"], c["qxsl"], n) + if const_expr(ti + 1 < len(tiles)): + regs = rows_loads(c, *qx, *tiles[ti + 1]) + stage_wqb(c, wqb_task, bops, t0, n, ropes[ti]) + stamp(c, 4) diff --git a/vllm/models/deepseek_v41/amd/mono/attention/gemv.py b/vllm/models/deepseek_v41/amd/mono/attention/gemv.py new file mode 100644 index 0000000000..9471a1e3f4 --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono/attention/gemv.py @@ -0,0 +1,234 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""MXFP8 GEMV of the mega kernel: 16 rows of a row-major e4m3 weight [N, K] +(vLLM's loaded layout, E8M0 scales in 32x32 blocks [N / 32, K / 32]) against +every token of a tile, its MXFP8 rows in LDS, on the gfx950 scaled MFMA. + +A task may cover a window of K (``kbeg`` .. ``kbeg + nsteps`` 128-steps): a +split-K part, its activations in LDS holding just that window. The weight +operands are loaded first (all in flight, before the task waits on its +activations); wave w takes the window's steps w, w + 8, ... (or, ``split``, a +contiguous range of waves 0 .. split - 1), its 16 x 16 partial in ``red``, +summed in wave order by ``row_sum``. Adapted from ATOM's +``attn_pre.gemv_fp8_*``, which reads aiter's (16, 16)-preshuffled layout. +""" + +import flydsl.expr as fx +from aiter.ops.flydsl.kernels import buffer_ops as bo +from flydsl.expr import const_expr, gpu, range_constexpr +from flydsl.expr.typing import T + +from .device import THREADS, WAVES, mfma_scaled, rsrc, tile_map, traced + +# a token's LDS row of MXFP8 words / E8M0 codes, padded by 16 B so 16 rows read +# by one MFMA operand do not share a bank group +LDS_PAD = 4 +RED = WAVES * 64 * 4 # one task's accumulators in ``red`` (f32) + + +def lds_row(words): + return words + LDS_PAD + + +def lds_at(u, row): + """Element u of a span of rows ``row`` wide, at the ``lds_row`` stride.""" + return u + u // row * LDS_PAD + + +def _steps(c, nsteps, split, tiled): + """This wave's steps of the window: [(local step, live)].""" + wave = c["wave"] + if const_expr(split is None): + per = (nsteps + WAVES - 1) // WAVES + ragged = nsteps % WAVES != 0 + return [ + (wave + WAVES * i, (wave + WAVES * i < nsteps) if ragged else None) + for i in range(per) + ] + assert split <= WAVES and nsteps % split == 0 + per = nsteps // split + part = wave % split if tiled else fx.min(wave, split - 1) + return [(part * per + i, None) for i in range(per)] + + +def _step_load(c, r_w, r_ws, k, row0, st, pred=None): + """Global step ``st``'s weight operand of rows row0 .. +16: two 16 B chunks a + lane (K bytes 16 j .. and 64 + 16 j .. of the step, row row0 + lane % 16) + and the E8M0 byte of its K block. ``pred`` false: masked (out of range: no + memory traffic, zeros).""" + lane = c["lane"] + j = lane // 16 + row = row0 + lane % 16 + base = row * k + 128 * st + 16 * j + lo = fx.Vector(bo.buffer_load(r_w, base // 4, vec_width=4, dtype=T.i32)) + hi = fx.Vector(bo.buffer_load(r_w, (base + 64) // 4, vec_width=4, dtype=T.i32)) + kg = k // 32 + wsi = (row0 // 32) * kg + st * 4 + j + word = fx.Int32(bo.buffer_load(r_ws, wsi // 4, vec_width=1, dtype=T.i32)) + return lo, hi, (word >> (wsi % 4 * 8)) & 0xFF + + +def gemv_loads( + c, w, ws, k, row0, split=None, tiled=False, kbeg=0, nsteps=None, pred=None +): + """The weight operands of rows row0 .. +16 (``tiled``: row0 + 16 (wave // + split)) over steps kbeg .. kbeg + nsteps for this wave, all in flight + (``pred`` false: masked off, no memory traffic).""" + nsteps = k // 128 if nsteps is None else nsteps + if const_expr(pred is None): + r_w, r_ws = rsrc(w), rsrc(ws) + else: + # a masked-off CTA reads through a zero-size buffer: zeros, no traffic + nb = pred.select(fx.Int64(0xFFFFFFFF), fx.Int64(0)) + r_w, r_ws = rsrc(w, nb), rsrc(ws, nb) + if const_expr(tiled): + row0 = row0 + 16 * (c["wave"] // split) + return [ + _step_load( + c, + r_w, + r_ws, + k, + row0, + kbeg + (st if live is None else fx.min(st, nsteps - 1)), + pred, + ) + for st, live in _steps(c, nsteps, split, tiled) + ] + + +def _step_mfma(c, nsteps, xl, xsl, st, ops, acc, live, rows): + lane = c["lane"] + lo, hi, sa = ops + x_row, s_row = lds_row(nsteps * 32), lds_row(nsteps * 4) + j = lane // 16 + col = fx.min(lane % 16, rows - 1) + if const_expr(live is not None): + st = fx.min(st, nsteps - 1) + av = fx.Vector.from_elements( + [lo[d] for d in range(4)] + [hi[d] for d in range(4)], fx.Int32 + ) + xlo = fx.Vector( + fx.ptr_load( + xl + (col * x_row + st * 32 + j * 4), + result_type=fx.Vector.make_type(4, fx.Int32), + ) + ) + xhi = fx.Vector( + fx.ptr_load( + xl + (col * x_row + st * 32 + 16 + j * 4), + result_type=fx.Vector.make_type(4, fx.Int32), + ) + ) + xv = fx.Vector.from_elements( + [xlo[d] for d in range(4)] + [xhi[d] for d in range(4)], fx.Int32 + ) + sb = fx.ptr_load(xsl + (col * s_row + st * 4 + j)) + if const_expr(live is not None): + # a step past the end adds zero at unit scale + av = fx.Vector.from_elements( + [live.select(av[d], fx.Int32(0)) for d in range(8)], fx.Int32 + ) + sa = live.select(sa, fx.Int32(127)) + sb = live.select(sb, fx.Int32(127)) + return mfma_scaled(av, xv, acc, sa, sb) + + +def gemv_mfmas(c, k, xl, xsl, red, ops, split=None, tiled=False, rows=16, nsteps=None): + """``gemv_loads``' operands against ``rows`` tokens' MXFP8 rows of the + window in LDS (``xl`` words, ``xsl`` codes, rows ``lds_row`` of the window) + -> ``red[wave][lane]`` (a 16 x 16 C: row 4 (lane // 16) + i, token lane % 16).""" + lane, wave = c["lane"], c["wave"] + nsteps = k // 128 if nsteps is None else nsteps + acc = fx.Vector.filled(4, 0.0, fx.Float32) + for i, (st, live) in enumerate(_steps(c, nsteps, split, tiled)): + acc = _step_mfma(c, nsteps, xl, xsl, st, ops[i], acc, live, rows) + if const_expr(split is not None and not tiled): + acc = fx.Vector.from_elements( + [(wave < split).select(acc[e], fx.Float32(0.0)) for e in range(4)], + fx.Float32, + ) + fx.ptr_store(acc, red + (wave * 64 + lane) * 4) + + +def row_sum(red, r, t, waves=None): + """Row r, token t of the waves' accumulators, summed in wave order.""" + tot = fx.Float32(0.0) + for w in waves if waves is not None else range(WAVES): + tot = tot + fx.ptr_load(red + ((w * 64 + 16 * (r // 4) + t) * 4 + r % 4)) + return tot + + +def span_items(width, n): + """A [n, width] span of 16 B items: rows a round of THREADS threads, rounds.""" + wq = width // 4 + rpr = THREADS // wq + assert rpr >= 1, width + return wq, rpr, (n + rpr - 1) // rpr + + +def rows_loads( + c, + words_base, + codes_base, + stride_w, + off_w, + words, + stride_g, + off_g, + groups, + t0, + n, + pred=None, +): + """Tokens t0 .. t0 + n's published MXFP8 row slices -- ``words`` i32 from + ``off_w`` and ``groups`` codes (an i32 each) from ``off_g``, rows + ``stride_w`` / ``stride_g`` i32 apart -- loaded (issued, not waited on), + 16 B a lane-item: a thread one column, rows a round apart. Plain loads: + their flags are seen, and nothing in the launch read these lines before. + ``pred`` false: masked off (zeros, no memory traffic).""" + tid = c["tid"] + nb = None if pred is None else pred.select(fx.Int64(0xFFFFFFFF), fx.Int64(0)) + out = [] + for base, stride, off, width in ( + (words_base, stride_w, off_w, words), + (codes_base, stride_g, off_g, groups), + ): + wq, rpr, rounds = span_items(width, n) + row0, col = tile_map(tid, wq) + a0 = col * 4 + off + t0 * stride + r = rsrc(base, nb) + regs = [] + for i in range(rounds): + row = row0 + rpr * i + if (i + 1) * rpr > n or rpr * wq < THREADS: + row = fx.min(row, n - 1) + regs.append( + fx.Vector( + bo.buffer_load(r, row * stride + a0, vec_width=4, dtype=T.i32) + ) + ) + out.append(regs) + return out + + +@traced +def _store_span(c, regs, width, dst, n): + tid = c["tid"] + wq, rpr, _ = span_items(width, n) + row0, col = tile_map(tid, wq) + lrow = lds_row(width) + base = dst + (row0 * lrow + col * 4) + for i in range_constexpr(len(regs)): + if const_expr((i + 1) * rpr <= n and rpr * wq == THREADS): + fx.ptr_store(regs[i], base + i * rpr * lrow) + else: + if (row0 + rpr * i < n) & (tid < rpr * wq): + fx.ptr_store(regs[i], base + i * rpr * lrow) + + +def rows_store(c, regs, words, groups, xl, xsl, n): + """``rows_loads``' registers into LDS rows 0 .. n (``lds_row`` strides).""" + _store_span(c, regs[0], words, xl, n) + _store_span(c, regs[1], groups, xsl, n) + gpu.barrier() diff --git a/vllm/models/deepseek_v41/amd/mono/attention/plan.py b/vllm/models/deepseek_v41/amd/mono/attention/plan.py new file mode 100644 index 0000000000..58865ac7d0 --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono/attention/plan.py @@ -0,0 +1,173 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Geometry of the DeepSeek-V4.1 attention mega kernel: model widths, a TP +rank's share of them, task counts and the scratch (mailbox) layout.""" + +from dataclasses import dataclass + +HIDDEN = 5120 +Q_RANK = 1280 +KV_DIM = 512 +QKV_ROWS = Q_RANK + KV_DIM # 1792 +HEAD_DIM = 512 +NOPE = 448 +ROPE = 64 +HALF = ROPE // 2 +N_HEADS = 64 +O_GROUPS = 8 +O_RANK = 1024 +WINDOW = 128 +TOPK = 512 +EPS = 1e-20 + +ROWS = 16 # GEMV rows a task +TILE = 16 # tokens a GEMV pass (an MFMA's N) +HEAD_TILE = 16 # heads an attention MFMA takes +WOA_ROWS = 32 # a wo_a task's rows: one MX group of its output +MAX_TOKENS = 48 + +# KV record (V4 fp8_ds_mla): 576 B data rows (448 e4m3 | 64 bf16), then 8 B of +# scales a token (7 UE8M0 + pad) after the block's rows +DATA = 576 +SCALE_BYTES = 8 + +# attention: a split is 16 keys; a token has at most TOPK + WINDOW keys +BK = 16 +SPLITS = 64 # split slots a (token, head tile): 1024 keys >= TOPK + WINDOW +COLS = 32 # PV columns a group (one MX group of the output) +GROUPS = HEAD_DIM // COLS # 16 + +PAIR = 8 # a mailbox (value, tag) pair + + +def cdiv(a, b): + return -(-a // b) + + +def round_up(a, b): + return cdiv(a, b) * b + + +@dataclass(frozen=True) +class Dims: + tp: int + + def __post_init__(self): + assert self.tp in (2, 4), self.tp + + @property + def heads(self) -> int: + return N_HEADS // self.tp + + @property + def head_tiles(self) -> int: + return self.heads // HEAD_TILE + + @property + def groups(self) -> int: + return O_GROUPS // self.tp + + @property + def group_k(self) -> int: + """wo_a's K a group: its eight heads' outputs.""" + return self.heads // self.groups * HEAD_DIM # 4096 + + @property + def o_rows(self) -> int: + return self.groups * O_RANK + + @property + def q_rows(self) -> int: + return self.heads * HEAD_DIM + + @property + def wqb_tasks(self) -> int: + return self.q_rows // ROWS + + @property + def woa_tasks(self) -> int: + return self.o_rows // WOA_ROWS + + @property + def wob_tasks(self) -> int: + return HIDDEN // ROWS # 320 + + +WQKV_TASKS = QKV_ROWS // ROWS # 112 +WQKV_KPARTS = 2 # wqkv's K split: 224 (row tile, K half) parts +KEYS = TOPK + WINDOW # a token's attention keys at most + + +def token_tiles(s): + return [(t0, min(TILE, s - t0)) for t0 in range(0, s, TILE)] + + +def tile_rows(s): + return min(s, TILE) + + +def pair_layout(items, start=0, align=256): + out, off = {}, start + for name, pairs in items: + off = round_up(off, align) + out[name] = (off, pairs * PAIR) + off += pairs * PAIR + return out + + +def byte_layout(items, start=0, align=256): + out, off = {}, start + for name, nbytes in items: + off = round_up(off, align) + out[name] = (off, nbytes) + off += nbytes + return out + + +def front_scratch(s: int, d: Dims) -> dict: + """K1's scratch regions -> (byte offset, bytes): published data (plain + words behind a flag) and mailbox pairs.""" + return byte_layout( + [ + ("x8", s * HIDDEN), # e4m3 words (S > 16: quantized once, shared) + ("x8s", s * HIDDEN // 32 * 4), # an i32 code a group + ("xrdy", WQKV_KPARTS * s * PAIR), # a flag a (token, K half) + ("qkvp", WQKV_KPARTS * s * QKV_ROWS * PAIR), + ("qx8", s * Q_RANK), # e4m3 words + ("qx8s", s * Q_RANK // 32 * 4), # an i32 code a group + ("qrdy", s * PAIR), + ] + ) + + +def back_scratch(s: int, d: Dims, start: int = 0) -> dict: + """K2's scratch regions, laid out past ``start``: published data (plain) + and flag pairs.""" + ck = ( + 64 if d.head_tiles == 2 and s * cdiv(KEYS, 64) <= 256 else 128 + ) # back.chunk_keys + nchunk = cdiv(KEYS, ck) + return byte_layout( + [ + ("part", s * nchunk * d.heads * HEAD_DIM * 2), # bf16 partials + ("pstat", s * nchunk * d.heads * 8), # (max, sum) f32 + ("srdy", s * nchunk * PAIR), + ("xo", s * d.q_rows), # e4m3 words + ("xos", s * d.q_rows // 32 * 4), # an i32 code a group + ("crdy", s * d.heads // 8 * PAIR), + ("x8b", s * d.o_rows), + ("x8bs", s * d.o_rows // 32 * 4), + ("ardy", d.woa_tasks * PAIR), + ], + start=start, + ) + + +def front_end(s: int, d: Dims) -> int: + """Where K2's regions start: past K1's, which share the launch tag.""" + return round_up(max(o + n for o, n in front_scratch(s, d).values()), 256) + + +def scratch_bytes(s: int, d: Dims) -> int: + b = back_scratch(s, d, start=front_end(s, d)) + return round_up(max(o + n for o, n in b.values()), 4096) diff --git a/vllm/models/deepseek_v41/amd/mono/common/__init__.py b/vllm/models/deepseek_v41/amd/mono/common/__init__.py new file mode 100644 index 0000000000..397642c6b4 --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono/common/__init__.py @@ -0,0 +1,9 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at 97359d6, MIT License, +# Copyright (C) 2026, Advanced Micro Devices, Inc.: +# atom/mono/__init__.py +"""The mono kernels' shared mechanisms: ``ops`` (wave / math primitives, the TP +group's peer buffers), ``sync`` (the tagged mailbox), ``mx`` (FP8 / MX +arithmetic), ``plan`` (execution model, layouts, build keys) and +``peer_memory`` (symmetric peer buffers).""" diff --git a/vllm/models/deepseek_v41/amd/mono/common/mx.py b/vllm/models/deepseek_v41/amd/mono/common/mx.py new file mode 100644 index 0000000000..9c228abdbf --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono/common/mx.py @@ -0,0 +1,78 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at 97359d6, MIT License, +# Copyright (C) 2026, Advanced Micro Devices, Inc.: +# atom/mono/device/mx.py +"""FP8 / MX arithmetic the mono kernels share: the E4M3 range, E8M0 scale +codes, the byte of a scale in aiter's ``shuffle_scale`` layout, an FP4 weight +tile's operand loads and the scaled MFMA that takes them. + +An E8M0 code is a biased exponent; ``pow2(code)`` is the scale it stands for.""" + +import flydsl.expr as fx +from aiter.ops.flydsl.kernels import buffer_ops as bo +from flydsl.expr import rocdl +from flydsl.expr.typing import T + +from vllm.models.deepseek_v41.amd.mono.common.ops import CM_NT + +FP8_MAX = 448.0 +# the f8f6f4 MFMA's operand formats (cbsz / blgp) +FP8, FP4 = 0, 4 +UNIT_SCALE = 127 # the E8M0 code of 1.0 + + +def clamp_fp8(v): + """``v`` clamped to the E4M3 range.""" + return fx.max(fx.min(v, fx.Float32(FP8_MAX)), fx.Float32(-FP8_MAX)) + + +def code_ceil(v): + """Biased exponent of the smallest power of two >= ``v`` (v >= 0, finite), + clamped to [1, 254]: the codes a scale may take.""" + bits = v.bitcast(fx.Int32) + e = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0).select( + fx.Int32(1), fx.Int32(0) + ) + return fx.max(fx.min(e, fx.Int32(254)), fx.Int32(1)) + + +def rcp_pow2(code): + """1 / ``pow2(code)``, exact (code in [1, 253]): x times it is x / scale + bit for bit, a multiply where the division is a dozen instructions.""" + return ((254 - code) << 23).bitcast(fx.Float32) + + +def scale_index(row, col, cols): + """Byte of e8m0 scale (row, col) in aiter's ``shuffle_scale`` layout: + ``view(rows/32, 2, 16, cols/8, 2, 4).permute(0, 3, 5, 2, 4, 1)``.""" + r32, a, b = row // 32, (row // 16) % 2, row % 16 + c8, d, e = col // 8, (col // 4) % 2, col % 4 + return ((((r32 * (cols // 8) + c8) * 4 + e) * 16 + b) * 2 + d) * 2 + a + + +def fp4_tile_load(r_w, base, rg, k, st, at): + """A lane's 16 B of an FP4 weight's (16, 16)-shuffled 16-row group ``rg`` + (K = ``k``) at 128-k step ``st``: dword ``base`` + the group's 2 k + 256 a + step + ``at`` (``lane * 4``: 32 fp4, K block lane / 16 of row lane % 16). + Nontemporal: a weight a step reads once.""" + return fx.Vector( + bo.buffer_load( + r_w, + base + rg * (2 * k) + st * 256 + at, + vec_width=4, + dtype=T.i32, + cache_modifier=CM_NT, + ) + ) + + +def mfma_scaled(a, b, c, sa, sb, a_fmt=FP8, b_fmt=FP8, sa_byte=0, sb_byte=0): + """One 16x16x128 scaled MFMA (gfx950 f8f6f4): A ``a`` (``a_fmt``) with E8M0 + ``sa`` (its byte ``sa_byte``), B ``b`` (``b_fmt``) with E8M0 ``sb`` (its + byte ``sb_byte``).""" + return fx.Vector( + rocdl.mfma_scale_f32_16x16x128_f8f6f4( + T.vec(4, T.f32), [a, b, c, a_fmt, b_fmt, sa_byte, sa, sb_byte, sb] + ) + ) diff --git a/vllm/models/deepseek_v41/amd/mono/common/ops.py b/vllm/models/deepseek_v41/amd/mono/common/ops.py new file mode 100644 index 0000000000..f9f7fb22ec --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono/common/ops.py @@ -0,0 +1,300 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at 97359d6, MIT License, +# Copyright (C) 2026, Advanced Micro Devices, Inc.: +# atom/mono/device/ops.py +# atom/mono/device/ranks.py +# atom/mono/device/stamps.py +"""Device-side primitives of the mono kernels: wave reductions and shuffles, the +approximate hardware math, bf16 / fp8 packing, buffer resources, the +execution-model constants and the kernel symbol naming. + +Ported from the FlyDSL GLM-5 layer kernel (ROCm/FlyDSL branch +``glm5-mla-moe-layer-monokernel`` @ 1e28fd8a, +``kernels/mla_moe_layer/shared_reuse_moe_kernel.py``). Only published flydsl APIs +and aiter's vendored ``buffer_ops`` / ``dpp_utils`` are used. + +Every function but ``kernel_symbol`` is traced inside a ``@flyc.kernel`` body. + +Raw boundaries kept, for want of an equivalent on the pinned FlyDSL: +``buffer_ops`` loads / stores (the kernels take raw Int64 addresses, no tensors +for ``make_buffer_tensor``); raw MFMA intrinsics (hand-laid operands: permlane +transposes, fp8 packing, MXFP4 scales); the LLVM intrinsics here (``expect``, +``s_memrealtime``, the approximate rsq / rcp / exp2) and the struct unpacking of +``permlane*_swap``. + +Also the TP group seen from inside a kernel (every rank's peer buffer, an +all-reduce's rank-ordered sum) and a timeline build's phase stamps.""" + +import flydsl.expr as fx +from aiter.ops.flydsl.kernels import buffer_ops as bo +from aiter.ops.flydsl.kernels.dpp_utils import update_dpp_i32 +from flydsl._mlir.dialects import llvm +from flydsl.compiler.ast_rewriter import ASTRewriter +from flydsl.expr import const_expr, range_constexpr, rocdl +from flydsl.expr import math as fmath +from flydsl.expr.typing import T, as_ir_value + +from vllm.models.deepseek_v41.amd.mono.common.plan import THREADS, WAVES + +# ---------------------------------------------------------------- primitives +# gfx95x cache policy bits (LLVM CPol): SC0 = 1, SC1 = 16. SC1 alone is device +# scope (past the per-XCD caches); SC0 | SC1 is system scope (peers over XGMI). +CM_DEV = 16 + + +CM_SYS = 17 + + +# NT = 2: nontemporal. A weight a layer reads once streams past the caches it +# would only evict the step's reused data from (the MoE at 48 rows: 143 -> 133 us) +CM_NT = 2 + + +def rsrc(addr, nbytes=None): + return bo.create_buffer_resource_from_addr(addr, num_records_bytes=nbytes) + + +def permlane_swap(off, x, y): + """v_permlane{32,16}_swap on registers (x, y): off 32 trades x's upper-half + lanes with y's lower-half lanes; off 16 trades x's odd 16-lane rows with y's + even rows (each within its half). -> (x', y').""" + swap = rocdl.permlane32_swap if off == 32 else rocdl.permlane16_swap + pr = swap( + llvm.StructType.get_literal([T.i32, T.i32]), + as_ir_value(fx.Int32(x)), + as_ir_value(fx.Int32(y)), + False, + False, + ) + return tuple(fx.Int32(llvm.extractvalue(T.i32, pr, [j])) for j in range(2)) + + +def lane_gather(v, src_lane): + """Lane ``src_lane``'s value of v, a source lane a lane (ds_bpermute).""" + return fx.Int32( + rocdl.ds_bpermute(T.i32, fx.Int32(src_lane) * 4, fx.Int32(v).ir_value()) + ) + + +def memrealtime(): + """s_memrealtime: the 100 MHz device clock (no FlyDSL wrapper).""" + return fx.Int64(llvm.call_intrinsic(T.i64, "llvm.amdgcn.s.memrealtime", [], [], [])) + + +def fresh(v): + """``v`` (a per-lane Int32) through an opaque move, here: the lane math built + on it is computed at its use, where LLVM would otherwise hoist it out of a task + loop or CSE it across stages into registers held for the whole kernel (spills, + and a kernel with scratch is refused: ``runtime.widths``).""" + return fx.Int32( + llvm.inline_asm( + T.i32, + [fx.Int32(v).ir_value()], + "v_mov_b32 $0, $1", + "=v,v", + has_side_effects=True, + ) + ) + + +def uniform(v): + return fx.Int32(rocdl.readfirstlane(T.i32, fx.Int32(v).ir_value())) + + +def _hw_f32(name, x): + return fx.Float32( + llvm.call_intrinsic(T.f32, name, [fx.Float32(x).ir_value()], [], []) + ) + + +def hw_rsq(x): + """v_rsq_f32 (~1 ulp, no libm range fixups).""" + return _hw_f32("llvm.amdgcn.rsq.f32", x) + + +def hw_rcp(x): + return _hw_f32("llvm.amdgcn.rcp.f32", x) + + +def hw_exp2(x): + return _hw_f32("llvm.amdgcn.exp2.f32", x) + + +def div_rn(x, d, r): + """``x / d`` rounded to nearest, bit-identical to the IEEE division, given + ``r = 1.0 / d`` from one true division per divisor (Markstein: a quotient + within 1 ulp plus its exact fma residual times RN(1/d) rounds correctly).""" + q = x * r + return fx.Float32(fmath.fma(fx.Float32(fmath.fma(-q, d, x)), r, q)) + + +def xshfl(v, off): + """Value of lane ``lane ^ off``: offsets 32 / 16 via v_permlane*_swap, in-row + offsets via DPP (VALU latency, not LDS). Powers of two only: any other + offset would take the lane ^ 1 path.""" + assert off in (1, 2, 4, 8, 16, 32), off + if off >= 16: + return v.shuffle_xor(off, 64) + is_f = isinstance(v, fx.Float32) + x = v.bitcast(fx.Int32) if is_f else fx.Int32(v) + if off == 8: + y = fx.Int32(update_dpp_i32(x, x, 0x118, 0xF, 0xC, False)) + y = fx.Int32(update_dpp_i32(y, x, 0x108, 0xF, 0x3, False)) + elif off == 4: + y = fx.Int32(update_dpp_i32(x, x, 0x114, 0xF, 0xA, False)) + y = fx.Int32(update_dpp_i32(y, x, 0x104, 0xF, 0x5, False)) + elif off == 2: + y = fx.Int32(update_dpp_i32(x, x, 0x4E, 0xF, 0xF, False)) + else: + y = fx.Int32(update_dpp_i32(x, x, 0xB1, 0xF, 0xF, False)) + return y.bitcast(fx.Float32) if is_f else y + + +def xred(v, off, op): + """op(v, lane ^ off) for a symmetric op; 32 / 16 take both halves of one swap.""" + if off < 16: + return op(v, xshfl(v, off)) + is_f = isinstance(v, fx.Float32) + x = v.bitcast(fx.Int32) if is_f else fx.Int32(v) + a, b = permlane_swap(off, x, x) + if is_f: + return op(a.bitcast(fx.Float32), b.bitcast(fx.Float32)) + return op(type(v)(a), type(v)(b)) + + +def _add(a, b): + return a + b + + +def butterfly(v, offsets, op=_add): + """Butterfly reduction over the given xor offsets, in the given order. + + The order is the numerical contract: the aiter per-head norm reduces 32 + lanes with offsets 16, 8, 4, 2, 1, and matching it keeps rstd bit-exact.""" + for off in offsets: + v = xred(v, off, op) + return v + + +def traced(fn): + """Apply flydsl's AST rewriting to a helper, so its ``while`` / ``if`` over + traced values lower to scf ops. Kernels rewrite only their own body.""" + return ASTRewriter.transform(fn) + + +LOAD_BATCH = 4 # rounds of a strided loop whose loads go out together + + +@traced +def batched_rounds(tid, n, load, use, batch=LOAD_BATCH): + """A CTA's strided loop over elements v = tid + THREADS i < n (n a Python + int): ``load(v)`` of ``batch`` rounds issued before any ``use(v, loaded)``. + A round whose ``use`` stores to global memory keeps the next round's loads + behind it (they may alias), so the rounds' loads would go out in series. A + round past n loads element n - 1 (in range) and uses nothing.""" + rounds = -(-n // THREADS) + for b0 in range_constexpr(0, rounds, batch): + idx = range(b0, min(b0 + batch, rounds)) + got = [load(fx.min(tid + THREADS * i, n - 1)) for i in idx] + for i, g in zip(idx, got): + v = tid + THREADS * i + if v < n: + use(v, g) + + +def mfma_bf16(a, b, c): + return fx.Vector(rocdl.mfma_f32_16x16x32_bf16(T.vec(4, T.f32), [a, b, c])) + + +def row_sum(red, r, t, waves=None): + """Row r, token t of the waves' 16x16 MFMA accumulators (wave w's at ``red`` + + (w * 64 + lane) * 4), summed in wave order from 0 (``waves``: which, in + order; all of them by default).""" + tot = fx.Float32(0.0) + for w in waves if waves is not None else range(WAVES): + tot = tot + fx.ptr_load(red + ((w * 64 + 16 * (r // 4) + t) * 4 + r % 4)) + return tot + + +def ballot(pred): + """The wave's lanes where ``pred`` holds, a bit a lane.""" + return fx.Int64(rocdl.ballot(T.i64, as_ir_value(pred))) + + +def popcount(x): + return fx.Int32(fx.ctpop(x)) + + +def lanes_below(bal, lane): + """How many of the lanes in ballot ``bal`` are below ``lane``.""" + return popcount(bal & ((fx.Int64(1) << fx.Int64(lane)) - 1)) + + +def load_ptr64(table, i): + """Entry ``i`` of an i64 table in memory, wave-uniform (each half through + ``uniform``): a peer's buffer base, a kernel argument block's pointer.""" + pw = fx.Vector(bo.buffer_load(rsrc(table), i * 2, vec_width=2, dtype=T.i32)) + return (fx.Int64(uniform(pw[1])) << 32) | fx.Int64(fx.Uint32(uniform(pw[0]))) + + +def bf16_round(a): + return fx.Float32(fx.Float32(a).to(fx.BFloat16)) + + +def bf16_pair(a, b): + """Two f32 -> one f32-typed word holding (bf16(a), bf16(b)).""" + return ( + fx.Vector.from_elements([a, b], fx.Float32) + .to(fx.BFloat16) + .bitcast(fx.Float32)[0] + ) + + +def bf_lo(w): + """The f32 of a packed bf16 pair word's low half.""" + return (w << 16).bitcast(fx.Float32) + + +def bf_hi(w): + """The f32 of a packed bf16 pair word's high half.""" + return (w & fx.Int32(-65536)).bitcast(fx.Float32) + + +def bf2_f32(w): + """Packed bf16 pair word -> (f32 low, f32 high).""" + return bf_lo(w), bf_hi(w) + + +def fp8_pack4(a, b, c, d): + """Four f32 (already in range) -> one dword of four E4M3 bytes, RNE.""" + lo = fx.Int32(rocdl.cvt_pk_fp8_f32(T.i32, a, b, fx.Int32(0), False)) + return fx.Int32(rocdl.cvt_pk_fp8_f32(T.i32, c, d, lo, True)) + + +# ---------------------------------------------------------------- TP ranks +def peer_bases(peers, tp): + """Every rank's peer buffer base from the ``peers`` table (an i64 a rank), + loaded once: a put's address must not wait on a load of its own.""" + return [load_ptr64(peers, p) for p in range(tp)] + + +def sum_partials(poll, own, pair_of, tp): + """The ``tp`` ranks' bf16 pair partials at ``pair_of(src)`` of the mailbox + region ``own``, summed in fp32 in rank order 0 .. tp - 1: one poll batch -> + (lo sum, hi sum). Every rank sums in the same order, so they agree.""" + ws = poll([(own, pair_of(src), 1) for src in range(tp)]) + acc0, acc1 = bf2_f32(ws[0][0]) + for src in range_constexpr(1, tp): + lo, hi = bf2_f32(ws[src][0]) + acc0 = acc0 + lo + acc1 = acc1 + hi + return acc0, acc1 + + +# ---------------------------------------------------------------- timeline stamps +def stamp(on, tls, tid, k): + # compile-time gate outside, traced condition inside: they cannot be one `and` + if const_expr(on): # noqa: SIM102 + if tid == 0: + fx.ptr_store(memrealtime(), tls + k) diff --git a/vllm/models/deepseek_v41/amd/mono/common/peer_memory.py b/vllm/models/deepseek_v41/amd/mono/common/peer_memory.py new file mode 100644 index 0000000000..58c95a567e --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono/common/peer_memory.py @@ -0,0 +1,104 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at 97359d6, MIT License, +# Copyright (C) 2026, Advanced Micro Devices, Inc.: +# atom/mono/runtime/consensus.py +# atom/mono/runtime/peer_memory.py +"""Symmetric peer memory for the in-kernel TP all-reduces, and the agreement every +rank must reach before anything collective (a rank that turned the mono layer +off alone would leave the others spinning on the GPU).""" + +import torch +import torch.distributed as dist +from aiter.ops.flydsl.quick_allreduce_int4_ipc import UncachedIpcHeap + + +# ---------------------------------------------------------------- TP agreement +class MonoUnsupported(Exception): + """The loaded model or runtime configuration is outside what mono serves.""" + + +def tp_agree(ok: bool, group) -> bool: + """True when every rank of ``group`` passed ``ok`` true. Every rank must call + it at the same point; ``group`` is a CPU (gloo) process group, or None for a + single rank.""" + if group is None: + return ok + flag = torch.tensor([int(ok)], dtype=torch.int32) + dist.all_reduce(flag, op=dist.ReduceOp.MIN, group=group) + return bool(flag.item()) + + +# ---------------------------------------------------------------- peer buffers +class _DeviceBytes: + """``__cuda_array_interface__`` of ``nbytes`` bytes at a raw device address, for a + non-owning torch view.""" + + def __init__(self, ptr: int, nbytes: int): + self.__cuda_array_interface__ = { + "shape": (nbytes,), + "typestr": "|u1", + "data": (ptr, False), + "version": 3, + } + + +class PeerBuffer: + """One uncached buffer per rank, every rank holding every peer's address + (``addresses``, int64, indexed by TP rank). What lives where, and why no rank + overwrites a peer's unread data, is the model's layout's to say. + + The constructor is collective over ``group`` (the handle exchange, then an + agreement that every rank mapped every peer): every rank must reach it, so a + runner reaches it only after ``tp_agree``. A failed mapping on any rank raises + ``MonoUnsupported`` on all of them. + """ + + def __init__( + self, + nbytes: int, + group, + rank: int, + npes: int, + device: torch.device, + ): + self.local = UncachedIpcHeap.alloc_uncached(nbytes) + # a view, not an owner: close() frees the memory under it + self.bytes = torch.as_tensor(_DeviceBytes(self.local, nbytes), device=device) + self.rank = rank + self._opened: list[int] = [] + addresses = [self.local] + if npes > 1: + handles = UncachedIpcHeap.gather_object_list_via_broadcast( + group, UncachedIpcHeap.get_mem_handle_bytes(self.local) + ) + addresses, failure = [], None + try: + for peer, handle in enumerate(handles): + if peer == rank: + addresses.append(self.local) + else: + base = UncachedIpcHeap.open_mem_handle(handle) + self._opened.append(base) + addresses.append(base) + except RuntimeError as err: # every rank must still reach the agreement + failure = err + if not tp_agree(failure is None, group): + self.close() + raise MonoUnsupported( + "peer memory handshake failed on " + f"{'this' if failure else 'another'}" + " TP rank" + ) from failure + self.addresses = torch.tensor(addresses, dtype=torch.int64, device=device) + + def close(self) -> None: + """Close the peer mappings and free the local buffer; idempotent. Only once no + kernel of any rank can still touch it (the runner is being dropped).""" + opened, self._opened = self._opened, [] + for base in opened: + UncachedIpcHeap.close_mem_handle(base) + if self.local: + self.bytes = None + UncachedIpcHeap.free_device_mem(self.local) + self.local = 0 diff --git a/vllm/models/deepseek_v41/amd/mono/common/plan.py b/vllm/models/deepseek_v41/amd/mono/common/plan.py new file mode 100644 index 0000000000..2b05d7a073 --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono/common/plan.py @@ -0,0 +1,135 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at 97359d6, MIT License, +# Copyright (C) 2026, Advanced Micro Devices, Inc.: +# atom/mono/plan/execution.py +# atom/mono/plan/layout.py +# atom/mono/plan/arith.py +# atom/mono/plan/shard.py +# atom/mono/plan/build_key.py +"""Plain-Python planning of the mono kernels: the execution model (one CTA a CU, +all resident), mailbox layouts, arithmetic that runs on ints and traced values +alike, TP shards, and build keys (every build parameter, and this package's +sources, in FlyDSL's JIT cache key).""" + +import enum +import functools +import hashlib +from dataclasses import dataclass, fields +from pathlib import Path + +# ---------------------------------------------------------------- execution model +BLOCKS = 256 +THREADS = 512 +WAVES = THREADS // 64 + + +def first_task(bid, base): + """CTA ``bid``'s first task of a stage placed from CTA ``base`` (any count of + CTAs past 0: placements that add task counts run past the grid), then every + BLOCKS-th.""" + return (bid - base % BLOCKS + BLOCKS) % BLOCKS + + +# ---------------------------------------------------------------- mailbox layouts +PAIR_BYTES = 8 # a (value, tag) pair + + +def pair_layout(items, align: int = PAIR_BYTES, start: int = 0) -> dict: + """``items`` (name, pairs) laid out in order: name -> (byte offset, bytes), + each region from an ``align``-byte boundary at or past ``start``.""" + out, off = {}, start + for name, pairs in items: + off = (off + align - 1) // align * align + out[name] = (off, pairs * PAIR_BYTES) + off += pairs * PAIR_BYTES + return out + + +# ---------------------------------------------------------------- arithmetic +def sel(pred, a, b): + """A if pred else b: a conditional on Python values, ``select`` on traced + ones (rewrapped in the traced operand's type).""" + if isinstance(pred, bool): + return a if pred else b + out = pred.select(a, b) + traced = [x for x in (a, b) if not isinstance(x, int)] + return type(traced[0])(out) if traced else out + + +def imin(a, b): + return sel(a < b, a, b) + + +def i32(x): + """``x``, asserted inside the device's Int32 range when it is a Python int.""" + if isinstance(x, int): + assert -(2**31) <= x < 2**31, x + return x + + +def cdiv(a, b): + """ceil(a / b) for a >= 0, b > 0.""" + return i32((a + b - 1) // b) + + +# ---------------------------------------------------------------- TP shards +class ShardError(ValueError): + """A width the TP size does not divide, or a TP size the model refuses.""" + + +@dataclass(frozen=True) +class Shard: + tp: int + + def __post_init__(self): + if self.tp < 1: + raise ShardError(f"TP {self.tp}") + + def split(self, full: int, what: str) -> int: + """``full`` / tp, refusing a remainder: a rank's share of ``what``.""" + if full % self.tp: + raise ShardError(f"{what} {full} is not divisible by TP {self.tp}") + return full // self.tp + + +def padded(n: int, align: int) -> int: + return cdiv(n, align) * align + + +# ---------------------------------------------------------------- build keys +_SCALARS = (int, float, bool, str, enum.Enum) +_ROOT = Path(__file__).resolve().parents[1] # the mono package + + +@functools.cache +def source_digest(*paths: str) -> str: + """sha256 of every ``.py`` file under ``paths`` (files or directories, + relative to the mono package), each file's path and contents: one model's + kernel sources, so an edit to another model's rebuilds none of these.""" + h = hashlib.sha256() + for rel in paths: + root = _ROOT / rel + files = [root] if root.is_file() else sorted(root.rglob("*.py")) + if not files: + raise FileNotFoundError(f"no kernel sources at {root}") + for f in files: + h.update(str(f.relative_to(_ROOT)).encode()) + h.update(f.read_bytes()) + return h.hexdigest()[:16] + + +def key_tuple(key, sources: str) -> tuple: + """``key``'s (field name, value) pairs, every value a scalar, then + ("sources", ``sources``): the model's ``source_digest``.""" + out = [] + for f in fields(key): + value = getattr(key, f.name) + if not isinstance(value, _SCALARS): + raise TypeError( + f"{type(key).__name__}.{f.name} = {value!r}: a build parameter must be" + " a scalar to reach the JIT cache key" + ) + out.append((f.name, value)) + out.append(("sources", sources)) + return tuple(out) diff --git a/vllm/models/deepseek_v41/amd/mono/common/sync.py b/vllm/models/deepseek_v41/amd/mono/common/sync.py new file mode 100644 index 0000000000..0220ce12c3 --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono/common/sync.py @@ -0,0 +1,177 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at 97359d6, MIT License, +# Copyright (C) 2026, Advanced Micro Devices, Inc.: +# atom/mono/device/sync.py +"""The tagged mailbox: hand-offs between the CTAs of one GPU (device scope) and +between GPUs (system scope), every address a named region (``Addr``).""" + +from dataclasses import dataclass +from enum import Enum + +import flydsl.expr as fx +from aiter.ops.flydsl.kernels import buffer_ops as bo +from flydsl.expr import const_expr, gpu, range_constexpr, rocdl +from flydsl.expr.typing import T + +from vllm.models.deepseek_v41.amd.mono.common.ops import ( + CM_DEV, + CM_SYS, + bf16_pair, + rsrc, + traced, +) + +POLL_MAX = 12 # mailbox specs polled per batch (bounds live registers) + + +@traced +def _spin(load_all, pending): + """Re-issue ``load_all`` until ``pending`` clears; a side-effecting op in the + loop keeps the loads from being hoisted.""" + v = load_all() + while pending(v): + rocdl.s_sleep(1) + v = load_all() + return v + + +class Space(Enum): + SCRATCH = "scratch" # this GPU's memory: device-scope hand-off between CTAs + PEER = "peer" # a peer-symmetric region: system-scope hand-off between GPUs + + +@dataclass(frozen=True) +class Addr: + """A mailbox region's address: the region's name and space, and the traced + address. The space sets the hand-off's scope: device for SCRATCH, system for + PEER.""" + + region: str + space: Space + value: object + + +def sreg(base, offset, region): + """Scratch region ``region`` at byte ``offset`` of ``base``.""" + return Addr(region, Space.SCRATCH, base + fx.Int64(offset)) + + +def preg(base, offset, region): + """Peer-symmetric region ``region`` at byte ``offset`` of a rank's buffer + ``base``.""" + return Addr(region, Space.PEER, base + fx.Int64(offset)) + + +# the store / load cache policy of a hand-off in each space +_SCOPE_CM = {Space.SCRATCH: CM_DEV, Space.PEER: CM_SYS} + + +class Mailbox: + """Tagged-pair hand-off between CTAs (device scope) and GPUs (system scope). + + Every 32-bit value is stored next to this launch's tag; a consumer + polls the payload until all tags match, so a hand-off costs one round trip + (no store drain, no separate flag). The mono layer's launches tag their + hand-offs with their epoch (``layer.py``), so a pair from an earlier step + never carries a live tag. Every address is an ``Addr``: its space picks the + scope. + """ + + def __init__(self, tag): + self.tag = tag + + def put(self, a, i, v): + bits = v.bitcast(fx.Int32) if isinstance(v, fx.Float32) else fx.Int32(v) + bo.buffer_store( + fx.Vector.from_elements([bits, self.tag], fx.Int32), + rsrc(a.value), + i * 2, + cache_modifier=_SCOPE_CM[a.space], + ) + + def put_bf(self, a, i, vs): + """Elements i .. i+len(vs) (2 or 4, i aligned) as packed bf16 pairs, one + store.""" + words = [] + for j in range_constexpr(len(vs) // 2): + words += [bf16_pair(vs[2 * j], vs[2 * j + 1]).bitcast(fx.Int32), self.tag] + bo.buffer_store( + fx.Vector.from_elements(words, fx.Int32), + rsrc(a.value), + i, + cache_modifier=_SCOPE_CM[a.space], + ) + + def put_words(self, a, i, words): + """Raw 32-bit words (1 or 2) at pairs i.. as (word, tag). Integer payloads + (fp8 words) must come this way: through an f32 value a NaN bit pattern is + not kept.""" + vals = [] + for w in words: + vals += [fx.Int32(w), self.tag] + bo.buffer_store( + fx.Vector.from_elements(vals, fx.Int32), + rsrc(a.value), + i * 2, + cache_modifier=_SCOPE_CM[a.space], + ) + + def poll(self, specs, batch=POLL_MAX): + """Batched poll of pairs ``specs`` = [(Addr, pair index, npairs in {1, 2})], + all of one space. + + All pairs are re-loaded together while any tag is stale, so a batch costs + one round trip after its last producer lands. Returns the value words.""" + if const_expr(len(specs) == 0): + return [] + if const_expr(len(specs) > batch): + return self.poll(specs[:batch], batch) + self.poll(specs[batch:], batch) + spaces = {a.space for a, _, _ in specs} + assert len(spaces) == 1, f"one poll batch spans {spaces}" + cm = _SCOPE_CM[spaces.pop()] + bases = [a.value for a, _, _ in specs] + tag = self.tag + + def load_all(): + words = [] + for b, (_, i, n) in zip(bases, specs): + w = fx.Vector( + bo.buffer_load( + rsrc(b), + fx.Int32(i) * 2, + vec_width=2 * n, + dtype=T.i32, + cache_modifier=cm, + ) + ) + words += [w[e] for e in range(2 * n)] + return fx.Vector.from_elements(words, fx.Int32) + + nw = sum(2 * n for _, _, n in specs) + + def pending(v): + bad = v[1] != tag + for e in range_constexpr(3, nw, 2): + bad = bad | (v[e] != tag) + return bad + + v = _spin(load_all, pending) + outs, e = [], 0 + for _, _, n in specs: + outs.append([v[e + 2 * q] for q in range(n)]) + e += 2 * n + return outs + + +@traced +def publish(put, a, i, v, who): + """This CTA's plain stores made readable by every CTA, then (thread ``who``) + the mailbox pair (``a``, i) := v. The stores are drained first: a flag may + not overtake the data it announces, and a workgroup-scope release fence does + not wait for them (LLVM's AMDGPU memory model omits vmcnt outside tgsplit + mode: the workgroup's waves share one CU, the readers here do not).""" + rocdl.s_waitcnt(vmcnt=0) + gpu.barrier() + if who: + put(a, i, v) diff --git a/vllm/models/deepseek_v41/amd/mono/layer.py b/vllm/models/deepseek_v41/amd/mono/layer.py new file mode 100644 index 0000000000..c1be9e3428 --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono/layer.py @@ -0,0 +1,940 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The mono decode layer's two launches (README): K1 = the attention seam +(``stages.seam``: slice, gate, the norm into vLLM's MXFP8) + the attention front +(``attention.front``); K2 = the attention back (``attention.back``, wo_b pushed +to every rank) + the FFN seam (the TP sum folded in) + the MoE (``stages.moe``, +its all-reduce in-kernel too). A layer whose attention vLLM runs takes the FFN +launch instead: K2 with the attention back replaced by a push of vLLM's +unreduced wo_b output (``seam.stage_push``). + +Tags: K1's hand-offs carry 2 e + 1, K2's 2 e + 2 (e: the launch pair's epoch, +moved on by K2's CTA 0 at its end), so K2's seam may reuse K1's seam regions. +Peer regions alternate by e's parity; every rank runs the same launches, so +the epochs agree across ranks. +""" + +from dataclasses import dataclass + +import flydsl.compiler as flyc +import flydsl.expr as fx +from aiter.ops.flydsl.kernels import buffer_ops as bo +from flydsl.expr import const_expr, gpu, range_constexpr, rocdl +from flydsl.expr.typing import Int32, Int64, T + +from .attention import back as mb_back +from .attention import front as mb_front +from .attention.device import BLOCKS, CM_DEV, THREADS, gstore, kernel_symbol, rsrc +from .attention.plan import HIDDEN, Dims, back_scratch, front_scratch +from .common.mx import clamp_fp8 +from .common.ops import butterfly, fp8_pack4, peer_bases, traced +from .common.plan import WAVES, first_task, key_tuple +from .common.sync import Mailbox, preg, publish, sreg +from .sources import SOURCES +from .stages import moe, seam +from .stages.dims import Dims as StageDims +from .stages.moe_shape import SORT_NETS, MoeBuild, route_shape + +_STREAM = fx.Stream(None) +MAX_TOKENS = 48 +X_WORDS = HIDDEN // 4 +X_GROUPS = HIDDEN // 32 +SLICES = seam.SLICES # 160 +# K1's MXFP8 x and the MoE's: plain words in both, the same size, so the two +# launches share the MoE's (placed later); any other name is one region, so a +# word is never a tag to one launch and data to another +SHARED_REGIONS = ("x8", "x8s") +COUNTER_WORDS = 2 * 256 # the MoE's ug queue and down counts, a word a slot +COUNTERS = ("ugq", "dq") +# the epoch buffer, in words: the epoch, a mark a CTA (back.EPOCH_MARKS), then on +# lines of their own the MoE's counters +EPOCH_COUNTERS = -(-(mb_back.EPOCH_MARKS + BLOCKS) // 64) * 64 +EPOCH_WORDS = EPOCH_COUNTERS + COUNTER_WORDS + + +@dataclass(frozen=True) +class MonoBuild: + tokens: int + tp: int + ratio: int # the layer's compress ratio (its attention's key sources) + timeline: bool = False + + +def _moe_key(key: MonoBuild) -> MoeBuild: + return MoeBuild(tokens=key.tokens, tp=key.tp) + + +def _front_key(key: MonoBuild): + return mb_front.FrontBuild(key.tokens, key.tp, key.ratio, key.timeline) + + +def _back_key(key: MonoBuild): + return mb_back.BackBuild(key.tokens, key.tp, key.ratio, key.timeline) + + +def scratch_layout(s: int, tp: int) -> dict: + """Every region of step width s's launches -> (byte offset, bytes), disjoint: + K1's front and seam, K2's back, the MoE's regions, the normed rows and their + flags. Each width has a scratch of its own (``DSV41MonoLayer.scratch``): + under another width's layout a mailbox pair's tag word could be plain data, + which a poll may take for a current tag. The MoE's counters, whose slots + outlive a step width, are in the epoch buffer.""" + d = Dims(tp) + out: dict = {} + off = 0 + + def place(regions): + nonlocal off + base = off + for name, (o, n) in regions.items(): + assert name not in out or name in SHARED_REGIONS, name + out[name] = (base + o, n) + off = max(off, base + max(o + n for o, n in regions.values())) + off = -(-off // 256) * 256 + + place(front_scratch(s, d)) + place(seam.scratch_layout(s)) + place(back_scratch(s, d, start=0)) + regions = moe.scratch_layout(_moe_key(MonoBuild(s, tp, 1))).items() + mr = {n: v for n, v in regions if n not in COUNTERS} + lo = min(o for o, _ in mr.values()) + place({n: (o - lo, n_) for n, (o, n_) in mr.items()}) + place({"normed": (0, s * HIDDEN * 2), "xrdy_moe": (s * HIDDEN * 2 + 256, s * 8)}) + return out + + +def scratch_bytes(s: int = MAX_TOKENS, tp: int = 2) -> int: + return max(o + n for o, n in scratch_layout(s, tp).values()) + + +def peer_half_bytes(tp: int) -> int: + """One parity's peer regions: the attention's partials, then the MoE's.""" + return seam.attn_peer_bytes(MAX_TOKENS, tp) + moe.peer_bytes(MAX_TOKENS, tp) + + +def _epoch(epoch): + return fx.Int32(bo.buffer_load(rsrc(epoch), 0, vec_width=1, dtype=T.i32)) + + +# ---------------------------------------------------------------- K1 + + +@traced +def store_x8(c, t, col, ys, live): + """The attention seam's norm output (``seam.stage_norm``'s 8 columns at + ``col``) -> vLLM's MXFP8 of each 32 (the attention's input quant) -> the + front's X8 / X8S, plain at device scope.""" + amax = abs(ys[0]) + for y in ys[1:]: + amax = fx.max(amax, abs(y)) + code = moe.vllm_mx_code(butterfly(amax, (1, 2), fx.max)) + mul = ((254 - code) << 23).bitcast(fx.Float32) + w0 = fp8_pack4(*[clamp_fp8(ys[i] * mul) for i in range(4)]) + w1 = fp8_pack4(*[clamp_fp8(ys[4 + i] * mul) for i in range(4)]) + if live: + bo.buffer_store( + fx.Vector.from_elements([w0, w1], fx.Int32), + rsrc(c["x8"]), + t * X_WORDS + col // 4, + cache_modifier=CM_DEV, + ) + if c["tid"] % 4 == 0: + bo.buffer_store( + code, rsrc(c["x8s"]), t * X_GROUPS + col // 32, cache_modifier=CM_DEV + ) + + +@traced +def publish_x8(c, t): + """Token t's two XRDY flags (one a wqkv K half) once its X8 / X8S landed.""" + rocdl.s_waitcnt(vmcnt=0) + gpu.barrier() + if c["tid"] < 2: + c["mb"].put(c["xrdy"], 2 * t + c["tid"], fx.Int32(1)) + gpu.barrier() + + +def build_mono_k1(key: MonoBuild): + s, tp = key.tokens, key.tp + assert 1 <= s <= MAX_TOKENS + layout = scratch_layout(s, tp) + fkey = _front_key(key) + FrontLds = mb_front.front_lds(s, Dims(tp)) + NORM0 = SLICES # the norms on the CTAs past the slices: the wqkv waits on them + GATE0 = SLICES + s + assert GATE0 + s <= BLOCKS + + @fx.struct + class SeamLds: + red: fx.Array[fx.Float32, WAVES * 2 * 64 * 4, 16] + rl: fx.Array[fx.Float32, s * seam.KT, 16] + fl: fx.Array[fx.Float32, seam.MIX * seam.KT, 16] + + @fx.union + class K1Lds: + seam: SeamLds + front: FrontLds # type: ignore[valid-type] + + name = kernel_symbol("dsv41_mono_k1", s=s, tp=tp, r=key.ratio, tl=key.timeline) + keyed = key_tuple(key, SOURCES) + + @flyc.kernel(name=name, known_block_size=[THREADS, 1, 1]) + def k1( + res_in: Int64, + pend: Int64, + post_in: Int64, + comb_in: Int64, + pre_in: Int64, + hc_fn: Int64, + hc_scale: Int64, + hc_base: Int64, + attn_w: Int64, + res_out: Int64, + post_out: Int64, + comb_out: Int64, + pre_out: Int64, + wqkv: Int64, + wqkv_s: Int64, + qn_w: Int64, + kvn_w: Int64, + wqb: Int64, + wqb_s: Int64, + cos_sin: Int64, + pos: Int64, + slot: Int64, + swa: Int64, + swa_stride: Int64, + swa_block: Int32, + q: Int64, + swa_idx: Int64, + swa_lens: Int64, + comp_block: Int32, + comp_bt: Int64, + bt_stride: Int32, + topk: Int64, + t2r: Int64, + kt: Int64, + klen: Int64, + scratch: Int64, + epoch: Int64, + tl: Int64, + ): + _ = keyed + bid = fx.block_idx.x + lds = fx.SharedAllocator().allocate(K1Lds) + sl = lds.seam.peek() + ep = _epoch(epoch) + tag = (ep << 1) + 1 + fargs = { + "x": 0, + "x_stride": 0, + "wqkv": wqkv, + "wqkv_s": wqkv_s, + "qn_w": qn_w, + "kvn_w": kvn_w, + "wqb": wqb, + "wqb_s": wqb_s, + "cos_sin": cos_sin, + "pos": pos, + "slot": slot, + "swa": swa, + "swa_stride": swa_stride, + "swa_block": swa_block, + "q": q, + "swa_idx": swa_idx, + "swa_lens": swa_lens, + "comp_block": comp_block, + "comp_bt": comp_bt, + "bt_stride": bt_stride, + "topk": topk, + "t2r": t2r, + "kt": kt, + "klen": klen, + } + flayout = {n: layout[n] for n in front_scratch(s, Dims(tp))} + c = mb_front.front_context(fkey, lds.front.peek(), fargs, scratch, tag, flayout) + amb = Mailbox(tag) + tid = fx.thread_idx.x + ca = { + "S": s, + "tid": tid, + "bid": bid, + "lane": tid % 64, + "wave": tid // 64, + "rl": sl.rl.ptr, + "fl": sl.fl.ptr, + "red": sl.red.ptr, + "put": amb.put, + "put_bf": amb.put_bf, + "put_words": amb.put_words, + "poll": amb.poll, + "args": { + "res_in": res_in, + "pend": pend, + "post_in": post_in, + "comb_in": comb_in, + "pre_in": pre_in, + "hc_fn": hc_fn, + "hc_scale": hc_scale, + "hc_base": hc_base, + "norm_w": attn_w, + "res_out": res_out, + "post_out": post_out, + "comb_out": comb_out, + "pre_out": pre_out, + }, + "d": StageDims(tp), + } + for region in ("lin", "pmix"): + ca[region] = sreg(scratch, layout[region][0], region) + if const_expr(key.timeline): + c["tl"], c["tl_points"] = tl, mb_front.FRONT_POINTS + + def seam_slice(): + for task in range(first_task(bid, 0), SLICES, BLOCKS): + seam.stage_slice(ca, task) + gpu.barrier() + + def seam_rest(): + for t in range(first_task(bid, NORM0), s, BLOCKS): + seam.stage_norm(ca, t, lambda *v: store_x8(c, *v)) + publish_x8(c, t) + for t in range(first_task(bid, GATE0), s, BLOCKS): + seam.stage_gate(ca, t) + gpu.barrier() + + mb_front.run_front( + c, fkey, bid, x8_given=True, before_wqkv=seam_slice, after_wqkv=seam_rest + ) + + @flyc.jit + def launch( + res_in: Int64, + pend: Int64, + post_in: Int64, + comb_in: Int64, + pre_in: Int64, + hc_fn: Int64, + hc_scale: Int64, + hc_base: Int64, + attn_w: Int64, + res_out: Int64, + post_out: Int64, + comb_out: Int64, + pre_out: Int64, + wqkv: Int64, + wqkv_s: Int64, + qn_w: Int64, + kvn_w: Int64, + wqb: Int64, + wqb_s: Int64, + cos_sin: Int64, + pos: Int64, + slot: Int64, + swa: Int64, + swa_stride: Int64, + swa_block: Int32, + q: Int64, + swa_idx: Int64, + swa_lens: Int64, + comp_block: Int32, + comp_bt: Int64, + bt_stride: Int32, + topk: Int64, + t2r: Int64, + kt: Int64, + klen: Int64, + scratch: Int64, + epoch: Int64, + tl: Int64, + stream: fx.Stream = _STREAM, + ): + _ = keyed + k1( + res_in, + pend, + post_in, + comb_in, + pre_in, + hc_fn, + hc_scale, + hc_base, + attn_w, + res_out, + post_out, + comb_out, + pre_out, + wqkv, + wqkv_s, + qn_w, + kvn_w, + wqb, + wqb_s, + cos_sin, + pos, + slot, + swa, + swa_stride, + swa_block, + q, + swa_idx, + swa_lens, + comp_block, + comp_bt, + bt_stride, + topk, + t2r, + kt, + klen, + scratch, + epoch, + tl, + ).launch(grid=(BLOCKS,), block=(THREADS,), stream=stream) + + return launch + + +# ---------------------------------------------------------------- K2 + + +@traced +def store_normed(normed, t, col, ys, live): + """The FFN seam's norm output -> the MoE input row (bf16, device scope: + this launch reads it).""" + if live: + for j in range_constexpr(8): + bo.buffer_store( + ys[j].to(fx.BFloat16), + rsrc(normed), + t * HIDDEN + col + j, + cache_modifier=CM_DEV, + ) + + +def _seam_lds(s): + @fx.struct + class SeamLds: + red: fx.Array[fx.Float32, WAVES * 2 * 64 * 4, 16] + rl: fx.Array[fx.Float32, s * seam.KT, 16] + fl: fx.Array[fx.Float32, seam.MIX * seam.KT, 16] + pl: fx.Array[fx.Float32, s * seam.COLS, 16] + + return SeamLds + + +@traced +def run_ffn(key: MonoBuild, lds, a: dict, amb, bases, own, rank, ep, part=None): + """K2's FFN half on this CTA: the FFN seam (the attention's TP partials -- + ``part`` pushed to every rank's ATTN region first -- summed and folded into + the residual), the MoE and its all-reduce into ``a["out"]``, then the + launch pair's epoch moved on.""" + s, tp = key.tokens, key.tp + layout = scratch_layout(s, tp) + mkey = _moe_key(key) + rs = route_shape(mkey) + GATE0 = SLICES + NORM0 = SLICES + s + bid = fx.block_idx.x + tid = fx.thread_idx.x + sl = lds.seam.peek() + scratch = a["scratch"] + seam_args = ( + "res_in", + "post_in", + "comb_in", + "pre_in", + "hc_fn", + "hc_scale", + "hc_base", + "res_out", + "post_out", + "comb_out", + "pre_out", + ) + ca = { + "S": s, + "tid": tid, + "bid": bid, + "lane": tid % 64, + "wave": tid // 64, + "rl": sl.rl.ptr, + "fl": sl.fl.ptr, + "red": sl.red.ptr, + "pend_lds": sl.pl.ptr, + "put": amb.put, + "put_bf": amb.put_bf, + "put_words": amb.put_words, + "poll": amb.poll, + "peer_addr": lambda p: bases[p], + "rank": rank, + "sym": own, + "args": {n: a[n] for n in seam_args} | {"norm_w": a["ffn_w"]}, + "d": StageDims(tp), + } + normed = scratch + fx.Int64(layout["normed"][0]) + for region in ("lin", "pmix"): + ca[region] = sreg(scratch, layout[region][0], region) + mlayout = {n: layout[n] for n in moe.scratch_layout(mkey) if n not in COUNTERS} + mlayout["xrdy"] = layout["xrdy_moe"] + margs = moe.moe_args( + normed, + a["gate_w"], + a["bias"], + a["w13"], + a["w13_s"], + a["w2"], + a["w2_s"], + a["sgu"], + a["sgu_s"], + a["sw2"], + a["sw2_s"], + a["out"], + ) + cb = moe.moe_context( + s, rs, lds.moe.peek(), amb, bases, mlayout, scratch, margs, rank, own + ) + cb["x_cm"], cb["x_ready"] = CM_DEV, True + cb["tag"] = ep & 255 + for i, region in enumerate(COUNTERS): + cb[region] = sreg(a["epoch"], 4 * (EPOCH_COUNTERS + i * moe.TAGS), region) + for task in range(first_task(bid, 0), SLICES, BLOCKS): + if const_expr(part is not None): + seam.stage_push(ca, task, part) + seam.stage_reduce(ca, task) + seam.stage_slice(ca, task) + for t in range(first_task(bid, GATE0), s, BLOCKS): + seam.stage_gate(ca, t) + for t in range(first_task(bid, NORM0), s, BLOCKS): + seam.stage_norm(ca, t, lambda *v: store_normed(normed, *v)) + publish(cb["put"], cb["xrdy"], t, 1, tid == 0) + gpu.barrier() + + # ---- the MoE, its all-reduce into ``out`` + moe.run_moe(cb, mkey, bid, 0, 0) + + def reset(): + # the MoE counter slot 128 launch pairs ahead: its last use long + # done, its next use far off + slot = (ep + 128) & 255 + for region in COUNTERS: + gstore(cb[region].value + fx.Int64(slot * 4), fx.Int32(0), words=1) + + mb_back.epoch_end({"bid": bid, "tid": tid}, a["epoch"], ep, reset) + + +def build_mono_k2(key: MonoBuild): + s, tp = key.tokens, key.tp + assert 1 <= s <= MAX_TOKENS + layout = scratch_layout(s, tp) + bkey = _back_key(key) + mkey = _moe_key(key) + rs = route_shape(mkey) + assert rs.experts // 64 in SORT_NETS + BackM = mb_back.back_lds_members(s, Dims(tp)) + SplitLds, GemvLds = BackM["split"], BackM["gemv"] + MoeLds = moe.moe_smem(s, rs, False) + half = peer_half_bytes(tp) + assert SLICES + 2 * s <= BLOCKS + SeamLds = _seam_lds(s) + + @fx.union + class K2Lds: + split: SplitLds # type: ignore[valid-type] + gemv: GemvLds # type: ignore[valid-type] + seam: SeamLds # type: ignore[valid-type] + moe: MoeLds # type: ignore[valid-type] + + name = kernel_symbol("dsv41_mono_k2", s=s, tp=tp, r=key.ratio, tl=key.timeline) + keyed = key_tuple(key, SOURCES) + + @flyc.kernel(name=name, known_block_size=[THREADS, 1, 1]) + def k2( + q: Int64, + swa: Int64, + swa_stride: Int64, + swa_block: Int32, + comp: Int64, + comp_stride: Int64, + comp_block: Int32, + kt: Int64, + klen: Int64, + pos: Int64, + sink: Int64, + qk_scale: Int32, + cos_sin: Int64, + woa: Int64, + woa_s: Int64, + wob: Int64, + wob_s: Int64, + zrec: Int64, + res_in: Int64, + post_in: Int64, + comb_in: Int64, + pre_in: Int64, + hc_fn: Int64, + hc_scale: Int64, + hc_base: Int64, + ffn_w: Int64, + res_out: Int64, + post_out: Int64, + comb_out: Int64, + pre_out: Int64, + gate_w: Int64, + bias: Int64, + w13: Int64, + w13_s: Int64, + w2: Int64, + w2_s: Int64, + sgu: Int64, + sgu_s: Int64, + sw2: Int64, + sw2_s: Int64, + out: Int64, + scratch: Int64, + sym: Int64, + peers: Int64, + rank: Int32, + epoch: Int64, + tl: Int64, + ): + _ = keyed + bid = fx.block_idx.x + lds = fx.SharedAllocator().allocate(K2Lds) + ep = _epoch(epoch) + tag = (ep << 1) + 2 + par = fx.Int64(ep & 1) * fx.Int64(half) + bases = [b + par for b in peer_bases(peers, tp)] + own = sym + par + bargs = { + "q": q, + "swa": swa, + "swa_stride": swa_stride, + "swa_block": swa_block, + "comp": comp, + "comp_stride": comp_stride, + "comp_block": comp_block, + "kt": kt, + "klen": klen, + "pos": pos, + "sink": sink, + "qk_scale": qk_scale, + "cos_sin": cos_sin, + "woa": woa, + "woa_s": woa_s, + "wob": wob, + "wob_s": wob_s, + "out": 0, + "out_stride": 0, + "zrec": zrec, + } + blayout = {n: layout[n] for n in back_scratch(s, Dims(tp), start=0)} + views = {"split": lds.split.peek(), "gemv": lds.gemv.peek()} + c = mb_back.back_context(bkey, views, bargs, scratch, tag, blayout) + amb = Mailbox(tag) + + def push(t, col, v0, v1): + # this rank's bf16 partial pair to every rank's ATTN region + for p in range_constexpr(tp): + dst = preg(bases[p], 0, "attn") + amb.put_bf(dst, 2 * seam.attn_region_pair(rank, t, s, col), [v0, v1]) + + c["wob_out"] = push + if const_expr(key.timeline): + c["tl"], c["tl_points"] = tl, mb_back.BACK_POINTS + mb_back.epoch_begin(c, epoch, ep) + mb_back.run_back(c, bkey, bid) + gpu.barrier() + + ffn = dict( + res_in=res_in, + post_in=post_in, + comb_in=comb_in, + pre_in=pre_in, + hc_fn=hc_fn, + hc_scale=hc_scale, + hc_base=hc_base, + ffn_w=ffn_w, + res_out=res_out, + post_out=post_out, + comb_out=comb_out, + pre_out=pre_out, + gate_w=gate_w, + bias=bias, + w13=w13, + w13_s=w13_s, + w2=w2, + w2_s=w2_s, + sgu=sgu, + sgu_s=sgu_s, + sw2=sw2, + sw2_s=sw2_s, + out=out, + scratch=scratch, + epoch=epoch, + ) + run_ffn(key, lds, ffn, amb, bases, own, rank, ep) + + @flyc.jit + def launch( + q: Int64, + swa: Int64, + swa_stride: Int64, + swa_block: Int32, + comp: Int64, + comp_stride: Int64, + comp_block: Int32, + kt: Int64, + klen: Int64, + pos: Int64, + sink: Int64, + qk_scale: Int32, + cos_sin: Int64, + woa: Int64, + woa_s: Int64, + wob: Int64, + wob_s: Int64, + zrec: Int64, + res_in: Int64, + post_in: Int64, + comb_in: Int64, + pre_in: Int64, + hc_fn: Int64, + hc_scale: Int64, + hc_base: Int64, + ffn_w: Int64, + res_out: Int64, + post_out: Int64, + comb_out: Int64, + pre_out: Int64, + gate_w: Int64, + bias: Int64, + w13: Int64, + w13_s: Int64, + w2: Int64, + w2_s: Int64, + sgu: Int64, + sgu_s: Int64, + sw2: Int64, + sw2_s: Int64, + out: Int64, + scratch: Int64, + sym: Int64, + peers: Int64, + rank: Int32, + epoch: Int64, + tl: Int64, + stream: fx.Stream = _STREAM, + ): + _ = keyed + k2( + q, + swa, + swa_stride, + swa_block, + comp, + comp_stride, + comp_block, + kt, + klen, + pos, + sink, + qk_scale, + cos_sin, + woa, + woa_s, + wob, + wob_s, + zrec, + res_in, + post_in, + comb_in, + pre_in, + hc_fn, + hc_scale, + hc_base, + ffn_w, + res_out, + post_out, + comb_out, + pre_out, + gate_w, + bias, + w13, + w13_s, + w2, + w2_s, + sgu, + sgu_s, + sw2, + sw2_s, + out, + scratch, + sym, + peers, + rank, + epoch, + tl, + ).launch(grid=(BLOCKS,), block=(THREADS,), stream=stream) + + return launch + + +# ---------------------------------------------------------------- the FFN launch + + +def build_mono_ffn(key: MonoBuild): + """The FFN launch of a layer whose attention vLLM ran (``part``: this rank's + unreduced wo_b output, bf16 [S, HIDDEN]): K2's FFN half, the partial pushed + to every rank first. It takes a launch pair's epoch, as K2 does.""" + s, tp = key.tokens, key.tp + assert 1 <= s <= MAX_TOKENS + rs = route_shape(_moe_key(key)) + assert rs.experts // 64 in SORT_NETS + MoeLds = moe.moe_smem(s, rs, False) + half = peer_half_bytes(tp) + assert SLICES + 2 * s <= BLOCKS + SeamLds = _seam_lds(s) + + @fx.union + class FfnLds: + seam: SeamLds # type: ignore[valid-type] + moe: MoeLds # type: ignore[valid-type] + + name = kernel_symbol("dsv41_mono_ffn", s=s, tp=tp) + keyed = key_tuple(key, SOURCES) + + @flyc.kernel(name=name, known_block_size=[THREADS, 1, 1]) + def kffn( + part: Int64, + res_in: Int64, + post_in: Int64, + comb_in: Int64, + pre_in: Int64, + hc_fn: Int64, + hc_scale: Int64, + hc_base: Int64, + ffn_w: Int64, + res_out: Int64, + post_out: Int64, + comb_out: Int64, + pre_out: Int64, + gate_w: Int64, + bias: Int64, + w13: Int64, + w13_s: Int64, + w2: Int64, + w2_s: Int64, + sgu: Int64, + sgu_s: Int64, + sw2: Int64, + sw2_s: Int64, + out: Int64, + scratch: Int64, + sym: Int64, + peers: Int64, + rank: Int32, + epoch: Int64, + ): + _ = keyed + lds = fx.SharedAllocator().allocate(FfnLds) + ep = _epoch(epoch) + par = fx.Int64(ep & 1) * fx.Int64(half) + bases = [b + par for b in peer_bases(peers, tp)] + mb_back.epoch_begin({"bid": fx.block_idx.x, "tid": fx.thread_idx.x}, epoch, ep) + a = dict( + res_in=res_in, + post_in=post_in, + comb_in=comb_in, + pre_in=pre_in, + hc_fn=hc_fn, + hc_scale=hc_scale, + hc_base=hc_base, + ffn_w=ffn_w, + res_out=res_out, + post_out=post_out, + comb_out=comb_out, + pre_out=pre_out, + gate_w=gate_w, + bias=bias, + w13=w13, + w13_s=w13_s, + w2=w2, + w2_s=w2_s, + sgu=sgu, + sgu_s=sgu_s, + sw2=sw2, + sw2_s=sw2_s, + out=out, + scratch=scratch, + epoch=epoch, + ) + amb = Mailbox((ep << 1) + 2) + run_ffn(key, lds, a, amb, bases, sym + par, rank, ep, part=part) + + @flyc.jit + def launch( + part: Int64, + res_in: Int64, + post_in: Int64, + comb_in: Int64, + pre_in: Int64, + hc_fn: Int64, + hc_scale: Int64, + hc_base: Int64, + ffn_w: Int64, + res_out: Int64, + post_out: Int64, + comb_out: Int64, + pre_out: Int64, + gate_w: Int64, + bias: Int64, + w13: Int64, + w13_s: Int64, + w2: Int64, + w2_s: Int64, + sgu: Int64, + sgu_s: Int64, + sw2: Int64, + sw2_s: Int64, + out: Int64, + scratch: Int64, + sym: Int64, + peers: Int64, + rank: Int32, + epoch: Int64, + stream: fx.Stream = _STREAM, + ): + _ = keyed + kffn( + part, + res_in, + post_in, + comb_in, + pre_in, + hc_fn, + hc_scale, + hc_base, + ffn_w, + res_out, + post_out, + comb_out, + pre_out, + gate_w, + bias, + w13, + w13_s, + w2, + w2_s, + sgu, + sgu_s, + sw2, + sw2_s, + out, + scratch, + sym, + peers, + rank, + epoch, + ).launch(grid=(BLOCKS,), block=(THREADS,), stream=stream) + + return launch diff --git a/vllm/models/deepseek_v41/amd/mono/runner.py b/vllm/models/deepseek_v41/amd/mono/runner.py new file mode 100644 index 0000000000..eb37990275 --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono/runner.py @@ -0,0 +1,437 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""DeepSeek-V4.1-Flash mono decode layer: host side. + +One ``DSV41MonoLayer`` a TP rank runs every mono layer of a decode step (M <= 48 +rows): two persistent launches a layer (``layer``) from the layer's +inputs at the attention seam to its outputs at the next one -- the MoE's +reduced output, the residual after the FFN seam and that seam's mixes -- with +both TP all-reduces inside the kernels (symmetric peer memory). A layer whose +attention vLLM runs (``ffn``) takes one launch from its unreduced attention +output on. Every argument is a device pointer: the launches can be captured in a +HIP graph. The kernels move their mailbox epoch on themselves; every rank runs +the same launches, so the ranks' epochs agree. +""" + +from __future__ import annotations + +import os +from dataclasses import dataclass +from pathlib import Path + +import torch + +from .attention.plan import HEAD_DIM, HIDDEN, KEYS, Dims +from .common.plan import BLOCKS +from .layer import ( + EPOCH_WORDS, + MAX_TOKENS, + MonoBuild, + build_mono_ffn, + build_mono_k1, + build_mono_k2, + peer_half_bytes, + scratch_bytes, +) +from .stages.dims import Dims as MoeDims +from .stages.moe_shape import EXPERTS + +__all__ = ["BLOCKS", "MAX_TOKENS", "AttnWeights", "DSV41MonoLayer", "MonoLayerWeights"] + +_LOG2E = 1.4426950408889634 +HC = 4 + + +def _ensure_writable_flydsl_cache() -> None: + """Aiter points FlyDSL at its bundled, read-only cache; new kernels need a + writable one.""" + cur = os.environ.get("FLYDSL_RUNTIME_CACHE_DIR") + if cur and os.access(cur, os.W_OK): + return + path = Path.home() / ".flydsl" / "cache" + path.mkdir(parents=True, exist_ok=True) + os.environ["FLYDSL_RUNTIME_CACHE_DIR"] = str(path) + + +def _check_tensors(owner, tag: str, want: dict) -> None: + """Each named tensor of ``owner`` has the (shape, dtype) the kernels read, + contiguous: their buffer loads are unbounded, so a mismatch would read + garbage rather than fault.""" + for name, (shape, dtype) in want.items(): + t = getattr(owner, name) + assert tuple(t.shape) == shape and t.dtype == dtype and t.is_contiguous(), ( + f"{tag} {name}: {tuple(t.shape)} {t.dtype}, want {shape} {dtype}" + ) + + +@dataclass +class AttnWeights: + """One layer's attention tensors on this rank, in vLLM's loaded layout.""" + + layer_id: int + wqkv: torch.Tensor # [1792, 5120] e4m3: [wq_a; wkv] + wqkv_scale: torch.Tensor # [56, 160] uint8 E8M0 (32 x 32 blocks) + q_norm: torch.Tensor # [1280] bf16 + kv_norm: torch.Tensor # [512] bf16 + wq_b: torch.Tensor # [H * 512, 1280] e4m3 + wq_b_scale: torch.Tensor # [H * 16, 40] uint8 + wo_a: torch.Tensor # [G * 1024, 4096] e4m3 + wo_a_scale: torch.Tensor # [G * 32, 128] uint8 + wo_b: torch.Tensor # [5120, G * 1024] e4m3 + wo_b_scale: torch.Tensor # [160, G * 32] uint8 + attn_sink: torch.Tensor # [>= H] f32 + cos_sin: torch.Tensor # [positions, 64] f32 (the layer's rope table) + ratio: int # 0 / 1 / 2: the layer's compress ratio + + def check(self, d: Dims) -> None: + h, g = d.heads, d.groups + want = { + "wqkv": ((1792, HIDDEN), torch.float8_e4m3fn), + "wqkv_scale": ((56, HIDDEN // 32), torch.uint8), + "wq_b": ((h * HEAD_DIM, 1280), torch.float8_e4m3fn), + "wq_b_scale": ((h * HEAD_DIM // 32, 40), torch.uint8), + "wo_a": ((g * 1024, d.group_k), torch.float8_e4m3fn), + "wo_a_scale": ((g * 32, d.group_k // 32), torch.uint8), + "wo_b": ((HIDDEN, g * 1024), torch.float8_e4m3fn), + "wo_b_scale": ((HIDDEN // 32, g * 32), torch.uint8), + } + _check_tensors(self, f"layer {self.layer_id}", want) + assert self.cos_sin.dtype == torch.float32 and self.cos_sin.shape[-1] == 64 + + +@dataclass +class MonoLayerWeights: + """One layer's tensors on this rank, in vLLM's loaded layout (``attn``: + None for a layer whose attention vLLM runs).""" + + attn: AttnWeights | None + hc_attn_fn: torch.Tensor # [24, 4 * 5120] f32 + hc_attn_scale: torch.Tensor # [3] f32 + hc_attn_base: torch.Tensor # [24] f32 + attn_norm: torch.Tensor # [5120] bf16 + hc_ffn_fn: torch.Tensor + hc_ffn_scale: torch.Tensor + hc_ffn_base: torch.Tensor + ffn_norm: torch.Tensor + gate_w: torch.Tensor # [384, 5120] bf16 + bias: torch.Tensor # [384] f32 (e_score_correction_bias) + w13: torch.Tensor # [384, 2 inter, 2560] fp4x2, aiter (16, 16) shuffle + w13_s: torch.Tensor # its scales, aiter shuffle_scale order + w2: torch.Tensor # [384, 5120, inter / 2] fp4x2 + w2_s: torch.Tensor + sgu: torch.Tensor # shared gate_up [2 inter, 5120] e4m3, row-major + sgu_s: torch.Tensor # [2 inter / 32, 160] E8M0 + sw2: torch.Tensor # shared down [5120, inter] e4m3 + sw2_s: torch.Tensor + + def check(self, tp: int) -> None: + """The MoE's tensors: 384 experts, TP-sharded intermediates (the routed + one padded as the loader pads it), AITER's A8W4 layout, 32 x 32 E8M0 + blocks for the shared expert.""" + d, fp4, e4m3, u8 = ( + MoeDims(tp), + torch.float4_e2m1fn_x2, + torch.float8_e4m3fn, + torch.uint8, + ) + want = { + "gate_w": ((EXPERTS, HIDDEN), torch.bfloat16), + "bias": ((EXPERTS,), torch.float32), + "w13": ((EXPERTS, 2 * d.inter, HIDDEN // 2), fp4), + "w13_s": ((EXPERTS * 2 * d.inter, HIDDEN // 32), u8), + "w2": ((EXPERTS, HIDDEN, d.inter // 2), fp4), + "w2_s": ((EXPERTS * HIDDEN, d.down_scale_cols), u8), + "sgu": ((2 * d.sh_inter, HIDDEN), e4m3), + "sgu_s": ((2 * d.sh_inter // 32, HIDDEN // 32), u8), + "sw2": ((HIDDEN, d.sh_inter), e4m3), + "sw2_s": ((HIDDEN // 32, d.sh_inter // 32), u8), + } + _check_tensors(self, "MoE", want) + + +class DSV41MonoLayer: + """The mono decode runner of one TP rank (``group``: its TP group, for the + peer-memory handle exchange; a gloo / CPU group).""" + + def __init__(self, tp: int, rank: int, group, device: torch.device | str = "cuda"): + _ensure_writable_flydsl_cache() + from .common.peer_memory import PeerBuffer + + self.tp, self.rank = tp, rank + self.d = Dims(tp) + dev = self.device = torch.device(device) + self._scratch: dict[int, torch.Tensor] = {} + # [epoch, -, -, -, a mark per CTA, the MoE's counters] + self.epoch = torch.zeros(EPOCH_WORDS, dtype=torch.int32, device=dev) + self.q = torch.zeros( + MAX_TOKENS, self.d.heads, HEAD_DIM, dtype=torch.bfloat16, device=dev + ) + self.kt = torch.zeros(MAX_TOKENS * KEYS, dtype=torch.int32, device=dev) + self.klen = torch.zeros(2 * MAX_TOKENS, dtype=torch.int32, device=dev) + self._dummy = torch.zeros(4, dtype=torch.int32, device=dev) + self._zero_rec = torch.zeros(1024, dtype=torch.uint8, device=dev) + # the attention seam's outputs, K1 -> K2 + self.res_mid = torch.zeros( + MAX_TOKENS, HC, HIDDEN, dtype=torch.bfloat16, device=dev + ) + self.post_a = torch.zeros(MAX_TOKENS, HC, dtype=torch.float32, device=dev) + self.comb_a = torch.zeros(MAX_TOKENS, HC, HC, dtype=torch.float32, device=dev) + self.pre_a = torch.zeros(MAX_TOKENS, HC, dtype=torch.float32, device=dev) + self.peer = PeerBuffer(2 * peer_half_bytes(tp), group, rank, tp, dev) + self.peer.bytes.zero_() + self._qk_scale = ( + torch.tensor([HEAD_DIM**-0.5 * _LOG2E], dtype=torch.float32) + .view(torch.int32) + .item() + ) + self._kernels: dict = {} + + def scratch(self, tokens: int) -> torch.Tensor: + """Step width ``tokens``'s scratch: a width's layout has memory of its own + (``layer.scratch_layout``). Allocated at the width's first step, which is + eager: vLLM runs every graph's batch eagerly before capturing it.""" + buf = self._scratch.get(tokens) + if buf is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + f"DSv4.1 mono decode: step width {tokens} first reached inside " + "a CUDA graph capture" + ) + buf = torch.zeros( + scratch_bytes(tokens, self.tp), dtype=torch.uint8, device=self.device + ) + self._scratch[tokens] = buf + return buf + + @staticmethod + def supports(tokens: int) -> bool: + return 1 <= tokens <= MAX_TOKENS + + def kernels(self, tokens: int, ratio: int): + key = (tokens, ratio) + if key not in self._kernels: + b = MonoBuild(tokens, self.tp, ratio) + self._kernels[key] = (build_mono_k1(b), build_mono_k2(b)) + return self._kernels[key] + + def ffn_kernel(self, tokens: int): + key = (tokens, "ffn") + if key not in self._kernels: + self._kernels[key] = build_mono_ffn(MonoBuild(tokens, self.tp, 0)) + return self._kernels[key] + + @staticmethod + def _outs(M: int, residual: torch.Tensor) -> tuple: + dev = residual.device + return ( + torch.empty(M, HIDDEN, dtype=torch.bfloat16, device=dev), + torch.empty_like(residual), + torch.empty(M, HC, 1, dtype=torch.float32, device=dev), + torch.empty(M, HC, HC, dtype=torch.float32, device=dev), + torch.empty(M, HC, dtype=torch.float32, device=dev), + ) + + def forward( + self, + w: MonoLayerWeights, + x: torch.Tensor, # [M, 5120] bf16: the previous FFN output (reduced) + residual: torch.Tensor, # [M, 4, 5120] bf16 + post_mix: torch.Tensor, # [M, 4, 1] f32 + res_mix: torch.Tensor, # [M, 4, 4] f32 + pre_mix: torch.Tensor, # [M, 4] f32 + positions: torch.Tensor, # [M] int64 + slot_mapping: torch.Tensor, # [M] int64 + swa_cache: torch.Tensor, # [blocks, block, 584] uint8 + swa_indices: torch.Tensor, # [>= M, 1, 128] int32 + swa_lens: torch.Tensor, # [>= M] int32 + token_to_req: torch.Tensor, # [>= M] int32 + topk_indices: torch.Tensor | None = None, + comp_cache: torch.Tensor | None = None, + comp_block_table: torch.Tensor | None = None, + outs: tuple | None = None, + ): + """-> (out, residual, post_mix, res_mix, pre_mix): the layer's outputs + (``outs``: buffers to write them into).""" + M = x.shape[0] + assert self.supports(M), M + a = w.attn + assert a is not None + ratio = a.ratio + k1, k2 = self.kernels(M, ratio) + st = torch.cuda.current_stream() + if outs is None: + outs = self._outs(M, residual) + out, res_out, post_out, comb_out, pre_out = outs + for t_ in (x, residual, post_mix, res_mix, pre_mix, positions, slot_mapping): + assert t_.is_contiguous() + assert positions.dtype == slot_mapping.dtype == torch.int64 + assert swa_indices.dtype == swa_lens.dtype == token_to_req.dtype == torch.int32 + assert swa_indices.stride(0) == 128 and swa_cache.dtype == torch.uint8 + assert swa_cache.shape[-1] == 584 and swa_cache.stride(1) == 584 + swa_block = swa_cache.shape[1] + if ratio: + assert topk_indices is not None and comp_cache is not None + assert comp_block_table is not None + assert topk_indices.dtype == torch.int32 and topk_indices.stride(0) == 512 + assert comp_cache.dtype == torch.uint8 and comp_cache.stride(1) == 584 + assert comp_block_table.dtype == torch.int32 + comp, comp_stride, comp_block = ( + comp_cache, + comp_cache.stride(0), + comp_cache.shape[1], + ) + bt, bt_stride, topk = ( + comp_block_table, + comp_block_table.stride(0), + topk_indices, + ) + else: + comp, comp_stride, comp_block = self._dummy, 0, 1 + bt, bt_stride, topk = self._dummy, 0, self._dummy + k1( + residual.data_ptr(), + x.data_ptr(), + post_mix.data_ptr(), + res_mix.data_ptr(), + pre_mix.data_ptr(), + w.hc_attn_fn.data_ptr(), + w.hc_attn_scale.data_ptr(), + w.hc_attn_base.data_ptr(), + w.attn_norm.data_ptr(), + self.res_mid.data_ptr(), + self.post_a.data_ptr(), + self.comb_a.data_ptr(), + self.pre_a.data_ptr(), + a.wqkv.data_ptr(), + a.wqkv_scale.data_ptr(), + a.q_norm.data_ptr(), + a.kv_norm.data_ptr(), + a.wq_b.data_ptr(), + a.wq_b_scale.data_ptr(), + a.cos_sin.data_ptr(), + positions.data_ptr(), + slot_mapping.data_ptr(), + swa_cache.data_ptr(), + swa_cache.stride(0), + swa_block, + self.q.data_ptr(), + swa_indices.data_ptr(), + swa_lens.data_ptr(), + comp_block, + bt.data_ptr(), + bt_stride, + topk.data_ptr(), + token_to_req.data_ptr(), + self.kt.data_ptr(), + self.klen.data_ptr(), + self.scratch(M).data_ptr(), + self.epoch.data_ptr(), + 0, + stream=st, + ) + k2( + self.q.data_ptr(), + swa_cache.data_ptr(), + swa_cache.stride(0), + swa_block, + comp.data_ptr(), + comp_stride, + comp_block, + self.kt.data_ptr(), + self.klen.data_ptr(), + positions.data_ptr(), + a.attn_sink.data_ptr(), + self._qk_scale, + a.cos_sin.data_ptr(), + a.wo_a.data_ptr(), + a.wo_a_scale.data_ptr(), + a.wo_b.data_ptr(), + a.wo_b_scale.data_ptr(), + self._zero_rec.data_ptr(), + self.res_mid.data_ptr(), + self.post_a.data_ptr(), + self.comb_a.data_ptr(), + self.pre_a.data_ptr(), + w.hc_ffn_fn.data_ptr(), + w.hc_ffn_scale.data_ptr(), + w.hc_ffn_base.data_ptr(), + w.ffn_norm.data_ptr(), + res_out.data_ptr(), + post_out.data_ptr(), + comb_out.data_ptr(), + pre_out.data_ptr(), + w.gate_w.data_ptr(), + w.bias.data_ptr(), + w.w13.data_ptr(), + w.w13_s.data_ptr(), + w.w2.data_ptr(), + w.w2_s.data_ptr(), + w.sgu.data_ptr(), + w.sgu_s.data_ptr(), + w.sw2.data_ptr(), + w.sw2_s.data_ptr(), + out.data_ptr(), + self.scratch(M).data_ptr(), + self.peer.local, + self.peer.addresses.data_ptr(), + self.rank, + self.epoch.data_ptr(), + 0, + stream=st, + ) + return outs + + def ffn( + self, + w: MonoLayerWeights, + part: torch.Tensor, # [M, 5120] bf16: this rank's unreduced wo_b output + residual: torch.Tensor, # [M, 4, 5120] bf16, after the attention seam + post_mix: torch.Tensor, # [M, 4, 1] f32, the attention seam's + res_mix: torch.Tensor, # [M, 4, 4] f32 + pre_mix: torch.Tensor, # [M, 4] f32 + outs: tuple | None = None, + ): + """-> (out, residual, post_mix, res_mix, pre_mix) of a layer whose + attention vLLM ran: the attention's TP reduce, the FFN seam, the MoE and + its all-reduce in one launch.""" + M = part.shape[0] + assert self.supports(M), M + for t_ in (part, residual, post_mix, res_mix, pre_mix): + assert t_.is_contiguous() + assert part.dtype == residual.dtype == torch.bfloat16 + if outs is None: + outs = self._outs(M, residual) + out, res_out, post_out, comb_out, pre_out = outs + self.ffn_kernel(M)( + part.data_ptr(), + residual.data_ptr(), + post_mix.data_ptr(), + res_mix.data_ptr(), + pre_mix.data_ptr(), + w.hc_ffn_fn.data_ptr(), + w.hc_ffn_scale.data_ptr(), + w.hc_ffn_base.data_ptr(), + w.ffn_norm.data_ptr(), + res_out.data_ptr(), + post_out.data_ptr(), + comb_out.data_ptr(), + pre_out.data_ptr(), + w.gate_w.data_ptr(), + w.bias.data_ptr(), + w.w13.data_ptr(), + w.w13_s.data_ptr(), + w.w2.data_ptr(), + w.w2_s.data_ptr(), + w.sgu.data_ptr(), + w.sgu_s.data_ptr(), + w.sw2.data_ptr(), + w.sw2_s.data_ptr(), + out.data_ptr(), + self.scratch(M).data_ptr(), + self.peer.local, + self.peer.addresses.data_ptr(), + self.rank, + self.epoch.data_ptr(), + stream=torch.cuda.current_stream(), + ) + return outs diff --git a/vllm/models/deepseek_v41/amd/mono/sources.py b/vllm/models/deepseek_v41/amd/mono/sources.py new file mode 100644 index 0000000000..49f34d8fed --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono/sources.py @@ -0,0 +1,8 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The mono layer's kernel sources, in every build's JIT cache key +(``common.plan.key_tuple``): an edit anywhere in this package rebuilds.""" + +from vllm.models.deepseek_v41.amd.mono.common.plan import source_digest + +SOURCES = source_digest(".") diff --git a/vllm/models/deepseek_v41/amd/mono/stages/__init__.py b/vllm/models/deepseek_v41/amd/mono/stages/__init__.py new file mode 100644 index 0000000000..ebb582c99b --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono/stages/__init__.py @@ -0,0 +1,4 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The V4.1 stages the mono layer adapts from ROCm/ATOM: the mHC seams +(``seam``), the MoE (``moe``) and the GEMV / LDS helpers they share (``gemv``).""" diff --git a/vllm/models/deepseek_v41/amd/mono/stages/dims.py b/vllm/models/deepseek_v41/amd/mono/stages/dims.py new file mode 100644 index 0000000000..6fd119c772 --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono/stages/dims.py @@ -0,0 +1,60 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at 97359d6, MIT License, +# Copyright (C) 2026, Advanced Micro Devices, Inc.: +# atom/models/deepseek_v41/mono/kernels/dims.py +"""A TP rank's share of V4.1's MoE widths: every constant the MoE stages derive +from the TP size, in one place (a build key's ``tp`` picks it).""" + +from dataclasses import dataclass + +from vllm.models.deepseek_v41.amd.mono.common.plan import Shard, padded + +# the model's full widths +MOE_INTER = 2304 # a routed expert's intermediate width +SHARED_INTER = 2304 # the shared expert's +# the loader pads a routed expert's rank share to a multiple of this (``FusedMoE`` +# ``pad_align``) with zero weights +MOE_PAD = 128 +ROWS = 16 # GEMV rows a task +UG_PART = 32 # intermediate columns an ug task part: one FP4 group + + +@dataclass(frozen=True) +class Dims: + tp: int + + def __post_init__(self): + shard = Shard(self.tp) + shard.split(MOE_INTER, "routed intermediate") + shard.split(SHARED_INTER, "shared intermediate") + + @property + def inter_real(self) -> int: + """A routed expert's real intermediate width on this rank.""" + return MOE_INTER // self.tp + + @property + def inter(self) -> int: + """... padded as the loader pads it (the rest zero weights).""" + return padded(self.inter_real, MOE_PAD) + + @property + def sh_inter(self) -> int: + """The shared expert's intermediate width on this rank.""" + return SHARED_INTER // self.tp + + @property + def shared_tasks(self) -> int: + return self.sh_inter // ROWS + + @property + def mid_words(self) -> int: + """MXFP8 words of a pick's routed intermediate.""" + return self.inter // 4 + + @property + def down_scale_cols(self) -> int: + """w2_s's e8m0 columns: inter / 32, padded to a multiple of 8 + (aiter ``shuffle_scale``).""" + return padded(self.inter // 32, 8) diff --git a/vllm/models/deepseek_v41/amd/mono/stages/gemv.py b/vllm/models/deepseek_v41/amd/mono/stages/gemv.py new file mode 100644 index 0000000000..31a19ea4c7 --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono/stages/gemv.py @@ -0,0 +1,294 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at 97359d6, MIT License, +# Copyright (C) 2026, Advanced Micro Devices, Inc.: +# atom/models/deepseek_v41/mono/kernels/attn_pre.py (its GEMV and LDS helpers). +"""The FP8 GEMV tiles of the seam and MoE stages and the helpers they share: +token tiles, LDS rows of MXFP8 words, scalar loads and the mailbox-to-LDS +copy.""" + +import flydsl.expr as fx +from aiter.ops.flydsl.kernels import buffer_ops as bo +from flydsl.expr import const_expr, range_constexpr +from flydsl.expr.typing import T + +from vllm.models.deepseek_v41.amd.mono.common.mx import ( + UNIT_SCALE, + mfma_scaled, +) +from vllm.models.deepseek_v41.amd.mono.common.ops import ( + CM_NT, + fresh, + rsrc, + traced, +) +from vllm.models.deepseek_v41.amd.mono.common.plan import WAVES +from vllm.models.deepseek_v41.amd.mono.common.sync import POLL_MAX + +ROWS = 16 # GEMV rows a task +TILE = 16 # an MFMA's N: the tokens a GEMV pass takes (``token_tiles``) + + +def ld_f32(ptr, i): + return fx.Float32(bo.buffer_load(rsrc(ptr), i, vec_width=1, dtype=T.f32)) + + +def ld_bf(ptr, i): + return fx.Float32( + fx.BFloat16(bo.buffer_load(rsrc(ptr), i, vec_width=1, dtype=T.bf16)) + ) + + +def token_tiles(s): + """A step's tokens in GEMV passes: [(first token, tokens)]. A GEMV loads its + weights once and runs every pass on them, an x tile in LDS at a time.""" + return [(t0, min(TILE, s - t0)) for t0 in range(0, s, TILE)] + + +def tile_rows(s): + """The LDS rows an x tile holds.""" + return min(s, TILE) + + +# a token's LDS row of MXFP8 words / E8M0 codes, 16 of which a MFMA B operand +# gathers, is padded by 16 B: at a multiple of 64 words they share a bank group +LDS_PAD = 4 + + +def lds_row(words): + return words + LDS_PAD + + +def lds_at(u, row): + """Element u of a span of rows ``row`` wide, at the ``lds_row`` stride.""" + return u + u // row * LDS_PAD + + +def lds_tail(k): + """Words a K = ``k`` row's last K step reads past its ``lds_row`` (selected + out): the last row's buffer must hold them.""" + return max(0, 32 * -(-k // 128) - lds_row(k // 4)) + + +def plus(t0, v): + """Token ``v`` of the tile from ``t0`` (a first tile adds nothing; ``t0`` + traced: a tile a CTA picks at run time).""" + return v if isinstance(t0, int) and t0 == 0 else v + t0 + + +# ---------------------------------------------------------------- GEMV +def _fp8_step_load(c, r_w, r_ws, k, rg, st, row_major=False): + """K step st's global operands of rows 16 rg ..: the two 64-column weight + chunks (the upper clamped into the row) and the raw scale word. The (16, 16) + shuffle lays a row group out in 32-column blocks of 16 rows, a chunk two of + them (lanes 0-31, 32-63); ``k`` % 64 = 32 (a last quarter step): the lanes + of a block past the row read its last block. ``row_major`` (vLLM's dense + layout, k % 64 == 0): the same registers from a row-major [N, K] weight -- + lane l row 16 rg + l % 16, K bytes 128 st + 16 (l // 16) .. and + 64.""" + lane = c["lane"] + j = lane // 16 + kg = k // 32 + if const_expr(row_major): + assert k % 64 == 0 + row = rg * 16 + lane % 16 + lo_at = row * (k // 4) + st * 32 + j * 4 + hi_at = row * (k // 4) + fx.min(2 * st + 1, k // 64 - 1) * 16 + j * 4 + lo = fx.Vector( + bo.buffer_load(r_w, lo_at, vec_width=4, dtype=T.i32, cache_modifier=CM_NT) + ) + hi = fx.Vector( + bo.buffer_load(r_w, hi_at, vec_width=4, dtype=T.i32, cache_modifier=CM_NT) + ) + wsi = (rg // 2) * kg + fx.min(st * 4 + j, kg - 1) + word = fx.Int32(bo.buffer_load(r_ws, wsi // 4, vec_width=1, dtype=T.i32)) + return lo, hi, (word >> (wsi % 4 * 8)) & 0xFF + if const_expr(k % 64 == 0): + lo_at = (rg * (k // 64) + 2 * st) * 256 + lane * 4 + else: + lo_at = rg * (k * 4) + fx.min(4 * st + lane // 32, kg - 1) * 128 + lane % 32 * 4 + lo = fx.Vector( + bo.buffer_load(r_w, lo_at, vec_width=4, dtype=T.i32, cache_modifier=CM_NT) + ) + if const_expr(k % 64 == 0): + hi_at = (rg * (k // 64) + fx.min(2 * st + 1, k // 64 - 1)) * 256 + lane * 4 + else: + hi_at = rg * (k * 4) + fx.min(2 * st + 1, k // 64 - 1) * 256 + lane * 4 + hi = fx.Vector( + bo.buffer_load(r_w, hi_at, vec_width=4, dtype=T.i32, cache_modifier=CM_NT) + ) + wsi = (rg // 2) * kg + fx.min(st * 4 + j, kg - 1) + word = fx.Int32(bo.buffer_load(r_ws, wsi // 4, vec_width=1, dtype=T.i32)) + return lo, hi, (word >> (wsi % 4 * 8)) & 0xFF + + +def _fp8_step_mfma(c, k, xl, xsl, st, ops, acc, live=None, rows=None): + """``acc`` through K step st's scaled MFMA: ``_fp8_step_load``'s operands + against every token's MXFP8 row in LDS. ``live`` false: a step past the end + (its operands those of the last step), zero at unit scale. ``k`` % 64 = 32: + a lane's lower block past the row is zero on both sides.""" + s, lane = c["S"] if rows is None else rows, c["lane"] + lo, hi, sa = ops + kg = k // 32 + x_row, s_row = lds_row(k // 4), lds_row(kg) + j = lane // 16 + col = fx.min(lane % 16, s - 1) + if const_expr(live is not None): + st = fx.min(st, (k + 127) // 128 - 1) + lo = [live.select(lo[d], fx.Int32(0)) for d in range(4)] + past = (2 * st + 1) * 64 >= k + if const_expr(live is not None): + past = past | ~live + av = fx.Vector.from_elements( + [lo[d] for d in range(4)] + [past.select(fx.Int32(0), hi[d]) for d in range(4)], + fx.Int32, + ) + grp = st * 4 + j + xlo = fx.Vector( + fx.ptr_load( + xl + (col * x_row + st * 32 + j * 4), + result_type=fx.Vector.make_type(4, fx.Int32), + ) + ) + xhi = fx.Vector( + fx.ptr_load( + xl + (col * x_row + st * 32 + 16 + j * 4), + result_type=fx.Vector.make_type(4, fx.Int32), + ) + ) + if const_expr(k % 64 != 0): + over = 4 * st + j // 2 >= kg + av = fx.Vector.from_elements( + [over.select(fx.Int32(0), av[d]) for d in range(4)] + + [av[4 + d] for d in range(4)], + fx.Int32, + ) + xlo = [over.select(fx.Int32(0), xlo[d]) for d in range(4)] + xv = fx.Vector.from_elements( + [xlo[d] for d in range(4)] + + [past.select(fx.Int32(0), xhi[d]) for d in range(4)], + fx.Int32, + ) + sb = fx.ptr_load(xsl + (col * s_row + fx.min(grp, kg - 1))) + unit = grp >= kg + if const_expr(live is not None): + unit = unit | ~live + sa = unit.select(fx.Int32(UNIT_SCALE), sa) + sb = unit.select(fx.Int32(UNIT_SCALE), sb) + return mfma_scaled(av, xv, acc, sa, sb) + + +def _gemv_steps(c, k, split, tiled=False): + """This wave's K steps (a step past the end marked, ``split`` waves past + split repeat the last one's; ``tiled``: wave w takes part w % split of tile + w / split): [(step, live)].""" + wave = c["wave"] + steps = (k + 127) // 128 + if const_expr(split is None): + per = (steps + WAVES - 1) // WAVES + ragged = steps % WAVES != 0 + return [ + (wave + WAVES * i, (wave + WAVES * i < steps) if ragged else None) + for i in range(per) + ] + assert split <= WAVES and steps % split == 0 + per = steps // split + part = wave % split if tiled else fx.min(wave, split - 1) + return [(part * per + i, None) for i in range(per)] + + +def gemv_fp8_loads(c, w, ws, k, rg, split=None, tiled=False, row_major=False): + """The weight operands of ``gemv_fp8``'s rows for this wave, every one in + flight at once (a step's latency each would put them in series): loads + only. A step past the end loads the last step. ``tiled`` (split x tiles = + WAVES): row groups rg .., one a group of ``split`` waves.""" + steps = (k + 127) // 128 + r_w, r_ws = rsrc(w), rsrc(ws) + if const_expr(tiled): + assert WAVES % split == 0 + rg = rg + c["wave"] // split + return [ + _fp8_step_load( + c, + r_w, + r_ws, + k, + rg, + st if live is None else fx.min(st, steps - 1), + row_major, + ) + for st, live in _gemv_steps(c, k, split, tiled) + ] + + +def gemv_fp8_mfmas(c, k, xl, xsl, red, ops, split=None, tiled=False, rows=None): + """``gemv_fp8_loads``' operands against the rows in LDS (``rows``: a token + tile's, every token's by default) -> + ``red[wave][lane]``; the scaled 16x16x128 MFMA, chunk kc in dwords 0-3 and kc + + 1 in 4-7. ``k`` a multiple of 32: a last half step has its upper chunk zeroed + on both sides at unit scale (a NaN byte past the row must not reach the + MFMA), its loads clamped into the row: the last row group's would read past + the tensor; a last quarter step its lower chunk's upper block too. A step + past the end adds zero at unit scale. + + The K order is the original kernel's, which ``row_sum`` completes: + wave w takes steps w, w + 8, ... (``split`` None: aiter's preshuffled group32 + GEMM, P0.5 microbenchmark 1), or waves 0 .. split - 1 one contiguous K range + each, the other waves zero (aiter's split-K bmm, whose last split sums the + partials in split order).""" + lane, wave = c["lane"], c["wave"] + acc = fx.Vector.filled(4, 0.0, fx.Float32) + for i, (st, live) in enumerate(_gemv_steps(c, k, split, tiled)): + acc = _fp8_step_mfma(c, k, xl, xsl, st, ops[i], acc, live, rows) + if const_expr(split is not None and not tiled): + acc = fx.Vector.from_elements( + [(wave < split).select(acc[e], fx.Float32(0.0)) for e in range(4)], + fx.Float32, + ) + fx.ptr_store(acc, red + (wave * 64 + lane) * 4) + + +# ---------------------------------------------------------------- mailbox -> LDS +@traced +def poll_copy(c, idx, width, spans, live=None): + """Mailbox spans [(region, pair_of, n, dst, row)] into LDS: element u (pair + ``pair_of(u)``) to dst + u, or ``lds_at(u, row)`` for a span of rows + ``row`` words wide. Thread ``idx`` of ``width`` takes u = idx, idx + + width, ..., POLL_MAX a round trip, each batch stored before the next is + polled (holding them all spills past S = 16); a thread past a span's end + re-reads its last element and stores nothing. + + ``live``: a span's elements the reader needs (traced, <= n), a span each; + a batch past them is not polled and an element past them neither polled + nor written (one past a step's rows is never written).""" + idx = fresh(idx) # its elements' indices not hoisted out of a task loop + if const_expr(live is None): + _poll_store_batches( + c, + [ + (region, pair_of, n, (dst, row), idx + width * i) + for region, pair_of, n, dst, row in spans + for i in range((n + width - 1) // width) + ], + ) + return + for (region, pair_of, n, dst, row), need in zip(spans, live): + chunks = [ + (region, pair_of, need, (dst, row), idx + width * i) + for i in range((n + width - 1) // width) + ] + for b0 in range_constexpr(0, len(chunks), POLL_MAX): + if chunks[b0][4] < need: + _poll_store_batches(c, chunks[b0 : b0 + POLL_MAX]) + + +@traced +def _poll_store_batches(c, chunks): + """``poll_copy``'s chunks [(region, pair_of, n, (dst, row), u)], POLL_MAX a + round trip, each batch stored before the next is polled.""" + for b0 in range_constexpr(0, len(chunks), POLL_MAX): + batch = chunks[b0 : b0 + POLL_MAX] + got = c["poll"]([(r, p(fx.min(u, n - 1)), 1) for r, p, n, _, u in batch]) + for (_, _, n, (dst, row), u), v in zip(batch, got): + at = u if row is None else lds_at(u, row) + if u < n: + fx.ptr_store(v[0], dst + at) diff --git a/vllm/models/deepseek_v41/amd/mono/stages/moe.py b/vllm/models/deepseek_v41/amd/mono/stages/moe.py new file mode 100644 index 0000000000..63f8947166 --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono/stages/moe.py @@ -0,0 +1,1635 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at 97359d6, MIT License, +# Copyright (C) 2026, Advanced Micro Devices, Inc.: +# atom/models/deepseek_v41/mono/kernels/moe.py +"""The mono layer's MoE stages: one V4.1 layer's FFN, from the FFN seam's normed +rows to the reduced output (its TP all-reduce in-kernel). + +One launch a layer and rank, ``BLOCKS`` x ``THREADS``, S = ``tokens`` rows, +E routed experts and top-k (``MoeBuild``: the target's 384 / 6, the DSpark +draft's 128 / 3). The rank holds every routed expert's 1/TP of the +intermediate width (TP-sharded MoE) and the shared expert's 1/TP (``Dims``). +Stages, with TP4's task counts: + + xq (S): the row as MXFP8 for the routed experts (aiter + ``fused_mx_quant_moe_sort``: amax seeded 1e-10) and for the shared + one (ATOM ``quantize_fp8``: amax floored 1e-4) -> X8R, X8 (plain + words, then the token's flag X8RDY) + router (E / 16 x 4, 16 rows x a quarter of K): bf16 GEMV -> the part's + f32 sum -> LOGIT + route (S): the parts summed in order -> bf16 logits; aiter + ``topk_gating_kernel_opt`` (sqrtsoftplus + bias, top-k, renormalized + x 1.5), one wave a token -> ROUTE + shared (36, 16 gate + 16 up rows): FP8 GEMV -> bf16 -> aiter silu_and_mul + (gate <= 10, up in [-10, 10], ocml exp) -> bf16 -> SMID + ug (|U| x 9, 2 x (32 gate + 32 up) rows of one expert of the union U of + every token's top-k -- the draft's |U| x 18 of 1 x (32 + 32): + ``RouteShape.ug_groups``; w13 gate / up interleaved a 16-row tile): + the MXFP4 x MXFP8 GEMV, a wave the gate and up 16-row tiles of one + 32-row scale block over a half (the draft's: a quarter) of K (each + B operand from LDS read once for both), the slices summed + pairwise; then its epilogue: + clamp, silu x up, bf16, + and the stage-2 input quant to MXFP8 (amax seeded 1e-10, the ceil + code) -> MID (plain words at device scope, then the task's flag, UGF) + down (320, 16 rows; a CTA's two run together): per expert of U the MXFP4 + x MXFP8 GEMV of its intermediate (x its routing weight -> bf16), + summed in each token's top-k order, rounded to bf16 at every add + (aiter's stage 2 adds bf16 atomically, in no fixed order); the + shared expert's FP8 down projection (quantize_fp8 of SMID) -> bf16; + routed + shared -> bf16 -> pushed to every rank's FFN region + reduce (160, 32 columns): the TP partials in rank order 0 .. TP - 1 + (``sum_partials``), fp32 -> bf16 -> ``out`` +""" + +import flydsl.expr as fx +from aiter.ops.flydsl.kernels import buffer_ops as bo +from flydsl.expr import const_expr, gpu, range_constexpr, rocdl +from flydsl.expr import math as fmath +from flydsl.expr.typing import T + +from vllm.models.deepseek_v41.amd.mono.common.mx import ( + FP4, + FP8, + FP8_MAX, + UNIT_SCALE, + clamp_fp8, + code_ceil, + fp4_tile_load, + mfma_scaled, + rcp_pow2, + scale_index, +) +from vllm.models.deepseek_v41.amd.mono.common.ops import ( + CM_DEV, + CM_NT, + _hw_f32, + ballot, + bf16_round, + bf_hi, + bf_lo, + butterfly, + div_rn, + fp8_pack4, + fresh, + hw_exp2, + hw_rcp, + lane_gather, + lanes_below, + popcount, + row_sum, + rsrc, + stamp, + sum_partials, + traced, +) +from vllm.models.deepseek_v41.amd.mono.common.plan import ( + BLOCKS, + THREADS, + WAVES, + first_task, + pair_layout, +) +from vllm.models.deepseek_v41.amd.mono.common.sync import preg, publish, sreg +from vllm.models.deepseek_v41.amd.mono.stages import moe_plan as mplan +from vllm.models.deepseek_v41.amd.mono.stages.dims import UG_PART + +# helpers the kernel calls are imported by name: FlyDSL keys a build by the +# source of the same-directory functions it names, never through a module +from vllm.models.deepseek_v41.amd.mono.stages.gemv import ( + ROWS, + TILE, + gemv_fp8_loads, + gemv_fp8_mfmas, + ld_f32, + lds_row, + lds_tail, + plus, + poll_copy, + tile_rows, + token_tiles, +) +from vllm.models.deepseek_v41.amd.mono.stages.moe_shape import ( + MASK_WORDS, + MoeBuild, + route_shape, +) +from vllm.models.deepseek_v41.amd.mono.stages.router import PARTS as ROUTER_PARTS +from vllm.models.deepseek_v41.amd.mono.stages.router import ( + router_tile_mfma, + router_weight_loads, + router_x_loads, +) +from vllm.models.deepseek_v41.amd.mono.stages.seam import HIDDEN, attn_peer_bytes + +# a pick tile is a token tile's MFMA N: ``gather_x8r`` fills one from the other +assert mplan.TILE == TILE +SCALE = 1.5 +LIMIT = 10.0 # swiglu_limit +LOG2E = 1.4426950408889634 +LN2 = 0.6931471805599453 +DOWN_TASKS = HIDDEN // ROWS # 320 +# (task, slot) scale blocks a down wave holds a round (``stage_down``): 3 took +# no less time, 4 spills +DOWN_R_MAX = 2 +# a down slot's K steps on this many independent MFMA accumulators +DOWN_CHAINS = 2 +# a reduce task's columns: at 128 a thread polled 6 pairs at 48 rows, each +# pass a round trip in series (32: 1-2 passes, S=48 -1 us) +RED_COLS = 32 +RED_TASKS = HIDDEN // RED_COLS # 160 +# CTA b polls route copy b % ROUTE_COPIES: every CTA on one small region +# serializes on its lines (S=24: -0.5 us) +ROUTE_COPIES = 8 +X8_WORDS = HIDDEN // 4 +TAGS = 256 # the ug queue's / down count's slots, a word each: ``c["tag"]`` picks one + +XQ0 = 0 +DOWN0 = 0 +RED0 = 0 + + +# ``timeline`` stamps: kernel start, then the end of each stage's task loop +TL_STAGES = ("xq", "router", "route", "shared", "ug", "down", "reduce") +TL_POINTS = 1 + len(TL_STAGES) + + +# the MoE stages' kernel arguments, in launch order +ARGS = ( + "x", + "gate_w", + "bias", + "w13", + "w13_s", + "w2", + "w2_s", + "sgu", + "sgu_s", + "sw2", + "sw2_s", + "out", +) + + +def scratch_layout(key: MoeBuild) -> dict[str, tuple[int, int]]: + tokens, experts, topk = key.tokens, key.experts, key.topk + rs = route_shape(key) + d = rs.dims + pairs = { + # the routed experts' MXFP8 x (aiter's fused_mx_quant_moe_sort) and the + # shared expert's (ATOM quantize_fp8): their amax floors differ + "x8r": tokens * X8_WORDS // 2, + "x8rs": tokens * HIDDEN // 32 // 2, + "x8": tokens * X8_WORDS // 2, + "x8s": tokens * HIDDEN // 32 // 2, + # a token's four MXFP8 regions written (plain words): ``stage_xq`` + "x8rdy": tokens, + "logit": tokens * experts * ROUTER_PARTS, + "route": ROUTE_COPIES * tokens * 2 * topk, + "smid": tokens * d.sh_inter, + # SMID quantize_fp8'd (``smid_quant``): the down stage's MXFP8 rows + "smq": tokens * d.sh_inter // 4, + "smc": tokens * d.sh_inter // 32, + # a pick a row, in the slot order (``RouteShape.order0``) + "mid": rs.picks * d.mid_words, + "mids": rs.picks * d.inter // 32, + # an ug task's done flag (``stage_ug``) + "ugf": (rs.max_tiles if rs.multi else rs.max_u) * rs.ug_parts, + # the ug task queue: a word a mailbox tag (a layer), counting up from the + # step's clear (``take_ug_task``); plain words, 2 a pair + "ugq": TAGS // 2, + # the down stage's arrival count, a word a tag (``take_down_rank``) + "dq": TAGS // 2, + } + return pair_layout(pairs.items()) + + +def peer_bytes(tokens: int, tp: int) -> int: + """FFN: ``[source rank][token][hidden / 2]`` bf16 pairs.""" + return tp * tokens * HIDDEN // 2 * 8 + + +def ffn_offset(tokens: int, tp: int) -> int: + """FFN's byte offset in the peer buffer: past the ATTN region.""" + return attn_peer_bytes(tokens, tp) + + +def region_pair(src, t, s, col): + return (src * s + t) * (HIDDEN // 2) + col // 2 + + +# ---------------------------------------------------------------- xq +@traced +def stage_xq(c, t): + """Token t, a group of 32 a thread.""" + tid = c["tid"] + a = c["args"] + # thread g < 160 owns the 32 values of group g; the rest mirror thread 159 + g = fx.min(tid, HIDDEN // 32 - 1) + mine = tid < HIDDEN // 32 + # 32 bf16 as four 16-byte loads (a load an element made the stage issue bound) + words = [ + fx.Vector( + bo.buffer_load( + rsrc(a["x"]), + (t * HIDDEN + g * 32) // 2 + 4 * q, + vec_width=4, + dtype=T.i32, + cache_modifier=c["x_cm"], + ) + ) + for q in range(4) + ] + v = [(bf_lo if i % 2 == 0 else bf_hi)(words[i // 8][i % 8 // 2]) for i in range(32)] + amax0 = abs(v[0]) + for y in v[1:]: + amax0 = fx.max(amax0, abs(y)) + # the routed experts' (aiter fused_mx_quant_moe_sort): amax floored 1e-10; + # the shared expert's: vLLM's mxfp8_e4m3_quantize (its MXFP8 linear) + if mine: + put_mxfp8_group(c, "x8r", "x8rs", t, g, v, fx.max(amax0, fx.Float32(1e-10))) + put_mxfp8_group_vllm(c, "x8", "x8s", t, g, v, amax0) + + +def put_mxfp8_group(c, words_mb, codes_mb, t, g, v, amax): + """Token t's group g of 32 values ``v`` as MXFP8 (the ceil code of + ``amax`` / 448, x / scale clamped, RNE) -> scratch ``words_mb`` / + ``codes_mb`` (plain words: the token's X8RDY announces them).""" + code = code_ceil(amax * fx.Float32(1.0 / FP8_MAX)) + inv = rcp_pow2(code) + words = [] + for w4 in range_constexpr(4): + ws = [ + fp8_pack4(*[clamp_fp8(v[8 * w4 + 4 * h + i] * inv) for i in range(4)]) + for h in range(2) + ] + words.append(ws) + for h in range_constexpr(2): + st_dev( + fx.Vector.from_elements(words[2 * h] + words[2 * h + 1], fx.Int32), + c[words_mb], + t * X8_WORDS + g * 8 + 4 * h, + ) + st_dev(fx.Int32(code), c[codes_mb], t * (HIDDEN // 32) + g) + + +FLT_MIN = 1.1754943508222875e-38 + + +def vllm_mx_code(amax): + """VLLM's ``mxfp8_e4m3_quantize`` scale code of a 32-group: the smallest + power of two >= RN(max(amax, FLT_MIN) / 448), clamped to [0, 254].""" + a = fx.max(amax, fx.Float32(FLT_MIN)) + q = div_rn(a, fx.Float32(FP8_MAX), fx.Float32(1.0 / FP8_MAX)) + bits = q.bitcast(fx.Int32) + e = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0).select( + fx.Int32(1), fx.Int32(0) + ) + return fx.min(fx.max(e, fx.Int32(0)), fx.Int32(254)) + + +def put_mxfp8_group_vllm(c, words_mb, codes_mb, t, g, v, amax): + """``put_mxfp8_group`` with vLLM's rule (``vllm_mx_code``; x * 2^(127 - code), + which code 0 keeps finite).""" + code = vllm_mx_code(amax) + inv = ((254 - code) << 23).bitcast(fx.Float32) + words = [] + for w4 in range_constexpr(4): + ws = [ + fp8_pack4(*[clamp_fp8(v[8 * w4 + 4 * h + i] * inv) for i in range(4)]) + for h in range(2) + ] + words.append(ws) + for h in range_constexpr(2): + st_dev( + fx.Vector.from_elements(words[2 * h] + words[2 * h + 1], fx.Int32), + c[words_mb], + t * X8_WORDS + g * 8 + 4 * h, + ) + st_dev(fx.Int32(code), c[codes_mb], t * (HIDDEN // 32) + g) + + +# ---------------------------------------------------------------- router +@traced +def stage_router(c, task): + """Router task (token group, row tile, K part): K part ``part`` of rows + 16 ``tile`` .. of the bf16 gate weight [experts, 5120] times the group's + tokens (``RouteShape.router_groups``), a token tile at a time -- the + decode router's tile (``router``), the + same tiles in the same order -> the part's f32 sum (``route`` sums the + parts in order, rounds to bf16). Every tile's x loads issue first: a task + over every token's tiles in series read their x serially (48 rows: a router + CTA 9 us, 8 of them x).""" + s, tid, red, lane, wave = c["S"], c["tid"], c["red"], c["lane"], c["wave"] + a = c["args"] + experts = c["rs"].experts + groups, per = c["rs"].router_groups + per_group = experts // ROWS * ROUTER_PARTS + tile, part = task % per_group // ROUTER_PARTS, task % ROUTER_PARTS + g0 = 0 if groups == 1 else task // per_group * per + # the group's token tiles: (first token, tokens), tokens traced past one group + tiles = [ + ( + plus(g0, j0), + min(TILE, s - j0) + if groups == 1 + else fx.min(fx.Int32(min(TILE, per - j0)), s - plus(g0, j0)), + ) + for j0 in range(0, per, TILE) + ] + wvs = router_weight_loads(a["gate_w"], tile, part, lane, wave, CM_NT) + xs = [router_x_loads(a["x"], part, n, lane, wave, c["x_cm"], t0) for t0, n in tiles] + for (t0, n), xvs in zip(tiles, xs): + router_tile_mfma(wvs, xvs, lane, wave, red) + gpu.barrier() + if tid < ROWS * n: + r = tid % ROWS + tl = tid // ROWS + v = row_sum(red, r, tl) + c["put"]( + c["logit"], + (plus(t0, tl) * experts + tile * ROWS + r) * ROUTER_PARTS + part, + v, + ) + gpu.barrier() + + +# ---------------------------------------------------------------- route +def _score(x): + """sqrt(log2(1 + exp2(x log2e)) ln2), x itself past 20 (topk_gating).""" + sp = fx.Float32( + (x > fx.Float32(20.0)).select( + x, + _hw_f32( + "llvm.amdgcn.log.f32", fx.Float32(1.0) + hw_exp2(x * fx.Float32(LOG2E)) + ) + * fx.Float32(LN2), + ) + ) + return fx.Float32(fmath.sqrt(sp)) + + +@traced +def stage_route(c, t): + """VLLM's topk_hash_softplus_sqrt for token t on wave 0 (f32 logits).""" + lane, wave = c["lane"], c["wave"] + experts, topk = c["rs"].experts, c["rs"].topk + a = c["args"] + if wave == 0: + ept = experts // 64 + # each logit's router part sums, added in part order (vLLM's gate: f32 out) + assert ROUTER_PARTS == 4 + got = c["poll"]( + [ + (c["logit"], (t * experts + lane + 64 * i) * ROUTER_PARTS + h, 2) + for i in range(ept) + for h in (0, 2) + ] + ) + logits = [] + for i in range_constexpr(ept): + parts = got[2 * i] + got[2 * i + 1] + tot = fx.Float32(0.0) + for p in range_constexpr(ROUTER_PARTS): + tot = tot + parts[p].bitcast(fx.Float32) + logits.append(tot) + vals, orig, idxs = [], [], [] + for i in range_constexpr(ept): + e = lane + 64 * i + sc = _score(logits[i]) + # aiter's NaN guards: a NaN score weighs 0 and never wins + orig.append(fx.Float32(fx.isnan(sc).select(fx.Float32(0.0), sc))) + vals.append(fx.max(sc + ld_f32(a["bias"], e), fx.Float32(float("-inf")))) + idxs.append(fx.Int32(e)) + for p, q in c["rs"].sort_net: + sw = vals[p] < vals[q] + vals[p], vals[q] = sw.select(vals[q], vals[p]), sw.select(vals[p], vals[q]) + orig[p], orig[q] = sw.select(orig[q], orig[p]), sw.select(orig[p], orig[q]) + idxs[p], idxs[q] = sw.select(idxs[q], idxs[p]), sw.select(idxs[p], idxs[q]) + cursor = fx.Int32(0) + total = fx.Float32(0.0) + my_w, my_id = fx.Float32(0.0), fx.Int32(0) + for k in range_constexpr(topk): + cv, ci, co = fx.Float32(float("-inf")), fx.Int32(0), fx.Float32(0.0) + for i in range_constexpr(ept): + at = cursor == i + cv, ci, co = ( + at.select(vals[i], cv), + at.select(idxs[i], ci), + at.select(orig[i], co), + ) + best = butterfly(cv, (32, 16, 8, 4, 2, 1), fx.max) + ball = ballot(cv == best) + win = fx.Int32( + (ball == fx.Int64(0)).select(fx.Int32(0), fx.Int32(fmath.cttz(ball))) + ) + win_id = fx.Int32(rocdl.readlane(T.i32, ci.ir_value(), win.ir_value())) + won = (cursor < ept) & (ci == win_id) + mine_o = fx.Float32(won.select(co, fx.Float32(0.0))) + w = fx.Int32( + rocdl.readlane( + T.i32, mine_o.bitcast(fx.Int32).ir_value(), (win_id & 63).ir_value() + ) + ).bitcast(fx.Float32) + cursor = cursor + won.select(1, 0) + total = total + w + my_w = (lane == k).select(w, my_w) + my_id = (lane == k).select(win_id, my_id) + scale = fx.Float32(SCALE) / fx.max(total, fx.Float32(1e-20)) + if lane < topk: + for r in range_constexpr(ROUTE_COPIES): + at = (r * c["S"] + t) * 2 * topk + lane + c["put"](c["route"], at, my_id) + c["put"](c["route"], at + topk, my_w * scale) + gpu.barrier() + + +@traced +def load_route(c, tab): + """Every token's (ids, weights) into LDS ``tab``, then the union U of the ids + (ascending): tab[2 S topk ..] = |U|, the experts of U, each (t, k)'s slot in + U, each expert's slot, the picks in slot order (a slot's ascending) and each + slot's first position (``RouteShape``). + + A pick a thread, no scan in series: each pick ORs its token's bit into its + expert's token bitmap, so an expert's pick count is its bitmap's popcount, + a slot's first position the counts of the experts below it (a wave's prefix + sum; slots follow the experts ascending), and a pick's rank in its slot the + bitmap's tokens below its own. The bitmaps are the same whatever order the + ORs land in, so every CTA derives the same pick order.""" + s, tid, lane, wave = c["S"], c["tid"], c["lane"], c["wave"] + rs = c["rs"] + topk, max_u, experts = rs.topk, rs.max_u, rs.experts + n = s * topk + assert s <= 32 * MASK_WORDS + # a thread's route words in one poll batch: a pass each was a round trip + # each in series past 256 picks (48 rows: 5 us from the last route to the + # table); a word past 2n re-reads the last and is not stored + idxs = [plus(pas, tid) for pas in range(0, 2 * n, THREADS)] + copy_at = c["bid"] % ROUTE_COPIES * (2 * n) + got = c["poll"]([(c["route"], copy_at + fx.min(idx, 2 * n - 1), 1) for idx in idxs]) + for i in range_constexpr(len(idxs)): + if idxs[i] < 2 * n: + fx.ptr_store(got[i][0], tab + idxs[i]) + for pas in range_constexpr(0, MASK_WORDS * experts, THREADS): + idx = plus(pas, tid) + if idx < MASK_WORDS * experts: + fx.ptr_store(fx.Int32(0), tab + (rs.mask0 + idx)) + gpu.barrier() + if tid < n: + t = tid // topk + eid = fx.ptr_load(tab + (t * 2 * topk + tid % topk)) + fx.llvm.atomic_or( + tab + (rs.mask0 + eid * MASK_WORDS + t // 32), + fx.Int32(1) << (t % 32), + syncscope="workgroup", + ) + gpu.barrier() + base = 2 * n + # a chunk of 64 experts a wave (wave 0 alone through 6 chunks in series + # was 3.3 us of the table's 5 at 48 rows): each chunk's own counts and + # prefixes, its totals through LDS ``red`` (free here), then each chunk + # past the totals of the chunks below it -- the serial walk's values + chunks = experts // 64 + assert chunks <= WAVES + sums = c["red"] + e = lane + 64 * fx.min(wave, chunks - 1) + count = fx.Int32(0) + for w in range_constexpr(MASK_WORDS): + word = fx.ptr_load(tab + (rs.mask0 + e * MASK_WORDS + w)) + count = count + popcount(word) + ball = ballot(count != 0) + # this chunk's inclusive prefixes, lanes ascending: picks, and pick tiles + incl = lane_prefix(count, lane) + tiles = mplan.pick_tiles(count) + tincl = lane_prefix(tiles, lane) if rs.multi else tiles + totals = [popcount(ball), lane_gather(incl, 63), lane_gather(tincl, 63)] + if (wave < chunks) & (lane == 0): + for k in range_constexpr(3): + fx.ptr_store(totals[k].bitcast(fx.Float32), sums + (wave * 3 + k)) + gpu.barrier() + below_c = [fx.Int32(0), fx.Int32(0), fx.Int32(0)] # n_u, first, tile0 + for j in range_constexpr(chunks - 1): + for k in range_constexpr(3): + v = fx.ptr_load(sums + (j * 3 + k)).bitcast(fx.Int32) + below_c[k] = below_c[k] + (j < wave).select(v, fx.Int32(0)) + n_u, first, tile0 = below_c + slot = n_u + lanes_below(ball, lane) + # e afresh: hoisted, its address held a VGPR across the prefix sums (spill) + e_at = fresh(e) + if wave < chunks: + fx.ptr_store(slot, tab + (rs.expert_slot0 + e_at)) + if count != 0: + fx.ptr_store(fx.Int32(e), tab + (base + 1 + slot)) + fx.ptr_store(first + incl - count, tab + (rs.first0 + slot)) + if const_expr(rs.multi): + # the slot's pick tiles: their first index, and each one's slot + tf = tile0 + tincl - tiles + if count != 0: + fx.ptr_store(tf, tab + (rs.tile_first0 + slot)) + for j in range_constexpr(rs.slot_tiles): + if j < tiles: + fx.ptr_store(slot, tab + (rs.tile_slot0 + tf + j)) + if (wave == chunks - 1) & (lane == 0): + n_all = n_u + totals[0] + fx.ptr_store(n_all, tab + base) + fx.ptr_store(fx.Int32(n), tab + (rs.first0 + n_all)) + if const_expr(rs.multi): + fx.ptr_store(tile0 + totals[2], tab + (rs.tile_first0 + n_all)) + gpu.barrier() + if tid < n: + t = tid // topk + eid = fx.ptr_load(tab + (t * 2 * topk + tid % topk)) + slot = fx.ptr_load(tab + (rs.expert_slot0 + eid)) + fx.ptr_store(slot, tab + (base + 1 + max_u + tid)) + # the tokens below t that picked this expert: the picks before tid in it + rank = fx.Int32(0) + for w in range_constexpr(MASK_WORDS): + word = fx.ptr_load(tab + (rs.mask0 + eid * MASK_WORDS + w)) + lo = t - 32 * w + keep = (lo >= 32).select( + fx.Int32(-1), + (lo > 0).select((fx.Int32(1) << fx.max(lo, 0)) - 1, fx.Int32(0)), + ) + rank = rank + popcount(word & keep) + pos = fx.ptr_load(tab + (rs.first0 + slot)) + rank + fx.ptr_store(tid, tab + (rs.order0 + pos)) + gpu.barrier() + + +def lane_prefix(v, lane): + """The wave's inclusive prefix sum of v, lanes ascending.""" + incl = v + for d in (1, 2, 4, 8, 16, 32): + moved = lane_gather(incl, fx.max(lane - d, 0)) + incl = incl + (lane >= d).select(moved, fx.Int32(0)) + return incl + + +# ---------------------------------------------------------------- shared ug +def _silu_mul(g, u): + """VLLM's silu_and_mul_with_clamp (the shared expert's): gate min limit, + up clamped; silu = g / (1 + expf(-g)) rounded to bf16, then times up.""" + g = bf16_round(fx.min(g, fx.Float32(LIMIT))) + u = fx.max(fx.min(u, fx.Float32(LIMIT)), fx.Float32(-LIMIT)) + act = bf16_round(g / (fx.Float32(1.0) + fx.Float32(fmath.exp(-g)))) + return act * u + + +@traced +def stage_shared(c, task, gops, t0, n): + """Gate rows 16 task .. (weights ``gops``), up rows 576 + 16 task .. against + tokens t0 .. t0 + n (their x in LDS; ``n`` traced: a tile a CTA picks at run + time) -> bf16 -> silu x up -> SMID; an odd task then quantizes its group of + 32 (``smid_quant``).""" + tid, red, d = c["tid"], c["red"], c["d"] + a = c["args"] + rows = None if isinstance(n, int) else n + gemv_fp8_mfmas(c, HIDDEN, c["xl"], c["xsl"], red, gops, rows=rows) + gpu.barrier() + gate = fx.Float32(0.0) + if tid < ROWS * n: + gate = row_sum(red, tid % ROWS, tid // ROWS) + gpu.barrier() + uops = gemv_fp8_loads( + c, a["sgu"], a["sgu_s"], HIDDEN, d.shared_tasks + task, row_major=True + ) + gemv_fp8_mfmas(c, HIDDEN, c["xl"], c["xsl"], red, uops, rows=rows) + gpu.barrier() + if tid < ROWS * n: + r = tid % ROWS + t = plus(t0, tid // ROWS) + up = row_sum(red, r, tid // ROWS) + v = fx.Float32(_silu_mul(bf16_round(gate), bf16_round(up)).to(fx.BFloat16)) + c["put"](c["smid"], t * d.sh_inter + task * ROWS + r, v) + fx.ptr_store(v, c["pair"] + tid) + gpu.barrier() + if task % 2 == 1: + smid_quant(c, task, t0, n) + gpu.barrier() + + +@traced +def smid_quant(c, task, t0, n): + """Odd shared task ``task``: its 16 columns (LDS ``pair``) and the even + task's (SMID) are a quantize_fp8 group of 32 (amax floored 1e-4, the ceil + code) -> SMQ / SMC, a thread a token. Once here, not in each of the 256 CTAs + whose down stage reads them: that took a down CTA 14 us at 48 rows.""" + tid, sh = c["tid"], c["d"].sh_inter + assert c["d"].shared_tasks % 2 == 0 # tasks pair into groups + if tid < n: + t = plus(t0, tid) + got = c["poll"]( + [(c["smid"], t * sh + (task - 1) * ROWS + r, 1) for r in range(ROWS)] + ) + vals = [got[r][0].bitcast(fx.Float32) for r in range(ROWS)] + [ + fx.ptr_load(c["pair"] + (tid * ROWS + r)) for r in range(ROWS) + ] + amax = abs(vals[0]) + for y in vals[1:]: + amax = fx.max(amax, abs(y)) + # vLLM's down_proj input quant (mxfp8_e4m3_quantize) + code = vllm_mx_code(amax) + inv = ((254 - code) << 23).bitcast(fx.Float32) + g = task // 2 + for w in range_constexpr(8): + word = fp8_pack4(*[clamp_fp8(vals[4 * w + i] * inv) for i in range(4)]) + c["put_words"](c["smq"], t * (sh // 4) + g * 8 + w, [word]) + c["put_words"](c["smc"], t * (sh // 32) + g, [code]) + + +# ---------------------------------------------------------------- routed ug +@traced +def gemv_fp4_pair(c, w, ws, e, n_rows, k, rg0, ks, xl, xsl, scale_cols, col): + """Tiles rg0 and rg0 + 1 (16 rows each, one 32-row scale block) of expert + e's (16, 16)-preshuffled MXFP4 weight [n_rows, k] times 16 MXFP8 rows in + LDS, over this wave's K slice ``ks`` (``ug_rounds``) -> both tiles' partial + C, a Vector of 8 (rows 4 (lane / 16) .., column lane % 16, whose LDS row is + ``col``). Each B operand read from LDS feeds both tiles' MFMAs: the waves of + a task share x, so a wave a tile over all of K read it once a tile. + """ + lane, wave = c["lane"], c["wave"] + sl = c["wsl"] + wave * wsl_words(c["d"]) + slices = c["rs"].ug_k_slices + rounds = ug_rounds(k, ks, slices) + accs = [fx.Vector.filled(4, 0.0, fx.Float32) for _ in range(2)] + # rounds whose loads are in flight past the one computing: a quarter's + # 2-step rounds two (the draft's 20 rows: -0.5 us) + ahead = slices // 2 + ops = [ + fp4_pair_loads(c, w, ws, e, n_rows, k, rg0, scale_cols, *rounds[q]) + for q in range(min(ahead, len(rounds))) + ] + for r in range_constexpr(len(rounds)): + sc0, st0, n, sc_steps = rounds[r] + sv, avs = ops[r] + fx.ptr_store(sv, sl + lane * (sc_steps // 2)) + sas = [ + [ + lds_u8( + sl, + scale_index( + t * 16 + lane % 16, (st0 + i) * 4 + lane // 16, scale_cols + ) + - 128 * sc0, + ) + for i in range(n) + ] + for t in range(2) + ] + accs = fp4_pair_mfma(c, k, xl, xsl, st0, n, avs, sas, accs, col) + if const_expr(r + ahead < len(rounds)): + ops.append( + fp4_pair_loads( + c, w, ws, e, n_rows, k, rg0, scale_cols, *rounds[r + ahead] + ) + ) + return fx.Vector.from_elements( + [accs[t][q] for t in range(2) for q in range(4)], fx.Float32 + ) + + +def ug_rounds(k, ks, slices): + """K slice ``ks`` of ``slices``' rounds (first step of the scales loaded, + first step, steps, steps of scales loaded); ``ks`` traced: the steps are. + Halves: the 8-step scale blocks 2 r + ks, the odd block left split 4 + 4, + so both run k / 256 steps. Quarters: steps 8 i + 2 ks, + 1 of each block + i, a round loading only their scales (the shuffled layout keeps 2 steps' + together: a dword a lane, where the whole block was 4).""" + blocks = k // 128 // 8 + if slices == 4: + return [(fx.Int32(8 * i) + ks * 2,) * 2 + (2, 2) for i in range(blocks)] + assert slices == 2 + rounds = [ + (fx.Int32(16 * r) + ks * 8, fx.Int32(16 * r) + ks * 8, 8, 8) + for r in range(blocks // 2) + ] + if blocks % 2: + last = (blocks - 1) * 8 + rounds.append((fx.Int32(last), fx.Int32(last) + ks * 4, 4, 8)) + return rounds + + +def fp4_pair_loads(c, w, ws, e, n_rows, k, rg0, scale_cols, sc0, st0, n, sc_steps): + """``gemv_fp4_pair``'s loads of a round: the scales of steps sc0 .. sc0 + + sc_steps (of the tiles' 32-row block: 128 bytes a step, sc_steps / 2 dwords + a lane) and both tiles' weight operands of steps st0 .. st0 + n. Loads + only.""" + lane = c["lane"] + block = (e * n_rows + rg0 * 16) // 32 * (scale_cols * 8) + words = sc_steps // 2 + sv = bo.buffer_load( + rsrc(ws), + block + 32 * sc0 + lane * words, + vec_width=words, + dtype=T.i32, + cache_modifier=CM_NT, + ) + sv = fx.Vector(sv) if words > 1 else fx.Int32(sv) + return sv, [fp4_tile_weights(c, w, e, n_rows, k, rg0 + t, st0, n) for t in range(2)] + + +def scale_block_loads(c, ws, first_row, scale_cols): + """The loads of the 32-row e8m0 block holding ``first_row``: in aiter's + shuffle_scale order a block is contiguous (both 16-row halves, all of K), so + a few 16-byte loads a lane fetch it, where a byte load a K step doubled the + stage's load instructions and bound it by their issue (2026-09-29, ug). A + lane past the block re-reads its last 16 bytes.""" + lane = c["lane"] + words = scale_cols * 8 + block = first_row // 32 * words + return [ + fx.Vector( + bo.buffer_load( + rsrc(ws), + block + fx.min(256 * i + lane * 4, words - 4), + vec_width=4, + dtype=T.i32, + cache_modifier=CM_NT, + ) + ) + for i in range((words + 255) // 256) + ] + + +def down_scale_words(d) -> int: + """A down tile's 32-row scale block in LDS: ``scale_block_store`` writes + whole 256-word rounds.""" + return down_lds_words(d.down_scale_cols * 8) + + +def wsl_words(d) -> int: + """A wave's LDS scale words: an ug batch's 256, or a down round's blocks -- + a (task, slot) each, ``DOWN_R_MAX`` at most (``stage_down``).""" + return max(256, DOWN_R_MAX * down_scale_words(d)) + + +def scale_block_store(c, svs, sl): + """``scale_block_loads``' words -> this wave's LDS block ``sl`` + (``down_scale_words``; a lane past the block writes past it, inside it).""" + lane = c["lane"] + for i in range_constexpr(len(svs)): + fx.ptr_store(svs[i], sl + (256 * i + lane * 4)) + + +def lds_scales(sl, local_row, lane, st0, blk, scale_cols): + """K steps st0 .. st0 + blk's scale bytes of this lane's row from a 32-row + scale block in LDS (``local_row`` its row in the block). An even step and + the next are bytes 2 apart of one word (``scale_index``'s d): a read a + pair, not a step (down: 3 reads a slot, not 5).""" + out, word = [], None + for i in range(blk): + idx = scale_index(local_row, (st0 + i) * 4 + lane // 16, scale_cols) + if (st0 + i) % 2 == 0 or word is None: + word = fx.ptr_load(sl + idx // 4) + out.append((word >> (idx % 4 * 8)) & 0xFF) + return out + + +def lds_u8(ptr, i): + return (fx.ptr_load(ptr + i // 4) >> (i % 4 * 8)) & 0xFF + + +def fp4_tile_weights(c, w, e, n_rows, k, rg, st0, blk, k_real=None): + """The weight operands of expert e's 16-row tile rg, K steps st0 .. st0 + + blk: loads only, nothing waits on them here. ``k_real``: a lane whose 32 + columns lie past it (zero weights) reads its row group's first lane's + address instead (no extra traffic) and holds zero.""" + lane = c["lane"] + base = e * n_rows * (k // 2) // 4 # dwords: expert e + out = [] + for i in range_constexpr(blk): + past = k_real is not None and ((st0 + i) * 128 + 96) >= k_real + if const_expr(past): + real = (st0 + i) * 128 + lane // 16 * 32 < k_real + at = real.select(lane * 4, lane % 16 * 4) + else: + at = lane * 4 + v = fp4_tile_load(rsrc(w), base, rg, k, st0 + i, at) + if const_expr(past): + v = fx.Vector.from_elements( + [real.select(v[q], fx.Int32(0)) for q in range(4)], fx.Int32 + ) + out.append(v) + return out + + +def fp4_pair_mfma(c, k, xl, xsl, st0, n, avs, sas, accs, col): + """``accs`` (two tiles') chained through ``n`` K steps from st0 against MXFP8 + rows in LDS (column lane % 16: row ``col``): a step's B operand loaded once + for both tiles' MFMAs.""" + lane = c["lane"] + x_row = lds_row(k // 4) + s_row = lds_row(k // 32) + j = lane // 16 + accs = list(accs) + for i in range_constexpr(n): + bv = fp8_operand( + [ + fx.ptr_load( + xl + (col * x_row + (st0 + i) * 32 + 16 * h + j * 4), + result_type=fx.Vector.make_type(4, fx.Int32), + ) + for h in range(2) + ] + ) + sb = fx.ptr_load(xsl + (col * s_row + (st0 + i) * 4 + j)) + for t in range_constexpr(2): + accs[t] = mfma_scaled(avs[t][i], bv, accs[t], sas[t][i], sb, FP4, FP8) + return accs + + +def fp8_operand(halves): + """A lane's MXFP8 MFMA operand of a K step: its two 16-byte halves (K + 16 j .. and 64 + 16 j .. of the step, j = lane / 16).""" + return fx.Vector.from_elements( + [fx.Vector(h).bitcast(fx.Int32)[d] for h in halves for d in range(4)], + fx.Int32, + ) + + +def down_lds_words(n): + """A down wave's LDS block of n words: ``scale_block_store`` writes whole + 256-word rounds.""" + return -(-n // 256) * 256 + + +def st_dev(v, region, i): + """Word i of scratch ``region`` (plain words, not tagged pairs) at device + scope: past the per-XCD L2, so another CTA reads it after the flag.""" + bo.buffer_store(v, rsrc(region.value), i, cache_modifier=CM_DEV) + + +@traced +def take_ug_task(c): + """This CTA's next ug task past the first round, from the launch's queue + (``ugq``, a word a tag, counting from 0): the CTAs that ran a shared, + router or route task come to it late and so take fewer, where a fixed + placement gave every CTA the same count and down waited on theirs. A + task's result is the same whichever CTA runs it.""" + if c["tid"] == 0: + at = c["ugq"].value + fx.Int64(c["tag"] * 4) + ptr = fx.inttoptr(fx.PointerType.get(T.i32, fx.AddressSpace.Global, 4), at) + task = fx.llvm.atomic_add(ptr, fx.Int32(1), syncscope="agent") + fx.ptr_store(fx.Int32(task), c["tab"] + c["rs"].queue0) + gpu.barrier() + return fx.ptr_load(c["tab"] + c["rs"].queue0) + + +@traced +def stage_ug(c, task): + """One token tile: slot task // ug_parts of U (``RouteShape.ug_parts``), the + ``ug_groups`` FP4 groups of intermediate columns from 32 ug_groups (task % + ug_parts) (``ug_wave_pair``); then silu x up, rounded to bf16, and the + MXFP8 requant of each token's 32 values of a group -> MID (the stage-1 + epilogue and the stage-2 input quant). Past one: ``stage_ug_tile``.""" + s, tid, lane, tab = c["S"], c["tid"], c["lane"], c["tab"] + rs = c["rs"] + ug_parts = rs.ug_parts + u = task // ug_parts + if u < c["rs"].n_union(tab, s): + e = c["rs"].union_expert(tab, s, u) + # the slot's picks, a column each (a column past them repeats the last); + # every token's row is in LDS: a column reads its pick's token's + first = c["rs"].slot_first(tab, u) + count = c["rs"].slot_first(tab, u + 1) - first + pick = c["rs"].pick_at(tab, first + fx.min(lane % 16, count - 1)) + ug_pass(c, e, task % ug_parts * rs.ug_groups, first, count, pick // rs.topk) + publish(c["put"], c["ugf"], task, 1, tid == 0) + gpu.barrier() + + +def ug_wave_pair(c, part0): + """This wave's share of an ug task from group ``part0``: group p of its + ``ug_groups`` (wave / (2 slices)), its 16-column half bb = wave % 2 -- w13's + gate tile 2 b and up tile 2 b + 1 (aiter GUGU interleaves them a 16-row + tile; b = 2 (part0 + p) + bb), one 32-row scale block -- over K slice ks = + wave / 2 % slices (``gemv_fp4_pair``): (p, 2 b, ks, whether the group is in + the real width).""" + wave, slices = c["wave"], c["rs"].ug_k_slices + p = wave // (2 * slices) + b = (part0 + p) * 2 + wave % 2 + return p, 2 * b, wave // 2 % slices, part0 + p < c["d"].inter_real // UG_PART + + +def ug_red_slot(rs, p, bb, ks, t): + """RED's 64-lane slot of an ug task's partial C: group p, half bb, K slice + ks, gate (t 0) or up (t 1).""" + return ((p * 2 + bb) * rs.ug_k_slices + ks) * 2 + t + + +def pairwise_sum(vs): + """Vs summed as a balanced tree, neighbours first: (v0 + v1) + (v2 + v3).""" + while len(vs) > 1: + vs = [vs[i] + vs[i + 1] for i in range(0, len(vs), 2)] + return vs[0] + + +def ug_tile_task(c, task): + """Past one token tile, ug task ``task``: its slot's expert, its first + group, and its pick tile's first pick and the picks from it.""" + tab, rs, s = c["tab"], c["rs"], c["S"] + ug_parts = rs.ug_parts + tile = task // ug_parts + u = rs.tile_slot(tab, tile) + first = mplan.tile_pick0(rs.slot_first(tab, u), tile, rs.tile_first(tab, u)) + count = rs.slot_first(tab, u + 1) - first + return rs.union_expert(tab, s, u), task % ug_parts * rs.ug_groups, first, count + + +@traced +def stage_ug_tile(c, task, total): + """Past one token tile: ug task ``task`` of ``total`` (a pick tile of its + slot), then the next task from the queue (``take_ug_task``). No next task's + weights prefetched under this one: held across its epilogue and gather they + spilled the merged K2 (S = 24 / 48), and without them it ran faster.""" + tid, lane = c["tid"], c["lane"] + e, part0, first, count = ug_tile_task(c, task) + gather_x8r(c, first, count) + ug_pass(c, e, part0, first, count, fx.min(lane % 16, count - 1)) + nxt = (total > BLOCKS).select(take_ug_task(c) + BLOCKS, total) + # the epilogue's MID stores done before the flag + rocdl.s_waitcnt(vmcnt=0) + gpu.barrier() + if tid == 0: + c["put"](c["ugf"], task, 1) + return nxt + + +@traced +def ug_pass(c, e, part0, first, count, col): + """An ug task's GEMV and epilogue for the slot's picks first .. first + + min(count, 16) (column lane % 16's x in LDS row ``col``).""" + tid, lane, wave, red = c["tid"], c["lane"], c["wave"], c["red"] + a, d, rs = c["args"], c["d"], c["rs"] + p, rg0, ks, real = ug_wave_pair(c, part0) + # a group past the real width (the padding to inter) has zero weights: + # its GEMV is zero unread, and the epilogue still writes its zero MID + acc = fx.Vector.filled(8, 0.0, fx.Float32) + if real: + acc = gemv_fp4_pair( + c, + a["w13"], + a["w13_s"], + e, + 2 * d.inter, + HIDDEN, + rg0, + ks, + c["xl"], + c["xsl"], + HIDDEN // 32, + col, + ) + for t in range_constexpr(2): + slot = ug_red_slot(rs, p, wave % 2, ks, t) + fx.ptr_store( + fx.Vector.from_elements([acc[4 * t + q] for q in range(4)], fx.Float32), + red + (slot * 64 + lane) * 4, + ) + gpu.barrier() + # thread (p, j, i): group p, the slot's pick j, intermediate column + # 32 part + i, a pass of THREADS at a time; a group's 32 threads stay + # lane-aligned + cols = rs.tile_picks + for pas in range_constexpr(0, rs.ug_groups * cols * UG_PART, THREADS): + ug_epilogue(c, part0, pas + tid, first, count) + + +@traced +def load_x8_rows(c, words_mb, codes_mb, xl, xsl, token_of, live): + """MXFP8 rows ``token_of(j)``, j < ``live`` (<= TILE; traced or int), of + scratch ``words_mb`` / ``codes_mb`` into LDS rows j of ``xl`` / ``xsl``: + the first ``live`` lanes wait on their token's X8RDY, then plain 16-byte + loads (half the bytes of the tagged pairs a poll would read).""" + tid = c["tid"] + if tid < TILE: + c["poll"]([(c["x8rdy"], token_of(fx.min(tid, live - 1)), 1)]) + gpu.barrier() + for mb, width, dst in ((words_mb, X8_WORDS, xl), (codes_mb, HIDDEN // 32, xsl)): + quads = width // 4 + for i in range_constexpr(-(-TILE * quads // THREADS)): + q = tid + THREADS * i + row = q // quads + if (q < TILE * quads) & (row < live): + v = fx.Vector( + bo.buffer_load( + rsrc(c[mb].value), + token_of(row) * width + q % quads * 4, + vec_width=4, + dtype=T.i32, + cache_modifier=CM_DEV, + ) + ) + fx.ptr_store(v, dst + (row * lds_row(width) + q % quads * 4)) + gpu.barrier() + + +def gather_x8r(c, first, count): + """The X8R rows and codes of the slot's picks first .. first + min(count, 16) + into LDS rows 0 .. (only those rows: an ug pass reads a pick's row, the + columns past them repeat the last one's).""" + tab, topk = c["tab"], c["rs"].topk + + def token(j): + return c["rs"].pick_at(tab, first + fx.min(j, count - 1)) // topk + + load_x8_rows(c, "x8r", "x8rs", c["x8rl"], c["x8rsl"], token, fx.min(count, TILE)) + + +@traced +def ug_epilogue(c, part0, idx, first, count): + """Thread ``idx``'s value of an ug task: silu x up of its (group, pick j, + column), the group's MXFP8 scale, its byte of MID (the pick's row, ``first`` + + j in the pick order).""" + lane, red, d, rs = c["lane"], c["red"], c["d"], c["rs"] + cols = rs.tile_picks + if idx < rs.ug_groups * cols * UG_PART: + p = idx // (cols * UG_PART) + part = part0 + p + t = idx % (cols * UG_PART) // UG_PART + i = idx % UG_PART + tile = i // 16 + r = i % 16 + at = (16 * (r // 4) + t) * 4 + r % 4 + g, up = [ + pairwise_sum( + [ + fx.ptr_load(red + ug_red_slot(rs, p, tile, ks, t) * 64 * 4 + at) + for ks in range(rs.ug_k_slices) + ] + ) + for t in range(2) + ] + gc = fx.min(g, fx.Float32(LIMIT)) + uc = fx.max(fx.min(up, fx.Float32(LIMIT)), fx.Float32(-LIMIT)) + sig = hw_rcp(fx.Float32(1.0) + hw_exp2(gc * fx.Float32(-LOG2E))) + # the stage-1 kernel stores the intermediate as bf16; the stage-2 input + # quant (fused_mx_quant_moe_sort) makes it MXFP8: amax floored 1e-10, + # the ceil code, x / scale clamped, RNE + v = bf16_round(gc * sig * uc) + amax = fx.max(butterfly(abs(v), (1, 2, 4, 8, 16), fx.max), fx.Float32(1e-10)) + code = code_ceil(amax * fx.Float32(1.0 / FP8_MAX)) + q8 = clamp_fp8(v * rcp_pow2(code)) + nb = [ + lane_gather(q8.bitcast(fx.Int32), fx.min(lane + q, 63)).bitcast(fx.Float32) + for q in range(1, 4) + ] + # 4 threads' bytes -> one word + word = fp8_pack4(q8, *nb) + row = first + t + if (i % 4 == 0) & (t < count): + st_dev(word, c["mid"], row * d.mid_words + part * 8 + i // 4) + if (i == 0) & (t < count): + st_dev(code, c["mids"], mid_code_index(d, row, part)) + + +# ---------------------------------------------------------------- down +@traced +def stage_down(c, tasks): + """Rows 16 task .. of the hidden output for each of ``tasks`` (a CTA's one or + two, run together: a second pass would double the CTA's time): routed + experts of U (a wave each, round robin), the shared expert, their sum -> + pushed to every rank.""" + s, wave, tab, d = c["S"], c["wave"], c["tab"], c["d"] + n = len(tasks) + nu = c["rs"].n_union(tab, s) + sl = c["wsl"] + wave * wsl_words(d) + sw = down_scale_words(d) + # slots of U a round: the loads of a round stay the same + slots = DOWN_R_MAX // n + per_round = WAVES * slots + # the shared expert's down projection first: its input is ready long + # before the routed results, which a CTA mostly waits on here; after the + # rounds it sat on the path past the last ug task + down_shared(c, tasks) + for rnd in range(0, (nu + per_round - 1) // per_round, 1): + # the round's slots of U (wave, wave + WAVES, ...), every load of them + # and of every task in flight before the first wait: the weights and + # their scale blocks (no ug result: before the round's ug flags), then + # the slots' picks' intermediate rows, each lane its pick's straight + # from MID (``mid_operands``). A slot past U loads the last slot again + # (L2 hits) and contributes nothing + us = [wave + WAVES * (slots * rnd + i) for i in range(slots)] + ucs = [fx.min(u, nu - 1) for u in us] + firsts = [c["rs"].slot_first(tab, u) for u in ucs] + svs, avs = down_weights(c, tasks, rnd, slots, nu) + # this round's slots' ug tasks done, not every slot of the wave: the + # round starts while later slots' ug run + c["poll"]([(c["ugf"], f, 1) for f in down_ug_flags(c, ucs)]) + counts = [c["rs"].slot_first(tab, u + 1) - f for u, f in zip(ucs, firsts)] + mids = [mid_operands(c, f, k) for f, k in zip(firsts, counts)] + for q in range_constexpr(len(svs)): + scale_block_store(c, svs[q], sl + q * sw) + for i in range_constexpr(slots): + for j in range_constexpr(n): + down_mfma( + c, + j, + tasks[j], + sl + (j * slots + i) * sw, + mids[i], + firsts[i], + counts[i], + avs[j * slots + i], + us[i] < nu, + ) + # a slot of more than 16 picks: its next pick tile's rows + if const_expr(s > TILE): + for j0 in range_constexpr(TILE, s, TILE): + if j0 < counts[i]: + more = mid_operands(c, firsts[i] + j0, counts[i] - j0) + for j in range_constexpr(n): + down_mfma( + c, + j, + tasks[j], + sl + (j * slots + i) * sw, + more, + firsts[i] + j0, + counts[i] - j0, + avs[j * slots + i], + us[i] < nu, + ) + gpu.barrier() + for t0, nn in token_tiles(s): + for j in range_constexpr(n): + down_combine(c, tasks[j], j, t0, nn) + + +@traced +def take_down_rank(c): + """This CTA's arrival rank at the down stage (DQ, a word a tag): which + CTAs take the extra tiles. A tile's result is the same whichever CTA + computes it.""" + if c["tid"] == 0: + at = c["dq"].value + fx.Int64(c["tag"] * 4) + ptr = fx.inttoptr(fx.PointerType.get(T.i32, fx.AddressSpace.Global, 4), at) + rank = fx.llvm.atomic_add(ptr, fx.Int32(1), syncscope="agent") + fx.ptr_store(fx.Int32(rank), c["tab"] + c["rs"].queue0) + gpu.barrier() + return fx.ptr_load(c["tab"] + c["rs"].queue0) + + +def down_weights(c, tasks, rnd, slots, nu): + """Round ``rnd``'s weights of ``stage_down``: each (task, slot)'s scale block + and weight operands, this wave's slots of U (a slot past U: the last's). + Loads only.""" + s, wave, tab, a, d = c["S"], c["wave"], c["tab"], c["args"], c["d"] + us = [wave + WAVES * (slots * rnd + i) for i in range(slots)] + es = [c["rs"].union_expert(tab, s, fx.min(u, nu - 1)) for u in us] + svs = [ + scale_block_loads(c, a["w2_s"], e * HIDDEN + task * 16, d.down_scale_cols) + for task in tasks + for e in es + ] + avs = [ + fp4_tile_weights( + c, a["w2"], e, HIDDEN, d.inter, task, 0, d.inter // 128, d.inter_real + ) + for task in tasks + for e in es + ] + return svs, avs + + +def mid_operands(c, first, count): + """The MFMA B operands of picks first .. first + min(count, 16) (column + lane % 16 the pick, past them the last): each K step's two 16-byte halves + of the lane's MID row and its group's code, loaded at device scope straight into + registers -- no LDS block a slot, so a round holds more slots.""" + lane, d = c["lane"], c["d"] + j = lane // 16 + row = first + fx.min(lane % 16, count - 1) + steps = d.inter // 128 + # the step's constant after an opaque lane word folds into the immediate + # offset; summed with j * 4 it was a hoisted VGPR a step (spill) + lane_word = fresh(row * d.mid_words + j * 4) + bvs = [] + for st in range_constexpr(steps): + halves = [] + for h in range_constexpr(2): + v = fx.Vector( + bo.buffer_load( + rsrc(c["mid"].value), + lane_word + (st * 32 + 16 * h), + vec_width=4, + dtype=T.i32, + cache_modifier=CM_DEV, + ) + ) + if const_expr(st * 128 + 64 * h + 48 >= d.inter_real): + # columns past the rank's real intermediate (the loader pads + # it with zero weights) are never written: stale bytes, which + # a NaN pattern would carry through the zero weights + past = st * 128 + 64 * h + 16 * j >= d.inter_real + v = fx.Vector.from_elements( + [past.select(fx.Int32(0), v[q]) for q in range(4)], fx.Int32 + ) + halves.append(v) + bvs.append(fp8_operand(halves)) + # this lane's codes, one a K step: contiguous in MIDS (``mid_code_index``), + # 4 a load -- a load a step was 5 of a slot's 16 load instructions, and the + # stage is load-issue bound (probe ``harness/down_probe.py``: -8 %) + base = mid_code_index(d, row, j) + sbs = [] + for w0 in range_constexpr(0, steps, 4): + n = min(4, steps - w0) + v = bo.buffer_load( + rsrc(c["mids"].value), + base + w0, + vec_width=n, + dtype=T.i32, + cache_modifier=CM_DEV, + ) + if const_expr(n == 1): + sbs.append(fx.Int32(v)) + else: + v = fx.Vector(v) + sbs += [fx.Int32(v[k]) for k in range(n)] + for st in range_constexpr(steps): + if const_expr((st + 1) * 128 > d.inter_real): + # a group past the real intermediate has no code written either + past = (4 * st + j) * 32 >= d.inter_real + sbs[st] = past.select(fx.Int32(UNIT_SCALE), sbs[st]) + return bvs, sbs + + +def mid_code_index(d, row, group): + """MIDS word of pick row ``row``'s MXFP8 group ``group`` (32 columns): a K + step's 4 groups apart, so the steps of one group % 4 lie together + (``mid_operands``' lane j reads groups j, j + 4, ..).""" + steps = d.inter // 128 + return row * (d.inter // 32) + group % 4 * steps + group // 4 + + +def down_ug_flags(c, ucs): + """This lane's ug flags of a down round's slots ``ucs``: flag f a (slot, + pick tile past one token tile, part), a tile past the slot's repeating its + last; lane l takes flags l, l + 64, .. (TP2's 18 parts a tile, 48 rows: + 108 flags a round).""" + rs = c["rs"] + ug_parts = rs.ug_parts + per_slot = rs.slot_tiles * ug_parts if rs.multi else ug_parts + n = len(ucs) * per_slot + return [ + _down_ug_flag(c, ucs, fx.min(c["lane"] + 64 * j, n - 1), per_slot) + for j in range(-(-n // 64)) + ] + + +def _down_ug_flag(c, ucs, f, per_slot): + rs, tab = c["rs"], c["tab"] + ug_parts = rs.ug_parts + mine = ucs[0] + for i in range_constexpr(1, len(ucs)): + mine = (f // per_slot == i).select(ucs[i], mine) + if const_expr(not rs.multi): + return mplan.ug_flag(mine, f % ug_parts, ug_parts) + tile0 = rs.tile_first(tab, mine) + tiles = rs.tile_first(tab, mine + 1) - tile0 + return mplan.ug_flag( + mplan.polled_tile(tile0, tiles, f % per_slot // ug_parts), + f % ug_parts, + ug_parts, + ) + + +@traced +def down_shared(c, tasks): + """Each of ``tasks``' rows of the shared expert's down projection + (its MXFP8 intermediate, ``load_smid_tile``) for every + token, bf16 -> LDS ``sdl[task][token][16 rows]`` (``down_combine``): + every task's weights in flight first.""" + s, tid, red, d = c["S"], c["tid"], c["red"], c["d"] + a = c["args"] + shops = [ + gemv_fp8_loads(c, a["sw2"], a["sw2_s"], d.sh_inter, task, row_major=True) + for task in tasks + ] + for t0, nn in token_tiles(s): + load_smid_tile(c, t0, nn, d.sh_inter // 32) + for j in range_constexpr(len(tasks)): + gemv_fp8_mfmas( + c, d.sh_inter, c["shared_xl"], c["shared_xsl"], red, shops[j], rows=nn + ) + gpu.barrier() + if tid < ROWS * nn: + v = row_sum(red, tid % ROWS, tid // ROWS) + at = (j * s + plus(t0, tid // ROWS)) * ROWS + tid % ROWS + fx.ptr_store(bf16_round(v).bitcast(fx.Int32), c["sdl"] + at) + gpu.barrier() + + +@traced +def down_combine(c, task, j, t0, n): + """Down task ``task`` (the CTA's j-th) for tokens t0 .. t0 + n: the routed + contributions summed in bf16 in each token's top-k order, plus the shared expert's + (``down_shared``), the sum pushed to every rank.""" + s, tid, d, tab = c["S"], c["tid"], c["d"], c["tab"] + topk = c["rs"].topk + if tid < ROWS * n: + r = tid % ROWS + tl = tid // ROWS + t = plus(t0, tl) + shared = fx.ptr_load(c["sdl"] + ((j * s + t) * ROWS + r)).bitcast(fx.Float32) + # aiter's stage 2 adds each bf16(v * w) into a bf16 buffer: a rounding + # at every add (its order is the atomics'; here top-k order) + routed = fx.Float32(0.0) + for k in range_constexpr(topk): + w = c["rs"].route_w(tab, s, t, k) + v = contrib_load(c, j, (t * topk + k) * 16 + r) + routed = bf16_round(routed + bf16_round(v * w)) + fx.ptr_store(bf16_round(routed + shared), c["pair"] + tid) + gpu.barrier() + if tid < ROWS * n // 2: + rp = tid % (ROWS // 2) + tl = tid // (ROWS // 2) + t = plus(t0, tl) + v0 = fx.ptr_load(c["pair"] + (tl * ROWS + 2 * rp)) + v1 = fx.ptr_load(c["pair"] + (tl * ROWS + 2 * rp + 1)) + col = task * ROWS + 2 * rp + for p in range_constexpr(d.tp): + dst = preg(c["peer_addr"](p), c["ffn_off"], "ffn") + c["put_bf"](dst, 2 * region_pair(c["rank"], t, s, col), [v0, v1]) + gpu.barrier() + + +@traced +def down_mfma(c, j, task, sl, mid, first, count, av, live): + """Down task ``task`` (the CTA's j-th) on a slot of U: the routed MFMA of its + picks first .. first + min(count, 16) (their B operands ``mid``: + ``mid_operands``; the scale block in ``sl``), each pick's contribution.""" + lane, tab, d = c["lane"], c["tab"], c["d"] + sas = lds_scales( + sl, task % 2 * 16 + lane % 16, lane, 0, d.inter // 128, d.down_scale_cols + ) + jj = lane % 16 + col = fx.min(jj, count - 1) + bvs, sbs = mid + # the K steps on DOWN_CHAINS accumulators in turn, summed in order: one + # chain put every FP8-rate MFMA's latency in series + accs = [fx.Vector.filled(4, 0.0, fx.Float32) for _ in range(DOWN_CHAINS)] + for i in range_constexpr(len(av)): + h = i % DOWN_CHAINS + accs[h] = mfma_scaled(av[i], bvs[i], accs[h], sas[i], sbs[i], FP4, FP8) + acc = accs[0] + for h in range_constexpr(1, DOWN_CHAINS): + acc = acc + accs[h] + pick = c["rs"].pick_at(tab, first + col) + for q in range_constexpr(4): + r = 4 * (lane // 16) + q + if live & (jj < count): + contrib_store(c, j, pick * 16 + r, acc[q]) + + +def contrib_store(c, j, i, v): + """Task j's routed contribution word i (``[pick (t, k)][16 rows]`` f32 bits).""" + fx.ptr_store(v.bitcast(fx.Int32), c["contrib"] + (j * c["rs"].picks * 16 + i)) + + +def contrib_load(c, j, i): + return fx.ptr_load(c["contrib"] + (j * c["rs"].picks * 16 + i)).bitcast(fx.Float32) + + +# ---------------------------------------------------------------- loaders + + +@traced +def load_smid_tile(c, t0, n, groups): + """The shared expert's MXFP8 intermediate of tokens t0 .. t0 + n (SMQ / + SMC, ``smid_quant``) into LDS ``shared_xl`` / ``shared_xsl`` rows 0 .., by every + thread.""" + tid, words = c["tid"], c["d"].sh_inter // 4 + poll_copy( + c, + tid, + THREADS, + [ + (c["smq"], lambda u: plus(t0 * words, u), n * words, c["shared_xl"], words), + ( + c["smc"], + lambda u: plus(t0 * groups, u), + n * groups, + c["shared_xsl"], + groups, + ), + ], + ) + gpu.barrier() + + +# ---------------------------------------------------------------- reduce +@traced +def stage_reduce(c, task): + """Columns 128 task ..: the TP ranks' partials, this rank's first, then + rank + 1, + 2, .., in fp32 -> bf16 -> ``out``.""" + s, tid = c["S"], c["tid"] + for pas in range_constexpr(0, s * RED_COLS // 2, THREADS): + reduce_pair(c, task, pas + tid) + + +@traced +def reduce_pair(c, task, idx): + """Thread ``idx``'s column pair of token idx / 64.""" + s, tp = c["S"], c["d"].tp + a = c["args"] + if idx < s * RED_COLS // 2: + t = idx // (RED_COLS // 2) + col = task * RED_COLS + 2 * (idx % (RED_COLS // 2)) + own = preg(c["sym"], c["ffn_off"], "ffn") + acc0, acc1 = sum_partials( + c["poll"], own, lambda src: region_pair(src, t, s, col), tp + ) + bo.buffer_store(acc0.to(fx.BFloat16), rsrc(a["out"]), t * HIDDEN + col) + bo.buffer_store(acc1.to(fx.BFloat16), rsrc(a["out"]), t * HIDDEN + col + 1) + + +def moe_smem(s, rs, timeline): + """The MoE stages' LDS (``timeline``: with its own stamp record).""" + + @fx.struct + class Smem: + # an ug task's partial C: 2 tiles a wave (``ug_red_slot``) + red: fx.Array[fx.Float32, 2 * WAVES * 64 * 4, 16] + tab: fx.Array[fx.Int32, rs.tab_words, 16] + pair: fx.Array[fx.Float32, ROWS * s, 16] + # one stage's buffers at a time (``stage_lds``) + stage: fx.Array[fx.Int32, sum(stage_lds(s, rs)["words"].values()), 16] + tls: fx.Array[fx.Int64, TL_POINTS if timeline else 1, 16] + + return Smem + + +def stage_lds(s, rs): + """Word offsets into ``Smem.stage`` of the buffers only one stage reads: the + shared stage's MXFP8 x tile; the waves' scale blocks (``wsl``, the ug and + down stages'), then the ug stage's routed MXFP8 x tile, or the down stage's + token tile of the shared expert's MXFP8 intermediate (``load_smid_tile``), + each down task's routed contributions (``contrib_store``) and its shared + expert rows (``down_shared``). A barrier ends every stage.""" + d = rs.dims + # the x tile the shared / ug stages hold (every token's up to one tile) + rows = tile_rows(s) + x8 = rows * lds_row(X8_WORDS) + groups = rows * lds_row(HIDDEN // 32) + wsl = WAVES * wsl_words(d) + shared_xl = rows * lds_row(d.sh_inter // 4) + lds_tail(d.sh_inter) + shared_xsl = rows * lds_row(d.sh_inter // 32) + contrib = 2 * rs.picks * 16 # two down tasks' at most + offsets = { + "x8l": 0, + "x8sl": x8, + "wsl": 0, + "x8rl": wsl, + "x8rsl": wsl + x8, + "shared_xl": wsl, + "shared_xsl": wsl + shared_xl, + "contrib": wsl + shared_xl + shared_xsl, + "sdl": wsl + shared_xl + shared_xsl + contrib, + } + words = { + "shared": x8 + groups, + "ug": wsl + x8 + groups, + "down": wsl + shared_xl + shared_xsl + contrib + 2 * s * ROWS, + } + return {"offsets": offsets, "words": {"all": max(words.values())}} + + +def moe_context(s, rs, lds, mb, bases, layout, scratch, args, rank, sym): + """The MoE stages' context: LDS, mailbox, scratch regions, arguments. Its lane + math starts from a ``fresh`` thread id: in K2 it is not the attention and + seam stages', held from theirs on.""" + tid = fresh(fx.thread_idx.x) + c = { + "S": s, + "tid": tid, + "bid": fx.block_idx.x, + "lane": tid % 64, + "wave": tid // 64, + "red": lds.red.ptr, + "tab": lds.tab.ptr, + "pair": lds.pair.ptr, + "put": mb.put, + "put_bf": mb.put_bf, + "put_words": mb.put_words, + "poll": mb.poll, + "tag": mb.tag, + "peer_addr": lambda p: bases[p], + "rank": rank, + "sym": sym, + "ffn_off": ffn_offset(s, rs.tp), + "rs": rs, + "d": rs.dims, + "args": args, + # x (``ffn_normed``) before this launch: plain loads, no readiness wait + "x_cm": 0, + "x_ready": False, + } + for name, off in stage_lds(s, rs)["offsets"].items(): + c[name] = lds.stage.ptr + off + for region, (off, _) in layout.items(): + c[region] = sreg(scratch, off, region) + return c + + +@traced +def run_moe(c, key, bid, tls, tl0): + """The MoE stages, in order; ``timeline`` stamps from point ``tl0`` + 1.""" + s, timeline, tid = key.tokens, key.timeline, c["tid"] + router_tasks = c["rs"].router_tasks + route0, shared0 = c["rs"].route0, c["rs"].shared0 + shared_tasks = c["d"].shared_tasks + ug_tasks = c["rs"].max_u * c["rs"].ug_parts + for t in range(first_task(bid, XQ0), s, BLOCKS): + await_x(c, [t]) + stage_xq(c, t) + publish(c["put"], c["x8rdy"], t, 1, tid == 0) + stamp(timeline, tls, tid, tl0 + 1) + multi = len(token_tiles(s)) > 1 + for task in range(first_task(bid, c["rs"].router0), router_tasks, BLOCKS): + await_x(c, list(range(s))) + stage_router(c, task) + stamp(timeline, tls, tid, tl0 + 2) + for t in range(first_task(bid, route0), s, BLOCKS): + stage_route(c, t) + stamp(timeline, tls, tid, tl0 + 3) + c["xl"], c["xsl"] = c["x8l"], c["x8sl"] + # a CTA's one shared unit (task, token tile) at most: straight-line code, + # whose tile loads LLVM does not hoist into one live set as it would out of + # a task loop (spills). A task's tiles in series on one CTA put three tiles' + # latency on the path down waits on (the weights read again a tile: L2 hits) + assert c["rs"].shared_units <= BLOCKS + unit = first_task(bid, shared0) + if unit < c["rs"].shared_units: + # the gate weights in flight before x's wait + a = c["args"] + gops = gemv_fp8_loads( + c, a["sgu"], a["sgu_s"], HIDDEN, unit % shared_tasks, row_major=True + ) + if const_expr(not multi): + load_x8_rows(c, "x8", "x8s", c["x8l"], c["x8sl"], lambda j: j, s) + stage_shared(c, unit, gops, 0, s) + else: + t0 = unit // shared_tasks * TILE + n = fx.min(fx.Int32(TILE), s - t0) + load_x8_rows(c, "x8", "x8s", c["x8l"], c["x8sl"], lambda j: t0 + j, n) + stage_shared(c, unit % shared_tasks, gops, t0, n) + stamp(timeline, tls, tid, tl0 + 4) + c["xl"], c["xsl"] = c["x8rl"], c["x8rsl"] + if const_expr(multi): + # past one token tile a task gathers its pick tile's rows itself. The + # first round placed (the shared units' CTAs take its last tasks); the + # rest, past a round, from the queue: a CTA a shared or router task + # delayed takes fewer of them + load_route(c, c["tab"]) + n_tiles = c["rs"].tile_first(c["tab"], c["rs"].n_union(c["tab"], s)) + ug_total = mplan.ug_tasks(n_tiles, c["rs"].ug_parts) + task = first_task(bid, c["rs"].ug0) + while task < ug_total: + task = stage_ug_tile(c, task, ug_total) + else: + # one token tile: every task placed (its x in LDS once a CTA) + loaded = fx.Int32(0) + for task in range(first_task(bid, c["rs"].ug0), ug_tasks, BLOCKS): + if loaded == 0: + # X8R (ready long before the routing) first: after the routing + # lands only its own poll remains + load_x8_rows(c, "x8r", "x8rs", c["x8rl"], c["x8rsl"], lambda j: j, s) + load_route(c, c["tab"]) + loaded = fx.Int32(1) + stage_ug(c, task) + stamp(timeline, tls, tid, tl0 + 5) + assert BLOCKS < DOWN_TASKS <= 2 * BLOCKS + # the route table is in LDS already past one token tile (every CTA loads it + # for the ug queue) or when every CTA ran an ug iteration (more ug task + # slots than CTAs), whose first loads it; else load it + if const_expr(not multi and ug_tasks <= BLOCKS): + load_route(c, c["tab"]) + # every CTA its tile; the first DOWN_TASKS - BLOCKS to get here an extra + # one too: the CTAs whose ug ended first (by bid, those ending last held + # two tiles and ended the stage) + down_task = first_task(bid, DOWN0) + rank = take_down_rank(c) + if rank < DOWN_TASKS - BLOCKS: + stage_down(c, [down_task, BLOCKS + rank]) + else: + stage_down(c, [down_task]) + stamp(timeline, tls, tid, tl0 + 6) + for task in range(first_task(bid, RED0), RED_TASKS, BLOCKS): + stage_reduce(c, task) + stamp(timeline, tls, tid, tl0 + 7) + + +def await_x(c, tokens): + """In a launch that also writes x (K2: the FFN seam's norm): tokens' rows + ready.""" + if const_expr(c["x_ready"]): + c["poll"]([(c["xrdy"], t, 1) for t in tokens]) + + +def moe_args(*values) -> dict: + """The MoE stages' arguments by name, in ``ARGS`` order.""" + assert len(values) == len(ARGS) + return dict(zip(ARGS, values)) diff --git a/vllm/models/deepseek_v41/amd/mono/stages/moe_plan.py b/vllm/models/deepseek_v41/amd/mono/stages/moe_plan.py new file mode 100644 index 0000000000..5b5aa6fcdd --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono/stages/moe_plan.py @@ -0,0 +1,54 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at 97359d6, MIT License, +# Copyright (C) 2026, Advanced Micro Devices, Inc.: +# atom/models/deepseek_v41/mono/moe_plan.py +"""Every per-step decision of which MoE ug flags (UGF) and intermediate rows +(MID) the ug stage writes and the down stage reads, defined once. + +The route table (``stage_route``) gives each expert with picks a slot of U, its +picks a run from ``slot_first``, and, past one token tile, its picks tiles of +``TILE`` from ``tile_first`` (each tile's slot ``tile_slot``). An ug task covers +a unit -- a slot, or past one token tile a pick tile -- and one of ``ug_parts`` parts +of the intermediate; its flag is ``ug_flag``. A down round polls each of its +slots' flags (``polled_tile`` its tiles) before it reads their MID rows. Called +by the kernel on traced Int32 and by the CPU property test on Python ints. +""" + +from vllm.models.deepseek_v41.amd.mono.common.plan import cdiv, imin + +TILE = 16 # picks a pick tile: an MFMA's N + + +def pick_tiles(count): + """A slot's pick tiles: its ``count`` picks, ``TILE`` a tile.""" + return cdiv(count, TILE) + + +def ug_parts(groups, ug_groups): + """A unit's ug parts: the expert's ``groups`` FP4 groups of its real width, + ``ug_groups`` a part. A part wholly in the padding is not run: MID's padded + columns stay the step's zeros, against zero weights in down.""" + return cdiv(groups, ug_groups) + + +def ug_flag(unit, part, ug_parts): + """The UGF pair of ug task (``unit``, ``part``), ``ug_parts`` parts a unit.""" + return unit * ug_parts + part + + +def ug_tasks(units, ug_parts): + """A step's ug tasks: every part of each of its ``units``.""" + return units * ug_parts + + +def tile_pick0(slot_first, tile, tile_first): + """Pick tile ``tile``'s first pick, its slot's picks from ``slot_first`` and + tiles from ``tile_first``.""" + return slot_first + (tile - tile_first) * TILE + + +def polled_tile(tile_first, tiles, j): + """The pick tile a down round polls as a slot's ``j``-th, the slot's + ``tiles`` tiles from ``tile_first``: the last past them.""" + return tile_first + imin(j, tiles - 1) diff --git a/vllm/models/deepseek_v41/amd/mono/stages/moe_shape.py b/vllm/models/deepseek_v41/amd/mono/stages/moe_shape.py new file mode 100644 index 0000000000..a07e88cfb0 --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono/stages/moe_shape.py @@ -0,0 +1,229 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at 97359d6, MIT License, +# Copyright (C) 2026, Advanced Micro Devices, Inc.: +# atom/models/deepseek_v41/mono/kernels/moe_shape.py +"""The MoE stages' build key and route shape: what they are built for (``MoeBuild``), +and the task counts, placements and LDS route table that follow from its +routing (``RouteShape``), with the table's readers. Plain Python but for the +readers; the kernel's stages are in ``moe``.""" + +from dataclasses import dataclass, field + +import flydsl.expr as fx + +from vllm.models.deepseek_v41.amd.mono.common.plan import BLOCKS, WAVES +from vllm.models.deepseek_v41.amd.mono.stages import moe_plan as mplan +from vllm.models.deepseek_v41.amd.mono.stages.dims import UG_PART, Dims +from vllm.models.deepseek_v41.amd.mono.stages.gemv import ROWS, TILE, token_tiles +from vllm.models.deepseek_v41.amd.mono.stages.router import PARTS as ROUTER_PARTS + +EXPERTS, TOPK = 384, 6 # the target's routing (a build's: ``MoeBuild``) +ROUTER0 = 8 # the router's first task: past the xq tasks (a token each) +MASK_WORDS = 2 # an expert's token bitmap in the route table: tokens <= 64 +# a route lane's experts (EXPERTS / 64) sorted descending, ties kept in index +# order: the network by their count +SORT_NETS = { + 6: [ + (0, 1), + (2, 3), + (4, 5), + (0, 2), + (1, 4), + (3, 5), + (0, 1), + (2, 3), + (4, 5), + (1, 2), + (3, 4), + (2, 3), + ], + 2: [(0, 1)], +} + + +@dataclass(frozen=True) +class MoeBuild: + tokens: int + # routed experts and a token's top-k: the target's, or the DSpark draft's + experts: int = EXPERTS + topk: int = TOPK + # the TP size: every width a rank holds (``Dims``) + tp: int = 4 + timeline: bool = field(default=False, metadata={"sym": "tl"}) + # FP4 groups an ug task (``RouteShape.ug_groups``) + ug_groups: int = 2 + + +@dataclass(frozen=True) +class RouteShape: + """A build's routing: ``experts`` routed experts, ``topk`` a token of + ``tokens``, the task counts and placements that follow, and the route table + ``load_route`` lays out in LDS: every token's (ids, weights), |U|, the + experts of U, each (t, k)'s slot in U, each expert's slot, the picks in slot + order and each slot's first position in it (past one token tile also each + slot's first pick tile and each pick tile's slot).""" + + experts: int + topk: int + tokens: int + tp: int + # FP4 groups (32 intermediate columns) an ug task, its waves splitting K in + # halves (2) or quarters (1). 1 is twice the tasks, so the CTAs' ug ends + # closer together: faster at the draft's widths (40 rows -4 us), slower at + # the target's (48 rows +3 us; the runners pick) + ug_groups: int = 2 + + @property + def dims(self) -> Dims: + return Dims(self.tp) + + @property + def ug_parts(self) -> int: + """An expert's ug tasks (``moe_plan.ug_parts``).""" + return mplan.ug_parts(self.dims.inter_real // UG_PART, self.ug_groups) + + @property + def ug_k_slices(self) -> int: + """The K slices of an ug task's waves: a wave a (group, 16-column half, + slice).""" + return WAVES // (2 * self.ug_groups) + + @property + def picks(self) -> int: + return self.tokens * self.topk + + @property + def max_u(self) -> int: + return min(self.picks, self.experts) + + @property + def router_groups(self) -> tuple[int, int]: + """``stage_router``'s token groups and their tokens: a group a token + tile, as many as one round of CTAs holds (a CTA running two serialized + them: the route waited on the 32 that did at 48 rows).""" + per_group = self.experts // ROWS * ROUTER_PARTS + groups = max(1, min(len(token_tiles(self.tokens)), BLOCKS // per_group)) + tokens = -(-self.tokens // groups) + assert (groups - 1) * tokens < self.tokens + return groups, tokens + + @property + def router_tasks(self) -> int: + # (token group, 16-row tile, K part): ``stage_router`` + return self.router_groups[0] * self.experts // ROWS * ROUTER_PARTS + + @property + def router0(self) -> int: + """Past the xq tasks (a token each from CTA 0).""" + return max(ROUTER0, self.tokens) + + @property + def route0(self) -> int: + return self.router0 + self.router_tasks + + @property + def shared0(self) -> int: + return self.route0 + max(8, self.tokens) + + @property + def shared_units(self) -> int: + """The shared expert's (task, token tile) units (``stage_shared``).""" + return self.dims.shared_tasks * len(token_tiles(self.tokens)) + + @property + def ug0(self) -> int: + """The first ug round's placement: the shared units' CTAs last.""" + return self.shared0 + self.shared_units + + @property + def sort_net(self): + return SORT_NETS[self.experts // 64] + + @property + def tab_words(self) -> int: + return self.tile_slot0 + self.max_tiles if self.multi else self.queue0 + 1 + + @property + def multi(self) -> bool: + """More than one token tile: an ug task is a (tile of a slot's picks, + part) (``stage_ug``).""" + return self.tokens > TILE + + @property + def slot_tiles(self) -> int: + """A slot's pick tiles at most (a token picks an expert once).""" + return len(token_tiles(self.tokens)) + + @property + def max_tiles(self) -> int: + """The pick tiles of U at most: a slot's partial tile each, then the + whole tiles, and never more than a tile a pick.""" + return min(self.picks, self.max_u + self.picks // TILE) + + @property + def tile_first0(self) -> int: + """Past one token tile: each slot's first pick tile, then their count.""" + return self.queue0 + 1 + + @property + def tile_slot0(self) -> int: + """Past one token tile: each pick tile's slot of U.""" + return self.tile_first0 + self.max_u + 1 + + @property + def queue0(self) -> int: + """The ug task this CTA took from the queue (``take_ug_task``).""" + return self.mask0 + MASK_WORDS * self.experts + + @property + def expert_slot0(self) -> int: + """The route table's per-expert slots (a slot where the expert is in U).""" + return 3 * self.picks + 1 + self.max_u + + @property + def order0(self) -> int: + """The picks (t topk + k) in slot order, a slot's ascending.""" + return self.expert_slot0 + self.experts + + @property + def first0(self) -> int: + """Each slot's first position in the pick order, then the pick count.""" + return self.order0 + self.picks + + @property + def mask0(self) -> int: + """Each expert's tokens that picked it, a bit a token (``MASK_WORDS``).""" + return self.first0 + self.max_u + 1 + + @property + def tile_picks(self) -> int: + """A slot's picks at most: a token picks an expert once.""" + return min(self.tokens, 16) + + def route_w(self, tab, s, t, k): + return fx.ptr_load(tab + (t * 2 * self.topk + self.topk + k)).bitcast( + fx.Float32 + ) + + def n_union(self, tab, s): + return fx.ptr_load(tab + 2 * s * self.topk) + + def union_expert(self, tab, s, u): + return fx.ptr_load(tab + (2 * s * self.topk + 1 + u)) + + def slot_first(self, tab, u): + return fx.ptr_load(tab + (self.first0 + u)) + + def pick_at(self, tab, i): + return fx.ptr_load(tab + (self.order0 + i)) + + def tile_first(self, tab, u): + return fx.ptr_load(tab + (self.tile_first0 + u)) + + def tile_slot(self, tab, i): + return fx.ptr_load(tab + (self.tile_slot0 + i)) + + +def route_shape(key: MoeBuild) -> RouteShape: + return RouteShape(key.experts, key.topk, key.tokens, key.tp, key.ug_groups) diff --git a/vllm/models/deepseek_v41/amd/mono/stages/router.py b/vllm/models/deepseek_v41/amd/mono/stages/router.py new file mode 100644 index 0000000000..a7d6f411a8 --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono/stages/router.py @@ -0,0 +1,86 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at 97359d6, MIT License, +# Copyright (C) 2026, Advanced Micro Devices, Inc.: +# atom/model_ops/deepseek_v41/router.py +"""V4.1's router GEMV tile for decode: bf16 x [T, 5120] times the gate weight +[384, 5120] -> f32 logits, in one fixed summation order. + +Rows 16 a tile, K in ``PARTS`` contiguous parts, 16x16x32 bf16 MFMA: in part +p wave w takes the part's K steps w, w + 8, ... (32 columns each) in that +order, and the eight wave partials are summed in wave order +(``common.ops.row_sum``); the MoE's route stage sums the parts in part +order. The MoE stages spread a tile's parts over PARTS CTAs, and past 16 tokens +a token tile a CTA too, which the order allows: a CTA alone cannot read a +tile, or every token's x, fast enough. +""" + +import flydsl.expr as fx +from aiter.ops.flydsl.kernels import buffer_ops as bo +from flydsl.expr import const_expr, range_constexpr +from flydsl.expr.typing import T + +from vllm.models.deepseek_v41.amd.mono.common.ops import mfma_bf16, rsrc +from vllm.models.deepseek_v41.amd.mono.common.plan import WAVES + +HIDDEN = 5120 +ROWS = 16 # a tile's rows: the MFMA's M +PARTS = 4 # contiguous K parts of a tile, summed in order + + +def _k0s(part, lane, wave): + """The K offsets of this lane's loads in part ``part`` (8 bf16 each).""" + per_wave = HIDDEN // 32 // WAVES // PARTS + assert per_wave * 32 * WAVES * PARTS == HIDDEN + return [ + (part * per_wave * WAVES + wave + WAVES * i) * 32 + 8 * (lane // 16) + for i in range(per_wave) + ] + + +def router_weight_loads(w, task, part, lane, wave, w_cm=0): + """``router_tile``'s weight operands (rows 16 task .., K part ``part``). + ``w_cm``: their load cache policy (nontemporal where nothing else of the + step reuses them).""" + row = task * ROWS + lane % 16 + return [ + fx.Vector( + bo.buffer_load( + rsrc(w), + (row * HIDDEN + k0) // 2, + vec_width=4, + dtype=T.i32, + cache_modifier=w_cm, + ) + ) + for k0 in _k0s(part, lane, wave) + ] + + +def router_x_loads(x, part, s, lane, wave, x_cm=0, t0=0): + """``router_tile``'s x operands: the s tokens from ``t0`` (column + lane % 16; ``t0`` traced or an int). ``x_cm``: their load cache policy + (device scope when another CTA of the same launch wrote x).""" + tok = fx.min(lane % 16, s - 1) + if const_expr(not isinstance(t0, int) or t0 != 0): + tok = tok + t0 + return [ + fx.Vector( + bo.buffer_load( + rsrc(x), + (tok * HIDDEN + k0) // 2, + vec_width=4, + dtype=T.i32, + cache_modifier=x_cm, + ) + ) + for k0 in _k0s(part, lane, wave) + ] + + +def router_tile_mfma(wvs, xvs, lane, wave, red): + """``router_tile``'s MFMAs, K steps in order -> ``red[wave][lane]``.""" + acc = fx.Vector.filled(4, 0.0, fx.Float32) + for i in range_constexpr(len(wvs)): + acc = mfma_bf16(wvs[i].bitcast(fx.BFloat16), xvs[i].bitcast(fx.BFloat16), acc) + fx.ptr_store(acc, red + (wave * 64 + lane) * 4) diff --git a/vllm/models/deepseek_v41/amd/mono/stages/seam.py b/vllm/models/deepseek_v41/amd/mono/stages/seam.py new file mode 100644 index 0000000000..82abee9f46 --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono/stages/seam.py @@ -0,0 +1,406 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at 97359d6, MIT License, +# Copyright (C) 2026, Advanced Micro Devices, Inc.: +# atom/models/deepseek_v41/mono/kernels/attn_pre.py (slice, gate, norm) and +# atom/models/deepseek_v41/mono/kernels/attn_post.py (the attention's TP reduce). +"""The mono layer's mHC seam stages, each a task table of a persistent launch +handing off through the tagged mailbox: + + slice (160, 32 hidden columns x 4 streams each): the owed post folded into + the residual (R_new, bf16 out), the collapsed layer input (mailbox + LIN) and the partial mix projection / sum of squares (mailbox PMIX) + gate (S): the 160 partials -> post / comb (Sinkhorn) and the next pre + norm (S): the RMSNorm -> the bf16 row, handed to the caller's store + reduce (160): the FFN seam's attention output, the TP ranks' partials + (pushed to every rank's ATTN region) summed in the all-reduce's + order -> ``pend_lds``, where its slice reads it + push (160): a rank's partial given in memory (an attention vLLM ran) + to every rank's ATTN region, ahead of the slice's reduce + +Every elementwise step follows aiter's ``mhc_fused_post_pre_delayed_rmsnorm`` +rounding; the mix projection and the sums of squares, split-K reductions, +differ only in their split and order. +""" + +import math + +import flydsl.expr as fx +from aiter.ops.flydsl.kernels import buffer_ops as bo +from flydsl.expr import const_expr, gpu, range_constexpr, rocdl +from flydsl.expr import math as fmath +from flydsl.expr.typing import T + +from vllm.models.deepseek_v41.amd.mono.common.ops import ( + batched_rounds, + bf16_round, + bf_hi, + bf_lo, + butterfly, + hw_exp2, + hw_rcp, + hw_rsq, + lane_gather, + rsrc, + sum_partials, + traced, +) +from vllm.models.deepseek_v41.amd.mono.common.plan import THREADS, WAVES, pair_layout +from vllm.models.deepseek_v41.amd.mono.common.sync import preg +from vllm.models.deepseek_v41.amd.mono.stages.gemv import ( + ld_bf, + ld_f32, + plus, + token_tiles, +) + +HC = 4 +HIDDEN = 5120 +K_MIX = HC * HIDDEN +MIX = HC * (HC + 2) # 24 projections: pre 4 | post 4 | comb 16 +EPS = 1e-20 +HC_EPS = 1e-6 +SINKHORN = 20 +POST_MULT = 2.0 +LOG2E = 1.0 / math.log(2.0) + +COLS = 32 # hidden columns a slice task +SLICES = HIDDEN // COLS # 160 +KT = HC * COLS # projection columns a slice task +PART = MIX + 1 # a slice's partials a token: the mixes, then the sum of squares +GATE_BLOCK = 8 # slices a gate thread sums + + +def scratch_layout(tokens: int) -> dict[str, tuple[int, int]]: + """Mailbox region -> (byte offset, bytes) of a seam; 8 B a pair.""" + pairs = { + "lin": tokens * HIDDEN // 2, # bf16 pairs + "pmix": SLICES * tokens * PART, + } + return pair_layout(pairs.items()) + + +def attn_peer_bytes(tokens: int, tp: int) -> int: + """ATTN: ``[source rank][token][hidden / 2]`` bf16 pairs.""" + return tp * tokens * HIDDEN // 2 * 8 + + +def attn_region_pair(src, t, s, col): + return (src * s + t) * (HIDDEN // 2) + col // 2 + + +def fma(a, b, c): + return fx.Float32(fmath.fma(a, b, c)) + + +def sigmoid(z): + """Triton's ``tl.sigmoid``: 1 / (1 + exp2(-z log2(e))).""" + return fx.Float32(1.0) / (fx.Float32(1.0) + hw_exp2(z * fx.Float32(-LOG2E))) + + +def lane_f32(v, src): + return lane_gather(v.bitcast(fx.Int32), src).bitcast(fx.Float32) + + +def pend_value(c, t, col, cc): + """The owed post's sublayer output at (token t, column col): read from + ``pend``, or where a caller staged it (``pend_lds``, ``[token][32]`` of this + task's columns: a reduction it did first).""" + if const_expr(c.get("pend_lds") is None): + return ld_bf(c["args"]["pend"], t * HIDDEN + col) + return fx.ptr_load(c["pend_lds"] + (t * COLS + cc)) + + +# ---------------------------------------------------------------- slice +@traced +def slice_input(c, task, idx): + """Thread ``idx``'s pair of the layer input (token idx / 16, columns + 32 task + 2 (idx % 16) ..).""" + s, rl = c["S"], c["rl"] + a = c["args"] + c0 = task * COLS + if idx < s * COLS // 2: + t = idx // (COLS // 2) + cc = 2 * (idx % (COLS // 2)) + outs = [] + for d in range_constexpr(2): + acc = fx.Float32(0.0) + for h in range_constexpr(HC): + rb = fx.ptr_load(rl + (t * KT + h * COLS + cc + d)) + acc = acc + ld_f32(a["pre_in"], t * HC + h) * rb + outs.append(acc) + c["put_bf"](c["lin"], t * HIDDEN + c0 + cc, outs) + + +def slice_at(c0, e): + """Element e of a slice task at column c0: (token, stream, hidden column).""" + return e // KT, (e % KT) // COLS, c0 + e % COLS + + +def slice_loads(c, c0, e): + """Element e's loads of ``stage_slice`` (e clamped into the step's).""" + a = c["args"] + t, o, col = slice_at(c0, e) + r = [ld_bf(a["res_in"], (t * HC + h) * HIDDEN + col) for h in range(HC)] + cm = [ld_f32(a["comb_in"], (t * HC + h) * HC + o) for h in range(HC)] + return r, cm, ld_f32(a["post_in"], t * HC + o) + + +def slice_row(c, c0, e, got): + """Element e of ``stage_slice``'s R_new from ``slice_loads``' values.""" + a = c["args"] + t, o, col = slice_at(c0, e) + r, cm, post = got + # aiter's mhc_fused_post_pre_delayed_rmsnorm: ((post x + c0 r0) + # + c1 r1) + c2 r2, each product rounded, then fma(c3, r3, ..) + v = post * pend_value(c, t, col, e % COLS) + (r[0] * cm[0]) + for h in range_constexpr(1, HC - 1): + v = v + cm[h] * r[h] + # bf16 R_new is what the collapse, the sums of squares and the + # mix projection all read + v = bf16_round(fma(cm[HC - 1], r[HC - 1], v)) + bo.buffer_store(v.to(fx.BFloat16), rsrc(a["res_out"]), (t * HC + o) * HIDDEN + col) + fx.ptr_store(v, c["rl"] + (t * KT + o * COLS + e % COLS)) + + +@traced +def stage_slice(c, task): + """R_new, the layer input and the partial projections of hidden columns + 32 task .. +32, all four streams, every token.""" + s, tid, lane, wave = c["S"], c["tid"], c["lane"], c["wave"] + rl, fl = c["rl"], c["fl"] + a = c["args"] + c0 = task * COLS + batched_rounds( + tid, + s * KT, + lambda e: slice_loads(c, c0, e), + lambda e, got: slice_row(c, c0, e, got), + ) + for i in range_constexpr((MIX * KT + THREADS - 1) // THREADS): + e = tid + THREADS * i + if e < MIX * KT: + j = e // KT + o = (e % KT) // COLS + # the original's two bf16 MFMAs take fn as hi = bf16(fn) and lo = + # bf16(fn - hi): their sum, exact in fp32, is the weight it applies + f = ld_f32(a["hc_fn"], j * K_MIX + o * HIDDEN + c0 + e % COLS) + hi = bf16_round(f) + fx.ptr_store(hi + bf16_round(f - hi), fl + e) + gpu.barrier() + # the layer input: bf16 R_new of the four streams by the incoming pre-mix, + # products rounded, summed in order (the Triton collapse, no contraction); + # two columns a thread, one bf16 pair a mailbox store + for pas in range_constexpr(0, s * COLS // 2, THREADS): + slice_input(c, task, plus(pas, tid)) + # sum of squares: token t on wave t % WAVES, two columns a lane + for t in range_constexpr(s): + if wave == t % WAVES: + x0 = fx.ptr_load(rl + (t * KT + 2 * lane)) + x1 = fx.ptr_load(rl + (t * KT + 2 * lane + 1)) + sq = butterfly(x0 * x0 + x1 * x1, (32, 16, 8, 4, 2, 1)) + if lane == 0: + c["put"](c["pmix"], (task * s + t) * PART + MIX, sq) + # mixes: fp32 MFMA 16x16x4, A = fn rows (two 16-row tiles), B = R_new^T, + # a token tile at a time + for t0, n in token_tiles(s): + slice_mix(c, task, t0, n) + + +@traced +def slice_mix(c, task, t0, n): + """The mix projection partials of tokens t0 .. t0 + n (``stage_slice``).""" + s, tid, lane, wave = c["S"], c["tid"], c["lane"], c["wave"] + rl, fl, red = c["rl"], c["fl"], c["red"] + tok = plus(t0, fx.min(lane % 16, n - 1)) + kw = KT // WAVES + cs = [fx.Vector.filled(4, 0.0, fx.Float32) for _ in range(2)] + for kk in range_constexpr(kw // 4): + colk = wave * kw + lane // 16 + 4 * kk + b = fx.ptr_load(rl + (tok * KT + colk)) + for h in range_constexpr(2): + row = h * 16 + lane % 16 + av = fx.ptr_load(fl + (fx.min(row, MIX - 1) * KT + colk)) + av = (row < MIX).select(av, fx.Float32(0.0)) + ops = rocdl._split_mfma_operands([av, b, cs[h], 0, 0, 0]) + cs[h] = fx.Vector(rocdl.mfma_f32_16x16x4f32(T.vec(4, T.f32), *ops).result) + for h in range_constexpr(2): + fx.ptr_store(cs[h], red + ((wave * 2 + h) * 64 + lane) * 4) + gpu.barrier() + if tid < MIX * n: + tl = tid // MIX + t = plus(t0, tl) + j = tid % MIX + h = j // 16 + ln = 16 * ((j % 16) // 4) + tl + tot = fx.Float32(0.0) + for w in range_constexpr(WAVES): + tot = tot + fx.ptr_load(red + (((w * 2 + h) * 64 + ln) * 4 + j % 4)) + c["put"](c["pmix"], (task * s + t) * PART + j, tot) + gpu.barrier() + + +# ---------------------------------------------------------------- gate +def _row(v, lane, q): + """Comb lanes 8..23 hold ``comb[r][c]`` at lane 8 + 4 r + c: element q of + the row.""" + return lane_f32(v, (lane - 8) // 4 * 4 + 8 + q) + + +def _col(v, lane, q): + return lane_f32(v, 8 + 4 * q + (lane - 8) % 4) + + +@traced +def stage_gate(c, t): + """Token t's 160 slices, in slice order, -> post / comb / next pre.""" + s, tid, lane, wave = c["S"], c["tid"], c["lane"], c["wave"] + red = c["red"] + a = c["args"] + groups = SLICES // GATE_BLOCK # 20 + if tid < groups * PART: + g = tid // PART + j = tid % PART + vals = c["poll"]( + [ + (c["pmix"], ((g * GATE_BLOCK + k) * s + t) * PART + j, 1) + for k in range(GATE_BLOCK) + ] + ) + acc = vals[0][0].bitcast(fx.Float32) + for k in range_constexpr(1, GATE_BLOCK): + acc = acc + vals[k][0].bitcast(fx.Float32) + fx.ptr_store(acc, red + (g * PART + j)) + gpu.barrier() + if wave == 0: + j = fx.min(lane, MIX - 1) + mix = fx.ptr_load(red + j) + sq = fx.ptr_load(red + MIX) + for g in range_constexpr(1, groups): + mix = mix + fx.ptr_load(red + (g * PART + j)) + sq = sq + fx.ptr_load(red + (g * PART + MIX)) + sc = [ld_f32(a["hc_scale"], i) for i in range(3)] + base = ld_f32(a["hc_base"], j) + # aiter's mhc_fused_post_pre_delayed_rmsnorm reduce: one rstd = + # rsq(fma(sum, 1/K, eps)), then each gate fma(mix rstd, scale, base) + v = mix * hw_rsq(fma(sq, fx.Float32(1.0 / K_MIX), fx.Float32(EPS))) + if lane < HC: + pre = sigmoid(fma(v, sc[0], base)) + bo.buffer_store(pre + fx.Float32(HC_EPS), rsrc(a["pre_out"]), t * HC + lane) + if (lane >= HC) & (lane < 2 * HC): + post = sigmoid(fma(v, sc[1], base)) + bo.buffer_store( + post * fx.Float32(POST_MULT), rsrc(a["post_out"]), t * HC + lane - HC + ) + cv = fma(v, sc[2], base) + row = [_row(cv, lane, q) for q in range(4)] + # Triton's exp: exp2 of (x - max) log2(e) + m = fx.max(fx.max(row[0], row[1]), fx.max(row[2], row[3])) + cv = hw_exp2((cv - m) * fx.Float32(LOG2E)) + # fast_dividef is x rcp(y); its first "+ eps" contracts into the fma + row = [_row(cv, lane, q) for q in range(4)] + cv = fma(cv, hw_rcp((row[0] + row[1]) + (row[2] + row[3])), fx.Float32(HC_EPS)) + col = [_col(cv, lane, q) for q in range(4)] + cv = cv * hw_rcp(((col[0] + col[1]) + (col[2] + col[3])) + fx.Float32(HC_EPS)) + for _ in range_constexpr(SINKHORN - 1): + row = [_row(cv, lane, q) for q in range(4)] + cv = cv * hw_rcp( + ((row[0] + row[1]) + (row[2] + row[3])) + fx.Float32(HC_EPS) + ) + col = [_col(cv, lane, q) for q in range(4)] + cv = cv * hw_rcp( + ((col[0] + col[1]) + (col[2] + col[3])) + fx.Float32(HC_EPS) + ) + if (lane >= 2 * HC) & (lane < MIX): + bo.buffer_store(cv, rsrc(a["comb_out"]), t * HC * HC + lane - 2 * HC) + gpu.barrier() + + +# ---------------------------------------------------------------- norm +@traced +def stage_norm(c, t, store): + """Token t's RMSNorm (rstd = rsq(fma(sum, 1/H, eps)), bf16((x rstd) w)); + 256 threads, 3 chunks of 8 at 8 t + 2048 c, each chunk's bf16 values handed + to ``store(t, col, ys, live)`` by every thread (``live``: a thread's chunk + in the row; threads 256.. mirror thread 255).""" + tid, lane, wave, red = c["tid"], c["lane"], c["wave"], c["red"] + a = c["args"] + tt = fx.min(tid, 255) + mine = tid < 256 + cols = [tt * 8 + 2048 * ch for ch in range(3)] + xs = [] + for ch in range_constexpr(3): + base = t * HIDDEN + fx.min(cols[ch], HIDDEN - 8) + words = c["poll"]([(c["lin"], base // 2 + k, 1) for k in range(4)]) + v = [] + for k in range_constexpr(4): + v += [bf_lo(words[k][0]), bf_hi(words[k][0])] + xs.append([(cols[ch] < HIDDEN).select(x, fx.Float32(0.0)) for x in v]) + acc = fx.Float32(0.0) + for ch in range_constexpr(3): + for x in xs[ch]: + acc = acc + x * x + acc = butterfly(acc, (1, 2, 4, 8, 16, 32)) + if (lane == 0) & mine: + fx.ptr_store(acc, red + wave) + gpu.barrier() + tot = butterfly(fx.ptr_load(red + lane % 4), (1, 2)) + r = hw_rsq(fma(tot, fx.Float32(1.0 / HIDDEN), fx.Float32(EPS))) + for ch in range_constexpr(3): + col = fx.min(cols[ch], HIDDEN - 8) + ys = [ + bf16_round((xs[ch][j] * r) * ld_bf(a["norm_w"], col + j)) for j in range(8) + ] + store(t, col, ys, mine & (cols[ch] < HIDDEN)) + gpu.barrier() + + +# ---------------------------------------------------------------- reduce +@traced +def stage_reduce(c, task): + """The attention output at this slice's 32 columns: the TP ranks' partials + in the all-reduce's order (``sum_partials``), in fp32 -> bf16 -> + ``pend_lds``.""" + s, tid = c["S"], c["tid"] + for pas in range_constexpr(0, s * COLS // 2, THREADS): + reduce_attn_pair(c, task, plus(pas, tid)) + gpu.barrier() + + +@traced +def reduce_attn_pair(c, task, idx): + """Thread ``idx``'s column pair of token idx / 16 (``stage_reduce``).""" + s = c["S"] + c0 = task * COLS + if idx < s * COLS // 2: + t = idx // (COLS // 2) + col = c0 + 2 * (idx % (COLS // 2)) + own = preg(c["sym"], 0, "attn") + acc0, acc1 = sum_partials( + c["poll"], + own, + lambda src: attn_region_pair(src, t, s, col), + c["d"].tp, + ) + lds = c["pend_lds"] + (t * COLS + col - c0) + fx.ptr_store(bf16_round(acc0), lds) + fx.ptr_store(bf16_round(acc1), lds + 1) + + +@traced +def stage_push(c, task, part): + """This rank's attention partial (``part``: bf16 [S, HIDDEN], wo_b's + unreduced output) at this slice's 32 columns, every token, to every rank's + ATTN region: the pairs ``stage_reduce`` polls, as the bits vLLM wrote.""" + s, tid, tp = c["S"], c["tid"], c["d"].tp + for pas in range_constexpr(0, s * COLS // 2, THREADS): + idx = plus(pas, tid) + if idx < s * COLS // 2: + t = idx // (COLS // 2) + col = task * COLS + 2 * (idx % (COLS // 2)) + w = bo.buffer_load( + rsrc(part), (t * HIDDEN + col) // 2, vec_width=1, dtype=T.i32 + ) + pair = attn_region_pair(c["rank"], t, s, col) + for p in range_constexpr(tp): + c["put_words"](preg(c["peer_addr"](p), 0, "attn"), pair, [w]) diff --git a/vllm/models/deepseek_v41/amd/mono_decode.py b/vllm/models/deepseek_v41/amd/mono_decode.py new file mode 100644 index 0000000000..b404a42de3 --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono_decode.py @@ -0,0 +1,376 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""DeepSeek-V4.1 mono decode layer on ROCm CDNA4 (``VLLM_ROCM_MONO_DECODE=1``). + +A decode step's backbone layer -- the attention seam, the attention, its TP +all-reduce, the FFN seam, the MoE and its all-reduce -- runs as two persistent +FlyDSL launches (``mono``, this directory) on the layer's loaded +weights, the TP reductions inside the kernels. A layer runs whole when it has +the standard seam and no compressor or indexer of its own (compress ratio 1 or +2) over ``fp8_ds_mla`` records. Any other backbone layer (the first, the Engram +layers, the compressor / indexer layers) keeps vLLM's attention seam and +attention and runs the rest as one launch, the FFN launch, from wo_b's +unreduced output. A step takes the kernels when it is decode only, at most +``MAX_ROWS`` rows (eager or a FULL CUDA graph: the model is not torch.compiled); +a whole layer also needs causal SWA windows. Anything else runs the layer as +before; the paths share every tensor at the layer boundary, so they interleave +freely. + +Every TP rank must take the same path at every layer (the kernels wait on each +other's pushes): the conditions read only metadata every rank holds alike. +""" + +from typing import TYPE_CHECKING, cast + +import torch + +import vllm.envs as envs +from vllm.config import CUDAGraphMode +from vllm.distributed import ( + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, + get_tp_group, +) +from vllm.forward_context import get_forward_context, is_forward_context_available +from vllm.logger import init_logger +from vllm.platforms import current_platform + +if TYPE_CHECKING: + from vllm.models.deepseek_v41.amd.model import DeepseekV4DecoderLayer + from vllm.models.deepseek_v41.amd.mono.runner import MonoLayerWeights + from vllm.models.deepseek_v41.amd.rocm import DeepseekV4ROCMAiterSparseSWAMetadata + from vllm.models.deepseek_v41.sparse_mla import DeepseekV4FlashMLAMetadata + +logger = init_logger(__name__) + +RECORD = 584 # an fp8_ds_mla KV record: 576 data bytes, 8 scale bytes +SWA_WIDTH = 128 # a causal decode token's window slots +MAX_ROWS = 48 # the kernels' step rows at most: 8 requests x (1 + 5 DSpark drafts) +# the kernels' ISA: CDNA4 (gfx950) scaled MFMA and MX formats +CDNA_VERSIONS = (4,) + +_runner = None + + +def _mono_runner(device: torch.device): + """The process's mono runner (one a TP rank, shared by every layer). Its + peer-memory handshake is collective over the TP group: every rank builds it + at the same layer of the same step.""" + global _runner + if _runner is None: + if torch.cuda.is_current_stream_capturing(): + # its peer memory is allocated and exchanged eagerly; vLLM warms + # every graph up eagerly before capturing it + raise RuntimeError( + "VLLM_ROCM_MONO_DECODE was first reached inside a CUDA graph " + "capture, before an eager decode step." + ) + from vllm.models.deepseek_v41.amd.mono.runner import MAX_TOKENS, DSV41MonoLayer + + assert MAX_TOKENS >= MAX_ROWS, MAX_TOKENS + _runner = DSV41MonoLayer( + get_tensor_model_parallel_world_size(), + get_tensor_model_parallel_rank(), + get_tp_group().cpu_group, + device, + ) + return _runner + + +def _record_view(kv: torch.Tensor, block: int) -> torch.Tensor: + """A layer's fp8_ds_mla cache as [blocks, block, RECORD] bytes (a block's + data rows, then its scale words: the kernels index both from it).""" + kv = kv if kv.dtype == torch.uint8 else kv.view(torch.uint8) + assert kv.stride(0) >= block * RECORD, (tuple(kv.shape), kv.stride(), block) + return torch.as_strided(kv, (kv.shape[0], block, RECORD), (kv.stride(0), RECORD, 1)) + + +def _scale_bytes(scale: torch.Tensor) -> torch.Tensor: + return scale if scale.dtype == torch.uint8 else scale.view(torch.uint8) + + +def _compute_units() -> int: + return current_platform.num_compute_units(torch.accelerator.current_device_index()) + + +def _unsupported_moe(layer: "DeepseekV4DecoderLayer", vllm_config) -> str | None: + """Why the kernels cannot run this deployment's MoE, or None. They read every + one of the 384 experts on each rank, TP-sharded, in AITER's A8W4 MXFP4 + layout.""" + from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import Mxfp4MoeBackend + + pc = vllm_config.parallel_config + if pc.enable_expert_parallel or pc.enable_eplb or pc.data_parallel_size > 1: + return "needs tensor parallelism only: no expert or data parallelism" + if envs.VLLM_ROCM_USE_AITER_MOE_A4W4_DSV4: + return "needs AITER's A8W4 MoE (VLLM_ROCM_USE_AITER_MOE_A4W4_DSV4 unset)" + experts = getattr(getattr(layer.ffn, "experts", None), "routed_experts", None) + if experts is None: + return None # not a routed MoE: the layer checks decline it + backend = getattr(experts.quant_method, "mxfp4_backend", None) + if backend != Mxfp4MoeBackend.AITER_MXFP4_BF16: + return f"needs the AITER_MXFP4_BF16 MoE backend, not {backend}" + return None + + +class MonoDecodeLayer: + """One decoder layer's mono path (``create``: None when the layer is not + eligible): the whole layer (``__call__``) or, with ``ffn_only``, the FFN + launch after vLLM's attention (``ffn``). Its weights are taken from the + layer at the first call, after loading.""" + + def __init__(self, ffn_only: bool = False) -> None: + self.ffn_only = ffn_only + self._weights: MonoLayerWeights | None = None + + @staticmethod + def create( + layer: "DeepseekV4DecoderLayer", vllm_config + ) -> "MonoDecodeLayer | None": + if not envs.VLLM_ROCM_MONO_DECODE: + return None + from vllm.platforms.rocm import get_cdna_version + + # the deployment: an explicit opt-in that cannot run is an error + why = None + if get_cdna_version() not in CDNA_VERSIONS: + why = f"needs CDNA{'/'.join(map(str, CDNA_VERSIONS))}" + elif get_tensor_model_parallel_world_size() not in (2, 4): + why = "needs tensor parallel size 2 or 4" + else: + why = _unsupported_moe(layer, vllm_config) + if why is None: + try: + from vllm.models.deepseek_v41.amd.mono.runner import BLOCKS + except ImportError as err: + why = f"needs FlyDSL and AITER's FlyDSL helpers ({err})" + else: + # every CTA of a launch stays resident: a partitioned GPU deadlocks + cus = _compute_units() + if cus < BLOCKS: + why = f"needs {BLOCKS} compute units, the GPU has {cus}" + if why is not None: + raise ValueError(f"VLLM_ROCM_MONO_DECODE {why}.") + # the layer: the kernels serve the backbone's seams and MoE + config = vllm_config.model_config.hf_config + attn, ffn = layer.attn, layer.ffn + if layer.use_sequence_parallel: + why = "sequence-parallel layer" + elif not layer.fuse_seam_norm: + why = "needs aiter's fused mHC seam" + elif attn.layer_id >= config.num_hidden_layers: + why = "draft layer" + elif ( + ffn.shared_experts is None + or ffn.gate.tid2eid is not None + or ffn.n_routed_experts != 384 + or ffn.n_activated_experts != 6 + or ffn.routed_scaling_factor != 1.5 + or ffn.swiglu_limit != 10.0 + or ffn.scoring_func != "sqrtsoftplus" + or not ffn.renormalize + ): + why = "MoE routing outside the kernels' (384 / 6, sqrtsoftplus, x1.5)" + if why is not None: + logger.info_once("DSv4.1 mono decode off for some layers: %s", why) + return None + # the whole layer: the kernels' attention takes the standard layers only + if layer.engram is not None: + why = "Engram" + elif attn.compressor is not None or attn.indexer is not None: + why = "a compressor / indexer" + elif attn.compress_ratio not in (1, 2): + why = f"compress ratio {attn.compress_ratio}" + elif attn.kv_mxfp8 or attn.kv_cache_dtype != "fp8_ds_mla": + why = f"a {attn.kv_cache_dtype} KV cache" + if why is not None: + logger.info_once( + "DSv4.1 mono decode: the FFN launch for layers with %s", why + ) + return MonoDecodeLayer(ffn_only=True) + logger.info_once("DSv4.1 mono decode: the whole standard backbone layers") + return MonoDecodeLayer() + + def weights(self, layer: "DeepseekV4DecoderLayer"): + if self._weights is not None: + return self._weights + from vllm.models.deepseek_v41.amd.mono.runner import MonoLayerWeights + + a, f = layer.attn, layer.ffn + e, sh = f.experts.routed_experts, f.shared_experts + assert sh is not None # ``create`` took only layers with one + self._weights = MonoLayerWeights( + attn=None if self.ffn_only else self._attn_weights(a), + hc_attn_fn=layer.hc_attn_fn, + hc_attn_scale=layer.hc_attn_scale, + hc_attn_base=layer.hc_attn_base, + attn_norm=layer.attn_norm.weight, + hc_ffn_fn=layer.hc_ffn_fn, + hc_ffn_scale=layer.hc_ffn_scale, + hc_ffn_base=layer.hc_ffn_base, + ffn_norm=layer.ffn_norm.weight, + gate_w=f.gate.weight, + bias=f.gate.e_score_correction_bias, + w13=e.w13_weight, + w13_s=e.w13_weight_scale, + w2=e.w2_weight, + w2_s=e.w2_weight_scale, + sgu=sh.gate_up_proj.weight, + sgu_s=_scale_bytes(sh.gate_up_proj.weight_scale), + sw2=sh.down_proj.weight, + sw2_s=_scale_bytes(sh.down_proj.weight_scale), + ) + self._weights.check(get_tensor_model_parallel_world_size()) + return self._weights + + @staticmethod + def _attn_weights(a): + from vllm.models.deepseek_v41.amd.mono.attention.plan import Dims + from vllm.models.deepseek_v41.amd.mono.runner import AttnWeights + + attn = AttnWeights( + layer_id=a.layer_id, + wqkv=a.fused_wqa_wkv.weight, + wqkv_scale=_scale_bytes(a.fused_wqa_wkv.weight_scale), + q_norm=a.q_norm.weight, + kv_norm=a.kv_norm.weight, + wq_b=a.wq_b.weight, + wq_b_scale=_scale_bytes(a.wq_b.weight_scale), + wo_a=a.wo_a.weight, + wo_a_scale=_scale_bytes(a.wo_a.weight_scale), + wo_b=a.wo_b.weight, + wo_b_scale=_scale_bytes(a.wo_b.weight_scale), + attn_sink=a.attn_sink, + cos_sin=a.rotary_emb.cos_sin_cache, + ratio=a.compress_ratio, + ) + attn.check(Dims(get_tensor_model_parallel_world_size())) + return attn + + @staticmethod + def _skip(why: str) -> None: + logger.debug_once("DSv4.1 mono decode skipped a step: %s", why) + return None + + @staticmethod + def _decode_step(layer: "DeepseekV4DecoderLayer", rows: int): + """(why this step takes the original path -- None for a decode-only step + the kernels take -- the layer's SWA metadata, its compressed cache's).""" + if not 1 <= rows <= MAX_ROWS: + return f"{rows} rows", None, None + if not is_forward_context_available(): + return "no forward context", None, None + fc = get_forward_context() + md = fc.attn_metadata + attn = layer.attn + swa, comp = None, None + if isinstance(md, dict): + swa = cast( + "DeepseekV4ROCMAiterSparseSWAMetadata | None", + md.get(attn.swa_cache_layer.prefix), + ) + if attn.compressed_cache_prefix is not None: + comp = cast( + "DeepseekV4FlashMLAMetadata | None", + md.get(attn.compressed_cache_prefix), + ) + why = None + if fc.cudagraph_runtime_mode not in (CUDAGraphMode.NONE, CUDAGraphMode.FULL): + # a PIECEWISE (breakable) capture is replayed for mixed batches too + why = f"{fc.cudagraph_runtime_mode} step" + elif not isinstance(md, dict): # a profile / dummy run, or DBO ubatches + why = "no attention metadata" + elif swa is None: + why = "layer metadata missing" + elif swa.num_prefills != 0 or swa.num_decodes == 0: + why = "not a decode-only step" + elif swa.num_decode_tokens > rows: + why = "more decode tokens than rows" + return why, swa, comp + + def __call__( + self, + layer: "DeepseekV4DecoderLayer", + x: torch.Tensor, + positions: torch.Tensor, + residual: torch.Tensor | None, + post_mix: torch.Tensor | None, + res_mix: torch.Tensor | None, + pre_mix: torch.Tensor | None, + ) -> tuple[torch.Tensor, ...] | None: + """The layer's outputs (x, residual, post_mix, res_mix, pre_mix) by the + mono kernels, or None when this step takes the original path (always for + an ``ffn_only`` layer).""" + if self.ffn_only: + return None + if residual is None or x.dim() != 2: + return self._skip("first layer") + why, swa, comp = self._decode_step(layer, x.shape[0]) + attn = layer.attn + tensors = (x, residual, post_mix, res_mix, pre_mix) + if why is None: + # the kernels' attention: causal SWA windows, the compressed cache + if comp is None: + why = "layer metadata missing" + elif swa.decode_swa_width != SWA_WIDTH: + why = f"SWA width {swa.decode_swa_width}" + elif ( + swa.decode_swa_indices is None + or swa.decode_swa_lens is None + or swa.token_to_req_indices is None + ): + why = "SWA decode metadata missing" + elif not all(t is not None and t.is_contiguous() for t in tensors): + why = "non-contiguous layer inputs" + if why is not None: + return self._skip(why) + assert swa is not None and comp is not None + runner = _mono_runner(x.device) + logger.info_once( + "DSv4.1 mono decode layer active (decode steps of <= %d rows)", MAX_ROWS + ) + out = runner.forward( + self.weights(layer), + x, + residual, + post_mix, + res_mix, + pre_mix, + positions, + swa.slot_mapping, + _record_view(attn.swa_cache_layer.kv_cache, swa.block_size), + swa.decode_swa_indices, + swa.decode_swa_lens, + swa.token_to_req_indices, + topk_indices=attn.topk_indices_buffer, + comp_cache=_record_view( + attn._compressed_kv_cache(), comp.block_size // attn.compress_ratio + ), + comp_block_table=comp.block_table, + ) + return out + + def ffn( + self, + layer: "DeepseekV4DecoderLayer", + part: torch.Tensor, + residual: torch.Tensor, + post_mix: torch.Tensor, + res_mix: torch.Tensor, + pre_mix: torch.Tensor, + ) -> tuple[torch.Tensor, ...] | None: + """The layer's outputs by the FFN launch from ``part``, wo_b's unreduced + output, and the attention seam's outputs; None when this step takes the + original path (the caller then reduces ``part`` itself).""" + why, _, _ = self._decode_step(layer, part.shape[0]) + tensors = (part, residual, post_mix, res_mix, pre_mix) + if why is None and not all(t.is_contiguous() for t in tensors): + why = "non-contiguous layer inputs" + if why is not None: + return self._skip(why) + runner = _mono_runner(part.device) + logger.info_once( + "DSv4.1 mono decode FFN launch active (decode steps of <= %d rows)", + MAX_ROWS, + ) + return runner.ffn(self.weights(layer), *tensors) diff --git a/vllm/models/deepseek_v41/amd/rocm.py b/vllm/models/deepseek_v41/amd/rocm.py index 3aa54f9b0e..3ba6eebcbd 100644 --- a/vllm/models/deepseek_v41/amd/rocm.py +++ b/vllm/models/deepseek_v41/amd/rocm.py @@ -929,7 +929,9 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention): def _wo_b_after_wo_a(self, zf: torch.Tensor) -> torch.Tensor: if self._wo_b_scale is not None and zf.dim() == 2: - return self._bpre_attn_gemm(self.wo_b.weight, self._wo_b_scale, zf, True) + return self._bpre_attn_gemm( + self.wo_b.weight, self._wo_b_scale, zf, self.wo_b.reduce_results + ) return self.wo_b(zf) def forward_mqa(