From a25c58279520f89167546c4c2fdd7ea6f78fe47c Mon Sep 17 00:00:00 2001 From: Shiyang Chen Date: Mon, 27 Jul 2026 13:42:53 -0700 Subject: [PATCH] W4A8R8 DSA: NVFP4 KV cache format for SM100 sparse decode Adds an NVFP4 KV-cache format to the SM100 sparse MLA decode kernel. Like every existing format it is dispatched by shape, inferred from the KV cache's bytes-per-token (its last dim) together with d_qk/d_v, so the op signature is unchanged and no caller has to be updated: d_qk=576, d_v=512 -> 656 = V32 (fp8), 352 = NVFP4 d_qk=512, d_v=512 -> 584 = MODEL1 nvfp4.fp8rope (352 B/token) [0,256) 512 x e2m1 NoPE, packed 2/byte [256,320) 64 x e4m3 RoPE, unscaled [320,352) 32 x e4m3 NoPE scale factors (one per 16 elements), stored permuted: the scale for element block s lives at byte 8*(s & 3) + (s >> 2), an 8x4 -> 4x8 transpose. A dequant thread needs blocks {4c+q : c=0..7} for fixed q, which in element order is 8 byte loads at stride 4 and after the transpose is one contiguous 8-byte load whose scales also convert two at a time via cvt.rn.f16x2.e4m3x2. Worth 10-12% of decode latency. vs 656 B/token for the existing V3.2 fp8 format: 1.9x more KV capacity at 4.89 effective bits per value. W4 = 4-bit NoPE, R8 = 8-bit RoPE; the kernel dequantizes to bf16 in smem, so Q/P and the MMAs stay bf16 (A16). Design follows the existing self-describing DS-MLA convention: scale factors live inline in the token record, so no per-tensor scale is needed. The RoPE part carries no block scale: e4m3's 4 exponent bits span the RoPE range unaided, whereas e2m1 (2 exponent bits, max 6) genuinely needs one. Kernel: the quantized RoPE and all scale factors are TMA-gathered as one "tail" box into a staging buffer (128 B-aligned per 4-token group), then a dedicated dequant warpgroup converts e2m1/e4m3 -> bf16 in smem. RoPE is dequantized first so it unblocks the QK-RoPE UTCMMA. MMAs stay bf16 with fp32 accumulation, unchanged. SM90 is not supported and rejects this format via the feature-check mechanism. Validated on B200: 4908/4908 cases in tests/test_flash_mla_sparse_decoding.py (4748 pre-existing + 160 new NVFP4 cases covering h_q 64/128, varlen, invalid indices, attention sink, corner cases). Signed-off-by: Shiyang Chen --- csrc/api/common.h | 12 ++ csrc/api/sparse_decode.h | 52 +++-- csrc/params.h | 5 +- csrc/sm100/decode/head64/config.h | 80 +++++++- .../instantiations/v32_nvfp4_fp8rope.cu | 8 + csrc/sm100/decode/head64/kernel.cuh | 191 +++++++++++++++--- csrc/sm100/helpers.h | 29 +++ flash_mla/flash_mla_interface.py | 7 +- setup.py | 9 +- tests/lib.py | 17 +- tests/quant.py | 128 +++++++++++- tests/test_flash_mla_sparse_decoding.py | 47 +++++ 12 files changed, 515 insertions(+), 70 deletions(-) create mode 100644 csrc/sm100/decode/head64/instantiations/v32_nvfp4_fp8rope.cu diff --git a/csrc/api/common.h b/csrc/api/common.h index ea89bfa21..f0546a27d 100644 --- a/csrc/api/common.h +++ b/csrc/api/common.h @@ -111,6 +111,18 @@ inline int int64_stride_to_int(int64_t orig_stride) { } \ } (); +// Like DISPATCH_MODEL_TYPE, but also covers the NVFP4 format (SM100-only kernel). +// Kept separate so that SM90 kernel templates are never instantiated for NVFP4. +#define DISPATCH_MODEL_TYPE_SM100(MODEL_TYPE, CONSTEXPR_NAME, ...) \ +[&] () { \ + if (MODEL_TYPE == ModelType::V32_NVFP4_FP8ROPE) { \ + static constexpr ModelType CONSTEXPR_NAME = ModelType::V32_NVFP4_FP8ROPE; \ + return __VA_ARGS__(); \ + } else { \ + return DISPATCH_MODEL_TYPE(MODEL_TYPE, CONSTEXPR_NAME, __VA_ARGS__); \ + } \ +} (); + // The following code is adapted from https://ykiko.me/en/articles/680412313/, which converts enum values to string names. template constexpr auto get_static_enum_name(){ diff --git a/csrc/api/sparse_decode.h b/csrc/api/sparse_decode.h index d6627cf51..355e2dff4 100644 --- a/csrc/api/sparse_decode.h +++ b/csrc/api/sparse_decode.h @@ -24,7 +24,10 @@ enum class DecodeFeatures : int { ATTN_SINK, TOPK_LENGTH, EXTRA_KVCACHE, - EXTRA_TOPK_LENGTH + EXTRA_TOPK_LENGTH, + + // NVFP4 (e2m1 + per-16 e4m3 scales) NoPE with e4m3 RoPE. SM100-only. + NVFP4_FP8ROPE_KVCACHE_FORMAT }; struct DecodeImplMeta { @@ -82,6 +85,7 @@ class Decode_Sm100_Head64_Impl : public DecodeImplBase { DecodeFeatures::HEAD_DIM_576, DecodeFeatures::V32_KVCACHE_FORMAT, DecodeFeatures::MODEL1_KVCACHE_FORMAT, + DecodeFeatures::NVFP4_FP8ROPE_KVCACHE_FORMAT, DecodeFeatures::ATTN_SINK, DecodeFeatures::TOPK_LENGTH, DecodeFeatures::EXTRA_KVCACHE, @@ -100,7 +104,7 @@ class Decode_Sm100_Head64_Impl : public DecodeImplBase { protected: void run_(const SparseAttnDecodeParams ¶ms, const std::vector &required_features) override { - DISPATCH_MODEL_TYPE(params.model_type, MODEL_TYPE, [&]() { + DISPATCH_MODEL_TYPE_SM100(params.model_type, MODEL_TYPE, [&]() { sm100::decode::head64::run_flash_splitkv_mla_fp8_sparse_kernel(params); }); } @@ -116,6 +120,7 @@ class Decode_Sm100_Head64x2_Impl : public DecodeImplBase { DecodeFeatures::HEAD_DIM_576, DecodeFeatures::V32_KVCACHE_FORMAT, DecodeFeatures::MODEL1_KVCACHE_FORMAT, + DecodeFeatures::NVFP4_FP8ROPE_KVCACHE_FORMAT, DecodeFeatures::ATTN_SINK, DecodeFeatures::TOPK_LENGTH, DecodeFeatures::EXTRA_KVCACHE, @@ -134,7 +139,7 @@ class Decode_Sm100_Head64x2_Impl : public DecodeImplBase { protected: void run_(const SparseAttnDecodeParams ¶ms, const std::vector &required_features) override { - DISPATCH_MODEL_TYPE(params.model_type, MODEL_TYPE, [&]() { + DISPATCH_MODEL_TYPE_SM100(params.model_type, MODEL_TYPE, [&]() { for (int start_head_idx = 0; start_head_idx < 128; start_head_idx += 64) { SparseAttnDecodeParams cur_params = params; cur_params.q += start_head_idx * params.stride_q_h_q; @@ -286,18 +291,34 @@ sparse_attn_decode_interface( // Check shape KU_CHECK_SHAPE(q, b, s_q, h_q, d_qk); + ModelType model_type; { - int bytes_per_token; + // Infer the quantized KV cache format from the KV cache's bytes-per-token + // (i.e. its last dim) + const int bytes_per_token = static_cast(kv.size(3)); if (d_qk == 576 && d_v == 512) { - // V3.2 style - bytes_per_token = 512 + 64*2 + (512/128)*4; + if (bytes_per_token == 512 + 64*2 + (512/128)*4) { + // V3.2 style, 656B/token: 512B e4m3 NoPE | 128B bf16 RoPE | 16B fp32 NoPE SF + model_type = ModelType::V32; + } else if (bytes_per_token == 352) { + // NVFP4 NoPE + FP8 RoPE, 352B/token: 256B e2m1 NoPE | 64B e4m3 RoPE (unscaled) | 32B e4m3 NoPE SF + // Must match KernelTemplate::BYTES_PER_TOKEN + model_type = ModelType::V32_NVFP4_FP8ROPE; + } else { + STD_TORCH_CHECK(false, "Cannot infer the sparse KV cache format: with d_qk == ", d_qk, " and d_v == ", d_v, + ", kv.size(-1) (bytes per token) must be 656 (fp8) or 352 (nvfp4 nope + fp8 rope), but got ", bytes_per_token); + } } else if (d_qk == 512 && d_v == 512) { - // MODEL1 style - bytes_per_token = 448 + 64*2 + (448/64)*1 + 1; + if (bytes_per_token == 448 + 64*2 + (448/64)*1 + 1) { + // MODEL1 style, 584B/token + model_type = ModelType::MODEL1; + } else { + STD_TORCH_CHECK(false, "Cannot infer the sparse KV cache format: with d_qk == ", d_qk, " and d_v == ", d_v, + ", kv.size(-1) (bytes per token) must be 584 (fp8), but got ", bytes_per_token); + } } else { - STD_TORCH_CHECK(false, "Unsupported head sizes for is_fp8_kvcache == True"); + STD_TORCH_CHECK(false, "Unsupported head sizes for sparse decoding: d_qk == ", d_qk, ", d_v == ", d_v); } - KU_CHECK_SHAPE(kv, num_blocks, page_block_size, h_kv, bytes_per_token); KU_CHECK_SHAPE(extra_kv, extra_num_blocks, extra_page_block_size, h_kv, bytes_per_token); STD_TORCH_CHECK(kv.stride(1) == bytes_per_token, "The whole block must be contiguous when is_fp8_cache is True for kv cache"); if (extra_kv.has_value()) { @@ -324,15 +345,6 @@ sparse_attn_decode_interface( } Tensor lse = torch::stable::new_empty(q, {b, s_q, h_q}, ScalarType::Float); - ModelType model_type; - if (d_qk == 576) { - model_type = ModelType::V32; - } else if (d_qk == 512) { - model_type = ModelType::MODEL1; - } else { - STD_TORCH_CHECK(false, "Unsupported d_qk: ", d_qk); - } - std::vector features; if (h_q == 64) { features.push_back(DecodeFeatures::HEAD_64); @@ -352,6 +364,8 @@ sparse_attn_decode_interface( features.push_back(DecodeFeatures::V32_KVCACHE_FORMAT); } else if (model_type == ModelType::MODEL1) { features.push_back(DecodeFeatures::MODEL1_KVCACHE_FORMAT); + } else if (model_type == ModelType::V32_NVFP4_FP8ROPE) { + features.push_back(DecodeFeatures::NVFP4_FP8ROPE_KVCACHE_FORMAT); } else { STD_TORCH_CHECK(false, "Unsupported model type: ", (int)model_type); } diff --git a/csrc/params.h b/csrc/params.h index 4433e8d40..b07851392 100644 --- a/csrc/params.h +++ b/csrc/params.h @@ -4,7 +4,10 @@ enum class ModelType { V32, - MODEL1 + MODEL1, + // V3.2 geometry (d_qk=576) with NVFP4 (e2m1, per-16 e4m3 scales) NoPE and + // e4m3 RoPE. SM100-only. See csrc/sm100/decode/head64/config.h for the layout. + V32_NVFP4_FP8ROPE }; struct __align__(4*8) DecodingSchedMeta { diff --git a/csrc/sm100/decode/head64/config.h b/csrc/sm100/decode/head64/config.h index 401f3acf8..140f1030f 100644 --- a/csrc/sm100/decode/head64/config.h +++ b/csrc/sm100/decode/head64/config.h @@ -30,17 +30,67 @@ enum NamedBarriers : uint32_t { template struct KernelTemplate { -static constexpr int D_Q = MODEL_TYPE == ModelType::V32 ? 576 : 512; +// NVFP4 format: V3.2 geometry, NoPE stored as e2m1 (2 elems/byte) with per-16-element +// e4m3 scale factors; RoPE stored as plain e4m3 with no scale factors at all — e4m3 has +// 4 exponent bits, so it spans the RoPE magnitude range unaided and a block scale would +// only cost bytes. +static constexpr bool IS_NVFP4 = MODEL_TYPE == ModelType::V32_NVFP4_FP8ROPE; +// "V32 geometry": d_qk = 576 = 512 NoPE + 64 RoPE, V = NoPE only +static constexpr bool IS_V32_GEOM = MODEL_TYPE == ModelType::V32 || IS_NVFP4; + +static constexpr int D_Q = IS_V32_GEOM ? 576 : 512; static constexpr int D_K = D_Q; static constexpr int D_V = 512; -static constexpr int D_NOPE = MODEL_TYPE == ModelType::V32 ? 512 : 448; +static constexpr int D_NOPE = IS_V32_GEOM ? 512 : 448; static constexpr int D_ROPE = 64; -static constexpr int QUANT_TILE_SIZE = MODEL_TYPE == ModelType::V32 ? 128 : 64; -static constexpr bool V_HAVE_ROPE = MODEL_TYPE == ModelType::V32 ? false : true; -static constexpr int NUM_SCALES_EACH_TOKEN = MODEL_TYPE == ModelType::V32 ? 4 : 8; // Padding is included -static constexpr int TMA_K_STRIDE = MODEL_TYPE == ModelType::V32 ? D_NOPE+2*D_ROPE+4*(D_NOPE/QUANT_TILE_SIZE) : D_NOPE+2*D_ROPE; // Stride of K's tensormap. This stride must 1) be a factor of the actual stride between tokens 2) large enough to cover the entire KV cache. Since TMA copy's coordinate can only be 32bit signed integers, this number must >= 128, perferrably >= 256. So we set this to 656 for V32 and 576 for MODEL1. Extra padding may be necessary for KV blocks. +static constexpr int QUANT_TILE_SIZE = IS_NVFP4 ? 16 : (MODEL_TYPE == ModelType::V32 ? 128 : 64); +static constexpr bool V_HAVE_ROPE = IS_V32_GEOM ? false : true; +static constexpr int NUM_SCALES_EACH_TOKEN = MODEL_TYPE == ModelType::V32 ? 4 : 8; // Padding is included. Unused for NVFP4 (scales live in the raw tail buffer). + +// Raw (quantized) byte counts per token +static constexpr int NOPE_RAW_BYTES = IS_NVFP4 ? D_NOPE/2 : D_NOPE; // e2m1 packs 2/byte +static constexpr int ROPE_RAW_BYTES = D_ROPE; // e4m3, 1 byte per element +// NVFP4 tail region = [64B e4m3 rope | 32B e4m3 nope SF], already a multiple of 16B. +// The tail is TMA-gathered as one box, so scales arrive together with the data they scale. +static constexpr int NVFP4_NUM_NOPE_SCALES = D_NOPE/16; // 32 +static constexpr int NVFP4_SF_NOPE_OFFSET = ROPE_RAW_BYTES; // within tail +// The 32 SF bytes are NOT stored in element-block order. A dequant thread owns the scale +// groups {4c + q : c = 0..7} for its own q = (idx_in_group/2) in [0, 4), which in element +// order are 8 bytes with stride 4. They are permuted at quantization time so that the scale +// for element block s (covering NoPE dims [16s, 16s+16)) lives at +// NVFP4_SF_NOPE_OFFSET + NVFP4_SF_BYTE(s), NVFP4_SF_BYTE(s) = 8*(s&3) + (s>>2) +// which maps thread q's eight scales onto the contiguous bytes [8q, 8q+8) -> one LDS.64. +// Keep this in lockstep with the quantizer (tests/quant.py and the vLLM cache writer). +static constexpr int nvfp4_sf_byte(int s) { return 8*(s & 3) + (s >> 2); } +// Thread q's scale groups {4c+q} must land on the contiguous bytes 8q+c. That also +// makes it a permutation of [0, 32), since both sides cover that range exactly once. +static constexpr bool nvfp4_sf_perm_ok() { + for (int q = 0; q < 4; ++q) + for (int c = 0; c < 8; ++c) + if (nvfp4_sf_byte(4*c + q) != 8*q + c) return false; + return true; +} +static_assert(nvfp4_sf_perm_ok()); +static constexpr int TAIL_BYTES = IS_NVFP4 ? ROPE_RAW_BYTES + NVFP4_NUM_NOPE_SCALES : 16; // 96; dummy 16 for non-NVFP4 +// SMEM staging stride for one tma_gather4 call (4 tokens' tails, packed): TMA requires the +// unswizzled SMEM destination to be 128B-aligned, so pad between 4-token groups. +static constexpr int TAIL_GROUP_STRIDE = ku::ceil(4*TAIL_BYTES, 128); // 384 +static constexpr int BYTES_PER_TOKEN = + MODEL_TYPE == ModelType::V32 ? D_NOPE + 2*D_ROPE + 4*(D_NOPE/128) : // 656 + MODEL_TYPE == ModelType::MODEL1 ? D_NOPE + 2*D_ROPE + 8 : // 584 (per-block scale suffix layout) + NOPE_RAW_BYTES + TAIL_BYTES; // NVFP4: 352 +static constexpr int TMA_K_STRIDE = MODEL_TYPE == ModelType::V32 ? D_NOPE+2*D_ROPE+4*(D_NOPE/QUANT_TILE_SIZE) : + MODEL_TYPE == ModelType::MODEL1 ? D_NOPE+2*D_ROPE : + BYTES_PER_TOKEN; // Stride of K's tensormap. This stride must 1) be a factor of the actual stride between tokens 2) large enough to cover the entire KV cache. Since TMA copy's coordinate can only be 32bit signed integers, this number must >= 128, perferrably >= 256. So we set this to 656 for V32, 576 for MODEL1, and BYTES_PER_TOKEN for NVFP4. Extra padding may be necessary for KV blocks. static_assert(D_NOPE + D_ROPE == D_Q); static_assert(V_HAVE_ROPE ? (D_NOPE + D_ROPE == D_V) : (D_NOPE == D_V)); +static_assert(!IS_NVFP4 || (TMA_K_STRIDE % 16 == 0 && TMA_K_STRIDE >= 256)); +static_assert(!IS_NVFP4 || BYTES_PER_TOKEN == 352); // Keep in sync with the bytes_per_token literal in csrc/api/sparse_decode.h +// The permuted SF layout is read 8 bytes at a time from the 128B-aligned raw_tail staging +// buffer, at (t/4)*TAIL_GROUP_STRIDE + (t%4)*TAIL_BYTES + NVFP4_SF_NOPE_OFFSET + 8q, so +// every term of that offset must be a multiple of 8. +static_assert(!IS_NVFP4 || (TAIL_GROUP_STRIDE % 8 == 0 && TAIL_BYTES % 8 == 0 && + NVFP4_SF_NOPE_OFFSET % 8 == 0 && NVFP4_NUM_NOPE_SCALES == 32)); static constexpr int B_H = 64; static constexpr int B_TOPK = 64; @@ -50,9 +100,9 @@ static constexpr int NUM_THREADS = 128*3; // 128 exp + 1/32 utcmma + 1/32 raw K static constexpr float MAX_INIT_VAL = -1e30f; // To avoid (-inf) - (-inf) = NaN static constexpr int D_Q_SW128 = 512; -static constexpr int D_Q_SW64 = MODEL_TYPE == ModelType::V32 ? 64 : 0; +static constexpr int D_Q_SW64 = IS_V32_GEOM ? 64 : 0; static_assert(D_Q_SW128 + D_Q_SW64 == D_Q); -static constexpr int K_ROPE_SW = MODEL_TYPE == ModelType::V32 ? 64 : 128; // RoPE part stored in SW64 (for V32) or SW128 (for MODEL1), in bytes +static constexpr int K_ROPE_SW = IS_V32_GEOM ? 64 : 128; // RoPE part stored in SW64 (for V32 geometry) or SW128 (for MODEL1), in bytes template< typename Shape_Q_SW128, typename TMA_Q_SW128, @@ -172,7 +222,13 @@ struct SharedMemoryPlan { array_aligned rope; // RoPE part, dequantized. SW64 in v32 mode, SW128 in MODEL1 mode } dequant[NUM_BUFS]; static_assert(sizeof(dequant) >= sizeof(bf16) * (B_H*D_Q)); // So that Q does not covers raw_nope - array_aligned raw_nope[NUM_BUFS]; // Raw (quantized) NoPE part + array_aligned raw_nope[NUM_BUFS]; // Raw (quantized) NoPE part + // NVFP4 only: raw tail region per token = [rope raw | nope SFs], + // TMA-gathered as one box per 4 tokens; groups of 4 packed tails are padded to + // TAIL_GROUP_STRIDE for TMA alignment. Token t lives at + // (t/4)*TAIL_GROUP_STRIDE + (t%4)*TAIL_BYTES. + // Dummy-sized (16B total, unused) for other formats. + array_aligned raw_tail[IS_NVFP4 ? NUM_BUFS : 1]; } kv; } u; union { @@ -182,13 +238,15 @@ struct SharedMemoryPlan { CUTE_ALIGNAS(16) float rowwise_max_buf[128]; char is_token_valid[NUM_INDEX_BUFS][B_TOPK/8]; int tma_coord[NUM_INDEX_BUFS][B_TOPK]; - e8m0 scales[NUM_INDEX_BUFS][B_TOPK][NUM_SCALES_EACH_TOKEN]; + e8m0 scales[NUM_INDEX_BUFS][B_TOPK][IS_NVFP4 ? 1 : NUM_SCALES_EACH_TOKEN]; // Unused for NVFP4 (scales live in raw_tail) array_aligned tmem_start_addr; transac_bar_t bar_last_store_done; transac_bar_t bar_q_tma, bar_q_utccp; - transac_bar_t bar_rope_ready[NUM_BUFS]; + transac_bar_t bar_rope_ready[NUM_BUFS]; // Non-NVFP4: rope TMA done (init 1, expect_tx). NVFP4: rope dequant done (init 128, arrived by the dequant warpgroup) transac_bar_t bar_nope_ready[NUM_BUFS]; transac_bar_t bar_raw_ready[NUM_BUFS], bar_raw_free[NUM_BUFS]; + // NVFP4 only: raw tail TMA done (init 1, expect_tx) / raw tail consumed by dequant WG (init 128) + transac_bar_t bar_rawtail_ready[NUM_BUFS], bar_rawtail_free[NUM_BUFS]; transac_bar_t bar_valid_coord_scale_ready[NUM_INDEX_BUFS], bar_valid_coord_scale_free[NUM_INDEX_BUFS]; transac_bar_t bar_qk_done[NUM_BUFS], bar_so_ready[NUM_BUFS], bar_sv_done[NUM_BUFS]; }; diff --git a/csrc/sm100/decode/head64/instantiations/v32_nvfp4_fp8rope.cu b/csrc/sm100/decode/head64/instantiations/v32_nvfp4_fp8rope.cu new file mode 100644 index 000000000..68f6e451b --- /dev/null +++ b/csrc/sm100/decode/head64/instantiations/v32_nvfp4_fp8rope.cu @@ -0,0 +1,8 @@ +#include "../kernel.cuh" + +namespace sm100::decode::head64 { + +template +void run_flash_splitkv_mla_fp8_sparse_kernel(const SparseAttnDecodeParams ¶ms); + +} diff --git a/csrc/sm100/decode/head64/kernel.cuh b/csrc/sm100/decode/head64/kernel.cuh index 7c46921cc..786a5983e 100644 --- a/csrc/sm100/decode/head64/kernel.cuh +++ b/csrc/sm100/decode/head64/kernel.cuh @@ -46,17 +46,23 @@ KernelTemplate plan.bar_q_tma.init(1); plan.bar_q_utccp.init(1); for (int i = 0; i < NUM_BUFS; ++i) { - plan.bar_rope_ready[i].init(1); - plan.bar_nope_ready[i].init(128); + // Non-NVFP4: bar_rope_ready is completed by the RoPE TMA (expect_tx). + // NVFP4: RoPE is dequantized by the 128-thread dequant warpgroup, which arrives here. + plan.bar_rope_ready[i].init(IS_NVFP4 ? 128 : 1); + plan.bar_nope_ready[i].init(128); plan.bar_raw_ready[i].init(1); plan.bar_raw_free[i].init(128); + plan.bar_rawtail_ready[i].init(1); + plan.bar_rawtail_free[i].init(128); plan.bar_qk_done[i].init(1); plan.bar_so_ready[i].init(128); plan.bar_sv_done[i].init(1); } for (int i = 0; i < NUM_INDEX_BUFS; ++i) { plan.bar_valid_coord_scale_ready[i].init(32); - plan.bar_valid_coord_scale_free[i].init(128+128+1+1); + // Arrivals: exp warpgroup (128) + dequant warpgroup (128, non-NVFP4 only: it + // reads scales from this pipeline) + raw NoPE producer (1) + rope producer (1) + plan.bar_valid_coord_scale_free[i].init(IS_NVFP4 ? 128+1+1 : 128+128+1+1); } cutlass::arch::fence_barrier_init(); } @@ -529,8 +535,8 @@ KernelTemplate // Mainloop CUTE_NO_UNROLL for (int block_idx = args.start_block_idx; block_idx < args.end_block_idx; ++block_idx) { - if constexpr (MODEL_TYPE == ModelType::V32) { - // V3.2: RoPE behaves like an extra block with size 64, so we can do RoPE first + if constexpr (IS_V32_GEOM) { + // V3.2 (and NVFP4): RoPE behaves like an extra block with size 64, so we can do RoPE first // QK RoPE plan.bar_rope_ready[rs.buf_idx].wait(rs.bar_phase); ku::tcgen05_after_thread_sync(); @@ -601,27 +607,33 @@ KernelTemplate ku::tma_gather4( block_idx >= args.num_orig_kv_blocks ? &tma_params.tensor_map_extra_kv_nope : &tma_params.tensor_map_kv_nope, plan.bar_raw_ready[rs.buf_idx], - plan.u.kv.raw_nope[rs.buf_idx].data() + D_NOPE*row, + plan.u.kv.raw_nope[rs.buf_idx].data() + NOPE_RAW_BYTES*row, 0, cur_indices, (int64_t)TMA::CacheHintSm90::EVICT_LAST ); cur_indices = nxt_cur_indices; } - plan.bar_raw_ready[rs.buf_idx].arrive_and_expect_tx(B_TOPK*D_NOPE*sizeof(e4m3)); + plan.bar_raw_ready[rs.buf_idx].arrive_and_expect_tx(B_TOPK*NOPE_RAW_BYTES); plan.bar_valid_coord_scale_free[rs.index_buf_idx].arrive(); rs.update(); } }); } else if (warp_idx == 6 && elect_one_sync()) { - // KV RoPE retrieval warp + // KV RoPE retrieval warp. + // Non-NVFP4: gathers the bf16 RoPE part directly into the MMA-consumed buffer. + // NVFP4: gathers the raw tail region (quantized rope + all scale factors) into + // the raw_tail staging buffer; the dequant warpgroup turns it into bf16 rope. run_main_loop([&](const MainLoopArgs &args) { plan.bar_q_utccp.wait(args.bar_phase_batch_rel); plan.bar_last_store_done.wait(args.bar_phase_batch_rel); CUTE_NO_UNROLL for (int block_idx = args.start_block_idx; block_idx < args.end_block_idx; ++block_idx) { plan.bar_valid_coord_scale_ready[rs.index_buf_idx].wait(rs.index_bar_phase); - if constexpr (MODEL_TYPE == ModelType::V32) { + if constexpr (IS_NVFP4) { + // raw_tail[buf] is consumed by the dequant warpgroup + plan.bar_rawtail_free[rs.buf_idx].wait(rs.bar_phase^1); + } else if constexpr (MODEL_TYPE == ModelType::V32) { plan.bar_qk_done[rs.buf_idx].wait(rs.bar_phase^1); } else { plan.bar_sv_done[rs.buf_idx].wait(rs.bar_phase^1); @@ -632,20 +644,35 @@ KernelTemplate for (int row = 0; row < B_TOPK; row += 4) { if (row+4 < B_TOPK) nxt_cur_indices = *(int4*)(plan.tma_coord[rs.index_buf_idx] + row + 4); - CUTE_UNROLL - for (int t = 0; t < D_ROPE/(K_ROPE_SW/2); ++t) { + if constexpr (IS_NVFP4) { ku::tma_gather4( block_idx >= args.num_orig_kv_blocks ? &tma_params.tensor_map_extra_kv_rope : &tma_params.tensor_map_kv_rope, - plan.bar_rope_ready[rs.buf_idx], - plan.u.kv.dequant[rs.buf_idx].rope.data() + (K_ROPE_SW/2)*row + t*B_TOPK*(K_ROPE_SW/2), - t*(K_ROPE_SW/2), + plan.bar_rawtail_ready[rs.buf_idx], + plan.u.kv.raw_tail[rs.buf_idx].data() + TAIL_GROUP_STRIDE*(row/4), + 0, cur_indices, (int64_t)TMA::CacheHintSm90::EVICT_LAST ); + } else { + CUTE_UNROLL + for (int t = 0; t < D_ROPE/(K_ROPE_SW/2); ++t) { + ku::tma_gather4( + block_idx >= args.num_orig_kv_blocks ? &tma_params.tensor_map_extra_kv_rope : &tma_params.tensor_map_kv_rope, + plan.bar_rope_ready[rs.buf_idx], + plan.u.kv.dequant[rs.buf_idx].rope.data() + (K_ROPE_SW/2)*row + t*B_TOPK*(K_ROPE_SW/2), + t*(K_ROPE_SW/2), + cur_indices, + (int64_t)TMA::CacheHintSm90::EVICT_LAST + ); + } } cur_indices = nxt_cur_indices; } - plan.bar_rope_ready[rs.buf_idx].arrive_and_expect_tx(B_TOPK*D_ROPE*sizeof(bf16)); + if constexpr (IS_NVFP4) { + plan.bar_rawtail_ready[rs.buf_idx].arrive_and_expect_tx(B_TOPK*TAIL_BYTES); + } else { + plan.bar_rope_ready[rs.buf_idx].arrive_and_expect_tx(B_TOPK*D_ROPE*sizeof(bf16)); + } plan.bar_valid_coord_scale_free[rs.index_buf_idx].arrive(); rs.update(); } @@ -654,7 +681,10 @@ KernelTemplate // Indices transformation warp // Responsible for generating: TMA coordinates, scale factors, and valid masks static_assert(B_TOPK == 64); - static constexpr int tma_coords_step_per_token = MODEL_TYPE == ModelType::V32 ? 656/TMA_K_STRIDE : 576/TMA_K_STRIDE; + static constexpr int tma_coords_step_per_token = + MODEL_TYPE == ModelType::V32 ? 656/TMA_K_STRIDE : + MODEL_TYPE == ModelType::MODEL1 ? 576/TMA_K_STRIDE : + BYTES_PER_TOKEN/TMA_K_STRIDE; int tma_coords_step_per_block = params.stride_kv_block / TMA_K_STRIDE; // must < 2G since k_batch_stride < 1T and TMA_K_STRIDE > 512 int tma_coords_step_per_extra_block = params.stride_extra_kv_block / TMA_K_STRIDE; uint8_t* k_scales_ptr = @@ -691,7 +721,7 @@ KernelTemplate plan.bar_valid_coord_scale_free[rs.index_buf_idx].wait(rs.index_bar_phase^1); int tma_coords[2]; - e8m0 scales[2*NUM_SCALES_EACH_TOKEN]; + [[maybe_unused]] e8m0 scales[2*NUM_SCALES_EACH_TOKEN]; char valid_mask = 0; CUTE_UNROLL for (int i = 0; i < 2; ++i) { @@ -701,7 +731,9 @@ KernelTemplate bool is_token_valid = my_indices[i] != -1 && (abs_pos+i < (IS_EXTRA_BLOCK?args.extra_topk_length:args.topk_length)); valid_mask |= is_token_valid << i; tma_coords[i] = is_token_valid ? block_idx*cur_tma_coords_step_per_block + idx_in_block*tma_coords_step_per_token : -1; // If the token is invalid because it topk position exceeds topk_length, we must manually fill tma_coords with -1 to avoid copying-in NaN. - if constexpr (MODEL_TYPE == ModelType::V32) { + if constexpr (IS_NVFP4) { + // Scales are part of the raw tail region and arrive via TMA; nothing to load here. + } else if constexpr (MODEL_TYPE == ModelType::V32) { int64_t offset = is_token_valid ? block_idx*cur_k_block_stride + idx_in_block*cur_k_row_stride : 0; float4 cur_scale_fp32 = __ldg((float4*)(cur_k_scales_ptr + offset)); e8m0 res[4]; @@ -718,7 +750,9 @@ KernelTemplate valid_mask <<= lane_idx%4*2; valid_mask |= __shfl_xor_sync(0xFFFFFFFF, valid_mask, 0x1); valid_mask |= __shfl_xor_sync(0xFFFFFFFF, valid_mask, 0x2); - if constexpr (MODEL_TYPE == ModelType::V32) { + if constexpr (IS_NVFP4) { + // No scale staging for NVFP4 + } else if constexpr (MODEL_TYPE == ModelType::V32) { *(uint64_t*)(plan.scales[rs.index_buf_idx] + lane_idx*2) = *(uint64_t*)scales; } else { *(__int128_t*)(plan.scales[rs.index_buf_idx] + lane_idx*2) = *(__int128_t*)scales; @@ -754,8 +788,9 @@ KernelTemplate Tensor nope0 = make_tensor(make_smem_ptr(plan.u.kv.dequant[0].nope.data()), SmemLayoutKTiles_SW128{}); bf16* nope0_base = &nope0(group_idx, idx_in_group*8); bf16* nope1_base = nope0_base + (plan.u.kv.dequant[1].nope.data() - plan.u.kv.dequant[0].nope.data()); - e4m3* raw_nope0_base = plan.u.kv.raw_nope[rs.buf_idx].data() + group_idx*D_NOPE + idx_in_group*8; - e4m3* raw_nope1_base = raw_nope0_base + B_H*D_NOPE; + // Each thread-step covers 8 elements: 8 e4m3 (8B) for fp8 formats, 8 e2m1 (4B) for NVFP4 + const uint8_t* raw_nope0_base = plan.u.kv.raw_nope[0].data() + group_idx*NOPE_RAW_BYTES + idx_in_group*(IS_NVFP4 ? 4 : 8); + const uint8_t* raw_nope1_base = raw_nope0_base + B_TOPK*NOPE_RAW_BYTES; run_main_loop([&](const MainLoopArgs &args) { // plan.bar_last_store_done.wait(args.bar_phase_batch_rel); // No need to wait since the raw nope producer must wait @@ -763,19 +798,102 @@ KernelTemplate CUTE_NO_UNROLL for (int block_idx = args.start_block_idx; block_idx < args.end_block_idx; ++block_idx) { + if constexpr (IS_NVFP4) { + // NVFP4: dequantize RoPE first (it is small and unblocks the QK-RoPE MMA), + // then NoPE. Scale factors are read from the raw tail staging buffer, which + // arrives via TMA together with the quantized rope data. + plan.bar_rawtail_ready[rs.buf_idx].wait(rs.bar_phase); + // sv_done of the previous round in this buf implies its qk_done, so both + // dequant[].rope and dequant[].nope are free to overwrite. + plan.bar_sv_done[rs.buf_idx].wait(rs.bar_phase^1); + const uint8_t* tail_base = plan.u.kv.raw_tail[rs.buf_idx].data(); + + // --- RoPE dequant: 2 threads per token, 32 dims each. + // Plain e4m3 with no scale factor (see config.h). --- + { + int token = idx_in_warpgroup / 2, half = idx_in_warpgroup % 2; + const uint8_t* tail = tail_base + (token/4)*TAIL_GROUP_STRIDE + (token%4)*TAIL_BYTES; + Tensor sRope = make_tensor(make_smem_ptr(plan.u.kv.dequant[rs.buf_idx].rope.data()), SmemLayoutKTiles_DualGemm_SW64<1>{}); + CUTE_UNROLL + for (int j = 0; j < 4; ++j) { + nv_bfloat162 out[4]; + ku::nve4m3x2 d[4]; + *(uint64_t*)d = *(const uint64_t*)(tail + half*32 + j*8); + CUTE_UNROLL + for (int i = 0; i < 4; ++i) + out[i] = fp8x2_to_bf16x2(d[i]); + *(__int128_t*)&sRope(half*B_TOPK + token, j*8) = *(__int128_t*)out; + } + cutlass::arch::fence_view_async_shared(); + plan.bar_rope_ready[rs.buf_idx].arrive(); + } + + // --- NoPE dequant --- + plan.bar_raw_ready[rs.buf_idx].wait(rs.bar_phase); + { + uint32_t cur_nope_base_uint_addr = cute::cast_smem_ptr_to_uint(rs.buf_idx == 0 ? nope0_base : nope1_base); + const uint8_t* raw_nope_base = rs.buf_idx == 0 ? raw_nope0_base : raw_nope1_base; + auto st_128b = [&](int local_row_idx, int local_col_idx, __int128_t &data) { + asm volatile ("st.weak.shared::cta.b128 [%0], %1;\n" + : + : "r"(cur_nope_base_uint_addr + 2*(local_row_idx*NUM_GROUPS*64 + local_col_idx*B_TOPK*64)), "q"(data) // 2 for sizeof(bf16) + ); + }; + CUTE_UNROLL + for (int local_row_idx = 0; local_row_idx < ROWS_PER_GROUP; ++local_row_idx) { + int row_idx = local_row_idx*NUM_GROUPS + group_idx; + // This thread covers elements [c*64 + idx_in_group*8, +8) -> scale group c*4 + q, + // where q = idx_in_group/2 and c = local_col_idx. The tail's 32 SF bytes are + // stored permuted (scale group s at byte 8*(s&3) + (s>>2), see + // NVFP4_SF_NOPE_OFFSET in config.h) so this thread's 8 scale groups are the + // contiguous bytes [8q, 8q+8) and can be fetched with one LDS.64. + const uint8_t* sf_ptr = tail_base + (row_idx/4)*TAIL_GROUP_STRIDE + (row_idx%4)*TAIL_BYTES + NVFP4_SF_NOPE_OFFSET + (idx_in_group/2)*COLS_PER_GROUP; + uint2 sf_raw = *(const uint2*)sf_ptr; // one LDS.64, byte c = scale for column c + // Convert them two at a time: one cvt.rn.f16x2.e4m3x2 handles a 16-bit + // slice, i.e. two consecutive columns, and the halves are then free + // operand modifiers (.H0_H0/.H1_H1) on the HMUL2 below. + nv_bfloat162 sf_pairs[COLS_PER_GROUP/2]; + CUTE_UNROLL + for (int k = 0; k < COLS_PER_GROUP/2; ++k) { + uint32_t w = (k < 2 ? sf_raw.x : sf_raw.y) >> (16*(k&1)); + __nv_fp8x2_e4m3 d; + d.__x = (__nv_fp8x2_storage_t)w; + sf_pairs[k] = fp8x2_to_bf16x2(d); + } + uint32_t cur_raw = *(const uint32_t*)(raw_nope_base + local_row_idx*NUM_GROUPS*NOPE_RAW_BYTES); + CUTE_UNROLL + for (int local_col_idx = 0; local_col_idx < COLS_PER_GROUP; ++local_col_idx) { + uint32_t raw = cur_raw; + if (local_col_idx+1 < COLS_PER_GROUP) + cur_raw = *(const uint32_t*)(raw_nope_base + local_row_idx*NUM_GROUPS*NOPE_RAW_BYTES + (local_col_idx+1)*(GROUP_SIZE*4)); + nv_bfloat16 sf = (local_col_idx & 1) ? sf_pairs[local_col_idx/2].y : sf_pairs[local_col_idx/2].x; + nv_bfloat162 out[4]; + fp4x8_to_bf16x8_with_scale(raw, sf, out); + st_128b(local_row_idx, local_col_idx, *(__int128_t*)out); + } + } + } + cutlass::arch::fence_view_async_shared(); + plan.bar_nope_ready[rs.buf_idx].arrive(); + plan.bar_raw_free[rs.buf_idx].arrive(); + plan.bar_rawtail_free[rs.buf_idx].arrive(); + rs.update(); + continue; + } + plan.bar_valid_coord_scale_ready[rs.index_buf_idx].wait(rs.index_bar_phase); plan.bar_raw_ready[rs.buf_idx].wait(rs.bar_phase); plan.bar_sv_done[rs.buf_idx].wait(rs.bar_phase^1); uint32_t cur_nope_base_uint_addr = cute::cast_smem_ptr_to_uint(rs.buf_idx == 0 ? nope0_base : nope1_base); - e4m3* raw_nope_base = rs.buf_idx == 0 ? raw_nope0_base : raw_nope1_base; + const uint8_t* raw_nope_base = rs.buf_idx == 0 ? raw_nope0_base : raw_nope1_base; auto st_128b = [&](int local_row_idx, int local_col_idx, __int128_t &data) { - asm volatile ("st.weak.shared::cta.b128 [%0], %1;\n" - : + asm volatile ("st.weak.shared::cta.b128 [%0], %1;\n" + : : "r"(cur_nope_base_uint_addr + 2*(local_row_idx*NUM_GROUPS*64 + local_col_idx*B_TOPK*64)), "q"(data) // 2 for sizeof(bf16) ); // We have this `asm volatile` here, otherwise the compiler generates ST.E instead of STS }; auto get_raw_fp8 = [&](int local_row_idx, int local_col_idx) -> uint64_t { - return *(uint64_t*)(raw_nope_base + local_row_idx*NUM_GROUPS*D_NOPE + local_col_idx*(GROUP_SIZE*8)); + return *(const uint64_t*)(raw_nope_base + local_row_idx*NUM_GROUPS*NOPE_RAW_BYTES + local_col_idx*(GROUP_SIZE*8)); }; // The following code suffers from a 2-way bank conflict when reading from SMEM. if constexpr (MODEL_TYPE == ModelType::V32) { @@ -863,9 +981,11 @@ void KernelTemplate::run(const SparseAttnDecodeParams ¶ms) { KU_ASSERT(params.d_qk == D_Q); KU_ASSERT(params.d_v == D_V); if constexpr (MODEL_TYPE == ModelType::MODEL1) { - constexpr int BYTES_PER_TOKEN = D_NOPE + 2*D_ROPE + 8; KU_ASSERT(params.stride_kv_row == BYTES_PER_TOKEN, "Each page block in KV cache must be contiguous for head64 sparse fp8 decoding attention in MODEL1"); // Each block must be contiguous } + if constexpr (IS_NVFP4) { + KU_ASSERT(params.stride_kv_row == BYTES_PER_TOKEN, "Each page block in KV cache must be contiguous (%d bytes per token) for head64 sparse NVFP4 decoding attention", BYTES_PER_TOKEN); + } auto shape_Q_SW128 = make_shape(B_H, D_Q, params.s_q, params.b); auto tma_Q_SW128 = cute::make_tma_copy( @@ -907,19 +1027,30 @@ void KernelTemplate::run(const SparseAttnDecodeParams ¶ms) { } auto get_nope_rope_tensormap = [&](bool is_extra, void* k_ptr, int num_blocks, int64_t k_batch_stride) -> std::pair { - static_assert(D_NOPE%8 == 0); + static_assert(D_NOPE%8 == 0 && NOPE_RAW_BYTES%8 == 0 && TAIL_BYTES%8 == 0); KU_ASSERT((int64_t)k_ptr % 16 == 0, "The base address of %sk_ptr (%p) must be 16B aligned for sparse fp8 attention on sm100f", is_extra?"extra_":"", k_ptr); KU_ASSERT(k_batch_stride % TMA_K_STRIDE == 0, "%sk_cache.stride(0) (%ld) must be a multiple of %d. Padding might be necessary", is_extra?"extra_":"", k_batch_stride, TMA_K_STRIDE); CUtensorMap tensor_map_kv_nope = ku::make_tensor_map( - {D_NOPE/8, (uint64_t)num_blocks * (k_batch_stride/TMA_K_STRIDE)}, + {NOPE_RAW_BYTES/8, (uint64_t)num_blocks * (k_batch_stride/TMA_K_STRIDE)}, {TMA_K_STRIDE}, - {D_NOPE/8, 1}, + {NOPE_RAW_BYTES/8, 1}, k_ptr, CUtensorMapDataType::CU_TENSOR_MAP_DATA_TYPE_INT64, CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_NONE, CUtensorMapL2promotion::CU_TENSOR_MAP_L2_PROMOTION_L2_128B ); // NOTE We combine 8 float8 into 1 int64 since boxdim cannot > 256 - CUtensorMap tensor_map_kv_rope = ku::make_tensor_map( + // For NVFP4, the second tensormap covers the raw "tail" region of each token + // ([rope raw | nope SFs], TAIL_BYTES), copied unswizzled as int64s. + // For V32/MODEL1 it covers the bf16 rope part, swizzled for direct MMA consumption. + CUtensorMap tensor_map_kv_rope = IS_NVFP4 ? ku::make_tensor_map( + {TAIL_BYTES/8, (uint64_t)num_blocks * (k_batch_stride/TMA_K_STRIDE)}, + {TMA_K_STRIDE}, + {TAIL_BYTES/8, 1}, + (uint8_t*)k_ptr + NOPE_RAW_BYTES, + CUtensorMapDataType::CU_TENSOR_MAP_DATA_TYPE_INT64, + CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_NONE, + CUtensorMapL2promotion::CU_TENSOR_MAP_L2_PROMOTION_L2_128B + ) : ku::make_tensor_map( {D_ROPE, (uint64_t)num_blocks * (k_batch_stride/TMA_K_STRIDE)}, {TMA_K_STRIDE}, {K_ROPE_SW/2, 1}, diff --git a/csrc/sm100/helpers.h b/csrc/sm100/helpers.h index a5666257d..cc695dd35 100644 --- a/csrc/sm100/helpers.h +++ b/csrc/sm100/helpers.h @@ -3,6 +3,7 @@ #include #include #include +#include #include "defines.h" @@ -32,4 +33,32 @@ nv_bfloat162 fp8x2_to_bf16x2_with_scale(__nv_fp8x2_e4m3 data, nv_bfloat16 scale) }; } +// Convert 2x fp8_e4m3 to 2x bf16 (no scaling). Exact: e4m3 values are a subset of bf16. +CUTE_DEVICE +nv_bfloat162 fp8x2_to_bf16x2(__nv_fp8x2_e4m3 data) { + return __float22bfloat162_rn((float2)data); +} + +// Convert 1x fp8_e4m3 (a scale factor) to bf16. Exact: e4m3 values are a subset of bf16. +CUTE_DEVICE +nv_bfloat16 fp8_e4m3_to_bf16(uint8_t data) { + __half_raw h = __nv_cvt_fp8_to_halfraw((__nv_fp8_storage_t)data, __NV_E4M3); + return __float2bfloat16_rn(__half2float(*(__half*)&h)); +} + +// Convert 8x fp4_e2m1 (packed in a uint32, low nibble = even element) to 8x bf16 with scaling. +// The e2m1*scale product is exactly representable in bf16 (<= 5 mantissa bits), so the +// bf16 multiply below is exact. +CUTE_DEVICE +void fp4x8_to_bf16x8_with_scale(uint32_t data, nv_bfloat16 scale, nv_bfloat162 out[4]) { + nv_bfloat162 scale2 = {scale, scale}; + CUTE_UNROLL + for (int i = 0; i < 4; ++i) { + // Native cvt.rn.f16x2.e2m1x2 on sm_100f + __half2_raw h2 = __nv_cvt_fp4x2_to_halfraw2((__nv_fp4x2_storage_t)(data >> (8*i)), __NV_E2M1); + float2 f2 = __half22float2(*(__half2*)&h2); + out[i] = __hmul2(__float22bfloat162_rn(f2), scale2); + } +} + } diff --git a/flash_mla/flash_mla_interface.py b/flash_mla/flash_mla_interface.py index a9a56eaa3..55643f196 100644 --- a/flash_mla/flash_mla_interface.py +++ b/flash_mla/flash_mla_interface.py @@ -101,6 +101,11 @@ def flash_mla_with_kvcache( - First 512 bytes: The "quantized NoPE" part, containing 512 float8_e4m3 values. - Next 16 bytes: Scale factors, containing 4 float32 values. The first float32 is the scale for the first 128 float8_e4m3 values, the second for the next 128, and so on. - Last 128 bytes: The "RoPE" part, containing 64 bfloat16 values. This part is not quantized for accuracy. + The quantized KV cache format is inferred from `head_dim` and the number of bytes per token, i.e. `k_cache.shape[-1]`. + Besides the 656-byte layout above, head_dim == 576 also accepts a 352-byte NVFP4 layout (SM100 only): + - First 256 bytes: The "quantized NoPE" part, containing 512 float4_e2m1 values. + - Next 64 bytes: The "RoPE" part, containing 64 float8_e4m3 values. This part is not scaled. + - Last 32 bytes: Scale factors for the NoPE part, containing 32 float8_e4m3 values, one per 16 float4_e2m1 values. Return: out: (batch_size, seq_len_q, num_heads_q, head_dim_v). @@ -121,7 +126,7 @@ def flash_mla_with_kvcache( # Sanity check. We only perform sanity check during the first invocation to save CPU time. if indices_in_kvcache is not None: assert not causal, "causal must be False when indices_in_kvcache is not None (i.e. sparse attention is enabled)" - + # Initialize the tile scheduler metadata during the first invocation. sched_meta.have_initialized = True sched_meta.config = FlashMLASchedMeta.Config( diff --git a/setup.py b/setup.py index d35400b0a..1c4915da7 100644 --- a/setup.py +++ b/setup.py @@ -111,6 +111,7 @@ def get_nvcc_thread_args(): # sm100 sparse decode "csrc/sm100/decode/head64/instantiations/v32.cu", "csrc/sm100/decode/head64/instantiations/model1.cu", + "csrc/sm100/decode/head64/instantiations/v32_nvfp4_fp8rope.cu", "csrc/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512.cu", ], extra_compile_args={ @@ -139,7 +140,13 @@ def get_nvcc_thread_args(): Path(this_dir) / "csrc" / "sm90", Path(this_dir) / "csrc" / "cutlass" / "include", Path(this_dir) / "csrc" / "cutlass" / "tools" / "util" / "include", - ], + ] + ( + # CUDA 13 relocated the CCCL headers (cuda/std/...) under + # include/cccl; nvcc injects this path itself but the host C++ + # compiler does not get it. + [Path(CUDA_HOME) / "include" / "cccl"] + if (Path(CUDA_HOME) / "include" / "cccl").exists() else [] + ), # Build against CPython's Limited API (abi3) so one wheel works across # multiple CPython versions, which is possible now that pybind11 is gone py_limited_api=True, diff --git a/tests/lib.py b/tests/lib.py index 139e130d6..d565a5127 100644 --- a/tests/lib.py +++ b/tests/lib.py @@ -41,6 +41,7 @@ class TestParam: have_attn_sink: bool = False have_topk_length: bool = False decode: Optional[ExtraTestParamForDecode] = None + kv_format: str = "fp8" # "fp8" | "nvfp4.fp8rope" (decode only) @dataclasses.dataclass class RawTestParamForDecode: @@ -70,6 +71,7 @@ class RawTestParamForDecode: check_correctness: bool = True num_runs: int = 10 seed: int = -1 + kv_format: str = "fp8" # "fp8" | "nvfp4.fp8rope" def to_test_param(self) -> TestParam: return TestParam( @@ -83,7 +85,8 @@ def to_test_param(self) -> TestParam: self.b, self.is_varlen, self.have_zero_seqlen_k, self.extra_s_k, self.extra_topk, self.block_size, self.extra_block_size, self.have_extra_topk_length - ) + ), + kv_format = self.kv_format ) @dataclasses.dataclass @@ -181,7 +184,10 @@ def quant_and_dequant_(self): Besides, the quantization error may be too large to be distinguished from wrong kernels, so we de-quantize kvcache here to mitigate quantization error """ fp8_kvcache_layout = None - if self.t.d_qk == 576: + if self.t.kv_format == "nvfp4.fp8rope": + assert self.t.d_qk == 576 + fp8_kvcache_layout = quant.FP8KVCacheLayout.NVFP4_FP8Rope + elif self.t.d_qk == 576: fp8_kvcache_layout = quant.FP8KVCacheLayout.V32_FP8Sparse elif self.t.d_qk == 512: assert self.abs_indices is not None @@ -330,7 +336,7 @@ def run_flash_mla_decode(p: TestParam, t: TestcaseForDecode, tile_scheduler_meta t.extra_kv_scope.get_kvcache_for_flash_mla() if t.extra_kv_scope is not None else None, t.extra_kv_scope.indices_in_kvcache if t.extra_kv_scope is not None else None, t.kv_scope.topk_length, - t.extra_kv_scope.topk_length if t.extra_kv_scope is not None and t.extra_kv_scope.topk_length is not None else None + t.extra_kv_scope.topk_length if t.extra_kv_scope is not None and t.extra_kv_scope.topk_length is not None else None, ) @@ -390,7 +396,10 @@ def get_num_retrieved_tokens(kv_scope: KVScope) -> int: num_retrieved_tokens = get_num_retrieved_tokens(t.kv_scope) + (get_num_retrieved_tokens(t.extra_kv_scope) if t.extra_kv_scope is not None else 0) compute_flop = 2 * p.h_q * num_attended_tokens * (p.d_qk + p.d_v) - kv_token_size = 656 if p.d_qk == 576 else 576 # Assume FP8 KV Cache + kv_token_size = { + "fp8": 656 if p.d_qk == 576 else 576, + "nvfp4.fp8rope": 352, + }[p.kv_format] mem_vol = sum([ 2 * b * p.s_q * p.h_q * p.d_qk, # Q num_retrieved_tokens * kv_token_size, # K diff --git a/tests/quant.py b/tests/quant.py index b92b53997..ccc6988f6 100644 --- a/tests/quant.py +++ b/tests/quant.py @@ -6,17 +6,102 @@ class FP8KVCacheLayout(enum.Enum): V32_FP8Sparse = 1 MODEL1_FP8Sparse = 2 + NVFP4_FP8Rope = 3 # NVFP4 (e2m1, per-16 e4m3 SF) NoPE + FP8 (e4m3, unscaled) RoPE, 352B/token def get_meta(self) -> Tuple[int, int, int, int, int]: # Return: (d, d_nope, d_rope, tile_size, num_tiles) return { FP8KVCacheLayout.V32_FP8Sparse: (576, 512, 64, 128, 4), - FP8KVCacheLayout.MODEL1_FP8Sparse: (512, 448, 64, 64, 7) + FP8KVCacheLayout.MODEL1_FP8Sparse: (512, 448, 64, 64, 7), + FP8KVCacheLayout.NVFP4_FP8Rope: (576, 512, 64, 16, 32), + }[self] + + def is_nvfp4(self) -> bool: + return self is FP8KVCacheLayout.NVFP4_FP8Rope + + def bytes_per_token(self) -> int: + return { + FP8KVCacheLayout.V32_FP8Sparse: 656, + FP8KVCacheLayout.MODEL1_FP8Sparse: 584, + FP8KVCacheLayout.NVFP4_FP8Rope: 352, }[self] def _cast_scale_inv_to_ue8m0(scales_inv: torch.Tensor, out_dtype = torch.float32) -> torch.Tensor: return torch.pow(2, torch.clamp_min(scales_inv, 1e-4).log2().ceil()).to(out_dtype) +# The 8 non-negative values representable in fp4 e2m1 +_E2M1_VALUES = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0] +# Midpoints between consecutive e2m1 values, and the round-to-nearest-EVEN winner at each midpoint +_E2M1_MIDPOINTS = [0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0] +_E2M1_TIE_CODES = [0, 2, 2, 4, 4, 6, 6] + +def _cast_to_e2m1_codes(x: torch.Tensor) -> torch.Tensor: + """Round-to-nearest-even quantization to e2m1. Returns uint8 nibble codes (sign<<3 | mag).""" + xf = x.float() + xa = xf.abs().clamp(max=6.0) + mids = torch.tensor(_E2M1_MIDPOINTS, device=x.device, dtype=torch.float32) + codes = torch.bucketize(xa, mids, right=True).to(torch.uint8) # x == midpoint goes UP here... + for mid, tie_code in zip(_E2M1_MIDPOINTS, _E2M1_TIE_CODES): # ...and is fixed to the even value here + codes = torch.where(xa == mid, torch.tensor(tie_code, dtype=torch.uint8, device=x.device), codes) + codes = codes | (xf < 0).to(torch.uint8) * 8 + return codes + +def _e2m1_codes_to_float(codes: torch.Tensor) -> torch.Tensor: + """Decode uint8 nibble codes (low 4 bits used) to float32 values.""" + table = torch.tensor(_E2M1_VALUES + [-v for v in _E2M1_VALUES], device=codes.device, dtype=torch.float32) + return table[codes.long() & 0xF] + +def _pack_e2m1(codes: torch.Tensor) -> torch.Tensor: + """Pack e2m1 nibble codes pairwise into bytes: low nibble = even element, high nibble = odd.""" + assert codes.shape[-1] % 2 == 0 + return (codes[..., 0::2] | (codes[..., 1::2] << 4)).to(torch.uint8) + +def _unpack_e2m1(packed: torch.Tensor) -> torch.Tensor: + """Inverse of _pack_e2m1. Returns nibble codes with last dim doubled.""" + lo = packed & 0xF + hi = (packed >> 4) & 0xF + return torch.stack([lo, hi], dim=-1).flatten(start_dim=-2) + +# --- NVFP4 scale-factor permutation ------------------------------------------------- +# The kernel's dequant warpgroup gives each thread 8 of a token's 32 scale factors: thread +# q (= idx_in_group/2, in [0,4)) owns element blocks {4c + q : c = 0..7}, which in element +# order are 8 bytes with stride 4. The on-wire tail stores them permuted so those 8 are +# contiguous and can be fetched with a single 8-byte load: +# scale for element block s -> byte 8*(s & 3) + (s >> 2) +# Keep this in lockstep with NVFP4_SF_NOPE_OFFSET/nvfp4_sf_byte in +# csrc/sm100/decode/head64/config.h and with the production writer in vLLM. +_NVFP4_SF_COLS = 8 # scale factors owned by one dequant thread (kernel COLS_PER_GROUP) +_NVFP4_SF_QUADS = 4 # distinct thread-quarter indices q (kernel GROUP_SIZE/2) + +def _nvfp4_permute_sf(sf: torch.Tensor) -> torch.Tensor: + """[..., 32] in element-block order -> [..., 32] in on-wire (kernel) order.""" + return sf.unflatten(-1, (_NVFP4_SF_COLS, _NVFP4_SF_QUADS)).transpose(-1, -2).flatten(-2) + +def _nvfp4_unpermute_sf(sf: torch.Tensor) -> torch.Tensor: + """Inverse of _nvfp4_permute_sf.""" + return sf.unflatten(-1, (_NVFP4_SF_QUADS, _NVFP4_SF_COLS)).transpose(-1, -2).flatten(-2) + +def _quant_tiles_e4m3_sf(x: torch.Tensor, tile_size: int, max_val: float): + """ + Per-`tile_size` quantization with e4m3 scale factors: sf = e4m3(amax/max_val) + rounded UP to the next representable e4m3 value (so that amax/float(sf) never + exceeds max_val — round-to-nearest can round down by up to 12.5% in e4m3's + subnormal range, saturating the largest values of a tile), q = x / float(sf). + Returns (x_scaled, sf) where x_scaled is float32 (not yet cast to the target + dtype) of the same shape as x, and sf is float8_e4m3fn of shape + (*x.shape[:-1], x.shape[-1]//tile_size). + """ + tiles = x.float().unflatten(-1, (-1, tile_size)) # [..., num_tiles, tile_size] + amax = tiles.abs().amax(dim=-1) + sf_target = torch.clamp_min(amax / max_val, 2.0**-9) + sf = sf_target.to(torch.float8_e4m3fn) + # Round up: positive e4m3 bit patterns are monotonic (0x7E = 448 is max finite) + sf_bits = sf.view(torch.uint8) + bump = (sf.float() < sf_target) & (sf_bits < 0x7E) + sf = torch.where(bump, (sf_bits + 1).view(torch.float8_e4m3fn), sf) + x_scaled = tiles / sf.float().unsqueeze(-1) + return x_scaled.flatten(-2), sf + def quantize_k_cache( input_k_cache: torch.Tensor, # (num_blocks, block_size, h_k, d) kvcache_layout: FP8KVCacheLayout, @@ -74,9 +159,29 @@ def quantize_k_cache( result = result.view(num_blocks, block_size, 1, -1) return result + elif kvcache_layout.is_nvfp4(): + bytes_per_token = kvcache_layout.bytes_per_token() + num_nope_sf = d_nope // tile_size # 32 + sf_nope_off = d_nope // 2 + d_rope # NoPE is e2m1 (2 values/byte), RoPE is e4m3 + + # Over-allocate one extra token row per block (mirroring the V32 layout above) so that + # any trailing TMA reads stay within valid memory. + result = torch.zeros((num_blocks, block_size+1, bytes_per_token), dtype=torch.uint8, device=input_k_cache.device)[:, :block_size, :] + + # NoPE: e2m1 with per-16 e4m3 scale factors + nope_scaled, nope_sf = _quant_tiles_e4m3_sf(input_k_cache[..., :d_nope], tile_size, 6.0) + result[..., :d_nope//2] = _pack_e2m1(_cast_to_e2m1_codes(nope_scaled)) + result[..., sf_nope_off:sf_nope_off+num_nope_sf] = _nvfp4_permute_sf(nope_sf.view(torch.uint8)) + + # RoPE: plain e4m3, no scale factor + result[..., d_nope//2:sf_nope_off] = input_k_cache[..., d_nope:].to(torch.float8_e4m3fn).view(torch.uint8) + + result = result.view(num_blocks, block_size, 1, -1) + return result + else: raise NotImplementedError(f"Unsupported kvcache_layout: {kvcache_layout}") - + def dequantize_k_cache( quant_k_cache: torch.Tensor, # (num_blocks, block_size, 1, bytes_per_token) @@ -115,7 +220,24 @@ def dequantize_k_cache( cur_nope = input_nope[..., tile_idx*tile_size:(tile_idx+1)*tile_size].to(torch.bfloat16) cur_scales = input_scale[:, :, tile_idx].to(torch.bfloat16).unsqueeze(-1) result[..., tile_idx*tile_size: (tile_idx+1)*tile_size] = cur_nope * cur_scales - + + elif kvcache_layout.is_nvfp4(): + # NOTE This must match the kernel's dequantization bit-for-bit. The kernel multiplies + # the (exactly-representable) data value with the bf16-converted e4m3 scale factor in + # bf16; since data*sf has at most 8 mantissa bits, a float32 multiply followed by a + # bf16 round-trip produces identical bits. + num_nope_sf = d_nope // tile_size + sf_nope_off = d_nope // 2 + d_rope + + quant_k_cache = quant_k_cache.view(torch.uint8).view(num_blocks, block_size, -1) + nope_vals = _e2m1_codes_to_float(_unpack_e2m1(quant_k_cache[..., :d_nope//2])) # [nb, bs, d_nope] fp32 + nope_sf = _nvfp4_unpermute_sf( + quant_k_cache[..., sf_nope_off:sf_nope_off+num_nope_sf]).view(torch.float8_e4m3fn).float() + result[..., :d_nope] = (nope_vals.unflatten(-1, (-1, tile_size)) * nope_sf.unsqueeze(-1)).flatten(-2).to(torch.bfloat16) + + # RoPE: plain e4m3, no scale factor + result[..., d_nope:] = quant_k_cache[..., d_nope//2:sf_nope_off].view(torch.float8_e4m3fn).to(torch.bfloat16) + else: raise NotImplementedError(f"Unsupported kvcache_layout: {kvcache_layout}") diff --git a/tests/test_flash_mla_sparse_decoding.py b/tests/test_flash_mla_sparse_decoding.py index 31aac4863..3f4c425fe 100644 --- a/tests/test_flash_mla_sparse_decoding.py +++ b/tests/test_flash_mla_sparse_decoding.py @@ -99,9 +99,52 @@ def gen_testcase() -> List[RawTestParam]: ] corner_cases.extend(cur_corner_cases) + # NVFP4 KV cache format (SM100 only, V3.2 geometry: d_qk = 576) + for h_q in [64, 128]: + for have_topk_len in [False, True]: + correctness_cases.extend([ + RawTestParam(b, h_q, s_q, 1, s_k, is_varlen, topk, + have_topk_length=have_topk_len, + enable_attn_sink=True, + block_size=block_size, + d_qk=576, + check_correctness=True, + num_runs=0, + kv_format="nvfp4.fp8rope") + for (s_k, topk, block_size) in [ + (512, 64, 2), + (512, 64, 64), + (512, 64, 69), + (1024, 576, 61), + (2046, 2048, 64), + ] + for b in [4, 74] + for s_q in [1, 3] + for is_varlen in ([True, False] if (b == 74 and not have_topk_len) else [True]) + ]) + corner_cases.extend([ + RawTestParam(b, h_q, 3, 1, s_k, True, topk, + is_all_indices_invalid=is_all_indices_invalid, + have_zero_seqlen_k=have_zero_seqlen_k, + enable_attn_sink=enable_attn_sink, + block_size=block_size, + d_qk=576, + check_correctness=True, + num_runs=0, + kv_format="nvfp4.fp8rope") + for (s_k, topk, block_size) in [(512, 64, 61), (650, 576, 53)] + for b in [4, 74] + for is_all_indices_invalid in [True, False] + for have_zero_seqlen_k in [True, False] + for enable_attn_sink in [True, False] + if (is_all_indices_invalid or have_zero_seqlen_k or enable_attn_sink) + ]) + base_and_bszs = [ # V3.2 (RawTestParam(0, 128, 2, 1, 32768, True, topk=2048, d_qk=576), [2, 64, 74, 128]), + # V3.2 shape with NVFP4 KV cache + (RawTestParam(0, 128, 2, 1, 32768, True, topk=2048, d_qk=576, kv_format="nvfp4.fp8rope"), [64, 128]), # MODEL1 CONFIG1 (RawTestParam(0, 64, 2, 1, 16384, True, topk=128, d_qk=512, extra_s_k=16384, extra_topk=512, block_size=256, extra_block_size=64), [2, 64, 74, 128, 74*2, 256]), # MODEL1 CONFIG2 @@ -121,6 +164,10 @@ def gen_testcase() -> List[RawTestParam]: RawTestParam(74*2, h_q, 2, 1, 32768, True, topk=16384, d_qk=d_qk) for h_q in [64, 128] for d_qk in [512, 576] + ] + [ + # Peak perf cases, NVFP4 KV cache + RawTestParam(74*2, h_q, 2, 1, 32768, True, topk=16384, d_qk=576, kv_format="nvfp4.fp8rope") + for h_q in [64, 128] ] return correctness_cases + corner_cases + performance_cases