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

[Model][Qwen4Exp] FP8 QSA main KV cache on Ampere (sm_80, sm_86) (#60602)

Browse Source 
Signed-off-by: kimseunghyun <hellonypcteam@gmail.com>
Signed-off-by: kimseunghyun <incongnito12@gmail.com>
Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg>
Co-authored-by: kimseunghyun <hellonypcteam@gmail.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Co-authored-by: Thien Tran <gau.nernst@yahoo.com.sg>
This commit is contained in:
trienebutdiene
2026-10-13 15:30:19 +08:00
committed by GitHub
co-authored by kimseunghyun Claude Opus 5.5 Thien Tran
parent 4e4eab73ed
commit bc3ddb198a
2 changed files with 67 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vllm/models/qwen4_exp/nvidia/ops/qsa.py
+52 -17
View File
@@ -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),
vllm/models/qwen4_exp/nvidia/qsa.py
+15 -6
View File
@@ -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: