fp8 unit tests, test refactor

Author: jimmyzho

csrc/fmha_v2_jit_binding.cu

@@ -28,11 +28,11 @@ void fmha_v2_run(ffi::TensorView q, ffi::TensorView k, ffi::TensorView v, ffi::T
                  ffi::TensorView workspace_buffer, size_t workspace_buffer_size_in_bytes,
                  Optional<ffi::TensorView> maybe_block_tables, int page_size,
                  ffi::TensorView seq_lens, ffi::TensorView cum_seq_lens_q,
-                 ffi::TensorView cum_seq_lens_kv, const std::string& input_layout_str, int max_q_len,
-                 int max_kv_len, int batch_size, int total_q_tokens, int total_kv_tokens,
-                 const std::string& mask_mode_str, float scale_softmax, float scale_bmm1, float scale_bmm2,
-                 int window_left, int chunked_attention_size, bool has_alibi,
-                 float softcapping_scale, ffi::TensorView scale_bmm2_d,
+                 ffi::TensorView cum_seq_lens_kv, const std::string& input_layout_str,
+                 int max_q_len, int max_kv_len, int batch_size, int total_q_tokens,
+                 int total_kv_tokens, const std::string& mask_mode_str, float scale_softmax,
+                 float scale_bmm1, float scale_bmm2, int window_left, int chunked_attention_size,
+                 bool has_alibi, float softcapping_scale, ffi::TensorView scale_bmm2_d,
                  Optional<ffi::TensorView> softmax_stats, Optional<ffi::TensorView> sinks);
 
 // FMHAv2 attention operator

csrc/fmha_v2_run.cu

@@ -291,7 +291,6 @@ static inline Data_type dltype_to_data_type(DLDataType dtype) {
   return DATA_TYPE_FP16;
 }
 
-
 static inline Attention_mask_type string_to_mask_type(const std::string& s) {
   if (s == "padding") return Attention_mask_type::PADDING;
   if (s == "causal") return Attention_mask_type::CAUSAL;
@@ -309,7 +308,6 @@ static inline Attention_input_layout string_to_input_layout(const std::string& s
   return Attention_input_layout::Q_PAGED_KV;  // default
 }
 
-
 void fmha_v2_run(
     ffi::TensorView q,  // [batch, s_q, num_heads, head_dim]
     ffi::TensorView k,  // [batch, s_kv, num_kv_heads, head_dim]
@@ -321,12 +319,11 @@ void fmha_v2_run(
     ffi::TensorView seq_lens,         // [batch]
     ffi::TensorView cum_seq_lens_q,   // [batch + 1]
     ffi::TensorView cum_seq_lens_kv,  // [batch + 1]
-    const std::string& input_layout_str,
-    int max_q_len, int max_kv_len, int batch_size, int total_q_tokens,
-    int total_kv_tokens,     // Totals from cum_seq_lens (computed in Python)
-    const std::string& mask_mode_str,
-    float scale_softmax, float scale_bmm1, float scale_bmm2, int window_left,
-    int chunked_attention_size, bool has_alibi, float softcapping_scale,
+    const std::string& input_layout_str, int max_q_len, int max_kv_len, int batch_size,
+    int total_q_tokens,
+    int total_kv_tokens,  // Totals from cum_seq_lens (computed in Python)
+    const std::string& mask_mode_str, float scale_softmax, float scale_bmm1, float scale_bmm2,
+    int window_left, int chunked_attention_size, bool has_alibi, float softcapping_scale,
     ffi::TensorView scale_bmm2_d,             // Pre-populated scale_bmm2 on device [1] int32
     Optional<ffi::TensorView> softmax_stats,  // Optional [batch, s_q, num_heads, 2] for (max, sum)
     Optional<ffi::TensorView> sinks) {
@@ -473,7 +470,8 @@ void fmha_v2_run(
   std::tie(warps_m, warps_n, warps_k) = get_warps(launch_params, sm, data_type, s, b, d, 2);
 
   // Debug output for warps
-  printf("DEBUG: get_warps returned warps_m=%zu, warps_n=%zu, warps_k=%zu\n", warps_m, warps_n, warps_k);
+  printf("DEBUG: get_warps returned warps_m=%zu, warps_n=%zu, warps_k=%zu\n", warps_m, warps_n,
+         warps_k);
   printf("DEBUG: launch_params: flash_attention=%d, warp_specialization=%d, use_tma=%d\n",
          launch_params.flash_attention, launch_params.warp_specialization, launch_params.use_tma);
   printf("DEBUG: data_type=%d, sm=%d, s=%zu, d=%zu\n", int(data_type), sm, s, d);

flashinfer/jit/attention/fmha_v2/fmha_library.py

@@ -1129,15 +1129,13 @@ def gen_cta_spec(spec):
     return api_code
 
 
-def generate_jit_sources(input_layout: str) -> list:
-    uri = "trtllm_fmha_v2"
+def generate_jit_sources(uri: str, input_layout: str, input_dtype: str, output_dtype: str) -> list:
     gen_directory = jit_env.FLASHINFER_GEN_SRC_DIR / uri
     source_paths = []
     specs_names = []
-    dtype_values = ["fp16", "bf16", "e4m3"]
     head_size_qk_values = [16, 32, 64, 128, 256, 512]
     head_size_qk_warpspec_values = [32, 40, 48, 64, 80, 96, 104, 128, 160, 192, 256]
-      
+
     # 0 means head_size_v = head_size_qk (required for flash_valid)
     head_size_v_values = [0]
     map_input_layout = {
@@ -1146,8 +1144,11 @@ def generate_jit_sources(input_layout: str) -> list:
         "separate_q_k_v": InputLayout.SEPARATE_Q_K_V,
         "contiguous_q_kv": InputLayout.CONTIGUOUS_Q_KV,
     }
+
     input_layout_values = [map_input_layout[input_layout.lower()]]
-    output_dtype_values = ["fp16", "bf16"]
+    dtype_values = [input_dtype]
+    output_dtype_values = [output_dtype] if output_dtype is not None else [None]
+
     is_mla_values = [False]
 
     enable_attn_logit_softcapping_values = [True, False]

flashinfer/jit/attention/modules.py

@@ -1936,17 +1936,28 @@ def gen_trtllm_fmha_v2_sm120_module(device: torch.device) -> JitSpec:
     )
 
 
-def gen_fmha_v2_module(input_layout: str) -> JitSpec:
-    uri = "trtllm_fmha_v2"
-
+def gen_fmha_v2_module(input_layout: str, input_dtype: torch.dtype, output_dtype: torch.dtype = None) -> JitSpec:
     # Setup generated source directory
+    if output_dtype is None:
+        output_dtype = input_dtype
+    
+    dtype_map = {
+        torch.float16: "fp16",
+        torch.bfloat16: "bf16",
+        torch.float8_e4m3fn: "e4m3",
+    }
+    input_dtype_str = dtype_map[input_dtype]
+    output_dtype_str = dtype_map[output_dtype] if output_dtype is not None else None
+
+    uri = f"trtllm_fmha_v2_{input_layout.lower()}_{input_dtype_str}_{output_dtype_str}"
+
     gen_directory = jit_env.FLASHINFER_GEN_SRC_DIR / uri
     gen_directory.mkdir(parents=True, exist_ok=True)
 
     # Source directories
     csrc_dir = jit_env.FLASHINFER_CSRC_DIR
     fmha_v2_src_dir = csrc_dir / "fmha_v2"
-    source_paths = generate_jit_sources(input_layout)
+    source_paths = generate_jit_sources(uri, input_layout, input_dtype_str, output_dtype_str)
 
     # copy static fmha_v2_run.cu
     static_run_path = csrc_dir / "fmha_v2_run.cu"

flashinfer/prefill.py

@@ -3873,8 +3873,8 @@ def fmha_v2_prefill_deepseek(
 
 
 @functools.cache
-def get_trtllm_fmha_v2_module(input_layout: str):
-    return gen_fmha_v2_module(input_layout).build_and_load()
+def get_trtllm_fmha_v2_module(input_layout: str, input_dtype: torch.dtype, output_dtype: torch.dtype = None):
+    return gen_fmha_v2_module(input_layout, input_dtype, output_dtype).build_and_load()
 
 
 @flashinfer_api
@@ -4018,6 +4018,16 @@ def trtllm_fmha_v2_prefill(
         elif len(qkv) == 3:
             input_layout = "SEPARATE_Q_K_V"
             query, k_cache, v_cache = qkv
+            if hasattr(torch, "float8_e4m3fn") and query.dtype == torch.float8_e4m3fn:
+                raise ValueError(
+                    "FP8 (e4m3) is not supported for the SEPARATE_Q_K_V input layout. "
+                    "Use PACKED_QKV, CONTIGUOUS_Q_KV, or Q_PAGED_KV layout instead."
+                )
+            if logits_soft_cap_scale is not None and logits_soft_cap_scale > 0:
+                raise ValueError(
+                    "Logits soft capping is not supported for the SEPARATE_Q_K_V input layout. "
+                    "Use PACKED_QKV, CONTIGUOUS_Q_KV, or Q_PAGED_KV layout instead."
+                )
 
         else:
             raise ValueError(
@@ -4102,7 +4112,7 @@ def trtllm_fmha_v2_prefill(
         logits_soft_cap_scale if logits_soft_cap_scale is not None else 0.0
     )
 
-    module = get_trtllm_fmha_v2_module(input_layout)
+    module = get_trtllm_fmha_v2_module(input_layout, query.dtype, o_dtype if query.dtype == torch.float8_e4m3fn else None)
     total_q_tokens = int(cum_seq_lens_q[-1].item())
     total_kv_tokens = int(cum_seq_lens_kv[-1].item())