Skip to content

Fp8 block scale quantization for fmha fwd #3330

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

Open
wants to merge 15 commits into
base: develop
Choose a base branch
from
+618 −83
Open

Fp8 block scale quantization for fmha fwd #3330

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
15 commits
Select commit Hold shift + click to select a range
6862777
add block scale parameters to kernel
ltqin Nov 28, 2025
7260af6
add block scale to kernel
ltqin Nov 30, 2025
39104a0
add smoke test
ltqin Dec 1, 2025
356c3c9
format
ltqin Dec 1, 2025
42b5aa4
Revert "format"
ltqin Dec 1, 2025
2954847
only format my code
ltqin Dec 1, 2025
49a280b
format py
ltqin Dec 1, 2025
fa90ef7
Merge branch 'develop' into ck_tile/fmha_fwd_block_scale
ltqin Dec 1, 2025
326d8c3
fix auto not allowd in function prototype
ltqin Dec 1, 2025
43feee2
Merge branch 'develop' into ck_tile/fmha_fwd_block_scale
ltqin Dec 1, 2025
4484de9
change instance tttt to ttff
ltqin Dec 2, 2025
2709f16
fix structured binding issue
ltqin Dec 2, 2025
259fd92
change s_acc elementwise op
ltqin Dec 3, 2025
0d462a2
async pipeline add block scale
ltqin Dec 3, 2025
01b232c
Merge branch 'develop' into ck_tile/fmha_fwd_block_scale
ltqin Dec 4, 2025
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
2 changes: 2 additions & 0 deletions example/ck_tile/01_fmha/codegen/cpp_symbol_map.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -65,11 +65,13 @@ def get_mask_check_map(mask: str):
QSCALE_MAP = {
"no": "ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE",
"pertensor": "ck_tile::BlockAttentionQuantScaleEnum::PERTENSOR",
"blockscale": "ck_tile::BlockAttentionQuantScaleEnum::BLOCKSCALE",
}

QSCALE_CHECK_MAP = {
"no": "quant_scale_enum::no_scale",
"pertensor": "quant_scale_enum::pertensor",
"blockscale": "quant_scale_enum::blockscale",
}

BIAS_MAP = {
Expand Down
7 changes: 5 additions & 2 deletions example/ck_tile/01_fmha/codegen/ops/fmha_fwd.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -747,7 +747,7 @@ def get_pipelines(dtype, hdim, hdim_v, receipt, mask_impl) -> List[FmhaFwdPipeli
# no need lse/dropout kernels
for logits, qscale, mask, bias in itertools.product(
["f"],
["no", "pertensor"],
["no", "pertensor", "blockscale"],
get_mask_map(mask_impl).keys(),
["no"],
):
Expand Down Expand Up @@ -838,7 +838,10 @@ def get_pipelines(dtype, hdim, hdim_v, receipt, mask_impl) -> List[FmhaFwdPipeli
elif dtype in ["fp8", "fp8bf16", "fp8fp32"]:
# no need lse/dropout kernels
for logits, qscale, mask, bias in itertools.product(
["f"], ["no", "pertensor"], get_mask_map(mask_impl).keys(), ["no"]
["f"],
["no", "pertensor", "blockscale"],
get_mask_map(mask_impl).keys(),
["no"],
):
pipelines.append(FmhaFwdPipeline("qr", "row", "f", "f", "f", "f", logits, bias, "f", "f", qscale, mask, "f", "f")) # fmt: skip
pipelines.append(FmhaFwdPipeline("qr", "row", "t", "t", "t", "t", logits, bias, "f", "f", qscale, mask, "f", "f")) # fmt: skip
Expand Down
26 changes: 26 additions & 0 deletions example/ck_tile/01_fmha/fmha_fwd.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -230,6 +230,8 @@ struct fmha_fwd_args
// array [batch + 1]. (Used with padding)
const void* cu_seqlen_k_ptr = nullptr; // Cumulative logical (excluding padding) sequence length
// array [batch + 1]. (Used with padding)
const void* bseqstart_q_ptr;
const void* bseqstart_k_ptr;

ck_tile::index_t seqlen_q;
ck_tile::index_t seqlen_k;
Expand All @@ -256,13 +258,19 @@ struct fmha_fwd_args
ck_tile::index_t nhead_stride_randval;
ck_tile::index_t nhead_stride_lse;
ck_tile::index_t nhead_stride_o;
ck_tile::index_t nhead_stride_q_descale;
ck_tile::index_t nhead_stride_k_descale;
ck_tile::index_t nhead_stride_v_descale;
ck_tile::index_t batch_stride_q;
ck_tile::index_t batch_stride_k;
ck_tile::index_t batch_stride_v;
ck_tile::index_t batch_stride_bias;
ck_tile::index_t batch_stride_randval;
ck_tile::index_t batch_stride_lse;
ck_tile::index_t batch_stride_o;
ck_tile::index_t batch_stride_q_descale;
ck_tile::index_t batch_stride_k_descale;
ck_tile::index_t batch_stride_v_descale;

ck_tile::index_t window_size_left;
ck_tile::index_t window_size_right;
Expand All @@ -274,6 +282,9 @@ struct fmha_fwd_args

std::variant<std::pair<uint64_t, uint64_t>, std::pair<const void*, const void*>>
drop_seed_offset;

ck_tile::index_t block_scale_m;
ck_tile::index_t block_scale_n;
};

struct fmha_fwd_pagedkv_args
Expand Down Expand Up @@ -592,6 +603,8 @@ auto fmha_fwd_create_kargs_and_grids(fmha_fwd_args args)
args.seqstart_k_ptr,
args.seqlen_q_ptr,
args.seqlen_k_ptr,
args.bseqstart_q_ptr,
args.bseqstart_k_ptr,
args.hdim_q,
args.hdim_v,
args.nhead_q,
Expand All @@ -611,13 +624,18 @@ auto fmha_fwd_create_kargs_and_grids(fmha_fwd_args args)
args.nhead_stride_randval,
args.nhead_stride_lse,
args.nhead_stride_o,
args.nhead_stride_q_descale,
args.nhead_stride_k_descale,
args.nhead_stride_v_descale,
args.window_size_left,
args.window_size_right,
args.mask_type,
args.min_seqlen_q,
args.p_drop,
args.s_randval,
args.drop_seed_offset,
args.block_scale_m,
args.block_scale_n,
args.cu_seqlen_q_ptr,
args.cu_seqlen_k_ptr);
}
Expand Down Expand Up @@ -654,19 +672,27 @@ auto fmha_fwd_create_kargs_and_grids(fmha_fwd_args args)
args.nhead_stride_randval,
args.nhead_stride_lse,
args.nhead_stride_o,
args.nhead_stride_q_descale,
args.nhead_stride_k_descale,
args.nhead_stride_v_descale,
args.batch_stride_q,
args.batch_stride_k,
args.batch_stride_v,
args.batch_stride_bias,
args.batch_stride_randval,
args.batch_stride_lse,
args.batch_stride_o,
args.batch_stride_q_descale,
args.batch_stride_k_descale,
args.batch_stride_v_descale,
args.window_size_left,
args.window_size_right,
args.mask_type,
args.p_drop,
args.s_randval,
args.drop_seed_offset,
args.block_scale_m,
args.block_scale_n,
args.cu_seqlen_q_ptr,
args.cu_seqlen_k_ptr);
}
Expand Down
Loading