Add Logging helper

[drisspg]

Stacked PRs:

• Potential API change #2326
• [Ai-assisted] CLC work stealing #2218
• ->Add Logging helper #2327

---

Add Logging helper

For provenance, with this:

import os
os.environ.setdefault("FA4_CLC", "1")

import torch
from flash_attn.cute.logging import set_fa_log_level
from flash_attn.cute.interface import flash_attn_func

device = "cuda"
dtype = torch.bfloat16

q = torch.randn(2, 256, 4, 128, device=device, dtype=dtype)
k = torch.randn(2, 256, 4, 128, device=device, dtype=dtype)
v = torch.randn(2, 256, 4, 128, device=device, dtype=dtype)

for lvl in range(4):
    set_fa_log_level(lvl)
    print(f"{'='*60}")
    print(f"FA_LOG_LEVEL = {lvl}")
    print(f"{'='*60}")
    out, _ = flash_attn_func(q, k, v)
    torch.cuda.synchronize()
    print()

Prints

python ../my_scripts/flex/logging_demo.py
============================================================
FA_LOG_LEVEL = 0
============================================================

============================================================
FA_LOG_LEVEL = 1
============================================================
[FA] use_clc_scheduler=True, TileScheduler=CLCDynamicTileScheduler, USE_2CTA=False

============================================================
FA_LOG_LEVEL = 2
============================================================
[FA] use_clc_scheduler=True, TileScheduler=CLCDynamicTileScheduler, USE_2CTA=False

============================================================
FA_LOG_LEVEL = 3
============================================================
[FA] use_clc_scheduler=True, TileScheduler=CLCDynamicTileScheduler, USE_2CTA=False
[CLC] query sm=0 cta=0 (m_blk=2,h=0,b=0,s=0) valid=0

[CLC] query sm=1 cta=0 (m_blk=2,h=0,b=0,s=0) valid=0

[CLC] query sm=147 cta=0 (m_blk=2,h=0,b=0,s=0) valid=0

[CLC] query sm=146 cta=0 (m_blk=2,h=0,b=0,s=0) valid=0

[CLC] query sm=145 cta=0 (m_blk=2,h=0,b=0,s=0) valid=0

[CLC] query sm=143 cta=0 (m_blk=2,h=0,b=0,s=0) valid=0

[CLC] query sm=144 cta=0 (m_blk=2,h=0,b=0,s=0) valid=0

[CLC] query sm=142 cta=0 (m_blk=2,h=0,b=0,s=0) valid=0

[CLC] map sm=0 cta=0/1 linear=6/8 (m_blk=0,h=2,b=1,s=0) valid=1

[CLC] map sm=147 cta=0/1 linear=5/8 (m_blk=0,h=1,b=1,s=0) valid=1

[CLC] map sm=142 cta=0/1 linear=0/8 (m_blk=0,h=0,b=0,s=0) valid=1

[CLC] map sm=1 cta=0/1 linear=7/8 (m_blk=0,h=3,b=1,s=0) valid=1

[CLC] map sm=146 cta=0/1 linear=4/8 (m_blk=0,h=0,b=1,s=0) valid=1

[CLC] map sm=144 cta=0/1 linear=2/8 (m_blk=0,h=2,b=0,s=0) valid=1

[CLC] map sm=143 cta=0/1 linear=1/8 (m_blk=0,h=1,b=0,s=0) valid=1

[CLC] map sm=145 cta=0/1 linear=3/8 (m_blk=0,h=3,b=0,s=0) valid=1

[CLC] pull sm=144 cta=0/1 linear=2/8 (m_blk=2,h=0,b=0,s=0) valid=0

[CLC] pull sm=146 cta=0/1 linear=2/8 (m_blk=2,h=0,b=0,s=0) valid=0

[CLC] pull sm=1 cta=0/1 linear=2/8 (m_blk=2,h=0,b=0,s=0) valid=0

[CLC] pull sm=0 cta=0/1 linear=2/8 (m_blk=2,h=0,b=0,s=0) valid=0

[CLC] pull sm=147 cta=0/1 linear=2/8 (m_blk=2,h=0,b=0,s=0) valid=0

[CLC] pull sm=145 cta=0/1 linear=2/8 (m_blk=2,h=0,b=0,s=0) valid=0

[CLC] pull sm=142 cta=0/1 linear=2/8 (m_blk=2,h=0,b=0,s=0) valid=0

[CLC] pull sm=143 cta=0/1 linear=2/8 (m_blk=2,h=0,b=0,s=0) valid=0


~/meta/clc clc *21 ❯                                                                        13s  clc dev@gpu-dev-796fc59d