From bc3ddb198a0d5eb864bf484c4719864b671cc93b Mon Sep 17 00:00:00 2001 From: trienebutdiene <59642603+kimseunghyun-kr@users.noreply.github.com> Date: Tue, 13 Oct 2026 15:30:19 +0800 Subject: [PATCH] [Model][Qwen4Exp] FP8 QSA main KV cache on Ampere (sm_80, sm_86) (#60602) Signed-off-by: kimseunghyun Signed-off-by: kimseunghyun Signed-off-by: Thien Tran Co-authored-by: kimseunghyun Co-authored-by: Claude Opus 5.5 Co-authored-by: Thien Tran --- vllm/models/qwen4_exp/nvidia/ops/qsa.py | 69 +++++++++++++++++++------ vllm/models/qwen4_exp/nvidia/qsa.py | 21 +++++--- 2 files changed, 67 insertions(+), 23 deletions(-) diff --git a/vllm/models/qwen4_exp/nvidia/ops/qsa.py b/vllm/models/qwen4_exp/nvidia/ops/qsa.py index 401048e02c..c7fe601449 100644 --- a/vllm/models/qwen4_exp/nvidia/ops/qsa.py +++ b/vllm/models/qwen4_exp/nvidia/ops/qsa.py @@ -22,6 +22,28 @@ def _is_sm120() -> bool: return current_platform.get_device_capability() == (12, 0) +@triton.jit +def _maybe_dequant(tile): + if tile.dtype == tl.uint8: + # e4m3 bytes below sm_89 (no Triton fp8e4nv pointer type): sign to bit + # 15, exp+mantissa to bits 13..7 gives the fp16 of x * 2^-8. Exact for + # every finite code; the two NaN codes decode to +-480 after the host + # scale instead of NaN. + code = tile.to(tl.uint16) + bits = ((code & 0x80) << 8) | ((code & 0x7F) << 7) + return bits.to(tl.float16, bitcast=True) + elif tile.dtype == tl.float8e4nv: + return tile.to(tl.float16) + else: + return tile + + +@lru_cache(maxsize=1) +def _is_sm80() -> bool: + """True on sm_80 (A100/A30): selects the wide FP8 byte-decode tiles.""" + return current_platform.get_device_capability() == (8, 0) + + @lru_cache(maxsize=1) def _is_sm90() -> bool: """True on sm_90 (H100/H200/H20): selects the sm_90 tuning table.""" @@ -69,7 +91,6 @@ def _qsa_sparse_paged_gqa_splitk_kernel( NUM_TILES: tl.constexpr, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, - IS_FP8: tl.constexpr, ) -> None: row = tl.program_id(0) kv_head = tl.program_id(1) @@ -141,6 +162,9 @@ def _qsa_sparse_paged_gqa_splitk_kernel( mask=valid[None, :], other=0.0, ) + # e4m3 -> Q dtype is exact; keep the QK dot in Q's dtype (fp8 QK + # measured slower here and less accurate). + keys = _maybe_dequant(keys).to(query.dtype) values = tl.load( v_cache_ptr + safe_page[:, None] * stride_v_block @@ -150,10 +174,10 @@ def _qsa_sparse_paged_gqa_splitk_kernel( mask=valid[:, None], other=0.0, ) - if IS_FP8: - # e4m3 -> Q dtype is exact; keep the QK dot in Q's dtype (fp8 QK - # measured slower here and less accurate). - keys = keys.to(query.dtype) + # Dequant V to fp16 (not bf16) for the PV dot: P <= 1 (online + # softmax) so fp16 has the range, its wider mantissa is more + # accurate, and the fp8->fp16 upcast with an fp16 PV dot is faster. + values = _maybe_dequant(values) scores = tl.dot(query, keys) # Scaling scores avoids re-quantizing a scaled query to BF16; for fp8 # caches the K dequant scale is already folded into softmax_scale on the @@ -165,11 +189,6 @@ def _qsa_sparse_paged_gqa_splitk_kernel( probabilities = tl.where( valid[None, :], tl.math.exp2(scores - next_max[:, None]), 0.0 ) - if IS_FP8: - # Dequant V to fp16 (not bf16) for the PV dot: P <= 1 (online - # softmax) so fp16 has the range, its wider mantissa is more - # accurate, and the fp8->fp16 upcast with an fp16 PV dot is faster. - values = values.to(tl.float16) accumulator = tl.dot( probabilities.to(values.dtype), values, @@ -585,6 +604,10 @@ def _select_config( BLOCK_N, target_splits, num_warps = 64, 4, 2 else: BLOCK_N, target_splits, num_warps = 64, 1, 2 + if is_fp8 and _is_sm80(): + # Wider tiles measured faster for the byte-decode path on A100; + # unmeasured on sm_86, which keeps the table's tiles. + BLOCK_N, num_warps = 128, 4 num_tiles = triton.cdiv(num_columns, BLOCK_N) # Never more splits than tiles, never empty. num_splits = min(target_splits, num_tiles) @@ -607,7 +630,9 @@ def qsa_sparse_paged_attention( ) -> torch.Tensor: """Run sparse GQA directly over paged BF16 or FP8-e4m3 K/V caches. - With fp8 caches, k_scale/v_scale are the layer's per-tensor dequant scales + An fp8 cache is a native e4m3 view on sm_89+ and the raw uint8 storage + below sm_89 (no Triton fp8e4nv pointer type there); the kernel dequantizes + either on load. k_scale/v_scale are the layer's per-tensor dequant scales as host floats (e.g. layer._k_scale_float/_v_scale_float): k_scale is pre-multiplied into the softmax scale and v_scale becomes the kernel's output scale, so the kernel needs no device scale buffers. @@ -636,15 +661,21 @@ def qsa_sparse_paged_attention( assert head_dim >= 16 and (head_dim & (head_dim - 1)) == 0 assert q.dtype == torch.bfloat16 assert k_cache.dtype == v_cache.dtype - is_fp8 = k_cache.dtype == torch.float8_e4m3fn + is_fp8 = k_cache.dtype != torch.bfloat16 if is_fp8: assert k_scale is not None and v_scale is not None # Host pre-multiply: fold the K dequant scale into the attention scale # and pass V's dequant scale as the kernel's output scale. softmax_scale = (head_dim**-0.5) * float(k_scale) output_scale = float(v_scale) + if k_cache.dtype == torch.uint8: + # _maybe_dequant's byte decode yields x * 2^-8; fold 2^8 into the + # scales so the decode needs no multiply. + softmax_scale *= 256.0 + output_scale *= 256.0 + else: + assert k_cache.dtype == torch.float8_e4m3fn else: - assert k_cache.dtype == torch.bfloat16 softmax_scale = head_dim**-0.5 output_scale = 1.0 assert logical_indices.dtype == block_table.dtype == torch.int32 @@ -731,7 +762,6 @@ def qsa_sparse_paged_attention( NUM_TILES=num_tiles, BLOCK_M=block_m, BLOCK_N=block_n, - IS_FP8=is_fp8, num_warps=partial_warps, num_stages=2, ) @@ -768,9 +798,15 @@ def warmup_qsa_sparse_paged_attention( """Compile every production-reachable split-K/merge specialization.""" head_dim = kv_cache.shape[-1] // 2 key_cache, value_cache = kv_cache.transpose(1, 2).split(head_dim, dim=-1) - # An fp8 cache is allocated as uint8 and viewed as e4m3 at attention time. + # An fp8 cache is allocated as uint8 storage; warm the dtype the impl hands + # over: an e4m3 view on sm_89+, the raw bytes below. is_fp8 = kv_cache.dtype == torch.uint8 - cache_dtype = torch.float8_e4m3fn if is_fp8 else key_cache.dtype + fp8_dtype = ( + torch.float8_e4m3fn + if current_platform.has_device_capability(89) + else torch.uint8 + ) + cache_dtype = fp8_dtype if is_fp8 else key_cache.dtype num_kv_heads = key_cache.shape[2] group_size = num_query_heads // num_kv_heads block_m = triton.next_power_of_2(group_size) @@ -874,7 +910,6 @@ def warmup_qsa_sparse_paged_attention( NUM_TILES=num_tiles, BLOCK_M=block_m, BLOCK_N=block_n, - IS_FP8=is_fp8, num_warps=warps, num_stages=2, grid=(num_rows, num_kv_heads, num_splits), diff --git a/vllm/models/qwen4_exp/nvidia/qsa.py b/vllm/models/qwen4_exp/nvidia/qsa.py index 9f0934d5a0..9fb36673a5 100644 --- a/vllm/models/qwen4_exp/nvidia/qsa.py +++ b/vllm/models/qwen4_exp/nvidia/qsa.py @@ -297,10 +297,12 @@ class Qwen4ExpQSAFlashAttentionImpl(FlashAttentionImpl): key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_size, dim=-1) k_scale = v_scale = None if self.kv_cache_dtype in ("fp8", "fp8_e4m3"): - # The cache is allocated as uint8; reinterpret the e4m3 bytes - # (same itemsize, so shape and strides are preserved). - key_cache = key_cache.view(torch.float8_e4m3fn) - value_cache = value_cache.view(torch.float8_e4m3fn) + # The cache is allocated as uint8; on sm_89+ reinterpret the e4m3 + # bytes (same itemsize, so shape and strides are preserved). Below + # sm_89 the kernel byte-decodes the raw uint8 storage instead. + if current_platform.has_device_capability(89): + key_cache = key_cache.view(torch.float8_e4m3fn) + value_cache = value_cache.view(torch.float8_e4m3fn) # Host-side per-tensor dequant scales (Python floats), as used by # other host-scale backends; folded into the kernel's scales. k_scale = layer._k_scale_float @@ -308,6 +310,7 @@ class Qwen4ExpQSAFlashAttentionImpl(FlashAttentionImpl): if query.dtype != torch.bfloat16 or key_cache.dtype not in ( torch.bfloat16, torch.float8_e4m3fn, + torch.uint8, ): raise NotImplementedError( "Qwen4Exp QSA requires BF16 Q and BF16 or FP8-e4m3 K/V" @@ -492,9 +495,15 @@ class Qwen4ExpQSAAttention(Qwen3NextAttention, AttentionLayerBase): ) # One launch does the indexer prepare, the main QK-norm/RoPE/gate and # the main K/V cache write (see QSAIndexer.forward); otherwise all of - # them take the separate kernels. + # them take the separate kernels. The fused writer receives an fp8 + # pointer, which Triton cannot lower below sm_89. self.use_fused_qsa_prepare = ( - self.use_fused_qk_norm_rope_gate and self.indexer.use_fused_pre_indexer + self.use_fused_qk_norm_rope_gate + and self.indexer.use_fused_pre_indexer + and ( + self.kv_cache_torch_dtype != torch.uint8 + or current_platform.has_device_capability(89) + ) ) self.fuse_indexer_projection = vllm_config.lora_config is None if self.fuse_indexer_projection: