| #!/usr/bin/env python3 | |
| """Count entries in the coverage_ignore_* lists in docs/source/conf.py. | |
| Uses AST parsing instead of importing conf.py, since importing it directly | |
| fails outside of a Sphinx build (module-level setup). | |
| """ | |
| import ast | |
| import sys | |
| CONF = "docs/source/conf.py" |
Scope: the FA2 kernels in csrc/flash_attn/ (not FA3 hopper/, not the FA4 flash_attn/cute/ code).
Pinned to commit 0bbb25a (all line numbers/links below are against this commit).
FA2 uses Philox RNG only for dropout, and only on the default CUDA generator. It depends on three ATen/PyTorch RNG facilities:
| #!/usr/bin/env python3 | |
| """FA2 dropout RNG: Exploring options in Python. Structured to parallel the real FA2 code. | |
| Mirrors FA2's actual plumbing: | |
| * a CUDA `fwd` that produces an `rng_state` [seed, offset] int64 tensor and a `bwd` that | |
| consumes it (so forward and backward apply the SAME dropout mask), and | |
| * a `torch.autograd.Function` wrapper that does `ctx.save_for_backward(rng_state)`, | |
| exactly like flash_attn/flash_attn_interface.py. | |
| Run: python agent_space/fa2_rng_graph_demo.py |
| # write a hook to flatten optimizer state_dict when saving | |
| def flatten_state_dict(optim, osd): | |
| flattened_sd = {} | |
| state = osd['state'] | |
| for idx, param_group in enumerate(osd['param_groups']): | |
| assert 'param_names' in param_group, "param names are required as they'll be used as keys" | |
| for param_name, param_id in zip(param_group['param_names'], param_group['params']): | |
| # add all the state | |
| if param_id in state: | |
| for key, value in state[param_id].items(): |
| import torch | |
| from torch.utils.weak import WeakTensorKeyDictionary | |
| from torch.utils._pytree import tree_map | |
| from torch.utils._python_dispatch import TorchDispatchMode | |
| evictable_tensors: WeakTensorKeyDictionary = {} | |
| class OffloadTensor(torch.Tensor): | |
| cuda_elem: torch.Tensor | |
| cpu_elem: torch.Tensor |
| from ctypes import c_void_p, c_long | |
| import torch | |
| import math | |
| import random | |
| from torch import empty_strided, as_strided, device | |
| from torch._inductor.codecache import AsyncCompile | |
| from torch._inductor.select_algorithm import extern_kernels | |
| aten = torch.ops.aten |