iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[SM100] Hdim 256 optimized (#52980)

Browse Source 
Signed-off-by: Simon Veitner <sveitner@redhat.com>
This commit is contained in:
simon-veitner-redhat
2026-08-25 20:28:40 -07:00
committed by GitHub
parent f6130145c8
commit 7156c63bef
7 changed files with 423 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files
cmake/external_projects/vllm_flash_attn.cmake
+1 -1
View File
@@ -39,7 +39,7 @@ else()
FetchContent_Declare(
vllm-flash-attn
GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
GIT_TAG f3e1a4f74c99145c0717709860bf765de1703779
GIT_TAG 4eb29359a51b4fe67c8c0f0c5e931236971aa492
GIT_PROGRESS TRUE
# Don't share the vllm-flash-attn build between build types
BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn
docs/design/attention_backends.md
+8
View File
@@ -127,6 +127,14 @@ Priority is **1 = highest** (tried first).
> **†** FlashInfer Native is the regular FlashInfer path. XQA is the SM90 decode path exposed through FlashInfer's TRTLLM decode API. trtllm-gen is used on SM100 and supports sinks. Disable XQA/trtllm-gen via `--attention-config.use_trtllm_attention=0`.
>
> **\*** Specify the FlashAttention version via `--attention-config.flash_attn_version=2`, `3`, or `4`. Default is FA4 on SM100+ (Blackwell), FA3 on SM90 (Hopper), FA2 otherwise.
>
> On Blackwell, when the FlashAttention backend is selected, `head_size=256` is served
> by a dedicated FA4 kernel that requires a KV cache block size of 128 (advertised
> automatically, so it is picked as long as `--block-size` is not pinned) and does not
> support logit soft capping, attention sinks, mm_prefix/R-SWA masking, DCP, or windowed
> encoder attention. Those configurations transparently fall back to FA2. A pinned
> `--block-size` that is not a multiple of 128 instead makes FlashAttention ineligible
> for such models, which is an error if the backend was requested explicitly.
## MiniMax M3 Sparse Attention Backends
tests/kernels/attention/test_attention_selector.py
+176 -1
View File
@@ -1,7 +1,8 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from unittest.mock import patch
from contextlib import contextmanager
from unittest.mock import MagicMock, patch
import pytest
import torch
@@ -14,6 +15,7 @@ from vllm.config import (
)
from vllm.platforms import current_platform
from vllm.platforms.cpu import CpuPlatform
from vllm.platforms.interface import DeviceCapability
if current_platform.is_cuda():
from vllm.platforms.cuda import CudaPlatform
@@ -25,6 +27,7 @@ if current_platform.is_rocm():
else:
RocmPlatform = None
from vllm.v1.attention.backend import AttentionType
from vllm.v1.attention.backends.registry import AttentionBackendEnum
from vllm.v1.attention.selector import _cached_get_attn_backend, get_attn_backend
@@ -555,3 +558,175 @@ def test_flash_attn_accepts_handled_fp8_variants(
# import order across earlier tests that patch vllm.platforms.current_platform.
monkeypatch.setattr(fa_utils_mod.current_platform, "is_xpu", lambda: True)
assert FlashAttentionBackend.supports_kv_cache_dtype(kv_cache_dtype)
blackwell_only = pytest.mark.skipif(
not current_platform.is_cuda(), reason="FA4 is CUDA-only"
)
@contextmanager
def _blackwell(vllm_config=None):
platform = MagicMock()
platform.is_xpu.return_value = False
platform.is_rocm.return_value = False
platform.get_device_capability.return_value = DeviceCapability(10, 0)
with (
patch("vllm.v1.attention.backends.fa_utils.current_platform", platform),
patch(
"vllm.vllm_flash_attn.flash_attn_interface.is_fa_version_supported",
return_value=True,
),
patch("vllm.config.get_current_vllm_config_or_none", return_value=vllm_config),
patch(
"vllm.v1.attention.backends.flash_attn.get_current_vllm_config_or_none",
return_value=vllm_config,
),
):
yield
def _hd256_config(
*,
is_mm_prefix_lm=False,
rswa_window=None,
dcp_size=1,
softcap=None,
head_size=256,
cache_dtype="auto",
):
vllm_config = MagicMock()
vllm_config.attention_config.flash_attn_version = None
vllm_config.model_config.is_mm_prefix_lm = is_mm_prefix_lm
vllm_config.model_config.rswa_window = rswa_window
vllm_config.model_config.hf_text_config.attn_logit_softcapping = softcap
vllm_config.model_config.get_head_size.return_value = head_size
vllm_config.cache_config.cache_dtype = cache_dtype
vllm_config.parallel_config.decode_context_parallel_size = dcp_size
return vllm_config
@blackwell_only
@pytest.mark.parametrize(
"kwargs,config_kwargs,expected",
[
({}, {}, 4),
({"supports_fa4_hd256": False}, {}, 2),
({"has_sinks": True}, {}, 2),
({"requires_softcap": True}, {}, 2),
# Larger block sizes normalize to 128.
({"kv_cache_block_size": 16}, {}, 2),
({"kv_cache_block_size": 64}, {}, 2),
({"kv_cache_block_size": 256}, {}, 4),
({}, {"is_mm_prefix_lm": True}, 2),
({}, {"rswa_window": 512}, 2),
({}, {"dcp_size": 2}, 2),
({}, {"softcap": 50.0}, 2),
({}, {"cache_dtype": "fp8"}, 2),
({"head_size": 128, "kv_cache_block_size": 16}, {}, 4),
({"head_size": 192, "head_size_v": 128, "kv_cache_block_size": 16}, {}, 4),
({"head_size": 256, "head_size_v": 128, "kv_cache_block_size": 16}, {}, 2),
({"head_size": 256, "head_size_v": 64}, {}, 2),
],
)
def test_fa4_hd256_fallback_matrix(kwargs, config_kwargs, expected):
from vllm.v1.attention.backends.fa_utils import get_flash_attn_version
kwargs = {
"head_size": 256,
"kv_cache_block_size": 128,
"supports_fa4_hd256": True,
**kwargs,
}
with _blackwell(_hd256_config(head_size=kwargs["head_size"], **config_kwargs)):
assert get_flash_attn_version(**kwargs) == expected
@blackwell_only
@pytest.mark.parametrize(
"backend_name,config_kwargs,expected",
[
("FLASH_ATTN", {}, 128),
("FLASH_ATTN", {"softcap": 50.0}, 16),
("FLASH_ATTN", {"head_size": 128}, 16),
("FLASH_ATTN", {"cache_dtype": "fp8"}, 16),
("FLASH_ATTN_DIFFKV", {}, 16),
],
)
def test_fa4_hd256_block_size_advertisement(
backend_name: str, config_kwargs: dict, expected: int
):
from vllm.v1.attention.backend import MultipleOf
from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
from vllm.v1.attention.backends.flash_attn_diffkv import (
FlashAttentionDiffKVBackend,
)
backend = {
"FLASH_ATTN": FlashAttentionBackend,
"FLASH_ATTN_DIFFKV": FlashAttentionDiffKVBackend,
}[backend_name]
vllm_config = _hd256_config(**config_kwargs)
with _blackwell(vllm_config):
(size,) = backend.get_supported_kernel_block_sizes()
preferred = backend.get_preferred_block_size(16)
assert preferred == expected
if expected == 128:
assert size == 128
else:
assert isinstance(size, MultipleOf) and size.base == expected
@blackwell_only
def test_fa4_hd256_mm_prefix_deselects_flash_attn():
from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
with _blackwell(_hd256_config(is_mm_prefix_lm=True)):
assert (
FlashAttentionBackend.supports_combination(
head_size=256,
dtype=torch.bfloat16,
kv_cache_dtype=None,
block_size=128,
use_mla=False,
has_sink=False,
use_sparse=False,
use_mm_prefix=True,
device_capability=DeviceCapability(10, 0),
)
is not None
)
@pytest.mark.skipif(
not current_platform.is_cuda()
or not current_platform.is_device_capability_family(100),
reason="requires a Blackwell GPU",
)
@pytest.mark.parametrize(
"attn_type,sliding_window,expected",
[
(AttentionType.DECODER, None, 4),
(AttentionType.DECODER, 512, 4),
# Local encoder attention lacks the required seqused tensors.
(AttentionType.ENCODER_ONLY, None, 4),
(AttentionType.ENCODER_ONLY, 512, 2),
],
)
def test_fa4_hd256_impl_selection(attn_type, sliding_window, expected):
from vllm.v1.attention.backends.flash_attn import FlashAttentionImpl
# Implementations are constructed before the platform updates the block size.
with set_current_vllm_config(VllmConfig()):
impl = FlashAttentionImpl(
num_heads=8,
head_size=256,
scale=0.0625,
num_kv_heads=8,
alibi_slopes=None,
sliding_window=sliding_window,
kv_cache_dtype="auto",
attn_type=attn_type,
)
assert impl.vllm_flash_attn_version == expected
assert impl.fa4_hd256 == (expected == 4)
tests/kernels/attention/test_flash_attn.py
+75
View File
@@ -215,3 +215,78 @@ def test_varlen_with_paged_kv(
torch.testing.assert_close(output, ref_output, atol=atol, rtol=rtol),
f"{torch.max(torch.abs(output - ref_output))}",
)
@pytest.mark.skipif(
not current_platform.is_device_capability_family(100),
reason="FA4 only has a dedicated head_size=256 kernel on Blackwell",
)
@pytest.mark.parametrize(
"seq_lens", [[(1, 1328), (5, 18), (129, 463)], [(1, 523), (1, 37), (1, 2011)]]
)
@pytest.mark.parametrize("num_heads", NUM_HEADS)
@pytest.mark.parametrize("sliding_window", SLIDING_WINDOWS)
@torch.inference_mode()
def test_fa4_hd256_paged_call_shape(
seq_lens: list[tuple[int, int]],
num_heads: tuple[int, int],
sliding_window: int | None,
) -> None:
"""Handles excess block-table columns with page-aligned ``max_seqlen_k``."""
from vllm.v1.attention.backends.fa_utils import FA4_HD256_PAGE_SIZE
if not is_fa_version_supported(4):
pytest.skip(f'FA4 unsupported: "{fa_version_unsupported_reason(4)}"')
torch.set_default_device("cuda")
set_random_seed(0)
head_size = 256
block_size = FA4_HD256_PAGE_SIZE
num_blocks = 2048
dtype = torch.bfloat16
num_query_heads, num_kv_heads = num_heads
query_lens = [x[0] for x in seq_lens]
kv_lens = [x[1] for x in seq_lens]
max_kv_len = max(kv_lens)
window_size = (sliding_window - 1, 0) if sliding_window is not None else (-1, -1)
scale = head_size**-0.5
query = torch.randn(sum(query_lens), num_query_heads, head_size, dtype=dtype)
key_cache = torch.randn(
num_blocks, block_size, num_kv_heads, head_size, dtype=dtype
)
value_cache = torch.randn_like(key_cache)
cu_query_lens = torch.tensor([0] + query_lens, dtype=torch.int32).cumsum(
dim=0, dtype=torch.int32
)
num_pages = (max_kv_len + block_size - 1) // block_size
block_tables = torch.randint(
0, num_blocks, (len(seq_lens), num_pages + 3), dtype=torch.int32
)
output = flash_attn_varlen_func(
q=query,
k=key_cache,
v=value_cache,
cu_seqlens_q=cu_query_lens,
seqused_k=torch.tensor(kv_lens, dtype=torch.int32),
max_seqlen_q=max(query_lens),
max_seqlen_k=num_pages * block_size,
softmax_scale=scale,
causal=True,
window_size=window_size,
block_table=block_tables[:, :num_pages],
num_splits=1,
fa_version=4,
)
ref_output = ref_paged_attn(
query=query,
key_cache=key_cache,
value_cache=value_cache,
query_lens=query_lens,
kv_lens=kv_lens,
block_tables=block_tables,
scale=scale,
sliding_window=sliding_window,
)
torch.testing.assert_close(output, ref_output, atol=1.5e-2, rtol=1e-2)
vllm/v1/attention/backends/fa_utils.py
+79 -14
View File
@@ -6,9 +6,13 @@ from typing import Any
import vllm.envs as envs
from vllm.logger import init_logger
from vllm.platforms import current_platform
from vllm.utils.torch_utils import is_quantized_kv_cache
logger = init_logger(__name__)
# The dedicated hd256 kernel requires one 128-token page per tile.
FA4_HD256_PAGE_SIZE = 128
# Track whether upstream flash-attn is available on ROCm.
# Set during module initialization and never modified afterwards.
# This module-level flag avoids repeated import attempts and ensures
@@ -72,6 +76,9 @@ def get_flash_attn_version(
head_size: int | None = None,
head_size_v: int | None = None,
has_sinks: bool = False,
requires_softcap: bool = False,
kv_cache_block_size: int | None = None,
supports_fa4_hd256: bool = False,
) -> int | None:
if current_platform.is_xpu():
return 2
@@ -168,26 +175,31 @@ def get_flash_attn_version(
)
fa_version = 2
# TODO: Restore the `requires_local_attention` restriction when FA4
# head-dim 256 is re-enabled.
if fa_version == 4 and device_capability.major >= 10 and head_size == 256:
logger.warning_once(
"FA4 on Blackwell is temporarily disabled for head_size=256, "
"defaulting to FA version 2."
)
fa_version = 2
if fa_version == 4 and uses_fa4_hd256_kernel(head_size, head_size_v):
if not supports_fa4_hd256:
fa_version = 2
elif (
reason := _fa4_hd256_fallback_reason(
has_sinks, requires_softcap, kv_cache_block_size, vllm_config
)
) is not None:
logger.warning_once(
"FA4's Blackwell head_size=256 kernel does not support %s, "
"defaulting to FA version 2.",
reason,
)
fa_version = 2
# FA4 on SM100 (Blackwell) has TMEM capacity limits that restrict
# supported head dimensions to ≤128. The 192/128 MLA prefill case is
# supported; 256 is temporarily disabled until upstream supports the
# required features. Development of symmetric 192, 384, and 512 support
# is tracked in https://github.com/Dao-AILab/flash-attention/issues/2456
# FA4 head dimensions on Blackwell are limited by TMEM capacity.
if (
fa_version == 4
and device_capability.major >= 10
and head_size is not None
and head_size > 128
and not (head_size == 192 and head_size_v == 128)
and not (
(head_size == 256 and head_size_v in (None, 256))
or (head_size == 192 and head_size_v == 128)
)
):
logger.warning_once(
"FA4 on Blackwell does not support head_size=%d due to TMEM "
@@ -209,6 +221,53 @@ def get_flash_attn_version(
return None
def uses_fa4_hd256_kernel(
head_size: int | None, head_size_v: int | None = None
) -> bool:
"""Return whether FA4 uses its dedicated hd256 kernel."""
if head_size != 256:
return False
if head_size_v is not None and head_size_v != 256:
return False
capability = current_platform.get_device_capability()
return capability is not None and capability.major in (10, 11)
def _fa4_hd256_fallback_reason(
has_sinks: bool,
requires_softcap: bool,
kv_cache_block_size: int | None,
vllm_config: Any,
) -> str | None:
model_config = vllm_config.model_config if vllm_config is not None else None
cache_config = vllm_config.cache_config if vllm_config is not None else None
if has_sinks:
return "attention sinks"
if requires_softcap or (
# Keep model-level and per-layer version selection consistent.
model_config is not None
and getattr(model_config.hf_text_config, "attn_logit_softcapping", None)
):
return "logits soft capping"
if cache_config is not None and is_quantized_kv_cache(cache_config.cache_dtype):
return f"quantized KV cache dtype {cache_config.cache_dtype}"
if kv_cache_block_size is not None and kv_cache_block_size % FA4_HD256_PAGE_SIZE:
# Larger blocks are split into 128-token kernel pages.
return f"a KV cache block size of {kv_cache_block_size}"
if model_config is not None:
if model_config.is_mm_prefix_lm:
return "mm_prefix bidirectional attention"
if model_config.rswa_window is not None:
return "R-SWA"
if (
vllm_config is not None
and vllm_config.parallel_config.decode_context_parallel_size > 1
):
# DCP reuses an unsliced block table with a non-page-aligned length.
return "decode context parallelism"
return None
def is_fa_version_supported(fa_version: int) -> bool:
try:
from vllm.vllm_flash_attn.flash_attn_interface import (
@@ -227,6 +286,9 @@ def flash_attn_supports_kv_cache_dtype(
head_size: int | None = None,
head_size_v: int | None = None,
has_sinks: bool = False,
requires_softcap: bool = False,
kv_cache_block_size: int | None = None,
supports_fa4_hd256: bool = False,
) -> bool:
if kv_cache_dtype == "fp8_e5m2":
return False
@@ -237,6 +299,9 @@ def flash_attn_supports_kv_cache_dtype(
head_size=head_size,
head_size_v=head_size_v,
has_sinks=has_sinks,
requires_softcap=requires_softcap,
kv_cache_block_size=kv_cache_block_size,
supports_fa4_hd256=supports_fa4_hd256,
)
return (fa_version == 3 and current_platform.is_device_capability_family(90)) or (
fa_version == 4 and current_platform.is_device_capability_family(100)
vllm/v1/attention/backends/flash_attn.py
+83 -6
View File
@@ -24,11 +24,13 @@ from vllm.v1.attention.backend import (
MultipleOf,
)
from vllm.v1.attention.backends.fa_utils import (
FA4_HD256_PAGE_SIZE,
flash_attn_supports_kv_cache_dtype,
flash_attn_supports_quant_query_input,
get_flash_attn_version,
is_fa_version_supported,
is_flash_attn_varlen_func_available,
uses_fa4_hd256_kernel,
)
from vllm.v1.attention.backends.utils import (
fill_mm_prefix_query_ranges,
@@ -85,15 +87,40 @@ class FlashAttentionBackend(AttentionBackend):
"fp8",
"fp8_e4m3",
]
head_size_v: int | None = None
@staticmethod
def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
@classmethod
def _get_fa4_hd256_block_size(cls) -> int | None:
vllm_config = get_current_vllm_config_or_none()
if vllm_config is None or vllm_config.model_config is None:
return None
head_size = vllm_config.model_config.get_head_size()
if (
uses_fa4_hd256_kernel(head_size, cls.head_size_v)
and get_flash_attn_version(
head_size=head_size,
head_size_v=cls.head_size_v,
supports_fa4_hd256=True,
)
== 4
):
return FA4_HD256_PAGE_SIZE
return None
@classmethod
def get_supported_kernel_block_sizes(cls) -> list[int | MultipleOf]:
if block_size := cls._get_fa4_hd256_block_size():
# Sliding-window specs select the smallest advertised size.
return [block_size]
return [MultipleOf(16)]
forward_includes_kv_cache_update: bool = False
@classmethod
def get_preferred_block_size(cls, default_block_size: int) -> int:
if block_size := cls._get_fa4_hd256_block_size():
return max(default_block_size, block_size)
if current_platform.is_xpu():
return max(default_block_size, 64)
return super().get_preferred_block_size(default_block_size)
@@ -194,12 +221,20 @@ class FlashAttentionBackend(AttentionBackend):
head_size=head_size,
head_size_v=head_size,
has_sinks=has_sink,
kv_cache_block_size=block_size,
supports_fa4_hd256=True,
)
):
return "FP8 KV cache requires FA3 on SM90 or FA4 on SM100"
if (
use_mm_prefix
and get_flash_attn_version(head_size=head_size, has_sinks=has_sink) != 4
and get_flash_attn_version(
head_size=head_size,
has_sinks=has_sink,
kv_cache_block_size=block_size,
supports_fa4_hd256=True,
)
!= 4
):
return (
"mm_prefix (PrefixLM bidirectional attention) requires "
@@ -409,6 +444,15 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad
self.max_num_splits = 0 # No upper bound on the number of splits.
self.aot_schedule = get_flash_attn_version() == 3
self.fa4_hd256 = uses_fa4_hd256_kernel(self.headdim) and (
get_flash_attn_version(
head_size=self.headdim,
kv_cache_block_size=self.block_size,
supports_fa4_hd256=True,
)
== 4
)
try:
from vllm.distributed.parallel_state import get_dcp_group
@@ -780,6 +824,9 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad
metadata.scheduler_metadata = self._store_scheduler_metadata(scheduler_metadata)
def use_cascade_attention(self, *args, **kwargs) -> bool:
if self.fa4_hd256:
# Cascade may use a non-page-aligned prefix length.
return False
return use_cascade_attention(*args, **kwargs)
@@ -823,11 +870,30 @@ class FlashAttentionImpl(AttentionImpl):
self.num_queries_per_kv = self.num_heads // self.num_kv_heads
self.attn_type = attn_type
vllm_config = get_current_vllm_config_or_none()
uses_kv_cache = attn_type not in (
AttentionType.ENCODER,
AttentionType.ENCODER_ONLY,
)
# The final KV cache block size is unavailable during construction.
self.vllm_flash_attn_version = get_flash_attn_version(
requires_alibi=alibi_slopes is not None,
head_size=head_size,
has_sinks=sinks is not None,
requires_softcap=bool(self.logits_soft_cap),
supports_fa4_hd256=True,
)
self.fa4_hd256 = self.vllm_flash_attn_version == 4 and uses_fa4_hd256_kernel(
head_size
)
if self.fa4_hd256 and not uses_kv_cache and sliding_window is not None:
# The hd256 kernel requires seqused_k for local attention.
logger.warning_once(
"FA4's Blackwell head_size=256 kernel does not support local "
"attention on encoder inputs, defaulting to FA version 2."
)
self.vllm_flash_attn_version = 2
self.fa4_hd256 = False
logger.info_once(
"Using FlashAttention version %s",
self.vllm_flash_attn_version,
@@ -843,6 +909,8 @@ class FlashAttentionImpl(AttentionImpl):
head_size=head_size,
head_size_v=head_size,
has_sinks=sinks is not None,
requires_softcap=bool(self.logits_soft_cap),
supports_fa4_hd256=True,
):
raise NotImplementedError(
f"FlashAttention does not support {self.kv_cache_dtype}"
@@ -861,7 +929,6 @@ class FlashAttentionImpl(AttentionImpl):
self.supports_quant_query_input = flash_attn_supports_quant_query_input()
vllm_config = get_current_vllm_config_or_none()
dcp_a2a = (
vllm_config is not None
and vllm_config.parallel_config.decode_context_parallel_size > 1
@@ -1083,6 +1150,15 @@ class FlashAttentionImpl(AttentionImpl):
)
causal = not has_window
num_splits = attn_metadata.max_num_splits
if self.fa4_hd256:
# hd256 requires page-aligned lengths, exact-width block
# tables, and no SplitKV.
num_pages = cdiv(max_seqlen_k, FA4_HD256_PAGE_SIZE)
max_seqlen_k = num_pages * FA4_HD256_PAGE_SIZE
block_table = block_table[:, :num_pages]
num_splits = 1
flash_attn_varlen_func(
q=query[:num_actual_tokens],
k=key_cache,
@@ -1104,7 +1180,7 @@ class FlashAttentionImpl(AttentionImpl):
k_descale=k_descale,
v_descale=v_descale,
dynamic_causal=dynamic_causal,
num_splits=attn_metadata.max_num_splits,
num_splits=num_splits,
s_aux=self.sinks,
mask_mod=rswa_mask_mod_fn or mm_mask_mod,
aux_tensors=rswa_aux or mm_aux,
@@ -1416,7 +1492,8 @@ class FlashAttentionImpl(AttentionImpl):
else None,
k_descale=layer._k_scale.expand(descale_shape), # type: ignore[operator]
v_descale=layer._v_scale.expand(descale_shape), # type: ignore[operator]
num_splits=1 if self.batch_invariant_enabled else 0,
# The hd256 kernel does not support SplitKV.
num_splits=1 if self.batch_invariant_enabled or self.fa4_hd256 else 0,
s_aux=self.sinks,
)
vllm/v1/attention/backends/flash_attn_diffkv.py
+1
View File
@@ -87,6 +87,7 @@ class FlashAttentionDiffKVImpl(FlashAttentionImpl):
head_size_v=FlashAttentionDiffKVBackend.head_size_v,
has_sinks=self.sinks is not None,
)
self.fa4_hd256 = False
def do_kv_cache_update(
self,