From 948f469bfcd016c26d3d96220141bc4dab02615b Mon Sep 17 00:00:00 2001 From: DefTruth Date: Thu, 9 Jul 2026 05:37:21 +0000 Subject: [PATCH] fix sage3(sm120) attn backend tensor layout --- src/cache_dit/attention/backends/sage3.py | 79 ++++------------------- 1 file changed, 12 insertions(+), 67 deletions(-) diff --git a/src/cache_dit/attention/backends/sage3.py b/src/cache_dit/attention/backends/sage3.py index 73056fd2..ca82be19 100644 --- a/src/cache_dit/attention/backends/sage3.py +++ b/src/cache_dit/attention/backends/sage3.py @@ -41,12 +41,14 @@ def _sage3_attention_forward_op( raise RuntimeError("SageAttention-3 backend is not available. " "Please install `sageattn3` to use it.") + # SageAttention-3 required [B,H,N,D] tensor layout for Q/K/V. out = sageattn3_blackwell( - q=query.clone(), - k=key.clone(), - v=value.clone(), + q=query.transpose(1, 2).contiguous(), # [B,N,H,D] -> [B,H,N,D] + k=key.transpose(1, 2).contiguous(), + v=value.transpose(1, 2).contiguous(), is_causal=is_causal, ) + out = out.transpose(1, 2).contiguous() return out @@ -89,10 +91,10 @@ def _sage3_attention( + delta_s (f32) → S (f32) # OK: exact delta correction → softmax → P (f32) # OK: fine-grained attn weights ╔══════════════╗ - ║ quantize() ║ # BUG: P (f32, [0,1]) → E2M1 + ║ quantize() ║ # NOTE: P (f32, [0,1]) → E2M1 ║ P → FP4 ║ # Only {0, 0.5, 1} usable in [0,1] ╚══════╤═══════╝ - P_fp4 x V_fp4 ──MMA──→ O (f32) # BROKEN: PV based on destroyed P + P_fp4 x V_fp4 ──MMA──→ O (f32) # NOTE: PV based on FP4 P The E2M1 FP4 format only provides 2 non-zero levels ({0.5, 1}) in the [0, 1] softmax range, effectively collapsing continuous attention @@ -106,65 +108,6 @@ def _sage3_attention( aware of this limitation. For cache-dit the preferred path is ``--attn sage`` (SAGE2, lossless) rather than ``--attn sage3`` until the upstream kernel is fixed to keep P at FP16 or FP8 precision. - - .. warning:: - **Fixing this requires a kernel-level refactor, not a simple patch.** - - The entire kernel data pipeline is deeply coupled around FP4 (E2M1) - block-scaled MMA. Upgrading P/V from FP4 to FP8 (E4M3) touches the - full chain from Python quantization down to PTX MMA instructions. - Estimated scope: **12-15 files, 1-2 weeks**. - - Full change chain (Python → CUDA → PTX):: - - Python: v → scale_and_quant_fp4_transpose(v) → v_fp4 (uint8, D×N/2) - → sfv (fp8_e4m3, D×N/16) - Python: v → scale_and_quant_fp8_transpose(v) → v_fp8 (uint8, D×N/4) ← new kernel - → sfv (not needed) - - api.cu: params.v_ptr, params.sfv_ptr → delete sfv_ptr - params.v_row_stride → update (FP8 stride vs FP4) - - mainloop_tma_ws.h: - SmemLayoutVt ← E2M1 smem selector → E4M3 smem selector - SmemLayoutSFVt ← block-scale SF → delete - TMA load V ← E2M1 element type → E4M3 element type - TMA load SFV ← block-scale SF → delete - copy_v_block() ← LDSM for E2M1 → LDSM for E4M3 - tOrVt ← FP4 reg tensor → FP8 reg tensor - tOrSFVt ← FP4 scale factors → delete - (P side also:) - quantize() ← P(f32)→E2M1+UE4M3 → P(f32)→E4M3 (no SF) - tOrP ← FP4 reg tensor → FP8 reg tensor - tOrSFP ← FP4 scale factors → delete - TiledMmaPV ← SM120 block-scaled FP4 → SM80 FP8 MMA - - kernel_traits.h: - LayoutP ← FP4 register layout → FP8 register layout - LayoutSFP ← SF layout → delete - SmemLayoutSFV* ← SF deduced layouts → delete - NumSFPV ← kBlockN/16 → delete - - blockscaled_layout.h: - SmemLayoutAtomSFV/Vt ← deduced → delete - - softmax_fused.h: - AbsMaxP ← per-block tracking → delete (FP8 no block-scale) - fp8_scalexfp4_* ← SF scale constants → delete - fp4_scale_log2 ← E2M1 offset → delete - - params.h: - sfv_ptr, sfv_*_stride → delete - - launch.h / static_switch.h: - run_mha_fwd_ template args → remove SFV params - - fp4_quantization_4d.cu: - +scaled_fp8_quant_trans kernel → cvt e4m3, transpose, pack - - api.py: - scale_and_quant_fp4_transpose → scale_and_quant_fp8_transpose - blockscaled_fp4_attn → remove sfv arg """ if attn_mask is not None: raise ValueError("`attn_mask` is not yet supported for SageAttention-3.") @@ -179,12 +122,14 @@ def _sage3_attention( if sageattn3_blackwell is None: raise RuntimeError("SageAttention-3 backend is not available. " "Please install `sageattn3` to use it.") + # SageAttention-3 required [B,H,N,D] tensor layout for Q/K/V. out = sageattn3_blackwell( - q=query.clone(), - k=key.clone(), - v=value.clone(), + q=query.transpose(1, 2).contiguous(), # [B,N,H,D] -> [B,H,N,D] + k=key.transpose(1, 2).contiguous(), + v=value.transpose(1, 2).contiguous(), is_causal=is_causal, ) + out = out.transpose(1, 2).contiguous() else: out = _context_parallel_attention( query,