Skip to content
  • Aether's avatar
    feat: implement real decode optimization (FlashAttn fast-path + uniform-decode) · c41ac299
    Aether authored
    Rewrite decode_performance_opt.py from empty stats shell to real optimization:
    1. FlashAttentionBackend.forward fast-path (x64/step):
       - Cache kv_cache.unbind(0) per step (shared across 64 layers)
       - Skip redundant assert/attn_type/fp8/dcp checks for decoder-only models
       - Cache descale_shape and sliding_window list
       - Safe fallback: validates assumptions on first call, falls back permanently
         if not decoder-only/non-fp8/single-GPU; try/except on every call
    2. _get_cumsum_and_arange uniform-decode fast-path:
       - Skip cumsum+repeat when all requests schedule 1 token
    3. Lightweight decode/prefill step counter
    
    Compliance: no scheduler logic change, no locked params, no semantic change.
    Optimized forward produces identical output as original (only caches invariants).
    c41ac299