@tduncan / Nvidia-Optimization-Stack-VLLM / commits / dff1bde84d

[Model] Qwen4Exp: fp8_e4m3 main KV cache on the QSA path (#55557)

Signed-off-by: Andrea Semeraro <andrea.semeraro@sezione1.it> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Jiangyun Zhu <riverclouds.zhu@qq.com>

semerandre committed Sep 16, 2026 at 17:07 UTC dff1bde84dd6e34a49c150116d0f212507280910

3 files changed +287 -45

tests/models/qwen4_exp/test_qsa_reference.py

+57 -18

@@ -203,7 +203,14 @@ def _qsa_sparse_paged_attention_reference(

203 block_table: torch.Tensor,

204 token_to_req: torch.Tensor,

205 softmax_scale: float,

206 + k_scale: float = 1.0,

207 + v_scale: float = 1.0,

208 ) -> torch.Tensor:

209 + """Dense reference for QSA sparse paged attention.

210 +

211 + Mirrors the kernel's dequant: fp8-e4m3 K/V caches are dequantized with the

212 + per-tensor k_scale/v_scale host floats; bf16 caches use unit scales.

213 + """

214 output = torch.zeros_like(q)

215 repeats = q.shape[1] // k_cache.shape[2]

216 page_size = k_cache.shape[1]

@@ -215,13 +222,15 @@ def _qsa_sparse_paged_attention_reference(

222 request = token_to_req[row].long()

223 pages = block_table[request, logical // page_size].long()

224 offsets = logical % page_size

218 - keys = k_cache[pages, offsets].repeat_interleave(repeats, dim=1)

219 - values = v_cache[pages, offsets].repeat_interleave(repeats, dim=1)

220 - scores = torch.einsum("hd,khd->hk", q[row].float(), keys.float())

221 - probabilities = torch.softmax(scores * softmax_scale, dim=-1)

222 - output[row] = torch.einsum("hk,khd->hd", probabilities, values.float()).to(

223 - q.dtype

225 + keys = (k_cache[pages, offsets].float() * k_scale).repeat_interleave(

226 + repeats, dim=1

227 )

228 + values = (v_cache[pages, offsets].float() * v_scale).repeat_interleave(

229 + repeats, dim=1

230 + )

231 + scores = torch.einsum("hd,khd->hk", q[row].float(), keys)

232 + probabilities = torch.softmax(scores * softmax_scale, dim=-1)

233 + output[row] = torch.einsum("hk,khd->hd", probabilities, values).to(q.dtype)

234 return output

235

236

@@ -906,22 +915,28 @@ def test_qsa_block_expansion_correctness() -> None:

915 "page_size",

916 "use_prefill_config",

917 "num_requests",

918 + "fp8",

919 ),

920 [

921 # Production page sizes from hybrid-cache block alignment: 784/800

922 # at TP4 and 1568/1600 at TP1/TP2 (no-MTP / MTP num_spec=3). Head

923 # splits are per-rank TP1/TP2/TP4; the largest batch runs both

924 # use_prefill_config variants.

915 - pytest.param(1, 24, 2, 1600, True, 2, id="tp1_r1"),

916 - pytest.param(16, 12, 1, 1600, True, 3, id="tp2_r16"),

917 - pytest.param(32, 6, 1, 800, True, 5, id="tp4_r32"),

918 - pytest.param(128, 24, 2, 1568, True, 7, id="tp1_r128"),

919 - pytest.param(257, 6, 1, 800, True, 13, id="tp4_r257"),

920 - pytest.param(513, 6, 1, 784, True, 17, id="tp4_r513"),

921 - pytest.param(700, 6, 1, 800, True, 23, id="tp4_r700"),

922 - pytest.param(1024, 24, 2, 1600, True, 33, id="tp1_r1024"),

923 - pytest.param(2048, 24, 2, 1600, True, 63, id="tp1_r2048_prefill"),

924 - pytest.param(2048, 24, 2, 1600, False, 63, id="tp1_r2048_uniform"),

925 + pytest.param(1, 24, 2, 1600, True, 2, False, id="tp1_r1"),

926 + pytest.param(16, 12, 1, 1600, True, 3, False, id="tp2_r16"),

927 + pytest.param(32, 6, 1, 800, True, 5, False, id="tp4_r32"),

928 + pytest.param(128, 24, 2, 1568, True, 7, False, id="tp1_r128"),

929 + pytest.param(257, 6, 1, 800, True, 13, False, id="tp4_r257"),

930 + pytest.param(513, 6, 1, 784, True, 17, False, id="tp4_r513"),

931 + pytest.param(700, 6, 1, 800, True, 23, False, id="tp4_r700"),

932 + pytest.param(1024, 24, 2, 1600, True, 33, False, id="tp1_r1024"),

933 + pytest.param(2048, 24, 2, 1600, True, 63, False, id="tp1_r2048_prefill"),

934 + pytest.param(2048, 24, 2, 1600, False, 63, False, id="tp1_r2048_uniform"),

935 + # fp8_e4m3 K/V caches on the TP1 head split.

936 + pytest.param(1, 24, 2, 1600, True, 2, True, id="tp1_r1_fp8"),

937 + pytest.param(128, 24, 2, 1568, True, 7, True, id="tp1_r128_fp8"),

938 + pytest.param(2048, 24, 2, 1600, True, 63, True, id="tp1_r2048_prefill_fp8"),

939 + pytest.param(2048, 24, 2, 1600, False, 63, True, id="tp1_r2048_uniform_fp8"),

940 ],

941 )

942 def test_qsa_sparse_paged_attention_correctness(

@@ -931,8 +946,18 @@ def test_qsa_sparse_paged_attention_correctness(

946 page_size: int,

947 use_prefill_config: bool,

948 num_requests: int,

949 + fp8: bool,

950 ) -> None:

951 + """QSA sparse paged attention matches the dense reference.

952 +

953 + fp8 only changes the K/V cache dtype (e4m3 with a per-tensor scale pair) and

954 + the scales; the reference dequantizes the same cache with those scales, so

955 + both paths compare the production kernel against the reference on identical

956 + inputs. fp8=True additionally covers the host-side scale folding.

957 + """

958 torch.manual_seed(2)

959 + # One QSA attention problem: bf16 Q and paged K/V, a packed selection with

960 + # the trailing count column, block table and row-to-request map.

961 head_dim = 256

962 num_selected_pages = 64

963 # Keep the newest page outside the synthetic top-k as causal headroom.

@@ -1013,8 +1038,18 @@ def test_qsa_sparse_paged_attention_correctness(

1038 indexer_budget,

1039 logical_indices,

1040 )

1016 - assert logical_indices.shape == (num_rows, selection_width + 1)

1017 - scale = q.shape[-1] ** -0.5

1041 +

1042 + scale = head_dim**-0.5

1043 +

1044 + if fp8:

1045 + # A fixed non-unit pair (k != v) exercises the host-side scale folding

1046 + # and catches a k/v swap; scales are host floats, as the layer exposes

1047 + # them. Stored values are the scaled ones, as reshape_and_cache does.

1048 + k_scale, v_scale = 0.5, 2.0

1049 + k_cache = (k_cache / k_scale).to(torch.float8_e4m3fn)

1050 + v_cache = (v_cache / v_scale).to(torch.float8_e4m3fn)

1051 + else:

1052 + k_scale, v_scale = 1.0, 1.0

1053

1054 actual = qsa_ops.qsa_sparse_paged_attention(

1055 q,

@@ -1024,6 +1059,8 @@ def test_qsa_sparse_paged_attention_correctness(

1059 block_table,

1060 token_to_req,

1061 use_prefill_config=use_prefill_config,

1062 + k_scale=k_scale,

1063 + v_scale=v_scale,

1064 output_gate=output_gate,

1065 )

1066 expected = _qsa_sparse_paged_attention_reference(

@@ -1034,6 +1071,8 @@ def test_qsa_sparse_paged_attention_correctness(

1071 block_table,

1072 token_to_req,

1073 scale,

1074 + k_scale=k_scale,

1075 + v_scale=v_scale,

1076 )

1077 expected = expected * torch.sigmoid(output_gate)

1078

vllm/models/qwen4_exp/nvidia/ops/qsa.py

+124 -16

@@ -4,15 +4,24 @@

5 from __future__ import annotations

7 +from functools import lru_cache

8 +

9 import torch

11 from vllm.model_executor.warmup.jit_warmup_triton_helper import (

12 TritonWarmupTensor,

13 triton_scalar_specialization_rep,

14 )

15 +from vllm.platforms import current_platform

16 from vllm.triton_utils import HAS_TRITON, tl, triton

19 +@lru_cache(maxsize=1)

20 +def _is_sm120() -> bool:

21 + """True on sm_120 (RTX PRO 6000 Blackwell): selects the sm_120 tuning table."""

22 + return current_platform.get_device_capability() == (12, 0)

23 +

24 +

25 @triton.jit(do_not_specialize=["num_rows", "num_requests"])

26 def _qsa_sparse_paged_gqa_splitk_kernel(

27 q_ptr,

@@ -24,6 +33,8 @@ def _qsa_sparse_paged_gqa_splitk_kernel(

33 partial_output_ptr,

34 partial_lse_ptr,

35 output_ptr,

36 + softmax_scale,

37 + output_scale,

38 output_gate_ptr,

39 stride_q_row,

40 stride_q_head,

@@ -52,6 +63,7 @@ def _qsa_sparse_paged_gqa_splitk_kernel(

63 NUM_TILES: tl.constexpr,

64 BLOCK_M: tl.constexpr,

65 BLOCK_N: tl.constexpr,

66 + IS_FP8: tl.constexpr,

67 ) -> None:

68 row = tl.program_id(0)

69 kv_head = tl.program_id(1)

@@ -81,7 +93,10 @@ def _qsa_sparse_paged_gqa_splitk_kernel(

93 max_value = tl.full((BLOCK_M,), -1.0e20, dtype=tl.float32)

94 normalizer = tl.zeros((BLOCK_M,), dtype=tl.float32)

95 accumulator = tl.zeros((BLOCK_M, HEAD_DIM), dtype=tl.float32)

84 - softmax_scale_log2: tl.constexpr = (HEAD_DIM**-0.5) * 1.4426950408889634

96 + # softmax_scale is the host-side attention scale (1/sqrt(head_dim), with the

97 + # fp8 K dequant scale already pre-multiplied in); convert to log2 units once

98 + # here for the exp2-based online softmax.

99 + score_scale = softmax_scale * 1.4426950408889634

100

101 tile_end = tl.minimum(NUM_TILES, tl.cdiv(tl.minimum(valid_count, TOPK), BLOCK_N))

102

@@ -129,15 +144,26 @@ def _qsa_sparse_paged_gqa_splitk_kernel(

144 mask=valid[:, None],

145 other=0.0,

146 )

147 + if IS_FP8:

148 + # e4m3 -> Q dtype is exact; keep the QK dot in Q's dtype (fp8 QK

149 + # measured slower here and less accurate).

150 + keys = keys.to(query.dtype)

151 scores = tl.dot(query, keys)

133 - # Scaling scores avoids re-quantizing a scaled query to BF16.

134 - scores *= softmax_scale_log2

152 + # Scaling scores avoids re-quantizing a scaled query to BF16; for fp8

153 + # caches the K dequant scale is already folded into softmax_scale on the

154 + # host.

155 + scores *= score_scale

156 scores = tl.where(valid[None, :], scores, -1.0e20)

157 next_max = tl.maximum(max_value, tl.max(scores, axis=1))

158 alpha = tl.math.exp2(max_value - next_max)

159 probabilities = tl.where(

160 valid[None, :], tl.math.exp2(scores - next_max[:, None]), 0.0

161 )

162 + if IS_FP8:

163 + # Dequant V to fp16 (not bf16) for the PV dot: P <= 1 (online

164 + # softmax) so fp16 has the range, its wider mantissa is more

165 + # accurate, and the fp8->fp16 upcast with an fp16 PV dot is faster.

166 + values = values.to(tl.float16)

167 accumulator = tl.dot(

168 probabilities.to(values.dtype),

169 values,

@@ -147,9 +173,14 @@ def _qsa_sparse_paged_gqa_splitk_kernel(

173 max_value = next_max

174

175 has_values = normalizer > 0

176 + # Fold the fp8 V dequant scale (output_scale, 1.0 for bf16) into a per-row

177 + # reciprocal normalizer, so the output is a per-row multiply rather than a

178 + # HEAD_DIM-wide scale. The split-K LSE below keeps the unscaled normalizer,

179 + # so the merge stays correct.

180 + inv_normalizer = output_scale / tl.maximum(normalizer, 1.0e-20)

181 normalized_output = tl.where(

182 has_values[:, None],

152 - accumulator / tl.maximum(normalizer[:, None], 1.0e-20),

183 + accumulator * inv_normalizer[:, None],

184 0.0,

185 )

186 output_mask = head_offsets[:, None] < GROUP_SIZE

@@ -439,18 +470,65 @@ def _compress_qsa_groups_kernel(

470 )

471

472

473 +def _select_sm120_config(

474 + base_programs: int, use_prefill_config: bool, is_fp8: bool

475 +) -> tuple[int, int, int]:

476 + """(block_n, target_splits, num_warps) retuned on sm_120 (RTX PRO 6000

477 + Blackwell), split by cache dtype: bf16 and fp8 favour different configs on

478 + sm_120, most visibly on the large-prefill region. Each entry is the fastest

479 + config that stays correct on its own path; the main win is more warps."""

480 + if is_fp8:

481 + if base_programs > 2048:

482 + return (32, 2, 1) if use_prefill_config else (64, 1, 2)

483 + if base_programs <= 24:

484 + return 64, 64, 8

485 + if base_programs <= 32:

486 + return 128, 8, 4

487 + if base_programs <= 64:

488 + return 64, 8, 4

489 + if base_programs <= 128:

490 + return 32, 4, 4

491 + if base_programs <= 256:

492 + return 128, 4, 4

493 + if base_programs <= 512:

494 + return 64, 4, 4

495 + return 32, 2, 1

496 + # bf16 K/V cache.

497 + if base_programs > 2048:

498 + return 64, 2, 2

499 + if base_programs <= 24:

500 + return 64, 64, 8

501 + if base_programs <= 64:

502 + return 64, 8, 8

503 + if base_programs <= 128:

504 + return 32, 8, 8

505 + if base_programs <= 256:

506 + return 32, 8, 1

507 + if base_programs <= 512:

508 + return 64, 4, 2

509 + return 32, 4, 1

510 +

511 +

512 def _select_config(

443 - num_rows: int, num_kv_heads: int, use_prefill_config: bool, num_columns: int

513 + num_rows: int,

514 + num_kv_heads: int,

515 + use_prefill_config: bool,

516 + num_columns: int,

517 + is_fp8: bool = False,

518 ) -> tuple[int, int, int, int]:

519 """Select (block_n, num_warps, num_tiles, num_splits) for the kernel.

520

447 - Tuned on GB300 for the Qwen3.8-Flash-Next TP1/TP2/TP4 shapes, keyed on

448 - base_programs = num_rows * num_kv_heads. The bp > 2048 region splits on

449 - use_prefill_config (capture-stable: at FULL-graph capture max_query_len is the

450 - uniform decode/verify length).

521 + Keyed on base_programs = num_rows * num_kv_heads. The bp > 2048 region splits

522 + on use_prefill_config (capture-stable: at FULL-graph capture max_query_len is

523 + the uniform decode/verify length). This default table was tuned on GB300;

524 + sm_120 (RTX PRO 6000 Blackwell) dispatches to _select_sm120_config instead.

525 """

526 base_programs = num_rows * num_kv_heads

453 - if base_programs > 2048:

527 + if _is_sm120():

528 + BLOCK_N, target_splits, num_warps = _select_sm120_config(

529 + base_programs, use_prefill_config, is_fp8

530 + )

531 + elif base_programs > 2048:

532 BLOCK_N, target_splits, num_warps = (

533 (32, 1, 1) if use_prefill_config else (64, 1, 2)

534 )

@@ -483,10 +561,17 @@ def qsa_sparse_paged_attention(

561 token_to_req: torch.Tensor,

562 use_prefill_config: bool,

563 out: torch.Tensor | None = None,

564 + k_scale: float | None = None,

565 + v_scale: float | None = None,

566 *,

567 output_gate: torch.Tensor,

568 ) -> torch.Tensor:

489 - """Run sparse GQA directly over paged BF16 K/V caches.

569 + """Run sparse GQA directly over paged BF16 or FP8-e4m3 K/V caches.

570 +

571 + With fp8 caches, k_scale/v_scale are the layer's per-tensor dequant scales

572 + as host floats (e.g. layer._k_scale_float/_v_scale_float): k_scale is

573 + pre-multiplied into the softmax scale and v_scale becomes the kernel's

574 + output scale, so the kernel needs no device scale buffers.

575

576 logical_indices is the PACKED selection buffer: [rows, selection_width + 1]

577 with the trailing column holding each row's valid-entry count (written by

@@ -510,7 +595,19 @@ def qsa_sparse_paged_attention(

595 raise ValueError("QSA sparse attention requires valid grouped-query heads")

596 head_dim = q.shape[2]

597 assert head_dim >= 16 and (head_dim & (head_dim - 1)) == 0

513 - assert q.dtype == k_cache.dtype == v_cache.dtype == torch.bfloat16

598 + assert q.dtype == torch.bfloat16

599 + assert k_cache.dtype == v_cache.dtype

600 + is_fp8 = k_cache.dtype == torch.float8_e4m3fn

601 + if is_fp8:

602 + assert k_scale is not None and v_scale is not None

603 + # Host pre-multiply: fold the K dequant scale into the attention scale

604 + # and pass V's dequant scale as the kernel's output scale.

605 + softmax_scale = (head_dim**-0.5) * float(k_scale)

606 + output_scale = float(v_scale)

607 + else:

608 + assert k_cache.dtype == torch.bfloat16

609 + softmax_scale = head_dim**-0.5

610 + output_scale = 1.0

611 assert logical_indices.dtype == block_table.dtype == torch.int32

612 assert token_to_req.dtype == torch.int32

613 assert q.device == k_cache.device == v_cache.device

@@ -535,7 +632,7 @@ def qsa_sparse_paged_attention(

632 block_m = triton.next_power_of_2(group_size)

633 selection_width = logical_indices.shape[1] - 1 # trailing column is the count

634 block_n, partial_warps, num_tiles, num_splits = _select_config(

538 - q.shape[0], k_cache.shape[2], use_prefill_config, selection_width

635 + q.shape[0], k_cache.shape[2], use_prefill_config, selection_width, is_fp8

636 )

637

638 # Split=1 writes output directly and compiles out all workspace accesses.

@@ -565,6 +662,8 @@ def qsa_sparse_paged_attention(

662 partial_output,

663 partial_lse,

664 out,

665 + softmax_scale,

666 + output_scale,

667 output_gate_view,

668 q.stride(0),

669 q.stride(1),

@@ -593,6 +692,7 @@ def qsa_sparse_paged_attention(

692 NUM_TILES=num_tiles,

693 BLOCK_M=block_m,

694 BLOCK_N=block_n,

695 + IS_FP8=is_fp8,

696 num_warps=partial_warps,

697 num_stages=2,

698 )

@@ -630,13 +730,18 @@ def warmup_qsa_sparse_paged_attention(

730

731 head_dim = kv_cache.shape[-1] // 2

732 key_cache, value_cache = kv_cache.transpose(1, 2).split(head_dim, dim=-1)

733 + # An fp8 cache is allocated as uint8 and viewed as e4m3 at attention time.

734 + is_fp8 = kv_cache.dtype == torch.uint8

735 + cache_dtype = torch.float8_e4m3fn if is_fp8 else key_cache.dtype

736 num_kv_heads = key_cache.shape[2]

737 group_size = num_query_heads // num_kv_heads

738 block_m = triton.next_power_of_2(group_size)

739

740 # Every config the dispatch can pick for this group size.

741 profiles = {

639 - _select_config(num_rows, num_kv_heads, use_prefill_config, selection_width)

742 + _select_config(

743 + num_rows, num_kv_heads, use_prefill_config, selection_width, is_fp8

744 + )

745 for num_rows in range(1, 8193)

746 for use_prefill_config in (False, True)

747 }

@@ -650,12 +755,12 @@ def warmup_qsa_sparse_paged_attention(

755 torch.bfloat16, shape=(num_rows, num_query_heads, head_dim)

756 )

757 k_cache_ptr = TritonWarmupTensor(

653 - key_cache.dtype,

758 + cache_dtype,

759 shape=tuple(key_cache.shape),

760 strides=tuple(key_cache.stride()),

761 )

762 v_cache_ptr = TritonWarmupTensor(

658 - value_cache.dtype,

763 + cache_dtype,

764 shape=tuple(value_cache.shape),

765 strides=tuple(value_cache.stride()),

766 )

@@ -701,6 +806,8 @@ def warmup_qsa_sparse_paged_attention(

806 partial_output_ptr,

807 partial_lse_ptr,

808 output_ptr,

809 + 1.0,

810 + 1.0,

811 output_gate_ptr,

812 row_stride,

813 head_stride,

@@ -729,6 +836,7 @@ def warmup_qsa_sparse_paged_attention(

836 NUM_TILES=num_tiles,

837 BLOCK_M=block_m,

838 BLOCK_N=block_n,

839 + IS_FP8=is_fp8,

840 num_warps=warps,

841 num_stages=2,

842 grid=(num_rows, num_kv_heads, num_splits),

vllm/models/qwen4_exp/nvidia/qsa.py

+106 -11

@@ -24,6 +24,7 @@ from vllm.model_executor.layers.quantization import QuantizationConfig

24 from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding, get_rope

25 from vllm.model_executor.models.qwen3_next import Qwen3NextAttention

26 from vllm.platforms import current_platform

27 +from vllm.platforms.interface import DeviceCapability

28 from vllm.transformers_utils.configs.qwen4_exp import (

29 Qwen4ExpTextConfig,

30 )

@@ -64,7 +65,38 @@ class Qwen4ExpQSAFlashAttentionBackend(FlashAttentionBackend):

65 """FullAttentionSpec backend used by the merged QSA owner."""

67 supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16]

67 - supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = ["auto", "bfloat16"]

68 + # fp8/fp8_e4m3: e4m3 bytes in a uint8 cache, written by reshape_and_cache

69 + # with the layer's per-tensor scales and dequantized on load inside the QSA

70 + # Triton kernel. flash-attn never runs over this cache, so its fp8 probe

71 + # does not apply (see supports_kv_cache_dtype and the impl constructor).

72 + supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [

73 + "auto",

74 + "bfloat16",

75 + "fp8",

76 + "fp8_e4m3",

77 + ]

78 +

79 + @classmethod

80 + def supports_kv_cache_dtype(cls, kv_cache_dtype: CacheDType | None) -> bool:

81 + return kv_cache_dtype is None or kv_cache_dtype in cls.supported_kv_cache_dtypes

82 +

83 + @classmethod

84 + def supports_combination(

85 + cls,

86 + head_size: int,

87 + dtype: torch.dtype,

88 + kv_cache_dtype: CacheDType | None,

89 + block_size: int | None,

90 + use_mla: bool,

91 + has_sink: bool,

92 + use_sparse: bool,

93 + use_mm_prefix: bool,

94 + device_capability: DeviceCapability,

95 + ) -> str | None:

96 + # QSA dequantizes the fp8 KV in its own Triton kernel and never runs

97 + # flash-attn over the quantized cache, so the parent's fp8-KV rejection

98 + # does not apply and every combination it is handed is accepted here.

99 + return None

100

101 @staticmethod

102 def get_name() -> str:

@@ -98,16 +130,53 @@ class Qwen4ExpQSAFlashAttentionImpl(FlashAttentionImpl):

130 supports_dcp: bool = False

131 supports_pcp: bool = False

132

101 - def __init__(self, *args, **kwargs) -> None:

102 - super().__init__(*args, **kwargs)

133 + def __init__(

134 + self,

135 + num_heads: int,

136 + head_size: int,

137 + scale: float,

138 + num_kv_heads: int,

139 + alibi_slopes: list[float] | None,

140 + sliding_window: int | None,

141 + kv_cache_dtype: str,

142 + logits_soft_cap: float | None = None,

143 + attn_type: AttentionType = AttentionType.DECODER,

144 + kv_sharing_target_layer_name: str | None = None,

145 + sinks: torch.Tensor | None = None,

146 + ) -> None:

147 + # The parent constructor probes flash-attn for quantized-KV support and

148 + # raises where it is unavailable (sm120), but QSA dequantizes fp8 inside

149 + # its own Triton kernel and never runs flash-attn over the cache. Hand

150 + # the parent "auto" for that probe and restore the real dtype afterwards:

151 + # the parent only uses it there, and do_kv_cache_update reads the

152 + # attribute at call time.

153 + real_kv_cache_dtype = kv_cache_dtype

154 + if kv_cache_dtype in ("fp8", "fp8_e4m3"):

155 + kv_cache_dtype = "auto"

156 + super().__init__(

157 + num_heads,

158 + head_size,

159 + scale,

160 + num_kv_heads,

161 + alibi_slopes,

162 + sliding_window,

163 + kv_cache_dtype,

164 + logits_soft_cap,

165 + attn_type,

166 + kv_sharing_target_layer_name,

167 + sinks,

168 + )

169 + self.kv_cache_dtype = real_kv_cache_dtype

170 if not is_flash_attn_varlen_func_available():

171 raise NotImplementedError("Qwen4Exp QSA requires FlashAttention")

172 if self.dcp_world_size != 1:

173 raise NotImplementedError(

174 "Qwen4Exp QSA does not support decode context parallelism"

175 )

109 - if self.kv_cache_dtype not in ("auto", "bfloat16"):

110 - raise NotImplementedError("Qwen4Exp QSA requires a BF16 main KV cache")

176 + if self.kv_cache_dtype not in ("auto", "bfloat16", "fp8", "fp8_e4m3"):

177 + raise NotImplementedError(

178 + "Qwen4Exp QSA requires a BF16 or FP8-e4m3 main KV cache"

179 + )

180 self.supports_quant_query_input = False

181

182 def forward_qsa(

@@ -144,8 +213,23 @@ class Qwen4ExpQSAFlashAttentionImpl(FlashAttentionImpl):

213 logical_indices = topk_buffer[:num_tokens]

214 token_to_req = token_to_req[:num_tokens]

215 key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_size, dim=-1)

147 - if key_cache.dtype != torch.bfloat16 or query.dtype != torch.bfloat16:

148 - raise NotImplementedError("Qwen4Exp QSA requires BF16 Q/K/V")

216 + k_scale = v_scale = None

217 + if self.kv_cache_dtype in ("fp8", "fp8_e4m3"):

218 + # The cache is allocated as uint8; reinterpret the e4m3 bytes

219 + # (same itemsize, so shape and strides are preserved).

220 + key_cache = key_cache.view(torch.float8_e4m3fn)

221 + value_cache = value_cache.view(torch.float8_e4m3fn)

222 + # Host-side per-tensor dequant scales (Python floats), as used by

223 + # other host-scale backends; folded into the kernel's scales.

224 + k_scale = layer._k_scale_float

225 + v_scale = layer._v_scale_float

226 + if query.dtype != torch.bfloat16 or key_cache.dtype not in (

227 + torch.bfloat16,

228 + torch.float8_e4m3fn,

229 + ):

230 + raise NotImplementedError(

231 + "Qwen4Exp QSA requires BF16 Q and BF16 or FP8-e4m3 K/V"

232 + )

233

234 from .ops.qsa import qsa_sparse_paged_attention

235

@@ -158,6 +242,8 @@ class Qwen4ExpQSAFlashAttentionImpl(FlashAttentionImpl):

242 token_to_req,

243 use_prefill_config,

244 output[:num_tokens],

245 + k_scale=k_scale,

246 + v_scale=v_scale,

247 output_gate=output_gate[:num_tokens],

248 )

249 return output

@@ -185,8 +271,15 @@ class Qwen4ExpQSAAttention(Qwen3NextAttention, AttentionLayerBase):

271 raise ValueError("Qwen4Exp QSA requires a paged KV cache")

272 if model_config.dtype != torch.bfloat16:

273 raise NotImplementedError("Qwen4Exp QSA currently requires BF16")

188 - if cache_config.cache_dtype not in ("auto", "bfloat16"):

189 - raise NotImplementedError("Qwen4Exp QSA requires a BF16 main KV cache")

274 + if cache_config.cache_dtype not in (

275 + "auto",

276 + "bfloat16",

277 + "fp8",

278 + "fp8_e4m3",

279 + ):

280 + raise NotImplementedError(

281 + "Qwen4Exp QSA requires a BF16 or FP8-e4m3 main KV cache"

282 + )

283 if getattr(quant_config, "kv_cache_scheme", None) is not None:

284 raise NotImplementedError("Qwen4Exp QSA does not support KV quantization")

285 parallel_config = vllm_config.parallel_config

@@ -285,8 +378,10 @@ class Qwen4ExpQSAAttention(Qwen3NextAttention, AttentionLayerBase):

378 self.kv_cache_torch_dtype = kv_cache_dtype_str_to_dtype(

379 self.kv_cache_dtype, model_config

380 )

288 - if self.kv_cache_torch_dtype != torch.bfloat16:

289 - raise NotImplementedError("Qwen4Exp QSA requires BF16 cache storage")

381 + if self.kv_cache_torch_dtype not in (torch.bfloat16, torch.uint8):

382 + raise NotImplementedError(

383 + "Qwen4Exp QSA requires BF16 or FP8-e4m3 (uint8) cache storage"

384 + )

385 self.kv_sharing_target_layer_name = None

386 self.kv_cache = torch.tensor([])

387 set_default_quant_scales(self, register_buffer=True)