Dml dev

flashinfer-jit-cache/build_backend.py

@@ -82,53 +82,28 @@ def _compile_jit_cache(output_dir: Path, verbose: bool = True):
     # Get the project root directory
     project_root = Path(__file__).parent.parent
 
-    # Ensure 3rdparty submodules are populated (may be empty in CI Docker images).
-    # Skip if submodules are already present or if git metadata is incomplete
-    # (e.g., Docker builds where .git points to a parent repo not in the context).
-    import subprocess
-
-    submodule_check_paths = [
-        project_root / "3rdparty" / "cutlass" / "include",
-        project_root / "3rdparty" / "spdlog" / "include",
-        project_root / "3rdparty" / "cccl" / "cub",
-    ]
-    if not all(p.exists() for p in submodule_check_paths):
-        result = subprocess.run(
-            ["git", "submodule", "update", "--init", "--recursive"],
-            cwd=str(project_root),
-            capture_output=True,
-        )
-        if result.returncode != 0:
-            missing = [str(p) for p in submodule_check_paths if not p.exists()]
-            if missing:
-                raise RuntimeError(
-                    f"git submodule update failed and submodules are missing: {missing}\n"
-                    f"git stderr: {result.stderr.decode().strip()}"
-                )
-
-    # Ensure flashinfer/data/ symlinks exist (normally created by the main
-    # package's build_backend, but jit-cache builds may not install the main
-    # package first). Use importlib to avoid name collision with this file.
-    import importlib.util
-
-    spec = importlib.util.spec_from_file_location(
-        "main_build_backend", project_root / "build_backend.py"
-    )
-    main_build_backend = importlib.util.module_from_spec(spec)
-    spec.loader.exec_module(main_build_backend)
-    main_build_backend._create_data_dir(use_symlinks=True)
-
-    from flashinfer import aot
+    build_profile = os.environ.get("FLASHINFER_AOT_BUILD_PROFILE")
+    if not build_profile:
+        if any(
+            os.environ.get(var)
+            for var in ("CI", "GITHUB_ACTIONS", "JENKINS_HOME", "JENKINS_URL")
+        ):
+            build_profile = "full"
+        else:
+            build_profile = "edge_fm"
+        os.environ["FLASHINFER_AOT_BUILD_PROFILE"] = build_profile
 
     # Set up build directory
     build_dir = project_root / "build" / "aot"
 
     # Use the centralized compilation function from aot.py
+    print(f"Using FLASHINFER_AOT_BUILD_PROFILE={build_profile}")
     aot.compile_and_package_modules(
         out_dir=output_dir,
         build_dir=build_dir,
         project_root=project_root,
-        config=None,  # Use default config
+        config=None,
+        profile=build_profile,
         verbose=verbose,
         skip_prebuilt=False,
     )

flashinfer/aot.py

@@ -21,6 +21,7 @@
 """
 
 import argparse
+import copy
 import os
 import shutil
 from itertools import product
@@ -110,6 +111,90 @@
 from .jit.xqa import gen_xqa_module, gen_xqa_module_mla
 
 
+def get_default_config():
+    """Get the full upstream AOT configuration."""
+    return {
+        "fa2_head_dim": [(64, 64), (128, 128), (256, 256)],
+        "fa3_head_dim": [(192, 128), (128, 128), (64, 64), (256, 256)],
+        "f16_dtype": [torch.float16, torch.bfloat16],
+        "f8_dtype": [torch.float8_e4m3fn],
+        "use_sliding_window": [False, True],
+        "use_logits_soft_cap": [False, True],
+        "add_comm": True,
+        "add_gemma": True,
+        "add_oai_oss": True,
+        "add_moe": True,
+        "add_act": True,
+        "add_misc": True,
+        "add_xqa": True,
+    }
+
+
+def get_edge_fm_fast_config():
+    """
+    Get the trimmed AOT configuration used for fast local builds in edge-fm.
+
+    This keeps only the kernels currently exercised in the repo:
+    - single/batch prefill/decode attention on common 64/128 head dims
+    - activation / norm / rope / page / sampling / topk helpers
+
+    It intentionally excludes heavyweight optional families such as XQA, MoE,
+    communication kernels, FP8 attention variants, and large head-dim matrices
+    that are not used in edge-fm today.
+    """
+    return {
+        "fa2_head_dim": [(64, 64), (128, 128)],
+        "fa3_head_dim": [(64, 64), (128, 128)],
+        "f16_dtype": [torch.float16, torch.bfloat16],
+        "f8_dtype": [],
+        "use_sliding_window": [False],
+        "use_logits_soft_cap": [False],
+        "add_comm": False,
+        "add_gemma": False,
+        "add_oai_oss": False,
+        "add_moe": False,
+        "add_act": True,
+        "add_misc": True,
+        "add_xqa": False,
+    }
+
+
+def normalize_build_profile(profile: Optional[str]) -> str:
+    profile_name = (profile or os.environ.get("FLASHINFER_AOT_BUILD_PROFILE") or "full")
+    profile_name = profile_name.strip().lower().replace("-", "_")
+    aliases = {
+        "default": "full",
+        "minimal": "edge_fm",
+        "fast": "edge_fm",
+        "dev": "edge_fm",
+        "edgefm": "edge_fm",
+    }
+    return aliases.get(profile_name, profile_name)
+
+
+def get_config_for_profile(profile: Optional[str]) -> dict:
+    normalized = normalize_build_profile(profile)
+    if normalized == "full":
+        return get_default_config()
+    if normalized == "edge_fm":
+        return get_edge_fm_fast_config()
+    raise ValueError(
+        f"Unknown FLASHINFER AOT build profile: {profile!r}. "
+        "Supported values: full, edge_fm."
+    )
+
+
+def resolve_build_config(
+    config: Optional[dict] = None,
+    profile: Optional[str] = None,
+) -> Tuple[str, dict]:
+    normalized = normalize_build_profile(profile)
+    final_config = copy.deepcopy(get_config_for_profile(normalized))
+    if config is not None:
+        final_config.update(config)
+    return normalized, final_config
+
+
 def gen_fa2(
     dtype_qo: torch.dtype,
     dtype_kv: torch.dtype,
@@ -750,6 +835,7 @@ def compile_and_package_modules(
     build_dir: Path,
     project_root: Path,
     config: dict = None,
+    profile: Optional[str] = None,
     verbose: bool = False,
     skip_prebuilt: bool = True,
 ) -> None:
@@ -764,11 +850,7 @@ def compile_and_package_modules(
         verbose: Whether to print verbose build output
         skip_prebuilt: Whether to skip pre-built modules
     """
-    # Start with default config and override with user config
-    final_config = get_default_config()
-    if config is not None:
-        final_config.update(config)
-    config = final_config
+    build_profile, config = resolve_build_config(config, profile)
     # Cuda Arch
     if "FLASHINFER_CUDA_ARCH_LIST" not in os.environ:
         raise RuntimeError("Please explicitly set env var FLASHINFER_CUDA_ARCH_LIST.")
@@ -797,6 +879,7 @@ def compile_and_package_modules(
         if out_dir is not None:
             print("  out_dir:", out_dir)
         print("  build_dir:", build_dir)
+        print("  build_profile:", build_profile)
         print("  fa2_head_dim:", config["fa2_head_dim"])
         print("  fa3_head_dim:", config["fa3_head_dim"])
         print("  f16_dtype:", config["f16_dtype"])
@@ -865,25 +948,6 @@ def parse_head_dim(head_dim: str) -> Tuple[int, int]:
     return qo, kv
 
 
-def get_default_config():
-    """Get default AOT configuration"""
-    return {
-        "fa2_head_dim": [(64, 64), (128, 128), (256, 256)],
-        "fa3_head_dim": [(192, 128), (128, 128), (64, 64), (256, 256)],
-        "f16_dtype": [torch.float16, torch.bfloat16],
-        "f8_dtype": [torch.float8_e4m3fn],
-        "use_sliding_window": [False, True],
-        "use_logits_soft_cap": [False, True],
-        "add_comm": True,
-        "add_gemma": True,
-        "add_oai_oss": True,
-        "add_moe": True,
-        "add_act": True,
-        "add_misc": True,
-        "add_xqa": True,
-    }
-
-
 def detect_sm_capabilities():
     """Detect SM capabilities"""
     compilation_context = CompilationContext()
@@ -914,9 +978,9 @@ def has_sm(compute: str, version: str) -> bool:
     }
 
 
-def register_default_modules() -> int:
+def register_default_modules(profile: Optional[str] = None) -> int:
     """Register the default set of modules"""
-    config = get_default_config()
+    _, config = resolve_build_config(profile=profile)
     sm_capabilities = detect_sm_capabilities()
 
     jit_specs = gen_all_modules(
@@ -942,6 +1006,11 @@ def main():
     parser = argparse.ArgumentParser(
         description="Ahead-of-Time (AOT) build all modules"
     )
+    parser.add_argument(
+        "--profile",
+        help="AOT build profile (full or edge_fm). "
+        "Aliases: dev/fast/minimal -> edge_fm",
+    )
     parser.add_argument("--out-dir", type=Path, help="Output directory")
     parser.add_argument("--build-dir", type=Path, help="Build directory")
     parser.add_argument(
@@ -995,7 +1064,7 @@ def main():
 
     # Start with default configuration
     project_root = Path(__file__).resolve().parents[1]
-    config = get_default_config()
+    build_profile, config = resolve_build_config(profile=args.profile)
     build_dir = jit_env.FLASHINFER_WORKSPACE_DIR
     out_dir: Optional[Path] = None
 
@@ -1038,6 +1107,7 @@ def main():
         build_dir=build_dir,
         project_root=project_root,
         config=config,
+        profile=build_profile,
         verbose=True,
         skip_prebuilt=False,
     )

include/flashinfer/attention/decode.cuh

@@ -270,6 +270,23 @@ __global__ void SingleDecodeWithKVCacheKernel(const __grid_constant__ Params par
   block.sync();
 
   uint32_t chunk_start = kv_chunk_idx * kv_chunk_size;
+  // Early-exit for blocks beyond the actual sequence length.  This happens
+  // when the grid is sized for max_kv_len (CUDA graph) but the real kv_len
+  // is shorter.  Write zero output and very-negative lse so MergeStates
+  // treats this chunk as having no contribution.
+  if (chunk_start >= seq_len) {
+    if (tz == 0) {
+      DTypeO* o_ptr = o + (kv_chunk_idx * num_qo_heads + qo_head_idx) * head_dim + tx * vec_size;
+#pragma unroll
+      for (uint32_t i = 0; i < vec_size; ++i) {
+        o_ptr[i] = DTypeO(0);
+      }
+      if (lse != nullptr && tx == 0) {
+        lse[kv_chunk_idx * num_qo_heads + qo_head_idx] = -5e4;
+      }
+    }
+    return;
+  }
   kv_chunk_size = min(kv_chunk_size, seq_len - chunk_start);
   uint32_t chunk_end = chunk_start + kv_chunk_size;
 
@@ -628,6 +645,15 @@ constexpr uint32_t get_heuristic_num_threads(uint32_t group_size, uint32_t sizeo
     } else {
       return 512U;
     }
+  } else if (group_size == 6U) {
+    // GQA=6 on sm80 otherwise falls back to a 96-thread CTA (bdz=1), which
+    // leaves very little room to hide the decode kernel's memory latency.
+    return 288U;
+  } else if (group_size == 7U) {
+    // GQA=7 has the same problem as GQA=6 on sm80: the default 128-thread CTA
+    // only yields bdz=1 for head_dim=128. Use a larger CTA so group_size=7
+    // can keep multiple z-slices in flight and avoid the slow fallback path.
+    return 336U;
   } else {
     return 128U;
   }
@@ -664,7 +690,9 @@ cudaError_t SingleDecodeWithKVCacheDispatched(Params params, typename Params::DT
   using DTypeO = typename Params::DTypeO;
   const uint32_t num_qo_heads = params.num_qo_heads;
   const uint32_t num_kv_heads = params.num_kv_heads;
-  const uint32_t seq_len = params.kv_len;
+  // When max_kv_len is set (CUDA graph mode), use it for grid sizing so that
+  // the grid topology stays fixed across replays with different actual kv_len.
+  const uint32_t seq_len = (params.max_kv_len > 0) ? params.max_kv_len : params.kv_len;
 
   constexpr uint32_t vec_size = std::max(16UL / sizeof(DTypeKV), HEAD_DIM / 32UL);
   constexpr uint32_t bdx = HEAD_DIM / vec_size;

include/flashinfer/attention/default_decode_params.cuh

@@ -50,6 +50,10 @@ struct SingleDecodeParams {
   float rope_rcp_scale;
   float rope_rcp_theta;
   uint32_t kv_chunk_size;
+  // CUDA graph support: when d_kv_len is set, kernel reads actual kv_len from
+  // this device pointer; max_kv_len is used on host to size the grid.
+  uint32_t* d_kv_len;
+  uint32_t max_kv_len;
 
   __device__ __host__ SingleDecodeParams()
       : q(nullptr),
@@ -70,7 +74,9 @@ struct SingleDecodeParams {
         sm_scale(0.0f),
         rope_rcp_scale(0.0f),
         rope_rcp_theta(0.0f),
-        kv_chunk_size(0) {}
+        kv_chunk_size(0),
+        d_kv_len(nullptr),
+        max_kv_len(0) {}
 
   __device__ __host__ SingleDecodeParams(DTypeQ* q, DTypeKV* k, DTypeKV* v, DTypeO* o,
                                          float* maybe_alibi_slopes, uint32_t seq_len,
@@ -96,12 +102,18 @@ struct SingleDecodeParams {
         sm_scale(sm_scale),
         rope_rcp_scale(1.f / rope_scale),
         rope_rcp_theta(1.f / rope_theta),
-        kv_chunk_size(0) {}
+        kv_chunk_size(0),
+        d_kv_len(nullptr),
+        max_kv_len(0) {}
 
   __host__ __device__ __forceinline__ uint32_t get_qo_len(uint32_t batch_idx) const { return 1; }
 
   __host__ __device__ __forceinline__ uint32_t get_kv_len(uint32_t batch_idx) const {
+#ifdef __CUDA_ARCH__
+    return d_kv_len ? *d_kv_len : kv_len;
+#else
     return kv_len;
+#endif
   }
 };
 

include/flashinfer/attention/variant_helper.cuh

@@ -16,6 +16,7 @@
 #ifndef FLASHINFER_ATTENTION_VARIANT_HELPER_H
 #define FLASHINFER_ATTENTION_VARIANT_HELPER_H
 
+#include <cuda_fp16.h>
 #include <cuda_runtime.h>
 
 #include <cstdint>
@@ -26,6 +27,17 @@ namespace flashinfer {
 
 DEFINE_HAS_MEMBER(v_scale)
 
+namespace attention_variant_detail {
+
+template <typename T>
+__device__ __forceinline__ float m_to_float(T value) {
+  return static_cast<float>(value);
+}
+
+__device__ __forceinline__ float m_to_float(half value) { return __half2float(value); }
+
+}  // namespace attention_variant_detail
+
 #define REGISTER_QUERY_TRANSFORM(params, q, ...)                                    \
   template <typename Params, typename T>                                            \
   __device__ __forceinline__ T QueryTransform(const Params& params, void* q_smem) { \
@@ -83,7 +95,8 @@ struct AttentionVariantBase {
   REGISTER_M_D_UPDATE(params, kv_tile_idx, qo_head_idx, m, d, scale, { return; })
 
   REGISTER_OUTPUT_TRANSFORM(params, output, batch_idx, qo_idx, qo_head_idx, m, d, scale, {
-    float d_rcp = (m != -math::inf) ? math::ptx_rcp(d) : 0.f;
+    const float m_value = attention_variant_detail::m_to_float(m);
+    float d_rcp = (m_value != -math::inf) ? math::ptx_rcp(d) : 0.f;
     float v_scale_val = get_v_scale(params);
     return output * d_rcp * v_scale_val;
   })