cpp_class() — pytorch Function Reference

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

Function python CoreTensor Registration calls 1 called by 1

Entity Profile

CoreTensor→ Registration→ cpp_class() — pytorch Function Reference

Dependency Diagram

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

Relationship Graph

Source Code

aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/generate_kernels.py lines 193–208

    def cpp_class(self) -> str:
        template_args = ", ".join(
            [
                f"cutlass::arch::Sm{self.sm_range[0]}",
                DTYPES[self.dtype],
                "true" if self.aligned else "false",
                "true" if self.apply_dropout else "false",
                "true" if self.preload_mmas else "false",
                str(self.block_i),
                str(self.block_j),
                str(self.max_k),
            ]
        )
        if self.keys_queries_aligned_to_blocksizes:
            template_args += ", true"
        return f"AttentionBackwardKernel<{template_args}>"