[None][fix] read attention layout by schema position

Signed-off-by: Balamurugan Marimuthu <246387390+bmarimuthu-nv@users.noreply.github.com>

Author: bmarimuthu-nv

tensorrt_llm/_torch/auto_deploy/custom_ops/attention/flashinfer_attention.py

@@ -565,14 +565,7 @@ def get_host_prepare_metadata_function(cls) -> Optional[PrepareMetadataHostCalla
     @classmethod
     def get_constants(cls, source_attn_node: Node) -> List[Constant]:
         # Sanity check: layout == "bsnd"
-        # Prefer kwargs; fall back to the final positional arg if it's a string.
-        layout = source_attn_node.kwargs.get("layout", None)
-        if (
-            layout is None
-            and len(source_attn_node.args) > 0
-            and isinstance(source_attn_node.args[-1], str)
-        ):
-            layout = source_attn_node.args[-1]
+        layout = extract_op_args(source_attn_node, "layout")[0]
         if layout != "bsnd":
             raise RuntimeError(
                 f"Expected torch_attention layout='bsnd' but got {layout!r} "

tensorrt_llm/_torch/auto_deploy/custom_ops/attention/torch_backend_attention.py

@@ -499,14 +499,7 @@ def get_dynamic_inputs(cls, source_attn_node: Node) -> List[Optional[Node]]:
     @classmethod
     def get_constants(cls, source_attn_node: Node) -> List[Constant]:
         # Sanity check: layout == "bsnd"
-        # Prefer kwargs; fall back to the final positional arg if it's a string.
-        layout = source_attn_node.kwargs.get("layout", None)
-        if (
-            layout is None
-            and len(source_attn_node.args) > 0
-            and isinstance(source_attn_node.args[-1], str)
-        ):
-            layout = source_attn_node.args[-1]
+        layout = extract_op_args(source_attn_node, "layout")[0]
         if layout != "bsnd":
             raise RuntimeError(
                 f"Expected torch_attention layout='bsnd' but got {layout!r} "

tensorrt_llm/_torch/auto_deploy/custom_ops/attention/triton_attention.py

@@ -388,14 +388,7 @@ def get_cache_initializers(
     @classmethod
     def get_constants(cls, source_attn_node: Node) -> List[Constant]:
         # Sanity check: layout == "bsnd"
-        # Prefer kwargs; fall back to the final positional arg if it's a string.
-        layout = source_attn_node.kwargs.get("layout", None)
-        if (
-            layout is None
-            and len(source_attn_node.args) > 0
-            and isinstance(source_attn_node.args[-1], str)
-        ):
-            layout = source_attn_node.args[-1]
+        layout = extract_op_args(source_attn_node, "layout")[0]
         if layout != "bsnd":
             raise RuntimeError(
                 f"Expected torch_attention layout='bsnd' but got {layout!r} "

tensorrt_llm/_torch/auto_deploy/custom_ops/attention/triton_paged_attention.py

@@ -1359,13 +1359,7 @@ def get_dynamic_inputs(cls, source_attn_node: Node) -> List[Optional[Node]]:
 
     @classmethod
     def get_constants(cls, source_attn_node: Node) -> List[Constant]:
-        layout = source_attn_node.kwargs.get("layout", None)
-        if (
-            layout is None
-            and len(source_attn_node.args) > 0
-            and isinstance(source_attn_node.args[-1], str)
-        ):
-            layout = source_attn_node.args[-1]
+        layout = extract_op_args(source_attn_node, "layout")[0]
         if layout != "bsnd":
             raise RuntimeError(
                 f"Expected torch_attention layout='bsnd' but got {layout!r} "

tensorrt_llm/_torch/auto_deploy/custom_ops/attention/trtllm_attention.py

@@ -639,13 +639,7 @@ def get_constants(cls, source_attn_node: Node) -> List[Constant]:
         from tensor shapes or SequenceInfo metadata at runtime.
         """
         # Sanity check: layout == "bsnd"
-        layout = source_attn_node.kwargs.get("layout", None)
-        if (
-            layout is None
-            and len(source_attn_node.args) > 0
-            and isinstance(source_attn_node.args[-1], str)
-        ):
-            layout = source_attn_node.args[-1]
+        layout = extract_op_args(source_attn_node, "layout")[0]
         if layout != "bsnd":
             raise RuntimeError(
                 f"Expected torch_attention layout='bsnd' but got {layout!r} "