From de26bc8636b97d7be09202867275a1870ec53cf1 Mon Sep 17 00:00:00 2001 From: Yan Wang Date: Mon, 15 Jun 2026 04:41:17 -0700 Subject: [PATCH] Add SM120 pure NVFP4 attention JIT path --- benchmarks/bench_nvfp4_attention_sm120.py | 442 ++++++++++ .../nvfp4_attention_sm120_binding.cu | 311 +++++++ .../nvfp4_attention_sm120_quantize.cu | 597 +++++++++++++ docs/api/attention.rst | 12 + flashinfer/__init__.py | 7 + flashinfer/aot.py | 3 + flashinfer/jit/__init__.py | 3 + flashinfer/jit/nvfp4_attention_sm120.py | 97 ++ flashinfer/nvfp4_attention_sm120.py | 421 +++++++++ .../trace/templates/nvfp4_attention_sm120.py | 92 ++ .../nvfp4_attention_sm120/api/launcher.h | 146 +++ .../nvfp4_attention_sm120/common/block_info.h | 67 ++ .../common/cute_extension.h | 511 +++++++++++ .../common/gemm_with_interleave.h | 72 ++ .../nvfp4_attention_sm120/common/params.h | 143 +++ .../common/static_switch.h | 87 ++ .../compute/consumer/delta_correction.cuh | 79 ++ .../compute/consumer/pv_gemm.cuh | 176 ++++ .../compute/consumer/qk_gemm.cuh | 155 ++++ .../compute/consumer/softmax.cuh | 774 ++++++++++++++++ .../compute/epilogue.cuh | 217 +++++ .../compute/epilogue/lse_writer.cuh | 142 +++ .../compute/epilogue/output_writer.cuh | 187 ++++ .../compute/mainloop.cuh | 833 ++++++++++++++++++ .../compute/producer/load_k.cuh | 128 +++ .../compute/producer/load_q.cuh | 101 +++ .../compute/producer/load_v.cuh | 106 +++ .../kernel/attention_kernel.h | 259 ++++++ .../nvfp4_attention_sm120/kernel/scheduler.h | 222 +++++ .../nvfp4_attention_sm120/kernel/traits.h | 241 +++++ .../primitives/barrier.cuh | 249 ++++++ .../primitives/pipeline.cuh | 132 +++ .../nvfp4_attention_sm120/primitives/tma.cuh | 108 +++ .../primitives/warpgroup.cuh | 132 +++ .../quantization/fp4_convert.cuh | 64 ++ .../quantization/fp4_layout.h | 143 +++ .../nvfp4_attention_sm120/utils/copy.cuh | 59 ++ .../nvfp4_attention_sm120/utils/layout.cuh | 61 ++ .../nvfp4_attention_sm120/utils/math.cuh | 31 + include/flashinfer/math.cuh | 116 +++ tests/attention/test_nvfp4_attention_sm120.py | 222 +++++ tests/conftest.py | 13 + 42 files changed, 7961 insertions(+) create mode 100644 benchmarks/bench_nvfp4_attention_sm120.py create mode 100644 csrc/nvfp4_attention_sm120/nvfp4_attention_sm120_binding.cu create mode 100644 csrc/nvfp4_attention_sm120/nvfp4_attention_sm120_quantize.cu create mode 100644 flashinfer/jit/nvfp4_attention_sm120.py create mode 100644 flashinfer/nvfp4_attention_sm120.py create mode 100644 flashinfer/trace/templates/nvfp4_attention_sm120.py create mode 100644 include/flashinfer/attention/sm120/nvfp4_attention_sm120/api/launcher.h create mode 100644 include/flashinfer/attention/sm120/nvfp4_attention_sm120/common/block_info.h create mode 100644 include/flashinfer/attention/sm120/nvfp4_attention_sm120/common/cute_extension.h create mode 100644 include/flashinfer/attention/sm120/nvfp4_attention_sm120/common/gemm_with_interleave.h create mode 100644 include/flashinfer/attention/sm120/nvfp4_attention_sm120/common/params.h create mode 100644 include/flashinfer/attention/sm120/nvfp4_attention_sm120/common/static_switch.h create mode 100644 include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/consumer/delta_correction.cuh create mode 100644 include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/consumer/pv_gemm.cuh create mode 100644 include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/consumer/qk_gemm.cuh create mode 100644 include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/consumer/softmax.cuh create mode 100644 include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/epilogue.cuh create mode 100644 include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/epilogue/lse_writer.cuh create mode 100644 include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/epilogue/output_writer.cuh create mode 100644 include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/mainloop.cuh create mode 100644 include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/producer/load_k.cuh create mode 100644 include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/producer/load_q.cuh create mode 100644 include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/producer/load_v.cuh create mode 100644 include/flashinfer/attention/sm120/nvfp4_attention_sm120/kernel/attention_kernel.h create mode 100644 include/flashinfer/attention/sm120/nvfp4_attention_sm120/kernel/scheduler.h create mode 100644 include/flashinfer/attention/sm120/nvfp4_attention_sm120/kernel/traits.h create mode 100644 include/flashinfer/attention/sm120/nvfp4_attention_sm120/primitives/barrier.cuh create mode 100644 include/flashinfer/attention/sm120/nvfp4_attention_sm120/primitives/pipeline.cuh create mode 100644 include/flashinfer/attention/sm120/nvfp4_attention_sm120/primitives/tma.cuh create mode 100644 include/flashinfer/attention/sm120/nvfp4_attention_sm120/primitives/warpgroup.cuh create mode 100644 include/flashinfer/attention/sm120/nvfp4_attention_sm120/quantization/fp4_convert.cuh create mode 100644 include/flashinfer/attention/sm120/nvfp4_attention_sm120/quantization/fp4_layout.h create mode 100644 include/flashinfer/attention/sm120/nvfp4_attention_sm120/utils/copy.cuh create mode 100644 include/flashinfer/attention/sm120/nvfp4_attention_sm120/utils/layout.cuh create mode 100644 include/flashinfer/attention/sm120/nvfp4_attention_sm120/utils/math.cuh create mode 100644 tests/attention/test_nvfp4_attention_sm120.py diff --git a/benchmarks/bench_nvfp4_attention_sm120.py b/benchmarks/bench_nvfp4_attention_sm120.py new file mode 100644 index 00000000000..7d9c09da503 --- /dev/null +++ b/benchmarks/bench_nvfp4_attention_sm120.py @@ -0,0 +1,442 @@ +""" +Copyright (c) 2026 by FlashInfer team. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +import argparse +import csv +import math +import statistics +import sys +from dataclasses import dataclass +from typing import Sequence + +import torch + + +DEFAULT_CONFIGS = ( + (4, 8, 4096, 128, False), + (1, 8, 32768, 128, False), +) +PER_BLOCK_MEAN = True +CSV_FIELDS = ( + "batch_size", + "num_heads", + "seq_len", + "head_dim", + "causal", + "dtype", + "attention_only_ms", + "attention_only_tflops", + "attention_only_cuda_graph", + "end_to_end_ms", + "end_to_end_attention_tflops", + "warmup", + "repeat", +) + + +def _patch_cutlass_dsl_operand_major_mode() -> None: + try: + import cutlass.cute as cute + from cutlass.cute.nvgpu.tcgen05 import OperandMajorMode + except ImportError: + return + if not hasattr(cute.nvgpu, "OperandMajorMode"): + cute.nvgpu.OperandMajorMode = OperandMajorMode + + +@dataclass(frozen=True) +class BenchConfig: + batch_size: int + num_heads: int + seq_len: int + head_dim: int + causal: bool + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Benchmark FlashInfer SM120 NVFP4 attention latency and TFLOPs/s." + ) + parser.add_argument( + "--batch-size", + "--batch_size", + type=int, + nargs="+", + default=None, + help="Batch size(s). Scalars are broadcast across other shape lists.", + ) + parser.add_argument( + "--num-heads", + "--num_heads", + type=int, + nargs="+", + default=None, + help="Number of attention heads.", + ) + parser.add_argument( + "--seq-len", + "--seq_len", + type=int, + nargs="+", + default=None, + help="Sequence length(s). Must be a multiple of 128.", + ) + parser.add_argument( + "--head-dim", + "--head_dim", + type=int, + nargs="+", + default=None, + help="Head dimension(s). SM120 NVFP4 attention supports 64 or 128.", + ) + causal_group = parser.add_mutually_exclusive_group() + causal_group.add_argument( + "--causal", + dest="causal", + action="store_true", + default=None, + help="Benchmark causal attention.", + ) + causal_group.add_argument( + "--no-causal", + dest="causal", + action="store_false", + help="Benchmark non-causal attention.", + ) + parser.add_argument( + "--dtype", + type=str, + choices=("float16", "bfloat16"), + default="bfloat16", + help="Input/output dtype before NVFP4 quantization.", + ) + parser.add_argument( + "--warmup", + type=int, + default=5, + help="Number of warmup iterations for each measured path.", + ) + parser.add_argument( + "--repeat", + type=int, + default=20, + help="Number of measured iterations for each measured path.", + ) + parser.add_argument( + "--no-attention-cuda-graph", + action="store_true", + help=( + "Measure attention-only with normal CUDA events. By default the " + "attention-only path uses CUDA Graph replay to report pure GPU " + "kernel time without TVM FFI launch overhead." + ), + ) + parser.add_argument( + "--save-results-to", + type=str, + default=None, + help="Optional path to save benchmark results as CSV.", + ) + return parser.parse_args() + + +def skip_unless_sm120() -> None: + if not torch.cuda.is_available(): + print("Skipping: NVFP4 attention SM120 benchmark requires CUDA.") + sys.exit(0) + + capability = torch.cuda.get_device_capability() + if capability != (12, 0): + print(f"Current device capability: {capability}.") + print( + "Skipping: NVFP4 attention SM120 benchmark requires compute capability (12, 0)." + ) + sys.exit(0) + + +def torch_dtype(dtype: str) -> torch.dtype: + if dtype == "float16": + return torch.float16 + if dtype == "bfloat16": + return torch.bfloat16 + raise ValueError(f"Unsupported dtype: {dtype}") + + +def expand_values(name: str, values: Sequence[int] | None, default: int) -> list[int]: + if values is None: + return [default] + if not values: + raise ValueError(f"{name} must not be empty") + return list(values) + + +def broadcast_shape_lists(values: dict[str, list[int]]) -> dict[str, list[int]]: + max_len = max(len(v) for v in values.values()) + out = {} + for name, vals in values.items(): + if len(vals) == 1: + out[name] = vals * max_len + elif len(vals) == max_len: + out[name] = vals + else: + raise ValueError( + f"{name} has {len(vals)} values, expected 1 or {max_len} values" + ) + return out + + +def build_configs(args: argparse.Namespace) -> list[BenchConfig]: + has_custom_shape = any( + arg is not None + for arg in (args.batch_size, args.num_heads, args.seq_len, args.head_dim) + ) + if not has_custom_shape: + return [ + BenchConfig( + batch_size=batch_size, + num_heads=num_heads, + seq_len=seq_len, + head_dim=head_dim, + causal=args.causal if args.causal is not None else causal, + ) + for batch_size, num_heads, seq_len, head_dim, causal in DEFAULT_CONFIGS + ] + + values = broadcast_shape_lists( + { + "batch_size": expand_values("batch_size", args.batch_size, 4), + "num_heads": expand_values("num_heads", args.num_heads, 8), + "seq_len": expand_values("seq_len", args.seq_len, 4096), + "head_dim": expand_values("head_dim", args.head_dim, 128), + } + ) + causal = args.causal if args.causal is not None else False + return [ + BenchConfig( + batch_size=values["batch_size"][idx], + num_heads=values["num_heads"][idx], + seq_len=values["seq_len"][idx], + head_dim=values["head_dim"][idx], + causal=causal, + ) + for idx in range(len(values["batch_size"])) + ] + + +def validate_config(config: BenchConfig) -> None: + if config.batch_size <= 0: + raise ValueError(f"batch_size must be positive, got {config.batch_size}") + if config.num_heads <= 0: + raise ValueError(f"num_heads must be positive, got {config.num_heads}") + if config.seq_len <= 0 or config.seq_len % 128 != 0: + raise ValueError( + f"seq_len must be positive and divisible by 128, got {config.seq_len}" + ) + if config.head_dim not in (64, 128): + raise ValueError(f"head_dim must be 64 or 128, got {config.head_dim}") + + +def attention_flops(config: BenchConfig) -> float: + factor = 2 if config.causal else 4 + return ( + factor + * config.batch_size + * config.num_heads + * config.seq_len + * config.seq_len + * config.head_dim + ) + + +def tflops_per_sec(config: BenchConfig, ms: float) -> float: + return attention_flops(config) / ms / 1e9 + + +def dtype_label(dtype: torch.dtype) -> str: + return str(dtype).removeprefix("torch.") + + +def median_gpu_ms( + fn, + warmup: int, + repeat: int, + use_cuda_graph: bool = False, + cold_l2_cache: bool = True, + num_iters_within_graph: int = 1, +) -> float: + _patch_cutlass_dsl_operand_major_mode() + import flashinfer.testing + + measurements = flashinfer.testing.bench_gpu_time( + fn, + dry_run_iters=warmup, + repeat_iters=repeat, + cold_l2_cache=cold_l2_cache, + use_cuda_graph=use_cuda_graph, + num_iters_within_graph=num_iters_within_graph, + ) + return statistics.median(measurements) + + +def bench_config( + config: BenchConfig, + dtype: torch.dtype, + warmup: int, + repeat: int, + attention_cuda_graph: bool, +) -> dict[str, object]: + _patch_cutlass_dsl_operand_major_mode() + import flashinfer + + validate_config(config) + torch.manual_seed(123) + + q = torch.randn( + config.batch_size, + config.num_heads, + config.seq_len, + config.head_dim, + dtype=dtype, + device="cuda", + ) + k = torch.randn_like(q) + v = torch.randn_like(q) + + sm_scale = 1.0 / math.sqrt(config.head_dim) + out = torch.empty_like(q) + lse = torch.empty( + config.batch_size, + config.num_heads, + config.seq_len, + dtype=torch.float32, + device="cuda", + ) + + quantized_qkv = flashinfer.nvfp4_attention_sm120_quantize_qkv( + q, k, v, per_block_mean=PER_BLOCK_MEAN + ) + flashinfer.nvfp4_attention_sm120_fwd( + *quantized_qkv, + sm_scale=sm_scale, + causal=config.causal, + per_block_mean=PER_BLOCK_MEAN, + out=out, + lse=lse, + out_dtype=dtype, + ) + torch.cuda.synchronize() + + def attention_only(): + return flashinfer.nvfp4_attention_sm120_fwd( + *quantized_qkv, + sm_scale=sm_scale, + causal=config.causal, + per_block_mean=PER_BLOCK_MEAN, + out=out, + lse=lse, + out_dtype=dtype, + ) + + def end_to_end(): + qkv = flashinfer.nvfp4_attention_sm120_quantize_qkv( + q, k, v, per_block_mean=PER_BLOCK_MEAN + ) + return flashinfer.nvfp4_attention_sm120_fwd( + *qkv, + sm_scale=sm_scale, + causal=config.causal, + per_block_mean=PER_BLOCK_MEAN, + out=out, + lse=lse, + out_dtype=dtype, + ) + + attention_only_ms = median_gpu_ms( + attention_only, + warmup, + repeat, + use_cuda_graph=attention_cuda_graph, + cold_l2_cache=not attention_cuda_graph, + num_iters_within_graph=1, + ) + end_to_end_ms = median_gpu_ms(end_to_end, warmup, repeat) + attention_only_tflops = tflops_per_sec(config, attention_only_ms) + end_to_end_tflops = tflops_per_sec(config, end_to_end_ms) + + print( + "nvfp4_attention_sm120 " + f"B={config.batch_size} H={config.num_heads} S={config.seq_len} " + f"D={config.head_dim} causal={config.causal} dtype={dtype}: " + f"attention_only={attention_only_ms:.3f} ms " + f"({attention_only_tflops:.3f} TFLOPs/s, " + f"cuda_graph={attention_cuda_graph}), " + f"end_to_end={end_to_end_ms:.3f} ms " + f"({end_to_end_tflops:.3f} attention-TFLOPs/s)" + ) + return { + "batch_size": config.batch_size, + "num_heads": config.num_heads, + "seq_len": config.seq_len, + "head_dim": config.head_dim, + "causal": config.causal, + "dtype": dtype_label(dtype), + "attention_only_ms": attention_only_ms, + "attention_only_tflops": attention_only_tflops, + "attention_only_cuda_graph": attention_cuda_graph, + "end_to_end_ms": end_to_end_ms, + "end_to_end_attention_tflops": end_to_end_tflops, + "warmup": warmup, + "repeat": repeat, + } + + +def write_csv(path: str, rows: list[dict[str, object]]) -> None: + with open(path, "w", newline="") as csvfile: + writer = csv.DictWriter(csvfile, fieldnames=CSV_FIELDS) + writer.writeheader() + writer.writerows(rows) + print(f"Results saved to: {path}") + + +def main() -> None: + args = parse_args() + skip_unless_sm120() + + if args.warmup < 0: + raise ValueError(f"warmup must be non-negative, got {args.warmup}") + if args.repeat <= 0: + raise ValueError(f"repeat must be positive, got {args.repeat}") + + dtype = torch_dtype(args.dtype) + configs = build_configs(args) + results = [] + for config in configs: + results.append( + bench_config( + config, + dtype, + args.warmup, + args.repeat, + attention_cuda_graph=not args.no_attention_cuda_graph, + ) + ) + + if args.save_results_to: + write_csv(args.save_results_to, results) + + +if __name__ == "__main__": + main() diff --git a/csrc/nvfp4_attention_sm120/nvfp4_attention_sm120_binding.cu b/csrc/nvfp4_attention_sm120/nvfp4_attention_sm120_binding.cu new file mode 100644 index 00000000000..3cb4cb470a9 --- /dev/null +++ b/csrc/nvfp4_attention_sm120/nvfp4_attention_sm120_binding.cu @@ -0,0 +1,311 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +#include "dlpack/dlpack.h" + +using tvm::ffi::TensorView; +namespace ffi = tvm::ffi; + +constexpr DLDataType dl_uint8 = DLDataType{kDLUInt, 8, 1}; +constexpr DLDataType dl_float16 = DLDataType{kDLFloat, 16, 1}; +constexpr DLDataType dl_float32 = DLDataType{kDLFloat, 32, 1}; +constexpr DLDataType dl_float8_e4m3fn = DLDataType{kDLFloat8_e4m3fn, 8, 1}; +constexpr DLDataType dl_bfloat16 = DLDataType{kDLBfloat, 16, 1}; + +#define CHECK_CUDA(x) \ + TVM_FFI_ICHECK_EQ(x.device().device_type, kDLCUDA) << #x " must be a CUDA tensor"; +#define CHECK_CONTIGUOUS(x) TVM_FFI_ICHECK(x.IsContiguous()) << #x " must be contiguous"; +#define CHECK_INPUT(x) \ + CHECK_CUDA(x); \ + CHECK_CONTIGUOUS(x) +#define CHECK_DIM(d, x) TVM_FFI_ICHECK_EQ(x.ndim(), d) << #x " must be a " #d "D tensor"; + +inline cudaStream_t get_stream(DLDevice device) { + return static_cast(TVMFFIEnvGetStream(device.device_type, device.device_id)); +} + +inline int64_t get_element_size(ffi::TensorView x) { + return (x.dtype().bits * x.dtype().lanes) / 8; +} + +namespace flashinfer { +namespace nvfp4_attention_sm120 { + +namespace { + +int64_t numel(TensorView x) { + int64_t n = 1; + for (int i = 0; i < x.ndim(); ++i) { + n *= x.size(i); + } + return n; +} + +void check_same_device(TensorView ref, TensorView x, const char* name) { + TVM_FFI_ICHECK_EQ(ref.device().device_type, x.device().device_type) + << name << " must be on the same device as q_fp4"; + TVM_FFI_ICHECK_EQ(ref.device().device_id, x.device().device_id) + << name << " must be on the same device as q_fp4"; +} + +int round_multiple(int x, int m) { return (x + m - 1) / m * m; } + +void set_params_fprop(Flash_fwd_params& params, TensorView q, TensorView k, TensorView v, + TensorView q_scale, TensorView k_scale, TensorView v_scale, + TensorView qk_correction, TensorView out, TensorView lse, float sm_scale, + bool causal, bool per_block_mean) { + params = {}; + + const int batch = static_cast(q.size(0)); + const int num_heads = static_cast(q.size(1)); + const int seq_len_q = static_cast(q.size(2)); + const int seq_len_k = static_cast(k.size(2)); + const int head_dim = static_cast(q.size(3) * 2); + + params.q_ptr = q.data_ptr(); + params.k_ptr = k.data_ptr(); + params.v_ptr = v.data_ptr(); + params.delta_s_ptr = qk_correction.data_ptr(); + params.sfq_ptr = q_scale.data_ptr(); + params.sfk_ptr = k_scale.data_ptr(); + params.sfv_ptr = v_scale.data_ptr(); + + params.q_row_stride = q.stride(-2) * 2; + params.k_row_stride = k.stride(-2) * 2; + params.v_row_stride = v.stride(-2) * 2; + params.q_head_stride = q.stride(-3) * 2; + params.k_head_stride = k.stride(-3) * 2; + params.v_head_stride = v.stride(-3) * 2; + params.q_batch_stride = q.stride(0) * 2; + params.k_batch_stride = k.stride(0) * 2; + params.v_batch_stride = v.stride(0) * 2; + + params.ds_row_stride = qk_correction.stride(-2); + params.ds_head_stride = qk_correction.stride(-3); + params.ds_batch_stride = qk_correction.stride(0); + + params.sfq_row_stride = q_scale.stride(-2); + params.sfk_row_stride = k_scale.stride(-2); + params.sfv_row_stride = v_scale.stride(-2); + params.sfq_head_stride = q_scale.stride(-3); + params.sfk_head_stride = k_scale.stride(-3); + params.sfv_head_stride = v_scale.stride(-3); + params.sfq_batch_stride = q_scale.stride(0); + params.sfk_batch_stride = k_scale.stride(0); + params.sfv_batch_stride = v_scale.stride(0); + + params.o_ptr = out.data_ptr(); + params.o_row_stride = out.stride(-2); + params.o_head_stride = out.stride(-3); + params.o_batch_stride = out.stride(0); + + params.cu_seqlens_q = nullptr; + params.cu_seqlens_k = nullptr; + params.seqused_k = nullptr; + params.p_ptr = nullptr; + params.softmax_lse_ptr = lse.data_ptr(); + + params.b = batch; + params.h = num_heads; + params.h_k = num_heads; + params.h_h_k_ratio = 1; + params.seqlen_q = seq_len_q; + params.seqlen_k = seq_len_k; + params.unpadded_seqlen_k = seq_len_k; + params.seqlen_q_rounded = round_multiple(seq_len_q, 128); + params.seqlen_k_rounded = round_multiple(seq_len_k, 128); + params.d = head_dim; + params.d_rounded = head_dim; + params.head_divmod = cutlass::FastDivmod(num_heads); + + params.scale_softmax = sm_scale; + params.scale_softmax_log2 = sm_scale * 1.4426950408889634f; + __half scale_softmax_log2_half = __float2half(params.scale_softmax_log2); + __half2 scale_softmax_log2_half2 = + __halves2half2(scale_softmax_log2_half, scale_softmax_log2_half); + params.scale_softmax_log2_half2 = reinterpret_cast(scale_softmax_log2_half2); + + params.p_dropout = 1.f; + params.p_dropout_in_uint8_t = 255; + params.rp_dropout = 1.f; + params.scale_softmax_rp_dropout = sm_scale; + + params.is_causal = causal; + params.per_block_mean = per_block_mean; + params.seqlen_s = per_block_mean ? seq_len_q : 128; + params.window_size_left = -1; + params.window_size_right = causal ? 0 : -1; + params.is_seqlens_k_cumulative = true; + params.is_bf16 = out.dtype() == dl_bfloat16; + params.tile_count_semaphore = nullptr; +} + +template +void run_mha_fwd_dispatch_dtype(Flash_fwd_params& params, cudaStream_t stream) { + using OType = std::conditional_t; + if (params.d == 64) { + ::nvfp4_attention::run_mha_fwd_, 64, OType>(params, + stream); + } else if (params.d == 128) { + ::nvfp4_attention::run_mha_fwd_, 128, OType>( + params, stream); + } else { + TVM_FFI_ICHECK(false) << "Unsupported head dimension " << params.d; + } +} + +void run_mha_fwd(Flash_fwd_params& params, cudaStream_t stream) { + if (params.is_bf16) { + run_mha_fwd_dispatch_dtype(params, stream); + } else { + run_mha_fwd_dispatch_dtype(params, stream); + } +} + +} // namespace + +void fwd(TensorView q_fp4, TensorView k_fp4, TensorView v_fp4_t, TensorView q_scale, + TensorView k_scale, TensorView v_scale_t, TensorView qk_correction, TensorView out, + TensorView lse, double sm_scale, bool causal, bool per_block_mean) { + CHECK_INPUT(q_fp4); + CHECK_INPUT(k_fp4); + CHECK_INPUT(v_fp4_t); + CHECK_INPUT(q_scale); + CHECK_INPUT(k_scale); + CHECK_INPUT(v_scale_t); + CHECK_INPUT(qk_correction); + CHECK_INPUT(out); + CHECK_INPUT(lse); + + CHECK_DIM(4, q_fp4); + CHECK_DIM(4, k_fp4); + CHECK_DIM(4, v_fp4_t); + CHECK_DIM(4, q_scale); + CHECK_DIM(4, k_scale); + CHECK_DIM(4, v_scale_t); + CHECK_DIM(4, qk_correction); + CHECK_DIM(4, out); + CHECK_DIM(3, lse); + + TVM_FFI_ICHECK_EQ(q_fp4.dtype(), dl_uint8) << "q_fp4 must be uint8 packed FP4"; + TVM_FFI_ICHECK_EQ(k_fp4.dtype(), dl_uint8) << "k_fp4 must be uint8 packed FP4"; + TVM_FFI_ICHECK_EQ(v_fp4_t.dtype(), dl_uint8) << "v_fp4_t must be uint8 packed FP4"; + TVM_FFI_ICHECK_EQ(q_scale.dtype(), dl_float8_e4m3fn) << "q_scale must be float8_e4m3fn"; + TVM_FFI_ICHECK_EQ(k_scale.dtype(), dl_float8_e4m3fn) << "k_scale must be float8_e4m3fn"; + TVM_FFI_ICHECK_EQ(v_scale_t.dtype(), dl_float8_e4m3fn) << "v_scale_t must be float8_e4m3fn"; + TVM_FFI_ICHECK_EQ(qk_correction.dtype(), dl_float32) << "qk_correction must be float32"; + TVM_FFI_ICHECK_EQ(lse.dtype(), dl_float32) << "lse must be float32"; + TVM_FFI_ICHECK(out.dtype() == dl_bfloat16 || out.dtype() == dl_float16) + << "out must be bfloat16 or float16"; + + ffi::CUDADeviceGuard device_guard(q_fp4.device().device_id); + cudaDeviceProp props; + cudaError_t status = cudaGetDeviceProperties(&props, q_fp4.device().device_id); + TVM_FFI_ICHECK(status == cudaSuccess) + << "cudaGetDeviceProperties failed: " << cudaGetErrorString(status); + TVM_FFI_ICHECK(props.major == 12 && props.minor == 0) + << "NVFP4 attention SM120 kernel requires compute capability 12.0"; + + const int64_t batch = q_fp4.size(0); + const int64_t num_heads = q_fp4.size(1); + const int64_t seq_len = q_fp4.size(2); + const int64_t head_dim = q_fp4.size(3) * 2; + + TVM_FFI_ICHECK(head_dim == 64 || head_dim == 128) << "head_dim must be 64 or 128"; + TVM_FFI_ICHECK_EQ(seq_len % 128, 0) << "seq_len must be a multiple of 128"; + + TVM_FFI_ICHECK_EQ(k_fp4.size(0), batch); + TVM_FFI_ICHECK_EQ(k_fp4.size(1), num_heads); + TVM_FFI_ICHECK_EQ(k_fp4.size(2), seq_len); + TVM_FFI_ICHECK_EQ(k_fp4.size(3), q_fp4.size(3)); + + TVM_FFI_ICHECK_EQ(v_fp4_t.size(0), batch); + TVM_FFI_ICHECK_EQ(v_fp4_t.size(1), num_heads); + TVM_FFI_ICHECK_EQ(v_fp4_t.size(2), head_dim); + TVM_FFI_ICHECK_EQ(v_fp4_t.size(3), seq_len / 2); + + TVM_FFI_ICHECK_EQ(q_scale.size(0), batch); + TVM_FFI_ICHECK_EQ(q_scale.size(1), num_heads); + TVM_FFI_ICHECK_EQ(q_scale.size(2), seq_len); + TVM_FFI_ICHECK_EQ(q_scale.size(3), head_dim / 16); + TVM_FFI_ICHECK_EQ(k_scale.size(0), batch); + TVM_FFI_ICHECK_EQ(k_scale.size(1), num_heads); + TVM_FFI_ICHECK_EQ(k_scale.size(2), seq_len); + TVM_FFI_ICHECK_EQ(k_scale.size(3), head_dim / 16); + TVM_FFI_ICHECK_EQ(v_scale_t.size(0), batch); + TVM_FFI_ICHECK_EQ(v_scale_t.size(1), num_heads); + TVM_FFI_ICHECK_EQ(v_scale_t.size(2), head_dim); + TVM_FFI_ICHECK_EQ(v_scale_t.size(3), seq_len / 16); + + TVM_FFI_ICHECK_EQ(qk_correction.size(0), batch); + TVM_FFI_ICHECK_EQ(qk_correction.size(1), num_heads); + TVM_FFI_ICHECK_EQ(qk_correction.size(2), per_block_mean ? seq_len : 128); + TVM_FFI_ICHECK_EQ(qk_correction.size(3), seq_len); + + TVM_FFI_ICHECK_EQ(out.size(0), batch); + TVM_FFI_ICHECK_EQ(out.size(1), num_heads); + TVM_FFI_ICHECK_EQ(out.size(2), seq_len); + TVM_FFI_ICHECK_EQ(out.size(3), head_dim); + TVM_FFI_ICHECK_EQ(lse.size(0), batch); + TVM_FFI_ICHECK_EQ(lse.size(1), num_heads); + TVM_FFI_ICHECK_EQ(lse.size(2), seq_len); + + check_same_device(q_fp4, k_fp4, "k_fp4"); + check_same_device(q_fp4, v_fp4_t, "v_fp4_t"); + check_same_device(q_fp4, q_scale, "q_scale"); + check_same_device(q_fp4, k_scale, "k_scale"); + check_same_device(q_fp4, v_scale_t, "v_scale_t"); + check_same_device(q_fp4, qk_correction, "qk_correction"); + check_same_device(q_fp4, out, "out"); + check_same_device(q_fp4, lse, "lse"); + + cudaStream_t stream = get_stream(q_fp4.device()); + + if (seq_len == 0) { + status = cudaMemsetAsync(out.data_ptr(), 0, numel(out) * get_element_size(out), stream); + TVM_FFI_ICHECK(status == cudaSuccess) + << "cudaMemsetAsync(out) failed: " << cudaGetErrorString(status); + status = cudaMemsetAsync(lse.data_ptr(), 0, numel(lse) * get_element_size(lse), stream); + TVM_FFI_ICHECK(status == cudaSuccess) + << "cudaMemsetAsync(lse) failed: " << cudaGetErrorString(status); + return; + } + + Flash_fwd_params params; + set_params_fprop(params, q_fp4, k_fp4, v_fp4_t, q_scale, k_scale, v_scale_t, qk_correction, out, + lse, static_cast(sm_scale), causal, per_block_mean); + run_mha_fwd(params, stream); +} + +} // namespace nvfp4_attention_sm120 +} // namespace flashinfer + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(fwd, flashinfer::nvfp4_attention_sm120::fwd); diff --git a/csrc/nvfp4_attention_sm120/nvfp4_attention_sm120_quantize.cu b/csrc/nvfp4_attention_sm120/nvfp4_attention_sm120_quantize.cu new file mode 100644 index 00000000000..cba9c046c6b --- /dev/null +++ b/csrc/nvfp4_attention_sm120/nvfp4_attention_sm120_quantize.cu @@ -0,0 +1,597 @@ +/* + * Copyright (c) 2025 by SageAttention team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include +#include +#include + +#include +#include +#include +#include +#include + +#include "tvm_ffi_utils.h" + +#define DISPATCH_DLPACK_DTYPE_TO_CTYPE_FP16_LOCAL(dlpack_dtype, c_type, ...) \ + if ((dlpack_dtype) == dl_float16) { \ + using c_type = half; \ + __VA_ARGS__ \ + } else if ((dlpack_dtype) == dl_bfloat16) { \ + using c_type = nv_bfloat16; \ + __VA_ARGS__ \ + } else { \ + TVM_FFI_ICHECK(false) << __PRETTY_FUNCTION__ << " failed to dispatch data type"; \ + } + +#define DISPATCH_HEAD_DIM(head_dim, HEAD_DIM, ...) \ + if (head_dim == 64) { \ + constexpr int HEAD_DIM = 64; \ + __VA_ARGS__ \ + } else if (head_dim == 128) { \ + constexpr int HEAD_DIM = 128; \ + __VA_ARGS__ \ + } else { \ + TVM_FFI_ICHECK(false) << "Unsupported head dim: " << int(head_dim); \ + } + +#define CHECK_QUANT_CUDA(x) TVM_FFI_ICHECK_EQ(x.device().device_type, kDLCUDA) +#define CHECK_QUANT_DTYPE(x, true_dtype) \ + TVM_FFI_ICHECK((x).dtype() == (true_dtype)) << #x " dtype mismatch" +#define CHECK_QUANT_DIMS(x, true_dim) TVM_FFI_ICHECK_EQ((x).ndim(), true_dim) << #x " rank mismatch" +#define CHECK_QUANT_LASTDIM_CONTIGUOUS(x) \ + TVM_FFI_ICHECK_EQ((x).stride(-1), 1) << #x " must be contiguous at the last dimension" + +namespace flashinfer { +namespace nvfp4_attention_sm120 { + +void check_shape(TensorView x, std::initializer_list shape, const char* name) { + TVM_FFI_ICHECK_EQ(x.ndim(), static_cast(shape.size())) << name << " rank mismatch"; + int i = 0; + for (int64_t expected : shape) { + TVM_FFI_ICHECK_EQ(x.size(i), expected) << name << " shape mismatch at dim " << i; + ++i; + } +} + +#define CHECK_QUANT_SHAPE(x, ...) check_shape((x), {__VA_ARGS__}, #x) + +constexpr int CVT_FP4_ELTS_PER_THREAD = 16; + +template +struct TypeConverter { + using Type = half2; +}; + +template <> +struct TypeConverter { + using Type = half; +}; + +template <> +struct TypeConverter { + using Type = half2; +}; + +template <> +struct TypeConverter<__nv_bfloat162> { + using Type = __nv_bfloat16; +}; + +template <> +struct TypeConverter<__nv_bfloat16> { + using Type = __nv_bfloat162; +}; + +template +struct PackedVec { + typename TypeConverter::Type elts[8]; +}; + +template +__global__ void scaled_fp4_quant_kernel(const T* input, uint8_t* output, uint8_t* output_sf, + int batch_size, int num_heads, int num_tokens, + int stride_bz_input, int stride_h_input, + int stride_seq_input, int stride_bz_output, + int stride_h_output, int stride_seq_output, + int stride_bz_output_sf, int stride_h_output_sf, + int stride_seq_output_sf) { + static_assert(std::is_same::value || std::is_same::value, + "Only half and bfloat16 input are supported"); + using PackedVec = PackedVec; + + const int batch_id = blockIdx.y; + const int head_id = blockIdx.z; + const int token_block_id = blockIdx.x; + + static_assert(CVT_FP4_ELTS_PER_THREAD == 8 || CVT_FP4_ELTS_PER_THREAD == 16, + "CVT_FP4_ELTS_PER_THREAD must be 8 or 16"); + static_assert(sizeof(PackedVec) == sizeof(T) * CVT_FP4_ELTS_PER_THREAD, + "Vec size is not matched."); + + constexpr uint32_t NUM_THREADS_PER_TOKEN = head_dim / CVT_FP4_ELTS_PER_THREAD; + + const int token_id = token_block_id * BLOCK_SIZE + threadIdx.x / NUM_THREADS_PER_TOKEN; + + int load_token_id; + if constexpr (!permute) { + load_token_id = token_id; + } else { + int local_token_id = threadIdx.x / NUM_THREADS_PER_TOKEN; + int local_token_id_residue = local_token_id % 32; + + load_token_id = token_block_id * BLOCK_SIZE + (local_token_id / 32) * 32 + + (local_token_id_residue / 8) * 2 + ((local_token_id_residue % 8) / 2) * 8 + + (local_token_id_residue % 8) % 2; + } + + PackedVec in_vec; + +#pragma unroll + for (int i = 0; i < CVT_FP4_ELTS_PER_THREAD / 2; i++) { + reinterpret_cast(in_vec.elts[i]) = 0; + } + + if (load_token_id < num_tokens) { + in_vec = reinterpret_cast( + input + batch_id * stride_bz_input + head_id * stride_h_input + + load_token_id * stride_seq_input + + (threadIdx.x % NUM_THREADS_PER_TOKEN) * CVT_FP4_ELTS_PER_THREAD)[0]; + } + + auto localMax = __habs2(in_vec.elts[0]); +#pragma unroll + for (int i = 1; i < CVT_FP4_ELTS_PER_THREAD / 2; i++) { + localMax = __hmax2(localMax, __habs2(in_vec.elts[i])); + } + + if constexpr (CVT_FP4_ELTS_PER_THREAD == 8) { + localMax = __hmax2(__shfl_xor_sync(0xffffffff, localMax, 1, 32), localMax); + } + + float vecMax = float(__hmax(localMax.x, localMax.y)); + + float SFValue = vecMax / 6.0f; + uint8_t SFValueFP8; + reinterpret_cast<__nv_fp8_e4m3&>(SFValueFP8) = __nv_fp8_e4m3(SFValue); + SFValue = float(reinterpret_cast<__nv_fp8_e4m3&>(SFValueFP8)); + + float SFValueInv = (SFValue == 0.0f) ? 0.0f : 1.0f / SFValue; + + float2 fp2Vals[CVT_FP4_ELTS_PER_THREAD / 2]; + +#pragma unroll + for (int i = 0; i < CVT_FP4_ELTS_PER_THREAD / 2; i++) { + if constexpr (std::is_same::value) { + fp2Vals[i] = __half22float2(in_vec.elts[i]); + } else { + fp2Vals[i] = __bfloat1622float2(in_vec.elts[i]); + } + fp2Vals[i].x = fp2Vals[i].x * SFValueInv; + fp2Vals[i].y = fp2Vals[i].y * SFValueInv; + } + + uint32_t e2m1Vals[CVT_FP4_ELTS_PER_THREAD / 8]; +#pragma unroll + for (int i = 0; i < CVT_FP4_ELTS_PER_THREAD / 8; i++) { + e2m1Vals[i] = flashinfer::math::fp32_vec_to_e2m1(fp2Vals + i * 4); + } + + if constexpr (CVT_FP4_ELTS_PER_THREAD == 8) { + reinterpret_cast(output + batch_id * stride_bz_output + head_id * stride_h_output + + token_id * stride_seq_output + + (threadIdx.x % NUM_THREADS_PER_TOKEN) * CVT_FP4_ELTS_PER_THREAD / + 2)[0] = e2m1Vals[0]; + } else { + reinterpret_cast(output + batch_id * stride_bz_output + head_id * stride_h_output + + token_id * stride_seq_output + + (threadIdx.x % NUM_THREADS_PER_TOKEN) * CVT_FP4_ELTS_PER_THREAD / + 2)[0] = reinterpret_cast(e2m1Vals)[0]; + } + + uint8_t* output_sf_save_base = output_sf + batch_id * stride_bz_output_sf + + head_id * stride_h_output_sf + + (token_id / 64) * 64 * stride_seq_output_sf; + uint32_t token_id_local = token_id % 64; + + if constexpr (CVT_FP4_ELTS_PER_THREAD == 16) { + uint32_t col_id_local = threadIdx.x % NUM_THREADS_PER_TOKEN; + uint32_t offset_local = (col_id_local / 4) * 256 + (token_id_local % 16) * 16 + + (token_id_local / 16) * 4 + (col_id_local % 4); + reinterpret_cast(output_sf_save_base + offset_local)[0] = SFValueFP8; + } else { + if (threadIdx.x % 2 == 0) { + uint32_t col_id_local = (threadIdx.x % NUM_THREADS_PER_TOKEN) / 2; + uint32_t offset_local = (col_id_local / 4) * 256 + (token_id_local % 16) * 16 + + (token_id_local / 16) * 4 + (col_id_local % 4); + reinterpret_cast(output_sf_save_base + offset_local)[0] = SFValueFP8; + } + } +} + +template +__global__ void scaled_fp4_quant_trans_kernel(const T* input, uint8_t* output, uint8_t* output_sf, + int batch_size, int num_heads, int num_tokens, + int stride_bz_input, int stride_h_input, + int stride_seq_input, int stride_bz_output, + int stride_h_output, int stride_d_output, + int stride_bz_output_sf, int stride_h_output_sf, + int stride_d_output_sf) { + static_assert(std::is_same::value || std::is_same::value, + "Only half and bfloat16 input are supported"); + using PackedVec = PackedVec; + + const int batch_id = blockIdx.y; + const int head_id = blockIdx.z; + const int token_block_id = blockIdx.x; + + static_assert(CVT_FP4_ELTS_PER_THREAD == 8 || CVT_FP4_ELTS_PER_THREAD == 16, + "CVT_FP4_ELTS_PER_THREAD must be 8 or 16"); + static_assert(sizeof(PackedVec) == sizeof(T) * CVT_FP4_ELTS_PER_THREAD, + "Vec size is not matched."); + + constexpr uint32_t NUM_THREADS_PER_TOKEN = head_dim / CVT_FP4_ELTS_PER_THREAD; + constexpr uint32_t NUM_THREADS_PER_SEQ = BLOCK_SIZE / CVT_FP4_ELTS_PER_THREAD; + + const int token_id = token_block_id * BLOCK_SIZE + threadIdx.x / NUM_THREADS_PER_TOKEN; + + PackedVec in_vec; + +#pragma unroll + for (int i = 0; i < CVT_FP4_ELTS_PER_THREAD / 2; i++) { + reinterpret_cast(in_vec.elts[i]) = 0; + } + + if (token_id < num_tokens) { + in_vec = reinterpret_cast( + input + batch_id * stride_bz_input + head_id * stride_h_input + + token_id * stride_seq_input + + (threadIdx.x % NUM_THREADS_PER_TOKEN) * CVT_FP4_ELTS_PER_THREAD)[0]; + } + + struct alignas(16) SharedInputStorage { + T data[BLOCK_SIZE * head_dim]; + }; + __shared__ SharedInputStorage shared_input_storage; + T* shared_input = shared_input_storage.data; + reinterpret_cast(shared_input)[threadIdx.x] = in_vec; + __syncthreads(); +#pragma unroll + for (int i = 0; i < CVT_FP4_ELTS_PER_THREAD / 2; i++) { + in_vec.elts[i].x = + shared_input[(threadIdx.x / NUM_THREADS_PER_SEQ) + + ((threadIdx.x % NUM_THREADS_PER_SEQ) * CVT_FP4_ELTS_PER_THREAD + 2 * i) * + head_dim]; + in_vec.elts[i].y = + shared_input[(threadIdx.x / NUM_THREADS_PER_SEQ) + + ((threadIdx.x % NUM_THREADS_PER_SEQ) * CVT_FP4_ELTS_PER_THREAD + 2 * i + 1) * + head_dim]; + } + + auto localMax = __habs2(in_vec.elts[0]); +#pragma unroll + for (int i = 1; i < CVT_FP4_ELTS_PER_THREAD / 2; i++) { + localMax = __hmax2(localMax, __habs2(in_vec.elts[i])); + } + + if constexpr (CVT_FP4_ELTS_PER_THREAD == 8) { + localMax = __hmax2(__shfl_xor_sync(0xffffffff, localMax, 1, 32), localMax); + } + + float vecMax = float(__hmax(localMax.x, localMax.y)); + + float SFValue = vecMax / 6.0f; + uint8_t SFValueFP8; + reinterpret_cast<__nv_fp8_e4m3&>(SFValueFP8) = __nv_fp8_e4m3(SFValue); + SFValue = float(reinterpret_cast<__nv_fp8_e4m3&>(SFValueFP8)); + + float SFValueInv = (SFValue == 0.0f) ? 0.0f : 1.0f / SFValue; + + float2 fp2Vals[CVT_FP4_ELTS_PER_THREAD / 2]; + +#pragma unroll + for (int i = 0; i < CVT_FP4_ELTS_PER_THREAD / 2; i++) { + if constexpr (std::is_same::value) { + fp2Vals[i] = __half22float2(in_vec.elts[i]); + } else { + fp2Vals[i] = __bfloat1622float2(in_vec.elts[i]); + } + fp2Vals[i].x = fp2Vals[i].x * SFValueInv; + fp2Vals[i].y = fp2Vals[i].y * SFValueInv; + } + + uint32_t e2m1Vals[CVT_FP4_ELTS_PER_THREAD / 8]; +#pragma unroll + for (int i = 0; i < CVT_FP4_ELTS_PER_THREAD / 8; i++) { + e2m1Vals[i] = flashinfer::math::fp32_vec_to_e2m1(fp2Vals + i * 4); + } + + if constexpr (CVT_FP4_ELTS_PER_THREAD == 8) { + reinterpret_cast(output + batch_id * stride_bz_output + head_id * stride_h_output + + (threadIdx.x / NUM_THREADS_PER_SEQ) * stride_d_output + + (token_block_id * BLOCK_SIZE + + (threadIdx.x % NUM_THREADS_PER_SEQ) * CVT_FP4_ELTS_PER_THREAD) / + 2)[0] = e2m1Vals[0]; + } else { + reinterpret_cast(output + batch_id * stride_bz_output + head_id * stride_h_output + + (threadIdx.x / NUM_THREADS_PER_SEQ) * stride_d_output + + (token_block_id * BLOCK_SIZE + + (threadIdx.x % NUM_THREADS_PER_SEQ) * CVT_FP4_ELTS_PER_THREAD) / + 2)[0] = reinterpret_cast(e2m1Vals)[0]; + } + + uint8_t* output_sf_save_base = output_sf + batch_id * stride_bz_output_sf + + head_id * stride_h_output_sf + + (threadIdx.x / NUM_THREADS_PER_SEQ / 64) * 64 * stride_d_output_sf; + uint32_t row_id_local = (threadIdx.x / NUM_THREADS_PER_SEQ) % 64; + + if constexpr (CVT_FP4_ELTS_PER_THREAD == 16) { + uint32_t col_id_local = + token_block_id * BLOCK_SIZE / CVT_FP4_ELTS_PER_THREAD + threadIdx.x % NUM_THREADS_PER_SEQ; + uint32_t offset_local = (col_id_local / 4) * 256 + (col_id_local % 4) + + (row_id_local / 16) * 4 + (row_id_local % 16) * 16; + reinterpret_cast(output_sf_save_base + offset_local)[0] = SFValueFP8; + } else { + if (threadIdx.x % 2 == 0) { + uint32_t col_id_local = token_block_id * BLOCK_SIZE / CVT_FP4_ELTS_PER_THREAD + + (threadIdx.x % NUM_THREADS_PER_SEQ) / 2; + uint32_t offset_local = (col_id_local / 4) * 256 + (col_id_local % 4) + + (row_id_local / 16) * 4 + (row_id_local % 16) * 16; + reinterpret_cast(output_sf_save_base + offset_local)[0] = SFValueFP8; + } + } +} + +void scaled_fp4_quant(TensorView input, TensorView output, TensorView output_sf, + int64_t tensor_layout) { + constexpr int BLOCK_SIZE = 128; + + CHECK_QUANT_CUDA(input); + CHECK_QUANT_CUDA(output); + CHECK_QUANT_CUDA(output_sf); + + CHECK_QUANT_LASTDIM_CONTIGUOUS(input); + CHECK_QUANT_LASTDIM_CONTIGUOUS(output); + CHECK_QUANT_LASTDIM_CONTIGUOUS(output_sf); + + CHECK_QUANT_DTYPE(output, dl_uint8); + CHECK_QUANT_DTYPE(output_sf, dl_float8_e4m3fn); + + CHECK_QUANT_DIMS(input, 4); + CHECK_QUANT_DIMS(output, 4); + CHECK_QUANT_DIMS(output_sf, 4); + + const int batch_size = input.size(0); + const int head_dim = input.size(3); + + const int stride_bz_input = input.stride(0); + const int stride_bz_output = output.stride(0); + const int stride_bz_output_sf = output_sf.stride(0); + + int num_tokens, num_heads; + int stride_seq_input, stride_seq_output, stride_seq_output_sf; + int stride_h_input, stride_h_output, stride_h_output_sf; + if (tensor_layout == 0) { + num_tokens = input.size(1); + num_heads = input.size(2); + stride_seq_input = input.stride(1); + stride_seq_output = output.stride(1); + stride_seq_output_sf = output_sf.stride(1); + stride_h_input = input.stride(2); + stride_h_output = output.stride(2); + stride_h_output_sf = output_sf.stride(2); + + CHECK_QUANT_SHAPE(output, batch_size, num_tokens, num_heads, head_dim / 2); + CHECK_QUANT_SHAPE(output_sf, batch_size, num_tokens, num_heads, head_dim / 16); + } else { + num_tokens = input.size(2); + num_heads = input.size(1); + stride_seq_input = input.stride(2); + stride_seq_output = output.stride(2); + stride_seq_output_sf = output_sf.stride(2); + stride_h_input = input.stride(1); + stride_h_output = output.stride(1); + stride_h_output_sf = output_sf.stride(1); + + CHECK_QUANT_SHAPE(output, batch_size, num_heads, num_tokens, head_dim / 2); + CHECK_QUANT_SHAPE(output_sf, batch_size, num_heads, num_tokens, head_dim / 16); + } + + auto input_dtype = input.dtype(); + cudaStream_t stream = get_stream(input.device()); + + DISPATCH_DLPACK_DTYPE_TO_CTYPE_FP16_LOCAL(input_dtype, c_type, { + DISPATCH_HEAD_DIM(head_dim, HEAD_DIM, { + dim3 block(BLOCK_SIZE * HEAD_DIM / CVT_FP4_ELTS_PER_THREAD, 1, 1); + dim3 grid((num_tokens + BLOCK_SIZE - 1) / BLOCK_SIZE, batch_size, num_heads); + + scaled_fp4_quant_kernel<<>>( + reinterpret_cast(input.data_ptr()), + reinterpret_cast(output.data_ptr()), + reinterpret_cast(output_sf.data_ptr()), batch_size, num_heads, num_tokens, + stride_bz_input, stride_h_input, stride_seq_input, stride_bz_output, stride_h_output, + stride_seq_output, stride_bz_output_sf, stride_h_output_sf, stride_seq_output_sf); + }); + }); +} + +void scaled_fp4_quant_permute(TensorView input, TensorView output, TensorView output_sf, + int64_t tensor_layout) { + constexpr int BLOCK_SIZE = 128; + + CHECK_QUANT_CUDA(input); + CHECK_QUANT_CUDA(output); + CHECK_QUANT_CUDA(output_sf); + + CHECK_QUANT_LASTDIM_CONTIGUOUS(input); + CHECK_QUANT_LASTDIM_CONTIGUOUS(output); + CHECK_QUANT_LASTDIM_CONTIGUOUS(output_sf); + + CHECK_QUANT_DTYPE(output, dl_uint8); + CHECK_QUANT_DTYPE(output_sf, dl_float8_e4m3fn); + + CHECK_QUANT_DIMS(input, 4); + CHECK_QUANT_DIMS(output, 4); + CHECK_QUANT_DIMS(output_sf, 4); + + const int batch_size = input.size(0); + const int head_dim = input.size(3); + + const int stride_bz_input = input.stride(0); + const int stride_bz_output = output.stride(0); + const int stride_bz_output_sf = output_sf.stride(0); + + int num_tokens, num_heads; + int stride_seq_input, stride_seq_output, stride_seq_output_sf; + int stride_h_input, stride_h_output, stride_h_output_sf; + if (tensor_layout == 0) { + num_tokens = input.size(1); + num_heads = input.size(2); + stride_seq_input = input.stride(1); + stride_seq_output = output.stride(1); + stride_seq_output_sf = output_sf.stride(1); + stride_h_input = input.stride(2); + stride_h_output = output.stride(2); + stride_h_output_sf = output_sf.stride(2); + + CHECK_QUANT_SHAPE(output, batch_size, ((num_tokens + BLOCK_SIZE - 1) / BLOCK_SIZE) * BLOCK_SIZE, + num_heads, head_dim / 2); + CHECK_QUANT_SHAPE(output_sf, batch_size, + ((num_tokens + BLOCK_SIZE - 1) / BLOCK_SIZE) * BLOCK_SIZE, num_heads, + head_dim / 16); + } else { + num_tokens = input.size(2); + num_heads = input.size(1); + stride_seq_input = input.stride(2); + stride_seq_output = output.stride(2); + stride_seq_output_sf = output_sf.stride(2); + stride_h_input = input.stride(1); + stride_h_output = output.stride(1); + stride_h_output_sf = output_sf.stride(1); + + CHECK_QUANT_SHAPE(output, batch_size, num_heads, + ((num_tokens + BLOCK_SIZE - 1) / BLOCK_SIZE) * BLOCK_SIZE, head_dim / 2); + CHECK_QUANT_SHAPE(output_sf, batch_size, num_heads, + ((num_tokens + BLOCK_SIZE - 1) / BLOCK_SIZE) * BLOCK_SIZE, head_dim / 16); + } + + auto input_dtype = input.dtype(); + cudaStream_t stream = get_stream(input.device()); + + DISPATCH_DLPACK_DTYPE_TO_CTYPE_FP16_LOCAL(input_dtype, c_type, { + DISPATCH_HEAD_DIM(head_dim, HEAD_DIM, { + constexpr int BLOCK_SIZE = 128; + dim3 block(BLOCK_SIZE * HEAD_DIM / CVT_FP4_ELTS_PER_THREAD, 1, 1); + dim3 grid((num_tokens + BLOCK_SIZE - 1) / BLOCK_SIZE, batch_size, num_heads); + + scaled_fp4_quant_kernel<<>>( + reinterpret_cast(input.data_ptr()), + reinterpret_cast(output.data_ptr()), + reinterpret_cast(output_sf.data_ptr()), batch_size, num_heads, num_tokens, + stride_bz_input, stride_h_input, stride_seq_input, stride_bz_output, stride_h_output, + stride_seq_output, stride_bz_output_sf, stride_h_output_sf, stride_seq_output_sf); + }); + }); +} + +void scaled_fp4_quant_trans(TensorView input, TensorView output, TensorView output_sf, + int64_t tensor_layout) { + constexpr int BLOCK_SIZE = 128; + + CHECK_QUANT_CUDA(input); + CHECK_QUANT_CUDA(output); + CHECK_QUANT_CUDA(output_sf); + + CHECK_QUANT_LASTDIM_CONTIGUOUS(input); + CHECK_QUANT_LASTDIM_CONTIGUOUS(output); + CHECK_QUANT_LASTDIM_CONTIGUOUS(output_sf); + + CHECK_QUANT_DTYPE(output, dl_uint8); + CHECK_QUANT_DTYPE(output_sf, dl_float8_e4m3fn); + + CHECK_QUANT_DIMS(input, 4); + CHECK_QUANT_DIMS(output, 4); + CHECK_QUANT_DIMS(output_sf, 4); + + const int batch_size = input.size(0); + const int head_dim = input.size(3); + + const int stride_bz_input = input.stride(0); + const int stride_bz_output = output.stride(0); + const int stride_bz_output_sf = output_sf.stride(0); + + int num_tokens, num_heads; + int stride_seq_input; + int stride_d_output, stride_d_output_sf; + int stride_h_input, stride_h_output, stride_h_output_sf; + if (tensor_layout == 0) { + num_tokens = input.size(1); + num_heads = input.size(2); + stride_seq_input = input.stride(1); + stride_d_output = output.stride(1); + stride_d_output_sf = output_sf.stride(1); + stride_h_input = input.stride(2); + stride_h_output = output.stride(2); + stride_h_output_sf = output_sf.stride(2); + + CHECK_QUANT_SHAPE(output, batch_size, head_dim, num_heads, + ((num_tokens + BLOCK_SIZE - 1) / BLOCK_SIZE) * BLOCK_SIZE / 2); + CHECK_QUANT_SHAPE(output_sf, batch_size, head_dim, num_heads, + ((num_tokens + BLOCK_SIZE - 1) / BLOCK_SIZE) * BLOCK_SIZE / 16); + } else { + num_tokens = input.size(2); + num_heads = input.size(1); + stride_seq_input = input.stride(2); + stride_d_output = output.stride(2); + stride_d_output_sf = output_sf.stride(2); + stride_h_input = input.stride(1); + stride_h_output = output.stride(1); + stride_h_output_sf = output_sf.stride(1); + + CHECK_QUANT_SHAPE(output, batch_size, num_heads, head_dim, + ((num_tokens + BLOCK_SIZE - 1) / BLOCK_SIZE) * BLOCK_SIZE / 2); + CHECK_QUANT_SHAPE(output_sf, batch_size, num_heads, head_dim, + ((num_tokens + BLOCK_SIZE - 1) / BLOCK_SIZE) * BLOCK_SIZE / 16); + } + + auto input_dtype = input.dtype(); + cudaStream_t stream = get_stream(input.device()); + + DISPATCH_DLPACK_DTYPE_TO_CTYPE_FP16_LOCAL(input_dtype, c_type, { + DISPATCH_HEAD_DIM(head_dim, HEAD_DIM, { + dim3 block(BLOCK_SIZE * HEAD_DIM / CVT_FP4_ELTS_PER_THREAD, 1, 1); + dim3 grid((num_tokens + BLOCK_SIZE - 1) / BLOCK_SIZE, batch_size, num_heads); + + scaled_fp4_quant_trans_kernel<<>>( + reinterpret_cast(input.data_ptr()), + reinterpret_cast(output.data_ptr()), + reinterpret_cast(output_sf.data_ptr()), batch_size, num_heads, num_tokens, + stride_bz_input, stride_h_input, stride_seq_input, stride_bz_output, stride_h_output, + stride_d_output, stride_bz_output_sf, stride_h_output_sf, stride_d_output_sf); + }); + }); +} + +} // namespace nvfp4_attention_sm120 +} // namespace flashinfer + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(scaled_fp4_quant, + flashinfer::nvfp4_attention_sm120::scaled_fp4_quant); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(scaled_fp4_quant_permute, + flashinfer::nvfp4_attention_sm120::scaled_fp4_quant_permute); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(scaled_fp4_quant_trans, + flashinfer::nvfp4_attention_sm120::scaled_fp4_quant_trans); diff --git a/docs/api/attention.rst b/docs/api/attention.rst index 6492f80fe5c..803c6a96ac6 100644 --- a/docs/api/attention.rst +++ b/docs/api/attention.rst @@ -120,6 +120,18 @@ single kernel launch. .. automethod:: __init__ +SM120 NVFP4 Attention +--------------------- + +.. currentmodule:: flashinfer.nvfp4_attention_sm120 + +.. autosummary:: + :toctree: ../generated + + nvfp4_attention_sm120_quantize_qkv + nvfp4_attention_sm120_fwd + + flashinfer.mla ============== diff --git a/flashinfer/__init__.py b/flashinfer/__init__.py index 94136ba26cc..ddd862d7bb7 100644 --- a/flashinfer/__init__.py +++ b/flashinfer/__init__.py @@ -124,6 +124,13 @@ from .norm import rmsnorm_quant as rmsnorm_quant from .norm import fused_rmsnorm_silu as fused_rmsnorm_silu from .norm import fused_qk_rmsnorm_rope as fused_qk_rmsnorm_rope +from . import nvfp4_attention_sm120 as nvfp4_attention_sm120 +from .nvfp4_attention_sm120 import ( + nvfp4_attention_sm120_fwd as nvfp4_attention_sm120_fwd, +) +from .nvfp4_attention_sm120 import ( + nvfp4_attention_sm120_quantize_qkv as nvfp4_attention_sm120_quantize_qkv, +) from .norm import ( fused_dit_residual_layernorm_scale_shift as fused_dit_residual_layernorm_scale_shift, ) diff --git a/flashinfer/aot.py b/flashinfer/aot.py index f762cd13ec3..900958ad4e1 100644 --- a/flashinfer/aot.py +++ b/flashinfer/aot.py @@ -59,6 +59,7 @@ ) from .jit.fp4_kv_dequantization import gen_fp4_kv_dequantization_module from .jit.fp4_kv_quantization import gen_fp4_kv_quantization_module +from .jit.nvfp4_attention_sm120 import gen_nvfp4_attention_sm120_module from .jit.fp8_quantization import gen_mxfp8_quantization_sm100_module from .jit.fused_moe import ( gen_cutlass_fused_moe_sm90_module, @@ -500,6 +501,8 @@ def gen_all_modules( add_oai_oss, ) ) + if has_sm120: + jit_specs.append(gen_nvfp4_attention_sm120_module()) if add_act: for act_name in act_func_def_str: diff --git a/flashinfer/jit/__init__.py b/flashinfer/jit/__init__.py index fe6de6db746..46021b7b768 100644 --- a/flashinfer/jit/__init__.py +++ b/flashinfer/jit/__init__.py @@ -97,6 +97,9 @@ from .fp4_kv_quantization import ( gen_fp4_kv_quantization_module as gen_fp4_kv_quantization_module, ) +from .nvfp4_attention_sm120 import ( + gen_nvfp4_attention_sm120_module as gen_nvfp4_attention_sm120_module, +) from .bgmv_moe import gen_bgmv_moe_module as gen_bgmv_moe_module from .bgmv_moe import load_bgmv_moe_module as load_bgmv_moe_module diff --git a/flashinfer/jit/nvfp4_attention_sm120.py b/flashinfer/jit/nvfp4_attention_sm120.py new file mode 100644 index 00000000000..1a2fdd6ee06 --- /dev/null +++ b/flashinfer/jit/nvfp4_attention_sm120.py @@ -0,0 +1,97 @@ +""" +Copyright (c) 2026 by FlashInfer team. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +import functools +from pathlib import Path + +from . import env as jit_env +from .core import JitSpec, gen_jit_spec, sm120a_nvcc_flags + + +_NVFP4_ATTENTION_SM120_MODULE_NAME = "nvfp4_attention_sm120" + +_NVFP4_ATTENTION_SM120_SOURCE_FILES = ( + "nvfp4_attention_sm120/nvfp4_attention_sm120_binding.cu", + "nvfp4_attention_sm120/nvfp4_attention_sm120_quantize.cu", +) + + +_NVFP4_ATTENTION_SM120_CUDA_FLAGS = [ + "-DFLASHINFER_ENABLE_F16", + "-DFLASHINFER_ENABLE_BF16", + "-U__CUDA_NO_HALF_OPERATORS__", + "-U__CUDA_NO_HALF_CONVERSIONS__", + "-U__CUDA_NO_BFLOAT16_OPERATORS__", + "-U__CUDA_NO_BFLOAT16_CONVERSIONS__", + "-U__CUDA_NO_BFLOAT162_OPERATORS__", + "-U__CUDA_NO_BFLOAT162_CONVERSIONS__", + "-U__CUDA_NO_NVFP4_OPERATORS__", + "-U__CUDA_NO_NVFP4_CONVERSIONS__", + "-DCUTLASS_DEBUG_TRACE_LEVEL=0", + "-DNDEBUG", + "-DQBLKSIZE=128", + "-DKBLKSIZE=128", + "-DCTA256", + "-DDQINRMEM", + "-DPINGPONG_MATH_ORDER", + "-DPINGPONG_EARLY_RELEASE_K", + "-DCAUSAL_DISABLE_QK_ORDER", + "--expt-relaxed-constexpr", + "--expt-extended-lambda", + "-lineinfo", +] + + +def _nvfp4_attention_sm120_source_path(source_file: str) -> Path: + package_data_path = jit_env.FLASHINFER_CSRC_DIR / source_file + if package_data_path.exists(): + return package_data_path + return _repo_root() / "csrc" / source_file + + +def _repo_root() -> Path: + return Path(__file__).resolve().parents[2] + + +def _nvfp4_attention_sm120_include_paths() -> list[Path]: + root = _repo_root() + candidates = [ + root / "include", + root / "csrc", + root / "3rdparty" / "cccl" / "cub", + root / "3rdparty" / "cccl" / "libcudacxx" / "include", + root / "3rdparty" / "cccl" / "thrust", + root / "3rdparty" / "cutlass" / "include", + root / "3rdparty" / "cutlass" / "tools" / "util" / "include", + root / "3rdparty" / "spdlog" / "include", + ] + return [path for path in candidates if path.exists()] + + +@functools.cache +def gen_nvfp4_attention_sm120_module() -> JitSpec: + source_paths = [ + _nvfp4_attention_sm120_source_path(source_file) + for source_file in _NVFP4_ATTENTION_SM120_SOURCE_FILES + ] + include_paths: list[str | Path] = [] + include_paths.extend(_nvfp4_attention_sm120_include_paths()) + return gen_jit_spec( + _NVFP4_ATTENTION_SM120_MODULE_NAME, + source_paths, + extra_cuda_cflags=sm120a_nvcc_flags + _NVFP4_ATTENTION_SM120_CUDA_FLAGS, + extra_include_paths=include_paths, + ) diff --git a/flashinfer/nvfp4_attention_sm120.py b/flashinfer/nvfp4_attention_sm120.py new file mode 100644 index 00000000000..46cdc20bc9a --- /dev/null +++ b/flashinfer/nvfp4_attention_sm120.py @@ -0,0 +1,421 @@ +""" +Copyright (c) 2026 by FlashInfer team. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +import functools +from typing import Optional, Tuple + +import torch + +from .api_logging import flashinfer_api +from .jit.nvfp4_attention_sm120 import gen_nvfp4_attention_sm120_module +from .trace.templates.nvfp4_attention_sm120 import ( + nvfp4_attention_sm120_fwd_trace, + nvfp4_attention_sm120_quantize_qkv_trace, +) +from .utils import supported_compute_capability + + +_TOKEN_BLOCK_SIZE = 128 +_SUPPORTED_HEAD_DIMS = (64, 128) +_SUPPORTED_QKV_DTYPES = (torch.float16, torch.bfloat16) +_SUPPORTED_OUT_DTYPES = (torch.float16, torch.bfloat16) + +_HND_LAYOUT = 1 + + +@functools.cache +def get_nvfp4_attention_sm120_module(): + return gen_nvfp4_attention_sm120_module().build_and_load() + + +def _check_cuda_contiguous(name: str, tensor: torch.Tensor) -> None: + if not tensor.is_cuda: + raise ValueError(f"{name} must be a CUDA tensor, got device={tensor.device}") + if not tensor.is_contiguous(): + raise ValueError(f"{name} must be contiguous, got stride={tensor.stride()}") + + +def _check_same_device( + name: str, + tensor: torch.Tensor, + ref_name: str, + ref: torch.Tensor, +) -> None: + if tensor.device != ref.device: + raise ValueError( + f"{name} must be on the same device as {ref_name}, " + f"got {tensor.device} and {ref.device}" + ) + + +def _pad_seq_len_to_128(x: torch.Tensor) -> torch.Tensor: + pad_len = (-x.shape[2]) % _TOKEN_BLOCK_SIZE + if pad_len == 0: + return x.contiguous() + return torch.nn.functional.pad(x, (0, 0, 0, pad_len), value=0).contiguous() + + +def _preprocess_qkv( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + per_block_mean: bool, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + for name, tensor in (("q", q), ("k", k), ("v", v)): + _check_cuda_contiguous(name, tensor) + if tensor.dtype not in _SUPPORTED_QKV_DTYPES: + raise ValueError( + f"{name} must have dtype torch.float16 or torch.bfloat16, " + f"got {tensor.dtype}" + ) + if tensor.ndim != 4: + raise ValueError( + f"{name} must have shape [batch, num_heads, seq_len, head_dim], " + f"got shape={tuple(tensor.shape)}" + ) + if q.shape != k.shape or q.shape != v.shape: + raise ValueError( + "q, k, and v must have the same shape, " + f"got q={tuple(q.shape)}, k={tuple(k.shape)}, v={tuple(v.shape)}" + ) + if q.dtype != k.dtype or q.dtype != v.dtype: + raise ValueError( + "q, k, and v must have the same dtype, " + f"got q={q.dtype}, k={k.dtype}, v={v.dtype}" + ) + + _check_same_device("k", k, "q", q) + _check_same_device("v", v, "q", q) + + batch, num_heads, seq_len, head_dim = q.shape + if head_dim not in _SUPPORTED_HEAD_DIMS: + raise ValueError(f"head_dim must be 64 or 128, got {head_dim}") + + k = k - k.mean(dim=-2, keepdim=True) + q, k, v = map(_pad_seq_len_to_128, (q, k, v)) + seq_len = q.shape[2] + + if per_block_mean: + num_groups = seq_len // _TOKEN_BLOCK_SIZE + q_grouped = q.reshape(batch, num_heads, num_groups, _TOKEN_BLOCK_SIZE, head_dim) + qm = q_grouped.mean(dim=3) + q = ( + (q_grouped - qm.unsqueeze(3)) + .reshape(batch, num_heads, seq_len, head_dim) + .contiguous() + ) + else: + qm = q.mean(dim=-2, keepdim=True) + q = (q - qm).contiguous() + + qk_correction = torch.matmul(qm, k.transpose(-2, -1)).to(torch.float32) + if per_block_mean: + qk_correction = qk_correction.repeat_interleave(_TOKEN_BLOCK_SIZE, dim=2) + else: + qk_correction = qk_correction.expand(-1, -1, _TOKEN_BLOCK_SIZE, -1) + qk_correction = qk_correction.contiguous() + return q.contiguous(), k.contiguous(), v.contiguous(), qk_correction + + +@supported_compute_capability([120]) +@flashinfer_api(trace=nvfp4_attention_sm120_quantize_qkv_trace) +def nvfp4_attention_sm120_quantize_qkv( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + per_block_mean: bool = True, +) -> Tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, +]: + r"""Preprocess and quantize dense Q/K/V tensors for SM120 NVFP4 attention. + + The input layout is ``[batch, num_heads, seq_len, head_dim]``. Inputs must be + contiguous CUDA tensors with the same shape, dtype, and device. The sequence + dimension is padded to a multiple of 128 before Q/K/V are quantized. + + Parameters + ---------- + q, k, v : torch.Tensor + Dense Q/K/V tensors with dtype ``torch.float16`` or ``torch.bfloat16``. + per_block_mean : bool, optional + Whether to center Q per 128-token block. When ``False``, Q is centered + once across the full sequence. + + Returns + ------- + Tuple[torch.Tensor, ...] + ``q_fp4``, ``k_fp4``, transposed ``v_fp4_t``, scale tensors + ``q_scale``, ``k_scale``, ``v_scale_t``, and the expanded FP32 QK correction. + """ + q_proc, k_proc, v_proc, qk_correction = _preprocess_qkv(q, k, v, per_block_mean) + batch, num_heads, seq_len, head_dim = q_proc.shape + + q_fp4 = torch.empty( + (batch, num_heads, seq_len, head_dim // 2), device=q.device, dtype=torch.uint8 + ) + k_fp4 = torch.empty_like(q_fp4) + v_fp4_t = torch.empty( + (batch, num_heads, head_dim, seq_len // 2), device=q.device, dtype=torch.uint8 + ) + q_scale = torch.empty( + (batch, num_heads, seq_len, head_dim // 16), + device=q.device, + dtype=torch.float8_e4m3fn, + ) + k_scale = torch.empty_like(q_scale) + v_scale_t = torch.empty( + (batch, num_heads, head_dim, seq_len // 16), + device=q.device, + dtype=torch.float8_e4m3fn, + ) + + module = get_nvfp4_attention_sm120_module() + module.scaled_fp4_quant(q_proc, q_fp4, q_scale, _HND_LAYOUT) + module.scaled_fp4_quant_permute(k_proc, k_fp4, k_scale, _HND_LAYOUT) + module.scaled_fp4_quant_trans(v_proc, v_fp4_t, v_scale_t, _HND_LAYOUT) + + return q_fp4, k_fp4, v_fp4_t, q_scale, k_scale, v_scale_t, qk_correction + + +def _check_inputs( + q_fp4: torch.Tensor, + k_fp4: torch.Tensor, + v_fp4_t: torch.Tensor, + q_scale: torch.Tensor, + k_scale: torch.Tensor, + v_scale_t: torch.Tensor, + qk_correction: torch.Tensor, + per_block_mean: bool, +) -> Tuple[int, int, int, int]: + for name, tensor in ( + ("q_fp4", q_fp4), + ("k_fp4", k_fp4), + ("v_fp4_t", v_fp4_t), + ("q_scale", q_scale), + ("k_scale", k_scale), + ("v_scale_t", v_scale_t), + ("qk_correction", qk_correction), + ): + _check_cuda_contiguous(name, tensor) + + for name, tensor in ( + ("k_fp4", k_fp4), + ("v_fp4_t", v_fp4_t), + ("q_scale", q_scale), + ("k_scale", k_scale), + ("v_scale_t", v_scale_t), + ("qk_correction", qk_correction), + ): + _check_same_device(name, tensor, "q_fp4", q_fp4) + + if ( + q_fp4.dtype != torch.uint8 + or k_fp4.dtype != torch.uint8 + or v_fp4_t.dtype != torch.uint8 + ): + raise ValueError("q_fp4, k_fp4, and v_fp4_t must be uint8 packed FP4 tensors") + if q_scale.dtype != torch.float8_e4m3fn or k_scale.dtype != torch.float8_e4m3fn: + raise ValueError("q_scale and k_scale must be torch.float8_e4m3fn tensors") + if v_scale_t.dtype != torch.float8_e4m3fn: + raise ValueError("v_scale_t must be a torch.float8_e4m3fn tensor") + if qk_correction.dtype != torch.float32: + raise ValueError("qk_correction must be a torch.float32 tensor") + + if q_fp4.ndim != 4: + raise ValueError( + "q_fp4 must have shape [batch, num_heads, seq_len, head_dim / 2]" + ) + if k_fp4.shape != q_fp4.shape: + raise ValueError( + f"k_fp4 shape {tuple(k_fp4.shape)} must match q_fp4 {tuple(q_fp4.shape)}" + ) + + batch, num_heads, seq_len, packed_head_dim = q_fp4.shape + head_dim = packed_head_dim * 2 + if head_dim not in _SUPPORTED_HEAD_DIMS: + raise ValueError(f"head_dim must be 64 or 128, got {head_dim}") + if seq_len % _TOKEN_BLOCK_SIZE != 0: + raise ValueError(f"seq_len must be padded to a multiple of 128, got {seq_len}") + if head_dim % 16 != 0: + raise ValueError(f"head_dim must be divisible by 16, got {head_dim}") + + expected_v = (batch, num_heads, head_dim, seq_len // 2) + if tuple(v_fp4_t.shape) != expected_v: + raise ValueError(f"v_fp4_t shape {tuple(v_fp4_t.shape)} must be {expected_v}") + + expected_sf_qk = (batch, num_heads, seq_len, head_dim // 16) + if tuple(q_scale.shape) != expected_sf_qk: + raise ValueError( + f"q_scale shape {tuple(q_scale.shape)} must be {expected_sf_qk}" + ) + if tuple(k_scale.shape) != expected_sf_qk: + raise ValueError( + f"k_scale shape {tuple(k_scale.shape)} must be {expected_sf_qk}" + ) + + expected_v_scale = (batch, num_heads, head_dim, seq_len // 16) + if tuple(v_scale_t.shape) != expected_v_scale: + raise ValueError( + f"v_scale_t shape {tuple(v_scale_t.shape)} must be {expected_v_scale}" + ) + + if ( + qk_correction.ndim != 4 + or qk_correction.shape[0] != batch + or qk_correction.shape[1] != num_heads + ): + raise ValueError( + "qk_correction must have shape [batch, num_heads, seq_len_s, seq_len]" + ) + expected_delta_groups = seq_len if per_block_mean else _TOKEN_BLOCK_SIZE + if qk_correction.shape[2] != expected_delta_groups: + raise ValueError( + f"qk_correction seq_len_s dimension must be {expected_delta_groups}, " + f"got {qk_correction.shape[2]}" + ) + if qk_correction.shape[-1] != seq_len: + raise ValueError( + f"qk_correction last dimension must be {seq_len}, got {qk_correction.shape[-1]}" + ) + + return batch, num_heads, seq_len, head_dim + + +@supported_compute_capability([120]) +@flashinfer_api(trace=nvfp4_attention_sm120_fwd_trace) +def nvfp4_attention_sm120_fwd( + q_fp4: torch.Tensor, + k_fp4: torch.Tensor, + v_fp4_t: torch.Tensor, + q_scale: torch.Tensor, + k_scale: torch.Tensor, + v_scale_t: torch.Tensor, + qk_correction: torch.Tensor, + sm_scale: Optional[float] = None, + causal: bool = False, + per_block_mean: bool = True, + out: Optional[torch.Tensor] = None, + lse: Optional[torch.Tensor] = None, + out_dtype: torch.dtype = torch.bfloat16, + softmax_scale: Optional[float] = None, +) -> Tuple[torch.Tensor, torch.Tensor]: + r"""Run SM120 NVFP4 attention on pre-quantized Q/K/V tensors. + + The packed tensors should be produced by + :func:`nvfp4_attention_sm120_quantize_qkv`. ``q_fp4`` and ``k_fp4`` use layout + ``[batch, num_heads, seq_len, head_dim / 2]``; ``v_fp4_t`` and ``v_scale_t`` are + stored transposed as ``[batch, num_heads, head_dim, packed_seq_len]``. + + Parameters + ---------- + q_fp4, k_fp4, v_fp4_t : torch.Tensor + Packed NVFP4 Q/K/V tensors. + q_scale, k_scale, v_scale_t : torch.Tensor + Per-vector FP8 scale factors for Q/K/V. + qk_correction : torch.Tensor + FP32 correction term returned by ``nvfp4_attention_sm120_quantize_qkv``. + sm_scale : Optional[float], optional + Scale applied to QK scores before softmax. Defaults to + ``1 / sqrt(head_dim)`` when omitted. + causal : bool, optional + Whether to apply a causal mask. + per_block_mean : bool, optional + Must match the value used by ``nvfp4_attention_sm120_quantize_qkv``. + out, lse : Optional[torch.Tensor], optional + Optional output and log-sum-exp buffers. + out_dtype : torch.dtype, optional + Output dtype used when ``out`` is not provided. + softmax_scale : Optional[float], optional + Deprecated alias for ``sm_scale``. + + Returns + ------- + Tuple[torch.Tensor, torch.Tensor] + Attention output and log-sum-exp tensor. + """ + per_block_mean = bool(per_block_mean) + batch, num_heads, seq_len, head_dim = _check_inputs( + q_fp4, + k_fp4, + v_fp4_t, + q_scale, + k_scale, + v_scale_t, + qk_correction, + per_block_mean, + ) + if sm_scale is not None and softmax_scale is not None: + raise ValueError("Specify only one of sm_scale or softmax_scale") + if sm_scale is None: + sm_scale = head_dim**-0.5 if softmax_scale is None else softmax_scale + + if out is None: + if out_dtype not in _SUPPORTED_OUT_DTYPES: + raise ValueError( + f"out_dtype must be torch.float16 or torch.bfloat16, got {out_dtype}" + ) + out = torch.empty( + (batch, num_heads, seq_len, head_dim), + device=q_fp4.device, + dtype=out_dtype, + ) + else: + _check_cuda_contiguous("out", out) + _check_same_device("out", out, "q_fp4", q_fp4) + if tuple(out.shape) != (batch, num_heads, seq_len, head_dim): + raise ValueError( + f"out shape {tuple(out.shape)} must be {(batch, num_heads, seq_len, head_dim)}" + ) + if out.dtype not in _SUPPORTED_OUT_DTYPES: + raise ValueError( + f"out must have dtype torch.float16 or torch.bfloat16, got {out.dtype}" + ) + + if lse is None: + lse = torch.empty( + (batch, num_heads, seq_len), device=q_fp4.device, dtype=torch.float32 + ) + else: + _check_cuda_contiguous("lse", lse) + _check_same_device("lse", lse, "q_fp4", q_fp4) + if tuple(lse.shape) != (batch, num_heads, seq_len): + raise ValueError( + f"lse shape {tuple(lse.shape)} must be {(batch, num_heads, seq_len)}" + ) + if lse.dtype != torch.float32: + raise ValueError(f"lse must have dtype torch.float32, got {lse.dtype}") + + get_nvfp4_attention_sm120_module().fwd( + q_fp4, + k_fp4, + v_fp4_t, + q_scale, + k_scale, + v_scale_t, + qk_correction, + out, + lse, + float(sm_scale), + bool(causal), + per_block_mean, + ) + return out, lse diff --git a/flashinfer/trace/templates/nvfp4_attention_sm120.py b/flashinfer/trace/templates/nvfp4_attention_sm120.py new file mode 100644 index 00000000000..6334aab34bb --- /dev/null +++ b/flashinfer/trace/templates/nvfp4_attention_sm120.py @@ -0,0 +1,92 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from ..template import Const, Scalar, Tensor, TraceTemplate, Var + + +nvfp4_attention_sm120_quantize_qkv_trace = TraceTemplate( + op_type="attention", + name_prefix="nvfp4_attention_sm120_quantize_qkv", + description="Preprocess and quantize dense Q/K/V tensors for the SM120 NVFP4 attention kernel.", + axes={ + "batch_size": Var(), + "num_heads": Var(), + "seq_len": Var(), + "head_dim": Const(), + "packed_head_dim": Var(), + "scale_head_dim": Var(), + "packed_seq_len": Var(), + "scale_seq_len": Var(), + "correction_seq_len": Var(), + }, + inputs={ + "q": Tensor(["batch_size", "num_heads", "seq_len", "head_dim"]), + "k": Tensor(["batch_size", "num_heads", "seq_len", "head_dim"]), + "v": Tensor(["batch_size", "num_heads", "seq_len", "head_dim"]), + "per_block_mean": Scalar("bool"), + }, + outputs={ + "q_fp4": Tensor(["batch_size", "num_heads", "seq_len", "packed_head_dim"]), + "k_fp4": Tensor(["batch_size", "num_heads", "seq_len", "packed_head_dim"]), + "v_fp4_t": Tensor(["batch_size", "num_heads", "head_dim", "packed_seq_len"]), + "q_scale": Tensor(["batch_size", "num_heads", "seq_len", "scale_head_dim"]), + "k_scale": Tensor(["batch_size", "num_heads", "seq_len", "scale_head_dim"]), + "v_scale_t": Tensor(["batch_size", "num_heads", "head_dim", "scale_seq_len"]), + "qk_correction": Tensor( + ["batch_size", "num_heads", "correction_seq_len", "seq_len"] + ), + }, + constraints=[ + "head_dim == 2 * packed_head_dim", + "head_dim == 16 * scale_head_dim", + ], + tags=["sm120", "nvfp4"], +) + + +nvfp4_attention_sm120_fwd_trace = TraceTemplate( + op_type="attention", + name_prefix="nvfp4_attention_sm120_fwd", + description="Run the SM120 NVFP4 attention forward kernel on pre-quantized Q/K/V tensors.", + axes={ + "batch_size": Var(), + "num_heads": Var(), + "seq_len": Var(), + "head_dim": Const(), + "packed_head_dim": Const(), + "scale_head_dim": Const(), + "packed_seq_len": Var(), + "scale_seq_len": Var(), + "correction_seq_len": Var(), + }, + inputs={ + "q_fp4": Tensor(["batch_size", "num_heads", "seq_len", "packed_head_dim"]), + "k_fp4": Tensor(["batch_size", "num_heads", "seq_len", "packed_head_dim"]), + "v_fp4_t": Tensor(["batch_size", "num_heads", "head_dim", "packed_seq_len"]), + "q_scale": Tensor(["batch_size", "num_heads", "seq_len", "scale_head_dim"]), + "k_scale": Tensor(["batch_size", "num_heads", "seq_len", "scale_head_dim"]), + "v_scale_t": Tensor(["batch_size", "num_heads", "head_dim", "scale_seq_len"]), + "qk_correction": Tensor( + ["batch_size", "num_heads", "correction_seq_len", "seq_len"] + ), + "sm_scale": Scalar("float32"), + "causal": Scalar("bool"), + "per_block_mean": Scalar("bool"), + }, + outputs={ + "out": Tensor(["batch_size", "num_heads", "seq_len", "head_dim"]), + "lse": Tensor(["batch_size", "num_heads", "seq_len"]), + }, + tags=["sm120", "nvfp4"], +) diff --git a/include/flashinfer/attention/sm120/nvfp4_attention_sm120/api/launcher.h b/include/flashinfer/attention/sm120/nvfp4_attention_sm120/api/launcher.h new file mode 100644 index 00000000000..3c302294ba6 --- /dev/null +++ b/include/flashinfer/attention/sm120/nvfp4_attention_sm120/api/launcher.h @@ -0,0 +1,146 @@ +/* + * Copyright (c) 2025 by SageAttention team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include + +#include "../common/params.h" +#include "../common/static_switch.h" +#include "../compute/epilogue/lse_writer.cuh" +#include "../compute/epilogue/output_writer.cuh" +#include "../compute/producer/load_k.cuh" +#include "../compute/producer/load_q.cuh" +#include "../compute/producer/load_v.cuh" +#include "../kernel/attention_kernel.h" +#include "../kernel/scheduler.h" +#include "../kernel/traits.h" +#include "cute/tensor.hpp" +#include "cutlass/cluster_launch.hpp" +#include "flashinfer/utils.cuh" + +namespace nvfp4_attention { + +template +void run_flash_fwd(Flash_fwd_params& params, cudaStream_t stream) { + using Element = typename Kernel_traits::Element; + using ElementSF = typename Kernel_traits::ElementSF; + using ElementOut = typename Kernel_traits::ElementOut; + using ElementDS = typename Kernel_traits::ElementDS; + using TileShape_MNK = typename Kernel_traits::TileShape_MNK; + using ClusterShape = typename Kernel_traits::ClusterShape_MNK; + + using CollectiveMainloop = nvfp4_attention::CollectiveMainloopFwd; + using CollectiveEpilogue = nvfp4_attention::CollectiveEpilogueFwd; + + using Scheduler = nvfp4_attention::StaticPersistentTileScheduler; + + typename CollectiveMainloop::Params mainloop_params = CollectiveMainloop::to_underlying_arguments( + {static_cast(params.q_ptr), + {params.seqlen_q, params.d, params.h, params.b}, + {params.q_row_stride, _1{}, params.q_head_stride, params.q_batch_stride}, + + static_cast(params.k_ptr), + {params.seqlen_k, params.d, params.h_k, params.b}, + {params.k_row_stride, _1{}, params.k_head_stride, params.k_batch_stride}, + {params.unpadded_seqlen_k, params.d, params.h_k, params.b}, + + static_cast(params.v_ptr), + {params.d, params.seqlen_k, params.h_k, params.b}, + {params.v_row_stride, _1{}, params.v_head_stride, params.v_batch_stride}, + + static_cast(params.sfq_ptr), + {params.seqlen_q, params.d, params.h, params.b}, + static_cast(params.sfk_ptr), + {params.seqlen_k, params.d, params.h_k, params.b}, + static_cast(params.sfv_ptr), + {params.d, params.seqlen_k, params.h_k, params.b}, + + static_cast(params.delta_s_ptr), + {params.seqlen_s, params.seqlen_k, params.h_k, params.b}, + {params.ds_row_stride, _1{}, params.ds_head_stride, params.ds_batch_stride}, + + params.scale_softmax_log2}); + + typename CollectiveEpilogue::Params epilogue_params = + CollectiveEpilogue::to_underlying_arguments({ + + static_cast(params.o_ptr), + {params.seqlen_q, params.d, params.h, params.b}, + {params.o_row_stride, _1{}, params.o_head_stride, params.o_batch_stride}, + + static_cast(params.softmax_lse_ptr), + {_1{}, params.seqlen_q, params.h * params.seqlen_q}, + }); + + int num_blocks_m = cutlass::ceil_div(params.seqlen_q, Kernel_traits::kBlockM); + num_blocks_m = cutlass::ceil_div(num_blocks_m, size<0>(ClusterShape{})) * size<0>(ClusterShape{}); + + typename Scheduler::Arguments scheduler_args = {num_blocks_m, params.h, params.b, nullptr, + params.is_causal}; + typename Scheduler::Params scheduler_params = Scheduler::to_underlying_arguments(scheduler_args); + + void* kernel = (void*)nvfp4_attention::attention_kernel_ws; + + int smem_size = sizeof(typename Kernel_traits::SharedStorage); + if (smem_size >= 48 * 1024) { + FLASHINFER_CUDA_CHECK( + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + } + + static constexpr int ctaSize = Kernel_traits::kNWarps * 32; + + params.m_block_divmod = cutlass::FastDivmod(num_blocks_m); + params.total_blocks = num_blocks_m * params.h * params.b; + + int device_id = 0; + FLASHINFER_CUDA_CHECK(cudaGetDevice(&device_id)); + int num_sms = 0; + FLASHINFER_CUDA_CHECK( + cudaDeviceGetAttribute(&num_sms, cudaDevAttrMultiProcessorCount, device_id)); + dim3 grid_dims = Scheduler::get_grid_dim(scheduler_args, num_sms); + dim3 block_dims(ctaSize); + dim3 cluster_dims(size<0>(ClusterShape{}), size<1>(ClusterShape{}), size<2>(ClusterShape{})); + + cutlass::ClusterLaunchParams launch_params{grid_dims, block_dims, cluster_dims, smem_size, + stream}; + + cutlass::launch_kernel_on_cluster(launch_params, kernel, params, mainloop_params, epilogue_params, + scheduler_params); + + FLASHINFER_CUDA_CHECK(cudaGetLastError()); +} + +template +void run_mha_fwd_(Flash_fwd_params& params, cudaStream_t stream) { + using DeltaSType = float; + + BOOL_SWITCH(params.is_causal, Is_causal, [&] { + BOOL_SWITCH(params.per_block_mean, per_block, [&] { + if constexpr (Headdim == 64 || Headdim == 128) { + static constexpr int kStages = 3; + static constexpr int kBlockN = 128; + run_flash_fwd< + Flash_fwd_kernel_traits, + Is_causal>(params, stream); + } else { + static_assert(Headdim == 64 || Headdim == 128, "Unsupported Headdim"); + } + }); + }); +} + +} // namespace nvfp4_attention diff --git a/include/flashinfer/attention/sm120/nvfp4_attention_sm120/common/block_info.h b/include/flashinfer/attention/sm120/nvfp4_attention_sm120/common/block_info.h new file mode 100644 index 00000000000..b3ea68bd705 --- /dev/null +++ b/include/flashinfer/attention/sm120/nvfp4_attention_sm120/common/block_info.h @@ -0,0 +1,67 @@ +/* + * Copyright (c) 2025 by SageAttention team. + * + * This code is based on code from FlashAttention3, https://github.com/Dao-AILab/flash-attention + * Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri + * Dao. Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file + * except in compliance with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include + +namespace flash { + +template +struct BlockInfo { + template + __device__ BlockInfo(const Params& params, const int bidb) + : sum_s_q(!Varlen || params.cu_seqlens_q == nullptr ? -1 : params.cu_seqlens_q[bidb]), + sum_s_k(!Varlen || params.cu_seqlens_k == nullptr || !params.is_seqlens_k_cumulative + ? -1 + : params.cu_seqlens_k[bidb]), + actual_seqlen_q(!Varlen || params.cu_seqlens_q == nullptr + ? params.seqlen_q + : params.cu_seqlens_q[bidb + 1] - sum_s_q) + + , + seqlen_k_cache(!Varlen || params.cu_seqlens_k == nullptr + ? params.seqlen_k + : (params.is_seqlens_k_cumulative + ? params.cu_seqlens_k[bidb + 1] - sum_s_k + : params.cu_seqlens_k[bidb])), + actual_seqlen_k(params.seqused_k + ? params.seqused_k[bidb] + : seqlen_k_cache + + (params.knew_ptr == nullptr ? 0 : params.seqlen_knew)) {} + + template + __forceinline__ __device__ index_t q_offset(const index_t batch_stride, const index_t row_stride, + const int bidb) const { + return sum_s_q == -1 ? bidb * batch_stride : uint32_t(sum_s_q) * row_stride; + } + + template + __forceinline__ __device__ index_t k_offset(const index_t batch_stride, const index_t row_stride, + const int bidb) const { + return sum_s_k == -1 ? bidb * batch_stride : uint32_t(sum_s_k) * row_stride; + } + + const int sum_s_q; + const int sum_s_k; + const int actual_seqlen_q; + + const int seqlen_k_cache; + const int actual_seqlen_k; +}; + +} // namespace flash diff --git a/include/flashinfer/attention/sm120/nvfp4_attention_sm120/common/cute_extension.h b/include/flashinfer/attention/sm120/nvfp4_attention_sm120/common/cute_extension.h new file mode 100644 index 00000000000..7a5f0671acb --- /dev/null +++ b/include/flashinfer/attention/sm120/nvfp4_attention_sm120/common/cute_extension.h @@ -0,0 +1,511 @@ +/* + * Copyright (c) 2025 by SageAttention team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "cute/arch/mma_sm120.hpp" +#include "cute/atom/mma_atom.hpp" +#include "cute/atom/mma_traits_sm120.hpp" +#include "cutlass/cutlass.h" +#include "cutlass/float8.h" +#include "cutlass/float_subbyte.h" + +namespace cute::SM120::BLOCKSCALED { + +using cutlass::float_e2m1_t; +using cutlass::float_ue4m3_t; + +template +CUTE_HOST_DEVICE static void fma_with_interleave( + float& d0, float& d1, float& d2, float& d3, float& d4, float& d5, float& d6, float& d7, + float& d8, float& d9, float& d10, float& d11, float& d12, float& d13, float& d14, float& d15, + uint32_t const& a0, uint32_t const& a1, uint32_t const& a2, uint32_t const& a3, + uint32_t const& b0, uint32_t const& b1, uint32_t const& b2, uint32_t const& b3, + uint32_t const& b4, uint32_t const& b5, uint32_t const& b6, uint32_t const& b7, float const& c0, + float const& c1, float const& c2, float const& c3, float const& c4, float const& c5, + float const& c6, float const& c7, float const& c8, float const& c9, float const& c10, + float const& c11, float const& c12, float const& c13, float const& c14, float const& c15, + cute::uint_bit_t<32> const& sfa0, cute::uint_bit_t<32> const& sfb0, GapFn&& gap_fn) { +#if defined(CUTE_ARCH_MXF4NVF4_4X_UE4M3_MMA_ENABLED) + static constexpr uint16_t tidA = 0, bidA = 0, bidB = 0; + static constexpr uint16_t tidB0 = 0, tidB1 = 1, tidB2 = 2, tidB3 = 3; + + asm volatile( + "mma.sync.aligned.kind::mxf4nvf4.block_scale.scale_vec::4X.m16n8k64.row.col.f32.e2m1.e2m1." + "f32.ue4m3 " + "{%0,%1,%2,%3},{%4,%5,%6,%7},{%8,%9},{%10,%11,%12,%13},{%14},{%15,%16},{%17},{%18,%19};\n" + : "=f"(d0), "=f"(d1), "=f"(d8), "=f"(d9) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1), "f"(c0), "f"(c1), "f"(c8), "f"(c9), + "r"(uint32_t(sfa0)), "h"(bidA), "h"(tidA), "r"(uint32_t(sfb0)), "h"(bidB), "h"(tidB0)); + + gap_fn(0); + + asm volatile( + "mma.sync.aligned.kind::mxf4nvf4.block_scale.scale_vec::4X.m16n8k64.row.col.f32.e2m1.e2m1." + "f32.ue4m3 " + "{%0,%1,%2,%3},{%4,%5,%6,%7},{%8,%9},{%10,%11,%12,%13},{%14},{%15,%16},{%17},{%18,%19};\n" + : "=f"(d2), "=f"(d3), "=f"(d10), "=f"(d11) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b2), "r"(b3), "f"(c2), "f"(c3), "f"(c10), "f"(c11), + "r"(uint32_t(sfa0)), "h"(bidA), "h"(tidA), "r"(uint32_t(sfb0)), "h"(bidB), "h"(tidB1)); + + gap_fn(1); + + asm volatile( + "mma.sync.aligned.kind::mxf4nvf4.block_scale.scale_vec::4X.m16n8k64.row.col.f32.e2m1.e2m1." + "f32.ue4m3 " + "{%0,%1,%2,%3},{%4,%5,%6,%7},{%8,%9},{%10,%11,%12,%13},{%14},{%15,%16},{%17},{%18,%19};\n" + : "=f"(d4), "=f"(d5), "=f"(d12), "=f"(d13) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b4), "r"(b5), "f"(c4), "f"(c5), "f"(c12), "f"(c13), + "r"(uint32_t(sfa0)), "h"(bidA), "h"(tidA), "r"(uint32_t(sfb0)), "h"(bidB), "h"(tidB2)); + + gap_fn(2); + + asm volatile( + "mma.sync.aligned.kind::mxf4nvf4.block_scale.scale_vec::4X.m16n8k64.row.col.f32.e2m1.e2m1." + "f32.ue4m3 " + "{%0,%1,%2,%3},{%4,%5,%6,%7},{%8,%9},{%10,%11,%12,%13},{%14},{%15,%16},{%17},{%18,%19};\n" + : "=f"(d6), "=f"(d7), "=f"(d14), "=f"(d15) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b6), "r"(b7), "f"(c6), "f"(c7), "f"(c14), "f"(c15), + "r"(uint32_t(sfa0)), "h"(bidA), "h"(tidA), "r"(uint32_t(sfb0)), "h"(bidB), "h"(tidB3)); +#else + CUTE_INVALID_CONTROL_PATH("fma_with_interleave requires CUTE_ARCH_MXF4NVF4_4X_UE4M3_MMA_ENABLED"); +#endif +} + +struct SM120_16x32x64_TN_VS_NVFP4 { + using DRegisters = float[16]; + using ARegisters = uint32_t[4]; + using BRegisters = uint32_t[8]; + using CRegisters = float[16]; + + static constexpr int SFBits = 32; + using RegTypeSF = cute::uint_bit_t; + + using SFARegisters = RegTypeSF[1]; + using SFBRegisters = RegTypeSF[1]; + + CUTE_HOST_DEVICE static void fma( + float& d0, float& d1, float& d2, float& d3, float& d4, float& d5, float& d6, float& d7, + float& d8, float& d9, float& d10, float& d11, float& d12, float& d13, float& d14, float& d15, + uint32_t const& a0, uint32_t const& a1, uint32_t const& a2, uint32_t const& a3, + uint32_t const& b0, uint32_t const& b1, uint32_t const& b2, uint32_t const& b3, + uint32_t const& b4, uint32_t const& b5, uint32_t const& b6, uint32_t const& b7, + float const& c0, float const& c1, float const& c2, float const& c3, float const& c4, + float const& c5, float const& c6, float const& c7, float const& c8, float const& c9, + float const& c10, float const& c11, float const& c12, float const& c13, float const& c14, + float const& c15, RegTypeSF const& sfa0, RegTypeSF const& sfb0) { + static constexpr uint16_t tidA = 0; + static constexpr uint16_t bidA = 0; + static constexpr uint16_t bidB = 0; + static constexpr uint16_t tidB0 = 0; + static constexpr uint16_t tidB1 = 1; + static constexpr uint16_t tidB2 = 2; + static constexpr uint16_t tidB3 = 3; + +#if defined(LATENCY_PROBE) && (LATENCY_PROBE > 0) +#if defined(LATENCY_PROBE_REAL) + +#define _PROBE_MUFU_BLOCK() \ + { \ + float _pd = 1.0f; \ + _Pragma("unroll") for (int _pi = 0; _pi < LATENCY_PROBE; ++_pi) { \ + _pd = fmaxf(_pd, __shfl_xor_sync(int32_t(-1), _pd, 1)); \ + } \ + asm volatile("" ::"f"(_pd)); \ + } +#else + +#define _PROBE_MUFU_BLOCK() \ + { \ + float _pd = 1.0f; \ + _Pragma("unroll") for (int _pi = 0; _pi < LATENCY_PROBE; ++_pi) { \ + asm volatile("ex2.approx.ftz.f32 %0, %0;" : "+f"(_pd)); \ + } \ + asm volatile("" ::"f"(_pd)); \ + } +#endif +#else +#define _PROBE_MUFU_BLOCK() +#endif + +#if defined(CUTE_ARCH_MXF4NVF4_4X_UE4M3_MMA_ENABLED) + + asm volatile( + "mma.sync.aligned.kind::mxf4nvf4.block_scale.scale_vec::4X.m16n8k64.row.col.f32.e2m1.e2m1." + "f32.ue4m3 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13}," + "{%14}," + "{%15, %16}," + "{%17}," + "{%18, %19};\n" + : "=f"(d0), "=f"(d1), "=f"(d8), "=f"(d9) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1), "f"(c0), "f"(c1), "f"(c8), "f"(c9), + "r"(uint32_t(sfa0)), "h"(bidA), "h"(tidA), "r"(uint32_t(sfb0)), "h"(bidB), "h"(tidB0)); + + _PROBE_MUFU_BLOCK() + + asm volatile( + "mma.sync.aligned.kind::mxf4nvf4.block_scale.scale_vec::4X.m16n8k64.row.col.f32.e2m1.e2m1." + "f32.ue4m3 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13}," + "{%14}," + "{%15, %16}," + "{%17}," + "{%18, %19};\n" + : "=f"(d2), "=f"(d3), "=f"(d10), "=f"(d11) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b2), "r"(b3), "f"(c2), "f"(c3), "f"(c10), + "f"(c11), "r"(uint32_t(sfa0)), "h"(bidA), "h"(tidA), "r"(uint32_t(sfb0)), "h"(bidB), + "h"(tidB1)); + + _PROBE_MUFU_BLOCK() + + asm volatile( + "mma.sync.aligned.kind::mxf4nvf4.block_scale.scale_vec::4X.m16n8k64.row.col.f32.e2m1.e2m1." + "f32.ue4m3 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13}," + "{%14}," + "{%15, %16}," + "{%17}," + "{%18, %19};\n" + : "=f"(d4), "=f"(d5), "=f"(d12), "=f"(d13) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b4), "r"(b5), "f"(c4), "f"(c5), "f"(c12), + "f"(c13), "r"(uint32_t(sfa0)), "h"(bidA), "h"(tidA), "r"(uint32_t(sfb0)), "h"(bidB), + "h"(tidB2)); + + _PROBE_MUFU_BLOCK() + + asm volatile( + "mma.sync.aligned.kind::mxf4nvf4.block_scale.scale_vec::4X.m16n8k64.row.col.f32.e2m1.e2m1." + "f32.ue4m3 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13}," + "{%14}," + "{%15, %16}," + "{%17}," + "{%18, %19};\n" + : "=f"(d6), "=f"(d7), "=f"(d14), "=f"(d15) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b6), "r"(b7), "f"(c6), "f"(c7), "f"(c14), + "f"(c15), "r"(uint32_t(sfa0)), "h"(bidA), "h"(tidA), "r"(uint32_t(sfb0)), "h"(bidB), + "h"(tidB3)); +#else + CUTE_INVALID_CONTROL_PATH( + "Attempting to use SM120::BLOCKSCALED::SM120_16x32x64_TN_VS_NVFP4 without " + "CUTE_ARCH_MXF4NVF4_4X_UE4M3_MMA_ENABLED"); +#endif + +#undef _PROBE_MUFU_BLOCK + } +}; + +struct SM120_16x64x64_TN_VS_NVFP4 { + using DRegisters = float[32]; + using ARegisters = uint32_t[4]; + using BRegisters = uint32_t[16]; + using CRegisters = float[32]; + + static constexpr int SFBits = 32; + using RegTypeSF = cute::uint_bit_t; + using SFARegisters = RegTypeSF[1]; + using SFBRegisters = RegTypeSF[1]; + + CUTE_HOST_DEVICE static void fma( + float& d0, float& d1, float& d2, float& d3, float& d4, float& d5, float& d6, float& d7, + float& d8, float& d9, float& d10, float& d11, float& d12, float& d13, float& d14, float& d15, + float& d16, float& d17, float& d18, float& d19, float& d20, float& d21, float& d22, + float& d23, float& d24, float& d25, float& d26, float& d27, float& d28, float& d29, + float& d30, float& d31, uint32_t const& a0, uint32_t const& a1, uint32_t const& a2, + uint32_t const& a3, uint32_t const& b0, uint32_t const& b1, uint32_t const& b2, + uint32_t const& b3, uint32_t const& b4, uint32_t const& b5, uint32_t const& b6, + uint32_t const& b7, uint32_t const& b8, uint32_t const& b9, uint32_t const& b10, + uint32_t const& b11, uint32_t const& b12, uint32_t const& b13, uint32_t const& b14, + uint32_t const& b15, float const& c0, float const& c1, float const& c2, float const& c3, + float const& c4, float const& c5, float const& c6, float const& c7, float const& c8, + float const& c9, float const& c10, float const& c11, float const& c12, float const& c13, + float const& c14, float const& c15, float const& c16, float const& c17, float const& c18, + float const& c19, float const& c20, float const& c21, float const& c22, float const& c23, + float const& c24, float const& c25, float const& c26, float const& c27, float const& c28, + float const& c29, float const& c30, float const& c31, RegTypeSF const& sfa0, + RegTypeSF const& sfb0) { + static constexpr uint16_t tidA = 0; + static constexpr uint16_t bidA = 0; + static constexpr uint16_t bidB = 0; + +#if defined(CUTE_ARCH_MXF4NVF4_4X_UE4M3_MMA_ENABLED) + + asm volatile( + "mma.sync.aligned.kind::mxf4nvf4.block_scale.scale_vec::4X.m16n8k64.row.col.f32.e2m1.e2m1." + "f32.ue4m3 " + "{%0,%1,%2,%3},{%4,%5,%6,%7},{%8,%9},{%10,%11,%12,%13},{%14},{%15,%16},{%17},{%18,%19};\n" + : "=f"(d0), "=f"(d1), "=f"(d16), "=f"(d17) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1), "f"(c0), "f"(c1), "f"(c16), + "f"(c17), "r"(uint32_t(sfa0)), "h"(bidA), "h"(tidA), "r"(uint32_t(sfb0)), "h"(bidB), + "h"(uint16_t(0))); + + asm volatile( + "mma.sync.aligned.kind::mxf4nvf4.block_scale.scale_vec::4X.m16n8k64.row.col.f32.e2m1.e2m1." + "f32.ue4m3 " + "{%0,%1,%2,%3},{%4,%5,%6,%7},{%8,%9},{%10,%11,%12,%13},{%14},{%15,%16},{%17},{%18,%19};\n" + : "=f"(d2), "=f"(d3), "=f"(d18), "=f"(d19) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b2), "r"(b3), "f"(c2), "f"(c3), "f"(c18), + "f"(c19), "r"(uint32_t(sfa0)), "h"(bidA), "h"(tidA), "r"(uint32_t(sfb0)), "h"(bidB), + "h"(uint16_t(1))); + + asm volatile( + "mma.sync.aligned.kind::mxf4nvf4.block_scale.scale_vec::4X.m16n8k64.row.col.f32.e2m1.e2m1." + "f32.ue4m3 " + "{%0,%1,%2,%3},{%4,%5,%6,%7},{%8,%9},{%10,%11,%12,%13},{%14},{%15,%16},{%17},{%18,%19};\n" + : "=f"(d4), "=f"(d5), "=f"(d20), "=f"(d21) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b4), "r"(b5), "f"(c4), "f"(c5), "f"(c20), + "f"(c21), "r"(uint32_t(sfa0)), "h"(bidA), "h"(tidA), "r"(uint32_t(sfb0)), "h"(bidB), + "h"(uint16_t(2))); + + asm volatile( + "mma.sync.aligned.kind::mxf4nvf4.block_scale.scale_vec::4X.m16n8k64.row.col.f32.e2m1.e2m1." + "f32.ue4m3 " + "{%0,%1,%2,%3},{%4,%5,%6,%7},{%8,%9},{%10,%11,%12,%13},{%14},{%15,%16},{%17},{%18,%19};\n" + : "=f"(d6), "=f"(d7), "=f"(d22), "=f"(d23) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b6), "r"(b7), "f"(c6), "f"(c7), "f"(c22), + "f"(c23), "r"(uint32_t(sfa0)), "h"(bidA), "h"(tidA), "r"(uint32_t(sfb0)), "h"(bidB), + "h"(uint16_t(3))); + + asm volatile( + "mma.sync.aligned.kind::mxf4nvf4.block_scale.scale_vec::4X.m16n8k64.row.col.f32.e2m1.e2m1." + "f32.ue4m3 " + "{%0,%1,%2,%3},{%4,%5,%6,%7},{%8,%9},{%10,%11,%12,%13},{%14},{%15,%16},{%17},{%18,%19};\n" + : "=f"(d8), "=f"(d9), "=f"(d24), "=f"(d25) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b8), "r"(b9), "f"(c8), "f"(c9), "f"(c24), + "f"(c25), "r"(uint32_t(sfa0)), "h"(bidA), "h"(tidA), "r"(uint32_t(sfb0)), "h"(bidB), + "h"(uint16_t(0))); + + asm volatile( + "mma.sync.aligned.kind::mxf4nvf4.block_scale.scale_vec::4X.m16n8k64.row.col.f32.e2m1.e2m1." + "f32.ue4m3 " + "{%0,%1,%2,%3},{%4,%5,%6,%7},{%8,%9},{%10,%11,%12,%13},{%14},{%15,%16},{%17},{%18,%19};\n" + : "=f"(d10), "=f"(d11), "=f"(d26), "=f"(d27) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b10), "r"(b11), "f"(c10), "f"(c11), "f"(c26), + "f"(c27), "r"(uint32_t(sfa0)), "h"(bidA), "h"(tidA), "r"(uint32_t(sfb0)), "h"(bidB), + "h"(uint16_t(1))); + + asm volatile( + "mma.sync.aligned.kind::mxf4nvf4.block_scale.scale_vec::4X.m16n8k64.row.col.f32.e2m1.e2m1." + "f32.ue4m3 " + "{%0,%1,%2,%3},{%4,%5,%6,%7},{%8,%9},{%10,%11,%12,%13},{%14},{%15,%16},{%17},{%18,%19};\n" + : "=f"(d12), "=f"(d13), "=f"(d28), "=f"(d29) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b12), "r"(b13), "f"(c12), "f"(c13), "f"(c28), + "f"(c29), "r"(uint32_t(sfa0)), "h"(bidA), "h"(tidA), "r"(uint32_t(sfb0)), "h"(bidB), + "h"(uint16_t(2))); + + asm volatile( + "mma.sync.aligned.kind::mxf4nvf4.block_scale.scale_vec::4X.m16n8k64.row.col.f32.e2m1.e2m1." + "f32.ue4m3 " + "{%0,%1,%2,%3},{%4,%5,%6,%7},{%8,%9},{%10,%11,%12,%13},{%14},{%15,%16},{%17},{%18,%19};\n" + : "=f"(d14), "=f"(d15), "=f"(d30), "=f"(d31) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b14), "r"(b15), "f"(c14), "f"(c15), "f"(c30), + "f"(c31), "r"(uint32_t(sfa0)), "h"(bidA), "h"(tidA), "r"(uint32_t(sfb0)), "h"(bidB), + "h"(uint16_t(3))); +#else + CUTE_INVALID_CONTROL_PATH( + "Attempting to use SM120::BLOCKSCALED::SM120_16x64x64_TN_VS without " + "CUTE_ARCH_MXF4NVF4_4X_UE4M3_MMA_ENABLED"); +#endif + } +}; + +} // namespace cute::SM120::BLOCKSCALED + +namespace cute { + +template <> +struct MMA_Traits { + using ValTypeA = uint4_t; + using ValTypeB = uint4_t; + + using ValTypeD = float; + using ValTypeC = float; + + using ValTypeSF = cutlass::float_ue4m3_t; + constexpr static int SFVecSize = 16; + + using Shape_MNK = Shape<_16, _32, _64>; + using ThrID = Layout<_32>; + + using ALayout = Layout, Shape<_8, _2, _2>>, + Stride, Stride<_16, _8, _512>>>; + + using BLayout = Layout, Shape<_8, _2, _4>>, + Stride, Stride<_32, _1024, _8>>>; + + using SFALayout = Layout, _64>, Stride, _16>>; + + using SFBLayout = Layout, _64>, Stride, _32>>; + + using CLayout = Layout, Shape, _2>>, + Stride, Stride, _8>>>; +}; + +template <> +struct MMA_Traits { + using ValTypeA = uint4_t; + using ValTypeB = uint4_t; + using ValTypeD = float; + using ValTypeC = float; + using ValTypeSF = cutlass::float_ue4m3_t; + constexpr static int SFVecSize = 16; + + using Shape_MNK = Shape<_16, _64, _64>; + using ThrID = Layout<_32>; + + using ALayout = Layout, Shape<_8, _2, _2>>, + Stride, Stride<_16, _8, _512>>>; + + using BLayout = Layout, Shape<_8, _2, _8>>, + Stride, Stride<_64, _2048, _8>>>; + + using SFALayout = Layout, _64>, Stride, _16>>; + + using SFBLayout = Layout, _64>, Stride, _8>>; + + using CLayout = Layout, Shape, _2>>, + Stride, Stride, _8>>>; +}; + +template +CUTE_HOST_DEVICE constexpr auto thrfrg_SFA(SFATensor&& sfatensor, + TiledMMA& mma) { + CUTE_STATIC_ASSERT_V(rank(sfatensor) >= Int<2>{}); + + using AtomShape_MNK = typename Atom::Shape_MNK; + using AtomLayoutSFA_TV = typename Atom::Traits::SFALayout; + + auto permutation_mnk = TiledPerm{}; + auto thr_layout_vmnk = mma.get_thr_layout_vmnk(); + + auto t_tile = make_tile(get<0>(permutation_mnk), get<2>(permutation_mnk)); + auto t_tensor = logical_divide(sfatensor, t_tile); + + auto a_tile = + make_tile(make_layout(size<0>(AtomShape_MNK{})), make_layout(size<2>(AtomShape_MNK{}))); + auto a_tensor = zipped_divide(t_tensor, a_tile); + + auto tv_tensor = a_tensor.compose(AtomLayoutSFA_TV{}, _); + + auto thr_tile = make_tile( + _, make_tile(make_layout(size<1>(thr_layout_vmnk)), make_layout(size<3>(thr_layout_vmnk)))); + auto thr_tensor = zipped_divide(tv_tensor, thr_tile); + + return thr_tensor; +} + +template +CUTE_HOST_DEVICE constexpr auto thrfrg_SFB(SFBTensor&& sfbtensor, + TiledMMA& mma) { + CUTE_STATIC_ASSERT_V(rank(sfbtensor) >= Int<2>{}); + + using AtomShape_MNK = typename Atom::Shape_MNK; + using AtomLayoutSFB_TV = typename Atom::Traits::SFBLayout; + + auto permutation_mnk = TiledPerm{}; + auto thr_layout_vmnk = mma.get_thr_layout_vmnk(); + + auto t_tile = make_tile(get<1>(permutation_mnk), get<2>(permutation_mnk)); + auto t_tensor = logical_divide(sfbtensor, t_tile); + + auto a_tile = + make_tile(make_layout(size<1>(AtomShape_MNK{})), make_layout(size<2>(AtomShape_MNK{}))); + auto a_tensor = zipped_divide(t_tensor, a_tile); + + auto tv_tensor = a_tensor.compose(AtomLayoutSFB_TV{}, _); + + auto thr_tile = make_tile( + _, make_tile(make_layout(size<2>(thr_layout_vmnk)), make_layout(size<3>(thr_layout_vmnk)))); + auto thr_tensor = zipped_divide(tv_tensor, thr_tile); + return thr_tensor; +} + +template +CUTE_HOST_DEVICE constexpr auto partition_SFA(SFATensor&& sfatensor, ThrMma& thread_mma) { + auto thr_tensor = make_tensor(static_cast(sfatensor).data(), + thrfrg_SFA(sfatensor.layout(), thread_mma)); + auto thr_vmnk = thread_mma.thr_vmnk_; + auto thr_vmk = make_coord(get<0>(thr_vmnk), make_coord(get<1>(thr_vmnk), get<3>(thr_vmnk))); + return thr_tensor(thr_vmk, make_coord(_, repeat(thr_tensor)>(_))); +} + +template +CUTE_HOST_DEVICE constexpr auto partition_fragment_SFA(SFATensor&& sfatensor, ThrMma& thread_mma) { + using ValTypeSF = typename ThrMma::Atom::Traits::ValTypeSF; + return make_fragment_like(partition_SFA(sfatensor, thread_mma)); +} + +template +CUTE_HOST_DEVICE constexpr auto partition_SFB(SFBTensor&& sfbtensor, ThrMma& thread_mma) { + auto thr_tensor = make_tensor(static_cast(sfbtensor).data(), + thrfrg_SFB(sfbtensor.layout(), thread_mma)); + auto thr_vmnk = thread_mma.thr_vmnk_; + auto thr_vnk = make_coord(get<0>(thr_vmnk), make_coord(get<2>(thr_vmnk), get<3>(thr_vmnk))); + return thr_tensor(thr_vnk, make_coord(_, repeat(thr_tensor)>(_))); +} + +template +CUTE_HOST_DEVICE constexpr auto partition_fragment_SFB(SFBTensor&& sfbtensor, ThrMma& thread_mma) { + using ValTypeSF = typename ThrMma::Atom::Traits::ValTypeSF; + return make_fragment_like(partition_SFB(sfbtensor, thread_mma)); +} + +template +CUTE_HOST_DEVICE constexpr auto get_layoutSFA_TV(TiledMma& mma) { + auto tile_shape_mnk = tile_shape(mma); + auto ref_A = make_layout(make_shape(size<0>(tile_shape_mnk), size<2>(tile_shape_mnk))); + auto thr_layout_vmnk = mma.get_thr_layout_vmnk(); + + auto atile = make_tile( + _, make_tile(make_layout(make_shape(size<1>(thr_layout_vmnk), size<2>(thr_layout_vmnk)), + make_stride(Int<1>{}, Int<0>{})), + _)); + + auto thridx_2_thrid = right_inverse(thr_layout_vmnk); + + return thrfrg_SFA(ref_A, mma).compose(atile, _).compose(thridx_2_thrid, _); +} + +template +CUTE_HOST_DEVICE constexpr auto get_layoutSFB_TV(TiledMma& mma) { + auto tile_shape_mnk = tile_shape(mma); + auto ref_B = make_layout(make_shape(size<1>(tile_shape_mnk), size<2>(tile_shape_mnk))); + auto thr_layout_vmnk = mma.get_thr_layout_vmnk(); + + auto btile = make_tile( + _, make_tile(make_layout(make_shape(size<1>(thr_layout_vmnk), size<2>(thr_layout_vmnk)), + make_stride(Int<0>{}, Int<1>{})), + _)); + + auto thridx_2_thrid = right_inverse(thr_layout_vmnk); + + return thrfrg_SFB(ref_B, mma).compose(btile, _).compose(thridx_2_thrid, _); +} + +} // namespace cute diff --git a/include/flashinfer/attention/sm120/nvfp4_attention_sm120/common/gemm_with_interleave.h b/include/flashinfer/attention/sm120/nvfp4_attention_sm120/common/gemm_with_interleave.h new file mode 100644 index 00000000000..559de70839a --- /dev/null +++ b/include/flashinfer/attention/sm120/nvfp4_attention_sm120/common/gemm_with_interleave.h @@ -0,0 +1,72 @@ +/* + * Copyright (c) 2025 by SageAttention team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * + * This file adapts cute::SM120::BLOCKSCALED::mma_unpack from CUTLASS/CuTe: + * Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. + * SPDX-License-Identifier: BSD-3-Clause + */ + +#pragma once + +#include "cute/atom/mma_atom.hpp" +#include "cute/tensor.hpp" +#include "cute_extension.h" + +namespace nvfp4_attention { + +using namespace cute; + +template +CUTE_HOST_DEVICE void mma_unpack_interleaved(MMA_Traits const&, Tensor& D, + Tensor const& A_zipped, + Tensor const& B_zipped, + Tensor const& C, GapFn&& gap_fn) { + using RegTypeD = typename remove_extent::type; + using RegTypeA = typename remove_extent::type; + using RegTypeB = typename remove_extent::type; + using RegTypeC = typename remove_extent::type; + using RegTypeSFA = typename remove_extent::type; + using RegTypeSFB = typename remove_extent::type; + + auto [A, SFA] = unzip_tensor(A_zipped); + auto [B, SFB] = unzip_tensor(B_zipped); + + Tensor rA = recast(A); + Tensor rB = recast(B); + Tensor rD = recast(D); + Tensor rC = recast(C); + Tensor rSFA = recast(filter_zeros(SFA)); + Tensor rSFB = recast(filter_zeros(SFB)); + + cute::SM120::BLOCKSCALED::fma_with_interleave( + rD(0), rD(1), rD(2), rD(3), rD(4), rD(5), rD(6), rD(7), rD(8), rD(9), rD(10), rD(11), rD(12), + rD(13), rD(14), rD(15), rA(0), rA(1), rA(2), rA(3), rB(0), rB(1), rB(2), rB(3), rB(4), rB(5), + rB(6), rB(7), rC(0), rC(1), rC(2), rC(3), rC(4), rC(5), rC(6), rC(7), rC(8), rC(9), rC(10), + rC(11), rC(12), rC(13), rC(14), rC(15), rSFA(0), rSFB(0), gap_fn); +} + +template +CUTE_HOST_DEVICE void gemm_interleaved(TiledMma const& tiled_mma, Tensor& C, + Tensor const& A, Tensor const& B, + GapFn&& gap_fn) { + using Traits = typename TiledMma::AtomThrID; + + using MMAOp = typename TiledMma::MMA_Atom_Arch; + mma_unpack_interleaved(MMA_Traits{}, C, A, B, C, static_cast(gap_fn)); +} + +} // namespace nvfp4_attention diff --git a/include/flashinfer/attention/sm120/nvfp4_attention_sm120/common/params.h b/include/flashinfer/attention/sm120/nvfp4_attention_sm120/common/params.h new file mode 100644 index 00000000000..144677712a0 --- /dev/null +++ b/include/flashinfer/attention/sm120/nvfp4_attention_sm120/common/params.h @@ -0,0 +1,143 @@ +/* + * Copyright (c) 2025 by SageAttention team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include + +#include + +#include "cutlass/fast_math.h" + +struct Qkv_params { + using index_t = int64_t; + + void* __restrict__ q_ptr; + void* __restrict__ k_ptr; + void* __restrict__ v_ptr; + void* __restrict__ delta_s_ptr; + + void* __restrict__ sfq_ptr; + void* __restrict__ sfk_ptr; + void* __restrict__ sfv_ptr; + + index_t q_batch_stride; + index_t k_batch_stride; + index_t v_batch_stride; + index_t q_row_stride; + index_t k_row_stride; + index_t v_row_stride; + index_t q_head_stride; + index_t k_head_stride; + index_t v_head_stride; + index_t ds_batch_stride; + index_t ds_row_stride; + index_t ds_head_stride; + + index_t sfq_batch_stride; + index_t sfk_batch_stride; + index_t sfv_batch_stride; + index_t sfq_row_stride; + index_t sfk_row_stride; + index_t sfv_row_stride; + index_t sfq_head_stride; + index_t sfk_head_stride; + index_t sfv_head_stride; + + int h, h_k; + + int h_h_k_ratio; +}; + +struct Flash_fwd_params : public Qkv_params { + void* __restrict__ o_ptr; + void* __restrict__ oaccum_ptr; + void* __restrict__ s_ptr; + + index_t o_batch_stride; + index_t o_row_stride; + index_t o_head_stride; + + void* __restrict__ p_ptr; + + void* __restrict__ softmax_lse_ptr; + void* __restrict__ softmax_lseaccum_ptr; + + int b, seqlen_q, seqlen_k, seqlen_knew, d, seqlen_q_rounded, seqlen_k_rounded, d_rounded, + rotary_dim, unpadded_seqlen_k; + cutlass::FastDivmod head_divmod, m_block_divmod; + int total_blocks; + int seqlen_s; + + float scale_softmax; + float scale_softmax_log2; + uint32_t scale_softmax_log2_half2; + + int* __restrict__ cu_seqlens_q; + int* __restrict__ cu_seqlens_k; + + int* __restrict__ seqused_k; + + int* __restrict__ blockmask; + + void* __restrict__ knew_ptr; + void* __restrict__ vnew_ptr; + + index_t knew_batch_stride; + index_t vnew_batch_stride; + index_t knew_row_stride; + index_t vnew_row_stride; + index_t knew_head_stride; + index_t vnew_head_stride; + + void* __restrict__ rotary_cos_ptr; + void* __restrict__ rotary_sin_ptr; + + int* __restrict__ cache_batch_idx; + + int* __restrict__ block_table; + index_t block_table_batch_stride; + int page_block_size; + + float p_dropout; + + uint8_t p_dropout_in_uint8_t; + + float rp_dropout; + float scale_softmax_rp_dropout; + + int window_size_left, window_size_right; + + uint64_t philox_args[2]; + + uint64_t* rng_state; + + bool is_bf16; + bool is_e4m3; + bool is_causal; + bool per_block_mean; + + bool is_seqlens_k_cumulative; + + bool is_rotary_interleaved; + + int num_splits; + + void* __restrict__ alibi_slopes_ptr; + index_t alibi_slopes_batch_stride; + + int* __restrict__ tile_count_semaphore; +}; diff --git a/include/flashinfer/attention/sm120/nvfp4_attention_sm120/common/static_switch.h b/include/flashinfer/attention/sm120/nvfp4_attention_sm120/common/static_switch.h new file mode 100644 index 00000000000..1d3753c4e15 --- /dev/null +++ b/include/flashinfer/attention/sm120/nvfp4_attention_sm120/common/static_switch.h @@ -0,0 +1,87 @@ +/* + * Copyright (c) 2025 by SageAttention team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#define BOOL_SWITCH(COND, CONST_NAME, ...) \ + [&] { \ + if (COND) { \ + constexpr static bool CONST_NAME = true; \ + return __VA_ARGS__(); \ + } else { \ + constexpr static bool CONST_NAME = false; \ + return __VA_ARGS__(); \ + } \ + }() + +#define PREC_SWITCH(PRECTYPE, ...) \ + [&] { \ + if (PRECTYPE == 1) { \ + using kPrecType = cutlass::half_t; \ + constexpr static bool kSoftFp16 = false; \ + constexpr static bool kHybrid = false; \ + return __VA_ARGS__(); \ + } else if (PRECTYPE == 2) { \ + using kPrecType = cutlass::float_e4m3_t; \ + constexpr static bool kSoftFp16 = false; \ + constexpr static bool kHybrid = false; \ + return __VA_ARGS__(); \ + } else if (PRECTYPE == 3) { \ + using kPrecType = cutlass::float_e4m3_t; \ + constexpr static bool kSoftFp16 = false; \ + constexpr static bool kHybrid = true; \ + return __VA_ARGS__(); \ + } else if (PRECTYPE == 4) { \ + using kPrecType = cutlass::float_e4m3_t; \ + constexpr static bool kSoftFp16 = true; \ + constexpr static bool kHybrid = false; \ + return __VA_ARGS__(); \ + } else { \ + __builtin_unreachable(); \ + } \ + }() + +#define HEADDIM_SWITCH(HEADDIM, ...) \ + [&] { \ + if (HEADDIM == 64) { \ + constexpr static int kHeadSize = 64; \ + return __VA_ARGS__(); \ + } else if (HEADDIM == 128) { \ + constexpr static int kHeadSize = 128; \ + return __VA_ARGS__(); \ + } else if (HEADDIM == 256) { \ + constexpr static int kHeadSize = 256; \ + return __VA_ARGS__(); \ + } else { \ + __builtin_unreachable(); \ + } \ + }() + +#define SEQLEN_SWITCH(USE_VAR_SEQ_LEN, SEQ_LEN_OUT_OF_BOUND_CHECK, ...) \ + [&] { \ + if (!USE_VAR_SEQ_LEN) { \ + if (SEQ_LEN_OUT_OF_BOUND_CHECK) { \ + using kSeqLenTraitsType = FixedSeqLenTraits; \ + return __VA_ARGS__(); \ + } else { \ + using kSeqLenTraitsType = FixedSeqLenTraits; \ + return __VA_ARGS__(); \ + } \ + } else { \ + using kSeqLenTraitsType = VarSeqLenTraits; \ + return __VA_ARGS__(); \ + } \ + }() diff --git a/include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/consumer/delta_correction.cuh b/include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/consumer/delta_correction.cuh new file mode 100644 index 00000000000..d28069e7b3c --- /dev/null +++ b/include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/consumer/delta_correction.cuh @@ -0,0 +1,79 @@ +/* + * Copyright (c) 2025 by SageAttention team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "cute/tensor.hpp" +#include "cutlass/cutlass.h" + +namespace nvfp4_attention { + +using namespace cute; + +template +struct DeltaSCorrection { + using TileShape_MNK = typename Traits::TileShape_MNK; + using SmemLayoutDS = typename Traits::SmemLayoutDS; + static constexpr bool BlockMean = Traits::BlockMean; + + static constexpr int kBlockM = get<0>(TileShape_MNK{}); + static constexpr int kBlockN = get<1>(TileShape_MNK{}); + + template + __device__ __forceinline__ static void add_delta_s(TensorAcc& acc, TensorSDS const& sDS, + PipelineStateK const& smem_pipe_read_k) { + auto tSsDS_stage = recast(sDS(_, _, smem_pipe_read_k.index())); + auto acc_float4 = recast(acc); + + int quad_id = (threadIdx.x % 4) * 2; + + for (int i = 0; i < 4; i++) { + auto num = quad_id + i * 8; + + float4 delta_s_0 = tSsDS_stage(make_coord(_0{}, _0{}), make_coord(num, _0{})); + float4 delta_s_1 = tSsDS_stage(make_coord(_0{}, _0{}), make_coord(num + 1, _0{})); + + acc_float4(make_coord(make_coord(_0{}, _0{}), _0{}), _0{}, i) = delta_s_0; + acc_float4(make_coord(make_coord(_0{}, _0{}), _1{}), _0{}, i) = delta_s_0; + acc_float4(make_coord(make_coord(_0{}, _1{}), _0{}), _0{}, i) = delta_s_1; + acc_float4(make_coord(make_coord(_0{}, _1{}), _1{}), _0{}, i) = delta_s_1; + } + } + + template + __device__ __forceinline__ static auto make_lambda(TensorSDS const& sDS, + PipelineStateK const& smem_pipe_read_k) { + return [&sDS, &smem_pipe_read_k](auto& acc) { add_delta_s(acc, sDS, smem_pipe_read_k); }; + } + + __device__ __forceinline__ static auto make_noop_lambda() { + return [](auto& acc) { + + }; + } +}; + +template +__device__ __forceinline__ auto make_delta_s_lambda(TensorSDS const& sDS, + PipelineStateK const& smem_pipe_read_k) { + if constexpr (UseDeltaS) { + return DeltaSCorrection::make_lambda(sDS, smem_pipe_read_k); + } else { + return DeltaSCorrection::make_noop_lambda(); + } +} + +} // namespace nvfp4_attention diff --git a/include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/consumer/pv_gemm.cuh b/include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/consumer/pv_gemm.cuh new file mode 100644 index 00000000000..b19186f02ab --- /dev/null +++ b/include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/consumer/pv_gemm.cuh @@ -0,0 +1,176 @@ +/* + * Copyright (c) 2025 by SageAttention team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "../../quantization/fp4_convert.cuh" +#include "cute/tensor.hpp" +#include "cutlass/cutlass.h" + +namespace nvfp4_attention { + +using cute::_; +using cute::_0; +using cute::_1; +using cute::_2; +using cute::_3; +using cute::_4; +using cute::_5; +using cute::_6; +using cute::_7; +using cute::copy; +using cute::get; +using cute::make_coord; +using cute::make_tensor_like; +using cute::make_zip_tensor; +using cute::recast; +using cute::size; +using cute::Tensor; + +template +struct PVGemmComputer { + using Element = typename Traits::Element; + using ElementSF = typename Traits::ElementSF; + using TileShape_MNK = typename Traits::TileShape_MNK; + using TiledMmaPV = typename Traits::TiledMmaPV; + using SmemCopyAtomKV = typename Traits::SmemCopyAtomKV; + using SmemCopyAtomSF = typename Traits::SmemCopyAtomSF; + using LayoutP = typename Traits::LayoutP; + using LayoutSFP = typename Traits::LayoutSFP; + + static constexpr int kBlockM = get<0>(TileShape_MNK{}); + static constexpr int kBlockN = get<1>(TileShape_MNK{}); + static constexpr int kBlockK = get<2>(TileShape_MNK{}); + + template + __device__ __forceinline__ static void copy_v_block( + SmemTiledCopyV const& smem_tiled_copy_V, SmemTiledCopySFV const& smem_tiled_copy_SFV, + TensorSsVt const& tOsVt, TensorSsSFVt const& tOsSFVt, TensorSrVt& tOrVt_copy_view, + TensorSrSFVt& tOrSFVt_copy_view, PipelineStateV const& smem_pipe_read_v, auto block_id) { + auto tOsVt_stage = tOsVt(_, _, _, smem_pipe_read_v.index()); + auto tOsSFVt_stage = tOsSFVt(_, _, _, smem_pipe_read_v.index()); + + copy(smem_tiled_copy_V, tOsVt_stage(_, _, block_id), tOrVt_copy_view(_, _, block_id)); + copy(smem_tiled_copy_SFV, tOsSFVt_stage(_, _, block_id), tOrSFVt_copy_view(_, _, block_id)); + } + + template + __device__ __forceinline__ static void quantize_p(TensorAcc const& acc_conversion_view, + TensorMaxP const& AbsMaxP, TensorRP& tOrP, + TensorRSFP& tOrSFP, int mma_k) { + Tensor AbsMaxP_stagek = AbsMaxP(_, make_coord(_, _, mma_k)); + Tensor acc_conversion_stagek = acc_conversion_view(_, _, mma_k); + + Tensor SFP = make_tensor_like(AbsMaxP_stagek.layout()); + Tensor SFP_uint32_view = recast(SFP); + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(AbsMaxP_stagek); i += 4) { + uint32_t& tmp = SFP_uint32_view(i / 4); + nvfp4_attention::packed_float_to_ue4m3(AbsMaxP_stagek(i), AbsMaxP_stagek(i + 1), + AbsMaxP_stagek(i + 2), AbsMaxP_stagek(i + 3), tmp); + } + + int const quad_id = threadIdx.x & 3; + uint32_t MASK = (0xFF00FF) << ((quad_id & 1) * 8); + Tensor tOrSFP_uint32_view = recast(tOrSFP(_, _, mma_k)); + Tensor tOrP_uint32_view = recast(tOrP(_, _, mma_k)); + + CUTLASS_PRAGMA_UNROLL + for (int mma_m = 0; mma_m < size<1>(tOrP); ++mma_m) { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < 4; ++i) { + nvfp4_attention::packed_float_to_e2m1(acc_conversion_stagek(make_coord(_0{}, i), mma_m), + acc_conversion_stagek(make_coord(_1{}, i), mma_m), + acc_conversion_stagek(make_coord(_2{}, i), mma_m), + acc_conversion_stagek(make_coord(_3{}, i), mma_m), + acc_conversion_stagek(make_coord(_4{}, i), mma_m), + acc_conversion_stagek(make_coord(_5{}, i), mma_m), + acc_conversion_stagek(make_coord(_6{}, i), mma_m), + acc_conversion_stagek(make_coord(_7{}, i), mma_m), + tOrP_uint32_view(i, mma_m)); + } + + uint32_t local_sfp = SFP_uint32_view(_0{}, _0{}, mma_m); + uint32_t peer_sfp = __shfl_xor_sync(int32_t(-1), local_sfp, 2); + if ((quad_id & 1) == 0) { + uint32_t sfp = (local_sfp & MASK) | ((peer_sfp & MASK) << 8); + tOrSFP_uint32_view(_0{}, mma_m) = sfp; + } else { + uint32_t sfp = (peer_sfp & MASK) | ((local_sfp & MASK) >> 8); + tOrSFP_uint32_view(_0{}, mma_m) = sfp; + } + } + } + + template + __device__ __forceinline__ static void compute_pv_gemm( + TiledMma const& tiled_mma_pv, TensorRP& tOrP, TensorRSFP& tOrSFP, TensorRVt const& tOrVt, + TensorRSFVt const& tOrSFVt, TensorRO& tOrO, TensorAccConv const& acc_conversion_view, + TensorMaxP const& AbsMaxP, SmemTiledCopyV const& smem_tiled_copy_V, + SmemTiledCopySFV const& smem_tiled_copy_SFV, TensorSsVt const& tOsVt, + TensorSsSFVt const& tOsSFVt, TensorRVtView& tOrVt_copy_view, + TensorRSFVtView& tOrSFVt_copy_view, PipelineStateV const& smem_pipe_read_v) { + CUTLASS_PRAGMA_UNROLL + for (int v_block = 0; v_block < size<2>(tOrP); ++v_block) { + cute::gemm(tiled_mma_pv, make_zip_tensor(tOrP(_, _, v_block), tOrSFP(_, _, v_block)), + make_zip_tensor(tOrVt(_, _, v_block), tOrSFVt(_, _, v_block)), tOrO); + + if (v_block < size<2>(tOrP) - 1) { + copy_v_block(smem_tiled_copy_V, smem_tiled_copy_SFV, tOsVt, tOsSFVt, tOrVt_copy_view, + tOrSFVt_copy_view, smem_pipe_read_v, v_block + 1); + + quantize_p(acc_conversion_view, AbsMaxP, tOrP, tOrSFP, v_block + 1); + } + } + } + + template + __device__ __forceinline__ static void run( + TiledMma const& tiled_mma_pv, TensorRP& tOrP, TensorRSFP& tOrSFP, TensorRVt const& tOrVt, + TensorRSFVt const& tOrSFVt, TensorRO& tOrO, TensorAccConv const& acc_conversion_view, + TensorMaxP const& AbsMaxP, SmemTiledCopyV const& smem_tiled_copy_V, + SmemTiledCopySFV const& smem_tiled_copy_SFV, TensorSsVt const& tOsVt, + TensorSsSFVt const& tOsSFVt, TensorRVtView& tOrVt_copy_view, + TensorRSFVtView& tOrSFVt_copy_view, PipelineV& pipeline_v, PipelineStateV& smem_pipe_read_v) { + auto barrier_token = pipeline_v.consumer_try_wait(smem_pipe_read_v); + pipeline_v.consumer_wait(smem_pipe_read_v, barrier_token); + + copy_v_block(smem_tiled_copy_V, smem_tiled_copy_SFV, tOsVt, tOsSFVt, tOrVt_copy_view, + tOrSFVt_copy_view, smem_pipe_read_v, _0{}); + + quantize_p(acc_conversion_view, AbsMaxP, tOrP, tOrSFP, 0); + + compute_pv_gemm(tiled_mma_pv, tOrP, tOrSFP, tOrVt, tOrSFVt, tOrO, acc_conversion_view, AbsMaxP, + smem_tiled_copy_V, smem_tiled_copy_SFV, tOsVt, tOsSFVt, tOrVt_copy_view, + tOrSFVt_copy_view, smem_pipe_read_v); + + pipeline_v.consumer_release(smem_pipe_read_v); + ++smem_pipe_read_v; + } +}; + +} // namespace nvfp4_attention diff --git a/include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/consumer/qk_gemm.cuh b/include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/consumer/qk_gemm.cuh new file mode 100644 index 00000000000..0d3ad96b94d --- /dev/null +++ b/include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/consumer/qk_gemm.cuh @@ -0,0 +1,155 @@ +/* + * Copyright (c) 2025 by SageAttention team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "../../utils/layout.cuh" +#include "cute/tensor.hpp" +#include "cutlass/cutlass.h" + +namespace nvfp4_attention { + +using cute::_; +using cute::_0; +using cute::copy; +using cute::get; +using cute::make_identity_tensor; +using cute::make_zip_tensor; +using cute::select; +using cute::size; +using cute::Tensor; + +template +struct QKGemmComputer { + using Element = typename Traits::Element; + using ElementSF = typename Traits::ElementSF; + using TileShape_MNK = typename Traits::TileShape_MNK; + using TiledMmaQK = typename Traits::TiledMmaQK; + using SmemCopyAtomQ = typename Traits::SmemCopyAtomQ; + using SmemCopyAtomKV = typename Traits::SmemCopyAtomKV; + using SmemCopyAtomSF = typename Traits::SmemCopyAtomSF; + + static constexpr int kBlockM = get<0>(TileShape_MNK{}); + static constexpr int kBlockN = get<1>(TileShape_MNK{}); + static constexpr int kBlockK = get<2>(TileShape_MNK{}); + + template + __device__ __forceinline__ static void copy_k_block( + SmemTiledCopyK const& smem_tiled_copy_K, SmemTiledCopySFK const& smem_tiled_copy_SFK, + TensorSsK const& tSsK, TensorSsSFK const& tSsSFK, TensorSrK& tSrK_copy_view, + TensorSrSFK& tSrSFK_copy_view, PipelineStateK const& smem_pipe_read_k, auto block_id) { + auto tSsK_stage = tSsK(_, _, _, smem_pipe_read_k.index()); + auto tSsSFK_stage = tSsSFK(_, _, _, smem_pipe_read_k.index()); + + copy(smem_tiled_copy_K, tSsK_stage(_, _, block_id), tSrK_copy_view(_, _, block_id)); + copy(smem_tiled_copy_SFK, tSsSFK_stage(_, _, block_id), tSrSFK_copy_view(_, _, block_id)); + } + + template + __device__ __forceinline__ static void compute_qk_gemm( + TiledMma const& tiled_mma_qk, TensorRQ const& tSrQ, TensorRSFQ const& tSrSFQ, + TensorRK const& tSrK, TensorRSFK const& tSrSFK, TensorRS& tSrS, + SmemTiledCopyK const& smem_tiled_copy_K, SmemTiledCopySFK const& smem_tiled_copy_SFK, + TensorSsK const& tSsK, TensorSsSFK const& tSsSFK, TensorRKView& tSrK_copy_view, + TensorRSFKView& tSrSFK_copy_view, PipelineStateK const& smem_pipe_read_k) { + CUTLASS_PRAGMA_UNROLL + for (int k_block = 0; k_block < size<2>(tSrQ); ++k_block) { + cute::gemm(tiled_mma_qk, make_zip_tensor(tSrQ(_, _, k_block), tSrSFQ(_, _, k_block)), + make_zip_tensor(tSrK(_, _, k_block), tSrSFK(_, _, k_block)), tSrS); + + if (k_block < size<2>(tSrQ) - 1) { + copy_k_block(smem_tiled_copy_K, smem_tiled_copy_SFK, tSsK, tSsSFK, tSrK_copy_view, + tSrSFK_copy_view, smem_pipe_read_k, k_block + 1); + } + } + } + + template + __device__ __forceinline__ static void apply_masking(TensorRS& tSrS, TensorCS const& tScS, + int n_block, int seqlen_k, + int unpadded_seqlen_k, int seqlen_q, + int m_block) { + auto col_limit_causal = [&](int row, int n_block_idx) { + return row + 1 + seqlen_k - n_block_idx * kBlockN - seqlen_q + m_block * kBlockM; + }; + + int const valid_cols = int(unpadded_seqlen_k - n_block * kBlockN); + if constexpr (!IsCausal) { + if (valid_cols >= kBlockN) { + return; + } + } else { + if (valid_cols >= kBlockN && col_limit_causal(0, n_block) >= kBlockN) { + return; + } + } + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(tSrS); ++i) { + int const col = nvfp4_attention::qk_acc_col_to_k_col(int(get<1>(tScS(i)))); + if constexpr (!IsCausal) { + if (col >= int(unpadded_seqlen_k - n_block * kBlockN)) { + tSrS(i) = -INFINITY; + } + } else { + int col_limit = + std::min(seqlen_k - n_block * kBlockN, col_limit_causal(int(get<0>(tScS(i))), n_block)); + if (col >= col_limit) { + tSrS(i) = -INFINITY; + } + } + } + } + + template + __device__ __forceinline__ static void run( + TiledMma const& tiled_mma_qk, TensorRQ const& tSrQ, TensorRSFQ const& tSrSFQ, + TensorRK const& tSrK, TensorRSFK const& tSrSFK, TensorRS& tSrS, + SmemTiledCopyK const& smem_tiled_copy_K, SmemTiledCopySFK const& smem_tiled_copy_SFK, + TensorSsK const& tSsK, TensorSsSFK const& tSsSFK, TensorRKView& tSrK_copy_view, + TensorRSFKView& tSrSFK_copy_view, PipelineK& pipeline_k, PipelineStateK& smem_pipe_read_k, + int n_block, int seqlen_k, int unpadded_seqlen_k, int seqlen_q, int m_block, + DeltaSFunc const& add_delta_s_func) { + auto barrier_token = pipeline_k.consumer_try_wait(smem_pipe_read_k); + pipeline_k.consumer_wait(smem_pipe_read_k, barrier_token); + + copy_k_block(smem_tiled_copy_K, smem_tiled_copy_SFK, tSsK, tSsSFK, tSrK_copy_view, + tSrSFK_copy_view, smem_pipe_read_k, _0{}); + + add_delta_s_func(tSrS); + + compute_qk_gemm(tiled_mma_qk, tSrQ, tSrSFQ, tSrK, tSrSFK, tSrS, smem_tiled_copy_K, + smem_tiled_copy_SFK, tSsK, tSsSFK, tSrK_copy_view, tSrSFK_copy_view, + smem_pipe_read_k); + + Tensor cS = cute::make_identity_tensor(select<0, 1>(TileShape_MNK{})); + Tensor tScS = tiled_mma_qk.get_thread_slice(threadIdx.x).partition_C(cS); + apply_masking(tSrS, tScS, n_block, seqlen_k, unpadded_seqlen_k, seqlen_q, m_block); + + pipeline_k.consumer_release(smem_pipe_read_k); + ++smem_pipe_read_k; + } +}; + +} // namespace nvfp4_attention diff --git a/include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/consumer/softmax.cuh b/include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/consumer/softmax.cuh new file mode 100644 index 00000000000..cb60eda3fb2 --- /dev/null +++ b/include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/consumer/softmax.cuh @@ -0,0 +1,774 @@ +/* + * Copyright (c) 2025 by SageAttention team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include + +#include "../../utils/layout.cuh" +#include "../../utils/math.cuh" +#include "cute/tensor.hpp" +#include "cutlass/numeric_types.h" + +namespace nvfp4_attention { + +using cute::clear; +using cute::copy; +using cute::fill; +using cute::flatten; +using cute::group_modes; +using cute::Int; +using cute::make_coord; +using cute::make_fragment_like; +using cute::make_tensor; +using cute::make_tensor_like; +using cute::Shape; +using cute::size; +using cute::Tensor; + +template +struct SoftmaxFused { + using TensorT = decltype(make_fragment_like(Shape>{})); + TensorT row_sum; + TensorT row_max; + TensorT scores_scale; + + static constexpr float fp8_scalexfp4_scale = 1.f / (448 * 6); + static constexpr float fp8_scalexfp4_scale_log2 = -11.392317422778762f; + static constexpr float fp4_scale = 1.f / 6.f; + static constexpr float fp4_scale_log2 = -2.584962500721156f; + static constexpr float AbsMaxPEps = 1.0e-8f; + static constexpr int RowReductionThr = 8; + + CUTLASS_DEVICE SoftmaxFused() {}; + + CUTLASS_DEVICE static float reduce_row_max_from_pairs(float value) { + CUTLASS_PRAGMA_UNROLL + for (int i = 2; i < RowReductionThr; i <<= 1) { + value = fmaxf(value, __shfl_xor_sync(int32_t(-1), value, i)); + } + return value; + } + +#if defined(FAST_RCP_ABSMAXP) + CUTLASS_DEVICE static float fast_rcp_approx(float x) { + float y; + asm volatile("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; + } +#endif + + CUTLASS_DEVICE static float safe_inv_absmax(float x) { + float denom = fmaxf(x, AbsMaxPEps); +#if defined(FAST_RCP_ABSMAXP) + return fast_rcp_approx(denom); +#else + return 1.0f / denom; +#endif + } + + template + CUTLASS_DEVICE auto online_softmax_with_quant(TensorAcc& acc, TensorMax& AbsMaxP, + const float softmax_scale_log2) { + Tensor acc_reduction_view = + make_tensor(acc.data(), nvfp4_attention::convert_to_reduction_layout(acc.layout())); + + Tensor acc_conversion_view = + make_tensor(acc.data(), nvfp4_attention::convert_to_conversion_layout(acc.layout())); + + auto temp1 = flatten(acc_conversion_view); + auto temp2 = group_modes<0, 2>(temp1); + auto acc_conversion_flatten = group_modes<1, 5>(temp2); + + if constexpr (FirstTile) { + fill(row_max, -INFINITY); + clear(row_sum); + fill(scores_scale, 1.f); + + CUTLASS_PRAGMA_UNROLL + for (int mi = 0; mi < size<0>(acc_reduction_view); mi++) { + CUTLASS_PRAGMA_UNROLL + for (int ni = 0; ni < size<1, 1>(acc_reduction_view); ni++) { + float local_max = -INFINITY; + CUTLASS_PRAGMA_UNROLL + for (int ei = 0; ei < size<1, 0>(acc_reduction_view); ei++) { + local_max = fmaxf(local_max, acc_reduction_view(mi, make_coord(ei, ni))); + } + + float max_recv = __shfl_xor_sync(int32_t(-1), local_max, 1); + AbsMaxP(mi, ni) = fmaxf(local_max, max_recv); + row_max(mi) = fmaxf(row_max(mi), AbsMaxP(mi, ni)); + } + + row_max(mi) = reduce_row_max_from_pairs(row_max(mi)); + + const float max_scaled = + InfCheck ? (row_max(mi) == -INFINITY + ? 0.f + : (row_max(mi) * softmax_scale_log2 + fp8_scalexfp4_scale_log2)) + : (row_max(mi) * softmax_scale_log2 + fp8_scalexfp4_scale_log2); + +#if defined(DIRECT_P_QUANT_SOFTMAX) + + CUTLASS_PRAGMA_UNROLL + for (int sfi = 0; sfi < size<1>(AbsMaxP); sfi++) { + float chunk_max = AbsMaxP(mi, sfi); + float sfp = 0.0f; + if constexpr (InfCheck) { + if (chunk_max == -INFINITY) { + CUTLASS_PRAGMA_UNROLL + for (int ei = 0; ei < size<1, 0>(acc_reduction_view); ei++) { + acc_reduction_view(mi, make_coord(ei, sfi)) = 0.0f; + } + AbsMaxP(mi, sfi) = 0.0f; + continue; + } + } + float chunk_scaled = chunk_max * softmax_scale_log2; + sfp = softmax_exp2(chunk_scaled - max_scaled + fp4_scale_log2); + AbsMaxP(mi, sfi) = sfp; + CUTLASS_PRAGMA_UNROLL + for (int ei = 0; ei < size<1, 0>(acc_reduction_view); ei++) { + float p = softmax_exp2(acc_reduction_view(mi, make_coord(ei, sfi)) * + softmax_scale_log2 - + chunk_scaled - fp4_scale_log2); + acc_reduction_view(mi, make_coord(ei, sfi)) = p; + row_sum(mi) += p * sfp; + } + } +#else + + CUTLASS_PRAGMA_UNROLL + for (int ni = 0; ni < size<1>(acc_reduction_view); ni++) { + float exp_val = + softmax_exp2(acc_reduction_view(mi, ni) * softmax_scale_log2 - max_scaled); + acc_reduction_view(mi, ni) = exp_val; +#if defined(FIRST_TILE_SUM_IN_EXP) + row_sum(mi) += exp_val; +#endif + } + + CUTLASS_PRAGMA_UNROLL + for (int sfi = 0; sfi < size<1>(AbsMaxP); sfi++) { +#if defined(SFP_FROM_EXP_MAX) + float local_exp_max = 0.0f; + CUTLASS_PRAGMA_UNROLL + for (int ei = 0; ei < size<1, 0>(acc_reduction_view); ei++) { + local_exp_max = fmaxf(local_exp_max, acc_reduction_view(mi, make_coord(ei, sfi))); + } + float peer_exp_max = __shfl_xor_sync(int32_t(-1), local_exp_max, 1); + AbsMaxP(mi, sfi) = fmaxf(local_exp_max, peer_exp_max) * fp4_scale; +#else + AbsMaxP(mi, sfi) = softmax_exp2(AbsMaxP(mi, sfi) * softmax_scale_log2 - + max_scaled + fp4_scale_log2); +#endif + } +#endif + } + +#if !defined(DIRECT_P_QUANT_SOFTMAX) && !defined(FIRST_TILE_SUM_IN_EXP) + + CUTLASS_PRAGMA_UNROLL + for (int mi = 0; mi < size<0>(acc_reduction_view); mi++) { + CUTLASS_PRAGMA_UNROLL + for (int ni = 0; ni < size<1>(acc_reduction_view); ni++) { + row_sum(mi) += acc_reduction_view(mi, ni); + } + } +#endif + } else { + Tensor scores_max_prev = make_fragment_like(row_max); + cute::copy(row_max, scores_max_prev); + + CUTLASS_PRAGMA_UNROLL + for (int mi = 0; mi < size<0>(acc_reduction_view); mi++) { + CUTLASS_PRAGMA_UNROLL + for (int ni = 0; ni < size<1, 1>(acc_reduction_view); ni++) { + float local_max = -INFINITY; + CUTLASS_PRAGMA_UNROLL + for (int ei = 0; ei < size<1, 0>(acc_reduction_view); ei++) { + local_max = fmaxf(local_max, acc_reduction_view(mi, make_coord(ei, ni))); + } + float max_recv = __shfl_xor_sync(int32_t(-1), local_max, 1); + AbsMaxP(mi, ni) = fmaxf(local_max, max_recv); + row_max(mi) = fmaxf(row_max(mi), AbsMaxP(mi, ni)); + } + + row_max(mi) = reduce_row_max_from_pairs(row_max(mi)); + + float scores_max_cur = + !InfCheck ? row_max(mi) : (row_max(mi) == -INFINITY ? 0.0f : row_max(mi)); + scores_scale(mi) = + softmax_exp2((scores_max_prev(mi) - scores_max_cur) * softmax_scale_log2); + + const float max_scaled = + InfCheck ? (row_max(mi) == -INFINITY + ? 0.f + : (row_max(mi) * softmax_scale_log2 + fp8_scalexfp4_scale_log2)) + : (row_max(mi) * softmax_scale_log2 + fp8_scalexfp4_scale_log2); + + row_sum(mi) = row_sum(mi) * scores_scale(mi); + +#if defined(DIRECT_P_QUANT_SOFTMAX) + + CUTLASS_PRAGMA_UNROLL + for (int sfi = 0; sfi < size<1>(AbsMaxP); sfi++) { + float chunk_max = AbsMaxP(mi, sfi); + float sfp = 0.0f; + if constexpr (InfCheck) { + if (chunk_max == -INFINITY) { + CUTLASS_PRAGMA_UNROLL + for (int ei = 0; ei < size<1, 0>(acc_reduction_view); ei++) { + acc_reduction_view(mi, make_coord(ei, sfi)) = 0.0f; + } + AbsMaxP(mi, sfi) = 0.0f; + continue; + } + } + float chunk_scaled = chunk_max * softmax_scale_log2; + sfp = softmax_exp2(chunk_scaled - max_scaled + fp4_scale_log2); + AbsMaxP(mi, sfi) = sfp; + CUTLASS_PRAGMA_UNROLL + for (int ei = 0; ei < size<1, 0>(acc_reduction_view); ei++) { + float p = softmax_exp2(acc_reduction_view(mi, make_coord(ei, sfi)) * + softmax_scale_log2 - + chunk_scaled - fp4_scale_log2); + acc_reduction_view(mi, make_coord(ei, sfi)) = p; + row_sum(mi) += p * sfp; + } + } +#else + + CUTLASS_PRAGMA_UNROLL + for (int ni = 0; ni < size<1>(acc_reduction_view); ni++) { + acc_reduction_view(mi, ni) = + softmax_exp2(acc_reduction_view(mi, ni) * softmax_scale_log2 - max_scaled); + row_sum(mi) += acc_reduction_view(mi, ni); + } + + CUTLASS_PRAGMA_UNROLL + for (int sfi = 0; sfi < size<1>(AbsMaxP); sfi++) { +#if defined(SFP_FROM_EXP_MAX) + float local_exp_max = 0.0f; + CUTLASS_PRAGMA_UNROLL + for (int ei = 0; ei < size<1, 0>(acc_reduction_view); ei++) { + local_exp_max = fmaxf(local_exp_max, acc_reduction_view(mi, make_coord(ei, sfi))); + } + float peer_exp_max = __shfl_xor_sync(int32_t(-1), local_exp_max, 1); + AbsMaxP(mi, sfi) = fmaxf(local_exp_max, peer_exp_max) * fp4_scale; +#else + AbsMaxP(mi, sfi) = softmax_exp2(AbsMaxP(mi, sfi) * softmax_scale_log2 - + max_scaled + fp4_scale_log2); +#endif + } +#endif + } + } + +#if !defined(DIRECT_P_QUANT_SOFTMAX) + +#if defined(SCALAR_INV_ABSMAXP) + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(AbsMaxP); ++i) { + const float inv_absmax = safe_inv_absmax(AbsMaxP(i)); + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < size<0>(acc_conversion_flatten); ++j) { + acc_conversion_flatten(j, i) *= inv_absmax; + } + } +#else + Tensor inv_AbsMaxP = make_tensor_like(AbsMaxP.layout()); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(inv_AbsMaxP); ++i) { + inv_AbsMaxP(i) = safe_inv_absmax(AbsMaxP(i)); + } + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(inv_AbsMaxP); ++i) { + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < size<0>(acc_conversion_flatten); ++j) { + acc_conversion_flatten(j, i) *= inv_AbsMaxP(i); + } + } +#endif +#endif + } + +#if defined(FP16_SOFTMAX) + + template + CUTLASS_DEVICE auto online_softmax_with_quant_fp16(TensorAcc& acc, TensorMax& AbsMaxP, + const float softmax_scale_log2) { + Tensor acc_reduction_view = + make_tensor(acc.data(), nvfp4_attention::convert_to_reduction_layout(acc.layout())); + Tensor acc_conversion_view = + make_tensor(acc.data(), nvfp4_attention::convert_to_conversion_layout(acc.layout())); + auto acc_conversion_flatten = + group_modes<1, 5>(group_modes<0, 2>(flatten(acc_conversion_view))); + + if constexpr (FirstTile) { + fill(row_max, -INFINITY); + clear(row_sum); + fill(scores_scale, 1.f); + + CUTLASS_PRAGMA_UNROLL + for (int mi = 0; mi < size<0>(acc_reduction_view); mi++) { + CUTLASS_PRAGMA_UNROLL + for (int ni = 0; ni < size<1, 1>(acc_reduction_view); ni++) { + float local_max = -INFINITY; + CUTLASS_PRAGMA_UNROLL + for (int ei = 0; ei < size<1, 0>(acc_reduction_view); ei++) { + local_max = fmaxf(local_max, acc_reduction_view(mi, make_coord(ei, ni))); + } + float max_recv = __shfl_xor_sync(int32_t(-1), local_max, 1); + AbsMaxP(mi, ni) = fmaxf(local_max, max_recv); + row_max(mi) = fmaxf(row_max(mi), AbsMaxP(mi, ni)); + } + + row_max(mi) = reduce_row_max_from_pairs(row_max(mi)); + + const float max_scaled = + InfCheck ? (row_max(mi) == -INFINITY + ? 0.f + : (row_max(mi) * softmax_scale_log2 + fp8_scalexfp4_scale_log2)) + : (row_max(mi) * softmax_scale_log2 + fp8_scalexfp4_scale_log2); + + CUTLASS_PRAGMA_UNROLL + for (int ni = 0; ni < size<1>(acc_reduction_view); ni++) { + acc_reduction_view(mi, ni) = + softmax_exp2(acc_reduction_view(mi, ni) * softmax_scale_log2 - max_scaled); + } + + CUTLASS_PRAGMA_UNROLL + for (int sfi = 0; sfi < size<1>(AbsMaxP); sfi++) { + AbsMaxP(mi, sfi) = softmax_exp2(AbsMaxP(mi, sfi) * softmax_scale_log2 - + max_scaled + fp4_scale_log2); + } + } + + CUTLASS_PRAGMA_UNROLL + for (int mi = 0; mi < size<0>(acc_reduction_view); mi++) { + CUTLASS_PRAGMA_UNROLL + for (int ni = 0; ni < size<1>(acc_reduction_view); ni++) { + row_sum(mi) += acc_reduction_view(mi, ni); + } + } + } else { + Tensor scores_max_prev = make_fragment_like(row_max); + cute::copy(row_max, scores_max_prev); + + CUTLASS_PRAGMA_UNROLL + for (int mi = 0; mi < size<0>(acc_reduction_view); mi++) { + CUTLASS_PRAGMA_UNROLL + for (int ni = 0; ni < size<1, 1>(acc_reduction_view); ni++) { + float local_max = -INFINITY; + CUTLASS_PRAGMA_UNROLL + for (int ei = 0; ei < size<1, 0>(acc_reduction_view); ei++) { + local_max = fmaxf(local_max, acc_reduction_view(mi, make_coord(ei, ni))); + } + float max_recv = __shfl_xor_sync(int32_t(-1), local_max, 1); + AbsMaxP(mi, ni) = fmaxf(local_max, max_recv); + row_max(mi) = fmaxf(row_max(mi), AbsMaxP(mi, ni)); + } + + row_max(mi) = reduce_row_max_from_pairs(row_max(mi)); + + float scores_max_cur = + !InfCheck ? row_max(mi) : (row_max(mi) == -INFINITY ? 0.0f : row_max(mi)); + scores_scale(mi) = + softmax_exp2((scores_max_prev(mi) - scores_max_cur) * softmax_scale_log2); + + const float max_scaled = + InfCheck ? (row_max(mi) == -INFINITY + ? 0.f + : (row_max(mi) * softmax_scale_log2 + fp8_scalexfp4_scale_log2)) + : (row_max(mi) * softmax_scale_log2 + fp8_scalexfp4_scale_log2); + + row_sum(mi) = row_sum(mi) * scores_scale(mi); + + CUTLASS_PRAGMA_UNROLL + for (int ni = 0; ni < size<1>(acc_reduction_view); ni++) { + acc_reduction_view(mi, ni) = + softmax_exp2(acc_reduction_view(mi, ni) * softmax_scale_log2 - max_scaled); + row_sum(mi) += acc_reduction_view(mi, ni); + } + + CUTLASS_PRAGMA_UNROLL + for (int sfi = 0; sfi < size<1>(AbsMaxP); sfi++) { + AbsMaxP(mi, sfi) = softmax_exp2(AbsMaxP(mi, sfi) * softmax_scale_log2 - + max_scaled + fp4_scale_log2); + } + } + } + +#if defined(SCALAR_INV_ABSMAXP) + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(AbsMaxP); ++i) { + const float inv_absmax = safe_inv_absmax(AbsMaxP(i)); + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < size<0>(acc_conversion_flatten); ++j) { + acc_conversion_flatten(j, i) *= inv_absmax; + } + } +#else + Tensor inv_AbsMaxP = make_tensor_like(AbsMaxP.layout()); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(inv_AbsMaxP); ++i) { + inv_AbsMaxP(i) = safe_inv_absmax(AbsMaxP(i)); + } + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(inv_AbsMaxP); ++i) { + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < size<0>(acc_conversion_flatten); ++j) { + acc_conversion_flatten(j, i) *= inv_AbsMaxP(i); + } + } +#endif + } +#endif + + template + CUTLASS_DEVICE void finalize(TensorAcc& o_store) { + Tensor o_store_reduction_view = + make_tensor(o_store.data(), convert_to_reduction_layout(o_store.layout())); + + CUTLASS_PRAGMA_UNROLL + for (int mi = 0; mi < size(row_max); ++mi) { + CUTLASS_PRAGMA_UNROLL + for (int i = 1; i < RowReductionThr; i <<= 1) { + float sum_recv = __shfl_xor_sync(int32_t(-1), row_sum(mi), i); + row_sum(mi) += sum_recv; + } + + float sum = row_sum(mi); + + float inv_sum = (sum == 0.f || sum != sum) ? 0.f : 1.f / sum; + + CUTLASS_PRAGMA_UNROLL + for (int ni = 0; ni < size<1>(o_store_reduction_view); ++ni) { + o_store_reduction_view(mi, ni) *= inv_sum; + } + } + } + + template + CUTLASS_DEVICE void rescale_o(TensorAcc& o_store, TensorAcc const& o_tmp) { + Tensor o_store_reduction_view = + make_tensor(o_store.data(), nvfp4_attention::convert_to_reduction_layout(o_store.layout())); + Tensor o_tmp_reduction_view = + make_tensor(o_tmp.data(), nvfp4_attention::convert_to_reduction_layout(o_tmp.layout())); + + CUTLASS_PRAGMA_UNROLL + for (int mi = 0; mi < size(row_max); ++mi) { + CUTLASS_PRAGMA_UNROLL + for (int ni = 0; ni < size<1>(o_store_reduction_view); ++ni) { + o_store_reduction_view(mi, ni) = + o_store_reduction_view(mi, ni) * scores_scale(mi) + o_tmp_reduction_view(mi, ni); + } + } + } + + template + CUTLASS_DEVICE void find_max_chunk(TensorAcc& acc, TensorMax& AbsMaxP, int ni) { + Tensor acc_rv = + make_tensor(acc.data(), nvfp4_attention::convert_to_reduction_layout(acc.layout())); + + CUTLASS_PRAGMA_UNROLL + for (int mi = 0; mi < size<0>(acc_rv); mi++) { + float chunk_max = -INFINITY; + CUTLASS_PRAGMA_UNROLL + for (int ei = 0; ei < size<1, 0>(acc_rv); ei++) { + chunk_max = fmaxf(chunk_max, acc_rv(mi, make_coord(ei, ni))); + } + float max_recv = __shfl_xor_sync(int32_t(-1), chunk_max, 1); + chunk_max = fmaxf(chunk_max, max_recv); + AbsMaxP(mi, ni) = chunk_max; + + row_max(mi) = fmaxf(row_max(mi), chunk_max); + row_max(mi) = reduce_row_max_from_pairs(row_max(mi)); + } + } + + template + CUTLASS_DEVICE void exp2_sum_chunk(TensorAcc& acc, TensorMax& AbsMaxP, int ni, + const float softmax_scale_log2) { + Tensor acc_rv = + make_tensor(acc.data(), nvfp4_attention::convert_to_reduction_layout(acc.layout())); + + CUTLASS_PRAGMA_UNROLL + for (int mi = 0; mi < size<0>(acc_rv); mi++) { + const float max_scaled = + InfCheck ? (row_max(mi) == -INFINITY + ? 0.f + : (row_max(mi) * softmax_scale_log2 + fp8_scalexfp4_scale_log2)) + : (row_max(mi) * softmax_scale_log2 + fp8_scalexfp4_scale_log2); + + CUTLASS_PRAGMA_UNROLL + for (int ei = 0; ei < size<1, 0>(acc_rv); ei++) { + float val = softmax_exp2(acc_rv(mi, make_coord(ei, ni)) * softmax_scale_log2 - + max_scaled); + acc_rv(mi, make_coord(ei, ni)) = val; + row_sum(mi) += val; + } + + AbsMaxP(mi, ni) = softmax_exp2(AbsMaxP(mi, ni) * softmax_scale_log2 - max_scaled + + fp4_scale_log2); + } + } + + template + CUTLASS_DEVICE void quantize_after_partial_softmax(TensorAcc& acc, TensorMax& AbsMaxP) { + Tensor acc_cv = + make_tensor(acc.data(), nvfp4_attention::convert_to_conversion_layout(acc.layout())); + auto temp1 = flatten(acc_cv); + auto temp2 = group_modes<0, 2>(temp1); + auto acc_flat = group_modes<1, 5>(temp2); + + Tensor inv_AbsMaxP = make_tensor_like(AbsMaxP.layout()); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(inv_AbsMaxP); ++i) { + inv_AbsMaxP(i) = safe_inv_absmax(AbsMaxP(i)); + } + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(inv_AbsMaxP); ++i) { + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < size<0>(acc_flat); ++j) { + acc_flat(j, i) *= inv_AbsMaxP(i); + } + } + } + + template + CUTLASS_DEVICE void chunked_softmax_fixed(TensorAcc& acc, TensorMax& AbsMaxP, bool is_first, + const float softmax_scale_log2, + TensorPrev const& scores_max_prev) { + Tensor acc_rv = + make_tensor(acc.data(), nvfp4_attention::convert_to_reduction_layout(acc.layout())); + + Tensor acc_cv = + make_tensor(acc.data(), nvfp4_attention::convert_to_conversion_layout(acc.layout())); + auto acc_cv_flat = group_modes<1, 5>(group_modes<0, 2>(flatten(acc_cv))); + + constexpr int MmaN = decltype(size<1, 1>(acc_rv))::value; + + if (is_first) { + fill(row_max, -INFINITY); + clear(row_sum); + fill(scores_scale, 1.f); + } + + CUTLASS_PRAGMA_UNROLL + for (int mi = 0; mi < size<0>(acc_rv); mi++) { + CUTLASS_PRAGMA_UNROLL + for (int ni = 0; ni < MmaN; ni++) { + float local_max = -INFINITY; + CUTLASS_PRAGMA_UNROLL + for (int ei = 0; ei < size<1, 0>(acc_rv); ei++) { + local_max = fmaxf(local_max, acc_rv(mi, make_coord(ei, ni))); + } + float max_recv = __shfl_xor_sync(int32_t(-1), local_max, 1); + AbsMaxP(mi, ni) = fmaxf(local_max, max_recv); + row_max(mi) = fmaxf(row_max(mi), AbsMaxP(mi, ni)); + } + row_max(mi) = reduce_row_max_from_pairs(row_max(mi)); + } + + if (!is_first) { + CUTLASS_PRAGMA_UNROLL + for (int mi = 0; mi < size<0>(acc_rv); mi++) { + float scores_max_cur = + !InfCheck ? row_max(mi) : (row_max(mi) == -INFINITY ? 0.0f : row_max(mi)); + scores_scale(mi) = + softmax_exp2((scores_max_prev(mi) - scores_max_cur) * softmax_scale_log2); + row_sum(mi) *= scores_scale(mi); + } + } + + CUTLASS_PRAGMA_UNROLL + for (int mi = 0; mi < size<0>(acc_rv); mi++) { + const float max_scaled = + InfCheck ? (row_max(mi) == -INFINITY + ? 0.f + : (row_max(mi) * softmax_scale_log2 + fp8_scalexfp4_scale_log2)) + : (row_max(mi) * softmax_scale_log2 + fp8_scalexfp4_scale_log2); + + CUTLASS_PRAGMA_UNROLL + for (int ni = 0; ni < size<1>(acc_rv); ni++) { + float val = softmax_exp2(acc_rv(mi, ni) * softmax_scale_log2 - max_scaled); + acc_rv(mi, ni) = val; + row_sum(mi) += val; + } + + CUTLASS_PRAGMA_UNROLL + for (int sfi = 0; sfi < size<1>(AbsMaxP); sfi++) { + AbsMaxP(mi, sfi) = softmax_exp2(AbsMaxP(mi, sfi) * softmax_scale_log2 - + max_scaled + fp4_scale_log2); + } + } + + Tensor inv_AbsMaxP = make_tensor_like(AbsMaxP.layout()); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(inv_AbsMaxP); ++i) { + inv_AbsMaxP(i) = safe_inv_absmax(AbsMaxP(i)); + } + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(inv_AbsMaxP); ++i) { + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < size<0>(acc_cv_flat); ++j) { + acc_cv_flat(j, i) *= inv_AbsMaxP(i); + } + } + } + + template + CUTLASS_DEVICE void exp2_sum_and_quantize(TensorAcc& acc, TensorMax& AbsMaxP, bool is_first, + const float softmax_scale_log2, + TensorPrev const& scores_max_prev) { + Tensor acc_rv = + make_tensor(acc.data(), nvfp4_attention::convert_to_reduction_layout(acc.layout())); + Tensor acc_cv = + make_tensor(acc.data(), nvfp4_attention::convert_to_conversion_layout(acc.layout())); + auto acc_cv_flat = group_modes<1, 5>(group_modes<0, 2>(flatten(acc_cv))); + + if (!is_first) { + CUTLASS_PRAGMA_UNROLL + for (int mi = 0; mi < size<0>(acc_rv); mi++) { + float scores_max_cur = + !InfCheck ? row_max(mi) : (row_max(mi) == -INFINITY ? 0.0f : row_max(mi)); + scores_scale(mi) = + softmax_exp2((scores_max_prev(mi) - scores_max_cur) * softmax_scale_log2); + row_sum(mi) *= scores_scale(mi); + } + } + + CUTLASS_PRAGMA_UNROLL + for (int mi = 0; mi < size<0>(acc_rv); mi++) { + const float max_scaled = + InfCheck ? (row_max(mi) == -INFINITY + ? 0.f + : (row_max(mi) * softmax_scale_log2 + fp8_scalexfp4_scale_log2)) + : (row_max(mi) * softmax_scale_log2 + fp8_scalexfp4_scale_log2); + + CUTLASS_PRAGMA_UNROLL + for (int ni = 0; ni < size<1>(acc_rv); ni++) { + float val = softmax_exp2(acc_rv(mi, ni) * softmax_scale_log2 - max_scaled); + acc_rv(mi, ni) = val; + row_sum(mi) += val; + } + + CUTLASS_PRAGMA_UNROLL + for (int sfi = 0; sfi < size<1>(AbsMaxP); sfi++) { + AbsMaxP(mi, sfi) = softmax_exp2(AbsMaxP(mi, sfi) * softmax_scale_log2 - + max_scaled + fp4_scale_log2); + } + } + + Tensor inv_AbsMaxP = make_tensor_like(AbsMaxP.layout()); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(inv_AbsMaxP); ++i) { + inv_AbsMaxP(i) = safe_inv_absmax(AbsMaxP(i)); + } + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(inv_AbsMaxP); ++i) { + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < size<0>(acc_cv_flat); ++j) { + acc_cv_flat(j, i) *= inv_AbsMaxP(i); + } + } + } + + template + CUTLASS_DEVICE void online_softmax_chunk(TensorAcc& acc, TensorMax& AbsMaxP, int ni, + const float softmax_scale_log2) { + Tensor acc_rv = + make_tensor(acc.data(), nvfp4_attention::convert_to_reduction_layout(acc.layout())); + + if constexpr (IsInit) { + fill(row_max, -INFINITY); + clear(row_sum); + fill(scores_scale, 1.f); + } + + CUTLASS_PRAGMA_UNROLL + for (int mi = 0; mi < size<0>(acc_rv); mi++) { + float chunk_max = -INFINITY; + CUTLASS_PRAGMA_UNROLL + for (int ei = 0; ei < size<1, 0>(acc_rv); ei++) { + chunk_max = fmaxf(chunk_max, acc_rv(mi, make_coord(ei, ni))); + } + float max_recv = __shfl_xor_sync(int32_t(-1), chunk_max, 1); + chunk_max = fmaxf(chunk_max, max_recv); + AbsMaxP(mi, ni) = chunk_max; + + float prev_max = row_max(mi); + row_max(mi) = fmaxf(row_max(mi), chunk_max); + row_max(mi) = reduce_row_max_from_pairs(row_max(mi)); + + const float max_scaled = + InfCheck ? (row_max(mi) == -INFINITY + ? 0.f + : (row_max(mi) * softmax_scale_log2 + fp8_scalexfp4_scale_log2)) + : (row_max(mi) * softmax_scale_log2 + fp8_scalexfp4_scale_log2); + + if constexpr (!IsInit) { + if (prev_max != row_max(mi)) { + scores_scale(mi) = softmax_exp2((prev_max - row_max(mi)) * softmax_scale_log2); + row_sum(mi) *= scores_scale(mi); + CUTLASS_PRAGMA_UNROLL + for (int prev_ni = 0; prev_ni < ni; prev_ni++) { + CUTLASS_PRAGMA_UNROLL + for (int ei = 0; ei < size<1, 0>(acc_rv); ei++) { + acc_rv(mi, make_coord(ei, prev_ni)) *= scores_scale(mi); + } + AbsMaxP(mi, prev_ni) *= scores_scale(mi); + } + } + } + + CUTLASS_PRAGMA_UNROLL + for (int ei = 0; ei < size<1, 0>(acc_rv); ei++) { + float val = softmax_exp2(acc_rv(mi, make_coord(ei, ni)) * softmax_scale_log2 - + max_scaled); + acc_rv(mi, make_coord(ei, ni)) = val; + row_sum(mi) += val; + } + + AbsMaxP(mi, ni) = softmax_exp2(AbsMaxP(mi, ni) * softmax_scale_log2 - max_scaled + + fp4_scale_log2); + } + } + + private: + __device__ __forceinline__ static float ptx_exp2(float x) { + float result; + asm volatile("ex2.approx.ftz.f32 %0, %1;" : "=f"(result) : "f"(x)); + return result; + } + + template + __device__ __forceinline__ static float softmax_exp2(float x) { + if (x <= -126.0f) { + return 0.0f; + } +#if defined(SOFTMAX_FMA_EXP2) + if constexpr (!InfCheck) { + return nvfp4_attention::exp2_fma_poly(x); + } +#endif + return ptx_exp2(x); + } +}; + +} // namespace nvfp4_attention diff --git a/include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/epilogue.cuh b/include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/epilogue.cuh new file mode 100644 index 00000000000..3a0271b1e52 --- /dev/null +++ b/include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/epilogue.cuh @@ -0,0 +1,217 @@ +/* + * Copyright (c) 2025 by SageAttention team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include +#include + +#include "../primitives/barrier.cuh" +#include "../utils/copy.cuh" +#include "../utils/math.cuh" +#include "cute/tensor.hpp" +#include "cutlass/gemm/collective/collective_builder.hpp" + +namespace nvfp4_attention { + +using namespace cute; + +template +CUTLASS_DEVICE void copy_with_bounds_check(TiledCopy tiled_copy, Tensor const& S, + Tensor& D, + Tensor const& identity_MN, + Tensor const& predicate_K, + const int max_MN = 0) { + CUTE_STATIC_ASSERT_V(rank(S) == Int<3>{}); + CUTE_STATIC_ASSERT_V(rank(D) == Int<3>{}); + CUTE_STATIC_ASSERT_V(size<0>(S) == size<0>(D)); + CUTE_STATIC_ASSERT_V(size<1>(S) == size<1>(D)); + CUTE_STATIC_ASSERT_V(size<2>(S) == size<2>(D)); + static_assert(!(Clear_OOB_MN && !Clear_OOB_K)); +#pragma unroll + for (int m = 0; m < size<1>(S); ++m) { + if (Is_even_MN || get<0>(identity_MN(0, m, 0)) < max_MN) { +#pragma unroll + for (int k = 0; k < size<2>(S); ++k) { + if (Is_even_K || predicate_K(k)) { + cute::copy(tiled_copy, S(_, m, k), D(_, m, k)); + } else if (Clear_OOB_K) { + cute::clear(D(_, m, k)); + } + } + } else if (Clear_OOB_MN) { + cute::clear(D(_, m, _)); + } + } +} + +template +struct CollectiveEpilogueFwd { + using Element = typename Ktraits::ElementOut; + static constexpr int kBlockM = Ktraits::kBlockM; + static constexpr int kBlockN = Ktraits::kBlockN; + static constexpr int kHeadDim = Ktraits::kHeadDim; + static constexpr int kStoreBlockM = Ktraits::kStoreBlockM; + using TileShape_MNK = Shape, Int, Int>; + using TileShape_O = Shape, Int>; + static constexpr int kNWarps = Ktraits::kNWarps; + static constexpr int kNThreads = kNWarps * cutlass::NumThreadsPerWarp; + static constexpr int NumMmaThreads = kNThreads - cutlass::NumThreadsPerWarpGroup; + + using GmemTiledCopyOTMA = cute::SM90_TMA_STORE; + + static constexpr int kGmemElemsPerLoad = sizeof(cute::uint128_t) / sizeof(Element); + static_assert(kHeadDim % kGmemElemsPerLoad == 0, + "kHeadDim must be a multiple of kGmemElemsPerLoad"); + static constexpr int kGmemThreadsPerRow = kHeadDim / kGmemElemsPerLoad; + static_assert(NumMmaThreads % kGmemThreadsPerRow == 0, + "NumMmaThreads must be a multiple of kGmemThreadsPerRow"); + using GmemLayoutAtom = + Layout, Int>, + Stride, _1>>; + using GmemTiledCopyO = + decltype(make_tiled_copy(Copy_Atom{}, GmemLayoutAtom{}, + Layout>>{})); + + using SmemLayoutO = typename Ktraits::SmemLayoutO; + + using SmemCopyAtomO = Copy_Atom; + using SharedStorage = cute::array_aligned>; + + using ShapeO = cute::Shape; + using StrideO = cute::Stride; + using StrideLSE = cute::Stride<_1, int64_t, int64_t>; + + using TMA_O = decltype(make_tma_copy(GmemTiledCopyOTMA{}, + make_tensor(make_gmem_ptr(static_cast(nullptr)), + repeat_like(StrideO{}, int32_t(0)), StrideO{}), + SmemLayoutO{}, TileShape_O{}, _1{})); + + struct Arguments { + Element* ptr_O; + ShapeO const shape_O; + StrideO const stride_O; + float* ptr_LSE; + StrideLSE const stride_LSE; + }; + + struct Params { + Element* ptr_O; + ShapeO const shape_O; + StrideO const stride_O; + float* ptr_LSE; + StrideLSE const stride_LSE; + TMA_O tma_store_O; + }; + + static Params to_underlying_arguments(Arguments const& args) { + Tensor mO = make_tensor(make_gmem_ptr(args.ptr_O), args.shape_O, args.stride_O); + TMA_O tma_store_O = make_tma_copy(GmemTiledCopyOTMA{}, mO, SmemLayoutO{}, TileShape_O{}, _1{}); + return {args.ptr_O, args.shape_O, args.stride_O, args.ptr_LSE, args.stride_LSE, tma_store_O}; + } + + CUTLASS_DEVICE + static void prefetch_tma_descriptors(Params const& epilogue_params) { + cute::prefetch_tma_descriptor(epilogue_params.tma_store_O.get_tma_descriptor()); + } + + template + CUTLASS_DEVICE void mma_store(SharedStorage& shared_storage, TiledMma tiled_mma, + FrgTensorO const& tOrO, int thread_idx, int wg_id = 0) { + using TiledMmaPV_Store = typename Ktraits::TiledMmaPV_Store; + static constexpr int NumMmaThreads = size(TiledMma{}); + TiledMmaPV_Store tiled_mma_pv_store; + int consumer_thread_idx_full = thread_idx + wg_id * NumMmaThreads; + + Tensor sO = cute::as_position_independent_swizzle_tensor( + make_tensor(make_smem_ptr(shared_storage.smem_o.begin()), SmemLayoutO{})); + auto smem_tiled_copy_O = make_tiled_copy_C(SmemCopyAtomO{}, tiled_mma_pv_store); + auto smem_thr_copy_O = smem_tiled_copy_O.get_thread_slice(consumer_thread_idx_full); + constexpr int numel = decltype(size(tOrO))::value; + cutlass::NumericArrayConverter convert_op; + + auto frag = convert_op(*reinterpret_cast*>(tOrO.data())); + auto tOrO_out = make_tensor(make_rmem_ptr(&frag), tOrO.layout()); + Tensor taccOrO = smem_thr_copy_O.retile_S(tOrO_out); + Tensor taccOsO = smem_thr_copy_O.partition_D(sO); + cute::copy(smem_tiled_copy_O, taccOrO, taccOsO); + cutlass::arch::fence_view_async_shared(); + } + + template + CUTLASS_DEVICE void tma_store(SharedStorage& shared_storage, Params const& epilogue_params, + WorkTileInfo work_tile_info, + SchedulerParams const& scheduler_params, int thread_idx, + int store_m_subtile = 0) { + auto [m_block, bidh, bidb] = work_tile_info.get_block_coord(scheduler_params); + Tensor sO = cute::as_position_independent_swizzle_tensor( + make_tensor(make_smem_ptr(shared_storage.smem_o.begin()), SmemLayoutO{})); + Tensor mO = epilogue_params.tma_store_O.get_tma_tensor(epilogue_params.shape_O); + constexpr int StoreTilesPerMBlock = kBlockM / kStoreBlockM; + Tensor gO = local_tile(mO(_, _, bidh, bidb), TileShape_O{}, + make_coord(m_block * StoreTilesPerMBlock + store_m_subtile, _0{})); + auto block_tma_O = epilogue_params.tma_store_O.get_slice(_0{}); + Tensor tOgO = block_tma_O.partition_D(gO); + Tensor tOsO = block_tma_O.partition_S(sO); + + cute::copy(epilogue_params.tma_store_O, tOsO, tOgO); + tma_store_arrive(); + } + + CUTLASS_DEVICE void store_tail() { tma_store_wait<0>(); } + + CUTLASS_DEVICE void store_zero(Params const& epilogue_params, int thread_idx, + cute::tuple const& block_coord) { + auto [m_block, bidh, bidb] = block_coord; + Tensor mO = make_tensor(make_gmem_ptr(epilogue_params.ptr_O), epilogue_params.shape_O, + epilogue_params.stride_O); + Tensor gO = + local_tile(mO(_, _, bidh, bidb), select<0, 2>(TileShape_MNK{}), make_coord(m_block, _0{})); + auto shape_LSE = select<0, 2, 3>(epilogue_params.shape_O); + Tensor mLSE = + make_tensor(make_gmem_ptr(epilogue_params.ptr_LSE), shape_LSE, epilogue_params.stride_LSE); + Tensor gLSE = local_tile(mLSE(_, bidh, bidb), Shape>{}, make_coord(m_block)); + + GmemTiledCopyO gmem_tiled_copy_O; + auto gmem_thr_copy_O = gmem_tiled_copy_O.get_thread_slice(thread_idx); + Tensor tOgO = gmem_thr_copy_O.partition_D(gO); + Tensor tOrO = make_fragment_like(tOgO); + clear(tOrO); + + Tensor cO = cute::make_identity_tensor(select<0, 2>(TileShape_MNK{})); + + Tensor tOcO = gmem_thr_copy_O.partition_D(cO); + Tensor tOpO = make_tensor(make_shape(size<2>(tOgO))); +#pragma unroll + for (int k = 0; k < size(tOpO); ++k) { + tOpO(k) = get<1>(tOcO(_0{}, _0{}, k)) < get<1>(epilogue_params.shape_O); + } + + copy_with_bounds_check( + gmem_tiled_copy_O, tOrO, tOgO, tOcO, tOpO, + get<0>(epilogue_params.shape_O) - m_block * kBlockM); + static_assert(kBlockM <= NumMmaThreads); + if (thread_idx < get<0>(shape_LSE) - m_block * kBlockM) { + gLSE(thread_idx) = INFINITY; + } + } +}; + +} // namespace nvfp4_attention diff --git a/include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/epilogue/lse_writer.cuh b/include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/epilogue/lse_writer.cuh new file mode 100644 index 00000000000..b68a0eafa95 --- /dev/null +++ b/include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/epilogue/lse_writer.cuh @@ -0,0 +1,142 @@ +/* + * Copyright (c) 2025 by SageAttention team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "cute/tensor.hpp" +#include "cutlass/cutlass.h" + +namespace nvfp4_attention { + +using namespace cute; + +template +struct LSEWriter { + using TileShape_MNK = typename Traits::TileShape_MNK; + + static constexpr int kBlockM = get<0>(TileShape_MNK{}); + static constexpr int kBlockN = get<1>(TileShape_MNK{}); + static constexpr int kHeadDim = get<2>(TileShape_MNK{}); + static constexpr int kNWarps = Traits::kNWarps; + static constexpr int kNThreads = kNWarps * cutlass::NumThreadsPerWarp; + static constexpr int NumMmaThreads = kNThreads - cutlass::NumThreadsPerWarpGroup; + + using ShapeLSE = cute::Shape; + using StrideLSE = cute::Stride<_1, int64_t, int64_t>; + + template + __device__ __forceinline__ static void write_lse(float* ptr_LSE, Shape const& shape_LSE, + Stride const& stride_LSE, + SoftmaxFused const& softmax_fused, + float softmax_scale_log2, + TiledMma const& tiled_mma, int thread_idx, + int m_block, int bidh, int bidb) { + Tensor mLSE = make_tensor(make_gmem_ptr(ptr_LSE), shape_LSE, stride_LSE); + Tensor gLSE = + local_tile(mLSE(_, bidh, bidb), cute::Shape>{}, make_coord(m_block)); + + auto const& row_max = softmax_fused.row_max; + auto const& row_sum = softmax_fused.row_sum; + + Tensor caccO = cute::make_identity_tensor(select<0, 2>(TileShape_MNK{})); + auto thread_mma = tiled_mma.get_thread_slice(thread_idx); + Tensor taccOcO = thread_mma.partition_C(caccO); + + static_assert(decltype(size<0, 0>(taccOcO))::value == 2); + static_assert(decltype(size<0, 1>(taccOcO))::value == 2); + + Tensor taccOcO_row = taccOcO(make_coord(_0{}, _), _, _0{}); + CUTE_STATIC_ASSERT_V(size(row_max) == size(taccOcO_row)); + + if (get<1>(taccOcO_row(_0{})) == 0) { + constexpr float log2_e = 1.44269504088896340736f; + constexpr float ln_2 = 0.69314718055994530942f; + +#pragma unroll + for (int mi = 0; mi < size(row_max); ++mi) { + const int row = get<0>(taccOcO_row(mi)); + + if (row < get<0>(shape_LSE) - m_block * kBlockM) { + float max_scaled = row_max(mi) * softmax_scale_log2 / log2_e; + float sum = row_sum(mi); + + float lse = (sum == 0.f || sum != sum) ? INFINITY : (max_scaled + logf(sum)); + + gLSE(row) = lse; + } + } + } + } + + template + __device__ __forceinline__ static void write_lse_infinity(float* ptr_LSE, ShapeO const& shape_O, + Stride const& stride_LSE, + int thread_idx, int m_block, int bidh, + int bidb) { + auto shape_LSE = select<0, 2, 3>(shape_O); + + Tensor mLSE = make_tensor(make_gmem_ptr(ptr_LSE), shape_LSE, stride_LSE); + Tensor gLSE = local_tile(mLSE(_, bidh, bidb), Shape>{}, make_coord(m_block)); + + static_assert(kBlockM <= NumMmaThreads); + + if (thread_idx < get<0>(shape_LSE) - m_block * kBlockM) { + gLSE(thread_idx) = INFINITY; + } + } + + template + __device__ __forceinline__ static void run(float* ptr_LSE, ShapeO const& shape_O, + ShapeLSE const& shape_LSE, Stride const& stride_LSE, + SoftmaxFused const& softmax_fused, + float softmax_scale_log2, TiledMma const& tiled_mma, + int thread_idx, int m_block, int bidh, int bidb, + bool is_valid_block) { + if (is_valid_block) { + write_lse(ptr_LSE, shape_LSE, stride_LSE, softmax_fused, softmax_scale_log2, tiled_mma, + thread_idx, m_block, bidh, bidb); + } else { + write_lse_infinity(ptr_LSE, shape_O, stride_LSE, thread_idx, m_block, bidh, bidb); + } + } +}; + +__device__ __forceinline__ float compute_lse_log2(float row_max, float row_sum, + float softmax_scale_log2) { + constexpr float log2_e = 1.44269504088896340736f; + + if (row_sum == 0.f || row_sum != row_sum) { + return INFINITY; + } + + float max_scaled = row_max * softmax_scale_log2; + float lse_log2 = max_scaled + log2f(row_sum); + + return lse_log2; +} + +__device__ __forceinline__ float log2_to_ln(float x_log2) { + constexpr float ln_2 = 0.69314718055994530942f; + return x_log2 * ln_2; +} + +__device__ __forceinline__ float ln_to_log2(float x_ln) { + constexpr float log2_e = 1.44269504088896340736f; + return x_ln * log2_e; +} + +} // namespace nvfp4_attention diff --git a/include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/epilogue/output_writer.cuh b/include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/epilogue/output_writer.cuh new file mode 100644 index 00000000000..84b049eb507 --- /dev/null +++ b/include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/epilogue/output_writer.cuh @@ -0,0 +1,187 @@ +/* + * Copyright (c) 2025 by SageAttention team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "cute/tensor.hpp" +#include "cutlass/arch/barrier.h" +#include "cutlass/cutlass.h" +#include "cutlass/numeric_conversion.h" + +namespace nvfp4_attention { + +using namespace cute; + +template +struct OutputWriter { + using Element = typename Traits::ElementOut; + using TileShape_MNK = typename Traits::TileShape_MNK; + using SmemLayoutO = typename Traits::SmemLayoutO; + using SmemCopyAtomO = typename Traits::SmemCopyAtomO; + using GmemTiledCopyOTMA = cute::SM90_TMA_STORE; + + static constexpr int kBlockM = get<0>(TileShape_MNK{}); + static constexpr int kBlockN = get<1>(TileShape_MNK{}); + static constexpr int kHeadDim = get<2>(TileShape_MNK{}); + static constexpr int kNWarps = Traits::kNWarps; + static constexpr int kNThreads = kNWarps * cutlass::NumThreadsPerWarp; + static constexpr int NumMmaThreads = kNThreads - cutlass::NumThreadsPerWarpGroup; + + using ShapeO = cute::Shape; + using StrideO = cute::Stride; + + using TMA_O = decltype(make_tma_copy(GmemTiledCopyOTMA{}, + make_tensor(make_gmem_ptr(static_cast(nullptr)), + repeat_like(StrideO{}, int32_t(0)), StrideO{}), + SmemLayoutO{}, select<0, 2>(TileShape_MNK{}), _1{})); + + template + __device__ __forceinline__ static void prefetch_tma_descriptor(TMA const& tma_store_O) { + cute::prefetch_tma_descriptor(tma_store_O.get_tma_descriptor()); + } + + template + __device__ __forceinline__ static void register_to_smem(SharedStorage& shared_storage, + TiledMma const& tiled_mma, + FrgTensorO const& tOrO, int thread_idx) { + Tensor sO = cute::as_position_independent_swizzle_tensor( + make_tensor(make_smem_ptr(shared_storage.smem_o.begin()), SmemLayoutO{})); + + auto smem_tiled_copy_O = make_tiled_copy_C(SmemCopyAtomO{}, tiled_mma); + auto smem_thr_copy_O = smem_tiled_copy_O.get_thread_slice(thread_idx); + + constexpr int numel = decltype(size(tOrO))::value; + cutlass::NumericArrayConverter convert_op; + + auto frag = convert_op(*reinterpret_cast*>(tOrO.data())); + auto tOrO_out = make_tensor(make_rmem_ptr(&frag), tOrO.layout()); + + Tensor taccOrO = smem_thr_copy_O.retile_S(tOrO_out); + Tensor taccOsO = smem_thr_copy_O.partition_D(sO); + cute::copy(smem_tiled_copy_O, taccOrO, taccOsO); + + cutlass::arch::fence_view_async_shared(); + } + + template + __device__ __forceinline__ static void smem_to_gmem(SharedStorage& shared_storage, + TMA const& tma_store_O, Shape const& shape_O, + Stride const& stride_O, int m_block, int bidh, + int bidb) { + Tensor sO = cute::as_position_independent_swizzle_tensor( + make_tensor(make_smem_ptr(shared_storage.smem_o.begin()), SmemLayoutO{})); + + Tensor mO = tma_store_O.get_tma_tensor(shape_O); + Tensor gO = + local_tile(mO(_, _, bidh, bidb), select<0, 2>(TileShape_MNK{}), make_coord(m_block, _0{})); + + auto block_tma_O = tma_store_O.get_slice(_0{}); + Tensor tOgO = block_tma_O.partition_D(gO); + Tensor tOsO = block_tma_O.partition_S(sO); + + cute::copy(tma_store_O, tOsO, tOgO); + + tma_store_arrive(); + } + + template + __device__ __forceinline__ static void run(SharedStorage& shared_storage, + TiledMma const& tiled_mma, FrgTensorO const& tOrO, + TMA const& tma_store_O, Shape const& shape_O, + Stride const& stride_O, int thread_idx, int m_block, + int bidh, int bidb) { + register_to_smem(shared_storage, tiled_mma, tOrO, thread_idx); + + smem_to_gmem(shared_storage, tma_store_O, shape_O, stride_O, m_block, bidh, bidb); + } + + __device__ __forceinline__ static void wait_all_stores() { tma_store_wait<0>(); } + + template + __device__ __forceinline__ static void store_zero(Element* ptr_O, Shape const& shape_O, + Stride const& stride_O, int thread_idx, + int m_block, int bidh, int bidb) { + Tensor mO = make_tensor(make_gmem_ptr(ptr_O), shape_O, stride_O); + Tensor gO = + local_tile(mO(_, _, bidh, bidb), select<0, 2>(TileShape_MNK{}), make_coord(m_block, _0{})); + + static constexpr int kGmemElemsPerLoad = sizeof(cute::uint128_t) / sizeof(Element); + static_assert(kHeadDim % kGmemElemsPerLoad == 0, + "kHeadDim must be a multiple of kGmemElemsPerLoad"); + static constexpr int kGmemThreadsPerRow = kHeadDim / kGmemElemsPerLoad; + static_assert(NumMmaThreads % kGmemThreadsPerRow == 0, + "NumMmaThreads must be a multiple of kGmemThreadsPerRow"); + + using GmemLayoutAtom = Layout< + cute::Shape, cute::Int>, + cute::Stride, cute::_1>>; + using GmemTiledCopyO = + decltype(make_tiled_copy(Copy_Atom{}, GmemLayoutAtom{}, + Layout>>{})); + + GmemTiledCopyO gmem_tiled_copy_O; + auto gmem_thr_copy_O = gmem_tiled_copy_O.get_thread_slice(thread_idx); + + Tensor tOgO = gmem_thr_copy_O.partition_D(gO); + Tensor tOrO = make_fragment_like(tOgO); + clear(tOrO); + + Tensor cO = cute::make_identity_tensor(select<0, 2>(TileShape_MNK{})); + Tensor tOcO = gmem_thr_copy_O.partition_D(cO); + Tensor tOpO = make_tensor(make_shape(size<2>(tOgO))); +#pragma unroll + for (int k = 0; k < size(tOpO); ++k) { + tOpO(k) = get<1>(tOcO(_0{}, _0{}, k)) < get<1>(shape_O); + } + + copy_with_predicate(gmem_tiled_copy_O, tOrO, tOgO, tOcO, tOpO, + get<0>(shape_O) - m_block * kBlockM); + } + + private: + template + __device__ __forceinline__ static void copy_with_predicate(TiledCopy const& tiled_copy, + Tensor1 const& src, Tensor2& dst, + Tensor3 const& coord, + Tensor4 const& pred, int max_m) { + CUTE_STATIC_ASSERT_V(rank(src) == Int<3>{}); + CUTE_STATIC_ASSERT_V(rank(dst) == Int<3>{}); + CUTE_STATIC_ASSERT_V(size<0>(src) == size<0>(dst)); + CUTE_STATIC_ASSERT_V(size<1>(src) == size<1>(dst)); + CUTE_STATIC_ASSERT_V(size<2>(src) == size<2>(dst)); + static_assert(!(Clear_OOB_MN && !Clear_OOB_K)); +#pragma unroll + for (int m = 0; m < size<1>(src); ++m) { + if (Is_even_MN || get<0>(coord(0, m, 0)) < max_m) { +#pragma unroll + for (int k = 0; k < size<2>(src); ++k) { + if (Is_even_K || pred(k)) { + cute::copy(tiled_copy, src(_, m, k), dst(_, m, k)); + } else if (Clear_OOB_K) { + cute::clear(dst(_, m, k)); + } + } + } else if (Clear_OOB_MN) { + cute::clear(dst(_, m, _)); + } + } + } +}; + +} // namespace nvfp4_attention diff --git a/include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/mainloop.cuh b/include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/mainloop.cuh new file mode 100644 index 00000000000..366f7947ff2 --- /dev/null +++ b/include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/mainloop.cuh @@ -0,0 +1,833 @@ +/* + * Copyright (c) 2025 by SageAttention team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include +#include +#include + +#include + +#include "../primitives/barrier.cuh" +#include "../quantization/fp4_convert.cuh" +#include "../utils/layout.cuh" +#include "../utils/math.cuh" +#include "cute/tensor.hpp" +#include "cutlass/gemm/collective/collective_builder.hpp" +#include "cutlass/pipeline/pipeline.hpp" +namespace nvfp4_attention { + +using namespace cute; + +template +struct CollectiveMainloopFwd { + using Element = typename Ktraits::Element; + using ElementSF = typename Ktraits::ElementSF; + using ElementDS = typename Ktraits::ElementDS; + + using TileShape_MNK = typename Ktraits::TileShape_MNK; + using ClusterShape = typename Ktraits::ClusterShape_MNK; + + static constexpr int kStages = Ktraits::kStages; + static constexpr int kHeadDim = Ktraits::kHeadDim; + static constexpr int BlockMean = Ktraits::BlockMean; + using GmemTiledCopy = typename Ktraits::GmemTiledCopy; + using SmemLayoutQ = typename Ktraits::SmemLayoutQ; + using SmemLayoutK = typename Ktraits::SmemLayoutK; + using SmemLayoutV = typename Ktraits::SmemLayoutV; + using SmemLayoutVt = typename Ktraits::SmemLayoutVt; + using SmemLayoutDS = typename Ktraits::SmemLayoutDS; + using SmemLayoutAtomDS = typename Ktraits::SmemLayoutAtomDS; + using LayoutDS = decltype(blocked_product( + SmemLayoutAtomDS{}, make_layout(make_shape(int32_t(0), int32_t(0), int32_t(0), int32_t(0)), + make_stride(int32_t(0), _1{}, int32_t(0), int32_t(0))))); + using ShapeQKV = cute::Shape; + using StrideQKV = cute::Stride; + using ShapeSF = cute::Shape; + using LayoutSF = typename Ktraits::LayoutSF; + using LayoutP = typename Ktraits::LayoutP; + using LayoutSFP = typename Ktraits::LayoutSFP; + using SfAtom = typename Ktraits::SfAtom; + using TMA_Q = + decltype(make_tma_copy(GmemTiledCopy{}, + make_tensor(make_gmem_ptr(static_cast(nullptr)), + repeat_like(StrideQKV{}, int32_t(0)), StrideQKV{}), + SmemLayoutQ{}, select<0, 2>(TileShape_MNK{}), _1{})); + + using TMA_KV = + decltype(make_tma_copy(GmemTiledCopy{}, + make_tensor(make_gmem_ptr(static_cast(nullptr)), + repeat_like(StrideQKV{}, int32_t(0)), StrideQKV{}), + take<0, 2>(SmemLayoutK{}), select<1, 2>(TileShape_MNK{}), _1{})); + + using TMA_Vt = decltype(make_tma_copy( + GmemTiledCopy{}, + make_tensor(make_gmem_ptr(static_cast(nullptr)), + repeat_like(StrideQKV{}, int32_t(0)), StrideQKV{}), + take<0, 2>(SmemLayoutVt{}), make_shape(shape<2>(TileShape_MNK{}), shape<1>(TileShape_MNK{})), + _1{})); + + using TMA_DS = decltype(make_tma_copy( + GmemTiledCopy{}, + make_tensor(make_gmem_ptr(static_cast(nullptr)), LayoutDS{}), + take<0, 2>(SmemLayoutDS{}), make_shape(shape<0>(TileShape_MNK{}), shape<1>(TileShape_MNK{})), + _1{})); + + using BlkScaledConfig = typename Ktraits::BlkScaledConfig; + using GmemTiledCopySF = typename Ktraits::GmemTiledCopySF; + using SmemLayoutSFQ = typename Ktraits::SmemLayoutSFQ; + using SmemLayoutSFK = typename Ktraits::SmemLayoutSFK; + using SmemLayoutSFV = typename Ktraits::SmemLayoutSFV; + using SmemLayoutSFVt = typename Ktraits::SmemLayoutSFVt; + + using TMA_SFQ = decltype(make_tma_copy( + GmemTiledCopySF{}, make_tensor(static_cast(nullptr), LayoutSF{}), + SmemLayoutSFQ{}, make_shape(shape<0>(TileShape_MNK{}), shape<2>(TileShape_MNK{})), _1{})); + + using TMA_SFKV = decltype(make_tma_copy( + GmemTiledCopySF{}, make_tensor(static_cast(nullptr), LayoutSF{}), + SmemLayoutSFK{}(_, _, cute::Int<0>{}), + make_shape(shape<1>(TileShape_MNK{}), shape<2>(TileShape_MNK{})), _1{})); + + using TMA_SFVt = decltype(make_tma_copy( + GmemTiledCopySF{}, make_tensor(static_cast(nullptr), LayoutSF{}), + SmemLayoutSFVt{}(_, _, cute::Int<0>{}), + make_shape(shape<2>(TileShape_MNK{}), shape<1>(TileShape_MNK{})), _1{})); + + using SmemCopyAtomQ = typename Ktraits::SmemCopyAtomQ; + using SmemCopyAtomKV = typename Ktraits::SmemCopyAtomKV; + using SmemCopyAtomSF = typename Ktraits::SmemCopyAtomSF; + using TiledMmaQK = typename Ktraits::TiledMmaQK; + using TiledMmaPV = typename Ktraits::TiledMmaPV; + static constexpr int NumMmaThreads = size(TiledMmaQK{}); + using MainloopPipeline = typename Ktraits::MainloopPipeline; + using PipelineParams = typename MainloopPipeline::Params; + using PipelineState = typename MainloopPipeline::PipelineState; + using MainloopPipelineQ = typename Ktraits::MainloopPipelineQ; + using PipelineParamsQ = typename Ktraits::PipelineParamsQ; + using PipelineStateQ = typename Ktraits::PipelineStateQ; + using EpilogueBarrier = typename Ktraits::EpilogueBarrier; + + static constexpr uint32_t TmaTransactionBytesQ = static_cast( + cutlass::bits_to_bytes(cosize((SmemLayoutSFQ{})) * cute::sizeof_bits_v) + + cutlass::bits_to_bytes(size((SmemLayoutQ{})) * sizeof_bits::value)); + + static constexpr uint32_t TmaTransactionBytesK = static_cast( + cutlass::bits_to_bytes(cosize(take<0, 2>(SmemLayoutSFK{})) * cute::sizeof_bits_v) + + cutlass::bits_to_bytes(cosize(take<0, 2>(SmemLayoutDS{})) * cute::sizeof_bits_v) + + cutlass::bits_to_bytes(size(take<0, 2>(SmemLayoutK{})) * sizeof_bits::value)); + + static constexpr uint32_t TmaTransactionBytesV = static_cast( + cutlass::bits_to_bytes(cosize(take<0, 2>(SmemLayoutSFVt{})) * + cute::sizeof_bits_v) + + cutlass::bits_to_bytes(size(take<0, 2>(SmemLayoutVt{})) * sizeof_bits::value)); + + struct Arguments { + Element const* ptr_Q; + ShapeQKV const shape_Q; + StrideQKV const stride_Q; + Element const* ptr_K; + ShapeQKV const shape_K; + StrideQKV const stride_K; + ShapeQKV const unpadded_shape_K; + Element const* ptr_Vt; + ShapeQKV const shape_Vt; + StrideQKV const stride_Vt; + ElementSF const* ptr_SFQ{nullptr}; + ShapeSF const shape_SFQ{}; + ElementSF const* ptr_SFK{nullptr}; + ShapeSF const shape_SFK{}; + ElementSF const* ptr_SFVt{nullptr}; + ShapeSF const shape_SFVt{}; + ElementDS const* ptr_ds; + ShapeQKV const shape_ds; + StrideQKV const stride_ds; + float const softmax_scale_log2; + }; + + struct Params { + ShapeQKV const shape_Q; + LayoutSF const layout_SFQ; + ShapeQKV const shape_K; + ShapeQKV const unpadded_shape_K; + LayoutSF const layout_SFK; + ShapeQKV const shape_Vt; + LayoutSF const layout_SFVt; + LayoutDS const layout_DS; + TMA_Q tma_load_Q; + TMA_SFQ tma_load_SFQ; + TMA_KV tma_load_K; + TMA_SFKV tma_load_SFK; + TMA_Vt tma_load_Vt; + TMA_SFVt tma_load_SFVt; + TMA_DS tma_load_DS; + float const softmax_scale_log2; + }; + + static Params to_underlying_arguments(Arguments const& args) { + Tensor mQ = make_tensor(make_gmem_ptr(args.ptr_Q), args.shape_Q, args.stride_Q); + TMA_Q tma_load_Q = + make_tma_copy(GmemTiledCopy{}, mQ, SmemLayoutQ{}, select<0, 2>(TileShape_MNK{}), _1{}); + Tensor mK = make_tensor(make_gmem_ptr(args.ptr_K), args.shape_K, args.stride_K); + TMA_KV tma_load_K = make_tma_copy(GmemTiledCopy{}, mK, SmemLayoutK{}(_, _, _0{}), + select<1, 2>(TileShape_MNK{}), _1{}); + Tensor mVt = make_tensor(make_gmem_ptr(args.ptr_Vt), args.shape_Vt, args.stride_Vt); + TMA_Vt tma_load_Vt = + make_tma_copy(GmemTiledCopy{}, mVt, SmemLayoutVt{}(_, _, _0{}), + make_shape(shape<2>(TileShape_MNK{}), shape<1>(TileShape_MNK{})), _1{}); + auto [Seqlen_Q, Seqlen_K, HeadNum, Batch] = args.shape_ds; + LayoutDS layout_ds = tile_to_shape( + SmemLayoutAtomDS{}, make_shape(Seqlen_Q, Seqlen_K, HeadNum, Batch), Step<_2, _1, _3, _4>{}); + Tensor mDS = make_tensor(make_gmem_ptr(args.ptr_ds), layout_ds); + TMA_DS tma_load_ds = + make_tma_copy(GmemTiledCopy{}, mDS, SmemLayoutDS{}(_, _, _0{}), + make_shape(shape<0>(TileShape_MNK{}), shape<1>(TileShape_MNK{})), _1{}); + LayoutSF layout_sfq = BlkScaledConfig::tile_atom_to_shape_SFQKV(args.shape_SFQ); + Tensor mSFQ = make_tensor(make_gmem_ptr(args.ptr_SFQ), layout_sfq); + TMA_SFQ tma_load_sfq = make_tma_copy( + GmemTiledCopySF{}, mSFQ, SmemLayoutSFQ{}, + make_shape(shape<0>(TileShape_MNK{}), shape<2>(TileShape_MNK{})), _1{}); + LayoutSF layout_sfk = BlkScaledConfig::tile_atom_to_shape_SFQKV(args.shape_SFK); + Tensor mSFK = make_tensor(make_gmem_ptr(args.ptr_SFK), layout_sfk); + TMA_SFKV tma_load_sfk = make_tma_copy( + GmemTiledCopySF{}, mSFK, SmemLayoutSFK{}(_, _, _0{}), + make_shape(shape<1>(TileShape_MNK{}), shape<2>(TileShape_MNK{})), _1{}); + LayoutSF layout_sfvt = BlkScaledConfig::tile_atom_to_shape_SFVt(args.shape_SFVt); + Tensor mSFVt = make_tensor(make_gmem_ptr(args.ptr_SFVt), layout_sfvt); + TMA_SFVt tma_load_sfvt = make_tma_copy( + GmemTiledCopySF{}, mSFVt, SmemLayoutSFVt{}(_, _, _0{}), + make_shape(shape<2>(TileShape_MNK{}), shape<1>(TileShape_MNK{})), _1{}); + return {args.shape_Q, layout_sfq, args.shape_K, args.unpadded_shape_K, + layout_sfk, args.shape_Vt, layout_sfvt, layout_ds, + tma_load_Q, tma_load_sfq, tma_load_K, tma_load_sfk, + tma_load_Vt, tma_load_sfvt, tma_load_ds, args.softmax_scale_log2}; + } + + CUTLASS_DEVICE + static void prefetch_tma_descriptors(Params const& mainloop_params) { + cute::prefetch_tma_descriptor(mainloop_params.tma_load_Q.get_tma_descriptor()); + cute::prefetch_tma_descriptor(mainloop_params.tma_load_K.get_tma_descriptor()); + cute::prefetch_tma_descriptor(mainloop_params.tma_load_Vt.get_tma_descriptor()); + cute::prefetch_tma_descriptor(mainloop_params.tma_load_SFQ.get_tma_descriptor()); + cute::prefetch_tma_descriptor(mainloop_params.tma_load_SFK.get_tma_descriptor()); + cute::prefetch_tma_descriptor(mainloop_params.tma_load_SFVt.get_tma_descriptor()); + cute::prefetch_tma_descriptor(mainloop_params.tma_load_DS.get_tma_descriptor()); + } + + CUTLASS_DEVICE + int get_n_block_max(Params const& mainloop_params, int m_block) { + static constexpr int kBlockM = get<0>(TileShape_MNK{}); + static constexpr int kBlockN = get<1>(TileShape_MNK{}); + int const seqlen_q = get<0>(mainloop_params.shape_Q); + int const seqlen_k = get<0>(mainloop_params.shape_K); + int n_block_max = cute::ceil_div(seqlen_k, kBlockN); + if constexpr (Is_causal) { + n_block_max = std::min( + n_block_max, cute::ceil_div((m_block + 1) * kBlockM + seqlen_k - seqlen_q, kBlockN)); + } + return n_block_max; + } + + template + CUTE_HOST_DEVICE constexpr auto thrfrg_SFA(SFATensor&& sfatensor, + TiledMMA& mma) { + CUTE_STATIC_ASSERT_V(rank(sfatensor) >= Int<2>{}); + + using AtomShape_MNK = typename Atom::Shape_MNK; + using AtomLayoutSFA_TV = typename Atom::Traits::SFALayout; + + auto permutation_mnk = TiledPerm{}; + auto thr_layout_vmnk = mma.get_thr_layout_vmnk(); + + auto t_tile = make_tile(get<0>(permutation_mnk), get<2>(permutation_mnk)); + auto t_tensor = logical_divide(sfatensor, t_tile); + + auto a_tile = + make_tile(make_layout(size<0>(AtomShape_MNK{})), make_layout(size<2>(AtomShape_MNK{}))); + auto a_tensor = zipped_divide(t_tensor, a_tile); + + auto tv_tensor = a_tensor.compose(AtomLayoutSFA_TV{}, _); + + auto thr_tile = make_tile( + _, make_tile(make_layout(size<1>(thr_layout_vmnk)), make_layout(size<3>(thr_layout_vmnk)))); + auto thr_tensor = zipped_divide(tv_tensor, thr_tile); + + return thr_tensor; + } + + template + CUTE_HOST_DEVICE constexpr auto thrfrg_SFB(SFBTensor&& sfbtensor, + TiledMMA& mma) { + CUTE_STATIC_ASSERT_V(rank(sfbtensor) >= Int<2>{}); + + using AtomShape_MNK = typename Atom::Shape_MNK; + using AtomLayoutSFB_TV = typename Atom::Traits::SFBLayout; + + auto permutation_mnk = TiledPerm{}; + auto thr_layout_vmnk = mma.get_thr_layout_vmnk(); + + auto t_tile = make_tile(get<1>(permutation_mnk), get<2>(permutation_mnk)); + auto t_tensor = logical_divide(sfbtensor, t_tile); + + auto a_tile = + make_tile(make_layout(size<1>(AtomShape_MNK{})), make_layout(size<2>(AtomShape_MNK{}))); + auto a_tensor = zipped_divide(t_tensor, a_tile); + + auto tv_tensor = a_tensor.compose(AtomLayoutSFB_TV{}, _); + + auto thr_tile = make_tile( + _, make_tile(make_layout(size<2>(thr_layout_vmnk)), make_layout(size<3>(thr_layout_vmnk)))); + auto thr_tensor = zipped_divide(tv_tensor, thr_tile); + return thr_tensor; + } + + template + CUTE_HOST_DEVICE constexpr auto partition_fragment_SFA(SFATensor&& sfatensor, + ThrMma& thread_mma) { + using ValTypeSF = typename ThrMma::Atom::Traits::ValTypeSF; + auto thr_tensor = make_tensor(static_cast(sfatensor).data(), + thrfrg_SFA(sfatensor.layout(), thread_mma)); + auto thr_vmnk = thread_mma.thr_vmnk_; + auto thr_vmk = make_coord(get<0>(thr_vmnk), make_coord(get<1>(thr_vmnk), get<3>(thr_vmnk))); + auto partition_SFA = thr_tensor(thr_vmk, make_coord(_, repeat(thr_tensor)>(_))); + return make_fragment_like(partition_SFA); + } + + template + CUTE_HOST_DEVICE constexpr auto partition_fragment_SFB(SFBTensor&& sfbtensor, + ThrMma& thread_mma) { + using ValTypeSF = typename ThrMma::Atom::Traits::ValTypeSF; + auto thr_tensor = make_tensor(static_cast(sfbtensor).data(), + thrfrg_SFB(sfbtensor.layout(), thread_mma)); + auto thr_vmnk = thread_mma.thr_vmnk_; + auto thr_vnk = make_coord(get<0>(thr_vmnk), make_coord(get<2>(thr_vmnk), get<3>(thr_vmnk))); + auto partition_SFB = thr_tensor(thr_vnk, make_coord(_, repeat(thr_tensor)>(_))); + return make_fragment_like(partition_SFB); + } + + template + CUTE_HOST_DEVICE constexpr auto get_layoutSFA_TV(TiledMma& mma) { + auto tile_shape_mnk = tile_shape(mma); + auto ref_A = make_layout(make_shape(size<0>(tile_shape_mnk), size<2>(tile_shape_mnk))); + auto thr_layout_vmnk = mma.get_thr_layout_vmnk(); + + auto atile = make_tile( + _, make_tile(make_layout(make_shape(size<1>(thr_layout_vmnk), size<2>(thr_layout_vmnk)), + make_stride(Int<1>{}, Int<0>{})), + _)); + + auto thridx_2_thrid = right_inverse(thr_layout_vmnk); + + return thrfrg_SFA(ref_A, mma).compose(atile, _).compose(thridx_2_thrid, _); + } + + template + CUTE_HOST_DEVICE constexpr auto get_layoutSFB_TV(TiledMma& mma) { + auto tile_shape_mnk = tile_shape(mma); + auto ref_B = make_layout(make_shape(size<1>(tile_shape_mnk), size<2>(tile_shape_mnk))); + auto thr_layout_vmnk = mma.get_thr_layout_vmnk(); + + auto btile = make_tile( + _, make_tile(make_layout(make_shape(size<1>(thr_layout_vmnk), size<2>(thr_layout_vmnk)), + make_stride(Int<0>{}, Int<1>{})), + _)); + + auto thridx_2_thrid = right_inverse(thr_layout_vmnk); + + return thrfrg_SFB(ref_B, mma).compose(btile, _).compose(thridx_2_thrid, _); + } + + template + CUTLASS_DEVICE void load(Params const& mainloop_params, SchedulerParams const& scheduler_params, + MainloopPipelineQ pipeline_q, MainloopPipeline pipeline_k, + MainloopPipeline pipeline_v, PipelineStateQ& smem_pipe_write_q, + PipelineState& smem_pipe_write_k, PipelineState& smem_pipe_write_v, + SharedStorage& shared_storage, WorkTileInfo work_tile_info, + int& work_idx, int& tile_count_semaphore) { + static constexpr int kBlockM = get<0>(TileShape_MNK{}); + static constexpr int kBlockN = get<1>(TileShape_MNK{}); + + auto [m_block, bidh, bidb] = work_tile_info.get_block_coord(scheduler_params); + + int n_block_max = get_n_block_max(mainloop_params, m_block); + + Tensor sQ = make_tensor(make_smem_ptr(shared_storage.smem_q.begin()), SmemLayoutQ{}); + Tensor sK = make_tensor(make_smem_ptr(shared_storage.smem_k.begin()), SmemLayoutK{}); + Tensor sVt = make_tensor(make_smem_ptr(shared_storage.smem_v.begin()), SmemLayoutVt{}); + Tensor sSFQ = make_tensor(make_smem_ptr(shared_storage.smem_SFQ.begin()), SmemLayoutSFQ{}); + Tensor sSFK = make_tensor(make_smem_ptr(shared_storage.smem_SFK.begin()), SmemLayoutSFK{}); + Tensor sSFVt = make_tensor(make_smem_ptr(shared_storage.smem_SFV.begin()), SmemLayoutSFVt{}); + Tensor sDS = make_tensor(make_smem_ptr(shared_storage.smem_ds.begin()), SmemLayoutDS{}); + + Tensor mQ = mainloop_params.tma_load_Q.get_tma_tensor(mainloop_params.shape_Q); + Tensor mK = mainloop_params.tma_load_K.get_tma_tensor(mainloop_params.shape_K); + Tensor mVt = mainloop_params.tma_load_Vt.get_tma_tensor(mainloop_params.shape_Vt); + Tensor mDS = mainloop_params.tma_load_DS.get_tma_tensor(shape(mainloop_params.layout_DS)); + Tensor mSFQ = mainloop_params.tma_load_SFQ.get_tma_tensor(shape(mainloop_params.layout_SFQ)); + Tensor mSFK = mainloop_params.tma_load_SFK.get_tma_tensor(shape(mainloop_params.layout_SFK)); + Tensor mSFVt = mainloop_params.tma_load_SFVt.get_tma_tensor(shape(mainloop_params.layout_SFVt)); + uint32_t block_rank_in_cluster = cute::block_rank_in_cluster(); + constexpr uint32_t cluster_shape_x = get<0>(ClusterShape()); + uint2 cluster_local_block_id = {block_rank_in_cluster % cluster_shape_x, + block_rank_in_cluster / cluster_shape_x}; + Tensor gQ = + local_tile(mQ(_, _, bidh, bidb), select<0, 2>(TileShape_MNK{}), make_coord(m_block, _0{})); + Tensor gK = + local_tile(mK(_, _, bidh, bidb), select<1, 2>(TileShape_MNK{}), make_coord(_, _0{})); + Tensor gVt = local_tile(mVt(_, _, bidh, bidb), + make_shape(shape<2>(TileShape_MNK{}), shape<1>(TileShape_MNK{})), + make_coord(_0{}, _)); + Tensor gDS = [&] { + if constexpr (BlockMean) { + return local_tile(mDS(_, _, bidh, bidb), select<0, 1>(TileShape_MNK{}), + make_coord(m_block, _)); + } else { + return local_tile(mDS(_, _, bidh, bidb), select<0, 1>(TileShape_MNK{}), + make_coord(_0{}, _)); + } + }(); + Tensor gSFQ = local_tile(mSFQ(_, _, bidh, bidb), select<0, 2>(TileShape_MNK{}), + make_coord(m_block, _0{})); + Tensor gSFK = + local_tile(mSFK(_, _, bidh, bidb), select<1, 2>(TileShape_MNK{}), make_coord(_, _0{})); + Tensor gSFVt = local_tile(mSFVt(_, _, bidh, bidb), + make_shape(shape<2>(TileShape_MNK{}), shape<1>(TileShape_MNK{})), + make_coord(_0{}, _)); + auto block_tma_q = mainloop_params.tma_load_Q.get_slice(_0{}); + Tensor tQgQ = block_tma_q.partition_S(gQ); + Tensor tQsQ = block_tma_q.partition_D(sQ); + auto block_tma_sfq = mainloop_params.tma_load_SFQ.get_slice(_0{}); + Tensor tQgSFQ = block_tma_sfq.partition_S(gSFQ); + Tensor tQsSFQ = block_tma_sfq.partition_D(sSFQ); + auto block_tma_k = mainloop_params.tma_load_K.get_slice(cluster_local_block_id.x); + Tensor tKgK = group_modes<0, 3>(block_tma_k.partition_S(gK)); + Tensor tKsK = group_modes<0, 3>(block_tma_k.partition_D(sK)); + auto block_tma_sfk = mainloop_params.tma_load_SFK.get_slice(cluster_local_block_id.x); + Tensor tKgSFK = group_modes<0, 3>(block_tma_sfk.partition_S(gSFK)); + Tensor tKsSFK = group_modes<0, 3>(block_tma_sfk.partition_D(sSFK)); + auto block_tma_vt = mainloop_params.tma_load_Vt.get_slice(cluster_local_block_id.x); + Tensor tVgVt = group_modes<0, 3>(block_tma_vt.partition_S(gVt)); + Tensor tVsVt = group_modes<0, 3>(block_tma_vt.partition_D(sVt)); + auto block_tma_sfvt = mainloop_params.tma_load_SFVt.get_slice(cluster_local_block_id.x); + Tensor tVgSFVt = group_modes<0, 3>(block_tma_sfvt.partition_S(gSFVt)); + Tensor tVsSFVt = group_modes<0, 3>(block_tma_sfvt.partition_D(sSFVt)); + auto block_tma_ds = mainloop_params.tma_load_DS.get_slice(cluster_local_block_id.x); + Tensor tDSgDS = group_modes<0, 3>(block_tma_ds.partition_S(gDS)); + Tensor tDSsDS = group_modes<0, 3>(block_tma_ds.partition_D(sDS)); + uint16_t mcast_mask_kv = 0; + + int n_block = n_block_max - 1; + int lane_predicate = cute::elect_one_sync(); + if (lane_predicate) { + pipeline_q.producer_acquire(smem_pipe_write_q); + copy(mainloop_params.tma_load_Q.with(*pipeline_q.producer_get_barrier(smem_pipe_write_q), 0), + tQgQ, tQsQ); + copy( + mainloop_params.tma_load_SFQ.with(*pipeline_q.producer_get_barrier(smem_pipe_write_q), 0), + tQgSFQ, tQsSFQ); + ++smem_pipe_write_q; + pipeline_k.producer_acquire(smem_pipe_write_k); + copy(mainloop_params.tma_load_K.with(*pipeline_k.producer_get_barrier(smem_pipe_write_k), + mcast_mask_kv), + tKgK(_, n_block), tKsK(_, smem_pipe_write_k.index())); + copy(mainloop_params.tma_load_SFK.with(*pipeline_k.producer_get_barrier(smem_pipe_write_k), + mcast_mask_kv), + tKgSFK(_, n_block), tKsSFK(_, smem_pipe_write_k.index())); + copy(mainloop_params.tma_load_DS.with(*pipeline_k.producer_get_barrier(smem_pipe_write_k), + mcast_mask_kv), + tDSgDS(_, n_block), tDSsDS(_, smem_pipe_write_k.index())); + ++smem_pipe_write_k; + pipeline_v.producer_acquire(smem_pipe_write_v); + copy(mainloop_params.tma_load_Vt.with(*pipeline_v.producer_get_barrier(smem_pipe_write_v), + mcast_mask_kv), + tVgVt(_, n_block), tVsVt(_, smem_pipe_write_v.index())); + copy(mainloop_params.tma_load_SFVt.with(*pipeline_v.producer_get_barrier(smem_pipe_write_v), + mcast_mask_kv), + tVgSFVt(_, n_block), tVsSFVt(_, smem_pipe_write_v.index())); + ++smem_pipe_write_v; + } + + n_block--; + if (lane_predicate) { +#pragma unroll 2 + for (; n_block >= 0; --n_block) { + pipeline_k.producer_acquire(smem_pipe_write_k); + copy(mainloop_params.tma_load_K.with(*pipeline_k.producer_get_barrier(smem_pipe_write_k), + mcast_mask_kv), + tKgK(_, n_block), tKsK(_, smem_pipe_write_k.index())); + copy(mainloop_params.tma_load_SFK.with(*pipeline_k.producer_get_barrier(smem_pipe_write_k), + mcast_mask_kv), + tKgSFK(_, n_block), tKsSFK(_, smem_pipe_write_k.index())); + copy(mainloop_params.tma_load_DS.with(*pipeline_k.producer_get_barrier(smem_pipe_write_k), + mcast_mask_kv), + tDSgDS(_, n_block), tDSsDS(_, smem_pipe_write_k.index())); + ++smem_pipe_write_k; + pipeline_v.producer_acquire(smem_pipe_write_v); + copy(mainloop_params.tma_load_Vt.with(*pipeline_v.producer_get_barrier(smem_pipe_write_v), + mcast_mask_kv), + tVgVt(_, n_block), tVsVt(_, smem_pipe_write_v.index())); + copy(mainloop_params.tma_load_SFVt.with(*pipeline_v.producer_get_barrier(smem_pipe_write_v), + mcast_mask_kv), + tVgSFVt(_, n_block), tVsSFVt(_, smem_pipe_write_v.index())); + ++smem_pipe_write_v; + } + } + ++work_idx; + } + + CUTLASS_DEVICE void load_tail(MainloopPipelineQ pipeline_q, MainloopPipeline pipeline_k, + MainloopPipeline pipeline_v, PipelineStateQ& smem_pipe_write_q, + PipelineState& smem_pipe_write_k, + PipelineState& smem_pipe_write_v) { + int lane_predicate = cute::elect_one_sync(); + + if (lane_predicate) { + pipeline_q.producer_tail(smem_pipe_write_q); + pipeline_k.producer_tail(smem_pipe_write_k); + pipeline_v.producer_tail(smem_pipe_write_v); + } + } + + struct NoOpRefill { + CUTLASS_DEVICE void refill_k(int) {} + CUTLASS_DEVICE void refill_v(int) {} + }; + + template + CUTLASS_DEVICE void mma(Params const& mainloop_params, MainloopPipelineQ pipeline_q, + MainloopPipeline pipeline_k, MainloopPipeline pipeline_v, + PipelineStateQ& smem_pipe_read_q, PipelineState& smem_pipe_read_k, + PipelineState& smem_pipe_read_v, FrgTensorO& tOrO_store, + SoftmaxFused& softmax_fused, int n_block_count, int thread_idx, + int work_idx, int m_block, int wg_id, SharedStorage& shared_storage, + MathOrderBarrier& math_order, TmaRefill tma_refill = {}) { + static_assert(is_rmem::value, "O tensor must be rmem resident."); + + static constexpr int kBlockM = get<0>(TileShape_MNK{}); + static constexpr int kBlockN = get<1>(TileShape_MNK{}); + static constexpr int kBlockK = get<2>(TileShape_MNK{}); + static constexpr int kBlockMPerWG = Ktraits::kBlockMPerWG; + + Tensor sQ_full = make_tensor(make_smem_ptr(shared_storage.smem_q.begin()), SmemLayoutQ{}); + Tensor sK = make_tensor(make_smem_ptr(shared_storage.smem_k.begin()), SmemLayoutK{}); + Tensor sVt = make_tensor(make_smem_ptr(shared_storage.smem_v.begin()), SmemLayoutVt{}); + Tensor sDS = make_tensor(make_smem_ptr(shared_storage.smem_ds.begin()), SmemLayoutDS{}); + Tensor sSFQ_full = make_tensor(make_smem_ptr(shared_storage.smem_SFQ.begin()), SmemLayoutSFQ{}); + Tensor sSFK = make_tensor(make_smem_ptr(shared_storage.smem_SFK.begin()), SmemLayoutSFK{}); + Tensor sSFVt = make_tensor(make_smem_ptr(shared_storage.smem_SFV.begin()), SmemLayoutSFVt{}); + + auto sQ = + local_tile(sQ_full, make_shape(Int{}, Int{}), make_coord(wg_id, 0)); + + TiledMmaQK tiled_mma_qk; + TiledMmaPV tiled_mma_pv; + auto thread_mma_qk = tiled_mma_qk.get_thread_slice(thread_idx); + auto thread_mma_pv = tiled_mma_pv.get_thread_slice(thread_idx); + + using TiledMmaQK_Full = typename Ktraits::TiledMmaQK_Full; + TiledMmaQK_Full tiled_mma_qk_full; + int consumer_thread_idx_full = thread_idx + wg_id * NumMmaThreads; + auto thread_mma_qk_full = tiled_mma_qk_full.get_thread_slice(consumer_thread_idx_full); + + Tensor tSrQ = thread_mma_qk.partition_fragment_A(sQ); + Tensor tSrK = thread_mma_qk.partition_fragment_B(sK(_, _, Int<0>{})); + Tensor tOrVt = thread_mma_pv.partition_fragment_B(sVt(_, _, Int<0>{})); + Tensor tOrP = make_tensor_like(LayoutP{}); + + Tensor tSrSFQ = partition_fragment_SFA(sSFQ_full, thread_mma_qk_full); + Tensor tSrSFK = partition_fragment_SFB(sSFK(_, _, Int<0>{}), thread_mma_qk); + Tensor tOrSFVt = partition_fragment_SFB(sSFVt(_, _, Int<0>{}), thread_mma_pv); + Tensor tOrSFP = make_tensor(LayoutSFP{}); + Tensor tOrSFP_flt = filter_zeros(tOrSFP); + + auto smem_tiled_copy_Q = make_tiled_copy_A(SmemCopyAtomQ{}, tiled_mma_qk); + auto smem_thr_copy_Q = smem_tiled_copy_Q.get_thread_slice(thread_idx); + Tensor tSsQ = smem_thr_copy_Q.partition_S(as_position_independent_swizzle_tensor(sQ)); + Tensor tSrQ_copy_view = smem_thr_copy_Q.retile_D(tSrQ); + + auto smem_tiled_copy_K = make_tiled_copy_B(SmemCopyAtomKV{}, tiled_mma_qk); + auto smem_thr_copy_K = smem_tiled_copy_K.get_thread_slice(thread_idx); + Tensor tSsK = smem_thr_copy_K.partition_S(as_position_independent_swizzle_tensor(sK)); + Tensor tSrK_copy_view = smem_thr_copy_K.retile_D(tSrK); + + auto smem_tiled_copy_V = make_tiled_copy_B(SmemCopyAtomKV{}, tiled_mma_pv); + auto smem_thr_copy_V = smem_tiled_copy_V.get_thread_slice(thread_idx); + Tensor tOsVt = smem_thr_copy_V.partition_S(as_position_independent_swizzle_tensor(sVt)); + Tensor tOrVt_copy_view = smem_thr_copy_V.retile_D(tOrVt); + + auto tile_shape_mnk = tile_shape(tiled_mma_qk); + + auto tile_shape_mnk_full = tile_shape(tiled_mma_qk_full); + auto smem_tiled_copy_SFQ = make_tiled_copy_impl( + SmemCopyAtomSF{}, get_layoutSFA_TV(tiled_mma_qk_full), + make_shape(size<0>(tile_shape_mnk_full), size<2>(tile_shape_mnk_full))); + auto smem_thr_copy_SFQ = smem_tiled_copy_SFQ.get_thread_slice(consumer_thread_idx_full); + Tensor tSsSFQ = + smem_thr_copy_SFQ.partition_S(as_position_independent_swizzle_tensor(sSFQ_full)); + Tensor tSrSFQ_copy_view = smem_thr_copy_SFQ.retile_D(tSrSFQ); + + auto smem_tiled_copy_SFK = + make_tiled_copy_impl(SmemCopyAtomSF{}, get_layoutSFB_TV(tiled_mma_qk), + make_shape(size<1>(tile_shape_mnk), size<2>(tile_shape_mnk))); + auto smem_thr_copy_SFK = smem_tiled_copy_SFK.get_thread_slice(thread_idx); + Tensor tSsSFK = smem_thr_copy_SFK.partition_S(as_position_independent_swizzle_tensor(sSFK)); + Tensor tSrSFK_copy_view = smem_thr_copy_SFK.retile_D(tSrSFK); + + auto smem_tiled_copy_SFV = + make_tiled_copy_impl(SmemCopyAtomSF{}, get_layoutSFB_TV(tiled_mma_pv), + make_shape(size<1>(tile_shape_mnk), size<2>(tile_shape_mnk))); + auto smem_thr_copy_SFV = smem_tiled_copy_SFV.get_thread_slice(thread_idx); + Tensor tOsSFVt = smem_thr_copy_SFV.partition_S(as_position_independent_swizzle_tensor(sSFVt)); + Tensor tOrSFVt_copy_view = smem_thr_copy_SFV.retile_D(tOrSFVt); + + auto consumer_wait = [](auto& pipeline, auto& smem_pipe_read) { + auto barrier_token = pipeline.consumer_try_wait(smem_pipe_read); + pipeline.consumer_wait(smem_pipe_read, barrier_token); + }; + + int const seqlen_q = get<0>(mainloop_params.shape_Q); + int const seqlen_k = get<0>(mainloop_params.shape_K); + int const unpadded_seqlen_k = get<0>(mainloop_params.unpadded_shape_K); + int const wg_m_offset = wg_id * kBlockMPerWG; + + auto copy_k_block = [&](auto block_id) { + auto tSsK_stage = tSsK(_, _, _, smem_pipe_read_k.index()); + auto tSsSFK_stage = tSsSFK(_, _, _, smem_pipe_read_k.index()); + copy(smem_tiled_copy_K, tSsK_stage(_, _, block_id), tSrK_copy_view(_, _, block_id)); + copy(smem_tiled_copy_SFK, tSsSFK_stage(_, _, block_id), tSrSFK_copy_view(_, _, block_id)); + }; + + auto copy_v_block = [&](auto block_id) { + auto tOsVt_stage = tOsVt(_, _, _, smem_pipe_read_v.index()); + auto tOsSFVt_stage = tOsSFVt(_, _, _, smem_pipe_read_v.index()); + copy(smem_tiled_copy_V, tOsVt_stage(_, _, block_id), tOrVt_copy_view(_, _, block_id)); + copy(smem_tiled_copy_SFV, tOsSFVt_stage(_, _, block_id), tOrSFVt_copy_view(_, _, block_id)); + }; + + auto add_delta_s = [&](auto& acc) { + auto acc_float4 = recast(acc); + int quad_id = (thread_idx % 4) * 2; + if constexpr (std::is_same_v) { + auto tSsDS_stage = recast(sDS(_, _, smem_pipe_read_k.index())); + for (int i = 0; i < 4; i++) { + auto num = quad_id + i * 8; + float4 delta_s_0 = tSsDS_stage(make_coord(_0{}, _0{}), make_coord(num, _0{})); + float4 delta_s_1 = tSsDS_stage(make_coord(_0{}, _0{}), make_coord(num + 1, _0{})); + acc_float4(make_coord(make_coord(_0{}, _0{}), _0{}), _0{}, i) = delta_s_0; + acc_float4(make_coord(make_coord(_0{}, _0{}), _1{}), _0{}, i) = delta_s_0; + acc_float4(make_coord(make_coord(_0{}, _1{}), _0{}), _0{}, i) = delta_s_1; + acc_float4(make_coord(make_coord(_0{}, _1{}), _1{}), _0{}, i) = delta_s_1; + } + } else { + using ElementDSVec = cutlass::Array; + auto tSsDS_stage = recast(sDS(_, _, smem_pipe_read_k.index())); + cutlass::NumericConverter convert; + for (int i = 0; i < 4; i++) { + auto num = quad_id + i * 8; + ElementDSVec ds0 = tSsDS_stage(make_coord(_0{}, _0{}), make_coord(num, _0{})); + ElementDSVec ds1 = tSsDS_stage(make_coord(_0{}, _0{}), make_coord(num + 1, _0{})); + float4 delta_s_0; + float4 delta_s_1; + delta_s_0.x = convert(ds0[0]); + delta_s_0.y = convert(ds0[1]); + delta_s_0.z = convert(ds0[2]); + delta_s_0.w = convert(ds0[3]); + delta_s_1.x = convert(ds1[0]); + delta_s_1.y = convert(ds1[1]); + delta_s_1.z = convert(ds1[2]); + delta_s_1.w = convert(ds1[3]); + acc_float4(make_coord(make_coord(_0{}, _0{}), _0{}), _0{}, i) = delta_s_0; + acc_float4(make_coord(make_coord(_0{}, _0{}), _1{}), _0{}, i) = delta_s_0; + acc_float4(make_coord(make_coord(_0{}, _1{}), _0{}), _0{}, i) = delta_s_1; + acc_float4(make_coord(make_coord(_0{}, _1{}), _1{}), _0{}, i) = delta_s_1; + } + } + }; + + Tensor tSrS = + partition_fragment_C(tiled_mma_qk, make_shape(Int{}, Int{})); + Tensor tSrS_converion_view = + make_tensor(tSrS.data(), nvfp4_attention::convert_to_conversion_layout(tSrS.layout())); + Tensor AbsMaxP = make_tensor_like(make_layout( + shape(group<1, 4>(flatten(tSrS_converion_view.layout()(make_coord(_0{}, _), _, _)))))); + + auto col_limit_causal = [&](int row, int n_block) { + return row + wg_m_offset + 1 + seqlen_k - n_block * kBlockN - seqlen_q + m_block * kBlockM; + }; + + auto apply_mask = [&](auto& tSrS_local, int n_block_local) { + int const valid_cols = int(unpadded_seqlen_k - n_block_local * kBlockN); + if constexpr (!Is_causal) { + if (valid_cols >= kBlockN) { + return; + } + } else { + if (valid_cols >= kBlockN && col_limit_causal(0, n_block_local) >= kBlockN) { + return; + } + } + + Tensor cS = cute::make_identity_tensor(make_shape(Int{}, Int{})); + Tensor tScS = thread_mma_qk.partition_C(cS); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(tSrS_local); ++i) { + int const col = nvfp4_attention::qk_acc_col_to_k_col(int(get<1>(tScS(i)))); + if constexpr (!Is_causal) { + if (col >= int(unpadded_seqlen_k - n_block_local * kBlockN)) { + tSrS_local(i) = -INFINITY; + } + } else { + if (col >= std::min(seqlen_k - n_block_local * kBlockN, + col_limit_causal(int(get<0>(tScS(i))), n_block_local))) { + tSrS_local(i) = -INFINITY; + } + } + } + }; + + auto quantize = [&](auto mma_k, auto acc_conversion_view) { + Tensor AbsMaxP_stagek = AbsMaxP(_, make_coord(_, _, mma_k)); + Tensor acc_conversion_stagek = acc_conversion_view(_, _, mma_k); + Tensor SFP = make_tensor_like(AbsMaxP_stagek.layout()); + Tensor SFP_uint32_view = recast(SFP); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(AbsMaxP_stagek); i += 4) { + uint32_t& tmp = SFP_uint32_view(i / 4); + nvfp4_attention::packed_float_to_ue4m3(AbsMaxP_stagek(i), AbsMaxP_stagek(i + 1), + AbsMaxP_stagek(i + 2), AbsMaxP_stagek(i + 3), tmp); + } + int const quad_id = threadIdx.x & 3; + uint32_t MASK = (0xFF00FF) << ((quad_id & 1) * 8); + Tensor tOrSFP_uint32_view = recast(tOrSFP(_, _, mma_k)); + Tensor tOrP_uint32_view = recast(tOrP(_, _, mma_k)); + CUTLASS_PRAGMA_UNROLL + for (int mma_m = 0; mma_m < size<1>(tOrP); ++mma_m) { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < 4; ++i) { + nvfp4_attention::packed_float_to_e2m1(acc_conversion_stagek(make_coord(_0{}, i), mma_m), + acc_conversion_stagek(make_coord(_1{}, i), mma_m), + acc_conversion_stagek(make_coord(_2{}, i), mma_m), + acc_conversion_stagek(make_coord(_3{}, i), mma_m), + acc_conversion_stagek(make_coord(_4{}, i), mma_m), + acc_conversion_stagek(make_coord(_5{}, i), mma_m), + acc_conversion_stagek(make_coord(_6{}, i), mma_m), + acc_conversion_stagek(make_coord(_7{}, i), mma_m), + tOrP_uint32_view(i, mma_m)); + } + uint32_t local_sfp = SFP_uint32_view(_0{}, _0{}, mma_m); + uint32_t peer_sfp = __shfl_xor_sync(int32_t(-1), local_sfp, 2); + if ((quad_id & 1) == 0) { + tOrSFP_uint32_view(_0{}, mma_m) = (local_sfp & MASK) | ((peer_sfp & MASK) << 8); + } else { + tOrSFP_uint32_view(_0{}, mma_m) = (peer_sfp & MASK) | ((local_sfp & MASK) >> 8); + } + } + }; + + consumer_wait(pipeline_q, smem_pipe_read_q); + copy(smem_tiled_copy_Q, tSsQ, tSrQ_copy_view); + copy(smem_tiled_copy_SFQ, tSsSFQ, tSrSFQ_copy_view); + pipeline_q.consumer_release(smem_pipe_read_q); + ++smem_pipe_read_q; + + bool is_first_compute = true; + +#pragma unroll 1 + for (int tile_idx = 0; tile_idx < n_block_count; ++tile_idx) { + int n_block = n_block_count - 1 - tile_idx; + + consumer_wait(pipeline_k, smem_pipe_read_k); + + Tensor tSrS_local = + partition_fragment_C(tiled_mma_qk, make_shape(Int{}, Int{})); + Tensor tSrS_local_cv = make_tensor( + tSrS_local.data(), nvfp4_attention::convert_to_conversion_layout(tSrS_local.layout())); + + CUTLASS_PRAGMA_UNROLL + for (int k_block = 0; k_block < size<2>(tSrK); ++k_block) { + copy_k_block(k_block); + } + add_delta_s(tSrS_local); + pipeline_k.consumer_release(smem_pipe_read_k); + ++smem_pipe_read_k; + + if constexpr (!Is_causal) { + math_order.arrive(); + } + + CUTLASS_PRAGMA_UNROLL + for (int k_block = 0; k_block < size<2>(tSrQ); ++k_block) { + cute::gemm(tiled_mma_qk, make_zip_tensor(tSrQ(_, _, k_block), tSrSFQ(_, _, k_block)), + make_zip_tensor(tSrK(_, _, k_block), tSrSFK(_, _, k_block)), tSrS_local); + } + + apply_mask(tSrS_local, n_block); + if (is_first_compute) { + softmax_fused.template online_softmax_with_quant( + tSrS_local, AbsMaxP, mainloop_params.softmax_scale_log2); + } else { + softmax_fused.template online_softmax_with_quant( + tSrS_local, AbsMaxP, mainloop_params.softmax_scale_log2); + } + + auto quantize_score = [&](auto mma_k) { quantize(mma_k, tSrS_local_cv); }; + + math_order.wait(); + consumer_wait(pipeline_v, smem_pipe_read_v); + copy_v_block(_0{}); + quantize_score(_0{}); + + if (is_first_compute) { + CUTLASS_PRAGMA_UNROLL + for (int v_block = 0; v_block < size<2>(tOrP); ++v_block) { + cute::gemm(tiled_mma_pv, make_zip_tensor(tOrP(_, _, v_block), tOrSFP(_, _, v_block)), + make_zip_tensor(tOrVt(_, _, v_block), tOrSFVt(_, _, v_block)), tOrO_store); + if (v_block < size<2>(tOrP) - 1) { + copy_v_block(v_block + 1); + quantize_score(v_block + 1); + } + } + is_first_compute = false; + } else { + Tensor tOrO = make_fragment_like(tOrO_store); + CUTLASS_PRAGMA_UNROLL + for (int v_block = 0; v_block < size<2>(tOrP); ++v_block) { + cute::gemm(tiled_mma_pv, make_zip_tensor(tOrP(_, _, v_block), tOrSFP(_, _, v_block)), + make_zip_tensor(tOrVt(_, _, v_block), tOrSFVt(_, _, v_block)), tOrO); + if (v_block < size<2>(tOrP) - 1) { + copy_v_block(v_block + 1); + quantize_score(v_block + 1); + } + } + softmax_fused.rescale_o(tOrO_store, tOrO); + } + + math_order.arrive(); + pipeline_v.consumer_release(smem_pipe_read_v); + ++smem_pipe_read_v; + + tma_refill.refill_v(tile_idx); + } + + softmax_fused.finalize(tOrO_store); + return; + } +}; + +} // namespace nvfp4_attention diff --git a/include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/producer/load_k.cuh b/include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/producer/load_k.cuh new file mode 100644 index 00000000000..ea0bb750417 --- /dev/null +++ b/include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/producer/load_k.cuh @@ -0,0 +1,128 @@ +/* + * Copyright (c) 2025 by SageAttention team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "cute/tensor.hpp" +#include "cutlass/cutlass.h" + +namespace nvfp4_attention { + +using cute::_; +using cute::_0; +using cute::copy; +using cute::get; +using cute::group_modes; +using cute::local_tile; +using cute::make_coord; +using cute::make_smem_ptr; +using cute::make_tensor; +using cute::select; +using cute::shape; + +template +struct KLoader { + using Element = typename Traits::Element; + using ElementSF = typename Traits::ElementSF; + using TileShape_MNK = typename Traits::TileShape_MNK; + using SmemLayoutK = typename Traits::SmemLayoutK; + using SmemLayoutSFK = typename Traits::SmemLayoutSFK; + using SmemLayoutDS = typename Traits::SmemLayoutDS; + + static constexpr int kBlockN = get<1>(TileShape_MNK{}); + static constexpr int kHeadDim = get<2>(TileShape_MNK{}); + static constexpr bool BlockMean = Traits::BlockMean; + + template + __device__ __forceinline__ static void load_and_stage( + const MainloopParams& mainloop_params, PipelineK& pipeline_k, + PipelineStateK& smem_pipe_write_k, const TensorGK& tKgK, const TensorSK& tKsK, + const TensorGSFK& tKgSFK, const TensorSSFK& tKsSFK, const TensorGDS& tDSgDS, + const TensorSDS& tDSsDS, int n_block, uint16_t mcast_mask_kv, bool lane_predicate) { + if (lane_predicate) { + pipeline_k.producer_acquire(smem_pipe_write_k); + + copy(mainloop_params.tma_load_K.with(*pipeline_k.producer_get_barrier(smem_pipe_write_k), + mcast_mask_kv), + tKgK(_, n_block), tKsK(_, smem_pipe_write_k.index())); + + copy(mainloop_params.tma_load_SFK.with(*pipeline_k.producer_get_barrier(smem_pipe_write_k), + mcast_mask_kv), + tKgSFK(_, n_block), tKsSFK(_, smem_pipe_write_k.index())); + + copy(mainloop_params.tma_load_DS.with(*pipeline_k.producer_get_barrier(smem_pipe_write_k), + mcast_mask_kv), + tDSgDS(_, n_block), tDSsDS(_, smem_pipe_write_k.index())); + + ++smem_pipe_write_k; + } + } + + template + __device__ __forceinline__ static auto prepare_tma_tensors(const MainloopParams& mainloop_params, + SharedStorage& shared_storage, + int m_block, int bidh, int bidb, + uint2 cluster_local_block_id) { + auto sK = make_tensor(make_smem_ptr(shared_storage.smem_k.begin()), SmemLayoutK{}); + auto sSFK = make_tensor(make_smem_ptr(shared_storage.smem_SFK.begin()), SmemLayoutSFK{}); + auto sDS = make_tensor(make_smem_ptr(shared_storage.smem_ds.begin()), SmemLayoutDS{}); + + auto mK = mainloop_params.tma_load_K.get_tma_tensor(mainloop_params.shape_K); + auto mSFK = mainloop_params.tma_load_SFK.get_tma_tensor(shape(mainloop_params.layout_SFK)); + auto mDS = mainloop_params.tma_load_DS.get_tma_tensor(shape(mainloop_params.layout_DS)); + + auto gK = local_tile(mK(_, _, bidh, bidb), select<1, 2>(TileShape_MNK{}), make_coord(_, _0{})); + auto gSFK = + local_tile(mSFK(_, _, bidh, bidb), select<1, 2>(TileShape_MNK{}), make_coord(_, _0{})); + + auto gDS = [&] { + if constexpr (BlockMean) { + return local_tile(mDS(_, _, bidh, bidb), select<0, 1>(TileShape_MNK{}), + make_coord(m_block, _)); + } else { + return local_tile(mDS(_, _, bidh, bidb), select<0, 1>(TileShape_MNK{}), + make_coord(_0{}, _)); + } + }(); + + auto block_tma_k = mainloop_params.tma_load_K.get_slice(cluster_local_block_id.x); + auto tKgK = group_modes<0, 3>(block_tma_k.partition_S(gK)); + auto tKsK = group_modes<0, 3>(block_tma_k.partition_D(sK)); + + auto block_tma_sfk = mainloop_params.tma_load_SFK.get_slice(cluster_local_block_id.x); + auto tKgSFK = group_modes<0, 3>(block_tma_sfk.partition_S(gSFK)); + auto tKsSFK = group_modes<0, 3>(block_tma_sfk.partition_D(sSFK)); + + auto block_tma_ds = mainloop_params.tma_load_DS.get_slice(cluster_local_block_id.x); + auto tDSgDS = group_modes<0, 3>(block_tma_ds.partition_S(gDS)); + auto tDSsDS = group_modes<0, 3>(block_tma_ds.partition_D(sDS)); + + return cute::make_tuple(tKgK, tKsK, tKgSFK, tKsSFK, tDSgDS, tDSsDS, block_tma_k, block_tma_sfk, + block_tma_ds); + } + + template + __device__ __forceinline__ static void prefetch_tma_descriptors( + const MainloopParams& mainloop_params) { + cute::prefetch_tma_descriptor(mainloop_params.tma_load_K.get_tma_descriptor()); + cute::prefetch_tma_descriptor(mainloop_params.tma_load_SFK.get_tma_descriptor()); + cute::prefetch_tma_descriptor(mainloop_params.tma_load_DS.get_tma_descriptor()); + } +}; + +} // namespace nvfp4_attention diff --git a/include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/producer/load_q.cuh b/include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/producer/load_q.cuh new file mode 100644 index 00000000000..cc78821ea2a --- /dev/null +++ b/include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/producer/load_q.cuh @@ -0,0 +1,101 @@ +/* + * Copyright (c) 2025 by SageAttention team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "cute/tensor.hpp" +#include "cutlass/cutlass.h" + +namespace nvfp4_attention { + +using cute::_; +using cute::_0; +using cute::_1; +using cute::copy; +using cute::get; +using cute::local_tile; +using cute::make_coord; +using cute::make_smem_ptr; +using cute::make_tensor; +using cute::select; +using cute::shape; + +template +struct QLoader { + using Element = typename Traits::Element; + using ElementSF = typename Traits::ElementSF; + using TileShape_MNK = typename Traits::TileShape_MNK; + using SmemLayoutQ = typename Traits::SmemLayoutQ; + using SmemLayoutSFQ = typename Traits::SmemLayoutSFQ; + + static constexpr int kBlockM = get<0>(TileShape_MNK{}); + static constexpr int kHeadDim = get<2>(TileShape_MNK{}); + + template + __device__ __forceinline__ static void load_and_stage( + const MainloopParams& mainloop_params, PipelineQ& pipeline_q, + PipelineStateQ& smem_pipe_write_q, const TensorGQ& tQgQ, const TensorSQ& tQsQ, + const TensorGSFQ& tQgSFQ, const TensorSSFQ& tQsSFQ, bool lane_predicate) { + if (lane_predicate) { + pipeline_q.producer_acquire(smem_pipe_write_q); + + copy(mainloop_params.tma_load_Q.with(*pipeline_q.producer_get_barrier(smem_pipe_write_q), 0), + tQgQ, tQsQ); + + copy( + mainloop_params.tma_load_SFQ.with(*pipeline_q.producer_get_barrier(smem_pipe_write_q), 0), + tQgSFQ, tQsSFQ); + + ++smem_pipe_write_q; + } + } + + template + __device__ __forceinline__ static auto prepare_tma_tensors(const MainloopParams& mainloop_params, + SharedStorage& shared_storage, + int m_block, int bidh, int bidb) { + auto sQ = make_tensor(make_smem_ptr(shared_storage.smem_q.begin()), SmemLayoutQ{}); + auto sSFQ = make_tensor(make_smem_ptr(shared_storage.smem_SFQ.begin()), SmemLayoutSFQ{}); + + auto mQ = mainloop_params.tma_load_Q.get_tma_tensor(mainloop_params.shape_Q); + auto mSFQ = mainloop_params.tma_load_SFQ.get_tma_tensor(shape(mainloop_params.layout_SFQ)); + + auto gQ = + local_tile(mQ(_, _, bidh, bidb), select<0, 2>(TileShape_MNK{}), make_coord(m_block, _0{})); + auto gSFQ = local_tile(mSFQ(_, _, bidh, bidb), select<0, 2>(TileShape_MNK{}), + make_coord(m_block, _0{})); + + auto block_tma_q = mainloop_params.tma_load_Q.get_slice(_0{}); + auto tQgQ = block_tma_q.partition_S(gQ); + auto tQsQ = block_tma_q.partition_D(sQ); + + auto block_tma_sfq = mainloop_params.tma_load_SFQ.get_slice(_0{}); + auto tQgSFQ = block_tma_sfq.partition_S(gSFQ); + auto tQsSFQ = block_tma_sfq.partition_D(sSFQ); + + return cute::make_tuple(tQgQ, tQsQ, tQgSFQ, tQsSFQ, block_tma_q, block_tma_sfq); + } + + template + __device__ __forceinline__ static void prefetch_tma_descriptors( + const MainloopParams& mainloop_params) { + cute::prefetch_tma_descriptor(mainloop_params.tma_load_Q.get_tma_descriptor()); + cute::prefetch_tma_descriptor(mainloop_params.tma_load_SFQ.get_tma_descriptor()); + } +}; + +} // namespace nvfp4_attention diff --git a/include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/producer/load_v.cuh b/include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/producer/load_v.cuh new file mode 100644 index 00000000000..3a8e82fe53f --- /dev/null +++ b/include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/producer/load_v.cuh @@ -0,0 +1,106 @@ +/* + * Copyright (c) 2025 by SageAttention team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "cute/tensor.hpp" +#include "cutlass/cutlass.h" + +namespace nvfp4_attention { + +using cute::_; +using cute::_0; +using cute::copy; +using cute::get; +using cute::group_modes; +using cute::local_tile; +using cute::make_coord; +using cute::make_shape; +using cute::make_smem_ptr; +using cute::make_tensor; +using cute::shape; + +template +struct VLoader { + using Element = typename Traits::Element; + using ElementSF = typename Traits::ElementSF; + using TileShape_MNK = typename Traits::TileShape_MNK; + using SmemLayoutV = typename Traits::SmemLayoutV; + using SmemLayoutSFV = typename Traits::SmemLayoutSFV; + + static constexpr int kBlockN = get<1>(TileShape_MNK{}); + static constexpr int kHeadDim = get<2>(TileShape_MNK{}); + + template + __device__ __forceinline__ static void load_and_stage( + const MainloopParams& mainloop_params, PipelineV& pipeline_v, + PipelineStateV& smem_pipe_write_v, const TensorGVt& tVgVt, const TensorSVt& tVsVt, + const TensorGSFVt& tVgSFVt, const TensorSSFVt& tVsSFVt, int n_block, uint16_t mcast_mask_kv, + bool lane_predicate) { + if (lane_predicate) { + pipeline_v.producer_acquire(smem_pipe_write_v); + + copy(mainloop_params.tma_load_Vt.with(*pipeline_v.producer_get_barrier(smem_pipe_write_v), + mcast_mask_kv), + tVgVt(_, n_block), tVsVt(_, smem_pipe_write_v.index())); + + copy(mainloop_params.tma_load_SFVt.with(*pipeline_v.producer_get_barrier(smem_pipe_write_v), + mcast_mask_kv), + tVgSFVt(_, n_block), tVsSFVt(_, smem_pipe_write_v.index())); + + ++smem_pipe_write_v; + } + } + + template + __device__ __forceinline__ static auto prepare_tma_tensors(const MainloopParams& mainloop_params, + SharedStorage& shared_storage, + int bidh, int bidb, + uint2 cluster_local_block_id) { + auto sVt = make_tensor(make_smem_ptr(shared_storage.smem_v.begin()), SmemLayoutV{}); + auto sSFVt = make_tensor(make_smem_ptr(shared_storage.smem_SFV.begin()), SmemLayoutSFV{}); + + auto mVt = mainloop_params.tma_load_Vt.get_tma_tensor(mainloop_params.shape_Vt); + auto mSFVt = mainloop_params.tma_load_SFVt.get_tma_tensor(shape(mainloop_params.layout_SFVt)); + + auto gVt = local_tile(mVt(_, _, bidh, bidb), + make_shape(shape<2>(TileShape_MNK{}), shape<1>(TileShape_MNK{})), + make_coord(_0{}, _)); + auto gSFVt = local_tile(mSFVt(_, _, bidh, bidb), + make_shape(shape<2>(TileShape_MNK{}), shape<1>(TileShape_MNK{})), + make_coord(_0{}, _)); + + auto block_tma_vt = mainloop_params.tma_load_Vt.get_slice(cluster_local_block_id.x); + auto tVgVt = group_modes<0, 3>(block_tma_vt.partition_S(gVt)); + auto tVsVt = group_modes<0, 3>(block_tma_vt.partition_D(sVt)); + + auto block_tma_sfvt = mainloop_params.tma_load_SFVt.get_slice(cluster_local_block_id.x); + auto tVgSFVt = group_modes<0, 3>(block_tma_sfvt.partition_S(gSFVt)); + auto tVsSFVt = group_modes<0, 3>(block_tma_sfvt.partition_D(sSFVt)); + + return cute::make_tuple(tVgVt, tVsVt, tVgSFVt, tVsSFVt, block_tma_vt, block_tma_sfvt); + } + + template + __device__ __forceinline__ static void prefetch_tma_descriptors( + const MainloopParams& mainloop_params) { + cute::prefetch_tma_descriptor(mainloop_params.tma_load_Vt.get_tma_descriptor()); + cute::prefetch_tma_descriptor(mainloop_params.tma_load_SFVt.get_tma_descriptor()); + } +}; + +} // namespace nvfp4_attention diff --git a/include/flashinfer/attention/sm120/nvfp4_attention_sm120/kernel/attention_kernel.h b/include/flashinfer/attention/sm120/nvfp4_attention_sm120/kernel/attention_kernel.h new file mode 100644 index 00000000000..59f8ea12d36 --- /dev/null +++ b/include/flashinfer/attention/sm120/nvfp4_attention_sm120/kernel/attention_kernel.h @@ -0,0 +1,259 @@ +/* + * Copyright (c) 2025 by SageAttention team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include +#include +#include +#include + +#include "../common/params.h" +#include "../compute/consumer/delta_correction.cuh" +#include "../compute/consumer/pv_gemm.cuh" +#include "../compute/consumer/qk_gemm.cuh" +#include "../compute/consumer/softmax.cuh" +#include "../compute/epilogue.cuh" +#include "../compute/epilogue/lse_writer.cuh" +#include "../compute/epilogue/output_writer.cuh" +#include "../compute/mainloop.cuh" +#include "../compute/producer/load_k.cuh" +#include "../compute/producer/load_q.cuh" +#include "../compute/producer/load_v.cuh" +#include "cute/tensor.hpp" +#include "cutlass/pipeline/pipeline.hpp" +#include "scheduler.h" +#include "traits.h" + +namespace nvfp4_attention { + +using namespace cute; + +#if defined(LAUNCH_BOUNDS_2) +#define NVFP4_ATTENTION_MIN_BLOCKS_PER_SM 2 +#else +#define NVFP4_ATTENTION_MIN_BLOCKS_PER_SM 1 +#endif + +#define CONSUMER_REG_ALLOC 232 + +template +__global__ void __launch_bounds__(Ktraits::kNWarps* cutlass::NumThreadsPerWarp, + NVFP4_ATTENTION_MIN_BLOCKS_PER_SM) + attention_kernel_ws( + CUTE_GRID_CONSTANT Flash_fwd_params const params, + CUTE_GRID_CONSTANT + typename CollectiveMainloopFwd::Params const mainloop_params, + CUTE_GRID_CONSTANT typename CollectiveEpilogueFwd::Params const epilogue_params, + CUTE_GRID_CONSTANT typename TileScheduler::Params const scheduler_params) { + using Element = typename Ktraits::Element; + using ElementAccum = typename Ktraits::ElementAccum; + using SoftType = ElementAccum; + using TileShape_MNK = typename Ktraits::TileShape_MNK; + using ClusterShape = typename Ktraits::ClusterShape_MNK; + + static constexpr int NumMmaThreads = size(typename Ktraits::TiledMmaQK{}); + static constexpr int NumCopyThreads = cutlass::NumThreadsPerWarpGroup; + static constexpr int NumConsumerGroups = 2; + static constexpr int kBlockM = Ktraits::kBlockM; + static constexpr int kBlockMPerWG = Ktraits::kBlockMPerWG; + + using CollectiveMainloop = CollectiveMainloopFwd; + using CollectiveEpilogue = CollectiveEpilogueFwd; + + using MainloopPipeline = typename Ktraits::MainloopPipeline; + using PipelineParams = typename MainloopPipeline::Params; + using PipelineState = typename MainloopPipeline::PipelineState; + using MainloopPipelineQ = typename Ktraits::MainloopPipelineQ; + using PipelineParamsQ = typename Ktraits::PipelineParamsQ; + using PipelineStateQ = typename Ktraits::PipelineStateQ; + using EpilogueBarrier = typename Ktraits::EpilogueBarrier; + + enum class WarpGroupRole { + Producer = 0, + Consumer0 = 1, + Consumer1 = 2, + Consumer2 = 3, + Consumer3 = 4 + }; + enum class ProducerWarpRole { Mainloop = 0, Epilogue = 1, Warp2 = 2, Warp3 = 3 }; + + extern __shared__ char shared_memory[]; + auto& shared_storage = *reinterpret_cast(shared_memory); + + int const lane_predicate = cute::elect_one_sync(); + int const warp_idx = cutlass::canonical_warp_idx_sync(); + int warp_group_idx = cutlass::canonical_warp_group_idx(); + int const warp_group_thread_idx = threadIdx.x % cutlass::NumThreadsPerWarpGroup; + int warp_idx_in_warp_group = warp_idx % cutlass::NumWarpsPerWarpGroup; + auto warp_group_role = WarpGroupRole(warp_group_idx); + auto producer_warp_role = ProducerWarpRole(warp_idx_in_warp_group); + + if (warp_idx == 0 && lane_predicate) { + CollectiveMainloop::prefetch_tma_descriptors(mainloop_params); + CollectiveEpilogue::prefetch_tma_descriptors(epilogue_params); + } + + static constexpr int NumAllConsumerThreads = NumConsumerGroups * NumMmaThreads; + + PipelineParams pipeline_params_v; + pipeline_params_v.transaction_bytes = CollectiveMainloop::TmaTransactionBytesV; + pipeline_params_v.role = warp_group_role == WarpGroupRole::Producer + ? MainloopPipeline::ThreadCategory::Producer + : MainloopPipeline::ThreadCategory::Consumer; + pipeline_params_v.is_leader = warp_group_thread_idx == 0; + pipeline_params_v.num_consumers = NumAllConsumerThreads; + + PipelineParams pipeline_params_k; + pipeline_params_k.transaction_bytes = CollectiveMainloop::TmaTransactionBytesK; + pipeline_params_k.role = warp_group_role == WarpGroupRole::Producer + ? MainloopPipeline::ThreadCategory::Producer + : MainloopPipeline::ThreadCategory::Consumer; + pipeline_params_k.is_leader = warp_group_thread_idx == 0; + pipeline_params_k.num_consumers = NumAllConsumerThreads; + + PipelineParamsQ pipeline_params_q; + pipeline_params_q.transaction_bytes = CollectiveMainloop::TmaTransactionBytesQ; + pipeline_params_q.role = warp_group_role == WarpGroupRole::Producer + ? MainloopPipelineQ::ThreadCategory::Producer + : MainloopPipelineQ::ThreadCategory::Consumer; + pipeline_params_q.is_leader = warp_group_thread_idx == 0; + pipeline_params_q.num_consumers = NumAllConsumerThreads; + + MainloopPipelineQ pipeline_q(shared_storage.pipeline_q, pipeline_params_q, ClusterShape{}); + MainloopPipeline pipeline_k(shared_storage.pipeline_k, pipeline_params_k, ClusterShape{}); + MainloopPipeline pipeline_v(shared_storage.pipeline_v, pipeline_params_v, ClusterShape{}); + + uint32_t epilogue_barrier_group_size_list[2] = {cutlass::NumThreadsPerWarp, + NumAllConsumerThreads}; + typename EpilogueBarrier::Params params_epilogue_barrier; + params_epilogue_barrier.group_id = (warp_group_role == WarpGroupRole::Producer); + params_epilogue_barrier.group_size_list = epilogue_barrier_group_size_list; + EpilogueBarrier barrier_o(shared_storage.barrier_o, params_epilogue_barrier); + + using MathOrderBarrier = typename Ktraits::MathOrderBarrier; + uint32_t math_order_group_sizes[2] = {cutlass::NumThreadsPerWarpGroup, + cutlass::NumThreadsPerWarpGroup}; + typename MathOrderBarrier::Params math_order_params; + + math_order_params.group_id = (warp_group_role == WarpGroupRole::Consumer1) ? 1 : 0; + math_order_params.group_size_list = math_order_group_sizes; + MathOrderBarrier math_order(shared_storage.math_order, math_order_params); + + CollectiveMainloop collective_mainloop; + CollectiveEpilogue collective_epilogue; + + __syncthreads(); + + if (warp_group_role == WarpGroupRole::Producer) { + cutlass::arch::warpgroup_reg_dealloc<24>(); + + TileScheduler scheduler; + + if (producer_warp_role == ProducerWarpRole::Mainloop) { + PipelineStateQ smem_pipe_write_q = cutlass::make_producer_start_state(); + PipelineState smem_pipe_write_k = cutlass::make_producer_start_state(); + PipelineState smem_pipe_write_v = cutlass::make_producer_start_state(); + + int work_idx = 0; + + for (auto work_tile_info = scheduler.get_initial_work(); + work_tile_info.is_valid(scheduler_params); + work_tile_info = scheduler.get_next_work(scheduler_params, work_tile_info)) { + int tile_count_semaphore = 0; + + collective_mainloop.load(mainloop_params, scheduler_params, pipeline_q, pipeline_k, + pipeline_v, smem_pipe_write_q, smem_pipe_write_k, + smem_pipe_write_v, shared_storage, work_tile_info, work_idx, + tile_count_semaphore); + + work_idx++; + } + + collective_mainloop.load_tail(pipeline_q, pipeline_k, pipeline_v, smem_pipe_write_q, + smem_pipe_write_k, smem_pipe_write_v); + } + + else if (producer_warp_role == ProducerWarpRole::Epilogue) { + for (auto work_tile_info = scheduler.get_initial_work(); + work_tile_info.is_valid(scheduler_params); + work_tile_info = scheduler.get_next_work(scheduler_params, work_tile_info)) { + barrier_o.wait(); + + collective_epilogue.tma_store(shared_storage, epilogue_params, work_tile_info, + scheduler_params, threadIdx.x); + + collective_epilogue.store_tail(); + + barrier_o.arrive(); + } + } + } + + else if (warp_group_idx >= 1 && warp_group_idx <= NumConsumerGroups) { + cutlass::arch::warpgroup_reg_alloc(); + + typename Ktraits::TiledMmaPV tiled_mma_pv; + TileScheduler scheduler{}; + + int consumer_thread_idx = threadIdx.x - NumCopyThreads; + int wg_id = warp_group_idx - 1; + int mma_thread_idx = consumer_thread_idx % NumMmaThreads; + + PipelineState smem_pipe_read_k, smem_pipe_read_v; + PipelineStateQ smem_pipe_read_q; + + int work_idx = 0; + + CUTLASS_PRAGMA_NO_UNROLL + for (auto work_tile_info = scheduler.get_initial_work(); + work_tile_info.is_valid(scheduler_params); + work_tile_info = scheduler.get_next_work(scheduler_params, work_tile_info)) { + Tensor tOrO = partition_fragment_C(tiled_mma_pv, + make_shape(Int{}, Int{})); + + nvfp4_attention::SoftmaxFused<2 * (2 * kBlockMPerWG / NumMmaThreads)> softmax_fused; + + auto block_coord = work_tile_info.get_block_coord(scheduler_params); + auto [m_block, bidh, bidb] = block_coord; + + int n_block_max = collective_mainloop.get_n_block_max(mainloop_params, m_block); + + if (Is_causal && n_block_max <= 0) { + collective_epilogue.store_zero(epilogue_params, consumer_thread_idx, block_coord); + continue; + } + + collective_mainloop.mma(mainloop_params, pipeline_q, pipeline_k, pipeline_v, smem_pipe_read_q, + smem_pipe_read_k, smem_pipe_read_v, tOrO, softmax_fused, n_block_max, + mma_thread_idx, work_idx, m_block, wg_id, shared_storage, math_order); + + barrier_o.wait(); + + collective_epilogue.mma_store(shared_storage, tiled_mma_pv, tOrO, mma_thread_idx, wg_id); + + barrier_o.arrive(); + + ++work_idx; + } + } +} + +#undef NVFP4_ATTENTION_MIN_BLOCKS_PER_SM +#undef CONSUMER_REG_ALLOC + +} // namespace nvfp4_attention diff --git a/include/flashinfer/attention/sm120/nvfp4_attention_sm120/kernel/scheduler.h b/include/flashinfer/attention/sm120/nvfp4_attention_sm120/kernel/scheduler.h new file mode 100644 index 00000000000..5e1fa00fd89 --- /dev/null +++ b/include/flashinfer/attention/sm120/nvfp4_attention_sm120/kernel/scheduler.h @@ -0,0 +1,222 @@ +/* + * Copyright (c) 2025 by SageAttention team. + * + * This code is based on code from FlashAttention3, https://github.com/Dao-AILab/flash-attention + * Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri + * Dao. Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file + * except in compliance with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "cute/tensor.hpp" +#include "cutlass/fast_math.h" + +namespace nvfp4_attention { + +class SingleTileScheduler { + public: + struct Arguments { + int const num_blocks_m; + int const num_head; + int const num_batch; + int const* tile_count_semaphore = nullptr; + bool const is_causal = false; + }; + + struct Params {}; + + static Params to_underlying_arguments(Arguments const& args) { return {}; } + + static dim3 get_grid_dim(Arguments const& args, int num_sm) { + return {uint32_t(args.num_blocks_m), uint32_t(args.num_head), uint32_t(args.num_batch)}; + } + + struct WorkTileInfo { + int M_idx = 0; + int H_idx = 0; + int B_idx = 0; + bool is_valid_tile = false; + + CUTLASS_DEVICE + bool is_valid(Params const& params) const { return is_valid_tile; } + + CUTLASS_DEVICE + cute::tuple get_block_coord(Params const& params) const { + return {M_idx, H_idx, B_idx}; + } + + CUTLASS_DEVICE + WorkTileInfo get_next_work(Params const& params) const { return {-1, -1, -1, false}; } + }; + + CUTLASS_DEVICE + WorkTileInfo get_initial_work() const { + return {int(blockIdx.x), int(blockIdx.y), int(blockIdx.z), true}; + } + + CUTLASS_DEVICE + WorkTileInfo get_next_work(Params const& params, WorkTileInfo const& current_work) const { + return {-1, -1, -1, false}; + } +}; + +class StaticPersistentTileScheduler { + public: + struct Arguments { + int const num_blocks_m; + int const num_head; + int const num_batch; + int const* tile_count_semaphore = nullptr; + bool const is_causal = false; + }; + + struct Params { + int total_blocks; + int num_blocks_m; + cutlass::FastDivmod m_block_divmod; + cutlass::FastDivmod head_divmod; + bool is_causal; + }; + + static Params to_underlying_arguments(Arguments const& args) { + return {args.num_blocks_m * args.num_head * args.num_batch, args.num_blocks_m, + cutlass::FastDivmod(args.num_blocks_m), cutlass::FastDivmod(args.num_head), + args.is_causal}; + } + + static dim3 get_grid_dim(Arguments const& args, int num_sm) { return {uint32_t(num_sm)}; } + + struct WorkTileInfo { + int tile_idx; + + CUTLASS_DEVICE + bool is_valid(Params const& params) const { return tile_idx < params.total_blocks; } + + CUTLASS_DEVICE + cute::tuple get_block_coord(Params const& params) const { + int m_block, bidh, bidb; + bidb = params.head_divmod.divmod(bidh, params.m_block_divmod.divmod(m_block, tile_idx)); + if (params.is_causal) { + m_block = params.num_blocks_m - 1 - m_block; + } + return {m_block, bidh, bidb}; + } + }; + + CUTLASS_DEVICE + WorkTileInfo get_initial_work() const { return {int(blockIdx.x)}; } + + CUTLASS_DEVICE + WorkTileInfo get_next_work(Params const& params, WorkTileInfo const& current_work) const { + return {current_work.tile_idx + int(gridDim.x)}; + } +}; + +class DynamicPersistentTileScheduler { + public: + struct Arguments { + int const num_blocks_m; + int const num_head; + int const num_batch; + int const* tile_count_semaphore; + bool const is_causal = false; + }; + + struct Params { + int const total_blocks; + int const num_blocks_m; + cutlass::FastDivmod const m_block_divmod; + cutlass::FastDivmod const head_divmod; + bool const is_causal; + }; + + static Params to_underlying_arguments(Arguments const& args) { + return {args.num_blocks_m * args.num_head * args.num_batch, args.num_blocks_m, + cutlass::FastDivmod(args.num_blocks_m), cutlass::FastDivmod(args.num_head), + args.is_causal}; + } + + static dim3 get_grid_dim(Arguments const& args, int num_sm) { return {uint32_t(num_sm)}; } + + using WorkTileInfo = StaticPersistentTileScheduler::WorkTileInfo; + + CUTLASS_DEVICE + WorkTileInfo get_initial_work() const { return {int(blockIdx.x)}; } + + CUTLASS_DEVICE + WorkTileInfo get_next_work(Params const& params, WorkTileInfo const& current_work) const { + return {current_work.tile_idx + int(gridDim.x)}; + } +}; + +class StaticPersistentTileSchedulerOld { + private: + int current_work_linear_idx_; + cutlass::FastDivmod m_block_divmod, head_divmod; + int const total_blocks; + + public: + struct WorkTileInfo { + int M_idx = 0; + int H_idx = 0; + int B_idx = 0; + bool is_valid_tile = false; + + CUTLASS_HOST_DEVICE + bool is_valid() const { return is_valid_tile; } + + CUTLASS_HOST_DEVICE + static WorkTileInfo invalid_work_tile() { return {-1, -1, -1, false}; } + }; + + public: + CUTLASS_DEVICE explicit StaticPersistentTileSchedulerOld( + cutlass::FastDivmod const& m_block_divmod_, cutlass::FastDivmod const& head_divmod_, + int const total_blocks_) + : m_block_divmod(m_block_divmod_), head_divmod(head_divmod_), total_blocks(total_blocks_) { +#if defined(__CUDA_ARCH__) + current_work_linear_idx_ = blockIdx.x; +#else + CUTLASS_ASSERT(false && "This line should never be reached"); +#endif + } + + CUTLASS_DEVICE + WorkTileInfo get_current_work() const { + return get_current_work_for_linear_idx(current_work_linear_idx_); + } + + CUTLASS_DEVICE + WorkTileInfo get_current_work_for_linear_idx(int linear_idx) const { + if (linear_idx >= total_blocks) { + return WorkTileInfo::invalid_work_tile(); + } + + int M_idx, H_idx, B_idx; + int quotient = m_block_divmod.divmod(M_idx, linear_idx); + B_idx = head_divmod.divmod(H_idx, quotient); + return {M_idx, H_idx, B_idx, true}; + } + + CUTLASS_DEVICE + void advance_to_next_work() { current_work_linear_idx_ += int(gridDim.x); } + + CUTLASS_DEVICE + WorkTileInfo fetch_next_work() { + WorkTileInfo new_work_tile_info; + advance_to_next_work(); + new_work_tile_info = get_current_work(); + return new_work_tile_info; + } +}; + +} // namespace nvfp4_attention diff --git a/include/flashinfer/attention/sm120/nvfp4_attention_sm120/kernel/traits.h b/include/flashinfer/attention/sm120/nvfp4_attention_sm120/kernel/traits.h new file mode 100644 index 00000000000..93d7980deae --- /dev/null +++ b/include/flashinfer/attention/sm120/nvfp4_attention_sm120/kernel/traits.h @@ -0,0 +1,241 @@ +/* + * Copyright (c) 2025 by SageAttention team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "../common/cute_extension.h" +#include "../primitives/barrier.cuh" +#include "../quantization/fp4_layout.h" +#include "cute/algorithm/copy.hpp" +#include "cute/atom/mma_atom.hpp" +#include "cute/tensor.hpp" +#include "cutlass/cutlass.h" +#include "cutlass/gemm/collective/collective_builder.hpp" +#include "cutlass/layout/layout.h" +#include "cutlass/numeric_types.h" +#include "cutlass/pipeline/pipeline.hpp" + +using namespace cute; + +namespace nvfp4_attention { + +template +struct SharedStorageQKVOwithSF : cute::aligned_struct<128, _0> { + alignas(1024) cute::ArrayEngine> smem_q; + + alignas(1024) cute::ArrayEngine> smem_k; + + cute::ArrayEngine> smem_SFQ; + cute::ArrayEngine> smem_SFK; + cute::ArrayEngine> smem_SFV; + + alignas(1024) cute::ArrayEngine> smem_ds; + + alignas(1024) cute::ArrayEngine> smem_v; + + alignas(1024) cute::ArrayEngine> smem_o; + + struct { + alignas(16) typename cutlass::PipelineTmaAsync<1>::SharedStorage pipeline_q; + alignas(16) typename cutlass::PipelineTmaAsync::SharedStorage pipeline_k; + alignas(16) typename cutlass::PipelineTmaAsync::SharedStorage pipeline_v; + alignas(16) + typename nvfp4_attention::OrderedSequenceBarrierVarGroupSize::SharedStorage + barrier_o; + + alignas(16) typename nvfp4_attention::OrderedSequenceBarrier<2, 2>::SharedStorage math_order; + int tile_count_semaphore; + }; +}; + +template , + typename ElementOut_ = cutlass::bfloat16_t, typename ElementDS_ = float> +struct Flash_fwd_kernel_traits { + static constexpr int kBlockM = kBlockM_; + static constexpr int kBlockN = kBlockN_; + static constexpr int kHeadDim = kHeadDim_; + static constexpr bool BlockMean = BlockMean_; + static constexpr int kStoreBlockM = kBlockM; + static constexpr bool SmoothQ = true; + + static_assert(kHeadDim % 32 == 0, "Head dim must be multiple of 32"); + static_assert(kBlockM == 64 || kBlockM == 128, "BlockM must be 64 or 128"); + + static constexpr int kNWarps = kBlockM == 128 ? 12 : 8; + static constexpr int kNThreads = kNWarps * cutlass::NumThreadsPerWarp; + + static constexpr int kBlockMPerWG = kBlockM / 2; + + static constexpr int kClusterM = kClusterM_; + static constexpr int kStages = kStages_; + static constexpr int EpiStages = 1; + + static constexpr int NumSFQK = kHeadDim / 16; + static constexpr int NumSFPV = kBlockN / 16; + static constexpr auto SFVectorSize = 16; + + using ElementSF = cutlass::float_ue4m3_t; + using Element = cutlass::float_e2m1_t; + using ElementAccum = float; + using ElementOut = ElementOut_; + using ElementDS = ElementDS_; + using index_t = int64_t; + + using TileShape_MNK = Shape, Int, Int>; + using ClusterShape_MNK = Shape<_1, _1, _1>; + + using PermTileM = Int; + using PermTileN = _32; + using PermTileK = Int; + + using ElementQMma = + decltype(cutlass::gemm::collective::detail::sm1xx_kernel_input_element_to_mma_input_element< + Element>()); + using ElementKMma = + decltype(cutlass::gemm::collective::detail::sm1xx_kernel_input_element_to_mma_input_element< + Element>()); + + using AtomLayoutMNK = Layout>; + + using TiledMmaQK = + decltype(cute::make_tiled_mma(cute::SM120::BLOCKSCALED::SM120_16x32x64_TN_VS_NVFP4{}, + AtomLayoutMNK{}, Tile{})); + + using TiledMmaPV = + decltype(cute::make_tiled_mma(cute::SM120::BLOCKSCALED::SM120_16x32x64_TN_VS_NVFP4{}, + AtomLayoutMNK{}, Tile{})); + + using AtomLayoutMNK_Full = Layout>; + using TiledMmaQK_Full = decltype(cute::make_tiled_mma( + cute::SM120::BLOCKSCALED::SM120_16x32x64_TN_VS_NVFP4{}, AtomLayoutMNK_Full{}, + Tile, PermTileN, PermTileK>{})); + using TiledMmaPV_Full = + decltype(cute::make_tiled_mma(cute::SM120::BLOCKSCALED::SM120_16x32x64_TN_VS_NVFP4{}, + AtomLayoutMNK_Full{}, Tile, _32, PermTileK>{})); + using AtomLayoutMNK_Store = Layout, _1, _1>>; + using TiledMmaPV_Store = decltype(cute::make_tiled_mma( + cute::SM120::BLOCKSCALED::SM120_16x32x64_TN_VS_NVFP4{}, AtomLayoutMNK_Store{}, + Tile, _32, PermTileK>{})); + + static constexpr int MMA_NSF = size<2>(typename TiledMmaQK::AtomShape_MNK{}) / SFVectorSize; + + using GmemTiledCopy = SM90_TMA_LOAD; + using GmemTiledCopySF = SM90_TMA_LOAD; + + using SmemLayoutAtomQ = decltype(cutlass::gemm::collective::detail::sm120_rr_smem_selector< + Element, decltype(size<2>(TileShape_MNK{}))>()); + using SmemLayoutAtomK = decltype(cutlass::gemm::collective::detail::sm120_rr_smem_selector< + Element, decltype(size<2>(TileShape_MNK{}))>()); + using SmemLayoutAtomV = decltype(cutlass::gemm::collective::detail::sm120_rr_smem_selector< + Element, decltype(size<2>(TileShape_MNK{}))>()); + using SmemLayoutAtomVt = decltype(cutlass::gemm::collective::detail::sm120_rr_smem_selector< + Element, decltype(size<1>(TileShape_MNK{}))>()); + + using SmemLayoutQ = decltype(tile_to_shape(SmemLayoutAtomQ{}, select<0, 2>(TileShape_MNK{}))); + + using SmemLayoutK = decltype(tile_to_shape( + SmemLayoutAtomK{}, + make_shape(shape<1>(TileShape_MNK{}), shape<2>(TileShape_MNK{}), Int{}))); + + using SmemLayoutV = decltype(tile_to_shape( + SmemLayoutAtomV{}, + make_shape(shape<1>(TileShape_MNK{}), shape<2>(TileShape_MNK{}), Int{}))); + + using SmemLayoutVt = decltype(tile_to_shape( + SmemLayoutAtomVt{}, + make_shape(shape<2>(TileShape_MNK{}), shape<1>(TileShape_MNK{}), Int{}))); + + using SmemLayoutAtomDS = Layout, Int>, Stride<_0, _1>>; + using SmemLayoutDS = decltype(tile_to_shape( + SmemLayoutAtomDS{}, + make_shape(shape<0>(TileShape_MNK{}), shape<1>(TileShape_MNK{}), Int{}))); + + using SmemCopyAtomQ = Copy_Atom; + using SmemCopyAtomKV = Copy_Atom; + using SmemCopyAtomSF = Copy_Atom, ElementSF>; + using SmemCopyAtomDS = Copy_Atom, ElementDS>; + + using BlkScaledConfig = nvfp4_attention::BlockScaledConfig; + using LayoutSF = typename BlkScaledConfig::LayoutSF; + using SfAtom = typename BlkScaledConfig::SfAtom; + + using SmemLayoutAtomSFQ = + decltype(BlkScaledConfig::deduce_smem_layoutSFQ(TiledMmaQK{}, TileShape_MNK{})); + + using SmemLayoutAtomSFK = + decltype(BlkScaledConfig::deduce_smem_layoutSFKV(TiledMmaQK{}, TileShape_MNK{})); + + using SmemLayoutAtomSFV = + decltype(BlkScaledConfig::deduce_smem_layoutSFKV(TiledMmaPV{}, TileShape_MNK{})); + + using SmemLayoutAtomSFVt = decltype(BlkScaledConfig::deduce_smem_layoutSFVt( + TiledMmaPV{}, Shape, Int, Int>{})); + + using LayoutSFP = + decltype(make_layout(make_shape(make_shape(_16{}, _4{}), _1{}, Int{}), + make_stride(make_stride(_0{}, _1{}), _0{}, _4{}))); + + using LayoutP = + decltype(make_layout(make_shape(make_shape(_8{}, _2{}, _2{}), _1{}, Int{}), + make_stride(make_stride(_1{}, _8{}, _16{}), _0{}, _32{}))); + + using SmemLayoutSFQ = + decltype(make_layout(shape(SmemLayoutAtomSFQ{}), stride(SmemLayoutAtomSFQ{}))); + + using SmemLayoutSFK = decltype(make_layout( + append(shape(SmemLayoutAtomSFK{}), Int{}), + append(stride(SmemLayoutAtomSFK{}), size(filter_zeros(SmemLayoutAtomSFK{}))))); + + using SmemLayoutSFV = decltype(make_layout( + append(shape(SmemLayoutAtomSFV{}), Int{}), + append(stride(SmemLayoutAtomSFV{}), size(filter_zeros(SmemLayoutAtomSFV{}))))); + + using SmemLayoutSFVt = decltype(make_layout( + append(shape(SmemLayoutAtomSFVt{}), Int{}), + append(stride(SmemLayoutAtomSFVt{}), size(filter_zeros(SmemLayoutAtomSFVt{}))))); + + using SmemLayoutAtomO = decltype(cutlass::gemm::collective::detail::ss_smem_selector< + GMMA::Major::K, ElementOut, Int, + decltype(cute::get<2>(TileShape_MNK{}))>()); + using SmemLayoutO = decltype(tile_to_shape( + SmemLayoutAtomO{}, make_shape(Int{}, Int{}), Step<_1, _2>{})); + + using SmemLayoutO_Half = decltype(tile_to_shape( + SmemLayoutAtomO{}, make_shape(Int{}, Int{}), Step<_1, _2>{})); + + using SharedStorage = + SharedStorageQKVOwithSF; + + using MainloopPipeline = typename cutlass::PipelineTmaAsync; + using PipelineState = typename cutlass::PipelineState; + + using MainloopPipelineQ = cutlass::PipelineTmaAsync<1>; + using PipelineParamsQ = typename MainloopPipelineQ::Params; + using PipelineStateQ = typename cutlass::PipelineState<1>; + + using EpilogueBarrier = + typename nvfp4_attention::OrderedSequenceBarrierVarGroupSize; + + using MathOrderBarrier = nvfp4_attention::OrderedSequenceBarrier<2, 2>; +}; + +} // namespace nvfp4_attention diff --git a/include/flashinfer/attention/sm120/nvfp4_attention_sm120/primitives/barrier.cuh b/include/flashinfer/attention/sm120/nvfp4_attention_sm120/primitives/barrier.cuh new file mode 100644 index 00000000000..58942bd8593 --- /dev/null +++ b/include/flashinfer/attention/sm120/nvfp4_attention_sm120/primitives/barrier.cuh @@ -0,0 +1,249 @@ +/* + * Copyright (c) 2025 by SageAttention team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "cutlass/arch/barrier.h" +#include "cutlass/pipeline/sm90_pipeline.hpp" + +namespace nvfp4_attention { + +enum class NamedBarriers { + QueryEmpty = 1, + WarpSpecializedConsumer = 2, + WarpSpecializedPingPongConsumer1 = 3, + WarpSpecializedPingPongConsumer2 = 4, + ProducerEnd = 5, + ConsumerEnd = 6, + EpilogueBarrier = 7 +}; + +template +class OrderedSequenceBarrier { + public: + static constexpr int SequenceDepth = SequenceDepth_; + static constexpr int SequenceLength = SequenceLength_; + + using Barrier = cutlass::arch::ClusterBarrier; + using PipelineState = cutlass::PipelineState; + + struct SharedStorage { + Barrier barrier_[SequenceDepth][SequenceLength]; + }; + + struct Params { + uint32_t group_id; + uint32_t* group_size_list; + }; + + private: + Params params_; + Barrier* barrier_ptr_; + PipelineState stage_; + + public: + OrderedSequenceBarrier() = delete; + OrderedSequenceBarrier(const OrderedSequenceBarrier&) = delete; + OrderedSequenceBarrier(OrderedSequenceBarrier&&) = delete; + OrderedSequenceBarrier& operator=(const OrderedSequenceBarrier&) = delete; + OrderedSequenceBarrier& operator=(OrderedSequenceBarrier&&) = delete; + ~OrderedSequenceBarrier() = default; + + CUTLASS_DEVICE + OrderedSequenceBarrier(SharedStorage& storage, Params const& params) + : params_(params), + barrier_ptr_(&storage.barrier_[0][0]), + + stage_({0, params.group_id == 0, 0}) { + int warp_idx = cutlass::canonical_warp_idx_sync(); + int lane_predicate = cute::elect_one_sync(); + + if (warp_idx == 0 && lane_predicate) { + for (int d = 0; d < SequenceDepth; ++d) { + for (int l = 0; l < SequenceLength; ++l) { + barrier_ptr_[d * SequenceLength + l].init(*(params.group_size_list + l)); + } + } + } + + cutlass::arch::fence_barrier_init(); + } + + CUTLASS_DEVICE + void wait() { get_barrier_for_current_stage(params_.group_id).wait(stage_.phase()); } + + CUTLASS_DEVICE + void arrive() { + int signalling_id = (params_.group_id + 1) % SequenceLength; + get_barrier_for_current_stage(signalling_id).arrive(); + ++stage_; + } + + CUTLASS_DEVICE + void advance() { ++stage_; } + + CUTLASS_DEVICE + Barrier& get_barrier_for_current_stage(int group_id) { + return barrier_ptr_[stage_.index() * SequenceLength + group_id]; + } +}; + +template +class OrderedSequenceBarrierStart { + public: + static constexpr int SequenceDepth = SequenceDepth_; + static constexpr int SequenceLength = SequenceLength_; + static constexpr int StartGroup = StartGroup_; + + using Barrier = cutlass::arch::ClusterBarrier; + using PipelineState = cutlass::PipelineState; + + struct SharedStorage { + Barrier barrier_[SequenceDepth][SequenceLength]; + }; + + struct Params { + uint32_t group_id; + uint32_t* group_size_list; + }; + + private: + Params params_; + Barrier* barrier_ptr_; + PipelineState stage_; + + public: + OrderedSequenceBarrierStart() = delete; + OrderedSequenceBarrierStart(const OrderedSequenceBarrierStart&) = delete; + OrderedSequenceBarrierStart(OrderedSequenceBarrierStart&&) = delete; + OrderedSequenceBarrierStart& operator=(const OrderedSequenceBarrierStart&) = delete; + OrderedSequenceBarrierStart& operator=(OrderedSequenceBarrierStart&&) = delete; + ~OrderedSequenceBarrierStart() = default; + + CUTLASS_DEVICE + OrderedSequenceBarrierStart(SharedStorage& storage, Params const& params) + : params_(params), + barrier_ptr_(&storage.barrier_[0][0]), + stage_({0, params.group_id != StartGroup, 0}) { + int warp_idx = cutlass::canonical_warp_idx_sync(); + int lane_predicate = cute::elect_one_sync(); + + if (warp_idx == 0 && lane_predicate) { + for (int d = 0; d < SequenceDepth; ++d) { + for (int l = 0; l < SequenceLength; ++l) { + barrier_ptr_[d * SequenceLength + l].init(*(params.group_size_list + l)); + } + } + } + + cutlass::arch::fence_barrier_init(); + } + + CUTLASS_DEVICE + void wait() { get_barrier_for_current_stage(params_.group_id).wait(stage_.phase()); } + + CUTLASS_DEVICE + void arrive() { + int signalling_id = (params_.group_id + 1) % SequenceLength; + get_barrier_for_current_stage(signalling_id).arrive(); + ++stage_; + } + + CUTLASS_DEVICE + void advance() { ++stage_; } + + CUTLASS_DEVICE + Barrier& get_barrier_for_current_stage(int group_id) { + return barrier_ptr_[stage_.index() * SequenceLength + group_id]; + } +}; + +template +struct OrderedSequenceBarrierVarGroupSizeSharedStorage { + using Barrier = cutlass::arch::ClusterBarrier; + Barrier barrier_[SequenceDepth][SequenceLength]; +}; + +template +class OrderedSequenceBarrierVarGroupSize { + public: + static constexpr int SequenceDepth = SequenceDepth_; + static constexpr int SequenceLength = SequenceLength_; + using Barrier = cutlass::arch::ClusterBarrier; + using SharedStorage = + OrderedSequenceBarrierVarGroupSizeSharedStorage; + + struct Params { + uint32_t group_id; + uint32_t* group_size_list; + }; + + private: + Params params_; + Barrier* barrier_ptr_; + cutlass::PipelineState stage_; + + static constexpr int Depth = SequenceDepth; + static constexpr int Length = SequenceLength; + + public: + OrderedSequenceBarrierVarGroupSize() = delete; + OrderedSequenceBarrierVarGroupSize(const OrderedSequenceBarrierVarGroupSize&) = delete; + OrderedSequenceBarrierVarGroupSize(OrderedSequenceBarrierVarGroupSize&&) = delete; + OrderedSequenceBarrierVarGroupSize& operator=(const OrderedSequenceBarrierVarGroupSize&) = delete; + OrderedSequenceBarrierVarGroupSize& operator=(OrderedSequenceBarrierVarGroupSize&&) = delete; + ~OrderedSequenceBarrierVarGroupSize() = default; + + CUTLASS_DEVICE + OrderedSequenceBarrierVarGroupSize(SharedStorage& storage, Params const& params) + : params_(params), + barrier_ptr_(&storage.barrier_[0][0]), + + stage_({0, params.group_id == 0, 0}) { + int warp_idx = cutlass::canonical_warp_idx_sync(); + int lane_predicate = cute::elect_one_sync(); + + if (warp_idx == 0 && lane_predicate) { + for (int d = 0; d < Depth; ++d) { + for (int l = 0; l < Length; ++l) { + barrier_ptr_[d * Length + l].init(*(params.group_size_list + l)); + } + } + } + cutlass::arch::fence_barrier_init(); + } + + CUTLASS_DEVICE + void wait() { get_barrier_for_current_stage(params_.group_id).wait(stage_.phase()); } + + CUTLASS_DEVICE + void arrive() { + int signalling_id = (params_.group_id + 1) % Length; + get_barrier_for_current_stage(signalling_id).arrive(); + ++stage_; + } + + CUTLASS_DEVICE + void advance() { ++stage_; } + + private: + CUTLASS_DEVICE + Barrier& get_barrier_for_current_stage(int group_id) { + return barrier_ptr_[stage_.index() * Length + group_id]; + } +}; + +} // namespace nvfp4_attention diff --git a/include/flashinfer/attention/sm120/nvfp4_attention_sm120/primitives/pipeline.cuh b/include/flashinfer/attention/sm120/nvfp4_attention_sm120/primitives/pipeline.cuh new file mode 100644 index 00000000000..e81ac76d5ca --- /dev/null +++ b/include/flashinfer/attention/sm120/nvfp4_attention_sm120/primitives/pipeline.cuh @@ -0,0 +1,132 @@ +/* + * Copyright (c) 2025 by SageAttention team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "cute/tensor.hpp" +#include "cutlass/pipeline/sm90_pipeline.hpp" + +namespace nvfp4_attention { + +using namespace cute; + +template +using PipelineState = cutlass::PipelineState; + +template +class ProducerConsumerPipeline { + public: + using Pipeline = cutlass::PipelineTmaAsync; + using PipelineState = cutlass::PipelineState; + + private: + Pipeline pipeline_; + + public: + template + __device__ __forceinline__ ProducerConsumerPipeline(SharedStorage& shared_storage) + : pipeline_(shared_storage.pipeline) {} + + __device__ __forceinline__ void producer_acquire(PipelineState& state) { + pipeline_.producer_acquire(state); + } + + __device__ __forceinline__ auto* producer_get_barrier(PipelineState& state) { + return pipeline_.producer_get_barrier(state); + } + + __device__ __forceinline__ void producer_commit(PipelineState& state) { + pipeline_.producer_commit(state); + } + + __device__ __forceinline__ void producer_tail(PipelineState& state) { + pipeline_.producer_tail(state); + } + + __device__ __forceinline__ auto consumer_try_wait(PipelineState& state) { + return pipeline_.consumer_try_wait(state); + } + + template + __device__ __forceinline__ void consumer_wait(PipelineState& state, + BarrierToken const& barrier_token) { + pipeline_.consumer_wait(state, barrier_token); + } + + __device__ __forceinline__ void consumer_wait(PipelineState& state) { + auto token = consumer_try_wait(state); + consumer_wait(state, token); + } + + __device__ __forceinline__ void consumer_release(PipelineState& state) { + pipeline_.consumer_release(state); + } +}; + +template +class MultiPipelineManager { + public: + using Pipeline = ProducerConsumerPipeline; + using PipelineState = cutlass::PipelineState; + + private: + Pipeline* pipelines_[NumPipelines]; + + public: + __device__ __forceinline__ MultiPipelineManager(Pipeline* pipelines[NumPipelines]) { + for (int i = 0; i < NumPipelines; ++i) { + pipelines_[i] = pipelines[i]; + } + } + + __device__ __forceinline__ Pipeline& get_pipeline(int idx) { return *pipelines_[idx]; } + + __device__ __forceinline__ void producer_tail_all(PipelineState states[NumPipelines]) { + for (int i = 0; i < NumPipelines; ++i) { + pipelines_[i]->producer_tail(states[i]); + } + } +}; + +template +__device__ __forceinline__ auto make_pipeline_state(int index = 0, bool phase = false, + int count = 0) { + return PipelineState{index, phase, count}; +} + +template +__device__ __forceinline__ void pipeline_producer_load(Pipeline& pipeline, State& state, + LoadFunc const& load_func) { + pipeline.producer_acquire(state); + + load_func(pipeline.producer_get_barrier(state)); + + pipeline.producer_commit(state); + ++state; +} + +template +__device__ __forceinline__ void pipeline_consumer_compute(Pipeline& pipeline, State& state, + ComputeFunc const& compute_func) { + pipeline.consumer_wait(state); + + compute_func(); + + pipeline.consumer_release(state); + ++state; +} + +} // namespace nvfp4_attention diff --git a/include/flashinfer/attention/sm120/nvfp4_attention_sm120/primitives/tma.cuh b/include/flashinfer/attention/sm120/nvfp4_attention_sm120/primitives/tma.cuh new file mode 100644 index 00000000000..5854a69270a --- /dev/null +++ b/include/flashinfer/attention/sm120/nvfp4_attention_sm120/primitives/tma.cuh @@ -0,0 +1,108 @@ +/* + * Copyright (c) 2025 by SageAttention team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "cute/tensor.hpp" +#include "cutlass/cutlass.h" + +namespace nvfp4_attention { + +using namespace cute; + +template +__device__ __forceinline__ void tma_load(TMADesc const& tma_desc, void* barrier, + uint16_t mcast_mask, TensorSrc const& src, + TensorDst& dst) { + cute::copy(tma_desc.with(barrier, mcast_mask), src, dst); +} + +template +__device__ __forceinline__ void tma_store(TMADesc const& tma_desc, TensorSrc const& src, + TensorDst& dst) { + cute::copy(tma_desc, src, dst); +} + +__device__ __forceinline__ void tma_store_arrive() { + asm volatile("cp.async.bulk.commit_group;\n" ::: "memory"); +} + +template +__device__ __forceinline__ void tma_store_wait() { + asm volatile("cp.async.bulk.wait_group %0;\n" ::"n"(N) : "memory"); +} + +template +__device__ __forceinline__ void prefetch_tma_descriptor(TMADesc const& tma_desc) { + cute::prefetch_tma_descriptor(tma_desc.get_tma_descriptor()); +} + +template +struct TMALoader { + TMADesc tma_desc_; + + __device__ __forceinline__ TMALoader(TMADesc const& tma_desc) : tma_desc_(tma_desc) {} + + __device__ __forceinline__ void prefetch() const { + nvfp4_attention::prefetch_tma_descriptor(tma_desc_); + } + + template + __device__ __forceinline__ void load(void* barrier, uint16_t mcast_mask, TensorSrc const& src, + TensorDst& dst) const { + nvfp4_attention::tma_load(tma_desc_, barrier, mcast_mask, src, dst); + } + + template + __device__ __forceinline__ void load(void* barrier, TensorSrc const& src, TensorDst& dst) const { + nvfp4_attention::tma_load(tma_desc_, barrier, 0, src, dst); + } +}; + +template +struct TMAStorer { + TMADesc tma_desc_; + + __device__ __forceinline__ TMAStorer(TMADesc const& tma_desc) : tma_desc_(tma_desc) {} + + __device__ __forceinline__ void prefetch() const { + nvfp4_attention::prefetch_tma_descriptor(tma_desc_); + } + + template + __device__ __forceinline__ void store(TensorSrc const& src, TensorDst& dst) const { + nvfp4_attention::tma_store(tma_desc_, src, dst); + } + + __device__ __forceinline__ void arrive() const { nvfp4_attention::tma_store_arrive(); } + + template + __device__ __forceinline__ void wait() const { + nvfp4_attention::tma_store_wait(); + } +}; + +template +__device__ __forceinline__ auto make_tma_loader(TMADesc const& tma_desc) { + return TMALoader(tma_desc); +} + +template +__device__ __forceinline__ auto make_tma_storer(TMADesc const& tma_desc) { + return TMAStorer(tma_desc); +} + +} // namespace nvfp4_attention diff --git a/include/flashinfer/attention/sm120/nvfp4_attention_sm120/primitives/warpgroup.cuh b/include/flashinfer/attention/sm120/nvfp4_attention_sm120/primitives/warpgroup.cuh new file mode 100644 index 00000000000..8a457b13ecd --- /dev/null +++ b/include/flashinfer/attention/sm120/nvfp4_attention_sm120/primitives/warpgroup.cuh @@ -0,0 +1,132 @@ +/* + * Copyright (c) 2025 by SageAttention team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "cute/tensor.hpp" +#include "cutlass/cutlass.h" + +namespace nvfp4_attention { + +using namespace cute; + +struct WarpGroupConstants { + static constexpr int kThreadsPerWarp = 32; + static constexpr int kWarpsPerGroup = 4; + static constexpr int kThreadsPerGroup = kThreadsPerWarp * kWarpsPerGroup; +}; + +struct WarpGroupIndex { + __device__ __forceinline__ static int lane_id() { + return threadIdx.x % WarpGroupConstants::kThreadsPerWarp; + } + + __device__ __forceinline__ static int warp_id() { + return threadIdx.x / WarpGroupConstants::kThreadsPerWarp; + } + + __device__ __forceinline__ static int warp_id_in_group() { + return warp_id() % WarpGroupConstants::kWarpsPerGroup; + } + + __device__ __forceinline__ static int warp_group_id() { + return warp_id() / WarpGroupConstants::kWarpsPerGroup; + } + + __device__ __forceinline__ static int thread_id_in_group() { + return warp_id_in_group() * WarpGroupConstants::kThreadsPerWarp + lane_id(); + } + + __device__ __forceinline__ static bool is_first_lane() { return lane_id() == 0; } + + __device__ __forceinline__ static bool is_first_thread_in_group() { + return thread_id_in_group() == 0; + } +}; + +struct WarpGroupSync { + __device__ __forceinline__ static void sync() { __syncwarp(); } + + template + __device__ __forceinline__ static T shuffle(T var, int src_lane) { + return __shfl_sync(0xffffffff, var, src_lane); + } + + template + __device__ __forceinline__ static T shuffle_xor(T var, int lane_mask) { + return __shfl_xor_sync(0xffffffff, var, lane_mask); + } + + __device__ __forceinline__ static float reduce_max(float val) { +#pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) { + val = fmaxf(val, shuffle_xor(val, mask)); + } + return val; + } + + __device__ __forceinline__ static float reduce_sum(float val) { +#pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) { + val += shuffle_xor(val, mask); + } + return val; + } +}; + +struct WarpGroupElect { + __device__ __forceinline__ static bool elect_one_in_warp() { return cute::elect_one_sync(); } + + __device__ __forceinline__ static bool elect_one_in_group() { + bool is_first = WarpGroupIndex::is_first_lane(); + + return is_first && (WarpGroupIndex::warp_id_in_group() == 0); + } +}; + +template +struct WarpGroupLayout { + __device__ __forceinline__ static auto get_thread_slice(TiledMMA const& tiled_mma, + int thread_idx) { + return tiled_mma.get_thread_slice(thread_idx); + } + + template + __device__ __forceinline__ static auto partition_accumulator(ThreadMMA const& thread_mma, + CoordTensor const& coord_tensor) { + return thread_mma.partition_C(coord_tensor); + } +}; + +template +struct WarpGroupMMA { + using Traits = TiledMMA; + + __device__ __forceinline__ static TiledMMA get_tiled_mma() { return TiledMMA{}; } + + __device__ __forceinline__ static auto get_thread_slice(int thread_idx) { + return get_tiled_mma().get_thread_slice(thread_idx); + } + + template + __device__ __forceinline__ static void gemm(TensorA const& A, TensorB const& B, TensorC& C) { + cute::gemm(get_tiled_mma(), A, B, C); + } +}; + +__device__ __forceinline__ int canonical_warp_idx() { return cutlass::canonical_warp_idx_sync(); } + +} // namespace nvfp4_attention diff --git a/include/flashinfer/attention/sm120/nvfp4_attention_sm120/quantization/fp4_convert.cuh b/include/flashinfer/attention/sm120/nvfp4_attention_sm120/quantization/fp4_convert.cuh new file mode 100644 index 00000000000..e2a0f9dbf6a --- /dev/null +++ b/include/flashinfer/attention/sm120/nvfp4_attention_sm120/quantization/fp4_convert.cuh @@ -0,0 +1,64 @@ +/* + * Copyright (c) 2025 by SageAttention team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include + +#include + +namespace nvfp4_attention { + +CUTLASS_DEVICE void packed_float_to_ue4m3(float const& f0, float const& f1, float const& f2, + float const& f3, uint32_t& out) { + out = flashinfer::math::fp32_vec_to_e4m3(f0, f1, f2, f3); +} + +CUTLASS_DEVICE void packed_float_to_e2m1(float const& f0, float const& f1, float const& f2, + float const& f3, float const& f4, float const& f5, + float const& f6, float const& f7, uint32_t& out) { + out = flashinfer::math::fp32_vec_to_e2m1(f0, f1, f2, f3, f4, f5, f6, f7); +} + +template +struct FP8E4M3Converter { + static_assert(N % 4 == 0, "N must be multiple of 4"); + + __device__ __forceinline__ static void convert(float const inputs[N], uint32_t outputs[N / 4]) { +#pragma unroll + for (int i = 0; i < N / 4; ++i) { + packed_float_to_ue4m3(inputs[i * 4 + 0], inputs[i * 4 + 1], inputs[i * 4 + 2], + inputs[i * 4 + 3], outputs[i]); + } + } +}; + +template +struct FP4E2M1Converter { + static_assert(N % 8 == 0, "N must be multiple of 8"); + + __device__ __forceinline__ static void convert(float const inputs[N], uint32_t outputs[N / 8]) { +#pragma unroll + for (int i = 0; i < N / 8; ++i) { + packed_float_to_e2m1(inputs[i * 8 + 0], inputs[i * 8 + 1], inputs[i * 8 + 2], + inputs[i * 8 + 3], inputs[i * 8 + 4], inputs[i * 8 + 5], + inputs[i * 8 + 6], inputs[i * 8 + 7], outputs[i]); + } + } +}; + +} // namespace nvfp4_attention diff --git a/include/flashinfer/attention/sm120/nvfp4_attention_sm120/quantization/fp4_layout.h b/include/flashinfer/attention/sm120/nvfp4_attention_sm120/quantization/fp4_layout.h new file mode 100644 index 00000000000..4ba3ae8dfe6 --- /dev/null +++ b/include/flashinfer/attention/sm120/nvfp4_attention_sm120/quantization/fp4_layout.h @@ -0,0 +1,143 @@ +/* + * Copyright (c) 2025 by SageAttention team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "cute/atom/mma_traits_sm100.hpp" +#include "cute/int_tuple.hpp" +#include "cutlass/layout/matrix.h" + +namespace nvfp4_attention { + +using namespace cute; + +template +struct BlockScaledBasicChunk { + using Blk_MN = _64; + + using Blk_SF = _4; + + using SfAtom = Layout, Shape, _4>>, + Stride, Stride<_0, _1>>>; +}; + +template +struct BlockScaledConfig { + static constexpr int SFVecSize = SFVecSize_; + static constexpr int MMA_NSF = 4; + + using BlkScaledChunk = BlockScaledBasicChunk; + using Blk_MN = _64; + using Blk_SF = _4; + + using mnBasicBlockShape = Shape<_16, _4>; + using mnBasicBlockStride = Stride<_16, _4>; + + using kBasicBlockShape = Shape, Int>; + using kBasicBlockStride = Stride<_0, _1>; + + using SfAtom = Layout, + Stride>; + + using LayoutSF = decltype(blocked_product( + SfAtom{}, make_layout(make_shape(int32_t(0), int32_t(0), int32_t(0), int32_t(0)), + make_stride(int32_t(0), _1{}, int32_t(0), int32_t(0))))); + + using Blk_Elems = decltype(Blk_MN{} * Blk_SF{}); + + using sSF_strideMN = decltype(prepend(Blk_Elems{}, mnBasicBlockStride{})); + + template + CUTE_HOST_DEVICE static constexpr auto tile_atom_to_shape_SFQKV(ProblemShape problem_shape) { + auto [Seqlen, Dim, HeadNum, Batch] = problem_shape; + return tile_to_shape(SfAtom{}, make_shape(Seqlen, Dim, HeadNum, Batch), Step<_2, _1, _3, _4>{}); + } + + template + CUTE_HOST_DEVICE static constexpr auto tile_atom_to_shape_SFVt(ProblemShape problem_shape) { + auto [Dim, Seqlen, HeadNum, Batch] = problem_shape; + return tile_to_shape(SfAtom{}, make_shape(Dim, Seqlen, HeadNum, Batch), Step<_2, _1, _3, _4>{}); + } + + template + CUTE_HOST_DEVICE static constexpr auto deduce_smem_layoutSFQ([[maybe_unused]] TiledMma tiled_mma, + TileShape_MNK tileshape_mnk) { + using sSFQ_shapeK = + decltype(prepend(make_shape(Blk_SF{} / Int{}, + size<2>(TileShape_MNK{}) / Int{} / Blk_SF{}), + kBasicBlockShape{})); + + using sSFQ_shapeM = decltype(prepend(size<0>(TileShape_MNK{}) / Blk_MN{}, mnBasicBlockShape{})); + + using sSFQ_strideM = sSF_strideMN; + using sSFQ_strideK = decltype(prepend( + make_stride(Int{}, size<0>(TileShape_MNK{}) / Blk_MN{} * Blk_Elems{}), + kBasicBlockStride{})); + + using sSFQ_shape = decltype(make_shape(sSFQ_shapeM{}, sSFQ_shapeK{})); + using sSFQ_stride = decltype(make_stride(sSFQ_strideM{}, sSFQ_strideK{})); + using SmemLayoutAtomSFQ = decltype(make_layout(sSFQ_shape{}, sSFQ_stride{})); + + return SmemLayoutAtomSFQ{}; + } + + template + CUTE_HOST_DEVICE static constexpr auto deduce_smem_layoutSFKV([[maybe_unused]] TiledMma tiled_mma, + TileShape_MNK tileshape_mnk) { + using sSFK_shapeK = + decltype(prepend(make_shape(Blk_SF{} / Int{}, + size<2>(TileShape_MNK{}) / Int{} / Blk_SF{}), + kBasicBlockShape{})); + + using sSFK_shapeN = decltype(prepend(size<1>(TileShape_MNK{}) / Blk_MN{}, mnBasicBlockShape{})); + + using sSFK_strideN = sSF_strideMN; + using sSFK_strideK = decltype(prepend( + make_stride(Int{}, size<1>(TileShape_MNK{}) / Blk_MN{} * Blk_Elems{}), + kBasicBlockStride{})); + + using sSFK_shape = decltype(make_shape(sSFK_shapeN{}, sSFK_shapeK{})); + using sSFK_stride = decltype(make_stride(sSFK_strideN{}, sSFK_strideK{})); + using SmemLayoutAtomSFK = decltype(make_layout(sSFK_shape{}, sSFK_stride{})); + + return SmemLayoutAtomSFK{}; + } + + template + CUTE_HOST_DEVICE static constexpr auto deduce_smem_layoutSFVt([[maybe_unused]] TiledMma tiled_mma, + TileShape_MNK tileshape_mnk) { + using sSFVt_shapeK = + decltype(prepend(make_shape(Blk_SF{} / Int{}, + size<2>(TileShape_MNK{}) / Int{} / Blk_SF{}), + kBasicBlockShape{})); + + using sSFVt_shapeN = + decltype(prepend(size<1>(TileShape_MNK{}) / Blk_MN{}, mnBasicBlockShape{})); + + using sSFVt_strideN = sSF_strideMN; + using sSFVt_strideK = decltype(prepend( + make_stride(Int{}, size<1>(TileShape_MNK{}) / Blk_MN{} * Blk_Elems{}), + kBasicBlockStride{})); + + using sSFVt_shape = decltype(make_shape(sSFVt_shapeN{}, sSFVt_shapeK{})); + using sSFVt_stride = decltype(make_stride(sSFVt_strideN{}, sSFVt_strideK{})); + using SmemLayoutAtomSFVt = decltype(make_layout(sSFVt_shape{}, sSFVt_stride{})); + + return SmemLayoutAtomSFVt{}; + } +}; + +} // namespace nvfp4_attention diff --git a/include/flashinfer/attention/sm120/nvfp4_attention_sm120/utils/copy.cuh b/include/flashinfer/attention/sm120/nvfp4_attention_sm120/utils/copy.cuh new file mode 100644 index 00000000000..ada02b42d8d --- /dev/null +++ b/include/flashinfer/attention/sm120/nvfp4_attention_sm120/utils/copy.cuh @@ -0,0 +1,59 @@ +/* + * Copyright (c) 2025 by SageAttention team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include + +#include "cute/tensor.hpp" + +namespace nvfp4_attention { + +using namespace cute; + +template +CUTLASS_DEVICE void copy(TiledCopy tiled_copy, Tensor const& S, + Tensor& D, Tensor const& identity_MN, + Tensor const& predicate_K, const int max_MN = 0) { + CUTE_STATIC_ASSERT_V(rank(S) == Int<3>{}); + CUTE_STATIC_ASSERT_V(rank(D) == Int<3>{}); + CUTE_STATIC_ASSERT_V(size<0>(S) == size<0>(D)); + CUTE_STATIC_ASSERT_V(size<1>(S) == size<1>(D)); + CUTE_STATIC_ASSERT_V(size<2>(S) == size<2>(D)); + + static_assert(!(Clear_OOB_MN && !Clear_OOB_K), "Cannot clear OOB_MN without clearing OOB_K"); + +#pragma unroll + for (int m = 0; m < size<1>(S); ++m) { + if (Is_even_MN || get<0>(identity_MN(0, m, 0)) < max_MN) { +#pragma unroll + for (int k = 0; k < size<2>(S); ++k) { + if (Is_even_K || predicate_K(k)) { + cute::copy(tiled_copy, S(_, m, k), D(_, m, k)); + } else if (Clear_OOB_K) { + cute::clear(D(_, m, k)); + } + } + } else if (Clear_OOB_MN) { + cute::clear(D(_, m, _)); + } + } +} + +} // namespace nvfp4_attention diff --git a/include/flashinfer/attention/sm120/nvfp4_attention_sm120/utils/layout.cuh b/include/flashinfer/attention/sm120/nvfp4_attention_sm120/utils/layout.cuh new file mode 100644 index 00000000000..ab33d0780d3 --- /dev/null +++ b/include/flashinfer/attention/sm120/nvfp4_attention_sm120/utils/layout.cuh @@ -0,0 +1,61 @@ +/* + * Copyright (c) 2025 by SageAttention team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "cute/tensor.hpp" + +namespace nvfp4_attention { + +using namespace cute; + +template +CUTLASS_DEVICE constexpr auto convert_to_reduction_layout(Layout mma_layout) { + static_assert(rank(mma_layout) == 3, "Mma Layout should be (MmaAtom, MmaM, MmaN)"); + static_assert(rank(get<0>(shape(mma_layout))) == 2, "MmaAtom should be (AtomN, AtomM)"); + + return make_layout(make_layout(get<0, 1>(mma_layout), get<1>(mma_layout)), + make_layout(get<0, 0>(mma_layout), get<2>(mma_layout))); +} + +template +CUTLASS_DEVICE constexpr auto convert_to_conversion_layout(Layout mma_layout) { + static_assert(rank(mma_layout) == 3, "Mma Layout should be (MmaAtom, MmaM, MmaN)"); + static_assert(rank(get<0>(shape(mma_layout))) == 2, "MmaAtom should be (AtomN, AtomM)"); + + constexpr int MmaAtomN = size<0, 0>(mma_layout); + constexpr int MmaAtomM = size<0, 1>(mma_layout); + constexpr int MmaM = size<1>(mma_layout); + constexpr int MmaN = size<2>(mma_layout); + + static_assert(MmaAtomN % 8 == 0, "MmaAtomN should be multiple of 8."); + static_assert(MmaAtomM == 2, "MmaAtomM should be 2."); + static_assert(MmaN % 2 == 0, "MmaN should be multiple of 2."); + + auto mma_n_division = zipped_divide(layout<2>(mma_layout), make_tile(_2{})); + return make_layout(make_layout(layout<0, 0>(mma_layout), + make_layout(layout<0, 1>(mma_layout), layout<0>(mma_n_division))), + layout<1>(mma_layout), layout<1>(mma_n_division)); +} + +CUTLASS_DEVICE constexpr int qk_acc_col_to_k_col(int col) { + int const col_in_mma = col & 31; + int const pair = col_in_mma >> 1; + int const k_pair = ((pair & 3) << 2) | (pair >> 2); + return (col & ~31) + (k_pair << 1) + (col_in_mma & 1); +} + +} // namespace nvfp4_attention diff --git a/include/flashinfer/attention/sm120/nvfp4_attention_sm120/utils/math.cuh b/include/flashinfer/attention/sm120/nvfp4_attention_sm120/utils/math.cuh new file mode 100644 index 00000000000..30bbf6cd409 --- /dev/null +++ b/include/flashinfer/attention/sm120/nvfp4_attention_sm120/utils/math.cuh @@ -0,0 +1,31 @@ +/* + * Copyright (c) 2025 by SageAttention team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include + +namespace nvfp4_attention { + +using flashinfer::math::add; +using flashinfer::math::exp2_fma_poly; +using flashinfer::math::fma; +using flashinfer::math::MaxOp; +using flashinfer::math::mul; +using flashinfer::math::ptx_exp2; +using flashinfer::math::SumOp; + +} // namespace nvfp4_attention diff --git a/include/flashinfer/math.cuh b/include/flashinfer/math.cuh index 27c6351e8ff..3b47b2cc186 100644 --- a/include/flashinfer/math.cuh +++ b/include/flashinfer/math.cuh @@ -19,6 +19,7 @@ #include #include +#include #include namespace flashinfer { @@ -45,6 +46,121 @@ __forceinline__ __device__ float ptx_exp2(float x) { return y; } +template +struct MaxOp { + __device__ __forceinline__ T operator()(T const& x, T const& y) { return x > y ? x : y; } +}; + +template <> +struct MaxOp { + __device__ __forceinline__ float operator()(float const& x, float const& y) { + return fmaxf(x, y); + } +}; + +template +struct SumOp { + __device__ __forceinline__ T operator()(T const& x, T const& y) { return x + y; } +}; + +__forceinline__ __device__ void add(float2& c, float2 const& a, float2 const& b) { + c.x = a.x + b.x; + c.y = a.y + b.y; +} + +__forceinline__ __device__ void mul(float2& c, float2 const& a, float2 const& b) { + c.x = a.x * b.x; + c.y = a.y * b.y; +} + +__forceinline__ __device__ void fma(float2& d, float2 const& a, float2 const& b, float2 const& c) { + d.x = fmaf(a.x, b.x, c.x); + d.y = fmaf(a.y, b.y, c.y); +} + +__forceinline__ __device__ float exp2_fma_poly(float x) { + float n = rintf(x); + float f = x - n; + + float poly = ((0.0771f * f + 0.2276f) * f + 0.6951f) * f + 1.0f; + + int n_int = __float2int_rn(n); + uint32_t poly_bits = __float_as_uint(poly); + poly_bits += static_cast(n_int) << 23; + return __uint_as_float(poly_bits); +} + +__forceinline__ __device__ uint32_t fp32_vec_to_e4m3(float const& f0, float const& f1, + float const& f2, float const& f3) { +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + uint32_t val; + asm volatile( + "{\n" + ".reg .b16 lo;\n" + ".reg .b16 hi;\n" + "cvt.rn.satfinite.e4m3x2.f32 lo, %2, %1;\n" + "cvt.rn.satfinite.e4m3x2.f32 hi, %4, %3;\n" + "mov.b32 %0, {lo, hi};\n" + "}" + : "=r"(val) + : "f"(f0), "f"(f1), "f"(f2), "f"(f3)); + return val; +#elif defined(__CUDA_ARCH__) + asm volatile("trap;"); + return 0; +#else + return 0; +#endif +} + +__forceinline__ __device__ uint32_t fp32_vec_to_e4m3(float const (&array)[4]) { + return fp32_vec_to_e4m3(array[0], array[1], array[2], array[3]); +} + +__forceinline__ __device__ uint32_t fp32_vec_to_e2m1(float const& f0, float const& f1, + float const& f2, float const& f3, + float const& f4, float const& f5, + float const& f6, float const& f7) { +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + uint32_t val; + asm volatile( + "{\n" + ".reg .b8 byte0;\n" + ".reg .b8 byte1;\n" + ".reg .b8 byte2;\n" + ".reg .b8 byte3;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte0, %2, %1;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte1, %4, %3;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte2, %6, %5;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte3, %8, %7;\n" + "mov.b32 %0, {byte0, byte1, byte2, byte3};\n" + "}" + : "=r"(val) + : "f"(f0), "f"(f1), "f"(f2), "f"(f3), "f"(f4), "f"(f5), "f"(f6), "f"(f7)); + return val; +#elif defined(__CUDA_ARCH__) + asm volatile("trap;"); + return 0; +#else + return 0; +#endif +} + +__forceinline__ __device__ uint32_t fp32_vec_to_e2m1(float const (&array)[8]) { + return fp32_vec_to_e2m1(array[0], array[1], array[2], array[3], array[4], array[5], array[6], + array[7]); +} + +__forceinline__ __device__ uint32_t fp32_vec_to_e2m1(float2 const (&array)[4]) { + return fp32_vec_to_e2m1(array[0].x, array[0].y, array[1].x, array[1].y, array[2].x, array[2].y, + array[3].x, array[3].y); +} + +__forceinline__ __device__ uint32_t fp32_vec_to_e2m1(float2 const* array) { + return fp32_vec_to_e2m1(array[0].x, array[0].y, array[1].x, array[1].y, array[2].x, array[2].y, + array[3].x, array[3].y); +} + /*! * \brief Wrapper of PTX lg2.approx instruction, which computes log2(x) * \param x input diff --git a/tests/attention/test_nvfp4_attention_sm120.py b/tests/attention/test_nvfp4_attention_sm120.py new file mode 100644 index 00000000000..974213276e3 --- /dev/null +++ b/tests/attention/test_nvfp4_attention_sm120.py @@ -0,0 +1,222 @@ +""" +Copyright (c) 2026 by FlashInfer team. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +import pytest +import torch +import torch.nn.functional as F + + +def _patch_cutlass_dsl_operand_major_mode(): + try: + import cutlass.cute as cute + from cutlass.cute.nvgpu.tcgen05 import OperandMajorMode + except ImportError: + return + if not hasattr(cute.nvgpu, "OperandMajorMode"): + cute.nvgpu.OperandMajorMode = OperandMajorMode + + +_patch_cutlass_dsl_operand_major_mode() + +import flashinfer +from flashinfer.utils import is_sm120a_supported + + +def _require_sm120(): + if not torch.cuda.is_available(): + pytest.skip("CUDA is required") + if not is_sm120a_supported(torch.device("cuda")): + pytest.skip("SM120 GPU is required") + + +def _pad_seq_len_to_128(x): + pad_len = (-x.shape[2]) % 128 + if pad_len == 0: + return x.contiguous() + return torch.nn.functional.pad(x, (0, 0, 0, pad_len), value=0).contiguous() + + +def _preprocess_qkv_ref(q, k, v): + k = k - k.mean(dim=-2, keepdim=True) + q, k, v = map(_pad_seq_len_to_128, (q, k, v)) + batch, num_heads, seq_len, head_dim = q.shape + q_grouped = q.reshape(batch, num_heads, seq_len // 128, 128, head_dim) + qm = q_grouped.mean(dim=3) + q = ( + (q_grouped - qm.unsqueeze(3)) + .reshape(batch, num_heads, seq_len, head_dim) + .contiguous() + ) + qk_correction = ( + torch.matmul(qm, k.transpose(-2, -1)) + .repeat_interleave(128, dim=2) + .to(torch.float32) + .contiguous() + ) + return q, k, v, qk_correction + + +def _reference_attention(q, k, v, causal): + q, k, v, qk_correction = _preprocess_qkv_ref(q, k, v) + sm_scale = q.shape[-1] ** -0.5 + scores = torch.matmul(q.float(), k.float().transpose(-2, -1)) * sm_scale + scores = scores + qk_correction * sm_scale + if causal: + seq_len = q.shape[2] + mask = torch.triu( + torch.ones(seq_len, seq_len, device=q.device, dtype=torch.bool), diagonal=1 + ) + scores.masked_fill_(mask, float("-inf")) + probs = torch.softmax(scores, dim=-1) + return torch.matmul(probs, v.float()).to(q.dtype) + + +def _run_nvfp4_attention_sm120_accuracy_case( + batch, + num_heads, + seq_len, + head_dim, + causal, + cos_threshold, + mean_abs_err_threshold, +): + _require_sm120() + + torch.manual_seed(42) + q = torch.randn( + (batch, num_heads, seq_len, head_dim), device="cuda", dtype=torch.bfloat16 + ) + k = torch.randn_like(q) + v = torch.randn_like(q) + + q_fp4, k_fp4, v_fp4_t, q_scale, k_scale, v_scale_t, qk_correction = ( + flashinfer.nvfp4_attention_sm120_quantize_qkv(q, k, v) + ) + + out, lse = flashinfer.nvfp4_attention_sm120_fwd( + q_fp4, + k_fp4, + v_fp4_t, + q_scale, + k_scale, + v_scale_t, + qk_correction, + sm_scale=head_dim**-0.5, + causal=causal, + ) + + torch.cuda.synchronize() + ref = _reference_attention(q, k, v, causal)[:, :, :seq_len, :] + out = out[:, :, :seq_len, :] + lse = lse[:, :, :seq_len] + + assert out.shape == ref.shape + assert lse.shape == (batch, num_heads, seq_len) + assert out.dtype == torch.bfloat16 + assert lse.dtype == torch.float32 + assert not torch.isnan(out).any() + assert not torch.isinf(out).any() + assert not torch.isnan(lse).any() + assert not torch.isinf(lse).any() + + mean_abs_err = (out.float() - ref.float()).abs().mean().item() + cos_sim = F.cosine_similarity( + out.float().reshape(1, -1), ref.float().reshape(1, -1) + ).item() + assert mean_abs_err <= mean_abs_err_threshold + assert cos_sim >= cos_threshold + + +@pytest.mark.parametrize( + ( + "batch", + "num_heads", + "seq_len", + "head_dim", + "causal", + "cos_threshold", + "mean_abs_err_threshold", + ), + [ + pytest.param(1, 4, 128, 64, False, 0.95, 0.08, id="s128-d64-noncausal"), + pytest.param(1, 4, 256, 128, False, 0.95, 0.06, id="s256-d128-noncausal"), + pytest.param(1, 4, 256, 128, True, 0.94, 0.09, id="s256-d128-causal"), + pytest.param(1, 1, 4096, 64, False, 0.95, 0.02, id="s4096-d64-noncausal"), + pytest.param(1, 1, 4096, 128, True, 0.95, 0.04, id="s4096-d128-causal"), + pytest.param(1, 1, 8192, 64, False, 0.95, 0.02, id="s8192-d64-noncausal"), + ], +) +@torch.inference_mode() +def test_nvfp4_attention_sm120_accuracy( + batch, + num_heads, + seq_len, + head_dim, + causal, + cos_threshold, + mean_abs_err_threshold, +): + _run_nvfp4_attention_sm120_accuracy_case( + batch, + num_heads, + seq_len, + head_dim, + causal, + cos_threshold, + mean_abs_err_threshold, + ) + + +@torch.inference_mode() +def test_nvfp4_attention_sm120_causal_mask_column_order(): + _require_sm120() + + seq_len = 128 + head_dim = 128 + q = torch.zeros((1, 1, seq_len, head_dim), device="cuda", dtype=torch.bfloat16) + k = torch.zeros_like(q) + v = torch.eye(seq_len, device="cuda", dtype=torch.bfloat16).reshape( + 1, 1, seq_len, head_dim + ) + + q_fp4, k_fp4, v_fp4_t, q_scale, k_scale, v_scale_t, qk_correction = ( + flashinfer.nvfp4_attention_sm120_quantize_qkv(q, k, v) + ) + out, _ = flashinfer.nvfp4_attention_sm120_fwd( + q_fp4, + k_fp4, + v_fp4_t, + q_scale, + k_scale, + v_scale_t, + qk_correction, + sm_scale=head_dim**-0.5, + causal=True, + ) + torch.cuda.synchronize() + out = out[0, 0].float() + + ref = torch.zeros((seq_len, head_dim), device="cuda") + for row in range(seq_len): + ref[row, : row + 1] = 1.0 / (row + 1) + + suffix_max = torch.stack( + [out[row, row + 1 :].abs().max() for row in range(seq_len - 1)] + ).max() + cos_sim = F.cosine_similarity(out.reshape(1, -1), ref.reshape(1, -1)).item() + + assert suffix_max <= 1e-5 + assert cos_sim >= 0.98 diff --git a/tests/conftest.py b/tests/conftest.py index 8011d43eb92..0cfa56d684b 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -9,6 +9,19 @@ from torch.torch_version import TorchVersion from torch.torch_version import __version__ as torch_version + +def _patch_cutlass_dsl_operand_major_mode(): + try: + import cutlass.cute as cute + from cutlass.cute.nvgpu.tcgen05 import OperandMajorMode + except ImportError: + return + if not hasattr(cute.nvgpu, "OperandMajorMode"): + cute.nvgpu.OperandMajorMode = OperandMajorMode + + +_patch_cutlass_dsl_operand_major_mode() + import flashinfer from flashinfer.jit import MissingJITCacheError