[ROCm][Perf][DSv4.1] Mono decode layer for DeepSeek-V4.1-Flash on MI355X (#60397)
Signed-off-by: Fangzhou Ai <fangzhou.ai@amd.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Fangzhou Ai
co-authored by
Claude Opus 5.5
Claude Opus 5.5
parent
3c802af3cf
commit
193922d6ff
@@ -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
|
||||
|
||||
@@ -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", {})))
|
||||
@@ -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))
|
||||
@@ -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)
|
||||
@@ -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: (
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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).
|
||||
@@ -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``."""
|
||||
@@ -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)."""
|
||||
@@ -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)
|
||||
@@ -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::<stem>_<params>``."""
|
||||
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)
|
||||
@@ -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)
|
||||
@@ -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()
|
||||
@@ -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)
|
||||
@@ -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)."""
|
||||
@@ -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]
|
||||
)
|
||||
)
|
||||
@@ -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)
|
||||
@@ -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
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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(".")
|
||||
@@ -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``)."""
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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])
|
||||
@@ -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)
|
||||
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user