From da646cdf831664f75cd2385a444446a7a8b0e201 Mon Sep 17 00:00:00 2001 From: moreh Date: Thu, 20 Aug 2026 06:42:38 +0000 Subject: [PATCH] [MI355X][ROCm][AMD] Fuse the DSA indexer QK prologue for GLM-5.x The DSA indexer spends five kernel launches per layer per step on pre-processing: LayerNorm on k, RoPE on q and k, FP8 quantization of q, folding the q scale into the indexer weights, and the FP8 k quantization plus paged K-cache write. At decode sizes this is launch-bound. AITER's indexer_qk_rope_quant_and_cache does all five in one launch. Wire it in behind a new default-off flag, VLLM_ROCM_USE_AITER_INDEXER_QK_FUSION; the kernel writes the indexer K cache itself, so the indexer op is built with skip_k_cache_insert=True from the same init-time predicate. Two data-movement optimizations ride along: the rope cos_sin_cache halves are registered once as strided views instead of being re-split every layer every step, and one zero-initialized q_fp8/weights pair is shared by a model's indexer layers. Signed-off-by: moreh Co-authored-by: Claude Opus 5 (1M context) --- .../kernels/benchmark_indexer_qk_fusion.py | 203 ++++++++++ .../test_rocm_aiter_indexer_qk_fusion.py | 379 ++++++++++++++++++ vllm/_aiter_ops.py | 17 + vllm/config/compilation.py | 1 + vllm/envs.py | 11 + vllm/model_executor/models/deepseek_mtp.py | 8 + vllm/model_executor/models/deepseek_v2.py | 191 ++++++++- .../v1/attention/ops/rocm_aiter_mla_sparse.py | 136 +++++++ 8 files changed, 944 insertions(+), 2 deletions(-) create mode 100644 benchmarks/kernels/benchmark_indexer_qk_fusion.py create mode 100644 tests/kernels/attention/test_rocm_aiter_indexer_qk_fusion.py diff --git a/benchmarks/kernels/benchmark_indexer_qk_fusion.py b/benchmarks/kernels/benchmark_indexer_qk_fusion.py new file mode 100644 index 000000000000..d943d8f192f7 --- /dev/null +++ b/benchmarks/kernels/benchmark_indexer_qk_fusion.py @@ -0,0 +1,203 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Fused vs unfused DSA indexer QK pre-processing on ROCm. + +The DSA indexer (DeepSeek-V3.2, GLM-5.x) runs five launches per layer per step: +LayerNorm(k), RoPE(q, k), per-token-group fp8 quant of q, folding the q scale +into the indexer weights, and the fp8 K quant + paged K-cache write. With +VLLM_ROCM_USE_AITER_INDEXER_QK_FUSION=1 one AITER kernel does all five. + +Usage: + python benchmarks/kernels/benchmark_indexer_qk_fusion.py + python benchmarks/kernels/benchmark_indexer_qk_fusion.py --num-tokens 1 8 64 +""" + +import argparse +import functools + +import torch + +from vllm import _custom_ops as ops +from vllm.model_executor.layers.quantization.utils.fp8_utils import ( + per_token_group_quant_fp8, +) +from vllm.platforms import current_platform +from vllm.triton_utils import triton +from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( + indexer_k_quant_and_cache_triton, +) + +HEAD_DIM = 128 +ROPE_DIM = 64 +N_HEAD = 32 +MAX_POS = 65536 +QUANT_BLOCK = 128 +EPS = 1e-6 +SCALE_FMT = "ue8m0" +WEIGHTS_SCALE = HEAD_DIM**-0.5 * N_HEAD**-0.5 + + +def _unfused( + q, k_raw, weights_raw, positions, cache, norm_w, norm_b, kv, slots, is_neox +): + num_tokens = q.shape[0] + k = torch.nn.functional.layer_norm( + k_raw.float(), (HEAD_DIM,), norm_w.float(), norm_b.float(), EPS + ).to(q.dtype) + q = q.clone() + ops.rotary_embedding( + positions, + q[..., :ROPE_DIM], + k[..., :ROPE_DIM].unsqueeze(1), + ROPE_DIM, + cache, + is_neox, + ) + q_fp8, q_scale = per_token_group_quant_fp8( + q.view(-1, HEAD_DIM), QUANT_BLOCK, column_major_scales=False, use_ue8m0=True + ) + _ = weights_raw.float() * q_scale.view(num_tokens, N_HEAD) * WEIGHTS_SCALE + indexer_k_quant_and_cache_triton(k, kv, slots, QUANT_BLOCK, SCALE_FMT) + + +def _fused( + q, + k_raw, + weights_raw, + positions, + cache, + norm_w, + norm_b, + kv, + slots, + q_out, + w_out, + is_neox, +): + from aiter import indexer_qk_rope_quant_and_cache + + half = ROPE_DIM // 2 + indexer_qk_rope_quant_and_cache( + q, + q_out, + weights_raw, + w_out, + k_raw, + kv, + slots, + norm_w, + norm_b, + positions, + cache[:, :half], + cache[:, half:], + EPS, + QUANT_BLOCK, + SCALE_FMT, + WEIGHTS_SCALE, + preshuffle=kv.shape[1] > 1, + is_neox=is_neox, + ) + + +def _time_us(fn) -> float: + ms = triton.testing.do_bench(fn, warmup=25, rep=100) + return ms * 1000.0 + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument( + "--num-tokens", type=int, nargs="+", default=[1, 8, 32, 64, 256, 1024] + ) + parser.add_argument("--block-size", type=int, default=64) + parser.add_argument( + "--is-neox", + action=argparse.BooleanOptionalAction, + default=False, + help="RoPE layout. GLM-5.x sets indexer_rope_interleave, i.e. is_neox " + "False; DeepSeek-V3.2 leaves it at the NeoX default.", + ) + parser.add_argument( + "--repeat", + type=int, + default=3, + help="Measure each point this many times and report the median. " + "Process-level state (kernel tuning caches, clocks) moves these numbers " + "more than do_bench's own variance does.", + ) + args = parser.parse_args() + + if not current_platform.is_rocm(): + raise SystemExit("ROCm only") + fp8 = current_platform.fp8_dtype() + dev, dt = "cuda", torch.bfloat16 + + def median(runs: list[float]) -> tuple[float, float, float]: + ordered = sorted(runs) + return ordered[len(ordered) // 2], ordered[0], ordered[-1] + + rows = [] + for num_tokens in args.num_tokens: + num_blocks = (num_tokens + args.block_size - 1) // args.block_size + 2 + q = torch.randn(num_tokens, N_HEAD, HEAD_DIM, device=dev, dtype=dt) + kw = torch.randn(num_tokens, HEAD_DIM + N_HEAD, device=dev, dtype=dt) + positions = torch.randint( + 0, MAX_POS, (num_tokens,), device=dev, dtype=torch.int64 + ) + norm_w = torch.randn(HEAD_DIM, device=dev, dtype=dt) + norm_b = torch.randn(HEAD_DIM, device=dev, dtype=dt) + cache = torch.randn(MAX_POS, ROPE_DIM, device=dev, dtype=dt) + kv = torch.zeros( + num_blocks, args.block_size, HEAD_DIM + 4, dtype=fp8, device=dev + ) + slots = torch.randperm( + num_blocks * args.block_size, device=dev, dtype=torch.int64 + )[:num_tokens] + q_out = torch.zeros((num_tokens, N_HEAD, HEAD_DIM), dtype=fp8, device=dev) + w_out = torch.zeros((num_tokens, N_HEAD), dtype=torch.float32, device=dev) + + unfused_fn = functools.partial( + _unfused, + q, + kw[:, :HEAD_DIM], + kw[:, HEAD_DIM:], + positions, + cache, + norm_w, + norm_b, + kv, + slots, + args.is_neox, + ) + fused_fn = functools.partial( + _fused, + q, + kw[:, :HEAD_DIM], + kw[:, HEAD_DIM:], + positions, + cache, + norm_w, + norm_b, + kv, + slots, + q_out, + w_out, + args.is_neox, + ) + unfused_runs = [_time_us(unfused_fn) for _ in range(args.repeat)] + fused_runs = [_time_us(fused_fn) for _ in range(args.repeat)] + rows.append((num_tokens, median(unfused_runs), median(fused_runs))) + + print( + f"{'tokens':>8} {'unfused us':>12} {'fused us':>10} {'speedup':>9}" + f" spread over {args.repeat} repeats" + ) + for num_tokens, (u_med, u_lo, u_hi), (f_med, f_lo, f_hi) in rows: + print( + f"{num_tokens:>8} {u_med:>12.2f} {f_med:>10.2f} {u_med / f_med:>8.2f}x" + f" unfused {u_lo:.1f}-{u_hi:.1f}, fused {f_lo:.1f}-{f_hi:.1f}" + ) + + +if __name__ == "__main__": + main() diff --git a/tests/kernels/attention/test_rocm_aiter_indexer_qk_fusion.py b/tests/kernels/attention/test_rocm_aiter_indexer_qk_fusion.py new file mode 100644 index 000000000000..b1d98bac55e8 --- /dev/null +++ b/tests/kernels/attention/test_rocm_aiter_indexer_qk_fusion.py @@ -0,0 +1,379 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit test for the ROCm AITER fused DSA indexer QK path. + +Compares ``torch.ops.vllm.rocm_aiter_indexer_qk_rope_quant_and_cache`` against +the unfused flow it replaces in ``Indexer.forward`` +(vllm/model_executor/models/deepseek_v2.py), built from the same primitives: + + k = layer_norm(k, k_norm.weight, k_norm.bias, eps) + ops.rotary_embedding(positions, q[..., :64], k[..., :64], 64, cache, neox) + q_fp8, q_scale = per_token_group_quant_fp8(q, 128, use_ue8m0=True) + weights_out = weights * q_scale * softmax_scale * n_head_scale + indexer_k_quant_and_cache_triton(k, kv_cache, slot_mapping, 128, "ue8m0") + +The fused kernel folds the q scale into ``weights_out``, so the comparison is +on the products the downstream MQA-logits kernels consume +(``q_fp8 * weights_out``) and on the dequantized indexer K cache. + +The two paths cannot be bit-equal, and the difference is in the q scale, not in +the RoPE: the kernel rounds the roped q through bf16 exactly as the unfused flow +does, but derives a plain fp32 scale from it, while per_token_group_quant_fp8 +rounds that scale to a power of two (use_ue8m0=True). A different scale moves +every element of the token by at most one fp8 code, in either direction, so +asserting closeness to the unfused path would penalise whichever path happens to +sit further from the truth. Both are instead scored against an fp64 golden of the +same math, and the fused path must be no less accurate than the unfused one, with +every element within one fp8 code of it. + +Speed is not asserted here: see benchmarks/kernels/benchmark_indexer_qk_fusion.py. +""" + +from types import SimpleNamespace + +import pytest +import torch + +from vllm import _custom_ops as ops +from vllm.model_executor.layers.quantization.utils.fp8_utils import ( + per_token_group_quant_fp8, +) +from vllm.platforms import current_platform + +_SKIP_NON_MI3XX = True +if current_platform.is_rocm(): + from vllm.platforms.rocm import on_mi3xx + + _SKIP_NON_MI3XX = not on_mi3xx() + +pytestmark = [ + pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm-specific tests"), + pytest.mark.skipif(_SKIP_NON_MI3XX, reason="MI300/MI350 only (aiter CK kernel)"), +] + +HEAD_DIM = 128 +ROPE_DIM = 64 +N_HEAD = 32 +MAX_POS = 65536 +QUANT_BLOCK = 128 +EPS = 1e-6 +SCALE_FMT = "ue8m0" +WEIGHTS_SCALE = HEAD_DIM**-0.5 * N_HEAD**-0.5 +PREFIX = "model.layers.0.self_attn.indexer.k_cache" + +# fp8-e4m3 keeps 3 mantissa bits, so one code step is at most 12.5% of an +# element: a value sitting on a quantization boundary flips to the neighbouring +# code under any reordering of the RoPE/LayerNorm arithmetic. +ONE_CODE_REL = 0.13 +# Slack on "no less accurate than unfused": both are fp8, so their errors +# against the golden are dominated by the same quantization step. +ACCURACY_SLACK = 1.10 +MAX_FP8_REL_L2 = 0.05 + + +def _require_aiter() -> None: + from vllm._aiter_ops import is_aiter_found_and_supported + + if not is_aiter_found_and_supported(): + pytest.skip("aiter is required for the fused indexer QK kernel") + + +def _indexer_metadata(slot_mapping: torch.Tensor): + from vllm.v1.attention.backends.mla.indexer import DeepseekV32IndexerMetadata + + num_tokens = int(slot_mapping.shape[0]) + return DeepseekV32IndexerMetadata( + seq_lens=torch.tensor([num_tokens], device=slot_mapping.device), + max_seq_len=num_tokens, + slot_mapping=slot_mapping, + num_decodes=0, + num_decode_tokens=0, + num_prefills=1, + num_prefill_tokens=num_tokens, + ) + + +def _inputs(num_tokens: int, block_size: int, capacity: int | None = None): + capacity = capacity or num_tokens + num_blocks = (capacity + block_size - 1) // block_size + 2 + dev, dt = "cuda", torch.bfloat16 + return SimpleNamespace( + q=torch.randn(capacity, N_HEAD, HEAD_DIM, device=dev, dtype=dt), + kw=torch.randn(capacity, HEAD_DIM + N_HEAD, device=dev, dtype=dt), + positions=torch.randint(0, MAX_POS, (capacity,), device=dev, dtype=torch.int64), + slots=torch.randperm(num_blocks * block_size, device=dev, dtype=torch.int64)[ + :num_tokens + ], + norm_weight=torch.randn(HEAD_DIM, device=dev, dtype=dt), + norm_bias=torch.randn(HEAD_DIM, device=dev, dtype=dt), + cos_sin_cache=torch.randn(MAX_POS, ROPE_DIM, device=dev, dtype=dt), + kv_cache=torch.zeros( + num_blocks, + block_size, + HEAD_DIM + 4, + dtype=current_platform.fp8_dtype(), + device=dev, + ), + ) + + +def _reference(t, kv_cache: torch.Tensor, is_neox: bool): + """The five launches the fused kernel replaces, in the same order.""" + from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( + indexer_k_quant_and_cache_triton, + ) + + num_tokens = t.q.shape[0] + k = torch.nn.functional.layer_norm( + t.kw[:, :HEAD_DIM].float(), + (HEAD_DIM,), + t.norm_weight.float(), + t.norm_bias.float(), + EPS, + ).to(t.q.dtype) + q = t.q.clone() + ops.rotary_embedding( + t.positions, + q[..., :ROPE_DIM], + k[..., :ROPE_DIM].unsqueeze(1), + ROPE_DIM, + t.cos_sin_cache, + is_neox, + ) + q_fp8, q_scale = per_token_group_quant_fp8( + q.view(-1, HEAD_DIM), QUANT_BLOCK, column_major_scales=False, use_ue8m0=True + ) + weights = ( + t.kw[:, HEAD_DIM:].float() * q_scale.view(num_tokens, N_HEAD) * WEIGHTS_SCALE + ) + indexer_k_quant_and_cache_triton(k, kv_cache, t.slots, QUANT_BLOCK, SCALE_FMT) + return q_fp8.view(num_tokens, N_HEAD, HEAD_DIM), weights + + +def _fused(t, kv_cache, is_neox, q_out=None, w_out=None, zero_outputs=False): + half = ROPE_DIM // 2 + if q_out is None: + q_out = torch.zeros( + t.q.shape, dtype=current_platform.fp8_dtype(), device=t.q.device + ) + if w_out is None: + w_out = torch.zeros(t.q.shape[:2], dtype=torch.float32, device=t.q.device) + torch.ops.vllm.rocm_aiter_indexer_qk_rope_quant_and_cache( + PREFIX, + kv_cache, + t.q, + t.kw[:, :HEAD_DIM], + t.kw[:, HEAD_DIM:], + t.positions, + t.cos_sin_cache[:, :half], + t.cos_sin_cache[:, half:], + t.norm_weight, + t.norm_bias, + q_out, + w_out, + EPS, + QUANT_BLOCK, + SCALE_FMT, + WEIGHTS_SCALE, + zero_outputs, + is_neox, + ) + return q_out, w_out + + +def _patch_context(monkeypatch, slot_mapping: torch.Tensor) -> None: + import vllm.v1.attention.ops.rocm_aiter_mla_sparse as mla_sparse + + monkeypatch.setattr( + mla_sparse, + "get_forward_context", + lambda: SimpleNamespace( + attn_metadata={PREFIX: _indexer_metadata(slot_mapping)} + ), + ) + + +def _golden(t, is_neox: bool) -> tuple[torch.Tensor, torch.Tensor]: + """fp64 LayerNorm + RoPE: the same math with no intermediate rounding.""" + half = ROPE_DIM // 2 + k = torch.nn.functional.layer_norm( + t.kw[:, :HEAD_DIM].double(), + (HEAD_DIM,), + t.norm_weight.double(), + t.norm_bias.double(), + EPS, + ) + cos = t.cos_sin_cache[t.positions, :half].double() + sin = t.cos_sin_cache[t.positions, half:].double() + + def rope(x: torch.Tensor) -> torch.Tensor: + c, s_ = cos, sin + while c.dim() < x.dim(): + c, s_ = c.unsqueeze(-2), s_.unsqueeze(-2) + pe, rest = x[..., :ROPE_DIM], x[..., ROPE_DIM:] + if is_neox: + x1, x2 = pe[..., :half], pe[..., half:] + roped = torch.cat([x1 * c - x2 * s_, x2 * c + x1 * s_], dim=-1) + else: + x1, x2 = pe[..., 0::2], pe[..., 1::2] + roped = torch.stack([x1 * c - x2 * s_, x2 * c + x1 * s_], dim=-1).flatten( + -2 + ) + return torch.cat([roped, rest], dim=-1) + + return rope(t.q.double()), rope(k) + + +def _rel_l2(got: torch.Tensor, golden: torch.Tensor) -> float: + return ((got.double() - golden).norm() / golden.norm().clamp(min=1e-12)).item() + + +def _assert_no_less_accurate( + fused: torch.Tensor, unfused: torch.Tensor, golden: torch.Tensor, what: str +) -> None: + err_fused, err_unfused = _rel_l2(fused, golden), _rel_l2(unfused, golden) + assert err_fused <= MAX_FP8_REL_L2, ( + f"{what}: fused error vs fp64 golden {err_fused:.3e} exceeds fp8 granularity" + ) + assert err_fused <= err_unfused * ACCURACY_SLACK, ( + f"{what}: fused is less accurate than unfused " + f"({err_fused:.3e} vs {err_unfused:.3e})" + ) + # And the two paths must still agree element-wise to within one fp8 code. + scale = torch.maximum(fused.abs(), unfused.abs()) + live = scale > 1e-3 * unfused.abs().max() + rel = (fused - unfused).abs()[live] / scale[live] + worst = rel.max().item() + assert worst <= ONE_CODE_REL, ( + f"{what}: {(rel > ONE_CODE_REL).sum().item()} of {rel.numel()} elements " + f"differ by more than one fp8 code (worst rel {worst:.4f})" + ) + + +def _dequant_cache(kv_cache: torch.Tensor, slots: torch.Tensor) -> torch.Tensor: + """Dequantize the rows `slots` point at, for either in-block layout.""" + num_blocks, block_size = kv_cache.shape[0], kv_cache.shape[1] + flat = kv_cache.view(num_blocks, -1) + values = flat[:, : block_size * HEAD_DIM] + scales = flat[:, block_size * HEAD_DIM :].contiguous().view(torch.float32) + tile = 16 + j = torch.arange(HEAD_DIM, device=kv_cache.device) + out = torch.empty( + (slots.shape[0], HEAD_DIM), dtype=torch.float32, device=kv_cache.device + ) + for i, slot in enumerate(slots.tolist()): + block_id, off = slot // block_size, slot % block_size + if block_size == 1: + idx = off * HEAD_DIM + j + else: + # 16x16-tiled in-block layout, mirroring the writer's SHUFFLE path. + idx = ( + (off // tile) * tile * HEAD_DIM + + (off % tile) * tile + + (j // tile) * tile * tile + + j % tile + ) + out[i] = ( + values[block_id, idx].view(kv_cache.dtype).float() * scales[block_id, off] + ) + return out + + +@pytest.mark.parametrize("num_tokens", [1, 7, 32, 257, 1023]) +@pytest.mark.parametrize("block_size", [1, 64]) +@pytest.mark.parametrize("is_neox", [True, False]) +@torch.inference_mode() +def test_fused_matches_unfused(monkeypatch, num_tokens, block_size, is_neox): + """One fused launch reproduces the five unfused launches it replaces.""" + _require_aiter() + torch.manual_seed(0) + t = _inputs(num_tokens, block_size) + kv_ref = torch.zeros_like(t.kv_cache) + + q_fp8_ref, weights_ref = _reference(t, kv_ref, is_neox) + _patch_context(monkeypatch, t.slots) + q_fp8, weights_out = _fused(t, t.kv_cache, is_neox) + q_golden, k_golden = _golden(t, is_neox) + + # Scale-fold invariant: the kernel may split the scale between q_fp8 and + # weights_out differently, only the product reaches the logits kernels. + weights_raw = t.kw[:, HEAD_DIM:].double() * WEIGHTS_SCALE + _assert_no_less_accurate( + q_fp8.float() * weights_out.unsqueeze(-1), + q_fp8_ref.float() * weights_ref.unsqueeze(-1), + q_golden * weights_raw.unsqueeze(-1), + "q_fp8 * weights_out", + ) + _assert_no_less_accurate( + _dequant_cache(t.kv_cache, t.slots), + _dequant_cache(kv_ref, t.slots), + k_golden, + "indexer K cache", + ) + + +@torch.inference_mode() +def test_unowned_and_padded_rows_stay_zero(monkeypatch): + """Rows the kernel skips must read as zero, not as stale data. + + The kernel early-returns on ``slot_mapping < 0`` - PAD_SLOT_ID marks both + CUDA-graph padding and, under context parallel, tokens this rank does not + own - and never touches rows past ``slot_mapping``. Decode reads + ``weights[:batch_size * next_n]``, which covers those rows, so they must be + zero for the padded logits they feed to stay finite. + """ + _require_aiter() + torch.manual_seed(0) + num_tokens, capacity = 64, 96 + t = _inputs(num_tokens, block_size=64, capacity=capacity) + t.slots[::4] = -1 + before = t.kv_cache.view(torch.uint8).clone() + + _patch_context(monkeypatch, t.slots) + q_fp8, weights_out = _fused(t, t.kv_cache, is_neox=True) + + skipped = t.slots < 0 + assert torch.all(q_fp8[:num_tokens][skipped].view(torch.uint8) == 0) + assert torch.all(weights_out[:num_tokens][skipped] == 0) + assert torch.all(q_fp8[num_tokens:].view(torch.uint8) == 0) + assert torch.all(weights_out[num_tokens:] == 0) + # Not vacuous: the owned rows were written, in both outputs and the cache. + assert not torch.all(q_fp8[:num_tokens][~skipped].view(torch.uint8) == 0) + assert not torch.equal(t.kv_cache.view(torch.uint8), before) + + +@torch.inference_mode() +def test_op_schema(monkeypatch): + """opcheck the custom op: fake impl for tracing, declared mutates_args.""" + _require_aiter() + from tests.kernels.utils import opcheck + + torch.manual_seed(0) + t = _inputs(num_tokens=8, block_size=64) + _patch_context(monkeypatch, t.slots) + half = ROPE_DIM // 2 + opcheck( + torch.ops.vllm.rocm_aiter_indexer_qk_rope_quant_and_cache, + ( + PREFIX, + t.kv_cache, + t.q, + t.kw[:, :HEAD_DIM], + t.kw[:, HEAD_DIM:], + t.positions, + t.cos_sin_cache[:, :half], + t.cos_sin_cache[:, half:], + t.norm_weight, + t.norm_bias, + torch.zeros( + t.q.shape, dtype=current_platform.fp8_dtype(), device=t.q.device + ), + torch.zeros(t.q.shape[:2], dtype=torch.float32, device=t.q.device), + EPS, + QUANT_BLOCK, + SCALE_FMT, + WEIGHTS_SCALE, + False, + True, + ), + ) diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index b6a504d0cfbe..55117e8002b5 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -14,6 +14,8 @@ from vllm.utils.import_utils import PlaceholderModule from vllm.utils.torch_utils import direct_register_custom_op from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( + rocm_aiter_indexer_qk_rope_quant_and_cache, + rocm_aiter_indexer_qk_rope_quant_and_cache_fake, rocm_aiter_sparse_attn_indexer, rocm_aiter_sparse_attn_indexer_fake, ) @@ -1685,6 +1687,7 @@ def get_moe_dispatch_policy(cls) -> int: _FMOE_ENABLED = envs.VLLM_ROCM_USE_AITER_MOE _MLA_ENABLED = envs.VLLM_ROCM_USE_AITER_MLA _MHA_ENABLED = envs.VLLM_ROCM_USE_AITER_MHA + _INDEXER_QK_FUSION_ENABLED = envs.VLLM_ROCM_USE_AITER_INDEXER_QK_FUSION _SHUFFLE_KV_CACHE_ENABLED = envs.VLLM_ROCM_SHUFFLE_KV_CACHE_LAYOUT _TRITON_UNIFIED_ATTN_ENABLED = envs.VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION # TODO: Consolidate under _LINEAR_ENABLED @@ -1718,6 +1721,7 @@ def refresh_env_variables(cls): cls._FMOE_ENABLED = envs.VLLM_ROCM_USE_AITER_MOE cls._MLA_ENABLED = envs.VLLM_ROCM_USE_AITER_MLA cls._MHA_ENABLED = envs.VLLM_ROCM_USE_AITER_MHA + cls._INDEXER_QK_FUSION_ENABLED = envs.VLLM_ROCM_USE_AITER_INDEXER_QK_FUSION cls._SHUFFLE_KV_CACHE_ENABLED = envs.VLLM_ROCM_SHUFFLE_KV_CACHE_LAYOUT cls._TRITON_UNIFIED_ATTN_ENABLED = envs.VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION cls._FP8BMM_ENABLED = envs.VLLM_ROCM_USE_AITER_FP8BMM @@ -1832,6 +1836,11 @@ def is_linear_fp8_enabled(cls) -> bool: def is_fused_moe_enabled(cls) -> bool: return cls._AITER_ENABLED and cls._FMOE_ENABLED + @classmethod + @if_aiter_supported + def is_indexer_qk_fusion_enabled(cls) -> bool: + return cls._AITER_ENABLED and cls._INDEXER_QK_FUSION_ENABLED + @classmethod @if_aiter_supported def is_fusion_moe_shared_experts_enabled(cls) -> bool: @@ -2221,6 +2230,14 @@ def register_ops_once() -> None: dispatch_key=current_platform.dispatch_key, ) + direct_register_custom_op( + op_name="rocm_aiter_indexer_qk_rope_quant_and_cache", + op_func=rocm_aiter_indexer_qk_rope_quant_and_cache, + mutates_args=["kv_cache", "q_fp8_out", "weights_out"], + fake_impl=rocm_aiter_indexer_qk_rope_quant_and_cache_fake, + dispatch_key=current_platform.dispatch_key, + ) + direct_register_custom_op( op_name="aiter_fp8_attn_wrapper", op_func=_rocm_aiter_fp8_attn_impl, diff --git a/vllm/config/compilation.py b/vllm/config/compilation.py index 3cd227d72ce4..eb7965eeaca1 100644 --- a/vllm/config/compilation.py +++ b/vllm/config/compilation.py @@ -774,6 +774,7 @@ class CompilationConfig: "vllm::olmo_hybrid_gdn_full_forward", "vllm::sparse_attn_indexer", "vllm::rocm_aiter_sparse_attn_indexer", + "vllm::rocm_aiter_indexer_qk_rope_quant_and_cache", "vllm::deepseek_v4_attention", "vllm::hpc_rope_norm_forward", ] diff --git a/vllm/envs.py b/vllm/envs.py index b49841d9fff0..bf6f357ff99b 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -143,6 +143,7 @@ VLLM_ROCM_USE_AITER_MLA: bool = True VLLM_ROCM_AITER_MLA_ASM_PADDING: Literal["auto", "gluon", "asm"] = "auto" VLLM_ROCM_USE_AITER_MHA: bool = True + VLLM_ROCM_USE_AITER_INDEXER_QK_FUSION: bool = False VLLM_ROCM_USE_AITER_FP4_ASM_GEMM: bool = False VLLM_ROCM_USE_AITER_TRITON_ROPE: bool = False VLLM_ROCM_USE_AITER_FP8BMM: bool = True @@ -1303,6 +1304,16 @@ def _resolve_rust_cli_path() -> str | None: "VLLM_ROCM_USE_AITER_MHA": lambda: ( os.getenv("VLLM_ROCM_USE_AITER_MHA", "True").lower() in ("true", "1") ), + # Whether to use the aiter fused indexer QK kernel + # (indexer_qk_rope_quant_and_cache) on DeepSeek sparse attention models + # (DeepSeek-V3.2, GLM-5.x): fuses the indexer's Q/K RoPE, K LayerNorm, + # FP8 quantization and K-cache write into one launch. Needs + # VLLM_ROCM_USE_AITER=1 on gfx942/gfx950. + # By default is disabled. + "VLLM_ROCM_USE_AITER_INDEXER_QK_FUSION": lambda: ( + os.getenv("VLLM_ROCM_USE_AITER_INDEXER_QK_FUSION", "False").lower() + in ("true", "1") + ), # Whether to use aiter fp4 gemm asm. # By default is disabled. "VLLM_ROCM_USE_AITER_FP4_ASM_GEMM": lambda: ( diff --git a/vllm/model_executor/models/deepseek_mtp.py b/vllm/model_executor/models/deepseek_mtp.py index 8ac43e0bb2a1..ab3d16b990c9 100644 --- a/vllm/model_executor/models/deepseek_mtp.py +++ b/vllm/model_executor/models/deepseek_mtp.py @@ -37,6 +37,7 @@ DeepseekV2DecoderLayer, DeepseekV2MixtureOfExperts, DeepseekV2MoE, + IndexerQKFusionBuffers, _try_load_fp8_indexer_wk, ) from .utils import ( @@ -90,8 +91,14 @@ def __init__(self, vllm_config: VllmConfig, prefix: str) -> None: dtype=torch.int32, device=self.device, ) + # Each MTP layer is its own buffer root: it always re-zeros its + # pair instead of inheriting the target model's unwritten rows. + indexer_qk_fusion_buffers = IndexerQKFusionBuffers.maybe_build( + vllm_config, config, self.device + ) else: topk_indices_buffer = None + indexer_qk_fusion_buffers = None self.shared_head = SharedHead( config=config, prefix=prefix, quant_config=quant_config @@ -101,6 +108,7 @@ def __init__(self, vllm_config: VllmConfig, prefix: str) -> None: prefix, config=self.config, topk_indices_buffer=topk_indices_buffer, + indexer_qk_fusion_buffers=indexer_qk_fusion_buffers, ) def forward( diff --git a/vllm/model_executor/models/deepseek_v2.py b/vllm/model_executor/models/deepseek_v2.py index d5521a1960b8..7d259d94ad34 100644 --- a/vllm/model_executor/models/deepseek_v2.py +++ b/vllm/model_executor/models/deepseek_v2.py @@ -78,6 +78,7 @@ scaled_dequantize, ) from vllm.model_executor.layers.rotary_embedding import get_rope +from vllm.model_executor.layers.rotary_embedding.base import RotaryEmbeddingBase from vllm.model_executor.layers.sparse_attn_indexer import ( SparseAttnIndexer, fused_indexer_q_rope_quant, @@ -97,11 +98,15 @@ ) from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors -from vllm.utils.torch_utils import direct_register_custom_op +from vllm.utils.torch_utils import _encode_layer_name, direct_register_custom_op from vllm.v1.attention.backend import AttentionBackend from vllm.v1.attention.backends.mla.indexer import ( DeepseekV32IndexerBackend, ) +from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( + get_indexer_rope_halves, + register_indexer_rope_halves, +) from vllm.v1.kv_cache_interface import KVCacheSpec, MLAAttentionSpec from .interfaces import ( @@ -459,6 +464,7 @@ def __init__( cache_config: CacheConfig | None = None, quant_config: QuantizationConfig | None = None, topk_indices_buffer: torch.Tensor | None = None, + indexer_qk_fusion_buffers: "IndexerQKFusionBuffers | None" = None, reduce_results: bool = True, prefix: str = "", ) -> None: @@ -480,6 +486,9 @@ def __init__( "topk_indices_buffer is not \ supported for DeepseekV2Attention" ) + assert indexer_qk_fusion_buffers is None, ( + "indexer_qk_fusion_buffers is not supported for DeepseekV2Attention" + ) if self.q_lora_rank is not None: self.q_a_proj = ReplicatedLinear( @@ -639,6 +648,71 @@ def get_attn_backend(self) -> type[AttentionBackend]: return DeepseekV32IndexerBackend +class IndexerQKFusionBuffers: + """Model-shared q_fp8/weights output pair for the fused indexer QK kernel. + + Allocated next to ``topk_indices_buffer`` and sized the same way, so the + address is stable before CUDA-graph capture. The first Indexer constructed + under the root claims the per-pass zero-fill; which rows stay unwritten + depends only on ``slot_mapping``, which every layer of a pass shares, so + zeroing once is enough. Takes the per-step fills from 2 per indexer layer + down to 2, plus the matching allocations. + """ + + def __init__( + self, capacity: int, n_heads: int, head_dim: int, device: torch.device + ): + # Zero-init is load-bearing: rows the kernel skips must read as zero. + self.q_fp8 = torch.zeros( + (capacity, n_heads, head_dim), + dtype=current_platform.fp8_dtype(), + device=device, + ) + self.weights = torch.zeros( + (capacity, n_heads), dtype=torch.float32, device=device + ) + self._zero_fill_owner: int | None = None + + @classmethod + def maybe_build( + cls, + vllm_config: VllmConfig, + config: DeepseekV2Config | DeepseekV3Config, + device: torch.device, + ) -> "IndexerQKFusionBuffers | None": + # Every indexer rope is a RotaryEmbeddingBase subclass, so custom-op + # enablement stands in for the per-layer is_inplace_rope. + if not Indexer.can_use_aiter_qk_fusion( + vllm_config, config, RotaryEmbeddingBase.enabled() + ): + return None + return cls( + vllm_config.scheduler_config.max_num_batched_tokens, + config.index_n_heads, + config.index_head_dim, + device, + ) + + def claim_zero_fill(self, layer_index: int) -> bool: + """True only for the first caller, which owns the per-pass zero-fill. + + The claimer must also be the first indexer to run under this root, or + the rows the kernel skips would keep the previous pass's values. Layers + that skip the indexer build none (deepseek_v2.py, `_skip_topk`), so + construction order is execution order; assert it, because the runtime + toggle at mla.py `not self.skip_topk` could break that in the future. + """ + if self._zero_fill_owner is None: + self._zero_fill_owner = layer_index + return True + assert layer_index > self._zero_fill_owner, ( + f"indexer at layer {layer_index} was built after the zero-fill " + f"owner at layer {self._zero_fill_owner}, so the owner may not run " + "first" + ) + return False + + class Indexer(nn.Module): def __init__( self, @@ -651,6 +725,7 @@ def __init__( topk_indices_buffer: torch.Tensor | None, prefix: str = "", is_inplace_rope: bool = False, + indexer_qk_fusion_buffers: IndexerQKFusionBuffers | None = None, ): super().__init__() self.vllm_config = vllm_config @@ -702,6 +777,16 @@ def __init__( from vllm.v1.attention.backends.mla.indexer import get_max_prefill_buffer_size self.max_total_seq_len = get_max_prefill_buffer_size(vllm_config) + # The fused path writes the K cache itself; one predicate decides both + # the writer and the skip so they cannot disagree. + self.use_aiter_qk_fusion = self.can_use_aiter_qk_fusion( + vllm_config, config, is_inplace_rope + ) + if self.use_aiter_qk_fusion: + logger.info_once( + "Fusing the DSA indexer QK pre-processing into " + "aiter.indexer_qk_rope_quant_and_cache" + ) self.indexer_op = SparseAttnIndexer( self.k_cache, self.quant_block_size, @@ -711,6 +796,16 @@ def __init__( self.max_model_len, self.max_total_seq_len, self.topk_indices_buffer, + skip_k_cache_insert=self.use_aiter_qk_fusion, + ) + self.indexer_qk_fusion_buffers = ( + indexer_qk_fusion_buffers if self.use_aiter_qk_fusion else None + ) + self.owns_qk_fusion_zero_fill = ( + self.indexer_qk_fusion_buffers is not None + and self.indexer_qk_fusion_buffers.claim_zero_fill( + extract_layer_index(prefix) + ) ) self.is_inplace_rope = is_inplace_rope @@ -723,13 +818,94 @@ def __init__( and self.scale_fmt is not None ) + @staticmethod + def can_use_aiter_qk_fusion( + vllm_config: VllmConfig, + config: DeepseekV2Config | DeepseekV3Config, + is_inplace_rope: bool, + ) -> bool: + """Whether the aiter fused indexer QK kernel serves this model. + + Replaces the in-place-rope path, on the DSA shapes the kernel is built + for (head_dim 128, rope_dim 64) and on the gfx942/gfx950 parts aiter + ships its CK build for. + + Context parallel must stay unfused: slot_mapping is PAD_SLOT_ID on + ranks that do not own a token, so the kernel would skip the row and + never write its query, which every rank needs to score its KV shard. + """ + if not current_platform.is_rocm(): + return False + from vllm.platforms.rocm import on_mi3xx + + parallel_config = vllm_config.parallel_config + return bool( + rocm_aiter_ops.is_indexer_qk_fusion_enabled() + and on_mi3xx() + and is_inplace_rope + and getattr(config, "index_head_dim", None) == 128 + and getattr(config, "qk_rope_head_dim", None) == 64 + and parallel_config.decode_context_parallel_size == 1 + and parallel_config.prefill_context_parallel_size == 1 + ) + def forward( self, hidden_states: torch.Tensor, qr: torch.Tensor, positions, rotary_emb ) -> torch.Tensor: q, _ = self.wq_b(qr) q = q.view(-1, self.n_head, self.head_dim) - if current_platform.is_rocm() and self.is_inplace_rope: + if self.use_aiter_qk_fusion: + # One GEMM, then split; the kernel reads both through their strides. + kw, _ = self.wk_weights_proj(hidden_states) + k = kw[:, : self.head_dim] + weights = kw[:, self.head_dim :] + + kv_cache = self.k_cache.kv_cache + cos, sin = get_indexer_rope_halves(rotary_emb, q.dtype) + + # Model-shared pair when one was handed down, a fresh zeroed pair + # otherwise; an init-time constant, so each layer traces one branch. + total = q.shape[0] + if self.indexer_qk_fusion_buffers is not None: + q_fp8 = self.indexer_qk_fusion_buffers.q_fp8 + weights_out = self.indexer_qk_fusion_buffers.weights + zero_outputs = self.owns_qk_fusion_zero_fill + else: + q_fp8 = torch.zeros( + (total, self.n_head, self.head_dim), + dtype=current_platform.fp8_dtype(), + device=q.device, + ) + weights_out = torch.zeros( + (total, self.n_head), dtype=torch.float32, device=q.device + ) + zero_outputs = False + + torch.ops.vllm.rocm_aiter_indexer_qk_rope_quant_and_cache( + _encode_layer_name(self.k_cache.prefix), + kv_cache, + q, + k, + weights, + positions, + cos, + sin, + # The aiter kernel reads the norm params in q's dtype. + self.k_norm.weight.to(q.dtype), + self.k_norm.bias.to(q.dtype), + q_fp8, + weights_out, + float(self.k_norm.eps), + self.quant_block_size, + self.scale_fmt, + float(self.softmax_scale * self.n_head_scale), + zero_outputs, + rotary_emb.is_neox_style, + ) + # K cache already written; indexer_op skips its insert, `k` is unused. + return self.indexer_op(hidden_states, q_fp8[:total], k, weights_out[:total]) + elif current_platform.is_rocm() and self.is_inplace_rope: # This path should works on all platform, will remove extra # branches in the future # This fast path relies on rotary_emb mutating q and k inplace. @@ -979,6 +1155,7 @@ def __init__( quant_config: QuantizationConfig | None = None, prefix: str = "", topk_indices_buffer: torch.Tensor | None = None, + indexer_qk_fusion_buffers: IndexerQKFusionBuffers | None = None, input_size: int | None = None, reduce_results: bool = True, non_causal_multi_token_decode: bool = False, @@ -1137,7 +1314,10 @@ def __init__( topk_indices_buffer, f"{prefix}.indexer", is_inplace_rope=self.indexer_rope_emb.enabled(), + indexer_qk_fusion_buffers=indexer_qk_fusion_buffers, ) + if self.indexer.use_aiter_qk_fusion: + register_indexer_rope_halves(self.indexer_rope_emb) else: self.indexer_rope_emb = None self.indexer = None @@ -1203,6 +1383,7 @@ def __init__( prefix: str, config: DeepseekV2Config | None = None, topk_indices_buffer: torch.Tensor | None = None, + indexer_qk_fusion_buffers: IndexerQKFusionBuffers | None = None, ) -> None: super().__init__() @@ -1266,6 +1447,7 @@ def __init__( quant_config=quant_config, prefix=f"{prefix}.self_attn", topk_indices_buffer=topk_indices_buffer, + indexer_qk_fusion_buffers=indexer_qk_fusion_buffers, reduce_results=not self.use_sequence_parallel_moe, ) @@ -1388,8 +1570,12 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): dtype=torch.int32, device=self.device, ) + indexer_qk_fusion_buffers = IndexerQKFusionBuffers.maybe_build( + vllm_config, config, self.device + ) else: topk_indices_buffer = None + indexer_qk_fusion_buffers = None if get_pp_group().is_first_rank: self.embed_tokens = VocabParallelEmbedding( @@ -1406,6 +1592,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): vllm_config=vllm_config, prefix=prefix, topk_indices_buffer=topk_indices_buffer, + indexer_qk_fusion_buffers=indexer_qk_fusion_buffers, ), prefix=f"{prefix}.layers", ) diff --git a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py index 0f55556c5dfb..8787fc208aed 100644 --- a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py @@ -907,6 +907,142 @@ def rocm_aiter_sparse_attn_indexer( return topk_indices_buffer +# ``cos_sin_cache`` is ``cat((cos, sin), dim=-1)`` and constant after init, but +# the kernel wants the halves separately. Splitting per forward would rewrite the +# whole table once per layer per step; these views cost nothing, since the kernel +# indexes by row stride and only needs the last dim contiguous. +_INDEXER_ROPE_COS_ATTR = "aiter_indexer_rope_cos" +_INDEXER_ROPE_SIN_ATTR = "aiter_indexer_rope_sin" + + +def register_indexer_rope_halves(rope: torch.nn.Module) -> None: + cache = getattr(rope, "cos_sin_cache", None) + if not isinstance(cache, torch.Tensor) or cache.ndim != 2: + return + half = cache.shape[-1] // 2 + rope.register_buffer(_INDEXER_ROPE_COS_ATTR, cache[:, :half], persistent=False) + rope.register_buffer(_INDEXER_ROPE_SIN_ATTR, cache[:, half:], persistent=False) + + +def get_indexer_rope_halves( + rope: torch.nn.Module, dtype: torch.dtype +) -> tuple[torch.Tensor, torch.Tensor]: + """cos/sin halves in ``dtype``, from the buffers registered at init. + + The fallback covers a cache rebuilt or re-typed after registration. + """ + cos = getattr(rope, _INDEXER_ROPE_COS_ATTR, None) + sin = getattr(rope, _INDEXER_ROPE_SIN_ATTR, None) + if cos is not None and sin is not None and cos.dtype == dtype: + return cos, sin + cache = rope.cos_sin_cache.to(dtype) + half = cache.shape[-1] // 2 + return cache[:, :half], cache[:, half:] + + +def rocm_aiter_indexer_qk_rope_quant_and_cache_fake( + k_cache_prefix: LayerNameType, + kv_cache: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + weights: torch.Tensor, + positions: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + norm_weight: torch.Tensor, + norm_bias: torch.Tensor, + q_fp8_out: torch.Tensor, + weights_out: torch.Tensor, + epsilon: float, + quant_block_size: int, + scale_fmt: str, + weights_scale: float, + zero_outputs: bool, + is_neox: bool, +) -> None: + return None + + +def rocm_aiter_indexer_qk_rope_quant_and_cache( + k_cache_prefix: LayerNameType, + kv_cache: torch.Tensor, + q: torch.Tensor, # [num_tokens, n_heads, head_dim], raw post-wq_b + k: torch.Tensor, # [num_tokens, head_dim], raw pre-norm/rope + weights: torch.Tensor, # [num_tokens, n_heads], raw post-weights_proj + positions: torch.Tensor, # [num_tokens] + cos: torch.Tensor, # [max_position, rope_dim // 2] + sin: torch.Tensor, # [max_position, rope_dim // 2] + norm_weight: torch.Tensor, # [head_dim], q dtype + norm_bias: torch.Tensor, # [head_dim], q dtype + q_fp8_out: torch.Tensor, # [>= num_tokens, n_heads, head_dim] fp8, written + weights_out: torch.Tensor, # [>= num_tokens, n_heads] fp32, written + epsilon: float, + quant_block_size: int, + scale_fmt: str, + weights_scale: float, + zero_outputs: bool, + is_neox: bool, +) -> None: + """Run aiter's fused DSA indexer QK kernel into the caller's buffers. + + One launch covers RoPE on q and k, LayerNorm on k, per-group FP8 + quantization of both, the q scale folded into ``weights_out``, and the + paged indexer K-cache write. + """ + from aiter import indexer_qk_rope_quant_and_cache + + attn_metadata = get_forward_context().attn_metadata + # Profiling / dummy run: no slot_mapping, so nothing to write. The caller's + # buffers are already shaped, which is all tracing and profiling need. + if not isinstance(attn_metadata, dict): + return + from vllm.utils.torch_utils import _resolve_layer_name + + k_cache_prefix = _resolve_layer_name(k_cache_prefix) + layer_attn_metadata = attn_metadata[k_cache_prefix] + assert isinstance(layer_attn_metadata, DeepseekV32IndexerMetadata) + slot_mapping = layer_attn_metadata.slot_mapping + num_tokens = slot_mapping.shape[0] + + # Outputs cover one row per input token, CUDA-graph padding included: decode + # reads weights[:batch_size * next_n], which can exceed num_tokens. Rows the + # kernel skips must read as zero, so the padded logits they feed stay finite. + # Slice here, not in the caller, to keep the mutated args base tensors. + total = q.shape[0] + # The kernel indexes q/k/weights by slot_mapping row, so extra slots would + # read past their ends. + assert num_tokens <= total, ( + f"slot_mapping has {num_tokens} rows but only {total} query rows" + ) + q_fp8_out = q_fp8_out[:total] + weights_out = weights_out[:total] + if zero_outputs: + q_fp8_out.zero_() + weights_out.zero_() + indexer_qk_rope_quant_and_cache( + q[:num_tokens], + q_fp8_out[:num_tokens], + weights[:num_tokens], + weights_out[:num_tokens], + k[:num_tokens], + kv_cache, + slot_mapping, + norm_weight, + norm_bias, + positions[:num_tokens], + cos, + sin, + epsilon, + quant_block_size, + scale_fmt, + weights_scale, + # Same layout predicate the readers use: tiled in-block for + # block_size > 1, flat otherwise. + preshuffle=kv_cache.shape[1] > 1, + is_neox=is_neox, + ) + + def _decode_e8m0_scales(scale: torch.Tensor) -> torch.Tensor: if scale.dtype == torch.float8_e8m0fnu: from vllm.model_executor.layers.quantization.utils.fp8_utils import (