cpp_class() — pytorch Function Reference

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

Function python ComputeKernels QuantizedCPU calls 1 called by 1

Entity Profile

ComputeKernels→ QuantizedCPU→ cpp_class() — pytorch Function Reference

Dependency Diagram

graph TD
  a19035f6_8189_b8f7_6a21_4dbaf2229779["cpp_class()"]
  a8e20f8a_af7c_0de9_f90c_676b733830c2["cpp_class()"]
  a8e20f8a_af7c_0de9_f90c_676b733830c2 -->|calls| a19035f6_8189_b8f7_6a21_4dbaf2229779
  a8e20f8a_af7c_0de9_f90c_676b733830c2["cpp_class()"]
  a19035f6_8189_b8f7_6a21_4dbaf2229779 -->|calls| a8e20f8a_af7c_0de9_f90c_676b733830c2
  style a19035f6_8189_b8f7_6a21_4dbaf2229779 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}>"