Skip to content

Add k_scale and v_scale to persistent attention #1322

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Add k_scale and v_scale to persistent attention #1322

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 10 additions & 0 deletions flashinfer/attention.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -130,6 +130,8 @@ def run(
kv_cache: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]],
out: Optional[torch.Tensor] = None,
lse: Optional[torch.Tensor] = None,
k_scale: Optional[torch.Tensor] = None,
v_scale: Optional[torch.Tensor] = None,
profiler_buffer: Optional[torch.Tensor] = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
if profiler_buffer is None:
Expand All @@ -147,6 +149,8 @@ def run(
head_dim_qk = q.shape[2]
if self._sm_scale is None:
self._sm_scale = 1.0 / math.sqrt(head_dim_qk)
if k_scale is not None:
self._sm_scale *= k_scale
Comment on lines +152 to +153
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Modifying the instance attribute self._sm_scale in-place can lead to incorrect behavior if run() is called multiple times, as the scaling factor will accumulate with each call. The effective scale should be calculated in a local variable within the run method.

        sm_scale = self._sm_scale
        if self._sm_scale is None:
            sm_scale = 1.0 / math.sqrt(head_dim_qk)
        if k_scale is not None:
            sm_scale *= k_scale

        # profiler_buffer is optional
        profiler_args = (profiler_buffer,) if self._use_profiler else ()

        self.module.run(
            self.float_workspace_buffer,
            self.int_workspace_buffer,
            self._plan_info,
            q,
            k_cache,
            v_cache,
            self._kv_indices,
            out,
            lse,
            self._mask_mode,
            TensorLayout[self._kv_layout].value,
            self._num_qo_heads,
            self._num_kv_heads,
            self._page_size,
            sm_scale,
            *profiler_args,
        )


# profiler_buffer is optional
profiler_args = (profiler_buffer,) if self._use_profiler else ()
Expand All @@ -169,5 +173,11 @@ def run(
self._sm_scale,
*profiler_args,
)
if v_scale is not None:
# TODO(Zihao): fused into kernel
if out.itemsize == 1:
out = (out.to(float) * v_scale).to(out.dtype)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

critical

When the out tensor is provided by the caller, reassigning it with out = ... breaks the contract, as the original tensor passed by the user will not be updated. The update must be performed in-place to ensure the caller's tensor is modified as expected.

                out.copy_((out.to(float) * v_scale).to(out.dtype))

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can you also implement fusing the multiply inside the kernel in the persistent kernel template?

Also, it will be good to specify an out_dtype, for quantized input such as fp8, we usually use bf16 as output data type. You can check

flashinfer/flashinfer/gemm.py

Lines 1198 to 1203 in 435f341

if out is None:
out = torch.empty(
(a.shape[0], b.shape[1]),
device=a.device,
dtype=out_dtype,
)
for reference.

Copy link
Contributor Author

@Edenzzzz Edenzzzz Aug 2, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In that case we need to pass in out_dtype during plan? Which can be troublesome

else:
out *= v_scale

return out, lse