cpp_class() — pytorch Function Reference

Architecture documentation for the cpp_class() function in generate_kernels.py from the pytorch codebase.

Function python CoreTensor Dispatcher calls 1 called by 1

Entity Profile

CoreTensor→ Dispatcher→ cpp_class() — pytorch Function Reference

Dependency Diagram

graph TD
  610c07fb_f141_ce09_a54c_1ab18c1ec52c["cpp_class()"]
  73f12026_6ca6_757e_7b90_dda04e7e72e4["cpp_class()"]
  73f12026_6ca6_757e_7b90_dda04e7e72e4 -->|calls| 610c07fb_f141_ce09_a54c_1ab18c1ec52c
  73f12026_6ca6_757e_7b90_dda04e7e72e4["cpp_class()"]
  610c07fb_f141_ce09_a54c_1ab18c1ec52c -->|calls| 73f12026_6ca6_757e_7b90_dda04e7e72e4
  style 610c07fb_f141_ce09_a54c_1ab18c1ec52c fill:#6366f1,stroke:#818cf8,color:#fff

Relationship Graph

Source Code

aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/generate_kernels.py lines 87–100

    def cpp_class(self) -> str:
        template_args = ", ".join(
            [
                DTYPES[self.dtype],
                f"cutlass::arch::Sm{self.sm_range[0]}",
                "true" if self.aligned else "false",
                str(self.q),
                str(self.k),
                str(self.max_k),
                "true" if self.supports_dropout else "false",
                "true" if self.supports_bias else "false",
            ]
        )
        return f"AttentionKernel<{template_args}>"