Star862
1359
代码
代码
Issues1721
Pull Requests2239
项目讨论
Wiki
分析
Star862
1359
已合并
flexattention: port flexattention template to 2.13 #44356
stonexxx创建于 8月11日
flexattention: port flexattention template to 2.13 #44356
已合并
stonexxx创建于 8月11日
从stonexxx/pytorchfa_213_origin合入到Ascend/pytorchmaster
共 12 个文件变更+7995-1543
Mtorch_npu/_inductor/__init__.py+23-0
@@ -112,6 +112,7 @@ def _load_triton_backend():
112 _register_npu_inductor_grouped_mm,112 _register_npu_inductor_grouped_mm,
113 _register_npu_inductor_mm,113 _register_npu_inductor_mm,
114 _validate_device,114 _validate_device,
115+ patch_flex_attention,
115 )116 )
116 from .lowering import make_reduction117 from .lowering import make_reduction
117 from .runtime import (118 from .runtime import (
@@ -125,6 +126,28 @@ def _load_triton_backend():
125 from .graph import patch_count_bytes126 from .graph import patch_count_bytes
126 from .autotune_process import patch_tuning_process127 from .autotune_process import patch_tuning_process
127 flex_attention._validate_device = _validate_device128 flex_attention._validate_device = _validate_device
129+ patch_flex_attention()
130+ 
131+ def _patch_flex_attention_singleton_sort():
132+ original = getattr(flex_attention, "_dense_to_ordered", None)
133+ if original is None or getattr(original, "_torch_npu_singleton_sort_patch", False):
134+ return
135+ 
136+ def _dense_to_ordered_npu_safe(dense_mask):
137+ if dense_mask.ndim > 0 and dense_mask.size(-1) == 1:
138+ dense_mask = dense_mask.to(dtype=torch.int32)
139+ num_blocks_in_row = dense_mask.sum(dim=-1)
140+ col_indices = torch.zeros_like(dense_mask, dtype=torch.int32)
141+ return (
142+ num_blocks_in_row.to(torch.int32, memory_format=torch.contiguous_format),
143+ col_indices.to(torch.int32, memory_format=torch.contiguous_format),
144+ )
145+ return original(dense_mask)
146+ 
147+ _dense_to_ordered_npu_safe._torch_npu_singleton_sort_patch = True
148+ flex_attention._dense_to_ordered = _dense_to_ordered_npu_safe
149+ 
150+ _patch_flex_attention_singleton_sort()
128 151 
129 def _inductor_register_backend_for_device():152 def _inductor_register_backend_for_device():
130 from .codegen.cpp_wrapper_npu import CppWrapperNpu153 from .codegen.cpp_wrapper_npu import CppWrapperNpu
Mtorch_npu/_inductor/codegen/scheduling.py+168-19
@@ -38,6 +38,7 @@ from ..lowering_fx import (
38)38)
39from ..fx_passes.utils.schedule_node_utils import is_multi_stream39from ..fx_passes.utils.schedule_node_utils import is_multi_stream
40from ..config import log40from ..config import log
41+from ..select_algorithm import NPUFlexAttentionDkdvTemplateBuffer
41 42 
42 43 
43def flatten_groups(nums):44def flatten_groups(nums):
@@ -50,6 +51,7 @@ def flatten_groups(nums):
50 res.append(i)51 res.append(i)
51 return res52 return res
52 53 
54+ 
53class NPUNoLinearTritonScheduling(TritonScheduling):55class NPUNoLinearTritonScheduling(TritonScheduling):
54 kernel_type = NPUTritonKernel56 kernel_type = NPUTritonKernel
55 57 
@@ -236,6 +238,14 @@ class NPUTritonScheduling(TritonScheduling):
236 self.scheduler.free_buffers()238 self.scheduler.free_buffers()
237 239 
238 def codegen_template(self, template_node, epilogue_nodes, only_gen_src_code=False):240 def codegen_template(self, template_node, epilogue_nodes, only_gen_src_code=False):
241+ if isinstance(template_node.node, NPUFlexAttentionDkdvTemplateBuffer):
242+ if only_gen_src_code:
243+ raise NotImplementedError(
244+ "source-only codegen is unsupported for composite dK/dV templates"
245+ )
246+ return self.codegen_flex_attention_dkdv_template(
247+ template_node, epilogue_nodes
248+ )
239 _, (numel, rnumel) = template_node.group249 _, (numel, rnumel) = template_node.group
240 assert rnumel == 1250 assert rnumel == 1
241 kernel, render = template_node.node.make_kernel_render(template_node.node)251 kernel, render = template_node.node.make_kernel_render(template_node.node)
@@ -244,21 +254,41 @@ class NPUTritonScheduling(TritonScheduling):
244 for node in [template_node, *epilogue_nodes]:254 for node in [template_node, *epilogue_nodes]:
245 node.mark_run()255 node.mark_run()
246 partial_code = render()256 partial_code = render()
247- with kernel.set_subgraph_body("<STORE_OUTPUT>"):257+ # Handle both legacy "<STORE_OUTPUT>" (without index) and torch
248- for node in epilogue_nodes:258+ # 2.10's "<STORE_OUTPUT_N>" (with index) subgraph body naming.
249- node.codegen(kernel.split_and_set_ranges(node.get_ranges()))259+ if "<STORE_OUTPUT>" in kernel.subgraph_bodies:
260+ with kernel.set_subgraph_body("<STORE_OUTPUT>"):
261+ for node in epilogue_nodes:
262+ node.codegen(kernel.split_and_set_ranges(node.get_ranges()))
263+ # torch 2.10+ creates indexed store_output subgraphs.
264+ if hasattr(kernel, "get_store_output_count"):
265+ num_store_subgraphs = kernel.get_store_output_count()
266+ for i in range(num_store_subgraphs):
267+ subgraph_name = kernel._get_store_output_subgraph_name(i)
268+ if subgraph_name in kernel.subgraph_bodies:
269+ with kernel.set_subgraph_body(subgraph_name):
270+ for node in epilogue_nodes:
271+ node.codegen(kernel.split_and_set_ranges(node.get_ranges()))
250 272 
251 if not isinstance(partial_code, str):273 if not isinstance(partial_code, str):
252 partial_code.finalize_hook("<DEF_KERNEL>")274 partial_code.finalize_hook("<DEF_KERNEL>")
253 partial_code.finalize_hook("<ARGDEFS>", strict=False)275 partial_code.finalize_hook("<ARGDEFS>", strict=False)
254 # finalize must be called after adding epilogue above276 # finalize must be called after adding epilogue above
255 with V.set_kernel_handler(kernel):277 with V.set_kernel_handler(kernel):
256- with kernel.set_subgraph_body("<STORE_OUTPUT>"):278+ # Finalize legacy "<STORE_OUTPUT>" (without index) if present.
257- if isinstance(partial_code, str):279+ has_store_output = "<STORE_OUTPUT>" in kernel.subgraph_bodies
258- src_code = partial_code280+ if has_store_output and not isinstance(partial_code, str):
259- else:281+ with kernel.set_subgraph_body("<STORE_OUTPUT>"):
260 partial_code.finalize_hook("<STORE_OUTPUT>")282 partial_code.finalize_hook("<STORE_OUTPUT>")
261- src_code = partial_code.code283+ # Now safe to extract the final source code.
284+ # Use finalize_remaining() to catch any leftover hooks (e.g.
285+ # torch 2.10's indexed "<STORE_OUTPUT_N>" hooks).
286+ if isinstance(partial_code, str):
287+ src_code = partial_code
288+ elif hasattr(partial_code, "finalize_remaining"):
289+ src_code = partial_code.finalize_remaining()
290+ else:
291+ src_code = partial_code.code
262 node_schedule = [template_node, *epilogue_nodes]292 node_schedule = [template_node, *epilogue_nodes]
263 293 
264 if config.benchmark_kernel:294 if config.benchmark_kernel:
@@ -286,6 +316,113 @@ class NPUTritonScheduling(TritonScheduling):
286 self.scheduler.free_buffers()316 self.scheduler.free_buffers()
287 return None317 return None
288 318 
319+ def codegen_flex_attention_dkdv_template(
320+ self, template_node, epilogue_nodes
321+ ):
322+ if epilogue_nodes:
323+ raise NotImplementedError(
324+ "epilogue fusion is unsupported for composite dK/dV templates"
325+ )
326+ composite = template_node.node
327+ 
328+ def render_source(kernel, render):
329+ with kernel:
330+ partial_code = render()
331+ if not isinstance(partial_code, str):
332+ partial_code.finalize_hook("<DEF_KERNEL>")
333+ partial_code.finalize_hook("<ARGDEFS>", strict=False)
334+ with V.set_kernel_handler(kernel):
335+ if not isinstance(partial_code, str):
336+ if "<STORE_OUTPUT>" in kernel.subgraph_bodies:
337+ partial_code.finalize_hook("<STORE_OUTPUT>")
338+ return partial_code.code
339+ return partial_code
340+ 
341+ legacy_kernel, legacy_render = composite.make_kernel_render(composite)
342+ legacy_source = render_source(legacy_kernel, legacy_render)
343+ legacy_kernel_name, _ = self.define_kernel(
344+ legacy_source, [template_node], legacy_kernel, None
345+ )
346+ 
347+ runtime_renderers = composite.runtime_renderer_factory(composite)
348+ tasklist_renderer = runtime_renderers["tasklist"]
349+ with tasklist_renderer.patch_runtime_args():
350+ tasklist_source = render_source(
351+ tasklist_renderer.kernel, tasklist_renderer.render
352+ )
353+ tasklist_kernel_name, _ = self.define_kernel(
354+ tasklist_source, [template_node], tasklist_renderer.kernel, None
355+ )
356+ 
357+ tasklist_no_split_renderer = runtime_renderers["tasklist_no_split"]
358+ with tasklist_no_split_renderer.patch_runtime_args():
359+ tasklist_no_split_source = render_source(
360+ tasklist_no_split_renderer.kernel,
361+ tasklist_no_split_renderer.render,
362+ )
363+ tasklist_no_split_kernel_name, _ = self.define_kernel(
364+ tasklist_no_split_source,
365+ [template_node],
366+ tasklist_no_split_renderer.kernel,
367+ None,
368+ )
369+ 
370+ reduce_renderer = runtime_renderers["reduce"]
371+ with reduce_renderer.patch_runtime_args():
372+ reduce_source = render_source(
373+ reduce_renderer.kernel, reduce_renderer.render
374+ )
375+ reduce_kernel_name, _ = self.define_kernel(
376+ reduce_source, [template_node], reduce_renderer.kernel, None
377+ )
378+ 
379+ template_node.mark_run()
380+ wrapper = V.graph.wrapper_code
381+ wrapper.write_triton_header_once()
382+ _, legacy_call_args, _, _ = legacy_kernel.args.python_argdefs()
383+ _, tasklist_call_args, _, _ = tasklist_renderer.python_argdefs()
384+ _, tasklist_no_split_call_args, _, _ = (
385+ tasklist_no_split_renderer.python_argdefs()
386+ )
387+ _, reduce_call_args, _, _ = reduce_renderer.python_argdefs()
388+ runtime_arg_names = {
389+ **tasklist_renderer.runtime_arg_names,
390+ **tasklist_no_split_renderer.runtime_arg_names,
391+ **reduce_renderer.runtime_arg_names,
392+ }
393+ wrapper.generate_flex_attention_dkdv_dispatch(
394+ dispatch_spec=composite.dispatch_spec,
395+ legacy_kernel_name=legacy_kernel_name,
396+ legacy_call_args=legacy_call_args,
397+ tasklist_kernel_name=tasklist_kernel_name,
398+ tasklist_call_args=tasklist_call_args,
399+ tasklist_no_split_kernel_name=tasklist_no_split_kernel_name,
400+ tasklist_no_split_call_args=tasklist_no_split_call_args,
401+ reduce_kernel_name=reduce_kernel_name,
402+ reduce_call_args=reduce_call_args,
403+ q_num_blocks_name=legacy_kernel.named_input_nodes[
404+ "Q_NUM_BLKS"
405+ ].get_name(),
406+ full_q_num_blocks_name=legacy_kernel.named_input_nodes[
407+ "FULL_Q_NUM_BLKS"
408+ ].get_name(),
409+ dk_name=legacy_kernel.named_input_nodes["DK"].get_name(),
410+ dv_name=legacy_kernel.named_input_nodes["DV"].get_name(),
411+ runtime_arg_names=runtime_arg_names,
412+ )
413+ 
414+ for kernel in (
415+ legacy_kernel,
416+ tasklist_renderer.kernel,
417+ tasklist_no_split_renderer.kernel,
418+ reduce_renderer.kernel,
419+ ):
420+ V.graph.removed_buffers |= kernel.removed_buffers
421+ V.graph.inplaced_to_remove |= kernel.inplaced_to_remove
422+ self.codegen_comment([template_node])
423+ self.scheduler.free_buffers()
424+ return None
425+ 
289 def codegen_combo_kernel(self, combo_kernel_node):426 def codegen_combo_kernel(self, combo_kernel_node):
290 subkernel_nodes = combo_kernel_node.get_subkernel_nodes()427 subkernel_nodes = combo_kernel_node.get_subkernel_nodes()
291 custom_part_algorithm = combo_kernel_node.use_custom_partition_algo428 custom_part_algorithm = combo_kernel_node.use_custom_partition_algo
@@ -382,17 +519,21 @@ class NPUTritonScheduling(TritonScheduling):
382 src_code = src_code.replace('TRACED_GRAPH_HASH', traced_graph_hash)519 src_code = src_code.replace('TRACED_GRAPH_HASH', traced_graph_hash)
383 src_code = src_code.replace('TRACED_GRAPH_DIR', npu_config.traced_fx_graph_cache)520 src_code = src_code.replace('TRACED_GRAPH_DIR', npu_config.traced_fx_graph_cache)
384 else:521 else:
385- fused_name = (522+ kernel_name = getattr(kernel, "_npu_codegen_kernel_name", None)
386- get_fused_kernel_name(node_schedule, config.triton.descriptive_names)523+ if kernel_name is None:
387- if config.triton.descriptive_names524+ fused_name = (
388- else ""525+ get_fused_kernel_name(node_schedule, config.triton.descriptive_names)
389- )526+ if config.triton.descriptive_names
390- if len(fused_name) > 35:527+ else ""
391- fused_name = fused_name[0:35]528+ )
392- kernel_category = get_kernel_category_by_source_code(src_code)[:3]529+ if len(fused_name) > 35:
393- kernel_name = "_".join(530+ fused_name = fused_name[0:35]
394- ["triton", kernel_category, fused_name, wrapper.next_kernel_suffix()]531+ kernel_category = get_kernel_category_by_source_code(src_code)[:3]
395- )532+ kernel_name = "_".join(
533+ ["triton", kernel_category, fused_name, wrapper.next_kernel_suffix()]
534+ )
535+ else:
536+ wrapper.next_kernel_suffix()
396 # use the original src_code as the key537 # use the original src_code as the key
397 wrapper.src_to_kernel[kernel_cache_key] = kernel_name538 wrapper.src_to_kernel[kernel_cache_key] = kernel_name
398 subs_name = kernel_name if config.triton.unique_kernel_names else "triton_"539 subs_name = kernel_name if config.triton.unique_kernel_names else "triton_"
@@ -586,6 +727,14 @@ class NPUTritonScheduling(TritonScheduling):
586 _, (numel2, rnumel2) = node2.group727 _, (numel2, rnumel2) = node2.group
587 why = WhyNoFuse(node1, node2)728 why = WhyNoFuse(node1, node2)
588 729 
730+ for node in (node1, node2):
731+ if node.is_template() and isinstance(
732+ node.get_template_node(),
733+ NPUFlexAttentionDkdvTemplateBuffer,
734+ ):
735+ why("composite dK/dV templates do not support fusion")
736+ return False
737+ 
589 if node1.is_split_scan() and not node2.is_split_scan():738 if node1.is_split_scan() and not node2.is_split_scan():
590 if node2.is_reduction():739 if node2.is_reduction():
591 why("Split scan cannot fuse with reductions")740 why("Split scan cannot fuse with reductions")
Mtorch_npu/_inductor/codegen/wrapper.py+170-1
@@ -15,7 +15,12 @@ from torch._inductor.utils import (
15from torch._inductor.virtualized import V15from torch._inductor.virtualized import V
16from torch._inductor.ir import GraphPartitionSignature, TorchBindObject, NoneLayout16from torch._inductor.ir import GraphPartitionSignature, TorchBindObject, NoneLayout
17from torch._dynamo.utils import counters17from torch._dynamo.utils import counters
18-from torch._inductor.codegen.common import DeferredLine, WorkspaceArg, IndentedBuffer18+from torch._inductor.codegen.common import (
19+ DeferredLine,
20+ DeferredLineBase,
21+ WorkspaceArg,
22+ IndentedBuffer,
23+)
19from torch._inductor.codegen.wrapper import BufferLike, WrapperLine24from torch._inductor.codegen.wrapper import BufferLike, WrapperLine
20from torch._inductor import ir25from torch._inductor import ir
21import torch_npu.npu.aclnn26import torch_npu.npu.aclnn
@@ -25,6 +30,26 @@ from torch_npu._inductor._aclgraph_update_plan import (
25 emit_inductor_aclgraph_update_plan_for_wrapper,30 emit_inductor_aclgraph_update_plan_for_wrapper,
26)31)
27from torch_npu._inductor.utils import resolve_npu_device_index32from torch_npu._inductor.utils import resolve_npu_device_index
33+from ..flex_attention_tasklist import DKDV_TASKLIST_HELPER_SOURCE
34+ 
35+ 
36+class _RuntimeHelperDefinitionsLine(DeferredLineBase):
37+ def __init__(self, definitions: IndentedBuffer, prefix: str = ""):
38+ super().__init__(prefix)
39+ self.definitions = definitions
40+ 
41+ def __call__(self):
42+ source = self.definitions.getvalue()
43+ if not source:
44+ return None
45+ if not self.line:
46+ return source
47+ return "\n".join(
48+ f"{self.line}{line}" if line else line for line in source.splitlines()
49+ )
50+ 
51+ def _new_line(self, line):
52+ return _RuntimeHelperDefinitionsLine(self.definitions, line)
28 53 
29 54 
30def _is_codegen_graph_partition_subgraph(wrapper) -> bool:55def _is_codegen_graph_partition_subgraph(wrapper) -> bool:
@@ -173,12 +198,156 @@ class NPUMultiOutputLine(MultiOutputLine):
173 198 
174class NPUPythonWrapperCodeGen(_NPUKernelCodegenMixin, PythonWrapperCodegen):199class NPUPythonWrapperCodeGen(_NPUKernelCodegenMixin, PythonWrapperCodegen):
175 def __init__(self):200 def __init__(self):
201+ self.runtime_helper_definitions = IndentedBuffer()
202+ self._flex_attention_dkdv_tasklist_helpers_written = False
176 super().__init__()203 super().__init__()
177 self.buffer_args_multi_stream_intent = {}204 self.buffer_args_multi_stream_intent = {}
178 self.buffer_define_multi_stream = {}205 self.buffer_define_multi_stream = {}
179 self.extern_node_intent_multi_stream = []206 self.extern_node_intent_multi_stream = []
180 self.pre_define_buffer = []207 self.pre_define_buffer = []
181 208 
209+ def write_async_compile_wait(self) -> None:
210+ super().write_async_compile_wait()
211+ self.prefix.writeline(
212+ _RuntimeHelperDefinitionsLine(self.runtime_helper_definitions)
213+ )
214+ 
215+ def write_flex_attention_dkdv_tasklist_helpers_once(self) -> None:
216+ if self._flex_attention_dkdv_tasklist_helpers_written:
217+ return
218+ self.runtime_helper_definitions.splice(DKDV_TASKLIST_HELPER_SOURCE)
219+ self.runtime_helper_definitions.writeline("")
220+ self._flex_attention_dkdv_tasklist_helpers_written = True
221+ 
222+ def generate_flex_attention_dkdv_dispatch(
223+ self,
224+ *,
225+ dispatch_spec,
226+ legacy_kernel_name,
227+ legacy_call_args,
228+ tasklist_kernel_name,
229+ tasklist_call_args,
230+ tasklist_no_split_kernel_name,
231+ tasklist_no_split_call_args,
232+ reduce_kernel_name,
233+ reduce_call_args,
234+ q_num_blocks_name,
235+ full_q_num_blocks_name,
236+ dk_name,
237+ dv_name,
238+ runtime_arg_names,
239+ ) -> None:
240+ self.write_flex_attention_dkdv_tasklist_helpers_once()
241+ suffix = self.next_kernel_suffix()
242+ prefix = f"dkdv_tasklist_{suffix}"
243+ use_tasklist = f"{prefix}_use_tasklist"
244+ max_sub = f"{prefix}_max_sub"
245+ work_items_tensor = f"{prefix}_work_items_t"
246+ task_offsets_tensor = f"{prefix}_task_offsets_t"
247+ split_bases_tensor = f"{prefix}_split_bases_t"
248+ dk_partial = f"{prefix}_dk_partial"
249+ dv_partial = f"{prefix}_dv_partial"
250+ num_split_bases = f"{prefix}_num_split_bases"
251+ 
252+ device = V.graph.get_current_device_or_throw()
253+ stream = PythonWrapperCodegen.write_get_raw_stream(
254+ self, device.index, V.graph
255+ )
256+ self.writeline(
257+ f"{use_tasklist}, {work_items_tensor}, {task_offsets_tensor}, "
258+ f"{split_bases_tensor}, {max_sub} = "
259+ f"_get_or_build_dkdv_task_list("
260+ f"{q_num_blocks_name}, {full_q_num_blocks_name}, "
261+ f"{dispatch_spec.batch_size}, "
262+ f"{dispatch_spec.num_kv_heads}, {dispatch_spec.num_kv_blocks}, "
263+ f"{dispatch_spec.sparse_kv_multiple}, "
264+ f"{dispatch_spec.launch_programs}, {dk_name}.device)"
265+ )
266+ self.writeline(f"if {use_tasklist}:")
267+ self.writeline(f" {num_split_bases} = {split_bases_tensor}.size(0)")
268+ self.writeline(f" if {num_split_bases} > 0:")
269+ self.writeline(
270+ f" {dk_partial} = torch.empty_strided("
271+ f"({max_sub}, *{dk_name}.size()), "
272+ f"({dispatch_spec.partial_dk_stride}, *{dk_name}.stride()), "
273+ f"dtype=torch.float32, device={dk_name}.device)"
274+ )
275+ self.writeline(f" {dk_partial}.zero_()")
276+ self.writeline(
277+ f" {dv_partial} = torch.empty_strided("
278+ f"({max_sub}, *{dv_name}.size()), "
279+ f"({dispatch_spec.partial_dv_stride}, *{dv_name}.stride()), "
280+ f"dtype=torch.float32, device={dv_name}.device)"
281+ )
282+ self.writeline(f" {dv_partial}.zero_()")
283+ self.writeline(" else:")
284+ self.writeline(f" {split_bases_tensor} = None")
285+ self.writeline(f" {dk_partial} = {dk_name}")
286+ self.writeline(f" {dv_partial} = {dv_name}")
287+ 
288+ runtime_values = {
289+ "work_items_t": work_items_tensor,
290+ "task_offsets_t": task_offsets_tensor,
291+ "split_bases_t": split_bases_tensor,
292+ "dk_partial": dk_partial,
293+ "dv_partial": dv_partial,
294+ }
295+ 
296+ def replace_runtime_args(call_args):
297+ replacements = {
298+ outer_name: runtime_values[wrapper_name]
299+ for outer_name, wrapper_name in runtime_arg_names.items()
300+ }
301+ return [replacements.get(str(arg), arg) for arg in call_args]
302+ 
303+ tasklist_args = replace_runtime_args(tasklist_call_args)
304+ tasklist_grid = [dispatch_spec.launch_programs, 1, 1]
305+ tasklist_args_text = ", ".join(
306+ self.prepare_triton_kernel_call([*tasklist_args, *tasklist_grid])
307+ )
308+ tasklist_no_split_args = replace_runtime_args(
309+ tasklist_no_split_call_args
310+ )
311+ tasklist_no_split_args_text = ", ".join(
312+ self.prepare_triton_kernel_call(
313+ [*tasklist_no_split_args, *tasklist_grid]
314+ )
315+ )
316+ self.writeline(f" if {num_split_bases} > 0:")
317+ self.writeline(
318+ f" {tasklist_kernel_name}.run("
319+ f"{tasklist_args_text}, stream={stream})"
320+ )
321+ self.writeline(" else:")
322+ self.writeline(
323+ f" {tasklist_no_split_kernel_name}.run("
324+ f"{tasklist_no_split_args_text}, stream={stream})"
325+ )
326+ 
327+ reduce_args = replace_runtime_args(reduce_call_args)
328+ reduce_args_text = ", ".join(
329+ self.prepare_triton_kernel_call(
330+ [*reduce_args, num_split_bases, 1, 1]
331+ )
332+ )
333+ self.writeline(f" if {num_split_bases} > 0:")
334+ self.writeline(
335+ f" {reduce_kernel_name}.run("
336+ f"{reduce_args_text}, stream={stream})"
337+ )
338+ self.writeline(
339+ f" del {work_items_tensor}, {task_offsets_tensor}, "
340+ f"{split_bases_tensor}, {dk_partial}, {dv_partial}"
341+ )
342+ self.writeline("else:")
343+ legacy_grid = [dispatch_spec.launch_programs, 1, 1]
344+ legacy_args_text = ", ".join(
345+ self.prepare_triton_kernel_call([*legacy_call_args, *legacy_grid])
346+ )
347+ self.writeline(
348+ f" {legacy_kernel_name}.run({legacy_args_text}, stream={stream})"
349+ )
350+ 
182 351 
183 @classmethod352 @classmethod
184 def _get_triton_info_kernel_cls(cls):353 def _get_triton_info_kernel_cls(cls):
Mtorch_npu/_inductor/config.py+161-0
@@ -1,6 +1,7 @@
1import logging1import logging
2import os # noqa: C1012import os # noqa: C101
3import re3import re
4+from typing import Optional
4import sys5import sys
5 6 
6import torch7import torch
@@ -179,6 +180,11 @@ inductor_config.triton.mix_order_reduction = False
179inductor_config.loop_reindexing_after_fusion = False180inductor_config.loop_reindexing_after_fusion = False
180 181 
181 182 
183+def _read_env_bool(name: str, default: str = "False") -> bool:
184+ value = os.environ.get(name, default)
185+ return value.strip().lower() in ("1", "true", "yes", "on")
186+ 
187+ 
182# Enable the SIMT Welford lowering for variance and layer normalization.188# Enable the SIMT Welford lowering for variance and layer normalization.
183# Keep it disabled by default while the new path is being rolled out.189# Keep it disabled by default while the new path is being rolled out.
184enable_welford = os.getenv("TORCHINDUCTOR_ENABLE_WELFORD", "0") == "1"190enable_welford = os.getenv("TORCHINDUCTOR_ENABLE_WELFORD", "0") == "1"
@@ -354,6 +360,9 @@ simt_default_warp_stacksize = 256 * 32
354default_nddma_switch = "1" if is_ascend950 else "0"360default_nddma_switch = "1" if is_ascend950 else "0"
355nddma_switch = os.getenv("TORCHINDUCTOR_NDDMA", default_nddma_switch) == "1"361nddma_switch = os.getenv("TORCHINDUCTOR_NDDMA", default_nddma_switch) == "1"
356enable_fast_gelu = os.getenv("TORCHINDUCTOR_ENABLE_FAST_GELU", "0") == "1"362enable_fast_gelu = os.getenv("TORCHINDUCTOR_ENABLE_FAST_GELU", "0") == "1"
363+enable_flex_attention_dq_before_scale_materialize = os.environ.get(
364+ "FLEX_ATTENTION_DQ_BEFORE_SCALE_MATERIALIZE", "1"
365+).lower() in ("1", "true", "yes")
357 366 
358aggresive_autotune = os.getenv("INDUCTOR_ASCEND_AGGRESSIVE_AUTOTUNE", "0").lower() in (367aggresive_autotune = os.getenv("INDUCTOR_ASCEND_AGGRESSIVE_AUTOTUNE", "0").lower() in (
359 "1",368 "1",
@@ -421,3 +430,155 @@ autotune_continue_on_failure = os.environ.get('TORCHINDUCTOR_NPU_BACKEND') == "d
421enable_fused_matmul_relu = _parse_bool_env(430enable_fused_matmul_relu = _parse_bool_env(
422 "TORCHINDUCTOR_ENABLE_FUSED_MATMUL_RELU", False431 "TORCHINDUCTOR_ENABLE_FUSED_MATMUL_RELU", False
423)432)
433+ 
434+FLEX_ATTENTION_NPU_COMPILE_HINT_KEYS = (
435+ "limit_auto_multi_buffer_buffer",
436+ "multibuffer",
437+ "unit_flag",
438+ "enable_ubuf_saving",
439+ "hfusion_enable_multiple_consumer_fusion",
440+ "enable_select_analysis",
441+ "limit_auto_multi_buffer_only_for_local_buffer",
442+ "limit_auto_multi_buffer_of_local_buffer",
443+ "set_workspace_multibuffer",
444+ "tile_mix_vector_loop",
445+ "tile_mix_cube_loop",
446+ "enable_dynamic_cv_pipeline",
447+ "intra_cache_num",
448+ "inter_cache_num",
449+ "enable_cross_if_fusion",
450+ "enable_buffer_insert_optimization",
451+ "enable_ub_refine_opt",
452+)
453+ 
454+ 
455+class flex_attention:
456+ enable_npu_optimization = False
457+ use_config_generator = True
458+ metadata_auto_infer = True
459+ flexattention_mask_out = True
460+ # Keep rollout disabled until generated outputcode and NPU numerics have
461+ # been reviewed. Unsupported graphs always retain the legacy dK/dV path.
462+ bwd_dkdv_tasklist = True
463+ 
464+ multibuffer = True
465+ unit_flag = True
466+ enable_ubuf_saving = True
467+ limit_auto_multi_buffer_buffer = "no-limit"
468+ hfusion_enable_multiple_consumer_fusion = True
469+ enable_select_analysis = False
470+ limit_auto_multi_buffer_only_for_local_buffer = False
471+ limit_auto_multi_buffer_of_local_buffer = "no-limit"
472+ set_workspace_multibuffer = 4
473+ tile_mix_vector_loop = 4
474+ tile_mix_cube_loop = 4
475+ enable_dynamic_cv_pipeline = False
476+ 
477+ bwd_dq_limit_auto_multi_buffer_of_local_buffer = "no-l0c"
478+ bwd_dkdv_limit_auto_multi_buffer_of_local_buffer = "no-l0c"
479+ 
480+ enable_buffer_insert_optimization = False
481+ enable_ub_refine_opt = False
482+ 
483+ @classmethod
484+ def _filter_compile_options_for_soc(cls, options: dict) -> dict:
485+ options = options.copy()
486+ if is_ascend950:
487+ options.pop("enable_dynamic_cv_pipeline", None)
488+ return options
489+ 
490+ @classmethod
491+ def _compile_options(cls, keys, overrides: Optional[dict] = None) -> dict:
492+ options = {key: getattr(cls, key) for key in keys}
493+ if overrides:
494+ options.update(overrides)
495+ return cls._filter_compile_options_for_soc(options)
496+ 
497+ @classmethod
498+ def get_npu_compile_hint_params(cls) -> dict:
499+ return cls._compile_options(FLEX_ATTENTION_NPU_COMPILE_HINT_KEYS)
500+ 
501+ @classmethod
502+ def get_sparse_mask_cvpipeline_compile_options(
503+ cls,
504+ *,
505+ enabled: bool,
506+ tile_mix_loop: int,
507+ enable_compile_hint: bool,
508+ ) -> dict:
509+ return cls._compile_options(
510+ (
511+ "enable_ubuf_saving",
512+ "unit_flag",
513+ "set_workspace_multibuffer",
514+ "limit_auto_multi_buffer_buffer",
515+ "hfusion_enable_multiple_consumer_fusion",
516+ ),
517+ overrides={
518+ "multibuffer": enabled,
519+ "limit_auto_multi_buffer_only_for_local_buffer": not enabled,
520+ "tile_mix_vector_loop": tile_mix_loop,
521+ "tile_mix_cube_loop": tile_mix_loop,
522+ "ENABLE_COMPILE_HINT": enable_compile_hint if enabled else False,
523+ "intra_cache_num": 3,
524+ "inter_cache_num": 2,
525+ "enable_cross_if_fusion": True,
526+ "enable_buffer_insert_optimization": True,
527+ "enable_ub_refine_opt": True,
528+ },
529+ )
530+ 
531+ @classmethod
532+ def get_bwd_dq_compile_options(cls) -> dict:
533+ return cls._compile_options(
534+ (
535+ "limit_auto_multi_buffer_buffer",
536+ "hfusion_enable_multiple_consumer_fusion",
537+ "enable_select_analysis",
538+ ),
539+ overrides={
540+ "limit_auto_multi_buffer_of_local_buffer": (
541+ cls.bwd_dq_limit_auto_multi_buffer_of_local_buffer
542+ ),
543+ "intra_cache_num": 3,
544+ "inter_cache_num": 2,
545+ },
546+ )
547+ 
548+ @classmethod
549+ def get_bwd_dkdv_compile_options(cls) -> dict:
550+ return cls._compile_options(
551+ (
552+ "limit_auto_multi_buffer_buffer",
553+ "hfusion_enable_multiple_consumer_fusion",
554+ "unit_flag",
555+ "enable_dynamic_cv_pipeline",
556+ ),
557+ overrides={
558+ "limit_auto_multi_buffer_of_local_buffer": (
559+ cls.bwd_dkdv_limit_auto_multi_buffer_of_local_buffer
560+ ),
561+ "intra_cache_num": 2,
562+ "inter_cache_num": 1,
563+ },
564+ )
565+ 
566+ 
567+flex_attention.bwd_dkdv_tasklist = _read_env_bool(
568+ "TORCHINDUCTOR_ASCEND_FLEX_ATTENTION_BWD_DKDV_TASKLIST",
569+ "1" if flex_attention.bwd_dkdv_tasklist else "0",
570+)
571+flex_attention.flexattention_mask_out = _read_env_bool(
572+ "TORCHINDUCTOR_FLEXATTENTION_MASKOUT",
573+ "1" if flex_attention.flexattention_mask_out else "0",
574+)
575+ 
576+ 
577+def apply_flex_attention_npu_params(config: dict, *, enable: bool) -> dict:
578+ config = config.copy()
579+ if enable:
580+ config.update(flex_attention.get_npu_compile_hint_params())
581+ config["ENABLE_COMPILE_HINT"] = True
582+ else:
583+ config["ENABLE_COMPILE_HINT"] = False
584+ return config
Atorch_npu/_inductor/flex_attention_tasklist.py+329-0
@@ -0,0 +1,329 @@
1+import inspect
2+import math
3+import re
4+import textwrap
5+from dataclasses import dataclass
6+ 
7+import torch
8+ 
9+ 
10+@dataclass(frozen=True)
11+class RuntimeTemplateArg:
12+ name: str
13+ dtype: torch.dtype
14+ rank: int
15+ wrapper_name: str
16+ 
17+ 
18+@dataclass(frozen=True)
19+class FlexAttentionDkdvDispatchSpec:
20+ launch_programs: int
21+ batch_size: int
22+ num_kv_heads: int
23+ num_kv_blocks: int
24+ sparse_kv_multiple: int
25+ sparse_kv_block_size: int
26+ block_n1: int
27+ partial_dk_stride: int
28+ partial_dv_stride: int
29+ 
30+ 
31+def is_dkdv_tasklist_codegen_compatible(
32+ *,
33+ cpp_wrapper,
34+ aot_mode,
35+ bq,
36+ bkv,
37+ sparse_z,
38+ sparse_hq,
39+ sparse_kv_block_size,
40+ block_n1,
41+ q_num_blocks_dtype,
42+ full_q_num_blocks_dtype,
43+ q_num_blocks_contiguous,
44+ full_q_num_blocks_contiguous,
45+ accum_dtype,
46+):
47+ static_dimensions = (bq, bkv, sparse_z, sparse_hq)
48+ return (
49+ not cpp_wrapper
50+ and not aot_mode
51+ and all(isinstance(value, int) for value in static_dimensions)
52+ and bq == 1
53+ and bq == bkv
54+ and sparse_z == bq
55+ and sparse_hq == 1
56+ and block_n1 > 0
57+ and sparse_kv_block_size % block_n1 == 0
58+ and q_num_blocks_dtype == torch.int32
59+ and full_q_num_blocks_dtype == torch.int32
60+ and q_num_blocks_contiguous
61+ and full_q_num_blocks_contiguous
62+ and accum_dtype == torch.float32
63+ )
64+ 
65+ 
66+def compute_dkdv_sparse_weights(q_num_blks, full_q_num_blks):
67+ weights = (q_num_blks + full_q_num_blks).reshape(-1)
68+ return [
69+ int(value)
70+ for value in weights.detach().to("cpu", dtype=torch.int64).tolist()
71+ ]
72+ 
73+ 
74+def should_use_dkdv_tasklist(
75+ w_sparse,
76+ batch_size,
77+ num_kv_heads,
78+ num_kv_blocks,
79+ sparse_kv_multiple,
80+ num_core,
81+):
82+ if (
83+ not w_sparse
84+ or batch_size != 1
85+ or num_core <= 0
86+ or sparse_kv_multiple <= 0
87+ ):
88+ return False
89+ 
90+ total_base = batch_size * num_kv_heads * num_kv_blocks
91+ total_weight = num_kv_heads * sum(w_sparse)
92+ if total_base == 0 or total_weight == 0:
93+ return False
94+ 
95+ mean_weight = total_weight / total_base
96+ full_rounds, tail_cores = divmod(total_base, num_core)
97+ has_significant_tail = (
98+ tail_cores > 0 and full_rounds <= 2 and tail_cores / num_core < 0.5
99+ )
100+ has_weight_imbalance = (
101+ tail_cores == 0 and max(w_sparse) / mean_weight > 1.5
102+ )
103+ return has_significant_tail or has_weight_imbalance
104+ 
105+ 
106+def bin_pack_dkdv_hkv_continuous(work_items, num_core):
107+ bins = [[] for _ in range(num_core)]
108+ bin_weights = [0.0] * num_core
109+ groups = {}
110+ for item in work_items:
111+ groups.setdefault(item[0], []).append(item)
112+ 
113+ for kv_head in sorted(groups):
114+ group = sorted(
115+ groups[kv_head], key=lambda item: item[5], reverse=True
116+ )
117+ for item in group:
118+ lightest = bin_weights.index(min(bin_weights))
119+ bins[lightest].append(item)
120+ bin_weights[lightest] += item[5]
121+ return bins
122+ 
123+ 
124+def build_dkdv_task_list(
125+ w_sparse,
126+ batch_size,
127+ num_kv_heads,
128+ num_kv_blocks,
129+ sparse_kv_multiple,
130+ num_core,
131+):
132+ target = max(
133+ num_kv_heads * sum(w_sparse) / max(num_core, 1),
134+ 1.0,
135+ )
136+ target_int = max(int(target), 1)
137+ weights_per_kv_block = [
138+ int(w_sparse[kv_block // sparse_kv_multiple])
139+ for kv_block in range(num_kv_blocks)
140+ ]
141+ template_items = []
142+ template_split_bases = []
143+ max_sub = 1
144+ for kv_block, weight in enumerate(weights_per_kv_block):
145+ if weight == 0:
146+ continue
147+ if weight <= target:
148+ template_items.append(
149+ (kv_block, 0, 1, 0, float(weight))
150+ )
151+ continue
152+ 
153+ split_count = max(1, math.ceil(weight / target_int))
154+ template_split_bases.append((kv_block, split_count))
155+ max_sub = max(max_sub, split_count)
156+ split_weight = weight / split_count
157+ for sub_id in range(split_count):
158+ template_items.append(
159+ (kv_block, sub_id, split_count, 1, split_weight)
160+ )
161+ 
162+ weighted_items = []
163+ split_bases = []
164+ for batch_idx in range(batch_size):
165+ for kv_head in range(num_kv_heads):
166+ weighted_items.extend(
167+ (kv_head, *item) for item in template_items
168+ )
169+ split_bases.extend(
170+ (kv_head, *item)
171+ for item in template_split_bases
172+ )
173+ 
174+ bins = bin_pack_dkdv_hkv_continuous(weighted_items, num_core)
175+ work_items = []
176+ task_offsets = [0]
177+ for bin_items in bins:
178+ work_items.extend(item[:5] for item in bin_items)
179+ task_offsets.append(len(work_items))
180+ return work_items, task_offsets, split_bases, max_sub
181+ 
182+ 
183+def get_or_build_dkdv_task_list(
184+ q_num_blks,
185+ full_q_num_blks,
186+ batch_size,
187+ num_kv_heads,
188+ num_kv_blocks,
189+ sparse_kv_multiple,
190+ num_core,
191+ device,
192+):
193+ cache_key = (
194+ batch_size,
195+ num_kv_heads,
196+ num_kv_blocks,
197+ sparse_kv_multiple,
198+ num_core,
199+ device,
200+ )
201+ try:
202+ q_num_blks_version = q_num_blks._version
203+ full_q_num_blks_version = full_q_num_blks._version
204+ except RuntimeError:
205+ q_num_blks_version = None
206+ full_q_num_blks_version = None
207+ 
208+ cache = getattr(q_num_blks, "_npu_dkdv_tasklist_cache", None)
209+ if cache is not None:
210+ entry = cache.get(cache_key)
211+ if (
212+ entry is not None
213+ and entry[0] is full_q_num_blks
214+ and entry[1] == q_num_blks_version
215+ and entry[2] == full_q_num_blks_version
216+ ):
217+ return entry[3]
218+ 
219+ weights = compute_dkdv_sparse_weights(q_num_blks, full_q_num_blks)
220+ use_tasklist = should_use_dkdv_tasklist(
221+ weights,
222+ batch_size,
223+ num_kv_heads,
224+ num_kv_blocks,
225+ sparse_kv_multiple,
226+ num_core,
227+ )
228+ if use_tasklist:
229+ work_items, task_offsets, split_bases, max_sub = (
230+ build_dkdv_task_list(
231+ weights,
232+ batch_size,
233+ num_kv_heads,
234+ num_kv_blocks,
235+ sparse_kv_multiple,
236+ num_core,
237+ )
238+ )
239+ if work_items:
240+ work_items_tensor = torch.tensor(
241+ work_items, dtype=torch.int32, device=device
242+ )
243+ else:
244+ work_items_tensor = torch.zeros(
245+ (0, 5), dtype=torch.int32, device=device
246+ )
247+ task_offsets_tensor = torch.tensor(
248+ task_offsets, dtype=torch.int32, device=device
249+ )
250+ if split_bases:
251+ split_bases_tensor = torch.tensor(
252+ split_bases, dtype=torch.int32, device=device
253+ )
254+ else:
255+ split_bases_tensor = torch.zeros(
256+ (0, 3), dtype=torch.int32, device=device
257+ )
258+ result = (
259+ True,
260+ work_items_tensor,
261+ task_offsets_tensor,
262+ split_bases_tensor,
263+ max_sub,
264+ )
265+ else:
266+ result = (False, None, None, None, 1)
267+ 
268+ if q_num_blks_version is not None and full_q_num_blks_version is not None:
269+ if cache is None:
270+ cache = {}
271+ setattr(q_num_blks, "_npu_dkdv_tasklist_cache", cache) # noqa: B010
272+ cache[cache_key] = (
273+ full_q_num_blks,
274+ q_num_blks_version,
275+ full_q_num_blks_version,
276+ result,
277+ )
278+ return result
279+ 
280+ 
281+def _generated_helper_source(function, generated_name, replacements=()):
282+ source = textwrap.dedent(inspect.getsource(function))
283+ definition = f"def {function.__name__}("
284+ source = source.replace(definition, f"def {generated_name}(", 1)
285+ for old_name, new_name in replacements:
286+ source = re.sub(
287+ rf"(?<![\w]){re.escape(old_name)}\(",
288+ f"{new_name}(",
289+ source,
290+ )
291+ return source
292+ 
293+ 
294+DKDV_TASKLIST_HELPER_SOURCE = "\n\n".join(
295+ (
296+ _generated_helper_source(
297+ compute_dkdv_sparse_weights, "_compute_dkdv_sparse_weights"
298+ ),
299+ _generated_helper_source(
300+ should_use_dkdv_tasklist, "_should_use_dkdv_tasklist"
301+ ),
302+ _generated_helper_source(
303+ bin_pack_dkdv_hkv_continuous,
304+ "_bin_pack_dkdv_hkv_continuous",
305+ ),
306+ _generated_helper_source(
307+ build_dkdv_task_list,
308+ "_build_dkdv_task_list",
309+ (
310+ (
311+ "bin_pack_dkdv_hkv_continuous",
312+ "_bin_pack_dkdv_hkv_continuous",
313+ ),
314+ ),
315+ ),
316+ _generated_helper_source(
317+ get_or_build_dkdv_task_list,
318+ "_get_or_build_dkdv_task_list",
319+ (
320+ (
321+ "compute_dkdv_sparse_weights",
322+ "_compute_dkdv_sparse_weights",
323+ ),
324+ ("should_use_dkdv_tasklist", "_should_use_dkdv_tasklist"),
325+ ("build_dkdv_task_list", "_build_dkdv_task_list"),
326+ ),
327+ ),
328+ )
329+)
Mtorch_npu/_inductor/kernel/__init__.py+1-1
@@ -1,4 +1,4 @@
1from .mm import _register_npu_inductor_mm, _register_npu_inductor_addmm1from .mm import _register_npu_inductor_mm, _register_npu_inductor_addmm
2from .bmm import _register_npu_inductor_bmm2from .bmm import _register_npu_inductor_bmm
3from .mm_grouped import _register_npu_inductor_grouped_mm3from .mm_grouped import _register_npu_inductor_grouped_mm
4-from .flex_attention import _register_npu_inductor_flex_attention, _validate_device4+from .flex_attention import _register_npu_inductor_flex_attention, _validate_device, patch_flex_attention
Atorch_npu/_inductor/kernel/flex_attention_config_generator.py+973-0
@@ -0,0 +1,973 @@
1+"""
2+Flex Attention Configuration Generator.
3+ 
4+This module provides FlexAttentionConfigGenerator class that dynamically
5+generates candidate configurations for Flex Attention kernel autotuning.
6+Similar to TileGenerator, it generates BLOCK_M/BLOCK_N combinations based
7+on input shapes and hardware constraints.
8+"""
9+ 
10+from dataclasses import dataclass
11+from enum import Enum
12+from typing import Optional, Union
13+ 
14+import torch
15+from torch._inductor import config as inductor_config
16+ 
17+from .. import config as npu_config
18+ 
19+log = npu_config.log
20+ 
21+ 
22+class FlexMode(Enum):
23+ """Operation mode for Flex Attention."""
24+ FWD = "fwd"
25+ BWD = "bwd"
26+ 
27+ 
28+@dataclass
29+class FlexAttentionConfig:
30+ """Configuration for Flex Attention kernel."""
31+ block_m: int
32+ block_n: int
33+ num_warps: int
34+ num_stages: int
35+ npu_params: Optional[dict] = None
36+ 
37+ def to_dict(self) -> dict:
38+ """Convert to dictionary format."""
39+ result = {
40+ "BLOCK_M": self.block_m,
41+ "BLOCK_N": self.block_n,
42+ "num_warps": self.num_warps,
43+ "num_stages": self.num_stages,
44+ }
45+ if self.npu_params:
46+ result.update(self.npu_params)
47+ return result
48+ 
49+ 
50+class FlexAttentionConfigGenerator:
51+ """
52+ Generate candidate configurations for Flex Attention kernel.
53+ 
54+ This class dynamically generates BLOCK_M/BLOCK_N combinations based on
55+ input shapes and hardware constraints, similar to how TileGenerator
56+ works for general NPU kernels.
57+ 
58+ Key features:
59+ 1. Constraint-aware: Only generates configs that satisfy SPARSE constraints
60+ 2. Performance-oriented: Considers wave efficiency and UB constraints
61+ 3. Extensible: Supports num_warps/num_stages/NPU parameter variations
62+ 
63+ Example:
64+ >>> generator = FlexAttentionConfigGenerator(
65+ ... query_shape=(1, 8, 1024, 128),
66+ ... key_shape=(1, 8, 1024, 128),
67+ ... sparse_q_block_size=128,
68+ ... sparse_kv_block_size=128,
69+ ... dtype=torch.float16,
70+ ... num_cube_core=24,
71+ ... mode=FlexMode.FWD,
72+ ... )
73+ >>> configs = generator.generate_configs()
74+ >>> print(len(configs)) # 10-20 configs
75+ """
76+ 
77+ BLOCK_SIZE_CANDIDATES = [256, 128, 64, 32, 16]
78+ 
79+ MAX_CONFIGS = 30
80+ 
81+ def __init__(
82+ self,
83+ query_shape: tuple,
84+ key_shape: tuple,
85+ sparse_q_block_size: int,
86+ sparse_kv_block_size: int,
87+ dtype: torch.dtype,
88+ num_cube_core: int,
89+ mode: FlexMode = FlexMode.FWD,
90+ ):
91+ """
92+ Initialize the configuration generator.
93+ 
94+ Args:
95+ query_shape: Shape of query tensor (batch, heads, seq_len_q, head_dim)
96+ key_shape: Shape of key tensor (batch, heads, seq_len_kv, head_dim)
97+ sparse_q_block_size: SPARSE_Q_BLOCK_SIZE constraint
98+ sparse_kv_block_size: SPARSE_KV_BLOCK_SIZE constraint
99+ dtype: Data type of tensors
100+ num_cube_core: Number of AICore (cube cores) available
101+ mode: Forward or backward mode
102+ """
103+ self.batch_size = query_shape[0]
104+ self.num_heads = query_shape[1]
105+ self.seq_len_q = query_shape[2]
106+ self.head_dim = query_shape[3]
107+ self.seq_len_kv = key_shape[2]
108+ 
109+ self.sparse_q_block_size = sparse_q_block_size
110+ self.sparse_kv_block_size = sparse_kv_block_size
111+ 
112+ self.dtype = dtype
113+ self.dtype_bytes = self._get_dtype_bytes(dtype)
114+ 
115+ self.num_cube_core = num_cube_core
116+ 
117+ self.mode = mode
118+ 
119+ self.valid_block_m = self._get_valid_block_sizes(sparse_q_block_size)
120+ self.valid_block_n = self._get_valid_block_sizes(sparse_kv_block_size)
121+ 
122+ self.configs: list[FlexAttentionConfig] = []
123+ 
124+ def _get_dtype_bytes(self, dtype: torch.dtype) -> int:
125+ """Get bytes per element for dtype."""
126+ dtype_bytes_map = {
127+ torch.float16: 2,
128+ torch.bfloat16: 2,
129+ torch.float32: 4,
130+ torch.int8: 1,
131+ torch.int16: 2,
132+ torch.int32: 4,
133+ }
134+ return dtype_bytes_map.get(dtype, 4)
135+ 
136+ def _get_valid_block_sizes(self, sparse_block_size: int) -> list[int]:
137+ """
138+ Get valid block sizes that divide SPARSE_BLOCK_SIZE.
139+ 
140+ For SPARSE_BLOCK_SIZE=128, returns: [128, 64, 32, 16]
141+ For SPARSE_BLOCK_SIZE=64, returns: [64, 32, 16]
142+ """
143+ valid_sizes = []
144+ for size in self.BLOCK_SIZE_CANDIDATES:
145+ if sparse_block_size % size == 0:
146+ valid_sizes.append(size)
147+ return valid_sizes
148+ 
149+ def generate_configs(self) -> list[dict]:
150+ """
151+ Generate all candidate configurations.
152+ 
153+ Returns:
154+ List of config dictionaries with BLOCK_M, BLOCK_N, num_warps, num_stages
155+ """
156+ self.configs = []
157+ 
158+ self._generate_block_combinations()
159+ self._filter_by_ub_constraint()
160+ self._add_npu_params()
161+ self._limit_config_count()
162+ 
163+ return [cfg.to_dict() for cfg in self.configs]
164+ 
165+ def _generate_block_combinations(self):
166+ """Generate BLOCK_M x BLOCK_N combinations based on mode."""
167+ if self.mode == FlexMode.FWD:
168+ self._generate_fwd_combinations()
169+ else:
170+ self._generate_bwd_combinations()
171+ 
172+ def _generate_fwd_combinations(self):
173+ """
174+ Generate forward pass configurations.
175+ 
176+ Strategy:
177+ 1. Start with safe default config (16, 16) for persistent mode compatibility
178+ 2. Add configs with different BLOCK_M/BLOCK_N ratios for autotuning
179+ 3. Consider wave efficiency (programs per AICore)
180+ """
181+ # Use (16, 16) as the default config to ensure compilation stability
182+ # This is especially important for persistent kernel mode which may have
183+ # stricter constraints on tile sizes
184+ default_m, default_n = 16, 16
185+ 
186+ if default_m in self.valid_block_m and default_n in self.valid_block_n:
187+ self.configs.append(FlexAttentionConfig(
188+ block_m=default_m,
189+ block_n=default_n,
190+ num_warps=4,
191+ num_stages=3,
192+ ))
193+ 
194+ seen = {(default_m, default_n)}
195+ 
196+ for block_m in self.valid_block_m:
197+ for block_n in self.valid_block_n:
198+ if (block_m, block_n) in seen:
199+ continue
200+ 
201+ programs_m = (self.seq_len_q + block_m - 1) // block_m
202+ total_programs = programs_m * self.batch_size * self.num_heads
203+ wave_efficiency = total_programs / self.num_cube_core
204+ 
205+ if wave_efficiency >= 0.5 or total_programs <= self.num_cube_core:
206+ self.configs.append(FlexAttentionConfig(
207+ block_m=block_m,
208+ block_n=block_n,
209+ num_warps=4,
210+ num_stages=3,
211+ ))
212+ seen.add((block_m, block_n))
213+ 
214+ def _generate_bwd_combinations(self):
215+ """
216+ Generate backward pass configurations.
217+ 
218+ Backward uses BLOCK_M1, BLOCK_N1, BLOCK_M2, BLOCK_N2.
219+ Constraint: BLOCK_N1 % BLOCK_M1 == 0
220+ """
221+ # Match the backward outputcode target tile first. dkdv and dq lowering
222+ # may still override their unused BLOCK_* dimensions independently.
223+ default_m, default_n = 64, 64
224+ if default_m in self.valid_block_m and default_n in self.valid_block_n:
225+ self.configs.append(FlexAttentionConfig(
226+ block_m=default_m,
227+ block_n=default_n,
228+ num_warps=4,
229+ num_stages=1,
230+ ))
231+ 
232+ seen = {(default_m, default_n)}
233+ 
234+ for block_m1 in self.valid_block_m:
235+ for block_n1 in self.valid_block_n:
236+ if block_n1 % block_m1 != 0:
237+ continue
238+ 
239+ if (block_m1, block_n1) in seen:
240+ continue
241+ 
242+ self.configs.append(FlexAttentionConfig(
243+ block_m=block_m1,
244+ block_n=block_n1,
245+ num_warps=4,
246+ num_stages=1,
247+ ))
248+ seen.add((block_m1, block_n1))
249+ 
250+ def _filter_by_ub_constraint(self):
251+ """
252+ Filter configs that exceed UB size limit.
253+ 
254+ UB usage estimation:
255+ - Q block: BLOCK_M * head_dim * dtype_bytes
256+ - K block: BLOCK_N * head_dim * dtype_bytes
257+ - V block: BLOCK_N * head_dim * dtype_bytes
258+ - Acc buffer: BLOCK_M * BLOCK_N * 4 (float32)
259+ """
260+ filtered = []
261+ 
262+ for cfg in self.configs:
263+ q_ub = cfg.block_m * self.head_dim * self.dtype_bytes
264+ k_ub = cfg.block_n * self.head_dim * self.dtype_bytes
265+ v_ub = cfg.block_n * self.head_dim * self.dtype_bytes
266+ acc_ub = cfg.block_m * cfg.block_n * 4
267+ 
268+ total_ub = q_ub + k_ub + v_ub + acc_ub
269+ 
270+ if total_ub <= npu_config.ub_size * 0.8:
271+ filtered.append(cfg)
272+ 
273+ self.configs = filtered if filtered else self.configs[:1]
274+ 
275+ def _add_npu_params(self):
276+ """
277+ Add NPU optimization parameters if enabled.
278+ 
279+ Similar to TileGenerator.tune_multibuffer()
280+ """
281+ log.info("[flex_attention] NPU optimization enabled: %s", npu_config.flex_attention.enable_npu_optimization)
282+ 
283+ if not npu_config.flex_attention.enable_npu_optimization:
284+ # Even when NPU optimization is disabled, we need to set ENABLE_COMPILE_HINT
285+ # to avoid NameError in kernel code
286+ for cfg in self.configs:
287+ cfg.npu_params = npu_config.apply_flex_attention_npu_params(
288+ cfg.npu_params or {},
289+ enable=False,
290+ )
291+ return
292+ 
293+ npu_params = npu_config.apply_flex_attention_npu_params(
294+ {},
295+ enable=True,
296+ )
297+ 
298+ log.debug("NPU parameters: %s", npu_params)
299+ 
300+ # Keep the original configs as a conservative fallback, then append the
301+ # NPU-tuned variants so autotuning can still fall back if the enhanced
302+ # path overflows UB or hits backend bugs.
303+ new_configs = []
304+ for cfg in self.configs:
305+ cfg.npu_params = npu_config.apply_flex_attention_npu_params(
306+ cfg.npu_params or {},
307+ enable=False,
308+ )
309+ new_configs.append(cfg)
310+ cfg_with_npu = FlexAttentionConfig(
311+ block_m=cfg.block_m,
312+ block_n=cfg.block_n,
313+ num_warps=cfg.num_warps,
314+ num_stages=cfg.num_stages,
315+ npu_params=npu_params.copy(),
316+ )
317+ new_configs.append(cfg_with_npu)
318+ self.configs = new_configs # Replace instead of append
319+ 
320+ def _limit_config_count(self):
321+ """
322+ Limit config count to avoid excessive autotuning time.
323+ 
324+ Strategy: Select configs with diverse BLOCK_M/BLOCK_N ratios
325+ """
326+ if len(self.configs) <= self.MAX_CONFIGS:
327+ return
328+ 
329+ ratio_groups: dict[float, list[FlexAttentionConfig]] = {}
330+ for cfg in self.configs:
331+ ratio = cfg.block_m / cfg.block_n
332+ if ratio not in ratio_groups:
333+ ratio_groups[ratio] = []
334+ ratio_groups[ratio].append(cfg)
335+ 
336+ selected = []
337+ per_group = max(1, self.MAX_CONFIGS // len(ratio_groups))
338+ for group in ratio_groups.values():
339+ selected.extend(group[:per_group])
340+ 
341+ self.configs = selected[:self.MAX_CONFIGS]
342+ 
343+ def calculate_wave_efficiency(self, block_m: int, block_n: int) -> tuple[int, float]:
344+ """
345+ Calculate wave efficiency for given block sizes.
346+ 
347+ Args:
348+ block_m: BLOCK_M size
349+ block_n: BLOCK_N size
350+ 
351+ Returns:
352+ Tuple of (waves, efficiency)
353+ """
354+ programs_m = (self.seq_len_q + block_m - 1) // block_m
355+ total_programs = programs_m * self.batch_size * self.num_heads
356+ 
357+ waves = (total_programs + self.num_cube_core - 1) // self.num_cube_core
358+ 
359+ efficiency = total_programs / (waves * self.num_cube_core) if waves > 0 else 0.0
360+ 
361+ return waves, efficiency
362+ 
363+def generate_fwd_configs(
364+ query_shape: tuple,
365+ key_shape: tuple,
366+ sparse_q_block_size: int,
367+ sparse_kv_block_size: int,
368+ dtype: torch.dtype,
369+ num_cube_core: int,
370+) -> list[dict]:
371+ """
372+ Convenience function to generate forward configs.
373+ 
374+ Args:
375+ query_shape: Shape of query tensor
376+ key_shape: Shape of key tensor
377+ sparse_q_block_size: SPARSE_Q_BLOCK_SIZE
378+ sparse_kv_block_size: SPARSE_KV_BLOCK_SIZE
379+ dtype: Data type
380+ num_cube_core: Number of AICore
381+ 
382+ Returns:
383+ List of config dictionaries
384+ """
385+ generator = FlexAttentionConfigGenerator(
386+ query_shape=query_shape,
387+ key_shape=key_shape,
388+ sparse_q_block_size=sparse_q_block_size,
389+ sparse_kv_block_size=sparse_kv_block_size,
390+ dtype=dtype,
391+ num_cube_core=num_cube_core,
392+ mode=FlexMode.FWD,
393+ )
394+ return generator.generate_configs()
395+ 
396+ 
397+def generate_bwd_configs(
398+ query_shape: tuple,
399+ key_shape: tuple,
400+ sparse_q_block_size: int,
401+ sparse_kv_block_size: int,
402+ dtype: torch.dtype,
403+ num_cube_core: int,
404+) -> list[dict]:
405+ """
406+ Convenience function to generate backward configs.
407+ 
408+ Args:
409+ query_shape: Shape of query tensor
410+ key_shape: Shape of key tensor
411+ sparse_q_block_size: SPARSE_Q_BLOCK_SIZE
412+ sparse_kv_block_size: SPARSE_KV_BLOCK_SIZE
413+ dtype: Data type
414+ num_cube_core: Number of AICore
415+ 
416+ Returns:
417+ List of config dictionaries
418+ """
419+ generator = FlexAttentionConfigGenerator(
420+ query_shape=query_shape,
421+ key_shape=key_shape,
422+ sparse_q_block_size=sparse_q_block_size,
423+ sparse_kv_block_size=sparse_kv_block_size,
424+ dtype=dtype,
425+ num_cube_core=num_cube_core,
426+ mode=FlexMode.BWD,
427+ )
428+ return generator.generate_configs()
429+ 
430+ 
431+def prefer_max_tiling_without_benchmark() -> bool:
432+ return (
433+ npu_config.flex_attention.use_config_generator
434+ and not getattr(inductor_config, "max_autotune", False)
435+ and not getattr(inductor_config, "max_autotune_gemm", False)
436+ and not getattr(npu_config, "aggresive_autotune", False)
437+ )
438+ 
439+ 
440+def _sort_fwd_candidate_configs_for_nobench(configs: list[dict]) -> list[dict]:
441+ return sorted(
442+ configs,
443+ key=lambda cfg: (
444+ int(cfg.get("BLOCK_M", 0)) * int(cfg.get("BLOCK_N", 0)),
445+ int(cfg.get("BLOCK_M", 0)),
446+ int(cfg.get("BLOCK_N", 0)),
447+ ),
448+ reverse=True,
449+ )
450+ 
451+ 
452+_FWD_MASK_IN_TILING_ORDER = (
453+ (128, 128),
454+ (128, 64),
455+ (64, 128),
456+ (64, 64),
457+ (32, 32),
458+)
459+ 
460+ 
461+def _build_fwd_mask_in_candidate_configs(
462+ configs: list[dict],
463+ *,
464+ sparse_q_block_size: int,
465+ sparse_kv_block_size: int,
466+) -> list[dict]:
467+ template = configs[0].copy() if configs else {
468+ "num_warps": 4,
469+ "num_stages": 1,
470+ }
471+ ordered_configs = []
472+ 
473+ for block_m, block_n in _FWD_MASK_IN_TILING_ORDER:
474+ if (
475+ block_m > sparse_q_block_size
476+ or block_n > sparse_kv_block_size
477+ or sparse_q_block_size % block_m != 0
478+ or sparse_kv_block_size % block_n != 0
479+ ):
480+ continue
481+ cfg = template.copy()
482+ cfg.update(
483+ {
484+ "BLOCK_M": block_m,
485+ "BLOCK_N": block_n,
486+ "num_warps": cfg.get("num_warps", 4),
487+ "num_stages": 1,
488+ }
489+ )
490+ ordered_configs.append(cfg)
491+ 
492+ if ordered_configs:
493+ return ordered_configs
494+ 
495+ return [
496+ cfg
497+ for cfg in configs
498+ if (
499+ sparse_q_block_size % int(cfg["BLOCK_M"]) == 0
500+ and sparse_kv_block_size % int(cfg["BLOCK_N"]) == 0
501+ )
502+ ]
503+ 
504+ 
505+def _sort_sparse_mask_candidate_configs_for_nobench(
506+ configs: list[dict[str, int]],
507+) -> list[dict[str, int]]:
508+ return sorted(
509+ configs,
510+ key=lambda cfg: (
511+ int(cfg["MASK_BLOCK_M"]) * int(cfg["MASK_BLOCK_N"]),
512+ int(cfg["MASK_BLOCK_M"]),
513+ int(cfg["MASK_BLOCK_N"]),
514+ ),
515+ reverse=True,
516+ )
517+ 
518+ 
519+def _get_default_fwd_config(dtype: torch.dtype, head_dim: int) -> dict:
520+ head_dim = int(head_dim)
521+ config = {
522+ "num_warps": 4,
523+ "num_stages": 3,
524+ }
525+ 
526+ if head_dim <= 256:
527+ config["BLOCK_M"] = 64
528+ config["BLOCK_N"] = 64
529+ elif dtype == torch.float32:
530+ config["BLOCK_M"] = 32
531+ config["BLOCK_N"] = 16
532+ else:
533+ config["BLOCK_M"] = 32
534+ config["BLOCK_N"] = 32
535+ 
536+ return config
537+ 
538+ 
539+def _tune_npu_params(configs: list[dict]) -> list[dict]:
540+ enable = npu_config.flex_attention.enable_npu_optimization
541+ if enable:
542+ npu_params = npu_config.flex_attention.get_npu_compile_hint_params()
543+ npu_config.log.info(
544+ "[flex_attention] NPU compile hint enabled with parameters: %s",
545+ npu_params,
546+ )
547+ log.debug("npu_params: %s", npu_params)
548+ 
549+ return [
550+ npu_config.apply_flex_attention_npu_params(config, enable=enable)
551+ for config in configs
552+ ]
553+ 
554+ 
555+def _build_single_fwd_config(
556+ dtype: torch.dtype,
557+ head_dim: int,
558+ sparse_q_block_size: int,
559+ sparse_kv_block_size: int,
560+) -> list[dict]:
561+ config = _get_default_fwd_config(dtype, head_dim)
562+ if sparse_q_block_size % config["BLOCK_M"] != 0:
563+ config["BLOCK_M"] = int(sparse_q_block_size)
564+ if sparse_kv_block_size % config["BLOCK_N"] != 0:
565+ config["BLOCK_N"] = int(sparse_kv_block_size)
566+ return _tune_npu_params([config])
567+ 
568+ 
569+def get_bwd_dq_compile_options() -> dict:
570+ return npu_config.flex_attention.get_bwd_dq_compile_options()
571+ 
572+ 
573+def get_bwd_dkdv_compile_options() -> dict:
574+ return npu_config.flex_attention.get_bwd_dkdv_compile_options()
575+ 
576+ 
577+def generate_fwd_candidate_configs(
578+ query_shape: tuple,
579+ key_shape: tuple,
580+ dtype: torch.dtype,
581+ sparse_q_block_size: int,
582+ sparse_kv_block_size: int,
583+ num_cube_core: int,
584+ head_dim: Optional[int] = None,
585+ mask_out: bool = True,
586+) -> list[dict]:
587+ """
588+ Generate candidate configs for forward flex attention.
589+ 
590+ This wrapper owns the generator/fallback policy so the lowering file only
591+ needs to pass ordinary Python values extracted from IR nodes.
592+ """
593+ if npu_config.flex_attention.use_config_generator:
594+ configs = generate_fwd_configs(
595+ query_shape=query_shape,
596+ key_shape=key_shape,
597+ sparse_q_block_size=sparse_q_block_size,
598+ sparse_kv_block_size=sparse_kv_block_size,
599+ dtype=dtype,
600+ num_cube_core=num_cube_core,
601+ )
602+ if prefer_max_tiling_without_benchmark():
603+ configs = _sort_fwd_candidate_configs_for_nobench(configs)
604+ if not mask_out:
605+ configs = _build_fwd_mask_in_candidate_configs(
606+ configs,
607+ sparse_q_block_size=sparse_q_block_size,
608+ sparse_kv_block_size=sparse_kv_block_size,
609+ )
610+ return configs
611+ 
612+ if head_dim is None:
613+ head_dim = int(query_shape[-1])
614+ configs = _build_single_fwd_config(
615+ dtype=dtype,
616+ head_dim=head_dim,
617+ sparse_q_block_size=sparse_q_block_size,
618+ sparse_kv_block_size=sparse_kv_block_size,
619+ )
620+ if not mask_out:
621+ configs = _build_fwd_mask_in_candidate_configs(
622+ configs,
623+ sparse_q_block_size=sparse_q_block_size,
624+ sparse_kv_block_size=sparse_kv_block_size,
625+ )
626+ return configs
627+ 
628+ 
629+def _flex_attention_sparse_mask_block_candidates(sparse_block_size: int) -> list[int]:
630+ sparse_block_size = int(sparse_block_size)
631+ if sparse_block_size <= 0:
632+ raise ValueError(f"sparse block size must be positive, got {sparse_block_size}")
633+ 
634+ min_mask_block = min(16, sparse_block_size)
635+ candidates = []
636+ mask_block = sparse_block_size
637+ while mask_block >= min_mask_block:
638+ candidates.append(mask_block)
639+ mask_block //= 2
640+ 
641+ for fallback_mask_block in (64, 32, 16):
642+ if fallback_mask_block <= sparse_block_size:
643+ candidates.append(fallback_mask_block)
644+ 
645+ unique_candidates = []
646+ seen = set()
647+ for mask_block in sorted(candidates, reverse=True):
648+ if mask_block in seen:
649+ continue
650+ if sparse_block_size % mask_block != 0:
651+ continue
652+ seen.add(mask_block)
653+ unique_candidates.append(mask_block)
654+ 
655+ return unique_candidates
656+ 
657+ 
658+def _flex_attention_sparse_mask_tiling_configs(
659+ sparse_q_block_size: int,
660+ sparse_kv_block_size: int,
661+) -> list[dict[str, int]]:
662+ sparse_q_block_size = int(sparse_q_block_size)
663+ sparse_kv_block_size = int(sparse_kv_block_size)
664+ if sparse_q_block_size <= 0:
665+ raise ValueError(
666+ f"SPARSE_Q_BLOCK_SIZE must be positive, got {sparse_q_block_size}"
667+ )
668+ if sparse_kv_block_size <= 0:
669+ raise ValueError(
670+ f"SPARSE_KV_BLOCK_SIZE must be positive, got {sparse_kv_block_size}"
671+ )
672+ 
673+ mask_block_m_candidates = _flex_attention_sparse_mask_block_candidates(
674+ sparse_q_block_size
675+ )
676+ mask_block_n_candidates = _flex_attention_sparse_mask_block_candidates(
677+ sparse_kv_block_size
678+ )
679+ 
680+ configs = []
681+ seen = set()
682+ candidate_pairs = (
683+ (mask_block_m, mask_block_n)
684+ for mask_block_m in mask_block_m_candidates
685+ for mask_block_n in mask_block_n_candidates
686+ )
687+ 
688+ for mask_block_m, mask_block_n in candidate_pairs:
689+ if (mask_block_m, mask_block_n) in seen:
690+ continue
691+ seen.add((mask_block_m, mask_block_n))
692+ configs.append(
693+ {
694+ "MASK_BLOCK_M": mask_block_m,
695+ "MASK_BLOCK_N": mask_block_n,
696+ "NUM_Q_SUB_BLOCKS": sparse_q_block_size // mask_block_m,
697+ "NUM_KV_SUB_BLOCKS": sparse_kv_block_size // mask_block_n,
698+ "num_warps": 4,
699+ "num_stages": 1,
700+ }
701+ )
702+ 
703+ return configs
704+ 
705+ 
706+def _get_default_sparse_mask_tiling_config(
707+ sparse_q_block_size: int,
708+ sparse_kv_block_size: int,
709+) -> dict[str, int]:
710+ sparse_q_block_size = int(sparse_q_block_size)
711+ sparse_kv_block_size = int(sparse_kv_block_size)
712+ return {
713+ "MASK_BLOCK_M": sparse_q_block_size,
714+ "MASK_BLOCK_N": sparse_kv_block_size,
715+ "NUM_Q_SUB_BLOCKS": 1,
716+ "NUM_KV_SUB_BLOCKS": 1,
717+ "num_warps": 4,
718+ "num_stages": 1,
719+ }
720+ 
721+ 
722+def build_sparse_mask_candidate_configs(
723+ sparse_q_block_size: int,
724+ sparse_kv_block_size: int,
725+) -> list[dict[str, int]]:
726+ """Generate sparse mask materialize kernel tiling candidates."""
727+ if npu_config.flex_attention.use_config_generator:
728+ configs = _flex_attention_sparse_mask_tiling_configs(
729+ sparse_q_block_size,
730+ sparse_kv_block_size,
731+ )
732+ if prefer_max_tiling_without_benchmark():
733+ configs = _sort_sparse_mask_candidate_configs_for_nobench(configs)
734+ return configs
735+ return [
736+ _get_default_sparse_mask_tiling_config(
737+ sparse_q_block_size,
738+ sparse_kv_block_size,
739+ )
740+ ]
741+ 
742+ 
743+def split_attention_block_n_candidates(
744+ base_block_n: int,
745+ min_block_n: int = 64,
746+) -> list[int]:
747+ base_block_n = int(base_block_n)
748+ min_block_n = int(min_block_n)
749+ if base_block_n <= 0:
750+ raise ValueError(f"base_block_n must be positive, got {base_block_n}")
751+ if min_block_n <= 0:
752+ raise ValueError(f"min_block_n must be positive, got {min_block_n}")
753+ 
754+ candidates: list[int] = []
755+ current = base_block_n
756+ while current >= min_block_n:
757+ if base_block_n % current == 0:
758+ candidates.append(current)
759+ current //= 2
760+ 
761+ if not candidates:
762+ candidates.append(base_block_n)
763+ return candidates
764+ 
765+ 
766+def _sparse_mask_attention_tile_mix_loop(block_n: int) -> int:
767+ block_n = int(block_n)
768+ if block_n >= 512:
769+ return 4
770+ if block_n >= 256:
771+ return 2
772+ if block_n >= 128:
773+ return 1
774+ return 0
775+ 
776+ 
777+def _sparse_mask_attention_cvpipeline_options(
778+ block_n: int,
779+ *,
780+ enabled: bool,
781+ enable_compile_hint: bool = False,
782+) -> dict[str, Union[int, bool, str]]:
783+ tile_mix_loop = _sparse_mask_attention_tile_mix_loop(block_n) if enabled else 0
784+ return npu_config.flex_attention.get_sparse_mask_cvpipeline_compile_options(
785+ enabled=enabled,
786+ tile_mix_loop=tile_mix_loop,
787+ enable_compile_hint=enable_compile_hint,
788+ )
789+ 
790+ 
791+def sparse_mask_attention_cvpipeline_config_variants(
792+ base_options: dict,
793+ *,
794+ block_n: int,
795+ enable_compile_hint: bool = False,
796+) -> list[dict]:
797+ variants = []
798+ for enabled in (True, False):
799+ variant = base_options.copy()
800+ variant.update(
801+ _sparse_mask_attention_cvpipeline_options(
802+ block_n,
803+ enabled=enabled,
804+ enable_compile_hint=enable_compile_hint,
805+ )
806+ )
807+ variants.append(variant)
808+ return variants
809+ 
810+ 
811+def is_bwd_config_compatible(
812+ cfg: dict,
813+ sparse_q_block_size: int,
814+ sparse_kv_block_size: int,
815+) -> bool:
816+ block_m1 = cfg["BLOCK_M1"]
817+ block_n1 = cfg["BLOCK_N1"]
818+ block_m2 = cfg["BLOCK_M2"]
819+ block_n2 = cfg["BLOCK_N2"]
820+ return (
821+ sparse_q_block_size % block_m1 == 0
822+ and sparse_kv_block_size % block_n1 == 0
823+ and sparse_q_block_size % block_m2 == 0
824+ and sparse_kv_block_size % block_n2 == 0
825+ )
826+ 
827+ 
828+def _convert_bwd_config_to_fused_mask_out_config(cfg: dict) -> dict:
829+ converted_cfg = {
830+ "BLOCK_M1": cfg["BLOCK_M"],
831+ "BLOCK_N1": cfg["BLOCK_N"],
832+ "BLOCK_M2": cfg["BLOCK_N"],
833+ "BLOCK_N2": cfg["BLOCK_M"],
834+ "num_warps": cfg["num_warps"],
835+ "num_stages": cfg["num_stages"],
836+ }
837+ for key, value in cfg.items():
838+ if key not in ("BLOCK_M", "BLOCK_N", "num_warps", "num_stages"):
839+ converted_cfg[key] = value
840+ return converted_cfg
841+ 
842+ 
843+def _start_bwd_mask_out_from_128x128_configs(
844+ configs: list[dict],
845+ *,
846+ sparse_q_block_size: int,
847+ sparse_kv_block_size: int,
848+) -> list[dict]:
849+ sparse_q_block_size = int(sparse_q_block_size)
850+ sparse_kv_block_size = int(sparse_kv_block_size)
851+ max_block = min(128, sparse_q_block_size, sparse_kv_block_size)
852+ 
853+ preferred_blocks = []
854+ block = max_block
855+ min_block = min(16, max_block)
856+ while block >= min_block:
857+ if sparse_q_block_size % block == 0 and sparse_kv_block_size % block == 0:
858+ preferred_blocks.append(block)
859+ block //= 2
860+ if not preferred_blocks:
861+ preferred_blocks.append(max_block)
862+ 
863+ template = configs[0].copy() if configs else {"num_warps": 4, "num_stages": 1}
864+ ordered_configs = []
865+ seen_configs = set()
866+ tiling_keys = ("BLOCK_M1", "BLOCK_N1", "BLOCK_M2", "BLOCK_N2")
867+ 
868+ for block in preferred_blocks:
869+ cfg = template.copy()
870+ cfg.update(
871+ {
872+ "BLOCK_M1": block,
873+ "BLOCK_N1": block,
874+ "BLOCK_M2": block,
875+ "BLOCK_N2": block,
876+ "num_warps": cfg.get("num_warps", 4),
877+ "num_stages": cfg.get("num_stages", 1),
878+ }
879+ )
880+ config_key = tuple(cfg.get(key) for key in tiling_keys)
881+ if config_key in seen_configs:
882+ continue
883+ seen_configs.add(config_key)
884+ ordered_configs.append(cfg)
885+ 
886+ for cfg in configs:
887+ config_key = tuple(cfg.get(key) for key in tiling_keys)
888+ if config_key in seen_configs:
889+ continue
890+ seen_configs.add(config_key)
891+ ordered_configs.append(cfg)
892+ return ordered_configs
893+ 
894+ 
895+def generate_bwd_fused_mask_out_candidate_configs(
896+ query_shape: tuple,
897+ key_shape: tuple,
898+ sparse_q_block_size: int,
899+ sparse_kv_block_size: int,
900+ dtype: torch.dtype,
901+ num_cube_core: int,
902+) -> list[dict]:
903+ """
904+ Generate candidate configs for the fused compact sparse mask-out backward path.
905+ 
906+ The fused backward kernel uses the split backward tiling names but runs as a
907+ single compact sparse mask-out template. Keep 128x128 square tiling first so
908+ the generated output_code remains aligned with the verified path.
909+ """
910+ base_configs = generate_bwd_configs(
911+ query_shape=query_shape,
912+ key_shape=key_shape,
913+ sparse_q_block_size=sparse_q_block_size,
914+ sparse_kv_block_size=sparse_kv_block_size,
915+ dtype=dtype,
916+ num_cube_core=num_cube_core,
917+ )
918+ configs = [
919+ _convert_bwd_config_to_fused_mask_out_config(cfg)
920+ for cfg in base_configs
921+ ]
922+ return _start_bwd_mask_out_from_128x128_configs(
923+ configs,
924+ sparse_q_block_size=sparse_q_block_size,
925+ sparse_kv_block_size=sparse_kv_block_size,
926+ )
927+ 
928+ 
929+def generate_bwd_split_mask_out_candidate_configs(
930+ query_shape: tuple,
931+ key_shape: tuple,
932+ sparse_q_block_size: int,
933+ sparse_kv_block_size: int,
934+ dtype: torch.dtype,
935+ num_cube_core: int,
936+) -> list[dict]:
937+ """Generate candidate configs for split DQ and DKDV backward mask-out kernels."""
938+ return generate_bwd_fused_mask_out_candidate_configs(
939+ query_shape=query_shape,
940+ key_shape=key_shape,
941+ sparse_q_block_size=sparse_q_block_size,
942+ sparse_kv_block_size=sparse_kv_block_size,
943+ dtype=dtype,
944+ num_cube_core=num_cube_core,
945+ )
946+ 
947+ 
948+def validate_benchmark_config() -> None:
949+ """
950+ Validate benchmark configuration before autotuning.
951+ 
952+ This function checks that required configurations are enabled for
953+ NPU optimized benchmark.
954+ 
955+ Note: This function now only warns instead of raising errors to avoid
956+ blocking execution. The actual benchmark will use fallback methods if
957+ configurations are not optimal.
958+ """
959+ aggresive_autotune = getattr(npu_config, 'aggresive_autotune', False)
960+ max_autotune = getattr(inductor_config, 'max_autotune', False)
961+ 
962+ if not aggresive_autotune:
963+ log.warning(
964+ "aggresive_autotune is False. NPU optimized benchmark is disabled. "
965+ "For optimal performance, set INDUCTOR_ASCEND_AGGRESSIVE_AUTOTUNE=1 environment variable. "
966+ "Continuing with fallback benchmark method."
967+ )
968+ 
969+ if not max_autotune:
970+ log.warning(
971+ "max_autotune is False, only default config will be used. "
972+ "Set TORCHINDUCTOR_MAX_AUTOTUNE=1 for multi-config autotuning."
973+ )
Atorch_npu/_inductor/kernel/flex_attention_metadata.py+1051-0
@@ -0,0 +1,1051 @@
1+from __future__ import annotations
2+ 
3+import logging
4+from typing import Any
5+ 
6+import torch
7+ 
8+from torch_npu._inductor import config as npu_config
9+ 
10+log = npu_config.log
11+ 
12+ 
13+def _metadata_auto_infer_enabled() -> bool:
14+ """Return whether metadata auto inference is enabled by config."""
15+ flex_attention_config = getattr(npu_config, "flex_attention", None)
16+ if flex_attention_config is None:
17+ return True
18+ return getattr(flex_attention_config, "metadata_auto_infer", True)
19+ 
20+ 
21+def _try_unwrap_tensor(value: Any) -> torch.Tensor | None:
22+ """Extract a torch.Tensor from wrapper objects that expose a data attribute."""
23+ if isinstance(value, torch.Tensor):
24+ return value
25+ 
26+ data_value = getattr(value, "data", None)
27+ if data_value is value or data_value is None:
28+ return None
29+ 
30+ return _try_unwrap_tensor(data_value)
31+ 
32+ 
33+_SPARSE_MASK_COMPACT_OPTION_KEYS = (
34+ "SPARSE_MASK_MAX_NORMAL_BLOCKS",
35+ "SPARSE_MASK_HEAD_SHARED",
36+ "SPARSE_MASK_HQ",
37+ "HAS_FULL_BLOCKS",
38+)
39+_BLOCK_SPARSE_SAFETY_OPTION_KEYS = (
40+ "NPU_ROWS_GUARANTEED_SAFE",
41+ "NPU_BLOCKS_ARE_CONTIGUOUS",
42+)
43+ 
44+ 
45+def _to_cpu_int_tensor(value: Any) -> torch.Tensor | None:
46+ """Best-effort conversion of a tensor-like value to a CPU int64 tensor."""
47+ tensor = _try_unwrap_tensor(value)
48+ if tensor is None:
49+ return None
50+ try:
51+ return tensor.detach().to("cpu", dtype=torch.int64)
52+ except Exception:
53+ return None
54+ 
55+ 
56+def _heads_share_used_block_entries(
57+ num_blocks: torch.Tensor | None, indices: torch.Tensor | None
58+) -> bool:
59+ """
60+ Return True when all heads have identical used sparse-block entries.
61+ 
62+ Only the valid prefix of each row is compared because entries after
63+ num_blocks[b, h, q] are undefined padding in BlockMask.
64+ """
65+ if num_blocks is None or indices is None:
66+ return False
67+ if num_blocks.ndim < 3 or indices.ndim < 4:
68+ return True
69+ 
70+ batch = int(num_blocks.shape[0])
71+ heads = int(num_blocks.shape[1])
72+ rows = int(num_blocks.shape[2])
73+ capacity = int(indices.shape[-1])
74+ if heads <= 1:
75+ return True
76+ 
77+ for b_idx in range(batch):
78+ for q_idx in range(rows):
79+ ref_count = int(num_blocks[b_idx, 0, q_idx].item())
80+ ref_count = max(0, min(ref_count, capacity))
81+ ref_indices = indices[b_idx, 0, q_idx, :ref_count]
82+ for h_idx in range(1, heads):
83+ cur_count = int(num_blocks[b_idx, h_idx, q_idx].item())
84+ cur_count = max(0, min(cur_count, capacity))
85+ if cur_count != ref_count:
86+ return False
87+ cur_indices = indices[b_idx, h_idx, q_idx, :cur_count]
88+ if not torch.equal(cur_indices, ref_indices):
89+ return False
90+ return True
91+ 
92+ 
93+def _infer_sparse_mask_compact_options(block_mask: Any) -> dict[str, Any]:
94+ """
95+ Infer compact sparse-mask materialization options from eager BlockMask metadata.
96+ 
97+ The options specialize the temporary mask buffer only. KV traversal still uses
98+ the original BlockMask tensors, so failures to inspect simply return no options
99+ and leave the existing uncompressed shape in place.
100+ """
101+ if block_mask is None:
102+ return {}
103+ 
104+ kv_num_blocks = _to_cpu_int_tensor(getattr(block_mask, "kv_num_blocks", None))
105+ kv_indices = _to_cpu_int_tensor(getattr(block_mask, "kv_indices", None))
106+ if kv_num_blocks is None or kv_indices is None:
107+ return {}
108+ if kv_num_blocks.numel() == 0 or kv_indices.ndim < 4:
109+ return {}
110+ 
111+ metadata_heads = int(kv_num_blocks.shape[1]) if kv_num_blocks.ndim >= 2 else 1
112+ capacity = int(kv_indices.shape[-1])
113+ max_normal_blocks = int(kv_num_blocks.max().item())
114+ max_normal_blocks = max(1, min(max_normal_blocks, capacity))
115+ 
116+ partial_heads_shared = _heads_share_used_block_entries(kv_num_blocks, kv_indices)
117+ 
118+ full_kv_num_blocks = _to_cpu_int_tensor(
119+ getattr(block_mask, "full_kv_num_blocks", None)
120+ )
121+ full_kv_indices = _to_cpu_int_tensor(getattr(block_mask, "full_kv_indices", None))
122+ has_full_blocks = bool(
123+ full_kv_num_blocks is not None
124+ and full_kv_num_blocks.numel() > 0
125+ and full_kv_num_blocks.max().item() > 0
126+ )
127+ 
128+ if full_kv_num_blocks is None and full_kv_indices is None:
129+ full_heads_shared = True
130+ elif full_kv_num_blocks is None or full_kv_indices is None:
131+ full_heads_shared = False
132+ else:
133+ full_heads_shared = _heads_share_used_block_entries(
134+ full_kv_num_blocks, full_kv_indices
135+ )
136+ 
137+ head_shared = partial_heads_shared and full_heads_shared
138+ options = {
139+ "SPARSE_MASK_MAX_NORMAL_BLOCKS": max_normal_blocks,
140+ "SPARSE_MASK_HEAD_SHARED": bool(head_shared),
141+ "SPARSE_MASK_HQ": 1 if head_shared else metadata_heads,
142+ }
143+ options["HAS_FULL_BLOCKS"] = has_full_blocks
144+ return options
145+ 
146+ 
147+def _precomputed_sparse_mask_compact_options(block_mask: Any) -> dict[str, Any]:
148+ """Read compact options cached on a BlockMask by the NPU patch, if present."""
149+ if block_mask is None:
150+ return {}
151+ options = getattr(block_mask, "_npu_flex_attention_kernel_options", None)
152+ if not isinstance(options, dict):
153+ return {}
154+ return {
155+ key: options[key]
156+ for key in _SPARSE_MASK_COMPACT_OPTION_KEYS
157+ if key in options
158+ }
159+ 
160+ 
161+def _precomputed_block_sparse_safety_options(block_mask: Any) -> dict[str, Any]:
162+ """Read exact safety diagnostics cached on a BlockMask by the NPU patch."""
163+ if block_mask is None:
164+ return {}
165+ options = getattr(block_mask, "_npu_flex_attention_kernel_options", None)
166+ if not isinstance(options, dict):
167+ return {}
168+ return {
169+ key: options[key]
170+ for key in _BLOCK_SPARSE_SAFETY_OPTION_KEYS
171+ if key in options
172+ }
173+ 
174+ 
175+def _apply_sparse_mask_compact_options(
176+ kernel_options: dict[str, Any],
177+ block_mask: Any,
178+ context: str,
179+ *,
180+ allow_tensor_analysis: bool,
181+) -> dict[str, Any]:
182+ """Merge cached or freshly inferred sparse-mask compact options."""
183+ updated = dict(kernel_options)
184+ compact_options = _precomputed_sparse_mask_compact_options(block_mask)
185+ missing_compact_options = any(
186+ key not in updated for key in _SPARSE_MASK_COMPACT_OPTION_KEYS
187+ )
188+ if allow_tensor_analysis and missing_compact_options:
189+ compact_options = {
190+ **_infer_sparse_mask_compact_options(block_mask),
191+ **compact_options,
192+ }
193+ 
194+ for key, value in compact_options.items():
195+ updated.setdefault(key, value)
196+ 
197+ if compact_options and log.isEnabledFor(logging.INFO):
198+ log.info(
199+ "[flex_attention][%s] sparse_mask_compact_options=%s final_hq=%s final_max_blocks=%s",
200+ context,
201+ compact_options,
202+ updated.get("SPARSE_MASK_HQ", "<unset>"),
203+ updated.get("SPARSE_MASK_MAX_NORMAL_BLOCKS", "<unset>"),
204+ )
205+ return updated
206+ 
207+ 
208+def _apply_precomputed_block_sparse_safety_options(
209+ kernel_options: dict[str, Any],
210+ block_mask: Any,
211+ context: str,
212+) -> dict[str, Any]:
213+ """Forward cached contiguity diagnostics while keeping row safety conservative."""
214+ updated = dict(kernel_options)
215+ safety_options = _precomputed_block_sparse_safety_options(block_mask)
216+ if "NPU_BLOCKS_ARE_CONTIGUOUS" in safety_options:
217+ updated.setdefault(
218+ "BLOCKS_ARE_CONTIGUOUS",
219+ bool(safety_options["NPU_BLOCKS_ARE_CONTIGUOUS"]),
220+ )
221+ if safety_options and log.isEnabledFor(logging.INFO):
222+ log.info(
223+ "[flex_attention][%s] cached_block_sparse_safety_diagnostics=%s "
224+ "forwarded_ROWS_GUARANTEED_SAFE=%s forwarded_BLOCKS_ARE_CONTIGUOUS=%s",
225+ context,
226+ safety_options,
227+ updated.get("ROWS_GUARANTEED_SAFE", "<unset>"),
228+ updated.get("BLOCKS_ARE_CONTIGUOUS", "<unset>"),
229+ )
230+ return updated
231+ 
232+ 
233+def _apply_disabled_metadata_defaults(kernel_options: dict[str, Any]) -> dict[str, Any]:
234+ """Use conservative kernel defaults when metadata auto inference is disabled."""
235+ updated = dict(kernel_options)
236+ updated.setdefault("ROWS_GUARANTEED_SAFE", False)
237+ updated.setdefault("BLOCKS_ARE_CONTIGUOUS", False)
238+ return updated
239+ 
240+ 
241+def _normalize_block_rows(kv_num_blocks: torch.Tensor, kv_indices: torch.Tensor) -> tuple[list[int], list[list[int]]]:
242+ """Convert block-sparse row counts and indices into CPU Python lists."""
243+ counts = kv_num_blocks.to("cpu", dtype=torch.int64).reshape(-1)
244+ indices = kv_indices.to("cpu", dtype=torch.int64).reshape(-1, kv_indices.shape[-1])
245+ row_counts = counts.tolist()
246+ row_indices = [indices[i, :count].tolist() for i, count in enumerate(row_counts)]
247+ return row_counts, row_indices
248+ 
249+ 
250+def _infer_blocks_are_contiguous_from_tensors(kv_num_blocks: Any, kv_indices: Any) -> bool | str:
251+ counts_tensor = _try_unwrap_tensor(kv_num_blocks)
252+ indices_tensor = _try_unwrap_tensor(kv_indices)
253+ if counts_tensor is None or indices_tensor is None:
254+ return "Unknown"
255+ 
256+ _, row_indices = _normalize_block_rows(counts_tensor, indices_tensor)
257+ for values in row_indices:
258+ if len(values) <= 1:
259+ continue
260+ if any((right - left) != 1 for left, right in zip(values, values[1:])):
261+ return False
262+ return True
263+ 
264+ 
265+def _layer2_fast_prefilter(block_mask: Any) -> tuple[bool, str]:
266+ """
267+ Perform Layer 2 safety check using to_dense() fast pre-filter.
268+ 
269+ This checks if every Sparse Q-block has at least one valid KV at the
270+ Sparse Block level (32x32 elements per cell). Cost: ~1ms per mask.
271+ 
272+ Args:
273+ block_mask: A PyTorch BlockMask object with to_dense() method
274+ 
275+ Returns:
276+ tuple[bool, str]: (is_safe, detail_message)
277+ - is_safe: True if all rows have >=1 valid cell, False otherwise
278+ - detail_message: Human-readable result for logging
279+ 
280+ Note:
281+ L2 has 40% false positive rate (granularity trap)!
282+ Must be combined with L3 whitelist for production safety.
283+ See v4.0 design doc: Level1_Level2_完整设计文档.md §2.1
284+ """
285+ try:
286+ import time as _time
287+ _t_start = _time.time()
288+ if log.isEnabledFor(logging.INFO):
289+ log.info(
290+ "[meta][L2] start block_mask_type=%s",
291+ type(block_mask).__name__,
292+ )
293+ 
294+ dense = block_mask.to_dense()
295+ dense_device = getattr(dense, "device", "<unknown>")
296+ dense_dtype = getattr(dense, "dtype", "<unknown>")
297+ dense_shape = tuple(dense.shape) if hasattr(dense, "shape") else "<unknown>"
298+ if log.isEnabledFor(logging.INFO):
299+ log.info(
300+ "[meta][L2] dense_ready shape=%s dense_device=%s dtype=%s elapsed=%.2fms",
301+ dense_shape,
302+ dense_device,
303+ dense_dtype,
304+ (_time.time() - _t_start) * 1000,
305+ )
306+ 
307+ b, h = 0, 0
308+ dense_2d = dense[b, h]
309+ 
310+ row_has_valid = dense_2d.any(dim=-1)
311+ l2_safe = row_has_valid.all().item()
312+ 
313+ _elapsed = (_time.time() - _t_start) * 1000
314+ 
315+ if l2_safe:
316+ valid_rows = row_has_valid.sum().item()
317+ total_rows = len(row_has_valid)
318+ detail = f"L2 SAFE ({valid_rows}/{total_rows} rows valid, {_elapsed:.2f}ms)"
319+ if log.isEnabledFor(logging.INFO):
320+ log.info(
321+ "[meta][L2] done safe=True valid_rows=%s total_rows=%s "
322+ "dense_device=%s elapsed=%.2fms",
323+ valid_rows,
324+ total_rows,
325+ dense_device,
326+ _elapsed,
327+ )
328+ return True, detail
329+ else:
330+ unsafe_rows = (~row_has_valid).nonzero(as_tuple=True)[0].tolist()
331+ detail = f"L2 UNSAFE ({len(unsafe_rows)} empty rows: {unsafe_rows[:10]}, {_elapsed:.2f}ms)"
332+ if log.isEnabledFor(logging.INFO):
333+ log.info(
334+ "[meta][L2] done safe=False unsafe_rows=%s total_rows=%s "
335+ "dense_device=%s elapsed=%.2fms",
336+ unsafe_rows[:20],
337+ len(row_has_valid),
338+ dense_device,
339+ _elapsed,
340+ )
341+ return False, detail
342+ 
343+ except Exception as exc:
344+ if log.isEnabledFor(logging.INFO):
345+ log.info( # noqa: G200
346+ "[meta][L2] done safe=False error=%s: %s",
347+ type(exc).__name__,
348+ exc,
349+ )
350+ return False, f"L2 ERROR: {type(exc).__name__}: {exc}"
351+ 
352+ 
353+def _get_critical_positions(seqlen: int, block_size: int = 32, sparse_block_size: int = 128) -> list[int]:
354+ """
355+ Generate critical boundary positions for Element-Level safety sampling.
356+ 
357+ Instead of scanning all N positions (O(N^2)), we sample key boundary
358+ locations where granularity traps are most likely to occur:
359+ - Sparse Block boundaries (multiples of 128)
360+ - Mask Block boundaries (multiples of 32)
361+ - First/last positions of each region
362+ - Segment boundaries if applicable
363+ 
364+ This reduces verification from O(N^2) to O(N) while maintaining high accuracy.
365+ Based on v3.0 empirical validation (526 positions sufficient for seqlen=8192).
366+ 
367+ Args:
368+ seqlen: Total sequence length
369+ block_size: Mask Block size (default 32)
370+ sparse_block_size: Sparse Block size (default 128)
371+ 
372+ Returns:
373+ list[int]: Sorted list of critical query positions to check
374+ """
375+ positions = set()
376+ 
377+ num_sparse_blocks = (seqlen + sparse_block_size - 1) // sparse_block_size
378+ num_blocks_per_sparse = sparse_block_size // block_size
379+ 
380+ for sb in range(num_sparse_blocks):
381+ sb_start = sb * sparse_block_size
382+ sb_end = min(sb_start + sparse_block_size, seqlen)
383+ 
384+ positions.add(sb_start)
385+ positions.add(min(sb_end - 1, seqlen - 1))
386+ 
387+ for bi in range(num_blocks_per_sparse):
388+ b_start = sb_start + bi * block_size
389+ b_end = min(b_start + block_size, seqlen)
390+ if b_start < seqlen:
391+ positions.add(b_start)
392+ if b_end - 1 < seqlen and b_end - 1 >= 0:
393+ positions.add(b_end - 1)
394+ 
395+ mid = (sb_start + sb_end) // 2
396+ if mid < seqlen:
397+ positions.add(mid)
398+ 
399+ for i in range(0, min(10, seqlen)):
400+ positions.add(i)
401+ 
402+ for i in range(max(0, seqlen - 10), seqlen):
403+ positions.add(i)
404+ 
405+ return sorted([p for p in positions if p < seqlen])
406+ 
407+ 
408+def has_any_valid_kv(mask_fn: Any, q_pos: int, seqlen: int, b: int = 0, h: int = 0) -> bool:
409+ """
410+ Check if a specific query position has at least one valid KV position.
411+ 
412+ This is the core Element-Level (1x1) safety check. For a given query,
413+ it scans all possible KVs to find at least one valid (q,k) pair according
414+ to the mask_mod function.
415+ 
416+ Args:
417+ mask_fn: The mask_mod function from create_block_mask()
418+ q_pos: Query position to check (element-level index)
419+ seqlen: Total sequence length
420+ b: Batch index (default 0)
421+ h: Head index (default 0)
422+ 
423+ Returns:
424+ bool: True if query has >=1 valid KV, False otherwise
425+ 
426+ Note:
427+ This is the definitive safety check that eliminates granularity traps.
428+ A return of False means this query would produce NaN in attention computation.
429+ """
430+ try:
431+ device = torch.device("cpu")
432+ b_idx = torch.tensor(b, device=device)
433+ h_idx = torch.tensor(h, device=device)
434+ q_idx = torch.tensor(q_pos, device=device)
435+ for kv_start in range(0, seqlen, 8192):
436+ kv_idx = torch.arange(kv_start, min(kv_start + 8192, seqlen), device=device)
437+ result = mask_fn(b_idx, h_idx, q_idx, kv_idx)
438+ result_tensor = (
439+ result.to(dtype=torch.bool)
440+ if isinstance(result, torch.Tensor)
441+ else torch.as_tensor(result, dtype=torch.bool, device=device)
442+ )
443+ if bool(result_tensor.any().item()):
444+ return True
445+ return False
446+ except Exception:
447+ return False
448+ 
449+ 
450+def _infer_block_mask_seq_lengths(block_mask: Any) -> tuple[int, int] | None:
451+ seq_lengths = getattr(block_mask, "seq_lengths", None)
452+ if (
453+ isinstance(seq_lengths, tuple)
454+ and len(seq_lengths) == 2
455+ and all(isinstance(length, int) for length in seq_lengths)
456+ ):
457+ return seq_lengths
458+ return None
459+ 
460+ 
461+def _infer_block_mask_batch_heads(block_mask: Any, counts: torch.Tensor | None) -> tuple[int, int]:
462+ if counts is not None and counts.ndim >= 3:
463+ return int(counts.shape[0]), int(counts.shape[1])
464+ shape = getattr(block_mask, "shape", None)
465+ if isinstance(shape, tuple) and len(shape) >= 4:
466+ return int(shape[0]), int(shape[1])
467+ return 1, 1
468+ 
469+ 
470+def _verify_rows_have_valid_kv_tensorized(
471+ mask_fn: Any,
472+ *,
473+ batch_size: int,
474+ num_heads: int,
475+ q_len: int,
476+ kv_len: int,
477+ device: torch.device,
478+ q_chunk_size: int = 256,
479+ kv_chunk_size: int = 8192,
480+ max_unsafe: int = 100,
481+) -> tuple[bool, list[str]]:
482+ import time as _time
483+ _t_start = _time.time()
484+ unsafe_locations: list[str] = []
485+ q_chunks_per_head = (q_len + q_chunk_size - 1) // q_chunk_size
486+ kv_chunks_per_q_chunk = (kv_len + kv_chunk_size - 1) // kv_chunk_size
487+ total_q_chunks = batch_size * num_heads * q_chunks_per_head
488+ progress_log_every = max(1, total_q_chunks // 16)
489+ q_chunks_done = 0
490+ result_device_warned = False
491+ 
492+ if log.isEnabledFor(logging.INFO):
493+ log.info(
494+ "[meta][L3] scan_start batch=%d heads=%d q_len=%d kv_len=%d "
495+ "device=%s q_chunk_size=%d kv_chunk_size=%d total_q_chunks=%d "
496+ "kv_chunks_per_q_chunk=%d max_unsafe=%d",
497+ batch_size,
498+ num_heads,
499+ q_len,
500+ kv_len,
501+ device,
502+ q_chunk_size,
503+ kv_chunk_size,
504+ total_q_chunks,
505+ kv_chunks_per_q_chunk,
506+ max_unsafe,
507+ )
508+ 
509+ for b_idx_value in range(batch_size):
510+ b_idx = torch.tensor(b_idx_value, device=device)
511+ for h_idx_value in range(num_heads):
512+ h_idx = torch.tensor(h_idx_value, device=device)
513+ for q_start in range(0, q_len, q_chunk_size):
514+ q_end = min(q_start + q_chunk_size, q_len)
515+ q_idx = torch.arange(q_start, q_end, device=device)[:, None]
516+ row_has_valid = torch.zeros(q_end - q_start, dtype=torch.bool, device=device)
517+ 
518+ for kv_start in range(0, kv_len, kv_chunk_size):
519+ kv_end = min(kv_start + kv_chunk_size, kv_len)
520+ kv_idx = torch.arange(kv_start, kv_end, device=device)[None, :]
521+ result = mask_fn(b_idx, h_idx, q_idx, kv_idx)
522+ result_tensor = (
523+ result.to(dtype=torch.bool)
524+ if isinstance(result, torch.Tensor)
525+ else torch.as_tensor(result, dtype=torch.bool, device=device)
526+ )
527+ result_device = getattr(result_tensor, "device", None)
528+ if (
529+ not result_device_warned
530+ and result_device is not None
531+ and str(result_device) != str(device)
532+ ):
533+ log.warning(
534+ "[meta][L3] mask_mod result device=%s differs from "
535+ "analysis device=%s; metadata verification may not be "
536+ "using the expected accelerated path",
537+ result_device,
538+ device,
539+ )
540+ result_device_warned = True
541+ if result_tensor.ndim == 0:
542+ result_tensor = result_tensor.expand(q_end - q_start, kv_end - kv_start)
543+ elif tuple(result_tensor.shape) != (q_end - q_start, kv_end - kv_start):
544+ result_tensor = torch.broadcast_to(
545+ result_tensor,
546+ (q_end - q_start, kv_end - kv_start),
547+ )
548+ row_has_valid |= result_tensor.any(dim=1)
549+ if bool(row_has_valid.all().item()):
550+ break
551+ 
552+ q_chunks_done += 1
553+ if (
554+ log.isEnabledFor(logging.INFO)
555+ and (
556+ q_chunks_done == 1
557+ or q_chunks_done % progress_log_every == 0
558+ or q_chunks_done == total_q_chunks
559+ )
560+ ):
561+ valid_rows = int(row_has_valid.sum().item())
562+ log.info(
563+ "[meta][L3] progress q_chunk=%d/%d b=%d h=%d "
564+ "q_range=[%d,%d) valid_rows=%d/%d unsafe_seen=%d "
565+ "elapsed=%.2fms device=%s",
566+ q_chunks_done,
567+ total_q_chunks,
568+ b_idx_value,
569+ h_idx_value,
570+ q_start,
571+ q_end,
572+ valid_rows,
573+ q_end - q_start,
574+ len(unsafe_locations),
575+ (_time.time() - _t_start) * 1000,
576+ device,
577+ )
578+ 
579+ if not bool(row_has_valid.all().item()):
580+ bad_rows = (~row_has_valid).nonzero(as_tuple=True)[0].to("cpu").tolist()
581+ for bad_row in bad_rows:
582+ unsafe_locations.append(
583+ f"b={b_idx_value},h={h_idx_value},q={q_start + int(bad_row)}"
584+ )
585+ if len(unsafe_locations) >= max_unsafe:
586+ if log.isEnabledFor(logging.INFO):
587+ log.info(
588+ "[meta][L3] stop max_unsafe=%d reached "
589+ "q_chunk=%d/%d elapsed=%.2fms device=%s",
590+ max_unsafe,
591+ q_chunks_done,
592+ total_q_chunks,
593+ (_time.time() - _t_start) * 1000,
594+ device,
595+ )
596+ return False, unsafe_locations
597+ 
598+ if log.isEnabledFor(logging.INFO):
599+ log.info(
600+ "[meta][L3] done safe=%s checked_q_chunks=%d/%d unsafe_count=%d "
601+ "elapsed=%.2fms device=%s",
602+ len(unsafe_locations) == 0,
603+ q_chunks_done,
604+ total_q_chunks,
605+ len(unsafe_locations),
606+ (_time.time() - _t_start) * 1000,
607+ device,
608+ )
609+ return len(unsafe_locations) == 0, unsafe_locations
610+ 
611+ 
612+def _verify_element_level_safety(block_mask: Any) -> tuple[bool, str]:
613+ """
614+ Perform TRUE Element-Level (1x1) online safety verification.
615+ 
616+ This replaces static whitelist lookup with dynamic runtime analysis.
617+ For each critical query position, it checks whether there exists at least
618+ one valid KV using the actual mask_mod function.
619+ 
620+ Based on v3.0 methodology (verify_all_masks_element_level.py).
621+ Samples ~526 critical boundary positions instead of full O(N^2) scan.
622+ 
623+ Args:
624+ block_mask: A PyTorch BlockMask object with mask_mod attribute
625+ 
626+ Returns:
627+ tuple[bool, str]: (is_safe, detail_message)
628+ - is_safe: True if ALL checked queries have valid KVs
629+ - detail_message: Human-readable result including unsafe count
630+ 
631+ Performance:
632+ - Uses tensorized q_chunk x kv_chunk row reductions.
633+ - Runs on the BlockMask tensor device, so eager NPU BlockMask creation
634+ can perform this check on NPU and cache the result.
635+ """
636+ import time as _time
637+ _t_start = _time.time()
638+ 
639+ try:
640+ mask_mod_fn = getattr(block_mask, 'mask_mod', None)
641+ kv_num_blocks_attr = getattr(block_mask, 'kv_num_blocks', None)
642+ 
643+ if mask_mod_fn is None or kv_num_blocks_attr is None:
644+ elapsed = (_time.time() - _t_start) * 1000
645+ return False, f"L3 ERROR: missing mask_mod/kv_num_blocks ({elapsed:.2f}ms)"
646+ 
647+ counts = _try_unwrap_tensor(kv_num_blocks_attr)
648+ if counts is None:
649+ elapsed = (_time.time() - _t_start) * 1000
650+ return False, f"L3 ERROR: cannot unwrap kv_num_blocks ({elapsed:.2f}ms)"
651+ 
652+ seq_lengths = _infer_block_mask_seq_lengths(block_mask)
653+ if seq_lengths is None:
654+ elapsed = (_time.time() - _t_start) * 1000
655+ return False, f"L3 ERROR: missing seq_lengths ({elapsed:.2f}ms)"
656+ q_len, kv_len = seq_lengths
657+ 
658+ if q_len <= 0 or kv_len <= 0 or q_len > 100000 or kv_len > 100000:
659+ elapsed = (_time.time() - _t_start) * 1000
660+ return False, f"L3 ERROR: invalid seq_lengths={seq_lengths} ({elapsed:.2f}ms)"
661+ 
662+ batch_size, num_heads = _infer_block_mask_batch_heads(block_mask, counts)
663+ if log.isEnabledFor(logging.INFO):
664+ log.info(
665+ "[meta][L3] start block_mask_type=%s batch=%d heads=%d "
666+ "q_len=%d kv_len=%d device=%s counts_shape=%s counts_dtype=%s",
667+ type(block_mask).__name__,
668+ batch_size,
669+ num_heads,
670+ q_len,
671+ kv_len,
672+ counts.device,
673+ tuple(counts.shape),
674+ counts.dtype,
675+ )
676+ if counts.device.type != "npu":
677+ log.warning(
678+ "[meta][L3] running on CPU device=%s; NPU acceleration is not "
679+ "active for metadata row-safety verification",
680+ counts.device,
681+ )
682+ is_safe, unsafe_locations = _verify_rows_have_valid_kv_tensorized(
683+ mask_mod_fn,
684+ batch_size=batch_size,
685+ num_heads=num_heads,
686+ q_len=q_len,
687+ kv_len=kv_len,
688+ device=counts.device,
689+ )
690+ 
691+ elapsed = (_time.time() - _t_start) * 1000
692+ 
693+ if is_safe:
694+ total_rows = batch_size * num_heads * q_len
695+ return True, f"L3 SAFE ({total_rows} rows verified exactly, {elapsed:.2f}ms)"
696+ else:
697+ return (
698+ False,
699+ f"L3 UNSAFE ({len(unsafe_locations)} rows have no valid KV: "
700+ f"{unsafe_locations[:20]}..., {elapsed:.2f}ms)"
701+ )
702+ 
703+ except Exception as exc:
704+ elapsed = (_time.time() - _t_start) * 1000
705+ return False, f"L3 EXCEPTION: {type(exc).__name__}: {exc} ({elapsed:.2f}ms)"
706+ 
707+ 
708+def infer_eager_block_mask_kernel_options(block_mask: Any) -> dict[str, Any]:
709+ """Infer and cache eager BlockMask options before Dynamo graph capture."""
710+ if block_mask is None:
711+ return {}
712+ 
713+ options = _infer_sparse_mask_compact_options(block_mask)
714+ if not _metadata_auto_infer_enabled():
715+ options["NPU_ROWS_GUARANTEED_SAFE"] = False
716+ options["NPU_BLOCKS_ARE_CONTIGUOUS"] = False
717+ return options
718+ 
719+ kv_num_blocks, kv_indices = _extract_block_sparse_tensors(block_mask)
720+ if kv_num_blocks is None or kv_indices is None:
721+ return options
722+ 
723+ rows_value, _ = _verify_element_level_safety(block_mask)
724+ options["NPU_ROWS_GUARANTEED_SAFE"] = rows_value
725+ 
726+ contiguous_value = _infer_blocks_are_contiguous_from_tensors(kv_num_blocks, kv_indices)
727+ if contiguous_value != "Unknown":
728+ options["NPU_BLOCKS_ARE_CONTIGUOUS"] = contiguous_value
729+ return options
730+ 
731+ 
732+def infer_block_sparse_metadata(kv_num_blocks: Any, kv_indices: Any, block_mask: Any = None) -> dict[str, Any]:
733+ """
734+ Infer metadata from block-sparse structures using v4.0 two-layer architecture.
735+ 
736+ Architecture (v4.0):
737+ - Layer 2: Fast Pre-filter via to_dense() (~1ms) - excludes obviously unsafe
738+ - Layer 3: Element-Level whitelist lookup (O(1)) - final safety decision
739+ 
740+ This replaces the previous overly optimistic "Level 0.5" approach that had
741+ 40% false positive rate (causing NaN crashes in production).
742+ 
743+ Args:
744+ kv_num_blocks: Row-level KV block counts tensor
745+ kv_indices: Row-level KV block index lists
746+ block_mask: (NEW) Optional BlockMask object for L2/L3 analysis
747+ 
748+ Returns:
749+ dict with keys:
750+ - rows_guaranteed_safe: bool/"Unknown" (final verdict)
751+ - blocks_are_contiguous: bool/"Unknown"
752+ - is_per_head_heterogeneous: bool
753+ - empty_row_risk_level: "low"/"medium"/"high"
754+ - l2_result: str (Layer 2 analysis detail for debugging)
755+ - l3_result: str (Layer 3 whitelist lookup detail for debugging)
756+ - safety_layer: str ("L2_FAIL"/"L3_WHITELIST"/"L3_CONSERVATIVE"/"LEGACY")
757+ """
758+ if not _metadata_auto_infer_enabled():
759+ return {
760+ "rows_guaranteed_safe": "Unknown",
761+ "blocks_are_contiguous": "Unknown",
762+ "is_per_head_heterogeneous": False,
763+ "empty_row_risk_level": "medium",
764+ "l2_result": "DISABLED_BY_CONFIG",
765+ "l3_result": "DISABLED_BY_CONFIG",
766+ "safety_layer": "DISABLED",
767+ }
768+ 
769+ if block_mask is None or not callable(getattr(block_mask, "to_dense", None)):
770+ log.warning(
771+ "[meta][v4.0] block_mask is missing to_dense(), falling back to legacy Level 0.5 method. "
772+ "This has 40% false positive rate! Pass block_mask for v4.0 safety."
773+ )
774+ counts_tensor = _try_unwrap_tensor(kv_num_blocks)
775+ indices_tensor = _try_unwrap_tensor(kv_indices)
776+ if counts_tensor is None or indices_tensor is None:
777+ return {
778+ "rows_guaranteed_safe": "Unknown",
779+ "blocks_are_contiguous": "Unknown",
780+ "is_per_head_heterogeneous": False,
781+ "empty_row_risk_level": "medium",
782+ "l2_result": "NO_BLOCK_MASK",
783+ "l3_result": "NO_BLOCK_MASK",
784+ "safety_layer": "LEGACY_FALLBACK",
785+ }
786+ 
787+ row_counts, row_indices = _normalize_block_rows(counts_tensor, indices_tensor)
788+ rows_guaranteed_safe_legacy = all(count > 0 for count in row_counts)
789+ blocks_are_contiguous = True
790+ for values in row_indices:
791+ if len(values) <= 1:
792+ continue
793+ if any((right - left) != 1 for left, right in zip(values, values[1:])):
794+ blocks_are_contiguous = False
795+ break
796+ 
797+ headwise_counts = counts_tensor.to("cpu", dtype=torch.int64)
798+ headwise_indices = indices_tensor.to("cpu", dtype=torch.int64)
799+ is_per_head_heterogeneous = not _heads_share_used_block_entries(
800+ headwise_counts, headwise_indices
801+ )
802+ 
803+ return {
804+ "rows_guaranteed_safe": rows_guaranteed_safe_legacy,
805+ "blocks_are_contiguous": blocks_are_contiguous,
806+ "is_per_head_heterogeneous": is_per_head_heterogeneous,
807+ "empty_row_risk_level": "low" if rows_guaranteed_safe_legacy else "high",
808+ "l2_result": "SKIPPED_NO_BLOCK_MASK",
809+ "l3_result": "SKIPPED_NO_BLOCK_MASK",
810+ "safety_layer": "LEGACY_DEPRECATED",
811+ }
812+ 
813+ l2_safe, l2_detail = _layer2_fast_prefilter(block_mask)
814+ 
815+ if not l2_safe:
816+ counts_tensor = _try_unwrap_tensor(kv_num_blocks)
817+ indices_tensor = _try_unwrap_tensor(kv_indices)
818+ if counts_tensor is not None and indices_tensor is not None:
819+ row_counts, row_indices = _normalize_block_rows(counts_tensor, indices_tensor)
820+ blocks_are_contiguous = True
821+ for values in row_indices:
822+ if len(values) <= 1:
823+ continue
824+ if any((right - left) != 1 for left, right in zip(values, values[1:])):
825+ blocks_are_contiguous = False
826+ break
827+ else:
828+ blocks_are_contiguous = "Unknown"
829+ 
830+ return {
831+ "rows_guaranteed_safe": False,
832+ "blocks_are_contiguous": blocks_are_contiguous,
833+ "is_per_head_heterogeneous": False,
834+ "empty_row_risk_level": "high",
835+ "l2_result": l2_detail,
836+ "l3_result": "SKIPPED_L2_FAIL",
837+ "safety_layer": "L2_FAIL_FAST_REJECT",
838+ }
839+ 
840+ l3_safe, l3_detail = _verify_element_level_safety(block_mask)
841+ 
842+ if l3_safe:
843+ counts_tensor = _try_unwrap_tensor(kv_num_blocks)
844+ indices_tensor = _try_unwrap_tensor(kv_indices)
845+ if counts_tensor is not None and indices_tensor is not None:
846+ row_counts, row_indices = _normalize_block_rows(counts_tensor, indices_tensor)
847+ blocks_are_contiguous = True
848+ for values in row_indices:
849+ if len(values) <= 1:
850+ continue
851+ if any((right - left) != 1 for left, right in zip(values, values[1:])):
852+ blocks_are_contiguous = False
853+ break
854+ else:
855+ blocks_are_contiguous = True
856+ 
857+ return {
858+ "rows_guaranteed_safe": True,
859+ "blocks_are_contiguous": blocks_are_contiguous,
860+ "is_per_head_heterogeneous": False,
861+ "empty_row_risk_level": "low",
862+ "l2_result": l2_detail,
863+ "l3_result": f"ELEMENT_VERIFIED_SAFE({l3_detail})",
864+ "safety_layer": "L3_ELEMENT_PASS",
865+ }
866+ else:
867+ counts_tensor = _try_unwrap_tensor(kv_num_blocks)
868+ indices_tensor = _try_unwrap_tensor(kv_indices)
869+ if counts_tensor is not None and indices_tensor is not None:
870+ row_counts, row_indices = _normalize_block_rows(counts_tensor, indices_tensor)
871+ blocks_are_contiguous = True
872+ for values in row_indices:
873+ if len(values) <= 1:
874+ continue
875+ if any((right - left) != 1 for left, right in zip(values, values[1:])):
876+ blocks_are_contiguous = False
877+ break
878+ else:
879+ blocks_are_contiguous = "Unknown"
880+ 
881+ return {
882+ "rows_guaranteed_safe": "Unknown",
883+ "blocks_are_contiguous": blocks_are_contiguous,
884+ "is_per_head_heterogeneous": False,
885+ "empty_row_risk_level": "medium",
886+ "l2_result": l2_detail,
887+ "l3_result": f"ELEMENT_UNSAFE({l3_detail})",
888+ "safety_layer": "L3_ELEMENT_FAIL",
889+ }
890+ 
891+ 
892+def build_flex_attention_metadata(kv_num_blocks: Any, kv_indices: Any, block_mask: Any = None) -> dict[str, Any]:
893+ """Build the minimal metadata contract used by flex attention lowering."""
894+ metadata = infer_block_sparse_metadata(kv_num_blocks, kv_indices, block_mask=block_mask)
895+ metadata.setdefault("primary_mask_type", "block_sparse")
896+ metadata.setdefault("is_segment_aware", "Unknown")
897+ metadata.setdefault("is_pad_safe", "Unknown")
898+ metadata.setdefault("causal_granularity", "Unknown")
899+ metadata.setdefault("is_approximate", False)
900+ metadata.setdefault("block_granularity_mode", "single_block_size")
901+ metadata.setdefault("has_cross_half_dependency", False)
902+ return metadata
903+ 
904+ 
905+def _has_block_sparse_kernel_option_override(kernel_options: dict[str, Any]) -> bool:
906+ """Check whether block-sparse kernel options are already fully specified."""
907+ return (
908+ "ROWS_GUARANTEED_SAFE" in kernel_options
909+ and "BLOCKS_ARE_CONTIGUOUS" in kernel_options
910+ )
911+ 
912+ 
913+def _extract_block_sparse_tensors(block_mask: Any) -> tuple[Any, Any]:
914+ """Read block-sparse tensors from a BlockMask-like object."""
915+ if block_mask is None:
916+ return None, None
917+ return getattr(block_mask, "kv_num_blocks", None), getattr(block_mask, "kv_indices", None)
918+ 
919+ 
920+def apply_kernel_options_from_metadata(
921+ kernel_options: dict[str, Any], metadata: dict[str, Any]
922+) -> dict[str, Any]:
923+ """Apply conservative kernel-option defaults derived from metadata."""
924+ updated = dict(kernel_options)
925+ rows_value = metadata.get("rows_guaranteed_safe", "Unknown")
926+ contiguous_value = metadata.get("blocks_are_contiguous", "Unknown")
927+ is_per_head_heterogeneous = metadata.get("is_per_head_heterogeneous", False)
928+ is_approximate = metadata.get("is_approximate", False)
929+ block_granularity_mode = metadata.get("block_granularity_mode", "single_block_size")
930+ empty_row_risk_level = metadata.get("empty_row_risk_level", "medium")
931+ has_cross_half_dependency = metadata.get("has_cross_half_dependency", False)
932+ causal_granularity = metadata.get("causal_granularity", "Unknown")
933+ 
934+ if "ROWS_GUARANTEED_SAFE" not in updated:
935+ if rows_value is False or empty_row_risk_level == "high":
936+ updated["ROWS_GUARANTEED_SAFE"] = False
937+ elif (
938+ rows_value is True
939+ and not is_per_head_heterogeneous
940+ and not is_approximate
941+ and block_granularity_mode != "multi_block_size"
942+ and empty_row_risk_level == "low"
943+ ):
944+ updated["ROWS_GUARANTEED_SAFE"] = True
945+ 
946+ if "BLOCKS_ARE_CONTIGUOUS" not in updated:
947+ if contiguous_value is False:
948+ updated["BLOCKS_ARE_CONTIGUOUS"] = False
949+ elif (
950+ contiguous_value is True
951+ and not is_per_head_heterogeneous
952+ and not is_approximate
953+ and block_granularity_mode != "multi_block_size"
954+ and not has_cross_half_dependency
955+ and causal_granularity != "mixed"
956+ ):
957+ updated["BLOCKS_ARE_CONTIGUOUS"] = True
958+ 
959+ return updated
960+ 
961+ 
962+def apply_kernel_options_from_eager_block_mask(
963+ kernel_options: dict[str, Any] | None,
964+ block_mask: Any,
965+ context: str = "eager",
966+ *,
967+ allow_tensor_analysis: bool = True,
968+) -> dict[str, Any]:
969+ """Infer kernel options from an eager BlockMask-like object when tensors are still available."""
970+ updated = {} if kernel_options is None else dict(kernel_options)
971+ updated = _apply_sparse_mask_compact_options(
972+ updated,
973+ block_mask,
974+ context,
975+ allow_tensor_analysis=allow_tensor_analysis,
976+ )
977+ updated = _apply_precomputed_block_sparse_safety_options(
978+ updated,
979+ block_mask,
980+ context,
981+ )
982+ if not _metadata_auto_infer_enabled():
983+ return _apply_disabled_metadata_defaults(updated)
984+ 
985+ if not allow_tensor_analysis:
986+ return updated
987+ if _has_block_sparse_kernel_option_override(updated):
988+ return updated
989+ 
990+ kv_num_blocks, kv_indices = _extract_block_sparse_tensors(block_mask)
991+ if kv_num_blocks is None or kv_indices is None:
992+ return updated
993+ 
994+ return apply_kernel_options_from_block_sparse_mask(
995+ updated,
996+ kv_num_blocks,
997+ kv_indices,
998+ block_mask=block_mask,
999+ context=context,
1000+ )
1001+ 
1002+ 
1003+def apply_kernel_options_from_block_sparse_mask(
1004+ kernel_options: dict[str, Any] | None,
1005+ kv_num_blocks: Any,
1006+ kv_indices: Any,
1007+ block_mask: Any = None,
1008+ context: str = "unknown"
1009+) -> dict[str, Any]:
1010+ """Infer block-sparse metadata and conservatively merge it into kernel options."""
1011+ updated = {} if kernel_options is None else dict(kernel_options)
1012+ if block_mask is not None:
1013+ updated = _apply_sparse_mask_compact_options(
1014+ updated,
1015+ block_mask,
1016+ context,
1017+ allow_tensor_analysis=True,
1018+ )
1019+ updated = _apply_precomputed_block_sparse_safety_options(
1020+ updated,
1021+ block_mask,
1022+ context,
1023+ )
1024+ if not _metadata_auto_infer_enabled():
1025+ return _apply_disabled_metadata_defaults(updated)
1026+ 
1027+ if _has_block_sparse_kernel_option_override(updated):
1028+ return updated
1029+ 
1030+ metadata = build_flex_attention_metadata(kv_num_blocks, kv_indices, block_mask=block_mask)
1031+ 
1032+ updated = apply_kernel_options_from_metadata(updated, metadata)
1033+ 
1034+ if log.isEnabledFor(logging.INFO):
1035+ log.info(
1036+ "[flex_attention][%s] rows_guaranteed_safe=%s blocks_are_contiguous=%s "
1037+ "is_per_head_heterogeneous=%s empty_row_risk_level=%s "
1038+ "ROWS_GUARANTEED_SAFE=%s BLOCKS_ARE_CONTIGUOUS=%s "
1039+ "safety_layer=%s l2_result=%s l3_result=%s",
1040+ context,
1041+ metadata.get("rows_guaranteed_safe", "Unknown"),
1042+ metadata.get("blocks_are_contiguous", "Unknown"),
1043+ metadata.get("is_per_head_heterogeneous", False),
1044+ metadata.get("empty_row_risk_level", "medium"),
1045+ updated.get("ROWS_GUARANTEED_SAFE", "<unset>"),
1046+ updated.get("BLOCKS_ARE_CONTIGUOUS", "<unset>"),
1047+ metadata.get("safety_layer", "UNKNOWN"),
1048+ metadata.get("l2_result", "N/A"),
1049+ metadata.get("l3_result", "N/A"),
1050+ )
1051+ return updated