quack() — pytorch Function Reference

Architecture documentation for the quack() function in kernels.py from the pytorch codebase.

Function python DynamoBenchmarks Microbenchmarking calls 1

Entity Profile

DynamoBenchmarks→ Microbenchmarking→ quack() — pytorch Function Reference

Dependency Diagram

graph TD
  5465167f_99c1_b00a_f75d_fc30dc46fccc["quack()"]
  d898eb25_c2ca_75ee_1aa8_482af50c99ca["compute_rstd()"]
  5465167f_99c1_b00a_f75d_fc30dc46fccc -->|calls| d898eb25_c2ca_75ee_1aa8_482af50c99ca
  style 5465167f_99c1_b00a_f75d_fc30dc46fccc fill:#6366f1,stroke:#818cf8,color:#fff

Relationship Graph

Source Code

benchmarks/dynamo/genai_layers/kernels.py lines 508–543

    def quack(self, args, kwargs=None) -> Any:
        if kwargs is not None:
            raise AssertionError(f"Expected kwargs to be None, but got {kwargs}")
        from quack.rmsnorm import _get_sm_count, _rmsnorm_bwd

        (
            x,
            w,
            dy,
        ) = args
        M, N = x.shape

        rstd = self.compute_rstd(x, eps=1e-6)
        dx = torch.empty_like(x)
        sm_count = _get_sm_count(x.size(1), x.device)
        dw_partial = torch.empty(
            sm_count, x.size(1), device=x.device, dtype=torch.float32
        )

        def quack_bwd():
            _rmsnorm_bwd(
                x,
                w,
                dy,
                rstd,
                dx,
                dw_partial,
                db_partial=None,
                dresidual_out=None,
                dresidual=None,
                sm_count=sm_count,
            )
            dw = dw_partial.sum(dim=0).to(w.dtype)
            return dx, dw

        return quack_bwd