Skip to content

Commit e78352a

Browse files
lmxyylmxyy
committed
make linter happy
1 parent 2bb718f commit e78352a

3 files changed

Lines changed: 226 additions & 70 deletions

File tree

‎nunchaku/models/attention_processor.py‎

Lines changed: 171 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,171 @@
1+
import math
2+
from typing import Optional, Tuple
3+
4+
import torch
5+
from torch.nn import functional as F
6+
7+
from ..ops.fused import fused_qkv_norm_rottary
8+
9+
10+
class NunchakuFA2Processor:
11+
12+
def __call__(
13+
self,
14+
attn,
15+
hidden_states: torch.Tensor,
16+
encoder_hidden_states: Optional[torch.Tensor] = None,
17+
attention_mask: Optional[torch.Tensor] = None,
18+
image_rotary_emb: Tuple[torch.Tensor, torch.Tensor] | torch.Tensor = None,
19+
**kwargs,
20+
) -> torch.Tensor | Tuple[torch.Tensor, torch.Tensor]:
21+
# Adapted from https://github.com/huggingface/diffusers/blob/50dea89dc6036e71a00bc3d57ac062a80206d9eb/src/diffusers/models/attention_processor.py#L2275
22+
if attention_mask is not None:
23+
raise NotImplementedError("attention_mask is not supported")
24+
25+
batch_size, _, channels = hidden_states.shape
26+
assert channels == self.heads * self.head_dim
27+
qkv = fused_qkv_norm_rottary(
28+
hidden_states,
29+
self.to_qkv,
30+
self.norm_q,
31+
self.norm_k,
32+
image_rotary_emb[0] if isinstance(image_rotary_emb, tuple) else image_rotary_emb,
33+
)
34+
35+
if self.added_kv_proj_dim is not None:
36+
assert encoder_hidden_states is not None
37+
assert isinstance(image_rotary_emb, tuple)
38+
qkv_context = fused_qkv_norm_rottary(
39+
encoder_hidden_states, self.add_qkv_proj, self.norm_added_q, self.norm_added_k, image_rotary_emb[1]
40+
)
41+
qkv = torch.cat([qkv_context, qkv], dim=1)
42+
43+
query, key, value = qkv.chunk(3, dim=-1)
44+
query = query.view(batch_size, -1, self.heads, self.head_dim).transpose(1, 2)
45+
key = key.view(batch_size, -1, self.heads, self.head_dim).transpose(1, 2)
46+
value = value.view(batch_size, -1, self.heads, self.head_dim).transpose(1, 2)
47+
hidden_states = F.scaled_dot_product_attention(
48+
query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
49+
)
50+
hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, self.heads * self.head_dim)
51+
hidden_states = hidden_states.to(query.dtype)
52+
if encoder_hidden_states is not None:
53+
encoder_hidden_states, hidden_states = (
54+
hidden_states[:, : encoder_hidden_states.shape[1]],
55+
hidden_states[:, encoder_hidden_states.shape[1] :],
56+
)
57+
# linear proj
58+
hidden_states = self.to_out[0](hidden_states)
59+
# dropout
60+
hidden_states = self.to_out[1](hidden_states)
61+
encoder_hidden_states = self.to_add_out(encoder_hidden_states)
62+
return hidden_states, encoder_hidden_states
63+
else:
64+
# for single transformer block, we split the proj_out into two linear layers
65+
hidden_states = self.to_out(hidden_states)
66+
return hidden_states
67+
68+
69+
class NunchakuFP16AttnProcessor:
70+
71+
def __init__(self, pad_size: int = 256):
72+
self.pad_size = pad_size
73+
74+
def __call__(
75+
self,
76+
attn,
77+
hidden_states: torch.Tensor,
78+
encoder_hidden_states: Optional[torch.Tensor] = None,
79+
attention_mask: Optional[torch.Tensor] = None,
80+
image_rotary_emb: Tuple[torch.Tensor, torch.Tensor] | torch.Tensor = None,
81+
**kwargs,
82+
) -> torch.Tensor | Tuple[torch.Tensor, torch.Tensor]:
83+
pad_size = self.pad_size
84+
85+
batch_size, _, channels = hidden_states.shape
86+
assert channels == self.heads * self.head_dim
87+
if encoder_hidden_states is None:
88+
num_tokens = hidden_states.shape[1]
89+
num_tokens_pad = math.ceil(num_tokens / pad_size) * pad_size
90+
query = torch.empty(
91+
batch_size,
92+
self.heads,
93+
num_tokens_pad,
94+
self.head_dim,
95+
dtype=torch.float16,
96+
device=hidden_states.device,
97+
)
98+
key = torch.empty_like(query)
99+
value = torch.empty_like(query)
100+
101+
assert torch.is_tensor(image_rotary_emb)
102+
fused_qkv_norm_rottary(
103+
hidden_states,
104+
self.to_qkv,
105+
self.norm_q,
106+
self.norm_k,
107+
image_rotary_emb,
108+
output=(query, key, value),
109+
num_tokens=num_tokens,
110+
)
111+
112+
else:
113+
num_txt_tokens = encoder_hidden_states.shape[1]
114+
num_img_tokens = hidden_states.shape[1]
115+
num_txt_tokens_pad = math.ceil(num_txt_tokens / pad_size) * pad_size
116+
num_img_tokens_pad = math.ceil(num_img_tokens / pad_size) * pad_size
117+
num_tokens_pad = num_txt_tokens_pad + num_img_tokens_pad
118+
query = torch.empty(
119+
batch_size,
120+
self.heads,
121+
num_tokens_pad,
122+
self.head_dim,
123+
dtype=torch.float16,
124+
device=hidden_states.device,
125+
)
126+
key = torch.empty_like(query)
127+
value = torch.empty_like(query)
128+
129+
assert isinstance(image_rotary_emb, tuple)
130+
fused_qkv_norm_rottary(
131+
hidden_states,
132+
self.to_qkv,
133+
self.norm_q,
134+
self.norm_k,
135+
image_rotary_emb[0],
136+
output=(query[:, :num_img_tokens_pad], key[:, :num_img_tokens_pad], value[:, :num_img_tokens_pad]),
137+
num_tokens=num_img_tokens,
138+
)
139+
fused_qkv_norm_rottary(
140+
encoder_hidden_states,
141+
self.add_qkv_proj,
142+
self.norm_added_q,
143+
self.norm_added_k,
144+
image_rotary_emb[1],
145+
output=(query[:, num_img_tokens_pad:], key[:, num_img_tokens_pad:], value[:, num_img_tokens_pad:]),
146+
num_tokens=num_txt_tokens,
147+
)
148+
attention_output = torch.empty(
149+
batch_size,
150+
num_tokens_pad,
151+
self.heads * self.head_dim,
152+
dtype=hidden_states.dtype,
153+
device=hidden_states.device,
154+
)
155+
attention_fp16(query, key, value, attention_output, self.head_dim ** (-0.5))
156+
157+
if encoder_hidden_states is not None:
158+
encoder_hidden_states, hidden_states = (
159+
hidden_states[:, : encoder_hidden_states.shape[1]],
160+
hidden_states[:, encoder_hidden_states.shape[1] :],
161+
)
162+
# linear proj
163+
hidden_states = self.to_out[0](hidden_states)
164+
# dropout
165+
hidden_states = self.to_out[1](hidden_states)
166+
encoder_hidden_states = self.to_add_out(encoder_hidden_states)
167+
return hidden_states, encoder_hidden_states
168+
else:
169+
# for single transformer block, we split the proj_out into two linear layers
170+
hidden_states = self.to_out(hidden_states)
171+
return hidden_states

‎nunchaku/models/transformers/transformer_flux_v2.py‎

Lines changed: 10 additions & 50 deletions
Original file line numberDiff line numberDiff line change
@@ -13,11 +13,11 @@
1313
from huggingface_hub import utils
1414
from torch import nn
1515
from torch.nn import GELU
16-
from torch.nn import functional as F
1716

18-
from ...ops.fused import fused_gelu_mlp, fused_qkv_norm_rottary
17+
from ...ops.fused import fused_gelu_mlp
1918
from ...utils import get_precision
2019
from ..attention import NunchakuFeedForward
20+
from ..attention_processor import NunchakuFA2Processor
2121
from ..embeddings import NunchakuFluxPosEmbed, pack_rotemb
2222
from ..linear import SVDQW4A4Linear
2323
from ..normalization import NunchakuAdaLayerNormZero, NunchakuAdaLayerNormZeroSingle
@@ -26,7 +26,7 @@
2626

2727

2828
class NunchakuFluxAttention(nn.Module):
29-
def __init__(self, flux_attention: FluxAttention, processor: str = "flashattn2", **kwargs):
29+
def __init__(self, flux_attention: FluxAttention, **kwargs):
3030
super(NunchakuFluxAttention, self).__init__()
3131

3232
self.head_dim = flux_attention.head_dim
@@ -65,7 +65,7 @@ def __init__(self, flux_attention: FluxAttention, processor: str = "flashattn2",
6565
self.add_qkv_proj = SVDQW4A4Linear.from_linear(add_qkv_proj, **kwargs)
6666
self.to_add_out = SVDQW4A4Linear.from_linear(flux_attention.to_add_out, **kwargs)
6767

68-
self.processor = processor
68+
self.processor = NunchakuFA2Processor()
6969

7070
def forward(
7171
self,
@@ -76,53 +76,13 @@ def forward(
7676
**kwargs,
7777
):
7878
# Adapted from [diffusers v0.34.0](https://github.com/huggingface/diffusers/blob/50dea89dc6036e71a00bc3d57ac062a80206d9eb/src/diffusers/models/attention_processor.py#L2275)
79-
if attention_mask is not None:
80-
raise NotImplementedError("attention_mask is not supported")
81-
82-
batch_size, _, channels = hidden_states.shape
83-
assert channels == self.heads * self.head_dim
84-
qkv = fused_qkv_norm_rottary(
85-
hidden_states,
86-
self.to_qkv,
87-
self.norm_q,
88-
self.norm_k,
89-
image_rotary_emb[0] if isinstance(image_rotary_emb, tuple) else image_rotary_emb,
90-
)
91-
92-
if self.added_kv_proj_dim is not None:
93-
assert encoder_hidden_states is not None
94-
assert isinstance(image_rotary_emb, tuple)
95-
qkv_context = fused_qkv_norm_rottary(
96-
encoder_hidden_states, self.add_qkv_proj, self.norm_added_q, self.norm_added_k, image_rotary_emb[1]
97-
)
98-
qkv = torch.cat([qkv_context, qkv], dim=1)
99-
100-
query, key, value = qkv.chunk(3, dim=-1)
101-
query = query.view(batch_size, -1, self.heads, self.head_dim).transpose(1, 2)
102-
key = key.view(batch_size, -1, self.heads, self.head_dim).transpose(1, 2)
103-
value = value.view(batch_size, -1, self.heads, self.head_dim).transpose(1, 2)
104-
105-
hidden_states = F.scaled_dot_product_attention(
106-
query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
79+
return self.processor(
80+
attn=self,
81+
hidden_states=hidden_states,
82+
encoder_hidden_states=encoder_hidden_states,
83+
attention_mask=attention_mask,
84+
image_rotary_emb=image_rotary_emb,
10785
)
108-
hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, self.heads * self.head_dim)
109-
hidden_states = hidden_states.to(query.dtype)
110-
111-
if encoder_hidden_states is not None:
112-
encoder_hidden_states, hidden_states = (
113-
hidden_states[:, : encoder_hidden_states.shape[1]],
114-
hidden_states[:, encoder_hidden_states.shape[1] :],
115-
)
116-
# linear proj
117-
hidden_states = self.to_out[0](hidden_states)
118-
# dropout
119-
hidden_states = self.to_out[1](hidden_states)
120-
encoder_hidden_states = self.to_add_out(encoder_hidden_states)
121-
return hidden_states, encoder_hidden_states
122-
else:
123-
# for single transformer block, we split the proj_out into two linear layers
124-
hidden_states = self.to_out(hidden_states)
125-
return hidden_states
12686

12787

12888
class NunchakuFluxTransformerBlock(FluxTransformerBlock):

‎nunchaku/ops/fused.py‎

Lines changed: 45 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -55,8 +55,9 @@ def fused_qkv_norm_rottary(
5555
norm_q: RMSNorm,
5656
norm_k: RMSNorm,
5757
rotary_emb: torch.Tensor,
58-
output: torch.Tensor | None = None,
59-
):
58+
output: torch.Tensor | tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
59+
num_tokens: int = 0,
60+
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
6061
assert isinstance(norm_q, RMSNorm)
6162
assert isinstance(norm_k, RMSNorm)
6263

@@ -67,21 +68,45 @@ def fused_qkv_norm_rottary(
6768
if output is None:
6869
output = torch.empty(quantized_x.shape[0], proj.out_features, dtype=x.dtype, device=x.device)
6970

70-
svdq_gemm_w4a4_cuda(
71-
act=quantized_x,
72-
wgt=proj.qweight,
73-
out=output,
74-
ascales=ascales,
75-
wscales=proj.wscales,
76-
lora_act_in=lora_act,
77-
lora_up=proj.proj_up,
78-
bias=proj.bias,
79-
fp4=proj.precision == "nvfp4",
80-
alpha=proj.wtscale,
81-
wcscales=proj.wcscales,
82-
norm_q=norm_q.weight,
83-
norm_k=norm_k.weight,
84-
rotary_emb=rotary_emb,
85-
)
86-
output = output.view(batch_size, seq_len, -1)
87-
return output
71+
if isinstance(output, tuple):
72+
assert len(output) == 3
73+
output_q, output_k, output_v = output
74+
svdq_gemm_w4a4_cuda(
75+
act=quantized_x,
76+
wgt=proj.qweight,
77+
ascales=ascales,
78+
wscales=proj.wscales,
79+
lora_act_in=lora_act,
80+
lora_up=proj.proj_up,
81+
bias=proj.bias,
82+
fp4=proj.precision == "nvfp4",
83+
alpha=proj.wtscale,
84+
wcscales=proj.wcscales,
85+
norm_q=norm_q.weight,
86+
norm_k=norm_k.weight,
87+
rotary_emb=rotary_emb,
88+
out_q=output_q,
89+
out_k=output_k,
90+
out_v=output_v,
91+
num_tokens=num_tokens,
92+
)
93+
return output_q, output_k, output_v
94+
else:
95+
svdq_gemm_w4a4_cuda(
96+
act=quantized_x,
97+
wgt=proj.qweight,
98+
out=output,
99+
ascales=ascales,
100+
wscales=proj.wscales,
101+
lora_act_in=lora_act,
102+
lora_up=proj.proj_up,
103+
bias=proj.bias,
104+
fp4=proj.precision == "nvfp4",
105+
alpha=proj.wtscale,
106+
wcscales=proj.wcscales,
107+
norm_q=norm_q.weight,
108+
norm_k=norm_k.weight,
109+
rotary_emb=rotary_emb,
110+
)
111+
output = output.view(batch_size, seq_len, -1)
112+
return output

0 commit comments

Comments
 (0)