1+ # Modified from the original SageATtention3 code
2+ """
3+ Copyright (c) 2025 by SageAttention team.
4+
5+ Licensed under the Apache License, Version 2.0 (the "License");
6+ you may not use this file except in compliance with the License.
7+ You may obtain a copy of the License at
8+
9+ http://www.apache.org/licenses/LICENSE-2.0
10+
11+ Unless required by applicable law or agreed to in writing, software
12+ distributed under the License is distributed on an "AS IS" BASIS,
13+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+ See the License for the specific language governing permissions and
15+ limitations under the License.
16+ """
17+ import torch
18+ import triton
19+ import triton .language as tl
20+ import torch .nn .functional as F
21+ from typing import Tuple
22+ from torch .nn .functional import scaled_dot_product_attention as sdpa
23+ import fp4attn_cuda
24+ import fp4quant_cuda
25+
26+ # Centralized block size configuration for sageattn_blackwell kernels
27+ # These should match the values in fastvideo/attention/backends/sageattn/blackwell/block_config.h
28+ BLOCK_M = 128 # Block size for M dimension (query sequence length)
29+ BLOCK_N = 128 # Block size for N dimension (key/value sequence length)
30+
31+
32+ @triton .jit
33+ def group_mean_kernel (
34+ q_ptr ,
35+ q_out_ptr ,
36+ qm_out_ptr ,
37+ B , H , L , D : tl .constexpr ,
38+ stride_qb , stride_qh , stride_ql , stride_qd ,
39+ stride_qmb , stride_qmh , stride_qml , stride_qmd ,
40+ GROUP_SIZE : tl .constexpr
41+ ):
42+ pid_b = tl .program_id (0 )
43+ pid_h = tl .program_id (1 )
44+ pid_group = tl .program_id (2 )
45+
46+ group_start = pid_group * GROUP_SIZE
47+ offsets = group_start + tl .arange (0 , GROUP_SIZE )
48+
49+ q_offsets = pid_b * stride_qb + pid_h * stride_qh + offsets [:, None ] * stride_ql + tl .arange (0 , D )[None , :] * stride_qd
50+ q_group = tl .load (q_ptr + q_offsets )
51+
52+ qm_group = tl .sum (q_group , axis = 0 ) / GROUP_SIZE
53+
54+ q_group = q_group - qm_group
55+ tl .store (q_out_ptr + q_offsets , q_group )
56+
57+ qm_offset = pid_b * stride_qmb + pid_h * stride_qmh + pid_group * stride_qml + tl .arange (0 , D ) * stride_qmd
58+ tl .store (qm_out_ptr + qm_offset , qm_group )
59+
60+
61+ def triton_group_mean (q : torch .Tensor ):
62+ B , H , L , D = q .shape
63+ GROUP_SIZE = BLOCK_M
64+ num_groups = L // GROUP_SIZE
65+
66+ q_out = torch .empty_like (q ) # [B, H, L, D]
67+ qm = torch .empty (B , H , num_groups , D , device = q .device , dtype = q .dtype )
68+
69+ grid = (B , H , num_groups )
70+
71+ group_mean_kernel [grid ](
72+ q , q_out , qm ,
73+ B , H , L , D ,
74+ q .stride (0 ), q .stride (1 ), q .stride (2 ), q .stride (3 ),
75+ qm .stride (0 ), qm .stride (1 ), qm .stride (2 ), qm .stride (3 ),
76+ GROUP_SIZE = GROUP_SIZE
77+ )
78+ return q_out , qm
79+
80+
81+ def preprocess_qkv (q : torch .Tensor , k : torch .Tensor , v : torch .Tensor , per_block_mean : bool = True , enable_smoothing_q : bool = False , enable_smoothing_k : bool = False ):
82+
83+ def pad_to_block_size (x ):
84+ L = x .size (2 )
85+ pad_len = (BLOCK_M - L % BLOCK_M ) % BLOCK_M
86+ if pad_len == 0 :
87+ return x .contiguous ()
88+ return F .pad (x , (0 , 0 , 0 , pad_len ), value = 0 ).contiguous ()
89+
90+ if enable_smoothing_k :
91+ k -= k .mean (dim = - 2 , keepdim = True )
92+ q , k , v = map (lambda x : pad_to_block_size (x ), [q , k , v ])
93+ if per_block_mean and enable_smoothing_q :
94+ q , qm = triton_group_mean (q )
95+ elif enable_smoothing_q :
96+ qm = q .mean (dim = - 2 , keepdim = True )
97+ q = q - qm
98+ if enable_smoothing_q :
99+ delta_s = torch .matmul (qm , k .transpose (- 2 , - 1 )).to (torch .float32 ).contiguous ()
100+ else : # used to disable q smoothing
101+ B , H , L , D = q .shape
102+ delta_s = torch .zeros ((B , H , L // BLOCK_M , k .shape [2 ]), device = q .device , dtype = torch .float32 )
103+
104+ return q , k , v , delta_s
105+
106+ def scale_and_quant_fp4 (x : torch .Tensor ) -> Tuple [torch .Tensor , torch .Tensor ]:
107+ assert x .ndim == 4
108+ B , H , N , D = x .shape
109+ packed_fp4 = torch .empty ((B , H , N , D // 2 ), device = x .device , dtype = torch .uint8 )
110+ fp8_scale = torch .empty ((B , H , N , D // 16 ), device = x .device , dtype = torch .float8_e4m3fn )
111+ fp4quant_cuda .scaled_fp4_quant (x , packed_fp4 , fp8_scale , 1 )
112+ return packed_fp4 , fp8_scale
113+
114+ def scale_and_quant_fp4_permute (x : torch .Tensor ) -> Tuple [torch .Tensor , torch .Tensor ]:
115+ assert x .ndim == 4
116+ B , H , N , D = x .shape
117+ packed_fp4 = torch .empty ((B , H , N , D // 2 ), device = x .device , dtype = torch .uint8 )
118+ fp8_scale = torch .empty ((B , H , N , D // 16 ), device = x .device , dtype = torch .float8_e4m3fn )
119+ fp4quant_cuda .scaled_fp4_quant_permute (x , packed_fp4 , fp8_scale , 1 )
120+ return packed_fp4 , fp8_scale
121+
122+ def scale_and_quant_fp4_transpose (x : torch .Tensor ) -> Tuple [torch .Tensor , torch .Tensor ]:
123+ assert x .ndim == 4
124+ B , H , N , D = x .shape
125+ packed_fp4 = torch .empty ((B , H , D , N // 2 ), device = x .device , dtype = torch .uint8 )
126+ fp8_scale = torch .empty ((B , H , D , N // 16 ), device = x .device , dtype = torch .float8_e4m3fn )
127+ fp4quant_cuda .scaled_fp4_quant_trans (x , packed_fp4 , fp8_scale , 1 )
128+ return packed_fp4 , fp8_scale
129+
130+ def blockscaled_fp4_attn (qlist : Tuple ,
131+ klist : Tuple ,
132+ vlist : Tuple ,
133+ delta_s : torch .Tensor ,
134+ KL : int ,
135+ is_causal : bool = False ,
136+ per_block_mean : bool = True ,
137+ is_bf16 : bool = True ,
138+ single_level_p_quant : bool = False
139+ ):
140+ softmax_scale = (qlist [0 ].shape [- 1 ] * 2 ) ** (- 0.5 )
141+ return fp4attn_cuda .fwd (qlist [0 ], klist [0 ], vlist [0 ], qlist [1 ], klist [1 ], vlist [1 ], delta_s , KL , None , softmax_scale , is_causal , per_block_mean , is_bf16 , single_level_p_quant )
142+
143+
144+ def sageattn_blackwell (q , k , v , attn_mask = None , is_causal = False , per_block_mean = True , single_level_p_quant = True , ** kwargs ):
145+ """
146+ SageAttention3 Blackwell kernel for FP4 attention.
147+
148+ Args:
149+ q: Query tensor [B, H, L, D]
150+ k: Key tensor [B, H, L, D]
151+ v: Value tensor [B, H, L, D]
152+ attn_mask: Attention mask (not used)
153+ is_causal: Whether to use causal masking
154+ per_block_mean: Whether to use per-block mean for Q smoothing
155+ single_level_p_quant: If True, use single-level quantization: s_P2, P̂_2 = φ(P̃) directly
156+ (standard per-block FP4 quantization like V, no s_P1).
157+ If False (default), use two-level quantization:
158+ s_P1 = rowmax(P̃)/(448×6), then s_P2, P̂_2 = φ(P̃/s_P1).
159+ **kwargs: Additional arguments (ignored)
160+
161+ Returns:
162+ Output tensor [B, H, L, D]
163+ """
164+ if q .size (- 1 ) >= 256 :
165+ print (f"Unsupported Headdim { q .size (- 1 )} )
166+ return sdpa (q , k , v , is_causal = is_causal )
167+ QL = q .size (2 )
168+ KL = k .size (2 )
169+ is_bf16 = q .dtype == torch .bfloat16
170+ q , k , v , delta_s = preprocess_qkv (q , k , v , per_block_mean )
171+ qlist_from_cuda = scale_and_quant_fp4 (q )
172+ klist_from_cuda = scale_and_quant_fp4_permute (k )
173+ vlist_from_cuda = scale_and_quant_fp4_transpose (v )
174+ o_fp4 = blockscaled_fp4_attn (
175+ qlist_from_cuda ,
176+ klist_from_cuda ,
177+ vlist_from_cuda ,
178+ delta_s ,
179+ KL ,
180+ is_causal ,
181+ per_block_mean ,
182+ is_bf16 ,
183+ single_level_p_quant
184+ )[0 ][:, :, :QL , :].contiguous ()
185+ return o_fp4
0 commit comments