Skip to content

[FlyDSL] Add FP4 LiteTopK prefill operator #5309

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
AMD-yanfeiwang wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

[FlyDSL] Add FP4 LiteTopK prefill operator #5309

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
40 changes: 40 additions & 0 deletions aiter/ops/flydsl/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,18 @@
)

_LAZY_IMPORTS = {
"FP4_LITETOPK_SUPPORTED_TOPKS": (
".pa_mqa_litetopk_fp4",
"FP4_LITETOPK_SUPPORTED_TOPKS",
),
"FP4LiteTopKResult": (
".pa_mqa_litetopk_fp4",
"FP4LiteTopKResult",
),
"FP4LiteTopKWorkspace": (
".pa_mqa_litetopk_fp4",
"FP4LiteTopKWorkspace",
),
"FP8_MQA_LOGITS_DEFAULT_VARIANT": (
".kernels.mqa_logits.fp8_mqa_logits",
"DEFAULT_VARIANT",
Expand All @@ -37,6 +49,26 @@
".kernels.mqa_logits.pa_mqa_logits_fp4_prefill",
"compute_varqlen_windows",
),
"allocate_fp4_litetopk_workspace": (
".pa_mqa_litetopk_fp4",
"allocate_fp4_litetopk_workspace",
),
"fp4_litetopk_workspace_nbytes": (
".pa_mqa_litetopk_fp4",
"fp4_litetopk_workspace_nbytes",
),
"fp4_litetopk_workspace_size": (
".pa_mqa_litetopk_fp4",
"fp4_litetopk_workspace_size",
),
"flydsl_pa_mqa_litetopk_fp4_prefill": (
".pa_mqa_litetopk_fp4",
"flydsl_pa_mqa_litetopk_fp4_prefill",
),
"prepare_fp4_litetopk_seed": (
".pa_mqa_litetopk_fp4",
"prepare_fp4_litetopk_seed",
),
"flydsl_flash_attn_func": (".fmha_kernels", "flydsl_flash_attn_func"),
"flydsl_fp8_mqa_logits": (
".kernels.mqa_logits.fp8_mqa_logits",
Expand Down Expand Up @@ -69,21 +101,29 @@
}

__all__ = [
"FP4_LITETOPK_SUPPORTED_TOPKS",
"FP8_MQA_LOGITS_DEFAULT_VARIANT",
"FP8_MQA_LOGITS_VARIANTS",
"FP4LiteTopKResult",
"FP4LiteTopKWorkspace",
"GateMode",
"allocate_fp4_litetopk_workspace",
"compute_varqlen_windows",
"flydsl_flash_attn_func",
"flydsl_fp8_mqa_logits",
"flydsl_hgemm",
"flydsl_mla_reduce_v1",
"flydsl_moe_stage1",
"flydsl_moe_stage2",
"flydsl_pa_mqa_litetopk_fp4_prefill",
"flydsl_pa_mqa_logits_fp4",
"flydsl_pa_mqa_logits_fp4_prefill",
"flydsl_pa_mqa_logits_fp4_varqlen",
"flydsl_preshuffle_gemm_a8",
"flydsl_qk_norm_rope_quant",
"fp4_litetopk_workspace_nbytes",
"fp4_litetopk_workspace_size",
"prepare_fp4_litetopk_seed",
]


Expand Down
27 changes: 27 additions & 0 deletions aiter/ops/flydsl/kernels/mqa_logits/pa_mqa_litetopk_fp4_common.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,27 @@
# SPDX-License-Identifier: MIT
# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.

import flydsl.expr as fx
from flydsl.expr import Array, Int32

FP4_LITETOPK_SUPPORTED_TOPKS = (512, 1024)


@fx.struct
class LiteTopKScanStorage:
histogram: Array[Int32, 256, 16]
scan: Array[Int32, 9, 16]
state: Array[Int32, 3, 4]


@fx.struct
class LiteTopKSeedStorage:
scores: Array[fx.Float32, 8192, 16]
histogram: Array[Int32, 256, 16]
scan: Array[Int32, 9, 16]
maxima: Array[fx.Float32, 8, 16]
neg_minima: Array[fx.Float32, 8, 16]
finite_counts: Array[Int32, 8, 16]
nonfinite_counts: Array[Int32, 8, 16]
calibration: Array[fx.Float32, 2, 16]
state: Array[Int32, 2, 4]
6 changes: 5 additions & 1 deletion aiter/ops/flydsl/kernels/mqa_logits/pa_mqa_logits_fp4_common.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,14 +15,18 @@
_NON_WRITER_LANE_OFF = 1 << 30


def _i32_buffer(ptr, width=1):
def _i32_buffer(ptr, width=1, elem_offset=None):
"""OOB-checked global i32 buffer-tensor over ``ptr`` (mirrors a max_size V#).

``width`` shapes it ``(N, width)`` so a per-``width`` row can be sliced and
vector-copied; ``width=1`` gives a flat tensor for scalar ``[idx]`` loads.
``elem_offset`` is folded into the 64-bit base pointer before creating the
descriptor, keeping subsequent V# offsets below the 32-bit hardware limit.
"""
src = fx.get_iter(ptr)
it = fx.recast_iter(fx.PointerType.get(T.i32, src.memspace, 4), src)
if elem_offset is not None:
it = fx.add_offset(it, fx.Int64(elem_offset))
if width == 1:
lay = fx.make_layout((1 << 30,), (1,))
else:
Expand Down
Loading
Loading