get_all() — pytorch Function Reference

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

Function python CoreTensor Dispatcher calls 1 called by 2

Entity Profile

CoreTensor→ Dispatcher→ get_all() — pytorch Function Reference

Dependency Diagram

graph TD
  94e5a4b8_623e_3785_583d_a91baa59116d["get_all()"]
  403a54ad_6077_f6c9_6d25_554b79570b27["get_all()"]
  403a54ad_6077_f6c9_6d25_554b79570b27 -->|calls| 94e5a4b8_623e_3785_583d_a91baa59116d
  aad87944_33f6_d3c8_e945_0c0de37e4fe0["main()"]
  aad87944_33f6_d3c8_e945_0c0de37e4fe0 -->|calls| 94e5a4b8_623e_3785_583d_a91baa59116d
  403a54ad_6077_f6c9_6d25_554b79570b27["get_all()"]
  94e5a4b8_623e_3785_583d_a91baa59116d -->|calls| 403a54ad_6077_f6c9_6d25_554b79570b27
  style 94e5a4b8_623e_3785_583d_a91baa59116d fill:#6366f1,stroke:#818cf8,color:#fff

Relationship Graph

Source Code

aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/generate_kernels.py lines 117–143

    def get_all(cls) -> list["FwdKernel"]:
        kernels: list[FwdKernel] = []
        for aligned, dtype, (sm, sm_max) in itertools.product(
            [True, False], DTYPES.keys(), itertools.pairwise(SM)
        ):
            # Remove some kernels we don't use
            if dtype == "bf16" and sm < 80:
                continue
            if not aligned and sm >= 80:
                continue
            for q, k, max_k in [
                (64, 64, 64),
                # We get better perf with 64x128 on A100
                (64 if sm > 75 else 32, 128, 128),
                (32, 128, 2**16),
            ]:
                kernels.append(
                    cls(
                        aligned=aligned,
                        dtype=dtype,
                        sm_range=(sm, sm_max),
                        q=q,
                        k=k,
                        max_k=max_k,
                    )
                )
        return kernels