About customized FA and h_w aspect ratio

[sihan-chen-yes]

Hi, thanks for the great post-training acceleration work for VAR!
I have two questions:
1.It seems the flash attention will not be applied if just using the standarad flash attn lib, could you provide the method to install your customized flash attention lib?
2.If changing the h_w div ratio to other values from 1.0, the acceleration seems not to work anymore... Does that meet the expectation?

Thanks in advance for your reply!

[csguoh]

Hi,
To install the flashattn lib, you may refer to the flashatn's github repo. For instance, you may clone this repo and install.
What do you mean by h_w div ratio in the 2nd question?

[sihan-chen-yes]

Sorry, forgot to clarify that I'm using infinity with FastVAR.
1.Yes. I installed the official flash attention lib, but it seems the code doesn't support this standard lib, it used a customized flash attention lib whose signature is different from the standard one...

        self.customized_flash_attn = customized_flash_attn and customized_kernel_installed
        if customized_flash_attn and not customized_kernel_installed:
            import inspect, warnings
            file_path = inspect.getsourcefile(flash_attn_func)
            line_number = inspect.getsourcelines(flash_attn_func)[1]
            info = (
                f'>>>>>> Customized FlashAttention2 is not installed or compiled, but specified in args by --flash=1. Set customized_flash_attn = False. <<<<<<\n'
                f'>>>>>> `flash_attn_func` is in [line {line_number}] [file {file_path}] <<<<<<\n'
                f'>>>>>> {flash_attn_func.__code__.co_varnames=} <<<<<<\n'
            )
            warnings.warn(info, ImportWarning)
            print(info, flush=True)

self.sa = FastVARSelfAttention(
            embed_dim=embed_dim, num_heads=num_heads, proj_drop=drop, tau=tau, cos_attn=cos_attn, customized_flash_attn=customized_flash_attn,
            use_flex_attn=use_flex_attn, batch_size=batch_size, pad_to_multiplier=pad_to_multiplier, rope2d_normalized_by_hw=rope2d_normalized_by_hw,
        )

        if self.using_flash:
            if attn_bias_or_two_vector is not None: # training
                kw = dict(VAR_visible_kvlen=attn_bias_or_two_vector[0], VAR_invisible_qlen=attn_bias_or_two_vector[1])
            else:                                   # inference (autoregressive sampling)
                kw = dict()
            oup = flash_attn_func(q.to(v.dtype), k.to(v.dtype), v, dropout_p=0, softmax_scale=self.scale, **kw).view(B, L, C)

2.For inference, we can specifiy the output resolution(h/w) for infinity, e.g., h_div_w = 1/1 # aspect ratio, height:width. It takes around 1.38s/img (11.8GB) on my RTX3090 if h_div_w is 1.0. However, when I change h_div_w to 9/16, the inference speed drops, the memory consuming also grows to the level of original infinity. It seems FastVAR not working anymore...