From 0428f927d7d42594ad68115330b4a3c6b8adc785 Mon Sep 17 00:00:00 2001 From: jiqing-feng Date: Wed, 15 Jul 2026 10:06:06 +0000 Subject: [PATCH 01/12] add cri Signed-off-by: jiqing-feng --- .../flash_attn_xpu/src/collective/fmha_fwd_common.hpp | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/flash-attn2/flash_attn_xpu/src/collective/fmha_fwd_common.hpp b/flash-attn2/flash_attn_xpu/src/collective/fmha_fwd_common.hpp index 0f7967478..961c35c80 100644 --- a/flash-attn2/flash_attn_xpu/src/collective/fmha_fwd_common.hpp +++ b/flash-attn2/flash_attn_xpu/src/collective/fmha_fwd_common.hpp @@ -316,13 +316,16 @@ struct FMHAFwdEpilogueTraits { copy_block_s2r(sA_max(_, k_blk, kr, a_tile), rA_kmax[kr]); } + using RowBase = typename ReduceFragARow::Base; rA_max = rA_kmax[0]; for (int kr = 1; kr < ReduceK{}; kr++) - cute::transform(rA_max, rA_kmax[kr], rA_max, cute::max_fn{}); + cute::transform(rA_max.tensor(), rA_kmax[kr].tensor(), + static_cast(rA_max), cute::max_fn{}); for (int kr = 0; kr < ReduceK{}; kr++) { cute::transform( - rA_max, rA_kmax[kr], rA_kmax[kr], [](auto gmax, auto kmax) { + rA_max.tensor(), rA_kmax[kr].tensor(), + static_cast(rA_kmax[kr]), [](auto gmax, auto kmax) { return sycl::native::exp2(kmax - gmax); }); } From 580ba9db89d545e0643646618928c1fb649487ae Mon Sep 17 00:00:00 2001 From: jiqing-feng Date: Fri, 17 Jul 2026 10:43:01 +0800 Subject: [PATCH 02/12] flash-attn2 xpu: add MXFP (block-scaled FP8/FP4) forward support (#2) * flash-attn2 xpu: add MXFP (block-scaled FP8/FP4) forward support Port the CRI MXFP path (e5m2/e4m3/e2m1) into the fused BMG+CRI build: - fmha_fwd_types/utils/fusion: scale args, dtype mapping, cumulative scale length - mainloop/kernel/FMHAConfig: thread scale tensors; block-scaled GEMM guarded by SYCL_INTEL_TARGET==35 with compilable BMG fallbacks - fmha_fwd dispatcher: unconditional host-side MXFP dispatch, gated to non-paged / non-rotary varlen+fix paths via if constexpr - flash_api / torch_binding: optional scale_q/k/v (+ cu_scale_q/kv) params - 16 explicit-instantiation TUs for hdim 32..512 x {varlen,fix} * flash-attn2: bump version to 3 for MXFP public API change * flash-attn2: validate V dtype per MXFP path (fp8 V=fp8, fp4 V=bf16) --- flash-attn2/build.toml | 18 +- flash-attn2/flash_attn_xpu/flash_api.cpp | 181 ++++++++++++--- .../src/collective/fmha_fusion.hpp | 1 + .../src/collective/fmha_fwd_mainloop_xe2.hpp | 216 +++++++++++++++++- .../src/flash_fwd_hdim128_fix_mxfp.cpp | 11 + .../src/flash_fwd_hdim128_varlen_mxfp.cpp | 11 + .../src/flash_fwd_hdim160_fix_mxfp.cpp | 11 + .../src/flash_fwd_hdim160_varlen_mxfp.cpp | 11 + .../src/flash_fwd_hdim192_fix_mxfp.cpp | 11 + .../src/flash_fwd_hdim192_varlen_mxfp.cpp | 11 + .../src/flash_fwd_hdim256_fix_mxfp.cpp | 11 + .../src/flash_fwd_hdim256_varlen_mxfp.cpp | 11 + .../src/flash_fwd_hdim32_fix_mxfp.cpp | 11 + .../src/flash_fwd_hdim32_varlen_mxfp.cpp | 11 + .../src/flash_fwd_hdim512_fix_mxfp.cpp | 11 + .../src/flash_fwd_hdim512_varlen_mxfp.cpp | 11 + .../src/flash_fwd_hdim64_fix_mxfp.cpp | 11 + .../src/flash_fwd_hdim64_varlen_mxfp.cpp | 11 + .../src/flash_fwd_hdim96_fix_mxfp.cpp | 11 + .../src/flash_fwd_hdim96_varlen_mxfp.cpp | 11 + flash-attn2/flash_attn_xpu/src/fmha_fwd.cpp | 24 +- flash-attn2/flash_attn_xpu/src/fmha_fwd.hpp | 88 ++++++- .../flash_attn_xpu/src/fmha_fwd_impl.hpp | 208 ++++++++++++++++- .../flash_attn_xpu/src/fmha_fwd_types.hpp | 17 +- flash-attn2/flash_attn_xpu/src/fmha_utils.hpp | 19 +- .../src/kernel/fmha_fwd_kernel_xe2.hpp | 133 +++++++++-- flash-attn2/torch-ext/torch_binding.cpp | 12 +- flash-attn2/torch-ext/torch_binding.h | 12 +- 28 files changed, 1018 insertions(+), 87 deletions(-) create mode 100644 flash-attn2/flash_attn_xpu/src/flash_fwd_hdim128_fix_mxfp.cpp create mode 100644 flash-attn2/flash_attn_xpu/src/flash_fwd_hdim128_varlen_mxfp.cpp create mode 100644 flash-attn2/flash_attn_xpu/src/flash_fwd_hdim160_fix_mxfp.cpp create mode 100644 flash-attn2/flash_attn_xpu/src/flash_fwd_hdim160_varlen_mxfp.cpp create mode 100644 flash-attn2/flash_attn_xpu/src/flash_fwd_hdim192_fix_mxfp.cpp create mode 100644 flash-attn2/flash_attn_xpu/src/flash_fwd_hdim192_varlen_mxfp.cpp create mode 100644 flash-attn2/flash_attn_xpu/src/flash_fwd_hdim256_fix_mxfp.cpp create mode 100644 flash-attn2/flash_attn_xpu/src/flash_fwd_hdim256_varlen_mxfp.cpp create mode 100644 flash-attn2/flash_attn_xpu/src/flash_fwd_hdim32_fix_mxfp.cpp create mode 100644 flash-attn2/flash_attn_xpu/src/flash_fwd_hdim32_varlen_mxfp.cpp create mode 100644 flash-attn2/flash_attn_xpu/src/flash_fwd_hdim512_fix_mxfp.cpp create mode 100644 flash-attn2/flash_attn_xpu/src/flash_fwd_hdim512_varlen_mxfp.cpp create mode 100644 flash-attn2/flash_attn_xpu/src/flash_fwd_hdim64_fix_mxfp.cpp create mode 100644 flash-attn2/flash_attn_xpu/src/flash_fwd_hdim64_varlen_mxfp.cpp create mode 100644 flash-attn2/flash_attn_xpu/src/flash_fwd_hdim96_fix_mxfp.cpp create mode 100644 flash-attn2/flash_attn_xpu/src/flash_fwd_hdim96_varlen_mxfp.cpp diff --git a/flash-attn2/build.toml b/flash-attn2/build.toml index 040e367fc..d93fae573 100644 --- a/flash-attn2/build.toml +++ b/flash-attn2/build.toml @@ -1,6 +1,6 @@ [general] name = "flash-attn2" -version = 2 + version = 3 license = "BSD-3-Clause" backends = [ "cpu", @@ -78,6 +78,22 @@ src = [ "flash_attn_xpu/src/flash_fwd_hdim256_kvcache_paged_bf16.cpp", "flash_attn_xpu/src/flash_fwd_hdim512_kvcache_paged_fp16.cpp", "flash_attn_xpu/src/flash_fwd_hdim512_kvcache_paged_bf16.cpp", + "flash_attn_xpu/src/flash_fwd_hdim32_varlen_mxfp.cpp", + "flash_attn_xpu/src/flash_fwd_hdim64_varlen_mxfp.cpp", + "flash_attn_xpu/src/flash_fwd_hdim96_varlen_mxfp.cpp", + "flash_attn_xpu/src/flash_fwd_hdim128_varlen_mxfp.cpp", + "flash_attn_xpu/src/flash_fwd_hdim160_varlen_mxfp.cpp", + "flash_attn_xpu/src/flash_fwd_hdim192_varlen_mxfp.cpp", + "flash_attn_xpu/src/flash_fwd_hdim256_varlen_mxfp.cpp", + "flash_attn_xpu/src/flash_fwd_hdim512_varlen_mxfp.cpp", + "flash_attn_xpu/src/flash_fwd_hdim32_fix_mxfp.cpp", + "flash_attn_xpu/src/flash_fwd_hdim64_fix_mxfp.cpp", + "flash_attn_xpu/src/flash_fwd_hdim96_fix_mxfp.cpp", + "flash_attn_xpu/src/flash_fwd_hdim128_fix_mxfp.cpp", + "flash_attn_xpu/src/flash_fwd_hdim160_fix_mxfp.cpp", + "flash_attn_xpu/src/flash_fwd_hdim192_fix_mxfp.cpp", + "flash_attn_xpu/src/flash_fwd_hdim256_fix_mxfp.cpp", + "flash_attn_xpu/src/flash_fwd_hdim512_fix_mxfp.cpp", "flash_attn_xpu/src/fmha_bwd_types.hpp", "flash_attn_xpu/src/fmha_bwd.hpp", "flash_attn_xpu/src/fmha_bwd_impl.hpp", diff --git a/flash-attn2/flash_attn_xpu/flash_api.cpp b/flash-attn2/flash_attn_xpu/flash_api.cpp index 8701dcc54..65b5ff092 100644 --- a/flash-attn2/flash_attn_xpu/flash_api.cpp +++ b/flash-attn2/flash_attn_xpu/flash_api.cpp @@ -47,7 +47,10 @@ mha_fwd( int window_size_right, const float softcap, const bool return_softmax, - std::optional gen_) { + std::optional gen_, + std::optional scale_q_ = std::nullopt, + std::optional scale_k_ = std::nullopt, + std::optional scale_v_ = std::nullopt) { auto device_idx = q.device().index(); compat::select_device(device_idx); @@ -65,22 +68,39 @@ mha_fwd( const int num_heads_k = k.size(2); auto q_dtype = q.dtype(); - TORCH_CHECK(q_dtype == torch::kFloat16 || q_dtype == torch::kBFloat16, - "FlashAttention only support fp16 and bf16 data type"); + const bool is_mxfp = (q_dtype == torch::kFloat8_e5m2 || q_dtype == torch::kFloat8_e4m3fn || + q_dtype == torch::kFloat4_e2m1fn_x2); + TORCH_CHECK(q_dtype == torch::kFloat16 || q_dtype == torch::kBFloat16 || is_mxfp, + "FlashAttention only support fp16, bf16, fp8_e5m2, fp8_e4m3fn and fp4_e2m1fn_x2 data type"); TORCH_CHECK(k.dtype() == q_dtype, "query and key must have the same dtype"); - TORCH_CHECK(v.dtype() == q_dtype, "query and value must have the same dtype"); TORCH_CHECK(k.size(-1) == head_size_og, "Key head dimension must match Query head dimension"); - TORCH_CHECK(v.size(-1) == head_size_og, "Value head dimension must match Query head dimension"); - TORCH_CHECK(head_size_og <= 256 || head_size_og == 512, - "FlashAttention XPU only supports head dimension up to 256 or exactly 512. Got: " + std::to_string(head_size_og)); + if (!is_mxfp) { + TORCH_CHECK(v.dtype() == q_dtype, "query and value must have the same dtype"); + TORCH_CHECK(v.size(-1) == head_size_og, "Value head dimension must match Query head dimension"); + } else if (q_dtype == torch::kFloat4_e2m1fn_x2) { + // MXFP4 (e2m1) consumes a BF16 value tensor; reject mismatches early + // instead of silently producing wrong results inside the kernel. + TORCH_CHECK(v.dtype() == torch::kBFloat16, + "For fp4_e2m1fn_x2 query, value must be bf16"); + } else { + // MXFP8 (e5m2 / e4m3fn) requires the value tensor to match the query fp8 dtype. + TORCH_CHECK(v.dtype() == q_dtype, + "For fp8 query, value must have the same fp8 dtype as query"); + } + // For FP4 (e2m1) the head dimension is packed (2 values per byte); unpack for padding math. + const int head_size_actual = (q_dtype == torch::kFloat4_e2m1fn_x2) ? head_size_og * 2 : head_size_og; + TORCH_CHECK(head_size_actual <= 256 || head_size_actual == 512, + "FlashAttention XPU only supports head dimension up to 256 or exactly 512. Got: " + std::to_string(head_size_actual)); CHECK_DEVICE(q); CHECK_DEVICE(k); CHECK_DEVICE(v); - // XPU requires head_size to be a multiple of 32 - const int head_size_padded = round_multiple(head_size_og, 32); - const bool needs_padding = (head_size_og != head_size_padded); + // XPU requires head_size to be a multiple of 32. + // FP4 packed inputs are consumed in their packed layout and are not padded. + const int head_size_padded = round_multiple(head_size_actual, 32); + const bool needs_padding = (head_size_actual != head_size_padded) + && (q_dtype != torch::kFloat4_e2m1fn_x2); at::Tensor q_padded = q; at::Tensor k_padded = k; @@ -88,21 +108,27 @@ mha_fwd( // Apply padding if needed if (needs_padding) { - const int pad_size = head_size_padded - head_size_og; + const int pad_size = head_size_padded - head_size_actual; q_padded = torch::nn::functional::pad(q, torch::nn::functional::PadFuncOptions({0, pad_size})); k_padded = torch::nn::functional::pad(k, torch::nn::functional::PadFuncOptions({0, pad_size})); v_padded = torch::nn::functional::pad(v, torch::nn::functional::PadFuncOptions({0, pad_size})); } + // Low-precision (FP8/FP4) inputs accumulate in float32 and produce a float32 output. + auto out_dtype = is_mxfp ? torch::kFloat32 : q_dtype; + at::Tensor out_padded; if (out_.has_value()) { auto out_val = out_.value(); - if (head_size_og != head_size_padded) { - const int pad_size = head_size_padded - head_size_og; + if (needs_padding) { + const int pad_size = head_size_padded - head_size_actual; out_padded = torch::nn::functional::pad(out_val, torch::nn::functional::PadFuncOptions({0, pad_size})); } else { out_padded = out_val; } + } else if (is_mxfp) { + out_padded = torch::zeros({batch_size, seqlen_q, num_heads, head_size_padded}, + q_padded.options().dtype(out_dtype)); } else { out_padded = torch::zeros_like(q_padded); } @@ -144,7 +170,7 @@ mha_fwd( // seqlen_q_rounded must be >= the kernel's Q tile size, otherwise the S_dmask // write in the mainloop will go out of bounds for rows in the padding region // of the tile. The Q tile size depends on head_size (see fmha_utils.hpp): - const int q_tile_size = (head_size_og <= 32) ? 64 : (head_size_og <= 128) ? 128 : 256; + const int q_tile_size = (head_size_actual <= 32) ? 64 : (head_size_actual <= 128) ? 128 : 256; const int seqlen_q_rounded = round_multiple(seqlen_q, q_tile_size); const int seqlen_k_rounded = round_multiple(seqlen_k, 128); at::Tensor S_dmask; @@ -155,6 +181,18 @@ mha_fwd( S_dmask = torch::empty({0}, q.options()); } + // MXFP (block-scaled) paths require per-tensor scale factors; contiguity is enforced. + std::optional scale_q = scale_q_.has_value() + ? std::optional(ensure_contiguous(scale_q_.value())) : std::nullopt; + std::optional scale_k = scale_k_.has_value() + ? std::optional(ensure_contiguous(scale_k_.value())) : std::nullopt; + std::optional scale_v = scale_v_.has_value() + ? std::optional(ensure_contiguous(scale_v_.value())) : std::nullopt; + if (is_mxfp) { + TORCH_CHECK(scale_q.has_value() && scale_k.has_value() && scale_v.has_value(), + "MXFP (fp8/fp4) forward currently requires scale_q, scale_k and scale_v tensors."); + } + cutlass_fmha_fwd_fix_impl( queue, q_padded, k_padded, v_padded, out_padded, @@ -164,12 +202,13 @@ mha_fwd( is_causal, is_local, p_dropout, philox_seed, philox_offset, nullptr, return_softmax ? S_dmask.data_ptr() : nullptr, - seqlen_q_rounded, seqlen_k_rounded); + seqlen_q_rounded, seqlen_k_rounded, + scale_q, scale_k, scale_v); - // Strip padding from output back to original head_size + // Strip padding from output back to original (unpacked) head_size at::Tensor out = needs_padding ? out_padded.index({torch::indexing::Slice(), torch::indexing::Slice(), - torch::indexing::Slice(), torch::indexing::Slice(0, head_size_og)}) + torch::indexing::Slice(), torch::indexing::Slice(0, head_size_actual)}) .contiguous() : ensure_contiguous(out_padded); @@ -386,7 +425,12 @@ mha_varlen_fwd( int window_size_right, const float softcap, const bool return_softmax, - std::optional gen_) { + std::optional gen_, + std::optional scale_q_ = std::nullopt, + std::optional scale_k_ = std::nullopt, + std::optional scale_v_ = std::nullopt, + std::optional cu_scale_q_ = std::nullopt, + std::optional cu_scale_kv_ = std::nullopt) { auto device_idx = q.device().index(); compat::select_device(device_idx); @@ -404,23 +448,40 @@ mha_varlen_fwd( const int batch_size = cu_seqlens_q.numel() - 1; auto q_dtype = q.dtype(); - TORCH_CHECK(q_dtype == torch::kFloat16 || q_dtype == torch::kBFloat16, - "FlashAttention only support fp16 and bf16 data type"); + const bool is_mxfp = (q_dtype == torch::kFloat8_e5m2 || q_dtype == torch::kFloat8_e4m3fn || + q_dtype == torch::kFloat4_e2m1fn_x2); + TORCH_CHECK(q_dtype == torch::kFloat16 || q_dtype == torch::kBFloat16 || is_mxfp, + "FlashAttention only support fp16, bf16, fp8_e5m2, fp8_e4m3fn and fp4_e2m1fn_x2 data type"); TORCH_CHECK(k.dtype() == q_dtype, "query and key must have the same dtype"); - TORCH_CHECK(v.dtype() == q_dtype, "query and value must have the same dtype"); TORCH_CHECK(k.size(-1) == head_size_og, "Key head dimension must match Query head dimension"); - TORCH_CHECK(v.size(-1) == head_size_og, "Value head dimension must match Query head dimension"); - TORCH_CHECK(head_size_og <= 256 || head_size_og == 512, - "FlashAttention XPU varlen only supports head dimension up to 256 or exactly 512. Got: " + std::to_string(head_size_og)); + if (!is_mxfp) { + TORCH_CHECK(v.dtype() == q_dtype, "query and value must have the same dtype"); + TORCH_CHECK(v.size(-1) == head_size_og, "Value head dimension must match Query head dimension"); + } else if (q_dtype == torch::kFloat4_e2m1fn_x2) { + // MXFP4 (e2m1) consumes a BF16 value tensor; reject mismatches early + // instead of silently producing wrong results inside the kernel. + TORCH_CHECK(v.dtype() == torch::kBFloat16, + "For fp4_e2m1fn_x2 query, value must be bf16"); + } else { + // MXFP8 (e5m2 / e4m3fn) requires the value tensor to match the query fp8 dtype. + TORCH_CHECK(v.dtype() == q_dtype, + "For fp8 query, value must have the same fp8 dtype as query"); + } + // For FP4 (e2m1) the head dimension is packed (2 values per byte); unpack for padding math. + const int head_size_actual = (q_dtype == torch::kFloat4_e2m1fn_x2) ? head_size_og * 2 : head_size_og; + TORCH_CHECK(head_size_actual <= 256 || head_size_actual == 512, + "FlashAttention XPU varlen only supports head dimension up to 256 or exactly 512. Got: " + std::to_string(head_size_actual)); CHECK_DEVICE(q); CHECK_DEVICE(k); CHECK_DEVICE(v); - // XPU requires head_size to be a multiple of 32 - const int head_size_padded = round_multiple(head_size_og, 32); - const bool needs_padding = (head_size_og != head_size_padded); - const int pad_size = head_size_padded - head_size_og; + // XPU requires head_size to be a multiple of 32. + // FP4 packed inputs are consumed in their packed layout and are not padded. + const int head_size_padded = round_multiple(head_size_actual, 32); + const bool needs_padding = (head_size_actual != head_size_padded) + && (q_dtype != torch::kFloat4_e2m1fn_x2); + const int pad_size = head_size_padded - head_size_actual; auto maybe_pad = [&](const at::Tensor& t) -> at::Tensor { return needs_padding @@ -432,8 +493,16 @@ mha_varlen_fwd( at::Tensor k_padded = maybe_pad(k); at::Tensor v_padded = maybe_pad(v); - at::Tensor out_padded = out_.has_value() ? maybe_pad(out_.value()) - : torch::zeros_like(q_padded); + auto out_dtype = is_mxfp ? torch::kFloat32 : q_dtype; + at::Tensor out_padded; + if (out_.has_value()) { + out_padded = maybe_pad(out_.value()); + } else if (is_mxfp) { + out_padded = torch::zeros({total_q, num_heads, head_size_padded}, + q_padded.options().dtype(out_dtype)); + } else { + out_padded = torch::zeros_like(q_padded); + } const bool is_local = (window_size_left != -1) || (window_size_right != -1); const bool is_paged = block_table_.has_value() && block_table_->defined(); @@ -461,7 +530,7 @@ mha_varlen_fwd( rng_state[1] = static_cast(philox_offset); } - const int q_tile_size = (head_size_og <= 32) ? 64 : (head_size_og <= 128) ? 128 : 256; + const int q_tile_size = (head_size_actual <= 32) ? 64 : (head_size_actual <= 128) ? 128 : 256; const int k_tile_size = (max_seqlen_q == 1 && !is_paged) ? 512 : 128; const int seqlen_q_rounded = round_multiple(max_seqlen_q, q_tile_size); const int seqlen_k_rounded = round_multiple(max_seqlen_k, k_tile_size); @@ -473,6 +542,22 @@ mha_varlen_fwd( S_dmask = torch::empty({0}, q.options()); } + // MXFP (block-scaled) scale factor tensors (optional; required for fp8/fp4). + std::optional scale_q = scale_q_.has_value() + ? std::optional(ensure_contiguous(scale_q_.value())) : std::nullopt; + std::optional scale_k = scale_k_.has_value() + ? std::optional(ensure_contiguous(scale_k_.value())) : std::nullopt; + std::optional scale_v = scale_v_.has_value() + ? std::optional(ensure_contiguous(scale_v_.value())) : std::nullopt; + std::optional cu_scale_q = cu_scale_q_.has_value() + ? std::optional(ensure_contiguous(cu_scale_q_.value())) : std::nullopt; + std::optional cu_scale_kv = cu_scale_kv_.has_value() + ? std::optional(ensure_contiguous(cu_scale_kv_.value())) : std::nullopt; + if (is_mxfp) { + TORCH_CHECK(scale_q.has_value() && scale_k.has_value() && scale_v.has_value(), + "MXFP (fp8/fp4) varlen forward currently requires scale_q, scale_k and scale_v tensors."); + } + cutlass_fmha_fwd_varlen_impl( queue, q_padded, k_padded, v_padded, out_padded, @@ -485,12 +570,13 @@ mha_varlen_fwd( true, is_paged, is_causal, is_local, p_dropout, philox_seed, philox_offset, nullptr, return_softmax ? S_dmask.data_ptr() : nullptr, - seqlen_q_rounded, seqlen_k_rounded); + seqlen_q_rounded, seqlen_k_rounded, + scale_q, scale_k, scale_v, cu_scale_q, cu_scale_kv); - // Strip padding from output back to original head_size + // Strip padding from output back to original (unpacked) head_size at::Tensor out = needs_padding ? out_padded.index({torch::indexing::Slice(), torch::indexing::Slice(), - torch::indexing::Slice(0, head_size_og)}) + torch::indexing::Slice(0, head_size_actual)}) .contiguous() : ensure_contiguous(out_padded); @@ -938,7 +1024,13 @@ mha_fwd( const int64_t window_size_right, const double softcap, const bool return_softmax, - c10::optional gen_) { + c10::optional gen_, + c10::optional scale_q_, + c10::optional scale_k_, + c10::optional scale_v_) { + auto to_std = [](const c10::optional& o) -> std::optional { + return o.has_value() ? std::optional(o.value()) : std::nullopt; + }; return FLASH_NAMESPACE::mha_fwd( q, k, @@ -952,7 +1044,10 @@ mha_fwd( static_cast(window_size_right), static_cast(softcap), return_softmax, - gen_ + gen_, + to_std(scale_q_), + to_std(scale_k_), + to_std(scale_v_) ); } @@ -1041,7 +1136,12 @@ mha_varlen_fwd( int64_t window_size_right, const double softcap, const bool return_softmax, - std::optional gen_) { + std::optional gen_, + std::optional scale_q_, + std::optional scale_k_, + std::optional scale_v_, + std::optional cu_scale_q_, + std::optional cu_scale_kv_) { return FLASH_NAMESPACE::mha_varlen_fwd( const_cast(q), k, @@ -1065,7 +1165,12 @@ mha_varlen_fwd( static_cast(window_size_right), static_cast(softcap), return_softmax, - gen_ + gen_, + scale_q_, + scale_k_, + scale_v_, + cu_scale_q_, + cu_scale_kv_ ); } diff --git a/flash-attn2/flash_attn_xpu/src/collective/fmha_fusion.hpp b/flash-attn2/flash_attn_xpu/src/collective/fmha_fusion.hpp index a87752588..6dcfeadbc 100644 --- a/flash-attn2/flash_attn_xpu/src/collective/fmha_fusion.hpp +++ b/flash-attn2/flash_attn_xpu/src/collective/fmha_fusion.hpp @@ -40,6 +40,7 @@ using namespace cute; struct VariableLength { int max_length; int* cumulative_length = nullptr; + int* cumulative_scale_length = nullptr; // For MXFP varlen scale indexing CUTE_HOST_DEVICE operator int() const { return max_length; diff --git a/flash-attn2/flash_attn_xpu/src/collective/fmha_fwd_mainloop_xe2.hpp b/flash-attn2/flash_attn_xpu/src/collective/fmha_fwd_mainloop_xe2.hpp index 8f7a5b1e7..aff675be5 100644 --- a/flash-attn2/flash_attn_xpu/src/collective/fmha_fwd_mainloop_xe2.hpp +++ b/flash-attn2/flash_attn_xpu/src/collective/fmha_fwd_mainloop_xe2.hpp @@ -18,6 +18,7 @@ #include "../philox.hpp" #include "./fmha_fwd_common.hpp" +#include "cutlass/gemm/collective/xe_common_blockscaled_mxfp.hpp" namespace cutlass::fmha { @@ -48,7 +49,12 @@ template < class TiledCopyQ_ = void, class TiledCopyK_ = void, class TiledCopyV_ = void, - bool HasRotary_ = false> + bool HasRotary_ = false, + bool UseScale_ = false, // MXFP: block-scaled MMA + bool F8kvF16mma_ = false, // MXFP: F8 KV with F16 MMA (scalar scale) + class TensorScaleQ_ = void, // MXFP: optional scale tensors + class TensorScaleK_ = void, + class TensorScaleV_ = void> struct FMHAFwdMainloopXe2 { static_assert( cutlass::detail::dependent_false, @@ -72,7 +78,12 @@ template < class TiledCopyQ_, class TiledCopyK_, class TiledCopyV_, - bool HasRotary_> + bool HasRotary_, + bool UseScale_, + bool F8kvF16mma_, + class TensorScaleQ_, + class TensorScaleK_, + class TensorScaleV_> struct FMHAFwdMainloopXe2< Xe2, CausalMask_, @@ -88,7 +99,12 @@ struct FMHAFwdMainloopXe2< TiledCopyQ_, TiledCopyK_, TiledCopyV_, - HasRotary_> { + HasRotary_, + UseScale_, + F8kvF16mma_, + TensorScaleQ_, + TensorScaleK_, + TensorScaleV_> { // Pull in common type aliases from the shared traits. using Traits = FMHAFwdMainloopTraits< @@ -128,6 +144,61 @@ struct FMHAFwdMainloopXe2< static constexpr bool PagedKV = PagedKV_; static constexpr bool HasRotary = HasRotary_; + // --- MXFP (block-scaled) support --------------------------------------- + using ElementQ = typename TensorQ::element_type; + static constexpr bool FP4Input = + cute::is_same_v; + static constexpr bool UseScale = UseScale_; + static constexpr bool F8kvF16mma = F8kvF16mma_; + + using TensorScaleQ = TensorScaleQ_; + using TensorScaleK = TensorScaleK_; + using TensorScaleV = TensorScaleV_; + using TensorScaleQ2D = conditional_t, void, + decltype(TensorScaleQ_{}(append>(make_coord(_, _), 0)))>; + using TensorScaleK2D = conditional_t, void, + decltype(TensorScaleK_{}(append>(make_coord(_, _), 0)))>; + using TensorScaleV2D = conditional_t, void, + decltype(TensorScaleV_{}(append>(make_coord(_, _), 0)))>; + + using ElementScaleQ = conditional_t, float, + typename TensorScaleQ::element_type>; + using ElementScaleK = conditional_t, float, + typename TensorScaleK::element_type>; + using ElementScaleV = conditional_t, float, + typename TensorScaleV::element_type>; + + // Block and MMA dimensions used for block-scale addressing. + static constexpr int BLK_Q = get<0>(TileShapeQK{}); + static constexpr int BLK_K = get<1>(TileShapeQK{}); + static constexpr int BLK_QK_D = get<2>(TileShapeQK{}); + static constexpr int BLK_P = get<0>(TileShapePV{}); + static constexpr int BLK_V = get<1>(TileShapePV{}); + static constexpr int BLK_PV_D = get<2>(TileShapePV{}); + + static constexpr int ATOM_Q = get<1>(typename TiledMMAQK::ThrLayoutVMNK{}.shape()); + static constexpr int ATOM_K = get<2>(typename TiledMMAQK::ThrLayoutVMNK{}.shape()); + static constexpr int ATOM_QK_D = get<3>(typename TiledMMAQK::ThrLayoutVMNK{}.shape()); + static constexpr int MMA_QK_D = get<2>(typename TiledMMAQK::Shape_MNK{}); + static constexpr int ATOM_P = get<1>(typename TiledMMAPV::ThrLayoutVMNK{}.shape()); + static constexpr int ATOM_V = get<2>(typename TiledMMAPV::ThrLayoutVMNK{}.shape()); + static constexpr int ATOM_PV_D = get<3>(typename TiledMMAPV::ThrLayoutVMNK{}.shape()); + static constexpr int MMA_PV_D = get<2>(typename TiledMMAPV::Shape_MNK{}); + + static constexpr int SG_Q = ceil_div(BLK_Q, ATOM_Q); + static constexpr int SG_K = ceil_div(BLK_K, ATOM_K); + static constexpr int SG_QK_D = ceil_div(BLK_QK_D, ATOM_QK_D); + static constexpr int SG_P = ceil_div(BLK_P, ATOM_P); + static constexpr int SG_V = ceil_div(BLK_V, ATOM_V); + static constexpr int SG_PV_D = ceil_div(BLK_PV_D, ATOM_PV_D); + + static constexpr auto GROUP_K = 32; + + using DefScaleType = cutlass::float_ue8m0_t; + using ElementScaleP = cute::conditional_t; + // ----------------------------------------------------------------------- + + // User-facing arguments struct Arguments { ElementS const scale; @@ -150,6 +221,9 @@ struct FMHAFwdMainloopXe2< const typename TensorQ::element_type* rotary_sin = nullptr; int rotary_dim = 0; bool is_rotary_interleaved = true; + // MXFP: scalar scales for F8 KV with F16 MMA path. + float scale_k_scalar = 1.0f; + float scale_v_scalar = 1.0f; }; struct LocalMaskFields { @@ -192,6 +266,9 @@ struct FMHAFwdMainloopXe2< paged; [[no_unique_address]] conditional_t rotary; + // MXFP: scalar scales for F8 KV with F16 MMA path. + float scale_k_scalar = 1.0f; + float scale_v_scalar = 1.0f; }; // SLM data @@ -231,6 +308,8 @@ struct FMHAFwdMainloopXe2< p.rotary = {args.rotary_cos, args.rotary_sin, args.rotary_dim, args.is_rotary_interleaved}; } + p.scale_k_scalar = args.scale_k_scalar; + p.scale_v_scalar = args.scale_v_scalar; return p; } @@ -260,7 +339,11 @@ struct FMHAFwdMainloopXe2< int head_q, int num_heads, int q_offset_sg, - int rotary_base) { + int rotary_base, + int l_coord = 0, + TensorScaleQ2D const& scaleQ_2D = TensorScaleQ2D{}, + TensorScaleK2D const& scaleK_2D = TensorScaleK2D{}, + TensorScaleV2D const& scaleV_2D = TensorScaleV2D{}) { using namespace sycl::ext::oneapi::this_work_item; auto tile_shape_v = @@ -444,9 +527,72 @@ struct FMHAFwdMainloopXe2< } } copy(copy_k, tKgK_cache(_, _, _, D), tKrK); - reorder(tQrQ, tSrQ); - reorder(tKrK, tSrK); - cute::gemm(mma_qk, tSrQ, tSrK, tSrS); + if constexpr (FP4Input) { + copy(tQrQ, tSrQ); + copy(tKrK, tSrK); + } else { + reorder(tQrQ, tSrQ); + reorder(tKrK, tSrK); + } + + if constexpr (UseScale) { +#if defined(SYCL_INTEL_TARGET) && (SYCL_INTEL_TARGET == 35) + // Block-scaled MMA with scale factors for Q*K. The block-scaled + // (BDPAS) atom and helpers only exist on CRI (SYCL_INTEL_TARGET==35). + // On BMG this kernel is instantiated (host references it) but never + // executed, so the body is replaced by an invalid-control-path stub. + const int q_coord = get<0>(blk_qv) * BLK_Q + + ((thr_id / intel::sg_size) / ATOM_K) * SG_Q; + const int k_coord = K * BLK_K + + ((thr_id % intel::sg_size) / ATOM_K) * SG_K; + + auto [tiled_copy_scaleQ, copy_iter_scaleQ, fragment_scaleQ] = + gemm::collective::make_scaled_copy(scaleQ_2D, q_coord, l_coord, size<4>(tKgK)); + auto [tiled_copy_scaleK, copy_iter_scaleK, fragment_scaleK] = + gemm::collective::make_scaled_copy(scaleK_2D, k_coord, l_coord, size<4>(tKgK)); + auto [gemm_qm_offsets, gemm_kn_offsets, gemm_qk_offsets, gemm_kk_offsets] = + gemm::collective::make_scaled_offsets< + decltype(size<1>(tSrQ.shape()))::value, + decltype(size<1>(tSrK.shape()))::value, + decltype(size<2>(tSrK.shape()))::value, + MMA_QK_D, + GROUP_K, + typename decltype(tiled_copy_scaleQ)::BlockShape, + typename decltype(tiled_copy_scaleK)::BlockShape>(); + + copy(tiled_copy_scaleQ, copy_iter_scaleQ(_, _, _, D), fragment_scaleQ); + copy(tiled_copy_scaleK, copy_iter_scaleK(_, _, _, D), fragment_scaleK); + + using scaleQSize = decltype(size(fragment_scaleQ)); + using scaleKSize = decltype(size(fragment_scaleK)); + + Tensor scaleQ_view = make_tensor( + recast>(fragment_scaleQ).data(), + make_layout(Shape<_1, decltype(size<1>(tSrQ.shape())), _1>{}, Stride<_1, _0, _0>{})); + Tensor scaleK_view = make_tensor( + recast>(fragment_scaleK).data(), + make_layout(Shape<_1, decltype(size<1>(tSrK.shape())), _1>{}, Stride<_1, _0, _0>{})); + + auto zipped_q = make_zip_tensor(tSrQ, scaleQ_view, gemm_qm_offsets, gemm_qk_offsets); + auto zipped_k = make_zip_tensor(tSrK, scaleK_view, gemm_kn_offsets, gemm_kk_offsets); + + cute::gemm(mma_qk, zipped_q, zipped_k, tSrS); +#else + CUTE_INVALID_CONTROL_PATH( + "Block-scaled (MXFP) Q*K path is only available on CRI " + "(SYCL_INTEL_TARGET==35)."); +#endif + } else { + if constexpr (F8kvF16mma) { + // F8 KV with F16 MMA path - apply scalar scale to K. + for (int i = 0; i < tSrK.size(); i++) + tSrK(i) = static_cast( + params.scale_k_scalar * static_cast(tSrK(i))); + } + cute::gemm(mma_qk, tSrQ, tSrK, tSrS); + } } auto cS_thread = thr_mma_qk.partition_C(gP_all(_, _, K)); @@ -580,7 +726,61 @@ struct FMHAFwdMainloopXe2< for (int VV = 0; VV < VTiles; VV++) { copy(copy_v, tVgV_cache(_, _, _, VV), tVrV); reorder(tVrV, tArV); - cute::gemm(mma_pv, tArP, tArV, tArA(_, _, _, VV)); + if constexpr (UseScale && !FP4Input) { +#if defined(SYCL_INTEL_TARGET) && (SYCL_INTEL_TARGET == 35) + // Block-scaled MMA with scale factors for P*V. CRI-only (BDPAS); + // on BMG this branch is instantiated but never runs. + const int v_coord = get<1>(blk_qv) * VTiles * BLK_V + VV * BLK_V + + ((thr_id % intel::sg_size) / ATOM_V) * SG_V; + + // P is a dummy scale, the same shape as V. + auto [tiled_copy_scaleP, copy_iter_scaleP, fragment_scaleP] = + gemm::collective::make_scaled_copy(scaleV_2D); + auto [tiled_copy_scaleV, copy_iter_scaleV, fragment_scaleV] = + gemm::collective::make_scaled_copy(scaleV_2D, v_coord, l_coord, blk_k1); + auto [gemm_p_offsets, gemm_v_offsets, gemm_pk_offsets, gemm_vk_offsets] = + gemm::collective::make_scaled_offsets< + decltype(size<1>(tArP.shape()))::value, + decltype(size<1>(tArV.shape()))::value, + decltype(size<2>(tArV.shape()))::value, + MMA_PV_D, + GROUP_K, + typename decltype(tiled_copy_scaleP)::BlockShape, + typename decltype(tiled_copy_scaleV)::BlockShape>(); + + fill(fragment_scaleP, ElementScaleV(1)); + copy(tiled_copy_scaleV, copy_iter_scaleV(_, _, _, K), fragment_scaleV); + + using scalePSize = decltype(size(fragment_scaleP)); + using scaleVSize = decltype(size(fragment_scaleV)); + + Tensor scaleV_view = make_tensor( + recast>(fragment_scaleV).data(), + make_layout(Shape<_1, decltype(size<1>(tArV.shape())), _1>{}, Stride<_1, _0, _0>{})); + Tensor scaleP_view = make_tensor( + recast>(fragment_scaleP).data(), + make_layout(Shape<_1, decltype(size<1>(tArP.shape())), _1>{}, Stride<_1, _0, _0>{})); + + auto zipped_p = make_zip_tensor(tArP, scaleP_view, gemm_p_offsets, gemm_pk_offsets); + auto zipped_v = make_zip_tensor(tArV, scaleV_view, gemm_v_offsets, gemm_vk_offsets); + + cute::gemm(mma_pv, zipped_p, zipped_v, tArA(_, _, _, VV)); +#else + CUTE_INVALID_CONTROL_PATH( + "Block-scaled (MXFP) P*V path is only available on CRI " + "(SYCL_INTEL_TARGET==35)."); +#endif + } else { + if constexpr (F8kvF16mma) { + // F8 KV with F16 MMA path - apply scalar scale to V. + for (int i = 0; i < tArV.size(); i++) + tArV(i) = static_cast( + params.scale_v_scalar * static_cast(tArV(i))); + } + cute::gemm(mma_pv, tArP, tArV, tArA(_, _, _, VV)); + } } if constexpr (PagedKV) { diff --git a/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim128_fix_mxfp.cpp b/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim128_fix_mxfp.cpp new file mode 100644 index 000000000..abe77f4bc --- /dev/null +++ b/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim128_fix_mxfp.cpp @@ -0,0 +1,11 @@ +#include "fmha_fwd_impl.hpp" + +// MXFP (block-scaled FP8/FP4) fixed mode: IsVarLen=0, IsPaged=0 +// Instantiates prefill + decode for all three element types (e5m2, e4m3, e2m1). + +template void policy_dispatch_mxfp_e5m2(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e5m2(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e4m3(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e4m3(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e2m1(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e2m1(sycl::queue&, const fmha_fwd_args_t&); diff --git a/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim128_varlen_mxfp.cpp b/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim128_varlen_mxfp.cpp new file mode 100644 index 000000000..f6dfec1a1 --- /dev/null +++ b/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim128_varlen_mxfp.cpp @@ -0,0 +1,11 @@ +#include "fmha_fwd_impl.hpp" + +// MXFP (block-scaled FP8/FP4) varlen mode: IsVarLen=1, IsPaged=0 +// Instantiates prefill + decode for all three element types (e5m2, e4m3, e2m1). + +template void policy_dispatch_mxfp_e5m2(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e5m2(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e4m3(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e4m3(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e2m1(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e2m1(sycl::queue&, const fmha_fwd_args_t&); diff --git a/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim160_fix_mxfp.cpp b/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim160_fix_mxfp.cpp new file mode 100644 index 000000000..c3c29808e --- /dev/null +++ b/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim160_fix_mxfp.cpp @@ -0,0 +1,11 @@ +#include "fmha_fwd_impl.hpp" + +// MXFP (block-scaled FP8/FP4) fixed mode: IsVarLen=0, IsPaged=0 +// Instantiates prefill + decode for all three element types (e5m2, e4m3, e2m1). + +template void policy_dispatch_mxfp_e5m2(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e5m2(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e4m3(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e4m3(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e2m1(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e2m1(sycl::queue&, const fmha_fwd_args_t&); diff --git a/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim160_varlen_mxfp.cpp b/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim160_varlen_mxfp.cpp new file mode 100644 index 000000000..ca18dd258 --- /dev/null +++ b/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim160_varlen_mxfp.cpp @@ -0,0 +1,11 @@ +#include "fmha_fwd_impl.hpp" + +// MXFP (block-scaled FP8/FP4) varlen mode: IsVarLen=1, IsPaged=0 +// Instantiates prefill + decode for all three element types (e5m2, e4m3, e2m1). + +template void policy_dispatch_mxfp_e5m2(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e5m2(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e4m3(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e4m3(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e2m1(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e2m1(sycl::queue&, const fmha_fwd_args_t&); diff --git a/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim192_fix_mxfp.cpp b/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim192_fix_mxfp.cpp new file mode 100644 index 000000000..371e3f1dd --- /dev/null +++ b/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim192_fix_mxfp.cpp @@ -0,0 +1,11 @@ +#include "fmha_fwd_impl.hpp" + +// MXFP (block-scaled FP8/FP4) fixed mode: IsVarLen=0, IsPaged=0 +// Instantiates prefill + decode for all three element types (e5m2, e4m3, e2m1). + +template void policy_dispatch_mxfp_e5m2(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e5m2(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e4m3(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e4m3(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e2m1(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e2m1(sycl::queue&, const fmha_fwd_args_t&); diff --git a/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim192_varlen_mxfp.cpp b/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim192_varlen_mxfp.cpp new file mode 100644 index 000000000..06d5f7aba --- /dev/null +++ b/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim192_varlen_mxfp.cpp @@ -0,0 +1,11 @@ +#include "fmha_fwd_impl.hpp" + +// MXFP (block-scaled FP8/FP4) varlen mode: IsVarLen=1, IsPaged=0 +// Instantiates prefill + decode for all three element types (e5m2, e4m3, e2m1). + +template void policy_dispatch_mxfp_e5m2(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e5m2(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e4m3(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e4m3(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e2m1(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e2m1(sycl::queue&, const fmha_fwd_args_t&); diff --git a/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim256_fix_mxfp.cpp b/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim256_fix_mxfp.cpp new file mode 100644 index 000000000..ae8814090 --- /dev/null +++ b/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim256_fix_mxfp.cpp @@ -0,0 +1,11 @@ +#include "fmha_fwd_impl.hpp" + +// MXFP (block-scaled FP8/FP4) fixed mode: IsVarLen=0, IsPaged=0 +// Instantiates prefill + decode for all three element types (e5m2, e4m3, e2m1). + +template void policy_dispatch_mxfp_e5m2(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e5m2(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e4m3(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e4m3(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e2m1(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e2m1(sycl::queue&, const fmha_fwd_args_t&); diff --git a/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim256_varlen_mxfp.cpp b/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim256_varlen_mxfp.cpp new file mode 100644 index 000000000..fb5641904 --- /dev/null +++ b/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim256_varlen_mxfp.cpp @@ -0,0 +1,11 @@ +#include "fmha_fwd_impl.hpp" + +// MXFP (block-scaled FP8/FP4) varlen mode: IsVarLen=1, IsPaged=0 +// Instantiates prefill + decode for all three element types (e5m2, e4m3, e2m1). + +template void policy_dispatch_mxfp_e5m2(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e5m2(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e4m3(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e4m3(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e2m1(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e2m1(sycl::queue&, const fmha_fwd_args_t&); diff --git a/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim32_fix_mxfp.cpp b/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim32_fix_mxfp.cpp new file mode 100644 index 000000000..d0650f8df --- /dev/null +++ b/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim32_fix_mxfp.cpp @@ -0,0 +1,11 @@ +#include "fmha_fwd_impl.hpp" + +// MXFP (block-scaled FP8/FP4) fixed mode: IsVarLen=0, IsPaged=0 +// Instantiates prefill + decode for all three element types (e5m2, e4m3, e2m1). + +template void policy_dispatch_mxfp_e5m2(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e5m2(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e4m3(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e4m3(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e2m1(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e2m1(sycl::queue&, const fmha_fwd_args_t&); diff --git a/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim32_varlen_mxfp.cpp b/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim32_varlen_mxfp.cpp new file mode 100644 index 000000000..e6476868a --- /dev/null +++ b/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim32_varlen_mxfp.cpp @@ -0,0 +1,11 @@ +#include "fmha_fwd_impl.hpp" + +// MXFP (block-scaled FP8/FP4) varlen mode: IsVarLen=1, IsPaged=0 +// Instantiates prefill + decode for all three element types (e5m2, e4m3, e2m1). + +template void policy_dispatch_mxfp_e5m2(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e5m2(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e4m3(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e4m3(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e2m1(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e2m1(sycl::queue&, const fmha_fwd_args_t&); diff --git a/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim512_fix_mxfp.cpp b/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim512_fix_mxfp.cpp new file mode 100644 index 000000000..5f9097589 --- /dev/null +++ b/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim512_fix_mxfp.cpp @@ -0,0 +1,11 @@ +#include "fmha_fwd_impl.hpp" + +// MXFP (block-scaled FP8/FP4) fixed mode: IsVarLen=0, IsPaged=0 +// Instantiates prefill + decode for all three element types (e5m2, e4m3, e2m1). + +template void policy_dispatch_mxfp_e5m2(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e5m2(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e4m3(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e4m3(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e2m1(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e2m1(sycl::queue&, const fmha_fwd_args_t&); diff --git a/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim512_varlen_mxfp.cpp b/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim512_varlen_mxfp.cpp new file mode 100644 index 000000000..b7ba4b1ed --- /dev/null +++ b/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim512_varlen_mxfp.cpp @@ -0,0 +1,11 @@ +#include "fmha_fwd_impl.hpp" + +// MXFP (block-scaled FP8/FP4) varlen mode: IsVarLen=1, IsPaged=0 +// Instantiates prefill + decode for all three element types (e5m2, e4m3, e2m1). + +template void policy_dispatch_mxfp_e5m2(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e5m2(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e4m3(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e4m3(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e2m1(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e2m1(sycl::queue&, const fmha_fwd_args_t&); diff --git a/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim64_fix_mxfp.cpp b/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim64_fix_mxfp.cpp new file mode 100644 index 000000000..b135019c4 --- /dev/null +++ b/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim64_fix_mxfp.cpp @@ -0,0 +1,11 @@ +#include "fmha_fwd_impl.hpp" + +// MXFP (block-scaled FP8/FP4) fixed mode: IsVarLen=0, IsPaged=0 +// Instantiates prefill + decode for all three element types (e5m2, e4m3, e2m1). + +template void policy_dispatch_mxfp_e5m2(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e5m2(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e4m3(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e4m3(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e2m1(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e2m1(sycl::queue&, const fmha_fwd_args_t&); diff --git a/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim64_varlen_mxfp.cpp b/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim64_varlen_mxfp.cpp new file mode 100644 index 000000000..803c3db80 --- /dev/null +++ b/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim64_varlen_mxfp.cpp @@ -0,0 +1,11 @@ +#include "fmha_fwd_impl.hpp" + +// MXFP (block-scaled FP8/FP4) varlen mode: IsVarLen=1, IsPaged=0 +// Instantiates prefill + decode for all three element types (e5m2, e4m3, e2m1). + +template void policy_dispatch_mxfp_e5m2(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e5m2(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e4m3(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e4m3(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e2m1(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e2m1(sycl::queue&, const fmha_fwd_args_t&); diff --git a/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim96_fix_mxfp.cpp b/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim96_fix_mxfp.cpp new file mode 100644 index 000000000..c3418a43d --- /dev/null +++ b/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim96_fix_mxfp.cpp @@ -0,0 +1,11 @@ +#include "fmha_fwd_impl.hpp" + +// MXFP (block-scaled FP8/FP4) fixed mode: IsVarLen=0, IsPaged=0 +// Instantiates prefill + decode for all three element types (e5m2, e4m3, e2m1). + +template void policy_dispatch_mxfp_e5m2(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e5m2(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e4m3(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e4m3(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e2m1(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e2m1(sycl::queue&, const fmha_fwd_args_t&); diff --git a/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim96_varlen_mxfp.cpp b/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim96_varlen_mxfp.cpp new file mode 100644 index 000000000..e95148fac --- /dev/null +++ b/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim96_varlen_mxfp.cpp @@ -0,0 +1,11 @@ +#include "fmha_fwd_impl.hpp" + +// MXFP (block-scaled FP8/FP4) varlen mode: IsVarLen=1, IsPaged=0 +// Instantiates prefill + decode for all three element types (e5m2, e4m3, e2m1). + +template void policy_dispatch_mxfp_e5m2(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e5m2(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e4m3(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e4m3(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e2m1(sycl::queue&, const fmha_fwd_args_t&); +template void policy_dispatch_mxfp_e2m1(sycl::queue&, const fmha_fwd_args_t&); diff --git a/flash-attn2/flash_attn_xpu/src/fmha_fwd.cpp b/flash-attn2/flash_attn_xpu/src/fmha_fwd.cpp index f4cb5ca6e..efca27d71 100644 --- a/flash-attn2/flash_attn_xpu/src/fmha_fwd.cpp +++ b/flash-attn2/flash_attn_xpu/src/fmha_fwd.cpp @@ -225,7 +225,12 @@ void cutlass_fmha_fwd_varlen_impl( void* rng_state, void* s_dmask, int seqlen_q_rounded, - int seqlen_k_rounded) { + int seqlen_k_rounded, + const std::optional& scale_q, + const std::optional& scale_k, + const std::optional& scale_v, + const std::optional& cu_scale_q, + const std::optional& cu_scale_kv) { int batch_size, num_heads_q, num_heads_kv, head_size; int total_seqlen_q, total_seqlen_k; int num_blocks, block_size, max_blocks_per_seq; @@ -309,6 +314,13 @@ void cutlass_fmha_fwd_varlen_impl( seqlen_q_rounded, seqlen_k_rounded}; + // MXFP block-scaled scale factor tensors (only set for FP8/FP4 paths). + if (scale_q.has_value()) args.scale_q = scale_q->data_ptr(); + if (scale_k.has_value()) args.scale_k = scale_k->data_ptr(); + if (scale_v.has_value()) args.scale_v = scale_v->data_ptr(); + if (cu_scale_q.has_value()) args.cu_scale_q = cu_scale_q->data_ptr(); + if (cu_scale_kv.has_value())args.cu_scale_kv = cu_scale_kv->data_ptr(); + const CutlassType cuType = aten_to_Cutlass_dtype(query); if (args.max_queries == 1) { dispatch_fwd_varlen_decode_by_head(queue, cuType, args, args.head_size); @@ -335,7 +347,10 @@ void cutlass_fmha_fwd_fix_impl( void* rng_state, void* s_dmask, int seqlen_q_rounded, - int seqlen_k_rounded) { + int seqlen_k_rounded, + const std::optional& scale_q, + const std::optional& scale_k, + const std::optional& scale_v) { const int batch_size = query.size(0); const int max_seqlen_q = query.size(1); const int num_heads_q = query.size(2); @@ -384,6 +399,11 @@ void cutlass_fmha_fwd_fix_impl( seqlen_q_rounded, seqlen_k_rounded}; + // MXFP block-scaled scale factor tensors (only set for FP8/FP4 paths). + if (scale_q.has_value()) args.scale_q = scale_q->data_ptr(); + if (scale_k.has_value()) args.scale_k = scale_k->data_ptr(); + if (scale_v.has_value()) args.scale_v = scale_v->data_ptr(); + const CutlassType cuType = aten_to_Cutlass_dtype(query); const int h = args.head_size; diff --git a/flash-attn2/flash_attn_xpu/src/fmha_fwd.hpp b/flash-attn2/flash_attn_xpu/src/fmha_fwd.hpp index bc1c7e42a..8ba7c9c09 100644 --- a/flash-attn2/flash_attn_xpu/src/fmha_fwd.hpp +++ b/flash-attn2/flash_attn_xpu/src/fmha_fwd.hpp @@ -56,8 +56,36 @@ void policy_dispatch_bf16( sycl::queue& queue, const fmha_fwd_args_t& args); -// Combined dispatch (delegates to fp16/bf16 based on cuType) +// MXFP (block-scaled FP8/FP4) dispatch functions. Unconditional: the host pass +// (always SYCL_INTEL_TARGET==20) must own these symbols so they are callable at +// runtime on CRI. The block-scaled device code is gated inside FMHAConfig and +// the mainloop; on the BMG device pass these kernels compile as unreachable +// placeholders (they are only ever executed on CRI hardware). +template +void policy_dispatch_mxfp_e5m2( + sycl::queue& queue, + const fmha_fwd_args_t& args); + +template +void policy_dispatch_mxfp_e4m3( + sycl::queue& queue, + const fmha_fwd_args_t& args); + +template +void policy_dispatch_mxfp_e2m1( + sycl::queue& queue, + const fmha_fwd_args_t& args); + +// Combined dispatch (delegates to fp16/bf16/mxfp based on cuType) // Defined inline in header so callers (fmha_fwd.cpp) can see the template body. +// MXFP is only supported on the varlen/fix, non-paged, non-rotary forward paths +// (matching the CRI reference implementation). The `if constexpr` guard ensures +// the MXFP dispatchers are only ODR-used (and therefore only instantiated) for +// those paths, so the paged / kvcache / rotary instantiations do not pull in the +// (expensive) MXFP kernels. template inline void policy_dispatch( @@ -66,9 +94,21 @@ inline void policy_dispatch( const fmha_fwd_args_t& args) { if (cuType == CutlassType::half) { policy_dispatch_fp16(queue, args); - } else { - policy_dispatch_bf16(queue, args); + return; + } + if constexpr (!IsPaged && !HasRotary) { + if (cuType == CutlassType::mx_float_e5m2) { + policy_dispatch_mxfp_e5m2(queue, args); + return; + } else if (cuType == CutlassType::mx_float_e4m3) { + policy_dispatch_mxfp_e4m3(queue, args); + return; + } else if (cuType == CutlassType::mx_float_e2m1) { + policy_dispatch_mxfp_e2m1(queue, args); + return; + } } + policy_dispatch_bf16(queue, args); } // Varlen prefill mode extern declarations (IsVarLen=1, IsPaged=0/1) @@ -139,6 +179,34 @@ EXTERN_DISPATCH_KVCACHE_PAGED(256) EXTERN_DISPATCH_KVCACHE_PAGED(512) #undef EXTERN_DISPATCH_KVCACHE_PAGED +// MXFP (block-scaled FP8/FP4) extern declarations. Only the varlen prefill, +// varlen decode and fixed non-paged / non-rotary forward paths are supported +// (matching the CRI reference). Each head dimension instantiates all three +// MXFP element types (e5m2, e4m3, e2m1). +#define EXTERN_DISPATCH_MXFP(HDIM) \ + extern template void policy_dispatch_mxfp_e5m2(sycl::queue&, const fmha_fwd_args_t&); \ + extern template void policy_dispatch_mxfp_e5m2(sycl::queue&, const fmha_fwd_args_t&); \ + extern template void policy_dispatch_mxfp_e5m2(sycl::queue&, const fmha_fwd_args_t&); \ + extern template void policy_dispatch_mxfp_e5m2(sycl::queue&, const fmha_fwd_args_t&); \ + extern template void policy_dispatch_mxfp_e4m3(sycl::queue&, const fmha_fwd_args_t&); \ + extern template void policy_dispatch_mxfp_e4m3(sycl::queue&, const fmha_fwd_args_t&); \ + extern template void policy_dispatch_mxfp_e4m3(sycl::queue&, const fmha_fwd_args_t&); \ + extern template void policy_dispatch_mxfp_e4m3(sycl::queue&, const fmha_fwd_args_t&); \ + extern template void policy_dispatch_mxfp_e2m1(sycl::queue&, const fmha_fwd_args_t&); \ + extern template void policy_dispatch_mxfp_e2m1(sycl::queue&, const fmha_fwd_args_t&); \ + extern template void policy_dispatch_mxfp_e2m1(sycl::queue&, const fmha_fwd_args_t&); \ + extern template void policy_dispatch_mxfp_e2m1(sycl::queue&, const fmha_fwd_args_t&); + +EXTERN_DISPATCH_MXFP(32) +EXTERN_DISPATCH_MXFP(64) +EXTERN_DISPATCH_MXFP(96) +EXTERN_DISPATCH_MXFP(128) +EXTERN_DISPATCH_MXFP(160) +EXTERN_DISPATCH_MXFP(192) +EXTERN_DISPATCH_MXFP(256) +EXTERN_DISPATCH_MXFP(512) +#undef EXTERN_DISPATCH_MXFP + void cutlass_fmha_fwd_varlen_impl( sycl::queue& queue, const at::Tensor& query, @@ -164,7 +232,13 @@ void cutlass_fmha_fwd_varlen_impl( void* rng_state = nullptr, void* s_dmask = nullptr, int seqlen_q_rounded = 0, - int seqlen_k_rounded = 0); + int seqlen_k_rounded = 0, + // MXFP (block-scaled) scale factor tensors (optional). + const std::optional& scale_q = std::nullopt, + const std::optional& scale_k = std::nullopt, + const std::optional& scale_v = std::nullopt, + const std::optional& cu_scale_q = std::nullopt, + const std::optional& cu_scale_kv = std::nullopt); void cutlass_fmha_fwd_fix_impl( sycl::queue& queue, @@ -184,7 +258,11 @@ void cutlass_fmha_fwd_fix_impl( void* rng_state = nullptr, void* s_dmask = nullptr, int seqlen_q_rounded = 0, - int seqlen_k_rounded = 0); + int seqlen_k_rounded = 0, + // MXFP (block-scaled) scale factor tensors (optional). + const std::optional& scale_q = std::nullopt, + const std::optional& scale_k = std::nullopt, + const std::optional& scale_v = std::nullopt); void cutlass_fmha_fwd_kvcache_impl( sycl::queue& queue, diff --git a/flash-attn2/flash_attn_xpu/src/fmha_fwd_impl.hpp b/flash-attn2/flash_attn_xpu/src/fmha_fwd_impl.hpp index e686170c8..d691e0d42 100644 --- a/flash-attn2/flash_attn_xpu/src/fmha_fwd_impl.hpp +++ b/flash-attn2/flash_attn_xpu/src/fmha_fwd_impl.hpp @@ -27,15 +27,22 @@ template < typename SubgroupLayoutQK, typename SubgroupLayoutPV_, int PipelineStages, + bool UseScale_ = false, // MXFP: block-scaled MMA (BDPAS) + bool F8kvF16mma_ = false, // MXFP: F8 KV with F16 MMA (scalar scale) typename ElementQ = bfloat16_t, typename ElementK = bfloat16_t, typename ElementV = bfloat16_t, typename ElementO = bfloat16_t, + typename ElementScale = float, // MXFP: block-scale element type typename MMAOperation_ = void, + typename MMAOperationPV_ = void, typename StrideQ = Stride, typename StrideK = Stride, typename StrideV = Stride<_1, int, int, int>, typename StrideO = Stride, + typename StrideScaleQ = Stride<_1, int, int, int>, + typename StrideScaleK = Stride<_1, int, int, int>, + typename StrideScaleV = Stride<_1, int, int, int>, typename GmemTiledCopyQ = void, typename GmemTiledCopyK = void, typename GmemTiledCopyV = void, @@ -43,9 +50,48 @@ template < struct FMHAConfig { static constexpr int SGTileQ = get<0>(shape_div(TileShapeQK{}, shape(SubgroupLayoutQK{})))(); + static constexpr bool UseScale = UseScale_; + static constexpr bool F8kvF16mma = F8kvF16mma_; + + template + static constexpr bool is_f8_v = + cute::is_any_of_v; + template + static constexpr bool is_f16_v = + cute::is_any_of_v; + +#if !(defined(SYCL_INTEL_TARGET) && (SYCL_INTEL_TARGET == 35)) + // BMG path: no block-scaled DPAS and no FP8/FP4 DPAS atoms. Any MXFP input + // (e5m2 / e4m3 / e2m1) falls back to a valid half_t DPAS atom so the kernel + // is compilable on BMG. These kernels are only ever executed on CRI, so the + // BMG fallback body is an (unreachable) placeholder. + using DefaultMMA = cute::conditional_t< + cute::is_same_v || + cute::is_same_v || + cute::is_same_v, + XE_DPAS_TT, + XE_DPAS_TT>; + using MMAOperationPV = cute::conditional_t< + is_void_v, + DefaultMMA, + MMAOperationPV_>; +#else + // CRI path: block-scaled BDPAS when UseScale, else plain DPAS. + using DefaultDpasOp = XE_DPAS_TT; + using DefaultBdpasOp = XE_BDPAS_TT; + using DefaultMMA = + cute::conditional_t; + // PV MMA: F8 V uses the block-scaled/DPAS path; other V (e.g. bf16 for + // MXFP4) uses a plain DPAS atom keyed on ElementV. + using MMAOperationPV = cute::conditional_t< + cute::is_same_v || + cute::is_same_v, + cute::conditional_t, DefaultMMA, MMAOperationPV_>, + XE_DPAS_TT>; +#endif using MMAOperation = cute::conditional_t< is_void_v, - XE_DPAS_TT, + DefaultMMA, MMAOperation_>; using SubgroupLayoutPV = cute::conditional_t< is_void_v, @@ -70,7 +116,7 @@ struct FMHAConfig { Layout, SubgroupLayoutQK>::TiledMMA; using TiledMMAPV = typename TiledMMAHelper< - MMA_Atom, + MMA_Atom, Layout, SubgroupLayoutPV>::TiledMMA; @@ -90,6 +136,15 @@ struct FMHAConfig { using TensorV = decltype(make_dummy_tensor(ElementV{}, StrideV{})); using TensorO = decltype(make_dummy_tensor(ElementO{}, StrideO{})); + // MXFP: block-scale dummy tensors (harmless float tensors when the + // config is not block-scaled; the mainloop ignores them unless UseScale). + using TensorScaleQ = + decltype(make_dummy_tensor(ElementScale{}, StrideScaleQ{})); + using TensorScaleK = + decltype(make_dummy_tensor(ElementScale{}, StrideScaleK{})); + using TensorScaleV = + decltype(make_dummy_tensor(ElementScale{}, StrideScaleV{})); + using MainloopDispatchPolicy = cutlass::fmha::Xe2; using CollectiveMainloop = cutlass::fmha::collective::FMHAFwdMainloopXe2< MainloopDispatchPolicy, @@ -106,7 +161,12 @@ struct FMHAConfig { GmemTiledCopyQ, GmemTiledCopyK, GmemTiledCopyV, - HasRotary>; + HasRotary, + UseScale, + F8kvF16mma, + TensorScaleQ, + TensorScaleK, + TensorScaleV>; using CollectiveEpilogue = cutlass::fmha::collective::FMHAFwdEpilogueXe2< CollectiveMainloop, @@ -191,7 +251,14 @@ struct FMHAConfig { args.knew_batch_stride, args.knew_head_stride, args.knew_row_stride, reinterpret_cast(args.vnew), args.vnew_batch_stride, args.vnew_head_stride, args.vnew_row_stride, - args.seqlen_knew}, + args.seqlen_knew, + reinterpret_cast(args.scale_q), + StrideScaleQ{}, + reinterpret_cast(args.scale_k), + StrideScaleK{}, + reinterpret_cast(args.scale_v), + StrideScaleV{}, + GROUP_SIZE}, {static_cast(args.sm_scale), args.window_size_left, args.window_size_right, args.p_dropout, @@ -205,7 +272,9 @@ struct FMHAConfig { reinterpret_cast(args.rotary_cos), reinterpret_cast(args.rotary_sin), args.rotary_dim, - args.is_rotary_interleaved}, + args.is_rotary_interleaved, + args.scale_k_scalar, + args.scale_v_scalar}, {}, hw_info}; @@ -282,6 +351,7 @@ void policy_dispatch_fp16(sycl::queue& queue, const fmha_fwd_args_t& args) { typename chunk_policy::SubgroupLayoutQK, void, PipelineStages, + false, false, half_t, half_t, half_t, half_t>; Config::template xe2_dispatch(queue, args); } @@ -296,8 +366,136 @@ void policy_dispatch_bf16(sycl::queue& queue, const fmha_fwd_args_t& args) { typename chunk_policy::SubgroupLayoutQK, void, PipelineStages, + false, false, bfloat16_t, bfloat16_t, bfloat16_t, bfloat16_t>; Config::template xe2_dispatch(queue, args); } +// --------------------------------------------------------------------------- +// MXFP (block-scaled FP8/FP4) dispatch. +// Prefill (SGTileQ >= 8) uses block-scaled BDPAS (UseScale=true); decode +// (SGTileQ < 8) falls back to plain DPAS with UseScale=false. +// +// These dispatchers are UNCONDITIONAL (no SYCL_INTEL_TARGET guard): the host +// pass is always SYCL_INTEL_TARGET==20, and the host must own these symbols so +// they are callable at runtime on CRI. The actual block-scaled device code is +// gated inside FMHAConfig (MMA-atom selection) and the mainloop GEMM sections +// via `#if SYCL_INTEL_TARGET==35`, with compilable BMG fallbacks. On the BMG +// device pass these kernels are emitted (as unreachable placeholders) but never +// executed, since MXFP is only run on CRI hardware. +// --------------------------------------------------------------------------- +template +void policy_dispatch_mxfp_e5m2(sycl::queue& queue, const fmha_fwd_args_t& args) { + constexpr int SGTileQ = get<0>(shape_div( + typename chunk_policy::ShapeQK{}, + shape(typename chunk_policy::SubgroupLayoutQK{})))(); + using ElementType = cutlass::mx_float8_t; + using ElementQ = typename ElementType::DataType; + using ElementK = typename ElementType::DataType; + using ElementV = typename ElementType::DataType; + using ElementScale = typename ElementType::ScaleFactorType; + if constexpr (SGTileQ >= 8) { + using Config = FMHAConfig< + typename chunk_policy::ShapeQK, + typename chunk_policy::ShapePV, + typename chunk_policy::ShapeOut, + typename chunk_policy::SubgroupLayoutQK, + void, + PipelineStages, + true, // UseScale (prefill) + false, // F8kvF16mma + ElementQ, ElementK, ElementV, float, ElementScale>; + Config::template xe2_dispatch(queue, args); + } else { + using Config = FMHAConfig< + typename chunk_policy::ShapeQK, + typename chunk_policy::ShapePV, + typename chunk_policy::ShapeOut, + typename chunk_policy::SubgroupLayoutQK, + void, + PipelineStages, + false, // UseScale (decode) + false, // F8kvF16mma + ElementQ, ElementK, ElementV, float, ElementScale>; + Config::template xe2_dispatch(queue, args); + } +} + +template +void policy_dispatch_mxfp_e4m3(sycl::queue& queue, const fmha_fwd_args_t& args) { + constexpr int SGTileQ = get<0>(shape_div( + typename chunk_policy::ShapeQK{}, + shape(typename chunk_policy::SubgroupLayoutQK{})))(); + using ElementType = cutlass::mx_float8_t; + using ElementQ = typename ElementType::DataType; + using ElementK = typename ElementType::DataType; + using ElementV = typename ElementType::DataType; + using ElementScale = typename ElementType::ScaleFactorType; + if constexpr (SGTileQ >= 8) { + using Config = FMHAConfig< + typename chunk_policy::ShapeQK, + typename chunk_policy::ShapePV, + typename chunk_policy::ShapeOut, + typename chunk_policy::SubgroupLayoutQK, + void, + PipelineStages, + true, // UseScale (prefill) + false, // F8kvF16mma + ElementQ, ElementK, ElementV, float, ElementScale>; + Config::template xe2_dispatch(queue, args); + } else { + using Config = FMHAConfig< + typename chunk_policy::ShapeQK, + typename chunk_policy::ShapePV, + typename chunk_policy::ShapeOut, + typename chunk_policy::SubgroupLayoutQK, + void, + PipelineStages, + false, // UseScale (decode) + false, // F8kvF16mma + ElementQ, ElementK, ElementV, float, ElementScale>; + Config::template xe2_dispatch(queue, args); + } +} + +template +void policy_dispatch_mxfp_e2m1(sycl::queue& queue, const fmha_fwd_args_t& args) { + constexpr int SGTileQ = get<0>(shape_div( + typename chunk_policy::ShapeQK{}, + shape(typename chunk_policy::SubgroupLayoutQK{})))(); + using ElementType = cutlass::mx_float4_t; + using ElementQ = typename ElementType::DataType; + using ElementK = typename ElementType::DataType; + using ElementV = bfloat16_t; // V uses BF16 for MXFP4 + using ElementScale = typename ElementType::ScaleFactorType; + if constexpr (SGTileQ >= 8) { + using Config = FMHAConfig< + typename chunk_policy::ShapeQK, + typename chunk_policy::ShapePV, + typename chunk_policy::ShapeOut, + typename chunk_policy::SubgroupLayoutQK, + void, + PipelineStages, + true, // UseScale (prefill) + false, // F8kvF16mma + ElementQ, ElementK, ElementV, float, ElementScale>; + Config::template xe2_dispatch(queue, args); + } else { + using Config = FMHAConfig< + typename chunk_policy::ShapeQK, + typename chunk_policy::ShapePV, + typename chunk_policy::ShapeOut, + typename chunk_policy::SubgroupLayoutQK, + void, + PipelineStages, + false, // UseScale (decode) + false, // F8kvF16mma + ElementQ, ElementK, ElementV, float, ElementScale>; + Config::template xe2_dispatch(queue, args); + } +} + // Combined policy_dispatch is now defined inline in fmha_fwd.hpp diff --git a/flash-attn2/flash_attn_xpu/src/fmha_fwd_types.hpp b/flash-attn2/flash_attn_xpu/src/fmha_fwd_types.hpp index 7730facac..4c19a37ad 100644 --- a/flash-attn2/flash_attn_xpu/src/fmha_fwd_types.hpp +++ b/flash-attn2/flash_attn_xpu/src/fmha_fwd_types.hpp @@ -11,11 +11,20 @@ struct fmha_fwd_args_t { void* block_table; void* cu_seqlens_q; void* cu_seqlens_k; + // Scale tensors for MXFP (block-scaled) operations + void* scale_q = nullptr; + void* scale_k = nullptr; + void* scale_v = nullptr; + void* cu_scale_q = nullptr; + void* cu_scale_kv = nullptr; int max_queries; int max_keys; int total_seqlen_q; int total_seqlen_k; float sm_scale; + // Scalar scales for F8 KV with F16 MMA path + float scale_k_scalar = 1.0f; + float scale_v_scalar = 1.0f; int batch_size; int num_heads_q; int num_heads_k; @@ -65,8 +74,14 @@ struct fmha_fwd_args_t { enum class CutlassType { half, - bfloat16 + bfloat16, + mx_float_e5m2, // MXFP8 E5M2 + mx_float_e4m3, // MXFP8 E4M3 + mx_float_e2m1 // MXFP4 E2M1 }; constexpr int PipelineStages_Decode = 1; constexpr int PipelineStages_Prefill = 2; + +// Scale group size for block-scaled (MXFP) operations +constexpr int GROUP_SIZE = 32; diff --git a/flash-attn2/flash_attn_xpu/src/fmha_utils.hpp b/flash-attn2/flash_attn_xpu/src/fmha_utils.hpp index f531800a9..a0d2ddea0 100644 --- a/flash-attn2/flash_attn_xpu/src/fmha_utils.hpp +++ b/flash-attn2/flash_attn_xpu/src/fmha_utils.hpp @@ -11,11 +11,28 @@ if (input.scalar_type() == torch::kBFloat16) { return CutlassType::bfloat16; } + if (input.scalar_type() == torch::kFloat8_e5m2) { + return CutlassType::mx_float_e5m2; + } + if (input.scalar_type() == torch::kFloat8_e4m3fn) { + return CutlassType::mx_float_e4m3; + } + if (input.scalar_type() == torch::kFloat4_e2m1fn_x2) { + return CutlassType::mx_float_e2m1; + } TORCH_INTERNAL_ASSERT( - false, "Current cutlass kernel only support half/bf16 data type."); + false, + "Current cutlass kernel only supports half/bf16/fp8_e5m2/fp8_e4m3/fp4_e2m1 data types."); return {}; // unreachable; silences compiler warning } +/// Helper to detect MXFP (block-scaled) types that require scale tensors. +[[nodiscard]] inline bool requires_scale_tensors(CutlassType cuType) { + return cuType == CutlassType::mx_float_e5m2 || + cuType == CutlassType::mx_float_e4m3 || + cuType == CutlassType::mx_float_e2m1; +} + using namespace cute; struct prefill_policy_head32 { diff --git a/flash-attn2/flash_attn_xpu/src/kernel/fmha_fwd_kernel_xe2.hpp b/flash-attn2/flash_attn_xpu/src/kernel/fmha_fwd_kernel_xe2.hpp index 5f78f6b9b..47c2cb15e 100644 --- a/flash-attn2/flash_attn_xpu/src/kernel/fmha_fwd_kernel_xe2.hpp +++ b/flash-attn2/flash_attn_xpu/src/kernel/fmha_fwd_kernel_xe2.hpp @@ -67,6 +67,13 @@ class XeFMHAFwdKernelXe2 { static constexpr bool HasDropout = CollectiveMainloop::HasDropout; static constexpr bool PagedKV = CollectiveMainloop::PagedKV; + // MXFP (block-scaled) support. + static constexpr bool UseScale = CollectiveMainloop::UseScale; + using ElementScale = typename CollectiveMainloop::ElementScaleQ; + using StrideScaleQ = decltype(stride(typename CollectiveMainloop::TensorScaleQ{})); + using StrideScaleK = decltype(stride(typename CollectiveMainloop::TensorScaleK{})); + using StrideScaleV = decltype(stride(typename CollectiveMainloop::TensorScaleV{})); + using MainloopSharedStorage = typename CollectiveMainloop::SharedStorage; using EpilogueSharedStorage = typename CollectiveEpilogue::SharedStorage; union SharedStorage { @@ -112,6 +119,14 @@ class XeFMHAFwdKernelXe2 { int64_t vnew_head_stride = 0; int64_t vnew_row_stride = 0; int seqlen_knew = 0; + // MXFP (block-scaled) scale tensors. + const ElementScale* scaleQ = nullptr; + StrideScaleQ dScaleQ{}; + const ElementScale* scaleK = nullptr; + StrideScaleK dScaleK{}; + const ElementScale* scaleV = nullptr; + StrideScaleV dScaleV{}; + int group_size = 32; }; using KernelParams = KernelArguments; @@ -482,28 +497,102 @@ class XeFMHAFwdKernelXe2 { int l_coord_kv = is_var_len ? 0 : (PagedKV ? 0 : bidx); CollectiveMainloop mainloop(params.mainloop, shared_storage.mainloop); - mainloop( - Q(_, _, head_q, l_coord_q), - K(_, _, head, l_coord_kv), - V(_, _, head, l_coord_kv), - tArA, - tA_max, - tA_sum, - blk_qv, - k_block0, - k_blocks, - total_blk, - thr_id, - seq_len, - seq_len_qo, - effective_seq_kv, - idx_b, - tile_row_idx, - rows_of_maxima, - head_q, - s.num_heads_q, - q_offset_sg, - p.cache_seqlens ? p.cache_seqlens[idx_b] : 0); + if constexpr (UseScale) { + // MXFP: build block-scale tensors mirroring the Q/K/V layout. + auto scale_q_dim = cute::ceil_div(s.head_size_qk, p.group_size); + auto scale_k_dim = cute::ceil_div(s.head_size_qk, p.group_size); + auto scale_v_dim = cute::ceil_div(total_seqlen_kv, p.group_size); + + auto shape_scale_Q = + make_shape(seq_len_qo, scale_q_dim, s.num_heads_q, batch_dim_q); + auto shape_scale_K = + make_shape(total_seqlen_kv, scale_k_dim, s.num_heads_kv, batch_dim_kv); + auto shape_scale_V = + make_shape(s.head_size_vo, scale_v_dim, s.num_heads_kv, batch_dim_kv); + + int offset_scaleQ = 0, offset_scaleK = 0, offset_scaleV = 0; + if constexpr (is_var_len) { + auto qo_cumulative = s.seq_len_qo.cumulative_length; + auto kv_cumulative = s.seq_len_kv.cumulative_length; + offset_scaleQ = s.num_heads_q * scale_q_dim * qo_cumulative[idx_b]; + offset_scaleK = s.num_heads_kv * scale_k_dim * kv_cumulative[idx_b]; + auto kv_scale_cumulative = s.seq_len_kv.cumulative_scale_length; + if (kv_scale_cumulative != nullptr) { + offset_scaleV = + s.num_heads_kv * s.head_size_vo * kv_scale_cumulative[idx_b]; + } else { + offset_scaleV = s.num_heads_kv * scale_v_dim * kv_cumulative[idx_b]; + } + } + + auto layout_scaleQ = is_var_len + ? make_ordered_layout(shape_scale_Q, Step<_0, _1, _2, _3>{}) + : make_layout(shape_scale_Q, p.dScaleQ); + auto layout_scaleK = is_var_len + ? make_ordered_layout(shape_scale_K, Step<_0, _1, _2, _3>{}) + : make_layout(shape_scale_K, p.dScaleK); + auto layout_scaleV = is_var_len + ? make_ordered_layout(shape_scale_V, Step<_0, _1, _2, _3>{}) + : make_layout(shape_scale_V, p.dScaleV); + + auto dcScaleQ = const_cast(p.scaleQ + offset_scaleQ); + auto dcScaleK = const_cast(p.scaleK + offset_scaleK); + auto dcScaleV = const_cast(p.scaleV + offset_scaleV); + + Tensor ScaleQ = make_tensor(make_gmem_ptr(dcScaleQ), layout_scaleQ); + Tensor ScaleK = make_tensor(make_gmem_ptr(dcScaleK), layout_scaleK); + Tensor ScaleV = make_tensor(make_gmem_ptr(dcScaleV), layout_scaleV); + + mainloop( + Q(_, _, head_q, l_coord_q), + K(_, _, head, l_coord_kv), + V(_, _, head, l_coord_kv), + tArA, + tA_max, + tA_sum, + blk_qv, + k_block0, + k_blocks, + total_blk, + thr_id, + seq_len, + seq_len_qo, + effective_seq_kv, + idx_b, + tile_row_idx, + rows_of_maxima, + head_q, + s.num_heads_q, + q_offset_sg, + p.cache_seqlens ? p.cache_seqlens[idx_b] : 0, + l_coord_q, + ScaleQ(_, _, head_q, l_coord_q), + ScaleK(_, _, head, l_coord_kv), + ScaleV(_, _, head, l_coord_kv)); + } else { + mainloop( + Q(_, _, head_q, l_coord_q), + K(_, _, head, l_coord_kv), + V(_, _, head, l_coord_kv), + tArA, + tA_max, + tA_sum, + blk_qv, + k_block0, + k_blocks, + total_blk, + thr_id, + seq_len, + seq_len_qo, + effective_seq_kv, + idx_b, + tile_row_idx, + rows_of_maxima, + head_q, + s.num_heads_q, + q_offset_sg, + p.cache_seqlens ? p.cache_seqlens[idx_b] : 0); + } if constexpr ( !is_empty_v && diff --git a/flash-attn2/torch-ext/torch_binding.cpp b/flash-attn2/torch-ext/torch_binding.cpp index b277ea46a..11f2656ed 100644 --- a/flash-attn2/torch-ext/torch_binding.cpp +++ b/flash-attn2/torch-ext/torch_binding.cpp @@ -27,7 +27,10 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { "int window_size_right, " "float softcap, " "bool return_softmax, " - "Generator? gen_) -> Tensor[]"); + "Generator? gen_, " + "Tensor? scale_q_=None, " + "Tensor? scale_k_=None, " + "Tensor? scale_v_=None) -> Tensor[]"); #if defined(CUDA_KERNEL) ops.impl("fwd", torch::kCUDA, &mha_fwd); #elif defined(XPU_KERNEL) @@ -55,7 +58,12 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { "int window_size_right, " "float softcap, " "bool return_softmax, " - "Generator? gen_) -> Tensor[]"); + "Generator? gen_, " + "Tensor? scale_q_=None, " + "Tensor? scale_k_=None, " + "Tensor? scale_v_=None, " + "Tensor? cu_scale_q_=None, " + "Tensor? cu_scale_kv_=None) -> Tensor[]"); #if defined(CUDA_KERNEL) ops.impl("varlen_fwd", torch::kCUDA, &mha_varlen_fwd); #elif defined(XPU_KERNEL) diff --git a/flash-attn2/torch-ext/torch_binding.h b/flash-attn2/torch-ext/torch_binding.h index 5dc073971..43bdbb5cf 100644 --- a/flash-attn2/torch-ext/torch_binding.h +++ b/flash-attn2/torch-ext/torch_binding.h @@ -17,7 +17,10 @@ mha_fwd( const int64_t window_size_right, const double softcap, const bool return_softmax, - c10::optional gen_); + c10::optional gen_, + c10::optional scale_q_ = c10::nullopt, + c10::optional scale_k_ = c10::nullopt, + c10::optional scale_v_ = c10::nullopt); std::vector mha_varlen_fwd( @@ -42,7 +45,12 @@ mha_varlen_fwd( int64_t window_size_right, const double softcap, const bool return_softmax, - std::optional gen_); + std::optional gen_, + std::optional scale_q_ = c10::nullopt, + std::optional scale_k_ = c10::nullopt, + std::optional scale_v_ = c10::nullopt, + std::optional cu_scale_q_ = c10::nullopt, + std::optional cu_scale_kv_ = c10::nullopt); std::vector mha_bwd(const torch::Tensor &dout, // batch_size x seqlen_q x num_heads, x multiple_of(head_size_og, 8) From 88606da03e9dad7127d842ad332144b7be098b21 Mon Sep 17 00:00:00 2001 From: jiqing-feng Date: Thu, 16 Jul 2026 20:21:37 -0700 Subject: [PATCH 03/12] flash-attn2: fix out_dtype ternary type mismatch (ScalarType vs TypeMeta) for MXFP --- flash-attn2/flash_attn_xpu/flash_api.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/flash-attn2/flash_attn_xpu/flash_api.cpp b/flash-attn2/flash_attn_xpu/flash_api.cpp index 65b5ff092..bb8dc351d 100644 --- a/flash-attn2/flash_attn_xpu/flash_api.cpp +++ b/flash-attn2/flash_attn_xpu/flash_api.cpp @@ -115,7 +115,7 @@ mha_fwd( } // Low-precision (FP8/FP4) inputs accumulate in float32 and produce a float32 output. - auto out_dtype = is_mxfp ? torch::kFloat32 : q_dtype; + auto out_dtype = is_mxfp ? torch::kFloat32 : c10::typeMetaToScalarType(q_dtype); at::Tensor out_padded; if (out_.has_value()) { @@ -493,7 +493,7 @@ mha_varlen_fwd( at::Tensor k_padded = maybe_pad(k); at::Tensor v_padded = maybe_pad(v); - auto out_dtype = is_mxfp ? torch::kFloat32 : q_dtype; + auto out_dtype = is_mxfp ? torch::kFloat32 : c10::typeMetaToScalarType(q_dtype); at::Tensor out_padded; if (out_.has_value()) { out_padded = maybe_pad(out_.value()); From a7b5fc21afa635adbe5950a8b09538f3123c9e6c Mon Sep 17 00:00:00 2001 From: jiqing-feng Date: Thu, 16 Jul 2026 20:27:05 -0700 Subject: [PATCH 04/12] flash-attn2: move MXFP scale fields to end of fmha_fwd_args_t to fix positional init narrowing --- .../flash_attn_xpu/src/fmha_fwd_types.hpp | 19 ++++++++++--------- 1 file changed, 10 insertions(+), 9 deletions(-) diff --git a/flash-attn2/flash_attn_xpu/src/fmha_fwd_types.hpp b/flash-attn2/flash_attn_xpu/src/fmha_fwd_types.hpp index 4c19a37ad..ccb481bad 100644 --- a/flash-attn2/flash_attn_xpu/src/fmha_fwd_types.hpp +++ b/flash-attn2/flash_attn_xpu/src/fmha_fwd_types.hpp @@ -11,20 +11,11 @@ struct fmha_fwd_args_t { void* block_table; void* cu_seqlens_q; void* cu_seqlens_k; - // Scale tensors for MXFP (block-scaled) operations - void* scale_q = nullptr; - void* scale_k = nullptr; - void* scale_v = nullptr; - void* cu_scale_q = nullptr; - void* cu_scale_kv = nullptr; int max_queries; int max_keys; int total_seqlen_q; int total_seqlen_k; float sm_scale; - // Scalar scales for F8 KV with F16 MMA path - float scale_k_scalar = 1.0f; - float scale_v_scalar = 1.0f; int batch_size; int num_heads_q; int num_heads_k; @@ -70,6 +61,16 @@ struct fmha_fwd_args_t { void* rotary_sin = nullptr; int rotary_dim = 0; bool is_rotary_interleaved = true; + + // Scale tensors for MXFP (block-scaled) operations + void* scale_q = nullptr; + void* scale_k = nullptr; + void* scale_v = nullptr; + void* cu_scale_q = nullptr; + void* cu_scale_kv = nullptr; + // Scalar scales for F8 KV with F16 MMA path + float scale_k_scalar = 1.0f; + float scale_v_scalar = 1.0f; }; enum class CutlassType { From 633fc1d491e6ba04e15936c441831925df0bcd5b Mon Sep 17 00:00:00 2001 From: jiqing-feng Date: Thu, 16 Jul 2026 20:34:59 -0700 Subject: [PATCH 05/12] flash-attn2: fix std::nullopt defaults on impl decls (include , drop defaults for incomplete at::Tensor) --- flash-attn2/flash_attn_xpu/src/fmha_fwd.hpp | 50 ++++++++++----------- 1 file changed, 24 insertions(+), 26 deletions(-) diff --git a/flash-attn2/flash_attn_xpu/src/fmha_fwd.hpp b/flash-attn2/flash_attn_xpu/src/fmha_fwd.hpp index 8ba7c9c09..c054d0552 100644 --- a/flash-attn2/flash_attn_xpu/src/fmha_fwd.hpp +++ b/flash-attn2/flash_attn_xpu/src/fmha_fwd.hpp @@ -1,5 +1,7 @@ #pragma once +#include + #include "fmha_fwd_types.hpp" namespace sycl { @@ -12,10 +14,6 @@ namespace at { class Tensor; } -namespace std { - template class optional; -} - struct prefill_policy_head32; struct prefill_policy_head64; struct prefill_policy_head96; @@ -226,19 +224,19 @@ void cutlass_fmha_fwd_varlen_impl( bool is_paged, bool is_causal, bool is_local, - float p_dropout = 0.0f, - uint64_t philox_seed = 0, - uint64_t philox_offset = 0, - void* rng_state = nullptr, - void* s_dmask = nullptr, - int seqlen_q_rounded = 0, - int seqlen_k_rounded = 0, + float p_dropout, + uint64_t philox_seed, + uint64_t philox_offset, + void* rng_state, + void* s_dmask, + int seqlen_q_rounded, + int seqlen_k_rounded, // MXFP (block-scaled) scale factor tensors (optional). - const std::optional& scale_q = std::nullopt, - const std::optional& scale_k = std::nullopt, - const std::optional& scale_v = std::nullopt, - const std::optional& cu_scale_q = std::nullopt, - const std::optional& cu_scale_kv = std::nullopt); + const std::optional& scale_q, + const std::optional& scale_k, + const std::optional& scale_v, + const std::optional& cu_scale_q, + const std::optional& cu_scale_kv); void cutlass_fmha_fwd_fix_impl( sycl::queue& queue, @@ -252,17 +250,17 @@ void cutlass_fmha_fwd_fix_impl( int window_size_right, bool is_causal, bool is_local, - float p_dropout = 0.0f, - uint64_t philox_seed = 0, - uint64_t philox_offset = 0, - void* rng_state = nullptr, - void* s_dmask = nullptr, - int seqlen_q_rounded = 0, - int seqlen_k_rounded = 0, + float p_dropout, + uint64_t philox_seed, + uint64_t philox_offset, + void* rng_state, + void* s_dmask, + int seqlen_q_rounded, + int seqlen_k_rounded, // MXFP (block-scaled) scale factor tensors (optional). - const std::optional& scale_q = std::nullopt, - const std::optional& scale_k = std::nullopt, - const std::optional& scale_v = std::nullopt); + const std::optional& scale_q, + const std::optional& scale_k, + const std::optional& scale_v); void cutlass_fmha_fwd_kvcache_impl( sycl::queue& queue, From 49437f33c36ba6e593c267d8ddbf688dd81581d2 Mon Sep 17 00:00:00 2001 From: jiqing-feng Date: Thu, 16 Jul 2026 22:17:00 -0700 Subject: [PATCH 06/12] flash-attn2: unify fmha_fwd.hpp on verified fix (ATen Tensor header + keep std::nullopt defaults) --- flash-attn2/flash_attn_xpu/src/fmha_fwd.hpp | 52 ++++++++++----------- 1 file changed, 25 insertions(+), 27 deletions(-) diff --git a/flash-attn2/flash_attn_xpu/src/fmha_fwd.hpp b/flash-attn2/flash_attn_xpu/src/fmha_fwd.hpp index c054d0552..eb7b86565 100644 --- a/flash-attn2/flash_attn_xpu/src/fmha_fwd.hpp +++ b/flash-attn2/flash_attn_xpu/src/fmha_fwd.hpp @@ -1,8 +1,10 @@ #pragma once +#include "fmha_fwd_types.hpp" + #include -#include "fmha_fwd_types.hpp" +#include namespace sycl { inline namespace _V1 { @@ -10,10 +12,6 @@ namespace sycl { } } -namespace at { - class Tensor; -} - struct prefill_policy_head32; struct prefill_policy_head64; struct prefill_policy_head96; @@ -224,19 +222,19 @@ void cutlass_fmha_fwd_varlen_impl( bool is_paged, bool is_causal, bool is_local, - float p_dropout, - uint64_t philox_seed, - uint64_t philox_offset, - void* rng_state, - void* s_dmask, - int seqlen_q_rounded, - int seqlen_k_rounded, + float p_dropout = 0.0f, + uint64_t philox_seed = 0, + uint64_t philox_offset = 0, + void* rng_state = nullptr, + void* s_dmask = nullptr, + int seqlen_q_rounded = 0, + int seqlen_k_rounded = 0, // MXFP (block-scaled) scale factor tensors (optional). - const std::optional& scale_q, - const std::optional& scale_k, - const std::optional& scale_v, - const std::optional& cu_scale_q, - const std::optional& cu_scale_kv); + const std::optional& scale_q = std::nullopt, + const std::optional& scale_k = std::nullopt, + const std::optional& scale_v = std::nullopt, + const std::optional& cu_scale_q = std::nullopt, + const std::optional& cu_scale_kv = std::nullopt); void cutlass_fmha_fwd_fix_impl( sycl::queue& queue, @@ -250,17 +248,17 @@ void cutlass_fmha_fwd_fix_impl( int window_size_right, bool is_causal, bool is_local, - float p_dropout, - uint64_t philox_seed, - uint64_t philox_offset, - void* rng_state, - void* s_dmask, - int seqlen_q_rounded, - int seqlen_k_rounded, + float p_dropout = 0.0f, + uint64_t philox_seed = 0, + uint64_t philox_offset = 0, + void* rng_state = nullptr, + void* s_dmask = nullptr, + int seqlen_q_rounded = 0, + int seqlen_k_rounded = 0, // MXFP (block-scaled) scale factor tensors (optional). - const std::optional& scale_q, - const std::optional& scale_k, - const std::optional& scale_v); + const std::optional& scale_q = std::nullopt, + const std::optional& scale_k = std::nullopt, + const std::optional& scale_v = std::nullopt); void cutlass_fmha_fwd_kvcache_impl( sycl::queue& queue, From eda130823b1b0319de4eba668383d8b4126965f2 Mon Sep 17 00:00:00 2001 From: jiqing-feng Date: Tue, 21 Jul 2026 00:57:59 -0700 Subject: [PATCH 07/12] flash-attn2: split MXFP kernels into spir64_gen-only component for tri-target Move the 16 block-scaled MXFP TUs out of flash_attn_xpu into a dedicated flash_attn_xpu_mxfp component compiled with -fsycl-targets=spir64_gen and -D__SYCL_TARGET_INTEL_GPU_CRI__=1. MXFP BDPAS AOT works only on cri, so the fp16/bf16 kernels stay multi-target (pvc/bmg/cri) while MXFP is spir64_gen-only and AOT-lowered to cri at link. All components link into one .so. --- flash-attn2/build.toml | 48 ++++++++++++++++++++++++++++-------------- 1 file changed, 32 insertions(+), 16 deletions(-) diff --git a/flash-attn2/build.toml b/flash-attn2/build.toml index d93fae573..b90a176c6 100644 --- a/flash-attn2/build.toml +++ b/flash-attn2/build.toml @@ -78,22 +78,6 @@ src = [ "flash_attn_xpu/src/flash_fwd_hdim256_kvcache_paged_bf16.cpp", "flash_attn_xpu/src/flash_fwd_hdim512_kvcache_paged_fp16.cpp", "flash_attn_xpu/src/flash_fwd_hdim512_kvcache_paged_bf16.cpp", - "flash_attn_xpu/src/flash_fwd_hdim32_varlen_mxfp.cpp", - "flash_attn_xpu/src/flash_fwd_hdim64_varlen_mxfp.cpp", - "flash_attn_xpu/src/flash_fwd_hdim96_varlen_mxfp.cpp", - "flash_attn_xpu/src/flash_fwd_hdim128_varlen_mxfp.cpp", - "flash_attn_xpu/src/flash_fwd_hdim160_varlen_mxfp.cpp", - "flash_attn_xpu/src/flash_fwd_hdim192_varlen_mxfp.cpp", - "flash_attn_xpu/src/flash_fwd_hdim256_varlen_mxfp.cpp", - "flash_attn_xpu/src/flash_fwd_hdim512_varlen_mxfp.cpp", - "flash_attn_xpu/src/flash_fwd_hdim32_fix_mxfp.cpp", - "flash_attn_xpu/src/flash_fwd_hdim64_fix_mxfp.cpp", - "flash_attn_xpu/src/flash_fwd_hdim96_fix_mxfp.cpp", - "flash_attn_xpu/src/flash_fwd_hdim128_fix_mxfp.cpp", - "flash_attn_xpu/src/flash_fwd_hdim160_fix_mxfp.cpp", - "flash_attn_xpu/src/flash_fwd_hdim192_fix_mxfp.cpp", - "flash_attn_xpu/src/flash_fwd_hdim256_fix_mxfp.cpp", - "flash_attn_xpu/src/flash_fwd_hdim512_fix_mxfp.cpp", "flash_attn_xpu/src/fmha_bwd_types.hpp", "flash_attn_xpu/src/fmha_bwd.hpp", "flash_attn_xpu/src/fmha_bwd_impl.hpp", @@ -143,6 +127,38 @@ src = [ "flash_attn_xpu/src/kernel/fmha_bwd_kernel.hpp", ] +# Block-scaled MXFP kernels are gated behind __SYCL_TARGET_INTEL_GPU_CRI__ and +# only build for cri, so compile them spir64_gen-only with that macro. The +# spir64_gen image is AOT-lowered to cri by the tri-target link. +[kernel.flash_attn_xpu_mxfp] +backend = "xpu" +depends = [ + "torch", + "sycl_tla", +] +sycl-flags = [ + "-fsycl-targets=spir64_gen", + "-D__SYCL_TARGET_INTEL_GPU_CRI__=1", +] +src = [ + "flash_attn_xpu/src/flash_fwd_hdim32_varlen_mxfp.cpp", + "flash_attn_xpu/src/flash_fwd_hdim64_varlen_mxfp.cpp", + "flash_attn_xpu/src/flash_fwd_hdim96_varlen_mxfp.cpp", + "flash_attn_xpu/src/flash_fwd_hdim128_varlen_mxfp.cpp", + "flash_attn_xpu/src/flash_fwd_hdim160_varlen_mxfp.cpp", + "flash_attn_xpu/src/flash_fwd_hdim192_varlen_mxfp.cpp", + "flash_attn_xpu/src/flash_fwd_hdim256_varlen_mxfp.cpp", + "flash_attn_xpu/src/flash_fwd_hdim512_varlen_mxfp.cpp", + "flash_attn_xpu/src/flash_fwd_hdim32_fix_mxfp.cpp", + "flash_attn_xpu/src/flash_fwd_hdim64_fix_mxfp.cpp", + "flash_attn_xpu/src/flash_fwd_hdim96_fix_mxfp.cpp", + "flash_attn_xpu/src/flash_fwd_hdim128_fix_mxfp.cpp", + "flash_attn_xpu/src/flash_fwd_hdim160_fix_mxfp.cpp", + "flash_attn_xpu/src/flash_fwd_hdim192_fix_mxfp.cpp", + "flash_attn_xpu/src/flash_fwd_hdim256_fix_mxfp.cpp", + "flash_attn_xpu/src/flash_fwd_hdim512_fix_mxfp.cpp", +] + [kernel.flash_attn_cpu_avx512] backend = "cpu" cxx-flags = [ From 1676429fc0d304ec66170358271d543ab7c18b9c Mon Sep 17 00:00:00 2001 From: jiqing-feng Date: Mon, 27 Jul 2026 19:48:52 -0700 Subject: [PATCH 08/12] megablocks: add MXFP8 support to the XPU grouped GEMM Port the block-scaled FP8 path from the Intel reference (607c031) so MoE weights can be stored as float8_e4m3fn/float8_e5m2 with E8M0 block scales, alongside the existing int4 and MXFP4 paths. MXFP8 reuses the generic block-scaled kernel (xe_gemm_4bits) rather than adding a new one. What is new: the is_B_mxfp8 type trait, the E8M0 scale decode branch, the 3D [E, N, K/group_size] scale pointer offset, the unpacked [E, N, K] weight layout, and the MXFP8LauncherCallER dispatch. The dispatch is mutually exclusive at compile time (is_B_mxfp8 requires ElementB in {e4m3, e5m2} *and* ElementS == uint8_t, while w8a16-fp8 uses a float ElementS and int4 a uint8_t ElementB), so is_B_block_scaled reduces to the previous is_B_4bits for every pre-existing path. Validated on the CRI simulator against an fp32 CPU reference: the three MXFP8 cases (e4m3, e5m2, e4m3+bias) land in the same residual range as the unquantized w16a16 baseline, and the w16a16 / w8a16-fp8 / int4 / MXFP4 paths are unchanged. --- .../grouped_gemm/grouped_gemm_interface.cpp | 6 +- .../grouped_gemm/grouped_gemm_interface.h | 3 +- .../csrc_xpu/grouped_gemm/xe_2/gemm_xe2.hpp | 4 +- .../grouped_gemm/xe_2/grouped_gemm_xe2.cpp | 6 +- .../grouped_gemm/xe_2/grouped_gemm_xe2.h | 3 +- .../grouped_gemm/xe_2/grouped_gemm_xe2.hpp | 16 ++++- .../xe_2/grouped_gemm_xe2_interface.hpp | 62 ++++++++++++++++++- .../torch-ext/megablocks/xpu_fused_moe.py | 34 ++++++---- megablocks/torch-ext/torch_binding.cpp | 2 +- 9 files changed, 113 insertions(+), 23 deletions(-) diff --git a/megablocks/csrc_xpu/grouped_gemm/grouped_gemm_interface.cpp b/megablocks/csrc_xpu/grouped_gemm/grouped_gemm_interface.cpp index bf6ccf6e8..b97adaa6a 100644 --- a/megablocks/csrc_xpu/grouped_gemm/grouped_gemm_interface.cpp +++ b/megablocks/csrc_xpu/grouped_gemm/grouped_gemm_interface.cpp @@ -13,7 +13,8 @@ torch::Tensor cutlass_grouped_gemm_interface( int64_t K, int64_t num_experts, bool is_B_int4, - bool is_B_mxfp4) { + bool is_B_mxfp4, + bool is_B_mxfp8) { // Use XE2 cutlass kernel return cutlass_grouped_gemm_xe2( ptr_A, @@ -26,5 +27,6 @@ torch::Tensor cutlass_grouped_gemm_interface( K, num_experts, is_B_int4, - is_B_mxfp4); + is_B_mxfp4, + is_B_mxfp8); } diff --git a/megablocks/csrc_xpu/grouped_gemm/grouped_gemm_interface.h b/megablocks/csrc_xpu/grouped_gemm/grouped_gemm_interface.h index b25b71d12..7dfcd1eb2 100644 --- a/megablocks/csrc_xpu/grouped_gemm/grouped_gemm_interface.h +++ b/megablocks/csrc_xpu/grouped_gemm/grouped_gemm_interface.h @@ -11,4 +11,5 @@ torch::Tensor cutlass_grouped_gemm_interface( int64_t K, int64_t num_experts, bool is_B_int4, - bool is_B_mxfp4); \ No newline at end of file + bool is_B_mxfp4, + bool is_B_mxfp8); \ No newline at end of file diff --git a/megablocks/csrc_xpu/grouped_gemm/xe_2/gemm_xe2.hpp b/megablocks/csrc_xpu/grouped_gemm/xe_2/gemm_xe2.hpp index af916201c..dbe1e09e1 100644 --- a/megablocks/csrc_xpu/grouped_gemm/xe_2/gemm_xe2.hpp +++ b/megablocks/csrc_xpu/grouped_gemm/xe_2/gemm_xe2.hpp @@ -366,7 +366,9 @@ CUTE_DEVICE void xe_gemm_4bits( scale = Scales [(n_tile_start + n_sg_start + sg_local_n) * group_num + group_idx]; - } else if constexpr (std::is_same_v) { + } else if constexpr (std::is_same_v || + std::is_same_v || + std::is_same_v) { uint32_t scale_u32 = Scales [(n_tile_start + n_sg_start + sg_local_n) * group_num + diff --git a/megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2.cpp b/megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2.cpp index 2d6c4b29a..d2a2f9714 100644 --- a/megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2.cpp +++ b/megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2.cpp @@ -13,7 +13,8 @@ torch::Tensor cutlass_grouped_gemm_xe2( int64_t K, int64_t num_experts, bool is_B_int4, - bool is_B_mxfp4) { + bool is_B_mxfp4, + bool is_B_mxfp8) { return MoE::cutlass_grouped_gemm_xe2_impl( ptr_A, ptr_B, @@ -25,5 +26,6 @@ torch::Tensor cutlass_grouped_gemm_xe2( K, num_experts, is_B_int4, - is_B_mxfp4); + is_B_mxfp4, + is_B_mxfp8); } \ No newline at end of file diff --git a/megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2.h b/megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2.h index feb722c6e..ea2d9a362 100644 --- a/megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2.h +++ b/megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2.h @@ -11,4 +11,5 @@ torch::Tensor cutlass_grouped_gemm_xe2( int64_t K, int64_t num_experts, bool is_B_int4, - bool is_B_mxfp4); \ No newline at end of file + bool is_B_mxfp4, + bool is_B_mxfp8); \ No newline at end of file diff --git a/megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2.hpp b/megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2.hpp index be7386409..0e8f6d458 100644 --- a/megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2.hpp +++ b/megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2.hpp @@ -90,7 +90,12 @@ CUTE_DEVICE void MoEGEMM( (!std::is_same_v); static constexpr bool is_B_mxfp4 = (std::is_same_v) && (std::is_same_v); + static constexpr bool is_B_mxfp8 = + (std::is_same_v || + std::is_same_v) && + (std::is_same_v); static constexpr bool is_B_4bits = std::is_same_v; + static constexpr bool is_B_block_scaled = is_B_4bits || is_B_mxfp8; auto item = sycl::ext::oneapi::this_work_item::get_nd_item<3>(); auto wg_tile = mma.tile_mnk(); @@ -146,6 +151,12 @@ CUTE_DEVICE void MoEGEMM( if constexpr (is_B_4bits) { ptr_Scales_curr_batch = const_cast(Scales) + B_offset * 2 / group_size; + } else if constexpr (is_B_mxfp8) { + ptr_Scales_curr_batch = + const_cast(Scales) + + static_cast(expert_id) * + static_cast(gemm_n) * + (static_cast(gemm_k) / group_size); } ElementBI* ptr_Bias_curr_batch = nullptr; if (Bias != static_cast(nullptr)) { @@ -161,6 +172,9 @@ CUTE_DEVICE void MoEGEMM( } else if constexpr (is_B_mxfp4) { return make_moe_tensor( reinterpret_cast(ptr_B_curr_batch), gemm_n, gemm_k); + } else if constexpr (is_B_mxfp8) { + return make_moe_tensor( + ptr_B_curr_batch, gemm_n, gemm_k); } else { return make_moe_tensor( ptr_B_curr_batch, gemm_n, gemm_k); @@ -174,7 +188,7 @@ CUTE_DEVICE void MoEGEMM( int m_coord = (group_m_id - pre_tiles); auto tile_coord = make_coord(m_coord, n_coord, _, 0); - if constexpr (is_B_4bits) { + if constexpr (is_B_block_scaled) { #define XE_GEMM_4BITS_CALLER(GroupSize) \ xe_gemm_4bits( \ A_tensor, \ diff --git a/megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2_interface.hpp b/megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2_interface.hpp index 6a0a6d312..5ab94efdd 100644 --- a/megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2_interface.hpp +++ b/megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2_interface.hpp @@ -50,11 +50,11 @@ #include "cutlass/tensor_ref.h" #include "cutlass/util/GPU_Clock.hpp" #include "cutlass/util/device_memory.h" +#include "cutlass/util/sycl_event_manager.hpp" #include "cutlass/util/initialize_block.hpp" #include "cutlass/util/reference/device/gemm_complex.h" #include "cutlass/util/reference/device/tensor_compare.h" #include "cutlass/util/reference/host/tensor_fill.h" -#include "cutlass/util/sycl_event_manager.hpp" #include "gemm_xe2_policy.hpp" #include "grouped_gemm_xe2.hpp" @@ -173,7 +173,8 @@ at::Tensor cutlass_grouped_gemm_xe2_impl( int64_t K, int64_t num_experts, bool is_B_int4, - bool is_B_mxfp4) { + bool is_B_mxfp4, + bool is_B_mxfp8) { auto& dpcpp_queue = at::xpu::getCurrentXPUStream(ptr_A.device().index()).queue(); auto A_dtype = ptr_A.dtype(); @@ -206,6 +207,10 @@ at::Tensor cutlass_grouped_gemm_xe2_impl( if (is_B_int4 || is_B_mxfp4) { B_K = ptr_B.size(2) * 2; B_N = ptr_B.size(1); + } else if (is_B_mxfp8) { + // MXFP8: weights stored as [E, N, K] (no packing), dtype is fp8 + B_K = ptr_B.size(2); + B_N = ptr_B.size(1); } int D_total_M = ptr_D.size(0); @@ -298,6 +303,59 @@ at::Tensor cutlass_grouped_gemm_xe2_impl( W4A16LauncherCallER(policy); } #undef W4A16LauncherCallER + } else if (is_B_mxfp8) { + // MXFP8: block-scaled FP8 with e8m0 scales + TORCH_CHECK(ptr_scales.has_value(), "mxfp8 grouped gemm must have scales"); + TORCH_CHECK(ptr_scales->is_contiguous(), "ptr_scales must be contiguous"); + TORCH_CHECK( + ptr_scales->dim() == 3, + "ptr_scales of mxfp8 must be 3D [num_experts, N, K // group_size]"); + TORCH_CHECK( + ptr_scales->size(0) == num_experts, + "ptr_scales.size(0) of mxfp8 must match num_experts"); + TORCH_CHECK( + ptr_scales->size(1) == N, + "ptr_scales.size(1) of mxfp8 must match N"); + TORCH_CHECK( + K % ptr_scales->size(2) == 0, + "K must be divisible by ptr_scales.size(2) for mxfp8"); + int group_num = ptr_scales->size(2); + group_size = K / group_num; + + TORCH_CHECK( + group_size == 32 || group_size == 64 || group_size == 128, + "mxfp8 group_size must be 32, 64 or 128"); + +#define MXFP8LauncherCallER(policy) \ + if (B_dtype == at::kFloat8_e4m3fn && A_dtype == at::kBFloat16) { \ + using scalar_t = bfloat16_t; \ + MoEGEMMLauncherCallER( \ + 'R', 'C', policy, scalar_t, float_e4m3_t, uint8_t); \ + } else if (B_dtype == at::kFloat8_e4m3fn && A_dtype == at::kHalf) { \ + using scalar_t = half_t; \ + MoEGEMMLauncherCallER( \ + 'R', 'C', policy, scalar_t, float_e4m3_t, uint8_t); \ + } else if (B_dtype == at::kFloat8_e5m2 && A_dtype == at::kBFloat16) { \ + using scalar_t = bfloat16_t; \ + MoEGEMMLauncherCallER( \ + 'R', 'C', policy, scalar_t, float_e5m2_t, uint8_t); \ + } else if (B_dtype == at::kFloat8_e5m2 && A_dtype == at::kHalf) { \ + using scalar_t = half_t; \ + MoEGEMMLauncherCallER( \ + 'R', 'C', policy, scalar_t, float_e5m2_t, uint8_t); \ + } + + if (A_avg_M <= 32) { + using policy = w8a16_policy_m_16; + MXFP8LauncherCallER(policy); + } else if (A_avg_M <= 128) { + using policy = w8a16_policy_m_32; + MXFP8LauncherCallER(policy); + } else { + using policy = w8a16_policy; + MXFP8LauncherCallER(policy); + } +#undef MXFP8LauncherCallER } else if (is_weight_fp8) { TORCH_CHECK(ptr_scales.has_value(), "w8a16 grouped gemm must have scales"); TORCH_CHECK(ptr_scales->is_contiguous(), "ptr_scales must be contiguous"); diff --git a/megablocks/torch-ext/megablocks/xpu_fused_moe.py b/megablocks/torch-ext/megablocks/xpu_fused_moe.py index 0440b484e..cf527bf0e 100644 --- a/megablocks/torch-ext/megablocks/xpu_fused_moe.py +++ b/megablocks/torch-ext/megablocks/xpu_fused_moe.py @@ -26,7 +26,7 @@ def _register_if_available(op_name, fn): _register_if_available( "cutlass_grouped_gemm_interface", - lambda ptr_A, ptr_B, ptr_scales, ptr_bias, ptr_D, expert_first_token_offset, N, K, num_experts, is_B_int4, is_B_mxfp4: ptr_D, + lambda ptr_A, ptr_B, ptr_scales, ptr_bias, ptr_D, expert_first_token_offset, N, K, num_experts, is_B_int4, is_B_mxfp4, is_B_mxfp8: ptr_D, ) _register_if_available( @@ -106,12 +106,13 @@ def exclusive_prefix_sum(arr): K=k, num_experts=num_experts, is_B_int4=False, - is_B_mxfp4=False) + is_B_mxfp4=False, + is_B_mxfp8=False) def cutlass_grouped_gemm_xe2(input_A, input_B, scales, bias, output, num_rows_per_expert, n, k, num_experts, is_B_int4, - is_B_mxfp4): + is_B_mxfp4, is_B_mxfp8=False): expert_first_token_offset = torch.cat([ torch.tensor([0], dtype=num_rows_per_expert.dtype, @@ -129,7 +130,8 @@ def cutlass_grouped_gemm_xe2(input_A, input_B, scales, bias, output, K=k, num_experts=num_experts, is_B_int4=is_B_int4, - is_B_mxfp4=is_B_mxfp4) + is_B_mxfp4=is_B_mxfp4, + is_B_mxfp8=is_B_mxfp8) def ceilDiv(a, b): @@ -206,7 +208,8 @@ def xpu_fused_moe(hidden_states, ep_size=1, is_fp8=False, is_int4=False, - is_mxfp4=False): + is_mxfp4=False, + is_mxfp8=False): ''' hidden_states: [num_rows, hidden_size] w13: [num_experts, 2*inter_size, hidden_size] @@ -228,6 +231,7 @@ def xpu_fused_moe(hidden_states, num_experts: int is_int4: bool is_mxfp4: bool + is_mxfp8: bool ''' # Resolve DTensors to local tensors before passing to custom ops @@ -266,7 +270,7 @@ def xpu_fused_moe(hidden_states, # 4bits support [E, N, K] # other types [E, K, N] - if not is_int4 and not is_mxfp4: + if not is_int4 and not is_mxfp4 and not is_mxfp8: if not hasattr(w13, 'xpu_fused_moe'): if needs_transpose: w13.data = w13.transpose(-1, -2).contiguous() @@ -389,7 +393,7 @@ def config_ws(name, size): ########### gemm1 ################## input_B = w13 - if not is_fp8 and not is_int4 and not is_mxfp4: + if not is_fp8 and not is_int4 and not is_mxfp4 and not is_mxfp8: ops.cutlass_grouped_gemm_interface( ptr_A=gemm1_input, ptr_B=input_B, @@ -401,7 +405,8 @@ def config_ws(name, size): K=hidden_size, num_experts=num_experts_per_node, is_B_int4=is_int4, - is_B_mxfp4=is_mxfp4) + is_B_mxfp4=is_mxfp4, + is_B_mxfp8=is_mxfp8) else: ops.cutlass_grouped_gemm_interface( ptr_A=gemm1_input, @@ -414,7 +419,8 @@ def config_ws(name, size): K=hidden_size, num_experts=num_experts_per_node, is_B_int4=is_int4, - is_B_mxfp4=is_mxfp4) + is_B_mxfp4=is_mxfp4, + is_B_mxfp8=is_mxfp8) # act act_output = torch.empty((num_moe_inputs, inter_size), @@ -435,7 +441,7 @@ def config_ws(name, size): gemm2_output = torch.empty((num_moe_inputs, hidden_size), dtype=hidden_states.dtype, device=hidden_states.device) - if not is_fp8 and not is_int4 and not is_mxfp4: + if not is_fp8 and not is_int4 and not is_mxfp4 and not is_mxfp8: ops.cutlass_grouped_gemm_interface( ptr_A=input_A, ptr_B=input_B, @@ -447,7 +453,8 @@ def config_ws(name, size): K=inter_size, num_experts=num_experts_per_node, is_B_int4=is_int4, - is_B_mxfp4=is_mxfp4) + is_B_mxfp4=is_mxfp4, + is_B_mxfp8=is_mxfp8) else: ops.cutlass_grouped_gemm_interface( ptr_A=input_A, @@ -460,7 +467,8 @@ def config_ws(name, size): K=inter_size, num_experts=num_experts_per_node, is_B_int4=is_int4, - is_B_mxfp4=is_mxfp4) + is_B_mxfp4=is_mxfp4, + is_B_mxfp8=is_mxfp8) ops.moe_gather(output, gemm2_output, topk_weights, permuted_row_to_unpermuted_row, @@ -609,6 +617,7 @@ def forward(self, x: torch.Tensor) -> tuple: is_fp8 = getattr(self.experts, "is_fp8", False) is_int4 = getattr(self.experts, "is_int4", False) is_mxfp4 = getattr(self.experts, "is_mxfp4", False) + is_mxfp8 = getattr(self.experts, "is_mxfp8", False) w13_scales = getattr(self.experts, "gate_up_proj_scales", None) w2_scales = getattr(self.experts, "down_proj_scales", None) @@ -650,6 +659,7 @@ def forward(self, x: torch.Tensor) -> tuple: is_fp8=is_fp8, is_int4=is_int4, is_mxfp4=is_mxfp4, + is_mxfp8=is_mxfp8, ) # All-reduce across EP group to combine partial expert outputs diff --git a/megablocks/torch-ext/torch_binding.cpp b/megablocks/torch-ext/torch_binding.cpp index 8ea4db20c..fff510eab 100644 --- a/megablocks/torch-ext/torch_binding.cpp +++ b/megablocks/torch-ext/torch_binding.cpp @@ -263,7 +263,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { "cutlass_grouped_gemm_interface(Tensor ptr_A, Tensor ptr_B, " "Tensor? ptr_scales, Tensor? ptr_bias, Tensor! ptr_D, " "Tensor expert_first_token_offset, int N, int K, int num_experts, " - "bool is_B_int4, bool is_B_mxfp4) -> Tensor"); + "bool is_B_int4, bool is_B_mxfp4, bool is_B_mxfp8) -> Tensor"); ops.impl("cutlass_grouped_gemm_interface", torch::kXPU, &cutlass_grouped_gemm_interface); } From 13e92b5d4fc784d13481870d8a46fef7848c76e2 Mon Sep 17 00:00:00 2001 From: jiqing-feng Date: Tue, 28 Jul 2026 01:47:05 -0700 Subject: [PATCH 09/12] megablocks: build the XPU grouped GEMM twice and dispatch CRI at runtime The Xe35 paths in sycl-tla (2D transpose/prefetch copies, mxfp8 fcvt dequant) are gated behind __SYCL_TARGET_INTEL_GPU_CRI__, which the compiler never defines on its own. Passing it globally would leak into the pvc and bmg images, so add a CRI-only component that compiles the grouped GEMM a second time with the macro forced on and spir64_gen as its sole target; the link step AOT-lowers that image to cri. Both images cover the same dtypes, so the host picks between them by reading the device IP version (pvc 12.x, bmg 20.x, cri 35.x) through a new XPUFeatures helper, mirroring the existing CPUFeatures convention. The two compilations expand the shared headers into different code, so wrap them in an inline namespace keyed on the macro and add an arch tag to GemmCuteName. Without this the launchers mangle identically in both translation units and the linker silently keeps one body. Collapse the former pass-through wrapper into a single templated entry point that static_asserts its Arch against SYCL_INTEL_TARGET. Verified on the built extension: Li35E kernels appear only in the cri device images, pvc and bmg carry none, and the BMG test suite still passes 134/134. --- megablocks/build.toml | 27 +++++++++++++++ .../grouped_gemm/grouped_gemm_interface.cpp | 11 +++++-- .../csrc_xpu/grouped_gemm/xe_2/gemm_xe2.hpp | 3 ++ .../grouped_gemm/xe_2/gemm_xe2_policy.hpp | 3 ++ .../grouped_gemm/xe_2/grouped_gemm_xe2.cpp | 30 ++--------------- .../grouped_gemm/xe_2/grouped_gemm_xe2.h | 9 ++++- .../grouped_gemm/xe_2/grouped_gemm_xe2.hpp | 2 ++ .../xe_2/grouped_gemm_xe2_cri.cpp | 6 ++++ .../xe_2/grouped_gemm_xe2_interface.hpp | 29 +++++++++++----- .../grouped_gemm/xe_2/xe2_target_ns.hpp | 16 +++++++++ megablocks/csrc_xpu/xpu_features.hpp | 33 +++++++++++++++++++ 11 files changed, 130 insertions(+), 39 deletions(-) create mode 100644 megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2_cri.cpp create mode 100644 megablocks/csrc_xpu/grouped_gemm/xe_2/xe2_target_ns.hpp create mode 100644 megablocks/csrc_xpu/xpu_features.hpp diff --git a/megablocks/build.toml b/megablocks/build.toml index 024eeeba3..327142582 100644 --- a/megablocks/build.toml +++ b/megablocks/build.toml @@ -91,6 +91,7 @@ depends = [ src = [ "csrc_xpu/utils.h", "csrc_xpu/dispatch_utils.h", + "csrc_xpu/xpu_features.hpp", "csrc_xpu/core/registration.h", "csrc_xpu/activation.cpp", "csrc_xpu/activation.h", @@ -111,6 +112,32 @@ src = [ "csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2_interface.hpp", "csrc_xpu/grouped_gemm/xe_2/gemm_xe2.hpp", "csrc_xpu/grouped_gemm/xe_2/gemm_xe2_policy.hpp", + "csrc_xpu/grouped_gemm/xe_2/xe2_target_ns.hpp", +] + +# The Xe35 block-scaled DPAS path in sycl-tla is gated behind +# __SYCL_TARGET_INTEL_GPU_CRI__, which the compiler never sets on its own, so +# build a CRI-only TU with the macro forced on and spir64_gen as the sole +# target. The link step AOT-lowers that image to cri; pvc/bmg keep the Xe20 TU +# above and the host dispatches by device IP version at runtime. +[kernel.megablocks_xpu_cri] +backend = "xpu" +depends = [ + "torch", + "sycl_tla", +] +sycl-flags = [ + "-fsycl-targets=spir64_gen", + "-D__SYCL_TARGET_INTEL_GPU_CRI__=1", +] +src = [ + "csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2_cri.cpp", + "csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2.h", + "csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2.hpp", + "csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2_interface.hpp", + "csrc_xpu/grouped_gemm/xe_2/gemm_xe2.hpp", + "csrc_xpu/grouped_gemm/xe_2/gemm_xe2_policy.hpp", + "csrc_xpu/grouped_gemm/xe_2/xe2_target_ns.hpp", ] [kernel.megablocks_cpu] diff --git a/megablocks/csrc_xpu/grouped_gemm/grouped_gemm_interface.cpp b/megablocks/csrc_xpu/grouped_gemm/grouped_gemm_interface.cpp index b97adaa6a..1962a8609 100644 --- a/megablocks/csrc_xpu/grouped_gemm/grouped_gemm_interface.cpp +++ b/megablocks/csrc_xpu/grouped_gemm/grouped_gemm_interface.cpp @@ -1,7 +1,10 @@ #include "grouped_gemm_interface.h" #include +#include "../xpu_features.hpp" #include "xe_2/grouped_gemm_xe2.h" +using megablocks::xpu::XPUFeatures; + torch::Tensor cutlass_grouped_gemm_interface( torch::Tensor ptr_A, torch::Tensor ptr_B, @@ -15,8 +18,12 @@ torch::Tensor cutlass_grouped_gemm_interface( bool is_B_int4, bool is_B_mxfp4, bool is_B_mxfp8) { - // Use XE2 cutlass kernel - return cutlass_grouped_gemm_xe2( + // The grouped GEMMs are built twice: an Xe20 image for pvc/bmg and an Xe35 + // image for CRI. Both cover the same dtypes, so pick the variant by device. + auto gemm = XPUFeatures::isXe35(ptr_A.device().index()) + ? MoE::cutlass_grouped_gemm_xe2<35> + : MoE::cutlass_grouped_gemm_xe2<20>; + return gemm( ptr_A, ptr_B, ptr_scales, diff --git a/megablocks/csrc_xpu/grouped_gemm/xe_2/gemm_xe2.hpp b/megablocks/csrc_xpu/grouped_gemm/xe_2/gemm_xe2.hpp index dbe1e09e1..7ab38637d 100644 --- a/megablocks/csrc_xpu/grouped_gemm/xe_2/gemm_xe2.hpp +++ b/megablocks/csrc_xpu/grouped_gemm/xe_2/gemm_xe2.hpp @@ -44,11 +44,13 @@ #include "cutlass/util/reference/device/gemm_complex.h" #include "cutlass/util/reference/device/tensor_compare.h" #include "cutlass/util/reference/host/tensor_fill.h" +#include "xe2_target_ns.hpp" #pragma clang diagnostic ignored "-Wpass-failed" #pragma clang diagnostic ignored "-Wdeprecated-declarations" namespace MoE { +inline namespace MEGABLOCKS_XE_TARGET_NS { using namespace cute; @@ -422,4 +424,5 @@ CUTE_DEVICE void xe_gemm_4bits( copy(copy_c, tCrC_final_sg_tensor, tCgC); } +} // inline namespace MEGABLOCKS_XE_TARGET_NS } // namespace MoE diff --git a/megablocks/csrc_xpu/grouped_gemm/xe_2/gemm_xe2_policy.hpp b/megablocks/csrc_xpu/grouped_gemm/xe_2/gemm_xe2_policy.hpp index 57f320725..aa53aeeba 100644 --- a/megablocks/csrc_xpu/grouped_gemm/xe_2/gemm_xe2_policy.hpp +++ b/megablocks/csrc_xpu/grouped_gemm/xe_2/gemm_xe2_policy.hpp @@ -2,8 +2,10 @@ #include "cute/atom/mma_atom.hpp" #include "cutlass/numeric_types.h" +#include "xe2_target_ns.hpp" namespace MoE { +inline namespace MEGABLOCKS_XE_TARGET_NS { using namespace cute; class xe_gemm_policy_base { @@ -85,4 +87,5 @@ class w4a16_policy_m_32 : public xe_gemm_policy_base { using SGLayout = Layout, Stride<_4, _1, _0>>; }; +} // inline namespace MEGABLOCKS_XE_TARGET_NS } // namespace MoE \ No newline at end of file diff --git a/megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2.cpp b/megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2.cpp index d2a2f9714..298d2fffa 100644 --- a/megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2.cpp +++ b/megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2.cpp @@ -1,31 +1,5 @@ -#include #include "grouped_gemm_xe2.h" #include "grouped_gemm_xe2_interface.hpp" -torch::Tensor cutlass_grouped_gemm_xe2( - torch::Tensor ptr_A, - torch::Tensor ptr_B, - const c10::optional& ptr_scales, - const c10::optional& ptr_bias, - torch::Tensor ptr_D, - torch::Tensor expert_first_token_offset, - int64_t N, - int64_t K, - int64_t num_experts, - bool is_B_int4, - bool is_B_mxfp4, - bool is_B_mxfp8) { - return MoE::cutlass_grouped_gemm_xe2_impl( - ptr_A, - ptr_B, - ptr_scales, - ptr_bias, - ptr_D, - expert_first_token_offset, - N, - K, - num_experts, - is_B_int4, - is_B_mxfp4, - is_B_mxfp8); -} \ No newline at end of file +// Xe20 image, used by pvc and bmg. +template decltype(MoE::cutlass_grouped_gemm_xe2<20>) MoE::cutlass_grouped_gemm_xe2<20>; diff --git a/megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2.h b/megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2.h index ea2d9a362..aba36a0db 100644 --- a/megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2.h +++ b/megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2.h @@ -1,5 +1,10 @@ #include +namespace MoE { + +// Instantiated once per target architecture: 20 for pvc/bmg, 35 for CRI. The +// two instantiations live in separate translation units with different flags. +template torch::Tensor cutlass_grouped_gemm_xe2( torch::Tensor ptr_A, torch::Tensor ptr_B, @@ -12,4 +17,6 @@ torch::Tensor cutlass_grouped_gemm_xe2( int64_t num_experts, bool is_B_int4, bool is_B_mxfp4, - bool is_B_mxfp8); \ No newline at end of file + bool is_B_mxfp8); + +} // namespace MoE diff --git a/megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2.hpp b/megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2.hpp index 0e8f6d458..21c374894 100644 --- a/megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2.hpp +++ b/megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2.hpp @@ -45,6 +45,7 @@ #pragma clang diagnostic ignored "-Wdeprecated-declarations" namespace MoE { +inline namespace MEGABLOCKS_XE_TARGET_NS { using namespace cute; template @@ -231,4 +232,5 @@ CUTE_DEVICE void MoEGEMM( } } +} // inline namespace MEGABLOCKS_XE_TARGET_NS } // namespace MoE diff --git a/megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2_cri.cpp b/megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2_cri.cpp new file mode 100644 index 000000000..fd9da9b3e --- /dev/null +++ b/megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2_cri.cpp @@ -0,0 +1,6 @@ +#include "grouped_gemm_xe2.h" +#include "grouped_gemm_xe2_interface.hpp" + +// Xe35 image. Built with -D__SYCL_TARGET_INTEL_GPU_CRI__=1 and +// -fsycl-targets=spir64_gen so sycl-tla emits the CRI block-scaled DPAS path. +template decltype(MoE::cutlass_grouped_gemm_xe2<35>) MoE::cutlass_grouped_gemm_xe2<35>; diff --git a/megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2_interface.hpp b/megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2_interface.hpp index 5ab94efdd..95307a20a 100644 --- a/megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2_interface.hpp +++ b/megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2_interface.hpp @@ -58,15 +58,18 @@ #include "gemm_xe2_policy.hpp" #include "grouped_gemm_xe2.hpp" +#include "xe2_target_ns.hpp" #pragma clang diagnostic ignored "-Wpass-failed" #pragma clang diagnostic ignored "-Wdeprecated-declarations" namespace MoE { +inline namespace MEGABLOCKS_XE_TARGET_NS { using namespace cute; -// type tag to define a unique sycl kernel name -template +// type tag to define a unique sycl kernel name. The trailing arch tag keeps the +// CRI (Xe35) and generic (Xe20) device images from sharing a kernel name. +template class GemmCuteName; template < @@ -135,7 +138,8 @@ void MoEGEMMLauncher( ElementD, layoutA, layoutB, - policy>>( + policy, + SYCL_INTEL_TARGET>>( sycl::nd_range<3>{global * local, local}, kernel_props, [=](auto) { MoE::MoEGEMM< GmemTiledCopyA, @@ -162,19 +166,28 @@ void MoEGEMMLauncher( EventManager::getInstance().addEvent(event); } -at::Tensor cutlass_grouped_gemm_xe2_impl( - at::Tensor& ptr_A, - at::Tensor& ptr_B, +} // inline namespace MEGABLOCKS_XE_TARGET_NS + +using namespace cute; + +// Entry point for one architecture, instantiated as <20> by the pvc/bmg +// translation unit and as <35> by the CRI one. +template +torch::Tensor cutlass_grouped_gemm_xe2( + torch::Tensor ptr_A, + torch::Tensor ptr_B, const c10::optional& ptr_scales, const c10::optional& ptr_bias, - at::Tensor& ptr_D, - at::Tensor& expert_first_token_offset, + torch::Tensor ptr_D, + torch::Tensor expert_first_token_offset, int64_t N, int64_t K, int64_t num_experts, bool is_B_int4, bool is_B_mxfp4, bool is_B_mxfp8) { + // Guards against the CRI translation unit losing -D__SYCL_TARGET_INTEL_GPU_CRI__. + static_assert(Arch == SYCL_INTEL_TARGET, "Arch must match SYCL_INTEL_TARGET"); auto& dpcpp_queue = at::xpu::getCurrentXPUStream(ptr_A.device().index()).queue(); auto A_dtype = ptr_A.dtype(); diff --git a/megablocks/csrc_xpu/grouped_gemm/xe_2/xe2_target_ns.hpp b/megablocks/csrc_xpu/grouped_gemm/xe_2/xe2_target_ns.hpp new file mode 100644 index 000000000..23ed0af3c --- /dev/null +++ b/megablocks/csrc_xpu/grouped_gemm/xe_2/xe2_target_ns.hpp @@ -0,0 +1,16 @@ +/***************************************************************************************** + * Copyright (C) 2026 Intel Corporation, All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + ****************************************************************************************/ + +#pragma once + +// These headers are compiled twice, once per target architecture, and sycl-tla +// expands them into different code each time. Keeping the two expansions in +// separate namespaces stops the linker from merging symbols that share a name +// but not a body. +#if defined(__SYCL_TARGET_INTEL_GPU_CRI__) +#define MEGABLOCKS_XE_TARGET_NS xe35 +#else +#define MEGABLOCKS_XE_TARGET_NS xe20 +#endif diff --git a/megablocks/csrc_xpu/xpu_features.hpp b/megablocks/csrc_xpu/xpu_features.hpp new file mode 100644 index 000000000..ad7a8bfc4 --- /dev/null +++ b/megablocks/csrc_xpu/xpu_features.hpp @@ -0,0 +1,33 @@ +/***************************************************************************************** + * Copyright (C) 2026 Intel Corporation, All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + ****************************************************************************************/ + +// XPU Device Feature Detection for MegaBlocks + +#pragma once + +#include +#include + +namespace megablocks { +namespace xpu { + +// XPU device feature detection, keyed on the device IP version reported by the +// driver (pvc 12.x, bmg 20.x, CRI 35.x). +class XPUFeatures { + public: + // True when the device is Xe35 (CRI) or newer, which selects the kernels + // built from the CRI-only translation unit. + static bool isXe35(c10::DeviceIndex device) { + return ipVersionMajor(device) >= 35; + } + + // Major component of the device IP version, e.g. 20 for "20.1.0". + static int ipVersionMajor(c10::DeviceIndex device) { + return std::atoi(at::xpu::getDeviceProperties(device)->version.c_str()); + } +}; + +} // namespace xpu +} // namespace megablocks From 4d77cb31e9a416bcb393dc89974c954f88f5fc34 Mon Sep 17 00:00:00 2001 From: jiqing-feng Date: Tue, 28 Jul 2026 22:20:50 -0700 Subject: [PATCH 10/12] flash-attn2: fix merge damage in the XPU/CPU torch bindings The merge of main left "Generator? gen_) -> Tensor[]" in front of the MXFP scale arguments in the fwd and varlen_fwd schemas. Adjacent string literals concatenate, so both schemas ended up with a return arrow in the middle and would have thrown at library load time. The same conflict also reinstated the kCUDA registrations, which main moved to the stable-ABI bindings in torch_binding_stable.cpp. They sit inside #if !defined(CUDA_KERNEL) here, so they were dead code; drop them and let the block start at XPU_KERNEL. --- flash-attn2/torch-ext/torch_binding.cpp | 12 ++++-------- 1 file changed, 4 insertions(+), 8 deletions(-) diff --git a/flash-attn2/torch-ext/torch_binding.cpp b/flash-attn2/torch-ext/torch_binding.cpp index 79ce566e4..b43290ca5 100644 --- a/flash-attn2/torch-ext/torch_binding.cpp +++ b/flash-attn2/torch-ext/torch_binding.cpp @@ -36,13 +36,11 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { "int window_size_right, " "float softcap, " "bool return_softmax, " - "Generator? gen_) -> Tensor[]" + "Generator? gen_, " "Tensor? scale_q_=None, " "Tensor? scale_k_=None, " "Tensor? scale_v_=None) -> Tensor[]"); -#if defined(CUDA_KERNEL) - ops.impl("fwd", torch::kCUDA, &mha_fwd); -#elif defined(XPU_KERNEL) +#if defined(XPU_KERNEL) ops.impl("fwd", torch::kXPU, &mha_fwd); #endif @@ -67,15 +65,13 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { "int window_size_right, " "float softcap, " "bool return_softmax, " - "Generator? gen_) -> Tensor[]" + "Generator? gen_, " "Tensor? scale_q_=None, " "Tensor? scale_k_=None, " "Tensor? scale_v_=None, " "Tensor? cu_scale_q_=None, " "Tensor? cu_scale_kv_=None) -> Tensor[]"); -#if defined(CUDA_KERNEL) - ops.impl("varlen_fwd", torch::kCUDA, &mha_varlen_fwd); -#elif defined(XPU_KERNEL) +#if defined(XPU_KERNEL) ops.impl("varlen_fwd", torch::kXPU, &mha_varlen_fwd); #elif defined(CPU_KERNEL) ops.impl("varlen_fwd", torch::kCPU, &mha_varlen_fwd); From a227b31d63bb4969f3d09bbc528e79f311cd6bc5 Mon Sep 17 00:00:00 2001 From: jiqing-feng Date: Wed, 29 Jul 2026 19:01:58 -0700 Subject: [PATCH 11/12] update Signed-off-by: jiqing-feng --- flash-attn2/flash_attn_xpu/flash_api.cpp | 13 ++++++ .../flash_attn2/flash_attn_interface.py | 41 ++++++++++++++++++- 2 files changed, 52 insertions(+), 2 deletions(-) diff --git a/flash-attn2/flash_attn_xpu/flash_api.cpp b/flash-attn2/flash_attn_xpu/flash_api.cpp index bb8dc351d..06c28f59a 100644 --- a/flash-attn2/flash_attn_xpu/flash_api.cpp +++ b/flash-attn2/flash_attn_xpu/flash_api.cpp @@ -120,6 +120,10 @@ mha_fwd( at::Tensor out_padded; if (out_.has_value()) { auto out_val = out_.value(); + // The kernel writes out_dtype elements; a mismatched buffer would be + // overrun (MXFP writes float32 into what may be a half-width tensor). + TORCH_CHECK(out_val.scalar_type() == out_dtype, + "out must have dtype ", out_dtype, ", got ", out_val.scalar_type()); if (needs_padding) { const int pad_size = head_size_padded - head_size_actual; out_padded = torch::nn::functional::pad(out_val, torch::nn::functional::PadFuncOptions({0, pad_size})); @@ -496,6 +500,10 @@ mha_varlen_fwd( auto out_dtype = is_mxfp ? torch::kFloat32 : c10::typeMetaToScalarType(q_dtype); at::Tensor out_padded; if (out_.has_value()) { + // The kernel writes out_dtype elements; a mismatched buffer would be + // overrun (MXFP writes float32 into what may be a half-width tensor). + TORCH_CHECK(out_.value().scalar_type() == out_dtype, + "out must have dtype ", out_dtype, ", got ", out_.value().scalar_type()); out_padded = maybe_pad(out_.value()); } else if (is_mxfp) { out_padded = torch::zeros({total_q, num_heads, head_size_padded}, @@ -506,6 +514,11 @@ mha_varlen_fwd( const bool is_local = (window_size_left != -1) || (window_size_right != -1); const bool is_paged = block_table_.has_value() && block_table_->defined(); + // policy_dispatch() only routes MXFP on the non-paged / non-rotary paths; any + // other combination silently falls through to the bf16 dispatcher and would + // reinterpret the fp8/fp4 bytes as bf16. Reject it here instead. + TORCH_CHECK(!(is_mxfp && is_paged), + "MXFP (fp8/fp4) forward does not support paged KV (block_table)."); q_padded = ensure_contiguous(q_padded); k_padded = ensure_contiguous(k_padded); diff --git a/flash-attn2/torch-ext/flash_attn2/flash_attn_interface.py b/flash-attn2/torch-ext/flash_attn2/flash_attn_interface.py index 0868572aa..d2abcd941 100644 --- a/flash-attn2/torch-ext/flash_attn2/flash_attn_interface.py +++ b/flash-attn2/torch-ext/flash_attn2/flash_attn_interface.py @@ -77,6 +77,31 @@ def round_multiple(x, m): return (x + m - 1) // m * m +# Block-scaled MXFP (FP8 e5m2/e4m3, FP4 e2m1) is selected by the Q dtype, matching +# the `is_mxfp` check in flash_api.cpp. +_FP4_DTYPE = getattr(torch, "float4_e2m1fn_x2", None) +_MXFP_DTYPES = tuple( + d + for d in ( + getattr(torch, "float8_e5m2", None), + getattr(torch, "float8_e4m3fn", None), + _FP4_DTYPE, + ) + if d is not None +) + + +def is_mxfp_dtype(dtype): + return dtype in _MXFP_DTYPES + + +def _mxfp_out_shape(q): + """Shape of the MXFP output: fp4_e2m1fn_x2 packs two values per byte, so the + kernel writes twice the head dimension that q carries.""" + head_dim = q.shape[-1] * 2 if q.dtype is _FP4_DTYPE else q.shape[-1] + return q.shape[:-1] + (head_dim,) + + @torch.library.custom_op( add_op_namespace_prefix("_flash_attn_forward"), mutates_args=(), @@ -131,7 +156,13 @@ def _flash_attn_forward_fake( q, k, v = [maybe_contiguous(x) for x in (q, k, v)] batch_size, seqlen_q, num_heads, head_size = q.shape seqlen_k = k.shape[1] - out = torch.empty_like(q) + # MXFP accumulates in and writes float32, so it cannot reuse q's metadata. + if is_mxfp_dtype(q.dtype): + out = torch.empty( + _mxfp_out_shape(q), dtype=torch.float32, device=q.device, layout=q.layout + ) + else: + out = torch.empty_like(q) softmax_lse = torch.empty( (batch_size, num_heads, seqlen_q), dtype=torch.float32, @@ -238,7 +269,13 @@ def _flash_attn_varlen_forward_fake( batch_size = cu_seqlens_q.numel() - 1 total_q, num_heads, _ = q.shape - out = torch.empty_like(q) + # MXFP accumulates in and writes float32, so it cannot reuse q's metadata. + if is_mxfp_dtype(q.dtype): + out = torch.empty( + _mxfp_out_shape(q), dtype=torch.float32, device=q.device, layout=q.layout + ) + else: + out = torch.empty_like(q) softmax_lse = torch.empty( (num_heads, total_q), dtype=torch.float32, device=q.device, layout=q.layout ) From b1274f6168a8ef7fc2ccf03e9214795cac7386f6 Mon Sep 17 00:00:00 2001 From: jiqing-feng Date: Sun, 2 Aug 2026 19:07:28 -0700 Subject: [PATCH 12/12] flash-attn2: revert the XPU MXFP forward support The MXFP (block-scaled FP8/FP4) forward path is unreachable in practice: - transformers has no entry point for it. After reverting the Python-side plumbing, flash_attn_interface.py is byte-identical to main, so nothing can pass the scale tensors in. - The CUDA backend has no MXFP support at all -- its stable-ABI schema stops at `Generator? gen_`. Keeping the XPU-only variant means the two backends disagree on the op schema, which is exactly what the PR review flagged. So drop the whole feature rather than ship a path nobody can call: the `flash_attn_xpu_mxfp` kernel component, its 16 translation units and all host/device MXFP plumbing. What stays is the `cute::transform` overload fix in fmha_fwd_common.hpp. That one is unrelated to MXFP -- it is required by the newer sycl-tla that megablocks pulls in for its CRI grouped GEMM, and sycl-tla is resolved once for the whole build rather than per kernel. It applies to pvc/bmg just as much as to cri. megablocks is untouched: its CRI support is independent of this revert. flash-attn2 is now 5 lines away from main. --- flash-attn2/build.toml | 32 --- flash-attn2/flash_attn_xpu/flash_api.cpp | 194 +++------------- .../src/collective/fmha_fusion.hpp | 1 - .../src/collective/fmha_fwd_mainloop_xe2.hpp | 216 +----------------- .../src/flash_fwd_hdim128_fix_mxfp.cpp | 11 - .../src/flash_fwd_hdim128_varlen_mxfp.cpp | 11 - .../src/flash_fwd_hdim160_fix_mxfp.cpp | 11 - .../src/flash_fwd_hdim160_varlen_mxfp.cpp | 11 - .../src/flash_fwd_hdim192_fix_mxfp.cpp | 11 - .../src/flash_fwd_hdim192_varlen_mxfp.cpp | 11 - .../src/flash_fwd_hdim256_fix_mxfp.cpp | 11 - .../src/flash_fwd_hdim256_varlen_mxfp.cpp | 11 - .../src/flash_fwd_hdim32_fix_mxfp.cpp | 11 - .../src/flash_fwd_hdim32_varlen_mxfp.cpp | 11 - .../src/flash_fwd_hdim512_fix_mxfp.cpp | 11 - .../src/flash_fwd_hdim512_varlen_mxfp.cpp | 11 - .../src/flash_fwd_hdim64_fix_mxfp.cpp | 11 - .../src/flash_fwd_hdim64_varlen_mxfp.cpp | 11 - .../src/flash_fwd_hdim96_fix_mxfp.cpp | 11 - .../src/flash_fwd_hdim96_varlen_mxfp.cpp | 11 - flash-attn2/flash_attn_xpu/src/fmha_fwd.cpp | 24 +- flash-attn2/flash_attn_xpu/src/fmha_fwd.hpp | 100 ++------ .../flash_attn_xpu/src/fmha_fwd_impl.hpp | 208 +---------------- .../flash_attn_xpu/src/fmha_fwd_types.hpp | 18 +- flash-attn2/flash_attn_xpu/src/fmha_utils.hpp | 19 +- .../src/kernel/fmha_fwd_kernel_xe2.hpp | 133 ++--------- .../flash_attn2/flash_attn_interface.py | 41 +--- flash-attn2/torch-ext/torch_binding.cpp | 12 +- flash-attn2/torch-ext/torch_binding.h | 12 +- 29 files changed, 96 insertions(+), 1090 deletions(-) delete mode 100644 flash-attn2/flash_attn_xpu/src/flash_fwd_hdim128_fix_mxfp.cpp delete mode 100644 flash-attn2/flash_attn_xpu/src/flash_fwd_hdim128_varlen_mxfp.cpp delete mode 100644 flash-attn2/flash_attn_xpu/src/flash_fwd_hdim160_fix_mxfp.cpp delete mode 100644 flash-attn2/flash_attn_xpu/src/flash_fwd_hdim160_varlen_mxfp.cpp delete mode 100644 flash-attn2/flash_attn_xpu/src/flash_fwd_hdim192_fix_mxfp.cpp delete mode 100644 flash-attn2/flash_attn_xpu/src/flash_fwd_hdim192_varlen_mxfp.cpp delete mode 100644 flash-attn2/flash_attn_xpu/src/flash_fwd_hdim256_fix_mxfp.cpp delete mode 100644 flash-attn2/flash_attn_xpu/src/flash_fwd_hdim256_varlen_mxfp.cpp delete mode 100644 flash-attn2/flash_attn_xpu/src/flash_fwd_hdim32_fix_mxfp.cpp delete mode 100644 flash-attn2/flash_attn_xpu/src/flash_fwd_hdim32_varlen_mxfp.cpp delete mode 100644 flash-attn2/flash_attn_xpu/src/flash_fwd_hdim512_fix_mxfp.cpp delete mode 100644 flash-attn2/flash_attn_xpu/src/flash_fwd_hdim512_varlen_mxfp.cpp delete mode 100644 flash-attn2/flash_attn_xpu/src/flash_fwd_hdim64_fix_mxfp.cpp delete mode 100644 flash-attn2/flash_attn_xpu/src/flash_fwd_hdim64_varlen_mxfp.cpp delete mode 100644 flash-attn2/flash_attn_xpu/src/flash_fwd_hdim96_fix_mxfp.cpp delete mode 100644 flash-attn2/flash_attn_xpu/src/flash_fwd_hdim96_varlen_mxfp.cpp diff --git a/flash-attn2/build.toml b/flash-attn2/build.toml index 237fac9cf..a10e2b3a6 100644 --- a/flash-attn2/build.toml +++ b/flash-attn2/build.toml @@ -133,38 +133,6 @@ src = [ "flash_attn_xpu/src/kernel/fmha_bwd_kernel.hpp", ] -# Block-scaled MXFP kernels are gated behind __SYCL_TARGET_INTEL_GPU_CRI__ and -# only build for cri, so compile them spir64_gen-only with that macro. The -# spir64_gen image is AOT-lowered to cri by the tri-target link. -[kernel.flash_attn_xpu_mxfp] -backend = "xpu" -depends = [ - "torch", - "sycl_tla", -] -sycl-flags = [ - "-fsycl-targets=spir64_gen", - "-D__SYCL_TARGET_INTEL_GPU_CRI__=1", -] -src = [ - "flash_attn_xpu/src/flash_fwd_hdim32_varlen_mxfp.cpp", - "flash_attn_xpu/src/flash_fwd_hdim64_varlen_mxfp.cpp", - "flash_attn_xpu/src/flash_fwd_hdim96_varlen_mxfp.cpp", - "flash_attn_xpu/src/flash_fwd_hdim128_varlen_mxfp.cpp", - "flash_attn_xpu/src/flash_fwd_hdim160_varlen_mxfp.cpp", - "flash_attn_xpu/src/flash_fwd_hdim192_varlen_mxfp.cpp", - "flash_attn_xpu/src/flash_fwd_hdim256_varlen_mxfp.cpp", - "flash_attn_xpu/src/flash_fwd_hdim512_varlen_mxfp.cpp", - "flash_attn_xpu/src/flash_fwd_hdim32_fix_mxfp.cpp", - "flash_attn_xpu/src/flash_fwd_hdim64_fix_mxfp.cpp", - "flash_attn_xpu/src/flash_fwd_hdim96_fix_mxfp.cpp", - "flash_attn_xpu/src/flash_fwd_hdim128_fix_mxfp.cpp", - "flash_attn_xpu/src/flash_fwd_hdim160_fix_mxfp.cpp", - "flash_attn_xpu/src/flash_fwd_hdim192_fix_mxfp.cpp", - "flash_attn_xpu/src/flash_fwd_hdim256_fix_mxfp.cpp", - "flash_attn_xpu/src/flash_fwd_hdim512_fix_mxfp.cpp", -] - [kernel.flash_attn_cpu_avx512] backend = "cpu" cxx-flags = [ diff --git a/flash-attn2/flash_attn_xpu/flash_api.cpp b/flash-attn2/flash_attn_xpu/flash_api.cpp index 06c28f59a..8701dcc54 100644 --- a/flash-attn2/flash_attn_xpu/flash_api.cpp +++ b/flash-attn2/flash_attn_xpu/flash_api.cpp @@ -47,10 +47,7 @@ mha_fwd( int window_size_right, const float softcap, const bool return_softmax, - std::optional gen_, - std::optional scale_q_ = std::nullopt, - std::optional scale_k_ = std::nullopt, - std::optional scale_v_ = std::nullopt) { + std::optional gen_) { auto device_idx = q.device().index(); compat::select_device(device_idx); @@ -68,39 +65,22 @@ mha_fwd( const int num_heads_k = k.size(2); auto q_dtype = q.dtype(); - const bool is_mxfp = (q_dtype == torch::kFloat8_e5m2 || q_dtype == torch::kFloat8_e4m3fn || - q_dtype == torch::kFloat4_e2m1fn_x2); - TORCH_CHECK(q_dtype == torch::kFloat16 || q_dtype == torch::kBFloat16 || is_mxfp, - "FlashAttention only support fp16, bf16, fp8_e5m2, fp8_e4m3fn and fp4_e2m1fn_x2 data type"); + TORCH_CHECK(q_dtype == torch::kFloat16 || q_dtype == torch::kBFloat16, + "FlashAttention only support fp16 and bf16 data type"); TORCH_CHECK(k.dtype() == q_dtype, "query and key must have the same dtype"); + TORCH_CHECK(v.dtype() == q_dtype, "query and value must have the same dtype"); TORCH_CHECK(k.size(-1) == head_size_og, "Key head dimension must match Query head dimension"); - if (!is_mxfp) { - TORCH_CHECK(v.dtype() == q_dtype, "query and value must have the same dtype"); - TORCH_CHECK(v.size(-1) == head_size_og, "Value head dimension must match Query head dimension"); - } else if (q_dtype == torch::kFloat4_e2m1fn_x2) { - // MXFP4 (e2m1) consumes a BF16 value tensor; reject mismatches early - // instead of silently producing wrong results inside the kernel. - TORCH_CHECK(v.dtype() == torch::kBFloat16, - "For fp4_e2m1fn_x2 query, value must be bf16"); - } else { - // MXFP8 (e5m2 / e4m3fn) requires the value tensor to match the query fp8 dtype. - TORCH_CHECK(v.dtype() == q_dtype, - "For fp8 query, value must have the same fp8 dtype as query"); - } - // For FP4 (e2m1) the head dimension is packed (2 values per byte); unpack for padding math. - const int head_size_actual = (q_dtype == torch::kFloat4_e2m1fn_x2) ? head_size_og * 2 : head_size_og; - TORCH_CHECK(head_size_actual <= 256 || head_size_actual == 512, - "FlashAttention XPU only supports head dimension up to 256 or exactly 512. Got: " + std::to_string(head_size_actual)); + TORCH_CHECK(v.size(-1) == head_size_og, "Value head dimension must match Query head dimension"); + TORCH_CHECK(head_size_og <= 256 || head_size_og == 512, + "FlashAttention XPU only supports head dimension up to 256 or exactly 512. Got: " + std::to_string(head_size_og)); CHECK_DEVICE(q); CHECK_DEVICE(k); CHECK_DEVICE(v); - // XPU requires head_size to be a multiple of 32. - // FP4 packed inputs are consumed in their packed layout and are not padded. - const int head_size_padded = round_multiple(head_size_actual, 32); - const bool needs_padding = (head_size_actual != head_size_padded) - && (q_dtype != torch::kFloat4_e2m1fn_x2); + // XPU requires head_size to be a multiple of 32 + const int head_size_padded = round_multiple(head_size_og, 32); + const bool needs_padding = (head_size_og != head_size_padded); at::Tensor q_padded = q; at::Tensor k_padded = k; @@ -108,31 +88,21 @@ mha_fwd( // Apply padding if needed if (needs_padding) { - const int pad_size = head_size_padded - head_size_actual; + const int pad_size = head_size_padded - head_size_og; q_padded = torch::nn::functional::pad(q, torch::nn::functional::PadFuncOptions({0, pad_size})); k_padded = torch::nn::functional::pad(k, torch::nn::functional::PadFuncOptions({0, pad_size})); v_padded = torch::nn::functional::pad(v, torch::nn::functional::PadFuncOptions({0, pad_size})); } - // Low-precision (FP8/FP4) inputs accumulate in float32 and produce a float32 output. - auto out_dtype = is_mxfp ? torch::kFloat32 : c10::typeMetaToScalarType(q_dtype); - at::Tensor out_padded; if (out_.has_value()) { auto out_val = out_.value(); - // The kernel writes out_dtype elements; a mismatched buffer would be - // overrun (MXFP writes float32 into what may be a half-width tensor). - TORCH_CHECK(out_val.scalar_type() == out_dtype, - "out must have dtype ", out_dtype, ", got ", out_val.scalar_type()); - if (needs_padding) { - const int pad_size = head_size_padded - head_size_actual; + if (head_size_og != head_size_padded) { + const int pad_size = head_size_padded - head_size_og; out_padded = torch::nn::functional::pad(out_val, torch::nn::functional::PadFuncOptions({0, pad_size})); } else { out_padded = out_val; } - } else if (is_mxfp) { - out_padded = torch::zeros({batch_size, seqlen_q, num_heads, head_size_padded}, - q_padded.options().dtype(out_dtype)); } else { out_padded = torch::zeros_like(q_padded); } @@ -174,7 +144,7 @@ mha_fwd( // seqlen_q_rounded must be >= the kernel's Q tile size, otherwise the S_dmask // write in the mainloop will go out of bounds for rows in the padding region // of the tile. The Q tile size depends on head_size (see fmha_utils.hpp): - const int q_tile_size = (head_size_actual <= 32) ? 64 : (head_size_actual <= 128) ? 128 : 256; + const int q_tile_size = (head_size_og <= 32) ? 64 : (head_size_og <= 128) ? 128 : 256; const int seqlen_q_rounded = round_multiple(seqlen_q, q_tile_size); const int seqlen_k_rounded = round_multiple(seqlen_k, 128); at::Tensor S_dmask; @@ -185,18 +155,6 @@ mha_fwd( S_dmask = torch::empty({0}, q.options()); } - // MXFP (block-scaled) paths require per-tensor scale factors; contiguity is enforced. - std::optional scale_q = scale_q_.has_value() - ? std::optional(ensure_contiguous(scale_q_.value())) : std::nullopt; - std::optional scale_k = scale_k_.has_value() - ? std::optional(ensure_contiguous(scale_k_.value())) : std::nullopt; - std::optional scale_v = scale_v_.has_value() - ? std::optional(ensure_contiguous(scale_v_.value())) : std::nullopt; - if (is_mxfp) { - TORCH_CHECK(scale_q.has_value() && scale_k.has_value() && scale_v.has_value(), - "MXFP (fp8/fp4) forward currently requires scale_q, scale_k and scale_v tensors."); - } - cutlass_fmha_fwd_fix_impl( queue, q_padded, k_padded, v_padded, out_padded, @@ -206,13 +164,12 @@ mha_fwd( is_causal, is_local, p_dropout, philox_seed, philox_offset, nullptr, return_softmax ? S_dmask.data_ptr() : nullptr, - seqlen_q_rounded, seqlen_k_rounded, - scale_q, scale_k, scale_v); + seqlen_q_rounded, seqlen_k_rounded); - // Strip padding from output back to original (unpacked) head_size + // Strip padding from output back to original head_size at::Tensor out = needs_padding ? out_padded.index({torch::indexing::Slice(), torch::indexing::Slice(), - torch::indexing::Slice(), torch::indexing::Slice(0, head_size_actual)}) + torch::indexing::Slice(), torch::indexing::Slice(0, head_size_og)}) .contiguous() : ensure_contiguous(out_padded); @@ -429,12 +386,7 @@ mha_varlen_fwd( int window_size_right, const float softcap, const bool return_softmax, - std::optional gen_, - std::optional scale_q_ = std::nullopt, - std::optional scale_k_ = std::nullopt, - std::optional scale_v_ = std::nullopt, - std::optional cu_scale_q_ = std::nullopt, - std::optional cu_scale_kv_ = std::nullopt) { + std::optional gen_) { auto device_idx = q.device().index(); compat::select_device(device_idx); @@ -452,40 +404,23 @@ mha_varlen_fwd( const int batch_size = cu_seqlens_q.numel() - 1; auto q_dtype = q.dtype(); - const bool is_mxfp = (q_dtype == torch::kFloat8_e5m2 || q_dtype == torch::kFloat8_e4m3fn || - q_dtype == torch::kFloat4_e2m1fn_x2); - TORCH_CHECK(q_dtype == torch::kFloat16 || q_dtype == torch::kBFloat16 || is_mxfp, - "FlashAttention only support fp16, bf16, fp8_e5m2, fp8_e4m3fn and fp4_e2m1fn_x2 data type"); + TORCH_CHECK(q_dtype == torch::kFloat16 || q_dtype == torch::kBFloat16, + "FlashAttention only support fp16 and bf16 data type"); TORCH_CHECK(k.dtype() == q_dtype, "query and key must have the same dtype"); + TORCH_CHECK(v.dtype() == q_dtype, "query and value must have the same dtype"); TORCH_CHECK(k.size(-1) == head_size_og, "Key head dimension must match Query head dimension"); - if (!is_mxfp) { - TORCH_CHECK(v.dtype() == q_dtype, "query and value must have the same dtype"); - TORCH_CHECK(v.size(-1) == head_size_og, "Value head dimension must match Query head dimension"); - } else if (q_dtype == torch::kFloat4_e2m1fn_x2) { - // MXFP4 (e2m1) consumes a BF16 value tensor; reject mismatches early - // instead of silently producing wrong results inside the kernel. - TORCH_CHECK(v.dtype() == torch::kBFloat16, - "For fp4_e2m1fn_x2 query, value must be bf16"); - } else { - // MXFP8 (e5m2 / e4m3fn) requires the value tensor to match the query fp8 dtype. - TORCH_CHECK(v.dtype() == q_dtype, - "For fp8 query, value must have the same fp8 dtype as query"); - } - // For FP4 (e2m1) the head dimension is packed (2 values per byte); unpack for padding math. - const int head_size_actual = (q_dtype == torch::kFloat4_e2m1fn_x2) ? head_size_og * 2 : head_size_og; - TORCH_CHECK(head_size_actual <= 256 || head_size_actual == 512, - "FlashAttention XPU varlen only supports head dimension up to 256 or exactly 512. Got: " + std::to_string(head_size_actual)); + TORCH_CHECK(v.size(-1) == head_size_og, "Value head dimension must match Query head dimension"); + TORCH_CHECK(head_size_og <= 256 || head_size_og == 512, + "FlashAttention XPU varlen only supports head dimension up to 256 or exactly 512. Got: " + std::to_string(head_size_og)); CHECK_DEVICE(q); CHECK_DEVICE(k); CHECK_DEVICE(v); - // XPU requires head_size to be a multiple of 32. - // FP4 packed inputs are consumed in their packed layout and are not padded. - const int head_size_padded = round_multiple(head_size_actual, 32); - const bool needs_padding = (head_size_actual != head_size_padded) - && (q_dtype != torch::kFloat4_e2m1fn_x2); - const int pad_size = head_size_padded - head_size_actual; + // XPU requires head_size to be a multiple of 32 + const int head_size_padded = round_multiple(head_size_og, 32); + const bool needs_padding = (head_size_og != head_size_padded); + const int pad_size = head_size_padded - head_size_og; auto maybe_pad = [&](const at::Tensor& t) -> at::Tensor { return needs_padding @@ -497,28 +432,11 @@ mha_varlen_fwd( at::Tensor k_padded = maybe_pad(k); at::Tensor v_padded = maybe_pad(v); - auto out_dtype = is_mxfp ? torch::kFloat32 : c10::typeMetaToScalarType(q_dtype); - at::Tensor out_padded; - if (out_.has_value()) { - // The kernel writes out_dtype elements; a mismatched buffer would be - // overrun (MXFP writes float32 into what may be a half-width tensor). - TORCH_CHECK(out_.value().scalar_type() == out_dtype, - "out must have dtype ", out_dtype, ", got ", out_.value().scalar_type()); - out_padded = maybe_pad(out_.value()); - } else if (is_mxfp) { - out_padded = torch::zeros({total_q, num_heads, head_size_padded}, - q_padded.options().dtype(out_dtype)); - } else { - out_padded = torch::zeros_like(q_padded); - } + at::Tensor out_padded = out_.has_value() ? maybe_pad(out_.value()) + : torch::zeros_like(q_padded); const bool is_local = (window_size_left != -1) || (window_size_right != -1); const bool is_paged = block_table_.has_value() && block_table_->defined(); - // policy_dispatch() only routes MXFP on the non-paged / non-rotary paths; any - // other combination silently falls through to the bf16 dispatcher and would - // reinterpret the fp8/fp4 bytes as bf16. Reject it here instead. - TORCH_CHECK(!(is_mxfp && is_paged), - "MXFP (fp8/fp4) forward does not support paged KV (block_table)."); q_padded = ensure_contiguous(q_padded); k_padded = ensure_contiguous(k_padded); @@ -543,7 +461,7 @@ mha_varlen_fwd( rng_state[1] = static_cast(philox_offset); } - const int q_tile_size = (head_size_actual <= 32) ? 64 : (head_size_actual <= 128) ? 128 : 256; + const int q_tile_size = (head_size_og <= 32) ? 64 : (head_size_og <= 128) ? 128 : 256; const int k_tile_size = (max_seqlen_q == 1 && !is_paged) ? 512 : 128; const int seqlen_q_rounded = round_multiple(max_seqlen_q, q_tile_size); const int seqlen_k_rounded = round_multiple(max_seqlen_k, k_tile_size); @@ -555,22 +473,6 @@ mha_varlen_fwd( S_dmask = torch::empty({0}, q.options()); } - // MXFP (block-scaled) scale factor tensors (optional; required for fp8/fp4). - std::optional scale_q = scale_q_.has_value() - ? std::optional(ensure_contiguous(scale_q_.value())) : std::nullopt; - std::optional scale_k = scale_k_.has_value() - ? std::optional(ensure_contiguous(scale_k_.value())) : std::nullopt; - std::optional scale_v = scale_v_.has_value() - ? std::optional(ensure_contiguous(scale_v_.value())) : std::nullopt; - std::optional cu_scale_q = cu_scale_q_.has_value() - ? std::optional(ensure_contiguous(cu_scale_q_.value())) : std::nullopt; - std::optional cu_scale_kv = cu_scale_kv_.has_value() - ? std::optional(ensure_contiguous(cu_scale_kv_.value())) : std::nullopt; - if (is_mxfp) { - TORCH_CHECK(scale_q.has_value() && scale_k.has_value() && scale_v.has_value(), - "MXFP (fp8/fp4) varlen forward currently requires scale_q, scale_k and scale_v tensors."); - } - cutlass_fmha_fwd_varlen_impl( queue, q_padded, k_padded, v_padded, out_padded, @@ -583,13 +485,12 @@ mha_varlen_fwd( true, is_paged, is_causal, is_local, p_dropout, philox_seed, philox_offset, nullptr, return_softmax ? S_dmask.data_ptr() : nullptr, - seqlen_q_rounded, seqlen_k_rounded, - scale_q, scale_k, scale_v, cu_scale_q, cu_scale_kv); + seqlen_q_rounded, seqlen_k_rounded); - // Strip padding from output back to original (unpacked) head_size + // Strip padding from output back to original head_size at::Tensor out = needs_padding ? out_padded.index({torch::indexing::Slice(), torch::indexing::Slice(), - torch::indexing::Slice(0, head_size_actual)}) + torch::indexing::Slice(0, head_size_og)}) .contiguous() : ensure_contiguous(out_padded); @@ -1037,13 +938,7 @@ mha_fwd( const int64_t window_size_right, const double softcap, const bool return_softmax, - c10::optional gen_, - c10::optional scale_q_, - c10::optional scale_k_, - c10::optional scale_v_) { - auto to_std = [](const c10::optional& o) -> std::optional { - return o.has_value() ? std::optional(o.value()) : std::nullopt; - }; + c10::optional gen_) { return FLASH_NAMESPACE::mha_fwd( q, k, @@ -1057,10 +952,7 @@ mha_fwd( static_cast(window_size_right), static_cast(softcap), return_softmax, - gen_, - to_std(scale_q_), - to_std(scale_k_), - to_std(scale_v_) + gen_ ); } @@ -1149,12 +1041,7 @@ mha_varlen_fwd( int64_t window_size_right, const double softcap, const bool return_softmax, - std::optional gen_, - std::optional scale_q_, - std::optional scale_k_, - std::optional scale_v_, - std::optional cu_scale_q_, - std::optional cu_scale_kv_) { + std::optional gen_) { return FLASH_NAMESPACE::mha_varlen_fwd( const_cast(q), k, @@ -1178,12 +1065,7 @@ mha_varlen_fwd( static_cast(window_size_right), static_cast(softcap), return_softmax, - gen_, - scale_q_, - scale_k_, - scale_v_, - cu_scale_q_, - cu_scale_kv_ + gen_ ); } diff --git a/flash-attn2/flash_attn_xpu/src/collective/fmha_fusion.hpp b/flash-attn2/flash_attn_xpu/src/collective/fmha_fusion.hpp index 6dcfeadbc..a87752588 100644 --- a/flash-attn2/flash_attn_xpu/src/collective/fmha_fusion.hpp +++ b/flash-attn2/flash_attn_xpu/src/collective/fmha_fusion.hpp @@ -40,7 +40,6 @@ using namespace cute; struct VariableLength { int max_length; int* cumulative_length = nullptr; - int* cumulative_scale_length = nullptr; // For MXFP varlen scale indexing CUTE_HOST_DEVICE operator int() const { return max_length; diff --git a/flash-attn2/flash_attn_xpu/src/collective/fmha_fwd_mainloop_xe2.hpp b/flash-attn2/flash_attn_xpu/src/collective/fmha_fwd_mainloop_xe2.hpp index aff675be5..8f7a5b1e7 100644 --- a/flash-attn2/flash_attn_xpu/src/collective/fmha_fwd_mainloop_xe2.hpp +++ b/flash-attn2/flash_attn_xpu/src/collective/fmha_fwd_mainloop_xe2.hpp @@ -18,7 +18,6 @@ #include "../philox.hpp" #include "./fmha_fwd_common.hpp" -#include "cutlass/gemm/collective/xe_common_blockscaled_mxfp.hpp" namespace cutlass::fmha { @@ -49,12 +48,7 @@ template < class TiledCopyQ_ = void, class TiledCopyK_ = void, class TiledCopyV_ = void, - bool HasRotary_ = false, - bool UseScale_ = false, // MXFP: block-scaled MMA - bool F8kvF16mma_ = false, // MXFP: F8 KV with F16 MMA (scalar scale) - class TensorScaleQ_ = void, // MXFP: optional scale tensors - class TensorScaleK_ = void, - class TensorScaleV_ = void> + bool HasRotary_ = false> struct FMHAFwdMainloopXe2 { static_assert( cutlass::detail::dependent_false, @@ -78,12 +72,7 @@ template < class TiledCopyQ_, class TiledCopyK_, class TiledCopyV_, - bool HasRotary_, - bool UseScale_, - bool F8kvF16mma_, - class TensorScaleQ_, - class TensorScaleK_, - class TensorScaleV_> + bool HasRotary_> struct FMHAFwdMainloopXe2< Xe2, CausalMask_, @@ -99,12 +88,7 @@ struct FMHAFwdMainloopXe2< TiledCopyQ_, TiledCopyK_, TiledCopyV_, - HasRotary_, - UseScale_, - F8kvF16mma_, - TensorScaleQ_, - TensorScaleK_, - TensorScaleV_> { + HasRotary_> { // Pull in common type aliases from the shared traits. using Traits = FMHAFwdMainloopTraits< @@ -144,61 +128,6 @@ struct FMHAFwdMainloopXe2< static constexpr bool PagedKV = PagedKV_; static constexpr bool HasRotary = HasRotary_; - // --- MXFP (block-scaled) support --------------------------------------- - using ElementQ = typename TensorQ::element_type; - static constexpr bool FP4Input = - cute::is_same_v; - static constexpr bool UseScale = UseScale_; - static constexpr bool F8kvF16mma = F8kvF16mma_; - - using TensorScaleQ = TensorScaleQ_; - using TensorScaleK = TensorScaleK_; - using TensorScaleV = TensorScaleV_; - using TensorScaleQ2D = conditional_t, void, - decltype(TensorScaleQ_{}(append>(make_coord(_, _), 0)))>; - using TensorScaleK2D = conditional_t, void, - decltype(TensorScaleK_{}(append>(make_coord(_, _), 0)))>; - using TensorScaleV2D = conditional_t, void, - decltype(TensorScaleV_{}(append>(make_coord(_, _), 0)))>; - - using ElementScaleQ = conditional_t, float, - typename TensorScaleQ::element_type>; - using ElementScaleK = conditional_t, float, - typename TensorScaleK::element_type>; - using ElementScaleV = conditional_t, float, - typename TensorScaleV::element_type>; - - // Block and MMA dimensions used for block-scale addressing. - static constexpr int BLK_Q = get<0>(TileShapeQK{}); - static constexpr int BLK_K = get<1>(TileShapeQK{}); - static constexpr int BLK_QK_D = get<2>(TileShapeQK{}); - static constexpr int BLK_P = get<0>(TileShapePV{}); - static constexpr int BLK_V = get<1>(TileShapePV{}); - static constexpr int BLK_PV_D = get<2>(TileShapePV{}); - - static constexpr int ATOM_Q = get<1>(typename TiledMMAQK::ThrLayoutVMNK{}.shape()); - static constexpr int ATOM_K = get<2>(typename TiledMMAQK::ThrLayoutVMNK{}.shape()); - static constexpr int ATOM_QK_D = get<3>(typename TiledMMAQK::ThrLayoutVMNK{}.shape()); - static constexpr int MMA_QK_D = get<2>(typename TiledMMAQK::Shape_MNK{}); - static constexpr int ATOM_P = get<1>(typename TiledMMAPV::ThrLayoutVMNK{}.shape()); - static constexpr int ATOM_V = get<2>(typename TiledMMAPV::ThrLayoutVMNK{}.shape()); - static constexpr int ATOM_PV_D = get<3>(typename TiledMMAPV::ThrLayoutVMNK{}.shape()); - static constexpr int MMA_PV_D = get<2>(typename TiledMMAPV::Shape_MNK{}); - - static constexpr int SG_Q = ceil_div(BLK_Q, ATOM_Q); - static constexpr int SG_K = ceil_div(BLK_K, ATOM_K); - static constexpr int SG_QK_D = ceil_div(BLK_QK_D, ATOM_QK_D); - static constexpr int SG_P = ceil_div(BLK_P, ATOM_P); - static constexpr int SG_V = ceil_div(BLK_V, ATOM_V); - static constexpr int SG_PV_D = ceil_div(BLK_PV_D, ATOM_PV_D); - - static constexpr auto GROUP_K = 32; - - using DefScaleType = cutlass::float_ue8m0_t; - using ElementScaleP = cute::conditional_t; - // ----------------------------------------------------------------------- - - // User-facing arguments struct Arguments { ElementS const scale; @@ -221,9 +150,6 @@ struct FMHAFwdMainloopXe2< const typename TensorQ::element_type* rotary_sin = nullptr; int rotary_dim = 0; bool is_rotary_interleaved = true; - // MXFP: scalar scales for F8 KV with F16 MMA path. - float scale_k_scalar = 1.0f; - float scale_v_scalar = 1.0f; }; struct LocalMaskFields { @@ -266,9 +192,6 @@ struct FMHAFwdMainloopXe2< paged; [[no_unique_address]] conditional_t rotary; - // MXFP: scalar scales for F8 KV with F16 MMA path. - float scale_k_scalar = 1.0f; - float scale_v_scalar = 1.0f; }; // SLM data @@ -308,8 +231,6 @@ struct FMHAFwdMainloopXe2< p.rotary = {args.rotary_cos, args.rotary_sin, args.rotary_dim, args.is_rotary_interleaved}; } - p.scale_k_scalar = args.scale_k_scalar; - p.scale_v_scalar = args.scale_v_scalar; return p; } @@ -339,11 +260,7 @@ struct FMHAFwdMainloopXe2< int head_q, int num_heads, int q_offset_sg, - int rotary_base, - int l_coord = 0, - TensorScaleQ2D const& scaleQ_2D = TensorScaleQ2D{}, - TensorScaleK2D const& scaleK_2D = TensorScaleK2D{}, - TensorScaleV2D const& scaleV_2D = TensorScaleV2D{}) { + int rotary_base) { using namespace sycl::ext::oneapi::this_work_item; auto tile_shape_v = @@ -527,72 +444,9 @@ struct FMHAFwdMainloopXe2< } } copy(copy_k, tKgK_cache(_, _, _, D), tKrK); - if constexpr (FP4Input) { - copy(tQrQ, tSrQ); - copy(tKrK, tSrK); - } else { - reorder(tQrQ, tSrQ); - reorder(tKrK, tSrK); - } - - if constexpr (UseScale) { -#if defined(SYCL_INTEL_TARGET) && (SYCL_INTEL_TARGET == 35) - // Block-scaled MMA with scale factors for Q*K. The block-scaled - // (BDPAS) atom and helpers only exist on CRI (SYCL_INTEL_TARGET==35). - // On BMG this kernel is instantiated (host references it) but never - // executed, so the body is replaced by an invalid-control-path stub. - const int q_coord = get<0>(blk_qv) * BLK_Q + - ((thr_id / intel::sg_size) / ATOM_K) * SG_Q; - const int k_coord = K * BLK_K + - ((thr_id % intel::sg_size) / ATOM_K) * SG_K; - - auto [tiled_copy_scaleQ, copy_iter_scaleQ, fragment_scaleQ] = - gemm::collective::make_scaled_copy(scaleQ_2D, q_coord, l_coord, size<4>(tKgK)); - auto [tiled_copy_scaleK, copy_iter_scaleK, fragment_scaleK] = - gemm::collective::make_scaled_copy(scaleK_2D, k_coord, l_coord, size<4>(tKgK)); - auto [gemm_qm_offsets, gemm_kn_offsets, gemm_qk_offsets, gemm_kk_offsets] = - gemm::collective::make_scaled_offsets< - decltype(size<1>(tSrQ.shape()))::value, - decltype(size<1>(tSrK.shape()))::value, - decltype(size<2>(tSrK.shape()))::value, - MMA_QK_D, - GROUP_K, - typename decltype(tiled_copy_scaleQ)::BlockShape, - typename decltype(tiled_copy_scaleK)::BlockShape>(); - - copy(tiled_copy_scaleQ, copy_iter_scaleQ(_, _, _, D), fragment_scaleQ); - copy(tiled_copy_scaleK, copy_iter_scaleK(_, _, _, D), fragment_scaleK); - - using scaleQSize = decltype(size(fragment_scaleQ)); - using scaleKSize = decltype(size(fragment_scaleK)); - - Tensor scaleQ_view = make_tensor( - recast>(fragment_scaleQ).data(), - make_layout(Shape<_1, decltype(size<1>(tSrQ.shape())), _1>{}, Stride<_1, _0, _0>{})); - Tensor scaleK_view = make_tensor( - recast>(fragment_scaleK).data(), - make_layout(Shape<_1, decltype(size<1>(tSrK.shape())), _1>{}, Stride<_1, _0, _0>{})); - - auto zipped_q = make_zip_tensor(tSrQ, scaleQ_view, gemm_qm_offsets, gemm_qk_offsets); - auto zipped_k = make_zip_tensor(tSrK, scaleK_view, gemm_kn_offsets, gemm_kk_offsets); - - cute::gemm(mma_qk, zipped_q, zipped_k, tSrS); -#else - CUTE_INVALID_CONTROL_PATH( - "Block-scaled (MXFP) Q*K path is only available on CRI " - "(SYCL_INTEL_TARGET==35)."); -#endif - } else { - if constexpr (F8kvF16mma) { - // F8 KV with F16 MMA path - apply scalar scale to K. - for (int i = 0; i < tSrK.size(); i++) - tSrK(i) = static_cast( - params.scale_k_scalar * static_cast(tSrK(i))); - } - cute::gemm(mma_qk, tSrQ, tSrK, tSrS); - } + reorder(tQrQ, tSrQ); + reorder(tKrK, tSrK); + cute::gemm(mma_qk, tSrQ, tSrK, tSrS); } auto cS_thread = thr_mma_qk.partition_C(gP_all(_, _, K)); @@ -726,61 +580,7 @@ struct FMHAFwdMainloopXe2< for (int VV = 0; VV < VTiles; VV++) { copy(copy_v, tVgV_cache(_, _, _, VV), tVrV); reorder(tVrV, tArV); - if constexpr (UseScale && !FP4Input) { -#if defined(SYCL_INTEL_TARGET) && (SYCL_INTEL_TARGET == 35) - // Block-scaled MMA with scale factors for P*V. CRI-only (BDPAS); - // on BMG this branch is instantiated but never runs. - const int v_coord = get<1>(blk_qv) * VTiles * BLK_V + VV * BLK_V + - ((thr_id % intel::sg_size) / ATOM_V) * SG_V; - - // P is a dummy scale, the same shape as V. - auto [tiled_copy_scaleP, copy_iter_scaleP, fragment_scaleP] = - gemm::collective::make_scaled_copy(scaleV_2D); - auto [tiled_copy_scaleV, copy_iter_scaleV, fragment_scaleV] = - gemm::collective::make_scaled_copy(scaleV_2D, v_coord, l_coord, blk_k1); - auto [gemm_p_offsets, gemm_v_offsets, gemm_pk_offsets, gemm_vk_offsets] = - gemm::collective::make_scaled_offsets< - decltype(size<1>(tArP.shape()))::value, - decltype(size<1>(tArV.shape()))::value, - decltype(size<2>(tArV.shape()))::value, - MMA_PV_D, - GROUP_K, - typename decltype(tiled_copy_scaleP)::BlockShape, - typename decltype(tiled_copy_scaleV)::BlockShape>(); - - fill(fragment_scaleP, ElementScaleV(1)); - copy(tiled_copy_scaleV, copy_iter_scaleV(_, _, _, K), fragment_scaleV); - - using scalePSize = decltype(size(fragment_scaleP)); - using scaleVSize = decltype(size(fragment_scaleV)); - - Tensor scaleV_view = make_tensor( - recast>(fragment_scaleV).data(), - make_layout(Shape<_1, decltype(size<1>(tArV.shape())), _1>{}, Stride<_1, _0, _0>{})); - Tensor scaleP_view = make_tensor( - recast>(fragment_scaleP).data(), - make_layout(Shape<_1, decltype(size<1>(tArP.shape())), _1>{}, Stride<_1, _0, _0>{})); - - auto zipped_p = make_zip_tensor(tArP, scaleP_view, gemm_p_offsets, gemm_pk_offsets); - auto zipped_v = make_zip_tensor(tArV, scaleV_view, gemm_v_offsets, gemm_vk_offsets); - - cute::gemm(mma_pv, zipped_p, zipped_v, tArA(_, _, _, VV)); -#else - CUTE_INVALID_CONTROL_PATH( - "Block-scaled (MXFP) P*V path is only available on CRI " - "(SYCL_INTEL_TARGET==35)."); -#endif - } else { - if constexpr (F8kvF16mma) { - // F8 KV with F16 MMA path - apply scalar scale to V. - for (int i = 0; i < tArV.size(); i++) - tArV(i) = static_cast( - params.scale_v_scalar * static_cast(tArV(i))); - } - cute::gemm(mma_pv, tArP, tArV, tArA(_, _, _, VV)); - } + cute::gemm(mma_pv, tArP, tArV, tArA(_, _, _, VV)); } if constexpr (PagedKV) { diff --git a/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim128_fix_mxfp.cpp b/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim128_fix_mxfp.cpp deleted file mode 100644 index abe77f4bc..000000000 --- a/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim128_fix_mxfp.cpp +++ /dev/null @@ -1,11 +0,0 @@ -#include "fmha_fwd_impl.hpp" - -// MXFP (block-scaled FP8/FP4) fixed mode: IsVarLen=0, IsPaged=0 -// Instantiates prefill + decode for all three element types (e5m2, e4m3, e2m1). - -template void policy_dispatch_mxfp_e5m2(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e5m2(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e4m3(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e4m3(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e2m1(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e2m1(sycl::queue&, const fmha_fwd_args_t&); diff --git a/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim128_varlen_mxfp.cpp b/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim128_varlen_mxfp.cpp deleted file mode 100644 index f6dfec1a1..000000000 --- a/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim128_varlen_mxfp.cpp +++ /dev/null @@ -1,11 +0,0 @@ -#include "fmha_fwd_impl.hpp" - -// MXFP (block-scaled FP8/FP4) varlen mode: IsVarLen=1, IsPaged=0 -// Instantiates prefill + decode for all three element types (e5m2, e4m3, e2m1). - -template void policy_dispatch_mxfp_e5m2(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e5m2(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e4m3(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e4m3(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e2m1(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e2m1(sycl::queue&, const fmha_fwd_args_t&); diff --git a/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim160_fix_mxfp.cpp b/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim160_fix_mxfp.cpp deleted file mode 100644 index c3c29808e..000000000 --- a/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim160_fix_mxfp.cpp +++ /dev/null @@ -1,11 +0,0 @@ -#include "fmha_fwd_impl.hpp" - -// MXFP (block-scaled FP8/FP4) fixed mode: IsVarLen=0, IsPaged=0 -// Instantiates prefill + decode for all three element types (e5m2, e4m3, e2m1). - -template void policy_dispatch_mxfp_e5m2(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e5m2(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e4m3(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e4m3(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e2m1(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e2m1(sycl::queue&, const fmha_fwd_args_t&); diff --git a/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim160_varlen_mxfp.cpp b/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim160_varlen_mxfp.cpp deleted file mode 100644 index ca18dd258..000000000 --- a/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim160_varlen_mxfp.cpp +++ /dev/null @@ -1,11 +0,0 @@ -#include "fmha_fwd_impl.hpp" - -// MXFP (block-scaled FP8/FP4) varlen mode: IsVarLen=1, IsPaged=0 -// Instantiates prefill + decode for all three element types (e5m2, e4m3, e2m1). - -template void policy_dispatch_mxfp_e5m2(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e5m2(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e4m3(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e4m3(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e2m1(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e2m1(sycl::queue&, const fmha_fwd_args_t&); diff --git a/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim192_fix_mxfp.cpp b/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim192_fix_mxfp.cpp deleted file mode 100644 index 371e3f1dd..000000000 --- a/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim192_fix_mxfp.cpp +++ /dev/null @@ -1,11 +0,0 @@ -#include "fmha_fwd_impl.hpp" - -// MXFP (block-scaled FP8/FP4) fixed mode: IsVarLen=0, IsPaged=0 -// Instantiates prefill + decode for all three element types (e5m2, e4m3, e2m1). - -template void policy_dispatch_mxfp_e5m2(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e5m2(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e4m3(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e4m3(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e2m1(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e2m1(sycl::queue&, const fmha_fwd_args_t&); diff --git a/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim192_varlen_mxfp.cpp b/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim192_varlen_mxfp.cpp deleted file mode 100644 index 06d5f7aba..000000000 --- a/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim192_varlen_mxfp.cpp +++ /dev/null @@ -1,11 +0,0 @@ -#include "fmha_fwd_impl.hpp" - -// MXFP (block-scaled FP8/FP4) varlen mode: IsVarLen=1, IsPaged=0 -// Instantiates prefill + decode for all three element types (e5m2, e4m3, e2m1). - -template void policy_dispatch_mxfp_e5m2(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e5m2(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e4m3(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e4m3(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e2m1(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e2m1(sycl::queue&, const fmha_fwd_args_t&); diff --git a/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim256_fix_mxfp.cpp b/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim256_fix_mxfp.cpp deleted file mode 100644 index ae8814090..000000000 --- a/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim256_fix_mxfp.cpp +++ /dev/null @@ -1,11 +0,0 @@ -#include "fmha_fwd_impl.hpp" - -// MXFP (block-scaled FP8/FP4) fixed mode: IsVarLen=0, IsPaged=0 -// Instantiates prefill + decode for all three element types (e5m2, e4m3, e2m1). - -template void policy_dispatch_mxfp_e5m2(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e5m2(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e4m3(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e4m3(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e2m1(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e2m1(sycl::queue&, const fmha_fwd_args_t&); diff --git a/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim256_varlen_mxfp.cpp b/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim256_varlen_mxfp.cpp deleted file mode 100644 index fb5641904..000000000 --- a/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim256_varlen_mxfp.cpp +++ /dev/null @@ -1,11 +0,0 @@ -#include "fmha_fwd_impl.hpp" - -// MXFP (block-scaled FP8/FP4) varlen mode: IsVarLen=1, IsPaged=0 -// Instantiates prefill + decode for all three element types (e5m2, e4m3, e2m1). - -template void policy_dispatch_mxfp_e5m2(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e5m2(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e4m3(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e4m3(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e2m1(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e2m1(sycl::queue&, const fmha_fwd_args_t&); diff --git a/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim32_fix_mxfp.cpp b/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim32_fix_mxfp.cpp deleted file mode 100644 index d0650f8df..000000000 --- a/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim32_fix_mxfp.cpp +++ /dev/null @@ -1,11 +0,0 @@ -#include "fmha_fwd_impl.hpp" - -// MXFP (block-scaled FP8/FP4) fixed mode: IsVarLen=0, IsPaged=0 -// Instantiates prefill + decode for all three element types (e5m2, e4m3, e2m1). - -template void policy_dispatch_mxfp_e5m2(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e5m2(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e4m3(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e4m3(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e2m1(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e2m1(sycl::queue&, const fmha_fwd_args_t&); diff --git a/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim32_varlen_mxfp.cpp b/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim32_varlen_mxfp.cpp deleted file mode 100644 index e6476868a..000000000 --- a/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim32_varlen_mxfp.cpp +++ /dev/null @@ -1,11 +0,0 @@ -#include "fmha_fwd_impl.hpp" - -// MXFP (block-scaled FP8/FP4) varlen mode: IsVarLen=1, IsPaged=0 -// Instantiates prefill + decode for all three element types (e5m2, e4m3, e2m1). - -template void policy_dispatch_mxfp_e5m2(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e5m2(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e4m3(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e4m3(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e2m1(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e2m1(sycl::queue&, const fmha_fwd_args_t&); diff --git a/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim512_fix_mxfp.cpp b/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim512_fix_mxfp.cpp deleted file mode 100644 index 5f9097589..000000000 --- a/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim512_fix_mxfp.cpp +++ /dev/null @@ -1,11 +0,0 @@ -#include "fmha_fwd_impl.hpp" - -// MXFP (block-scaled FP8/FP4) fixed mode: IsVarLen=0, IsPaged=0 -// Instantiates prefill + decode for all three element types (e5m2, e4m3, e2m1). - -template void policy_dispatch_mxfp_e5m2(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e5m2(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e4m3(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e4m3(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e2m1(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e2m1(sycl::queue&, const fmha_fwd_args_t&); diff --git a/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim512_varlen_mxfp.cpp b/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim512_varlen_mxfp.cpp deleted file mode 100644 index b7ba4b1ed..000000000 --- a/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim512_varlen_mxfp.cpp +++ /dev/null @@ -1,11 +0,0 @@ -#include "fmha_fwd_impl.hpp" - -// MXFP (block-scaled FP8/FP4) varlen mode: IsVarLen=1, IsPaged=0 -// Instantiates prefill + decode for all three element types (e5m2, e4m3, e2m1). - -template void policy_dispatch_mxfp_e5m2(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e5m2(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e4m3(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e4m3(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e2m1(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e2m1(sycl::queue&, const fmha_fwd_args_t&); diff --git a/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim64_fix_mxfp.cpp b/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim64_fix_mxfp.cpp deleted file mode 100644 index b135019c4..000000000 --- a/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim64_fix_mxfp.cpp +++ /dev/null @@ -1,11 +0,0 @@ -#include "fmha_fwd_impl.hpp" - -// MXFP (block-scaled FP8/FP4) fixed mode: IsVarLen=0, IsPaged=0 -// Instantiates prefill + decode for all three element types (e5m2, e4m3, e2m1). - -template void policy_dispatch_mxfp_e5m2(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e5m2(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e4m3(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e4m3(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e2m1(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e2m1(sycl::queue&, const fmha_fwd_args_t&); diff --git a/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim64_varlen_mxfp.cpp b/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim64_varlen_mxfp.cpp deleted file mode 100644 index 803c3db80..000000000 --- a/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim64_varlen_mxfp.cpp +++ /dev/null @@ -1,11 +0,0 @@ -#include "fmha_fwd_impl.hpp" - -// MXFP (block-scaled FP8/FP4) varlen mode: IsVarLen=1, IsPaged=0 -// Instantiates prefill + decode for all three element types (e5m2, e4m3, e2m1). - -template void policy_dispatch_mxfp_e5m2(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e5m2(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e4m3(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e4m3(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e2m1(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e2m1(sycl::queue&, const fmha_fwd_args_t&); diff --git a/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim96_fix_mxfp.cpp b/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim96_fix_mxfp.cpp deleted file mode 100644 index c3418a43d..000000000 --- a/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim96_fix_mxfp.cpp +++ /dev/null @@ -1,11 +0,0 @@ -#include "fmha_fwd_impl.hpp" - -// MXFP (block-scaled FP8/FP4) fixed mode: IsVarLen=0, IsPaged=0 -// Instantiates prefill + decode for all three element types (e5m2, e4m3, e2m1). - -template void policy_dispatch_mxfp_e5m2(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e5m2(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e4m3(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e4m3(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e2m1(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e2m1(sycl::queue&, const fmha_fwd_args_t&); diff --git a/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim96_varlen_mxfp.cpp b/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim96_varlen_mxfp.cpp deleted file mode 100644 index e95148fac..000000000 --- a/flash-attn2/flash_attn_xpu/src/flash_fwd_hdim96_varlen_mxfp.cpp +++ /dev/null @@ -1,11 +0,0 @@ -#include "fmha_fwd_impl.hpp" - -// MXFP (block-scaled FP8/FP4) varlen mode: IsVarLen=1, IsPaged=0 -// Instantiates prefill + decode for all three element types (e5m2, e4m3, e2m1). - -template void policy_dispatch_mxfp_e5m2(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e5m2(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e4m3(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e4m3(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e2m1(sycl::queue&, const fmha_fwd_args_t&); -template void policy_dispatch_mxfp_e2m1(sycl::queue&, const fmha_fwd_args_t&); diff --git a/flash-attn2/flash_attn_xpu/src/fmha_fwd.cpp b/flash-attn2/flash_attn_xpu/src/fmha_fwd.cpp index efca27d71..f4cb5ca6e 100644 --- a/flash-attn2/flash_attn_xpu/src/fmha_fwd.cpp +++ b/flash-attn2/flash_attn_xpu/src/fmha_fwd.cpp @@ -225,12 +225,7 @@ void cutlass_fmha_fwd_varlen_impl( void* rng_state, void* s_dmask, int seqlen_q_rounded, - int seqlen_k_rounded, - const std::optional& scale_q, - const std::optional& scale_k, - const std::optional& scale_v, - const std::optional& cu_scale_q, - const std::optional& cu_scale_kv) { + int seqlen_k_rounded) { int batch_size, num_heads_q, num_heads_kv, head_size; int total_seqlen_q, total_seqlen_k; int num_blocks, block_size, max_blocks_per_seq; @@ -314,13 +309,6 @@ void cutlass_fmha_fwd_varlen_impl( seqlen_q_rounded, seqlen_k_rounded}; - // MXFP block-scaled scale factor tensors (only set for FP8/FP4 paths). - if (scale_q.has_value()) args.scale_q = scale_q->data_ptr(); - if (scale_k.has_value()) args.scale_k = scale_k->data_ptr(); - if (scale_v.has_value()) args.scale_v = scale_v->data_ptr(); - if (cu_scale_q.has_value()) args.cu_scale_q = cu_scale_q->data_ptr(); - if (cu_scale_kv.has_value())args.cu_scale_kv = cu_scale_kv->data_ptr(); - const CutlassType cuType = aten_to_Cutlass_dtype(query); if (args.max_queries == 1) { dispatch_fwd_varlen_decode_by_head(queue, cuType, args, args.head_size); @@ -347,10 +335,7 @@ void cutlass_fmha_fwd_fix_impl( void* rng_state, void* s_dmask, int seqlen_q_rounded, - int seqlen_k_rounded, - const std::optional& scale_q, - const std::optional& scale_k, - const std::optional& scale_v) { + int seqlen_k_rounded) { const int batch_size = query.size(0); const int max_seqlen_q = query.size(1); const int num_heads_q = query.size(2); @@ -399,11 +384,6 @@ void cutlass_fmha_fwd_fix_impl( seqlen_q_rounded, seqlen_k_rounded}; - // MXFP block-scaled scale factor tensors (only set for FP8/FP4 paths). - if (scale_q.has_value()) args.scale_q = scale_q->data_ptr(); - if (scale_k.has_value()) args.scale_k = scale_k->data_ptr(); - if (scale_v.has_value()) args.scale_v = scale_v->data_ptr(); - const CutlassType cuType = aten_to_Cutlass_dtype(query); const int h = args.head_size; diff --git a/flash-attn2/flash_attn_xpu/src/fmha_fwd.hpp b/flash-attn2/flash_attn_xpu/src/fmha_fwd.hpp index eb7b86565..bc1c7e42a 100644 --- a/flash-attn2/flash_attn_xpu/src/fmha_fwd.hpp +++ b/flash-attn2/flash_attn_xpu/src/fmha_fwd.hpp @@ -2,16 +2,20 @@ #include "fmha_fwd_types.hpp" -#include - -#include - namespace sycl { inline namespace _V1 { class queue; } } +namespace at { + class Tensor; +} + +namespace std { + template class optional; +} + struct prefill_policy_head32; struct prefill_policy_head64; struct prefill_policy_head96; @@ -52,36 +56,8 @@ void policy_dispatch_bf16( sycl::queue& queue, const fmha_fwd_args_t& args); -// MXFP (block-scaled FP8/FP4) dispatch functions. Unconditional: the host pass -// (always SYCL_INTEL_TARGET==20) must own these symbols so they are callable at -// runtime on CRI. The block-scaled device code is gated inside FMHAConfig and -// the mainloop; on the BMG device pass these kernels compile as unreachable -// placeholders (they are only ever executed on CRI hardware). -template -void policy_dispatch_mxfp_e5m2( - sycl::queue& queue, - const fmha_fwd_args_t& args); - -template -void policy_dispatch_mxfp_e4m3( - sycl::queue& queue, - const fmha_fwd_args_t& args); - -template -void policy_dispatch_mxfp_e2m1( - sycl::queue& queue, - const fmha_fwd_args_t& args); - -// Combined dispatch (delegates to fp16/bf16/mxfp based on cuType) +// Combined dispatch (delegates to fp16/bf16 based on cuType) // Defined inline in header so callers (fmha_fwd.cpp) can see the template body. -// MXFP is only supported on the varlen/fix, non-paged, non-rotary forward paths -// (matching the CRI reference implementation). The `if constexpr` guard ensures -// the MXFP dispatchers are only ODR-used (and therefore only instantiated) for -// those paths, so the paged / kvcache / rotary instantiations do not pull in the -// (expensive) MXFP kernels. template inline void policy_dispatch( @@ -90,21 +66,9 @@ inline void policy_dispatch( const fmha_fwd_args_t& args) { if (cuType == CutlassType::half) { policy_dispatch_fp16(queue, args); - return; + } else { + policy_dispatch_bf16(queue, args); } - if constexpr (!IsPaged && !HasRotary) { - if (cuType == CutlassType::mx_float_e5m2) { - policy_dispatch_mxfp_e5m2(queue, args); - return; - } else if (cuType == CutlassType::mx_float_e4m3) { - policy_dispatch_mxfp_e4m3(queue, args); - return; - } else if (cuType == CutlassType::mx_float_e2m1) { - policy_dispatch_mxfp_e2m1(queue, args); - return; - } - } - policy_dispatch_bf16(queue, args); } // Varlen prefill mode extern declarations (IsVarLen=1, IsPaged=0/1) @@ -175,34 +139,6 @@ EXTERN_DISPATCH_KVCACHE_PAGED(256) EXTERN_DISPATCH_KVCACHE_PAGED(512) #undef EXTERN_DISPATCH_KVCACHE_PAGED -// MXFP (block-scaled FP8/FP4) extern declarations. Only the varlen prefill, -// varlen decode and fixed non-paged / non-rotary forward paths are supported -// (matching the CRI reference). Each head dimension instantiates all three -// MXFP element types (e5m2, e4m3, e2m1). -#define EXTERN_DISPATCH_MXFP(HDIM) \ - extern template void policy_dispatch_mxfp_e5m2(sycl::queue&, const fmha_fwd_args_t&); \ - extern template void policy_dispatch_mxfp_e5m2(sycl::queue&, const fmha_fwd_args_t&); \ - extern template void policy_dispatch_mxfp_e5m2(sycl::queue&, const fmha_fwd_args_t&); \ - extern template void policy_dispatch_mxfp_e5m2(sycl::queue&, const fmha_fwd_args_t&); \ - extern template void policy_dispatch_mxfp_e4m3(sycl::queue&, const fmha_fwd_args_t&); \ - extern template void policy_dispatch_mxfp_e4m3(sycl::queue&, const fmha_fwd_args_t&); \ - extern template void policy_dispatch_mxfp_e4m3(sycl::queue&, const fmha_fwd_args_t&); \ - extern template void policy_dispatch_mxfp_e4m3(sycl::queue&, const fmha_fwd_args_t&); \ - extern template void policy_dispatch_mxfp_e2m1(sycl::queue&, const fmha_fwd_args_t&); \ - extern template void policy_dispatch_mxfp_e2m1(sycl::queue&, const fmha_fwd_args_t&); \ - extern template void policy_dispatch_mxfp_e2m1(sycl::queue&, const fmha_fwd_args_t&); \ - extern template void policy_dispatch_mxfp_e2m1(sycl::queue&, const fmha_fwd_args_t&); - -EXTERN_DISPATCH_MXFP(32) -EXTERN_DISPATCH_MXFP(64) -EXTERN_DISPATCH_MXFP(96) -EXTERN_DISPATCH_MXFP(128) -EXTERN_DISPATCH_MXFP(160) -EXTERN_DISPATCH_MXFP(192) -EXTERN_DISPATCH_MXFP(256) -EXTERN_DISPATCH_MXFP(512) -#undef EXTERN_DISPATCH_MXFP - void cutlass_fmha_fwd_varlen_impl( sycl::queue& queue, const at::Tensor& query, @@ -228,13 +164,7 @@ void cutlass_fmha_fwd_varlen_impl( void* rng_state = nullptr, void* s_dmask = nullptr, int seqlen_q_rounded = 0, - int seqlen_k_rounded = 0, - // MXFP (block-scaled) scale factor tensors (optional). - const std::optional& scale_q = std::nullopt, - const std::optional& scale_k = std::nullopt, - const std::optional& scale_v = std::nullopt, - const std::optional& cu_scale_q = std::nullopt, - const std::optional& cu_scale_kv = std::nullopt); + int seqlen_k_rounded = 0); void cutlass_fmha_fwd_fix_impl( sycl::queue& queue, @@ -254,11 +184,7 @@ void cutlass_fmha_fwd_fix_impl( void* rng_state = nullptr, void* s_dmask = nullptr, int seqlen_q_rounded = 0, - int seqlen_k_rounded = 0, - // MXFP (block-scaled) scale factor tensors (optional). - const std::optional& scale_q = std::nullopt, - const std::optional& scale_k = std::nullopt, - const std::optional& scale_v = std::nullopt); + int seqlen_k_rounded = 0); void cutlass_fmha_fwd_kvcache_impl( sycl::queue& queue, diff --git a/flash-attn2/flash_attn_xpu/src/fmha_fwd_impl.hpp b/flash-attn2/flash_attn_xpu/src/fmha_fwd_impl.hpp index d691e0d42..e686170c8 100644 --- a/flash-attn2/flash_attn_xpu/src/fmha_fwd_impl.hpp +++ b/flash-attn2/flash_attn_xpu/src/fmha_fwd_impl.hpp @@ -27,22 +27,15 @@ template < typename SubgroupLayoutQK, typename SubgroupLayoutPV_, int PipelineStages, - bool UseScale_ = false, // MXFP: block-scaled MMA (BDPAS) - bool F8kvF16mma_ = false, // MXFP: F8 KV with F16 MMA (scalar scale) typename ElementQ = bfloat16_t, typename ElementK = bfloat16_t, typename ElementV = bfloat16_t, typename ElementO = bfloat16_t, - typename ElementScale = float, // MXFP: block-scale element type typename MMAOperation_ = void, - typename MMAOperationPV_ = void, typename StrideQ = Stride, typename StrideK = Stride, typename StrideV = Stride<_1, int, int, int>, typename StrideO = Stride, - typename StrideScaleQ = Stride<_1, int, int, int>, - typename StrideScaleK = Stride<_1, int, int, int>, - typename StrideScaleV = Stride<_1, int, int, int>, typename GmemTiledCopyQ = void, typename GmemTiledCopyK = void, typename GmemTiledCopyV = void, @@ -50,48 +43,9 @@ template < struct FMHAConfig { static constexpr int SGTileQ = get<0>(shape_div(TileShapeQK{}, shape(SubgroupLayoutQK{})))(); - static constexpr bool UseScale = UseScale_; - static constexpr bool F8kvF16mma = F8kvF16mma_; - - template - static constexpr bool is_f8_v = - cute::is_any_of_v; - template - static constexpr bool is_f16_v = - cute::is_any_of_v; - -#if !(defined(SYCL_INTEL_TARGET) && (SYCL_INTEL_TARGET == 35)) - // BMG path: no block-scaled DPAS and no FP8/FP4 DPAS atoms. Any MXFP input - // (e5m2 / e4m3 / e2m1) falls back to a valid half_t DPAS atom so the kernel - // is compilable on BMG. These kernels are only ever executed on CRI, so the - // BMG fallback body is an (unreachable) placeholder. - using DefaultMMA = cute::conditional_t< - cute::is_same_v || - cute::is_same_v || - cute::is_same_v, - XE_DPAS_TT, - XE_DPAS_TT>; - using MMAOperationPV = cute::conditional_t< - is_void_v, - DefaultMMA, - MMAOperationPV_>; -#else - // CRI path: block-scaled BDPAS when UseScale, else plain DPAS. - using DefaultDpasOp = XE_DPAS_TT; - using DefaultBdpasOp = XE_BDPAS_TT; - using DefaultMMA = - cute::conditional_t; - // PV MMA: F8 V uses the block-scaled/DPAS path; other V (e.g. bf16 for - // MXFP4) uses a plain DPAS atom keyed on ElementV. - using MMAOperationPV = cute::conditional_t< - cute::is_same_v || - cute::is_same_v, - cute::conditional_t, DefaultMMA, MMAOperationPV_>, - XE_DPAS_TT>; -#endif using MMAOperation = cute::conditional_t< is_void_v, - DefaultMMA, + XE_DPAS_TT, MMAOperation_>; using SubgroupLayoutPV = cute::conditional_t< is_void_v, @@ -116,7 +70,7 @@ struct FMHAConfig { Layout, SubgroupLayoutQK>::TiledMMA; using TiledMMAPV = typename TiledMMAHelper< - MMA_Atom, + MMA_Atom, Layout, SubgroupLayoutPV>::TiledMMA; @@ -136,15 +90,6 @@ struct FMHAConfig { using TensorV = decltype(make_dummy_tensor(ElementV{}, StrideV{})); using TensorO = decltype(make_dummy_tensor(ElementO{}, StrideO{})); - // MXFP: block-scale dummy tensors (harmless float tensors when the - // config is not block-scaled; the mainloop ignores them unless UseScale). - using TensorScaleQ = - decltype(make_dummy_tensor(ElementScale{}, StrideScaleQ{})); - using TensorScaleK = - decltype(make_dummy_tensor(ElementScale{}, StrideScaleK{})); - using TensorScaleV = - decltype(make_dummy_tensor(ElementScale{}, StrideScaleV{})); - using MainloopDispatchPolicy = cutlass::fmha::Xe2; using CollectiveMainloop = cutlass::fmha::collective::FMHAFwdMainloopXe2< MainloopDispatchPolicy, @@ -161,12 +106,7 @@ struct FMHAConfig { GmemTiledCopyQ, GmemTiledCopyK, GmemTiledCopyV, - HasRotary, - UseScale, - F8kvF16mma, - TensorScaleQ, - TensorScaleK, - TensorScaleV>; + HasRotary>; using CollectiveEpilogue = cutlass::fmha::collective::FMHAFwdEpilogueXe2< CollectiveMainloop, @@ -251,14 +191,7 @@ struct FMHAConfig { args.knew_batch_stride, args.knew_head_stride, args.knew_row_stride, reinterpret_cast(args.vnew), args.vnew_batch_stride, args.vnew_head_stride, args.vnew_row_stride, - args.seqlen_knew, - reinterpret_cast(args.scale_q), - StrideScaleQ{}, - reinterpret_cast(args.scale_k), - StrideScaleK{}, - reinterpret_cast(args.scale_v), - StrideScaleV{}, - GROUP_SIZE}, + args.seqlen_knew}, {static_cast(args.sm_scale), args.window_size_left, args.window_size_right, args.p_dropout, @@ -272,9 +205,7 @@ struct FMHAConfig { reinterpret_cast(args.rotary_cos), reinterpret_cast(args.rotary_sin), args.rotary_dim, - args.is_rotary_interleaved, - args.scale_k_scalar, - args.scale_v_scalar}, + args.is_rotary_interleaved}, {}, hw_info}; @@ -351,7 +282,6 @@ void policy_dispatch_fp16(sycl::queue& queue, const fmha_fwd_args_t& args) { typename chunk_policy::SubgroupLayoutQK, void, PipelineStages, - false, false, half_t, half_t, half_t, half_t>; Config::template xe2_dispatch(queue, args); } @@ -366,136 +296,8 @@ void policy_dispatch_bf16(sycl::queue& queue, const fmha_fwd_args_t& args) { typename chunk_policy::SubgroupLayoutQK, void, PipelineStages, - false, false, bfloat16_t, bfloat16_t, bfloat16_t, bfloat16_t>; Config::template xe2_dispatch(queue, args); } -// --------------------------------------------------------------------------- -// MXFP (block-scaled FP8/FP4) dispatch. -// Prefill (SGTileQ >= 8) uses block-scaled BDPAS (UseScale=true); decode -// (SGTileQ < 8) falls back to plain DPAS with UseScale=false. -// -// These dispatchers are UNCONDITIONAL (no SYCL_INTEL_TARGET guard): the host -// pass is always SYCL_INTEL_TARGET==20, and the host must own these symbols so -// they are callable at runtime on CRI. The actual block-scaled device code is -// gated inside FMHAConfig (MMA-atom selection) and the mainloop GEMM sections -// via `#if SYCL_INTEL_TARGET==35`, with compilable BMG fallbacks. On the BMG -// device pass these kernels are emitted (as unreachable placeholders) but never -// executed, since MXFP is only run on CRI hardware. -// --------------------------------------------------------------------------- -template -void policy_dispatch_mxfp_e5m2(sycl::queue& queue, const fmha_fwd_args_t& args) { - constexpr int SGTileQ = get<0>(shape_div( - typename chunk_policy::ShapeQK{}, - shape(typename chunk_policy::SubgroupLayoutQK{})))(); - using ElementType = cutlass::mx_float8_t; - using ElementQ = typename ElementType::DataType; - using ElementK = typename ElementType::DataType; - using ElementV = typename ElementType::DataType; - using ElementScale = typename ElementType::ScaleFactorType; - if constexpr (SGTileQ >= 8) { - using Config = FMHAConfig< - typename chunk_policy::ShapeQK, - typename chunk_policy::ShapePV, - typename chunk_policy::ShapeOut, - typename chunk_policy::SubgroupLayoutQK, - void, - PipelineStages, - true, // UseScale (prefill) - false, // F8kvF16mma - ElementQ, ElementK, ElementV, float, ElementScale>; - Config::template xe2_dispatch(queue, args); - } else { - using Config = FMHAConfig< - typename chunk_policy::ShapeQK, - typename chunk_policy::ShapePV, - typename chunk_policy::ShapeOut, - typename chunk_policy::SubgroupLayoutQK, - void, - PipelineStages, - false, // UseScale (decode) - false, // F8kvF16mma - ElementQ, ElementK, ElementV, float, ElementScale>; - Config::template xe2_dispatch(queue, args); - } -} - -template -void policy_dispatch_mxfp_e4m3(sycl::queue& queue, const fmha_fwd_args_t& args) { - constexpr int SGTileQ = get<0>(shape_div( - typename chunk_policy::ShapeQK{}, - shape(typename chunk_policy::SubgroupLayoutQK{})))(); - using ElementType = cutlass::mx_float8_t; - using ElementQ = typename ElementType::DataType; - using ElementK = typename ElementType::DataType; - using ElementV = typename ElementType::DataType; - using ElementScale = typename ElementType::ScaleFactorType; - if constexpr (SGTileQ >= 8) { - using Config = FMHAConfig< - typename chunk_policy::ShapeQK, - typename chunk_policy::ShapePV, - typename chunk_policy::ShapeOut, - typename chunk_policy::SubgroupLayoutQK, - void, - PipelineStages, - true, // UseScale (prefill) - false, // F8kvF16mma - ElementQ, ElementK, ElementV, float, ElementScale>; - Config::template xe2_dispatch(queue, args); - } else { - using Config = FMHAConfig< - typename chunk_policy::ShapeQK, - typename chunk_policy::ShapePV, - typename chunk_policy::ShapeOut, - typename chunk_policy::SubgroupLayoutQK, - void, - PipelineStages, - false, // UseScale (decode) - false, // F8kvF16mma - ElementQ, ElementK, ElementV, float, ElementScale>; - Config::template xe2_dispatch(queue, args); - } -} - -template -void policy_dispatch_mxfp_e2m1(sycl::queue& queue, const fmha_fwd_args_t& args) { - constexpr int SGTileQ = get<0>(shape_div( - typename chunk_policy::ShapeQK{}, - shape(typename chunk_policy::SubgroupLayoutQK{})))(); - using ElementType = cutlass::mx_float4_t; - using ElementQ = typename ElementType::DataType; - using ElementK = typename ElementType::DataType; - using ElementV = bfloat16_t; // V uses BF16 for MXFP4 - using ElementScale = typename ElementType::ScaleFactorType; - if constexpr (SGTileQ >= 8) { - using Config = FMHAConfig< - typename chunk_policy::ShapeQK, - typename chunk_policy::ShapePV, - typename chunk_policy::ShapeOut, - typename chunk_policy::SubgroupLayoutQK, - void, - PipelineStages, - true, // UseScale (prefill) - false, // F8kvF16mma - ElementQ, ElementK, ElementV, float, ElementScale>; - Config::template xe2_dispatch(queue, args); - } else { - using Config = FMHAConfig< - typename chunk_policy::ShapeQK, - typename chunk_policy::ShapePV, - typename chunk_policy::ShapeOut, - typename chunk_policy::SubgroupLayoutQK, - void, - PipelineStages, - false, // UseScale (decode) - false, // F8kvF16mma - ElementQ, ElementK, ElementV, float, ElementScale>; - Config::template xe2_dispatch(queue, args); - } -} - // Combined policy_dispatch is now defined inline in fmha_fwd.hpp diff --git a/flash-attn2/flash_attn_xpu/src/fmha_fwd_types.hpp b/flash-attn2/flash_attn_xpu/src/fmha_fwd_types.hpp index ccb481bad..7730facac 100644 --- a/flash-attn2/flash_attn_xpu/src/fmha_fwd_types.hpp +++ b/flash-attn2/flash_attn_xpu/src/fmha_fwd_types.hpp @@ -61,28 +61,12 @@ struct fmha_fwd_args_t { void* rotary_sin = nullptr; int rotary_dim = 0; bool is_rotary_interleaved = true; - - // Scale tensors for MXFP (block-scaled) operations - void* scale_q = nullptr; - void* scale_k = nullptr; - void* scale_v = nullptr; - void* cu_scale_q = nullptr; - void* cu_scale_kv = nullptr; - // Scalar scales for F8 KV with F16 MMA path - float scale_k_scalar = 1.0f; - float scale_v_scalar = 1.0f; }; enum class CutlassType { half, - bfloat16, - mx_float_e5m2, // MXFP8 E5M2 - mx_float_e4m3, // MXFP8 E4M3 - mx_float_e2m1 // MXFP4 E2M1 + bfloat16 }; constexpr int PipelineStages_Decode = 1; constexpr int PipelineStages_Prefill = 2; - -// Scale group size for block-scaled (MXFP) operations -constexpr int GROUP_SIZE = 32; diff --git a/flash-attn2/flash_attn_xpu/src/fmha_utils.hpp b/flash-attn2/flash_attn_xpu/src/fmha_utils.hpp index a0d2ddea0..f531800a9 100644 --- a/flash-attn2/flash_attn_xpu/src/fmha_utils.hpp +++ b/flash-attn2/flash_attn_xpu/src/fmha_utils.hpp @@ -11,28 +11,11 @@ if (input.scalar_type() == torch::kBFloat16) { return CutlassType::bfloat16; } - if (input.scalar_type() == torch::kFloat8_e5m2) { - return CutlassType::mx_float_e5m2; - } - if (input.scalar_type() == torch::kFloat8_e4m3fn) { - return CutlassType::mx_float_e4m3; - } - if (input.scalar_type() == torch::kFloat4_e2m1fn_x2) { - return CutlassType::mx_float_e2m1; - } TORCH_INTERNAL_ASSERT( - false, - "Current cutlass kernel only supports half/bf16/fp8_e5m2/fp8_e4m3/fp4_e2m1 data types."); + false, "Current cutlass kernel only support half/bf16 data type."); return {}; // unreachable; silences compiler warning } -/// Helper to detect MXFP (block-scaled) types that require scale tensors. -[[nodiscard]] inline bool requires_scale_tensors(CutlassType cuType) { - return cuType == CutlassType::mx_float_e5m2 || - cuType == CutlassType::mx_float_e4m3 || - cuType == CutlassType::mx_float_e2m1; -} - using namespace cute; struct prefill_policy_head32 { diff --git a/flash-attn2/flash_attn_xpu/src/kernel/fmha_fwd_kernel_xe2.hpp b/flash-attn2/flash_attn_xpu/src/kernel/fmha_fwd_kernel_xe2.hpp index 47c2cb15e..5f78f6b9b 100644 --- a/flash-attn2/flash_attn_xpu/src/kernel/fmha_fwd_kernel_xe2.hpp +++ b/flash-attn2/flash_attn_xpu/src/kernel/fmha_fwd_kernel_xe2.hpp @@ -67,13 +67,6 @@ class XeFMHAFwdKernelXe2 { static constexpr bool HasDropout = CollectiveMainloop::HasDropout; static constexpr bool PagedKV = CollectiveMainloop::PagedKV; - // MXFP (block-scaled) support. - static constexpr bool UseScale = CollectiveMainloop::UseScale; - using ElementScale = typename CollectiveMainloop::ElementScaleQ; - using StrideScaleQ = decltype(stride(typename CollectiveMainloop::TensorScaleQ{})); - using StrideScaleK = decltype(stride(typename CollectiveMainloop::TensorScaleK{})); - using StrideScaleV = decltype(stride(typename CollectiveMainloop::TensorScaleV{})); - using MainloopSharedStorage = typename CollectiveMainloop::SharedStorage; using EpilogueSharedStorage = typename CollectiveEpilogue::SharedStorage; union SharedStorage { @@ -119,14 +112,6 @@ class XeFMHAFwdKernelXe2 { int64_t vnew_head_stride = 0; int64_t vnew_row_stride = 0; int seqlen_knew = 0; - // MXFP (block-scaled) scale tensors. - const ElementScale* scaleQ = nullptr; - StrideScaleQ dScaleQ{}; - const ElementScale* scaleK = nullptr; - StrideScaleK dScaleK{}; - const ElementScale* scaleV = nullptr; - StrideScaleV dScaleV{}; - int group_size = 32; }; using KernelParams = KernelArguments; @@ -497,102 +482,28 @@ class XeFMHAFwdKernelXe2 { int l_coord_kv = is_var_len ? 0 : (PagedKV ? 0 : bidx); CollectiveMainloop mainloop(params.mainloop, shared_storage.mainloop); - if constexpr (UseScale) { - // MXFP: build block-scale tensors mirroring the Q/K/V layout. - auto scale_q_dim = cute::ceil_div(s.head_size_qk, p.group_size); - auto scale_k_dim = cute::ceil_div(s.head_size_qk, p.group_size); - auto scale_v_dim = cute::ceil_div(total_seqlen_kv, p.group_size); - - auto shape_scale_Q = - make_shape(seq_len_qo, scale_q_dim, s.num_heads_q, batch_dim_q); - auto shape_scale_K = - make_shape(total_seqlen_kv, scale_k_dim, s.num_heads_kv, batch_dim_kv); - auto shape_scale_V = - make_shape(s.head_size_vo, scale_v_dim, s.num_heads_kv, batch_dim_kv); - - int offset_scaleQ = 0, offset_scaleK = 0, offset_scaleV = 0; - if constexpr (is_var_len) { - auto qo_cumulative = s.seq_len_qo.cumulative_length; - auto kv_cumulative = s.seq_len_kv.cumulative_length; - offset_scaleQ = s.num_heads_q * scale_q_dim * qo_cumulative[idx_b]; - offset_scaleK = s.num_heads_kv * scale_k_dim * kv_cumulative[idx_b]; - auto kv_scale_cumulative = s.seq_len_kv.cumulative_scale_length; - if (kv_scale_cumulative != nullptr) { - offset_scaleV = - s.num_heads_kv * s.head_size_vo * kv_scale_cumulative[idx_b]; - } else { - offset_scaleV = s.num_heads_kv * scale_v_dim * kv_cumulative[idx_b]; - } - } - - auto layout_scaleQ = is_var_len - ? make_ordered_layout(shape_scale_Q, Step<_0, _1, _2, _3>{}) - : make_layout(shape_scale_Q, p.dScaleQ); - auto layout_scaleK = is_var_len - ? make_ordered_layout(shape_scale_K, Step<_0, _1, _2, _3>{}) - : make_layout(shape_scale_K, p.dScaleK); - auto layout_scaleV = is_var_len - ? make_ordered_layout(shape_scale_V, Step<_0, _1, _2, _3>{}) - : make_layout(shape_scale_V, p.dScaleV); - - auto dcScaleQ = const_cast(p.scaleQ + offset_scaleQ); - auto dcScaleK = const_cast(p.scaleK + offset_scaleK); - auto dcScaleV = const_cast(p.scaleV + offset_scaleV); - - Tensor ScaleQ = make_tensor(make_gmem_ptr(dcScaleQ), layout_scaleQ); - Tensor ScaleK = make_tensor(make_gmem_ptr(dcScaleK), layout_scaleK); - Tensor ScaleV = make_tensor(make_gmem_ptr(dcScaleV), layout_scaleV); - - mainloop( - Q(_, _, head_q, l_coord_q), - K(_, _, head, l_coord_kv), - V(_, _, head, l_coord_kv), - tArA, - tA_max, - tA_sum, - blk_qv, - k_block0, - k_blocks, - total_blk, - thr_id, - seq_len, - seq_len_qo, - effective_seq_kv, - idx_b, - tile_row_idx, - rows_of_maxima, - head_q, - s.num_heads_q, - q_offset_sg, - p.cache_seqlens ? p.cache_seqlens[idx_b] : 0, - l_coord_q, - ScaleQ(_, _, head_q, l_coord_q), - ScaleK(_, _, head, l_coord_kv), - ScaleV(_, _, head, l_coord_kv)); - } else { - mainloop( - Q(_, _, head_q, l_coord_q), - K(_, _, head, l_coord_kv), - V(_, _, head, l_coord_kv), - tArA, - tA_max, - tA_sum, - blk_qv, - k_block0, - k_blocks, - total_blk, - thr_id, - seq_len, - seq_len_qo, - effective_seq_kv, - idx_b, - tile_row_idx, - rows_of_maxima, - head_q, - s.num_heads_q, - q_offset_sg, - p.cache_seqlens ? p.cache_seqlens[idx_b] : 0); - } + mainloop( + Q(_, _, head_q, l_coord_q), + K(_, _, head, l_coord_kv), + V(_, _, head, l_coord_kv), + tArA, + tA_max, + tA_sum, + blk_qv, + k_block0, + k_blocks, + total_blk, + thr_id, + seq_len, + seq_len_qo, + effective_seq_kv, + idx_b, + tile_row_idx, + rows_of_maxima, + head_q, + s.num_heads_q, + q_offset_sg, + p.cache_seqlens ? p.cache_seqlens[idx_b] : 0); if constexpr ( !is_empty_v && diff --git a/flash-attn2/torch-ext/flash_attn2/flash_attn_interface.py b/flash-attn2/torch-ext/flash_attn2/flash_attn_interface.py index d2abcd941..0868572aa 100644 --- a/flash-attn2/torch-ext/flash_attn2/flash_attn_interface.py +++ b/flash-attn2/torch-ext/flash_attn2/flash_attn_interface.py @@ -77,31 +77,6 @@ def round_multiple(x, m): return (x + m - 1) // m * m -# Block-scaled MXFP (FP8 e5m2/e4m3, FP4 e2m1) is selected by the Q dtype, matching -# the `is_mxfp` check in flash_api.cpp. -_FP4_DTYPE = getattr(torch, "float4_e2m1fn_x2", None) -_MXFP_DTYPES = tuple( - d - for d in ( - getattr(torch, "float8_e5m2", None), - getattr(torch, "float8_e4m3fn", None), - _FP4_DTYPE, - ) - if d is not None -) - - -def is_mxfp_dtype(dtype): - return dtype in _MXFP_DTYPES - - -def _mxfp_out_shape(q): - """Shape of the MXFP output: fp4_e2m1fn_x2 packs two values per byte, so the - kernel writes twice the head dimension that q carries.""" - head_dim = q.shape[-1] * 2 if q.dtype is _FP4_DTYPE else q.shape[-1] - return q.shape[:-1] + (head_dim,) - - @torch.library.custom_op( add_op_namespace_prefix("_flash_attn_forward"), mutates_args=(), @@ -156,13 +131,7 @@ def _flash_attn_forward_fake( q, k, v = [maybe_contiguous(x) for x in (q, k, v)] batch_size, seqlen_q, num_heads, head_size = q.shape seqlen_k = k.shape[1] - # MXFP accumulates in and writes float32, so it cannot reuse q's metadata. - if is_mxfp_dtype(q.dtype): - out = torch.empty( - _mxfp_out_shape(q), dtype=torch.float32, device=q.device, layout=q.layout - ) - else: - out = torch.empty_like(q) + out = torch.empty_like(q) softmax_lse = torch.empty( (batch_size, num_heads, seqlen_q), dtype=torch.float32, @@ -269,13 +238,7 @@ def _flash_attn_varlen_forward_fake( batch_size = cu_seqlens_q.numel() - 1 total_q, num_heads, _ = q.shape - # MXFP accumulates in and writes float32, so it cannot reuse q's metadata. - if is_mxfp_dtype(q.dtype): - out = torch.empty( - _mxfp_out_shape(q), dtype=torch.float32, device=q.device, layout=q.layout - ) - else: - out = torch.empty_like(q) + out = torch.empty_like(q) softmax_lse = torch.empty( (num_heads, total_q), dtype=torch.float32, device=q.device, layout=q.layout ) diff --git a/flash-attn2/torch-ext/torch_binding.cpp b/flash-attn2/torch-ext/torch_binding.cpp index b43290ca5..76dec209a 100644 --- a/flash-attn2/torch-ext/torch_binding.cpp +++ b/flash-attn2/torch-ext/torch_binding.cpp @@ -36,10 +36,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { "int window_size_right, " "float softcap, " "bool return_softmax, " - "Generator? gen_, " - "Tensor? scale_q_=None, " - "Tensor? scale_k_=None, " - "Tensor? scale_v_=None) -> Tensor[]"); + "Generator? gen_) -> Tensor[]"); #if defined(XPU_KERNEL) ops.impl("fwd", torch::kXPU, &mha_fwd); #endif @@ -65,12 +62,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { "int window_size_right, " "float softcap, " "bool return_softmax, " - "Generator? gen_, " - "Tensor? scale_q_=None, " - "Tensor? scale_k_=None, " - "Tensor? scale_v_=None, " - "Tensor? cu_scale_q_=None, " - "Tensor? cu_scale_kv_=None) -> Tensor[]"); + "Generator? gen_) -> Tensor[]"); #if defined(XPU_KERNEL) ops.impl("varlen_fwd", torch::kXPU, &mha_varlen_fwd); #elif defined(CPU_KERNEL) diff --git a/flash-attn2/torch-ext/torch_binding.h b/flash-attn2/torch-ext/torch_binding.h index c4caa976e..5efb34783 100644 --- a/flash-attn2/torch-ext/torch_binding.h +++ b/flash-attn2/torch-ext/torch_binding.h @@ -19,10 +19,7 @@ mha_fwd( const int64_t window_size_right, const double softcap, const bool return_softmax, - c10::optional gen_, - c10::optional scale_q_ = c10::nullopt, - c10::optional scale_k_ = c10::nullopt, - c10::optional scale_v_ = c10::nullopt); + c10::optional gen_); std::vector mha_varlen_fwd( @@ -47,12 +44,7 @@ mha_varlen_fwd( int64_t window_size_right, const double softcap, const bool return_softmax, - std::optional gen_, - std::optional scale_q_ = c10::nullopt, - std::optional scale_k_ = c10::nullopt, - std::optional scale_v_ = c10::nullopt, - std::optional cu_scale_q_ = c10::nullopt, - std::optional cu_scale_kv_ = c10::nullopt); + std::optional gen_); std::vector mha_bwd(const torch::Tensor &dout, // batch_size x seqlen_q x num_heads, x multiple_of(head_size_og, 8)