[CK Tile][FMHA] Decouple StreamingLLM and GPT-OSS sink into independent compile-time modes

projects/composablekernel/example/ck_tile/01_fmha/CMakeLists.txt

@@ -43,13 +43,23 @@ set_directory_properties(PROPERTIES CMAKE_CONFIGURE_DEPENDS "${CODE_GEN_SCRIPTS}
 list(JOIN INST_TARGETS , FMHA_TARGETS_ARG)
 
 string(REPLACE ";" "," FMHA_FWD_APIS "${FMHA_FWD_ENABLE_APIS}")
+set(FMHA_FWD_OPTDIM "32,64,80,128,256" CACHE STRING
+  "comma-separated list of hdim values to optimize for")
+set(FMHA_FWD_FILTER "" CACHE STRING
+  "fnmatch filter pattern for fwd kernel instances. Empty = no filter.")
+set(FMHA_FWD_SINK_MODES "none" CACHE STRING
+  "comma-separated list of sink modes to generate (none,stream,gptoss,both). Default: none.")
+
 set(FMHA_FWD_CODE_GEN_COMMON_ARGS
   ${CMAKE_CURRENT_LIST_DIR}/generate.py
   --targets ${FMHA_TARGETS_ARG}
   --api ${FMHA_FWD_APIS}
-  --optdim 32,64,80,128,256
-  # --filter fmha_fwd...
+  --optdim ${FMHA_FWD_OPTDIM}
+  --sink ${FMHA_FWD_SINK_MODES}
 )
+if(FMHA_FWD_FILTER)
+  list(APPEND FMHA_FWD_CODE_GEN_COMMON_ARGS --filter ${FMHA_FWD_FILTER})
+endif()
 set(FMHA_BWD_CODE_GEN_COMMON_ARGS
   ${CMAKE_CURRENT_LIST_DIR}/generate.py
   --targets ${FMHA_TARGETS_ARG}
@@ -98,6 +108,7 @@ add_custom_command(
   --output_dir ${CMAKE_CURRENT_BINARY_DIR}
   DEPENDS ${CODE_GEN_SCRIPTS}
   COMMENT "Generate CK Tile FMHA FWD kernels"
+  VERBATIM
 )
 
 add_custom_command(

projects/composablekernel/example/ck_tile/01_fmha/codegen/ops/fmha_batch_prefill.py

@@ -828,7 +828,11 @@ def write_blobs(
     receipt,
     optdim_list,
     mask_impl,
+    sink_modes=("none",),
 ) -> None:
+    # sink_modes is intentionally not forwarded: batch_prefill does not yet support
+    # StreamLLM/GPT-OSS sink kernels. The parameter exists only for API uniformity
+    # with other fwd handlers. Non-"none" modes are silently treated as "none".
     api_pool, kernels = get_fwd_blobs(kernel_filter, receipt, optdim_list, mask_impl)
     for kernel in kernels:
         write_single_fwd_kernel(kernel, output_dir)
@@ -842,7 +846,9 @@ def list_blobs(
     receipt,
     optdim_list,
     mask_impl,
+    sink_modes=("none",),
 ) -> None:
+    # sink_modes is intentionally not forwarded: see write_blobs for rationale.
     with file_path.open("a") as f:
         _, kernels = get_fwd_blobs(kernel_filter, receipt, optdim_list, mask_impl)
         for kernel in kernels:

projects/composablekernel/example/ck_tile/01_fmha/codegen/ops/fmha_fwd_appendkv.py

@@ -493,7 +493,11 @@ def write_blobs(
     receipt,
     optdim_list,
     mask_impl,
+    sink_modes=("none",),
 ) -> None:
+    # sink_modes is intentionally not forwarded: appendkv does not support sink kernels.
+    # The parameter exists only for API uniformity with other fwd handlers.
+    # Non-"none" modes are silently treated as "none".
     api_pool, kernels = get_fwd_appendkv_blobs(
         targets, kernel_filter, receipt, mask_impl, optdim_list
     )
@@ -509,11 +513,16 @@ def list_blobs(
     receipt,
     optdim_list,
     mask_impl,
+    sink_modes=("none",),
 ) -> None:
+    # sink_modes is intentionally not forwarded: see write_blobs for rationale.
     with file_path.open("a") as f:
         _, kernels = get_fwd_appendkv_blobs(
             targets, kernel_filter, receipt, mask_impl, optdim_list
         )
         for kernel in kernels:
             f.write((file_path.parent / GEN_DIR / kernel.filename).as_posix() + "\n")
-        f.write((file_path.parent / GEN_DIR / FMHA_FWD_APPENDKV_API_FILENAME).as_posix() + "\n")
+        f.write(
+            (file_path.parent / GEN_DIR / FMHA_FWD_APPENDKV_API_FILENAME).as_posix()
+            + "\n"
+        )

projects/composablekernel/example/ck_tile/01_fmha/codegen/ops/fmha_fwd_splitkv.py

@@ -32,6 +32,9 @@
     FMHA_FWD_API_PER_ARCH,
     FMHA_FWD_API_PER_DTYPE,
     FMHA_FWD_API_PER_HDIM_CASE,
+    SINK_MODE_MAP,
+    SINK_MODE_DISPATCH_MAP,
+    SINK_NAME_MAP,
 )
 
 
@@ -74,7 +77,7 @@
                                                      kHasUnevenSplits,
                                                      kMergeNumHeadGroupsSeqLenQ,
                                                      {F_occupancy},
-                                                     {F_sink}>;
+                                                     {F_sink_mode}>;
 
 using fmha_pipeline_problem = ck_tile::BlockFmhaFwdSplitKVPipelineProblem<
     typename FmhaFwdTypeConfig<fmha_dtype_{F_idx}>::QDataType,
@@ -118,7 +121,7 @@
 }} // anonymous namespace
 
 using trait_{F_idx} = fmha_fwd_splitkv_traits_<{F_hdim}, {F_dtype}, {F_mode}, {F_bm0}, {F_bn0}, {F_bk0}, {F_bn1}, {F_bk1}, {F_bk0max}, {F_vlayout},
-                        {F_pipeline_enum}, {F_logits}, fmha_mask_{F_idx}, {F_bias}, {F_lse}, {F_squant}, {F_pagedkv}, {F_sink}, {F_spad}, {F_skpad}, {F_dpad},
+                        {F_pipeline_enum}, {F_logits}, fmha_mask_{F_idx}, {F_bias}, {F_lse}, {F_squant}, {F_pagedkv}, {F_sink_mode}, {F_spad}, {F_skpad}, {F_dpad},
                         {F_dvpad}>;
 
 #pragma clang diagnostic push
@@ -280,8 +283,8 @@
 """
 
 FMHA_FWD_SPLITKV_API_INNER_DISPATCH = """{F_if}((t.is_group_mode == {F_mode}) && (t.is_v_rowmajor == {F_vlayout}) && (t.has_logits_soft_cap == {F_logits}) && ({F_mask_check}) && (t.bias_type == {F_bias_check}) && (t.do_fp8_static_quant == {F_squant}) &&
-        ((a.block_table_ptr != nullptr) == {F_pagedkv}) && (t.has_sink == {F_sink}) && ({F_scheck}) && ({F_skcheck}) && ({F_dcheck}) && ({F_dvcheck})) {{
-    using traits_ = fmha_fwd_splitkv_traits_<{F_hdim}, {F_dtype}, {F_mode}, {F_bm0}, {F_bn0}, {F_bk0}, {F_bn1}, {F_bk1}, {F_bk0max}, {F_vlayout}, {F_pipeline_enum}, {F_logits}, {F_mask}, {F_bias}, true, {F_squant}, {F_pagedkv},{F_sink}, {F_spad}, {F_skpad}, {F_dpad}, {F_dvpad}>;
+        ((a.block_table_ptr != nullptr) == {F_pagedkv}) && (t.has_stream_sink == {F_stream_sink}) && (t.has_gptoss_sink == {F_gptoss_sink}) && ({F_scheck}) && ({F_skcheck}) && ({F_dcheck}) && ({F_dvcheck})) {{
+    using traits_ = fmha_fwd_splitkv_traits_<{F_hdim}, {F_dtype}, {F_mode}, {F_bm0}, {F_bn0}, {F_bk0}, {F_bn1}, {F_bk1}, {F_bk0max}, {F_vlayout}, {F_pipeline_enum}, {F_logits}, {F_mask}, {F_bias}, true, {F_squant}, {F_pagedkv},{F_sink_mode}, {F_spad}, {F_skpad}, {F_dpad}, {F_dvpad}>;
 
     // get combine kernel tile sizes
     using OaccDataType = typename FmhaFwdTypeConfig<{F_dtype}>::OaccDataType;
@@ -427,7 +430,7 @@ class FmhaFwdSplitKVPipeline:
     F_lse: str  #
     F_squant: str  #
     F_pagedkv: str  # t/f
-    F_sink: str  # t/f
+    F_sink: str  # "none" / "stream" / "gptoss" / "both"
     F_mask: str  # value from MASK_MAP
 
     @property
@@ -488,10 +491,7 @@ def pad_name() -> str:
             n += "_pagedkv"
         else:
             n += "_npagedkv"
-        if self.F_sink == "t":
-            n += "_sink"
-        else:
-            n += "_nsink"
+        n += SINK_NAME_MAP[self.F_sink]
         return n
 
 
@@ -574,7 +574,9 @@ def api(self) -> str:
                             F_lse=BOOL_MAP[trait.lse],
                             F_squant=BOOL_MAP[trait.squant],
                             F_pagedkv=BOOL_MAP[trait.pagedkv],
-                            F_sink=BOOL_MAP[trait.sink],
+                            F_sink_mode=SINK_MODE_MAP[trait.sink],
+                            F_stream_sink=SINK_MODE_DISPATCH_MAP[trait.sink][0],
+                            F_gptoss_sink=SINK_MODE_DISPATCH_MAP[trait.sink][1],
                             F_scheck=trait.scheck,
                             F_skcheck=trait.skcheck,
                             F_dcheck=trait.dcheck,
@@ -675,7 +677,7 @@ def template(self) -> str:
             F_squant=BOOL_MAP[self.F_pipeline.F_squant],
             F_pagedkv=BOOL_MAP[self.F_pipeline.F_pagedkv],
             F_occupancy=self.F_tile.F_occupancy,
-            F_sink=BOOL_MAP[self.F_pipeline.F_sink],
+            F_sink_mode=SINK_MODE_MAP[self.F_pipeline.F_sink],
             F_pipeline_enum=PIPELINE_ENUM_MAP[self.F_pipeline.tag],
             F_mask=get_mask_map(self.mask_impl)[self.F_pipeline.F_mask],
             F_mode=MODE_MAP[self.F_mode],
@@ -740,7 +742,9 @@ def filename(self) -> str:
 
 class KernelComponentFactoryBase:
     @staticmethod
-    def get_pipelines(dtype, hdim, mask_impl) -> List[FmhaFwdSplitKVPipeline]:
+    def get_pipelines(
+        dtype, hdim, mask_impl, sink_modes=("none",)
+    ) -> List[FmhaFwdSplitKVPipeline]:
         # this function will populate a list possible pipelines
         # TODO: the order of List matters! the later in this list will be also be checked later
         # TODO: currently for qr pipeline, let "t" padding to appear later!!
@@ -754,7 +758,7 @@ def get_pipelines(dtype, hdim, mask_impl) -> List[FmhaFwdSplitKVPipeline]:
                 get_mask_map(mask_impl).keys(),
                 BIAS_MAP.keys(),
                 ["t", "f"],
-                ["t", "f"],
+                sink_modes,
             ):
                 pipelines.append(Pipeline("qr", "row", "f", "t", "f", "f", logits, bias, "t", squant, pagedkv, sink, mask))  # fmt: skip
                 pipelines.append(Pipeline("qr", "row", "t", "f", "f", "f", logits, bias, "t", squant, pagedkv, sink, mask))  # fmt: skip
@@ -764,8 +768,8 @@ def get_pipelines(dtype, hdim, mask_impl) -> List[FmhaFwdSplitKVPipeline]:
             for logits, mask, bias in itertools.product(
                 ["t", "f"], get_mask_map(mask_impl).keys(), BIAS_MAP.keys()
             ):
-                pipelines.append(Pipeline("qr", "row", "f", "f", "f", "f", logits, bias, "t", squant, "f", "f", mask))  # fmt: skip
-                pipelines.append(Pipeline("qr", "row", "t", "t", "f", "f", logits, bias, "t", squant, "f", "f", mask))  # fmt: skip
+                pipelines.append(Pipeline("qr", "row", "f", "f", "f", "f", logits, bias, "t", squant, "f", "none", mask))  # fmt: skip
+                pipelines.append(Pipeline("qr", "row", "t", "t", "f", "f", logits, bias, "t", squant, "f", "none", mask))  # fmt: skip
         elif dtype in ["fp8fp16", "fp8bf16"]:
             # TODO
             None
@@ -891,7 +895,12 @@ def get_factory(target: str):
 
 
 def get_fwd_splitkv_blobs(
-    targets: List[str], kernel_filter: Optional[str], receipt, mask_impl, optdim_list
+    targets: List[str],
+    kernel_filter: Optional[str],
+    receipt,
+    mask_impl,
+    optdim_list,
+    sink_modes=("none",),
 ) -> List[FmhaFwdSplitKVKernel]:
     Kernel = FmhaFwdSplitKVKernel
 
@@ -907,7 +916,7 @@ def get_fwd_splitkv_blobs(
         for hdim_str, mode in itertools.product(d.keys(), MODE_MAP.keys()):
             tile = d[hdim_str]
             hdim = int(hdim_str)
-            for pipeline in factory.get_pipelines(dtype, hdim, mask_impl):
+            for pipeline in factory.get_pipelines(dtype, hdim, mask_impl, sink_modes):
                 if mode == "group":
                     if pipeline.F_spad != "t" or pipeline.F_skpad != "t":
                         # in group mode, spad/skpad must be true, since we can't predict if seqlen of current batch need pad or not
@@ -942,7 +951,7 @@ def get_fwd_splitkv_blobs(
                     # FlashAttention splitkv paths use softcap-disabled kernels only.
                     cond &= pipeline.F_logits == "f"
                     cond &= pipeline.F_squant == "f"
-                    cond &= pipeline.F_sink == "f"
+                    cond &= pipeline.F_sink == "none"
                     if not cond:
                         continue
                 # PyTorch integration
@@ -952,7 +961,7 @@ def get_fwd_splitkv_blobs(
                     cond &= pipeline.F_bias in ["no", "bias"]
                     cond &= pipeline.F_squant == "f"
                     cond &= mode == "batch"
-                    cond &= pipeline.F_sink == "f"
+                    cond &= pipeline.F_sink == "none"
                     if not cond:
                         continue
                 # Aiter(mha_varlen_fwd) integration
@@ -1058,6 +1067,7 @@ def write_blobs(
     receipt,
     optdim_list,
     mask_impl,
+    sink_modes=("none",),
 ) -> None:
     filter_list = filter_list.split("@")
     filter_list.extend([""] * (2 - len(filter_list)))
@@ -1068,7 +1078,7 @@ def write_blobs(
     for kernel in combine_kernels:
         write_single_kernel(kernel, output_dir)
     kernels = get_fwd_splitkv_blobs(
-        targets, filter_list[1], receipt, mask_impl, optdim_list
+        targets, filter_list[1], receipt, mask_impl, optdim_list, sink_modes
     )
     for kernel in kernels:
         write_single_kernel(kernel, output_dir)
@@ -1129,6 +1139,7 @@ def list_blobs(
     receipt,
     optdim_list,
     mask_impl,
+    sink_modes=("none",),
 ) -> None:
     filter_list = filter_list.split("@")
     filter_list.extend([""] * (2 - len(filter_list)))
@@ -1140,7 +1151,7 @@ def list_blobs(
         for kernel in kernels:
             f.write((file_path.parent / GEN_DIR / kernel.filename).as_posix() + "\n")
         kernels = get_fwd_splitkv_blobs(
-            targets, filter_list[1], receipt, mask_impl, optdim_list
+            targets, filter_list[1], receipt, mask_impl, optdim_list, sink_modes
         )
         for kernel in kernels:
             f.write((file_path.parent / GEN_DIR / kernel.filename).as_posix() + "\n")