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

[Bugfix] speedup nvfp4 kv for FMHA (#55031)

Browse Source 
Signed-off-by: Shiyang Chen <shiychen@nvidia.com>
Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
This commit is contained in:
sychen52
2026-09-08 17:25:42 +00:00
committed by GitHub
co-authored by Cyrus Leung
parent bcca76e7db
commit 1b2c591cd0
2 changed files with 67 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/compile/passes/test_fusion_attn.py
+55 -19
View File
@@ -119,12 +119,14 @@ class AttentionQuantPatternModel(torch.nn.Module):
max_blocks = (max(batch_spec.seq_lens) + self.block_size - 1) // self.block_size
num_blocks = batch_size * max_blocks
spec = AttentionSpec(
block_size=self.block_size,
num_kv_heads=self.num_kv_heads,
head_size=self.head_size,
dtype=self.attn.kv_cache_torch_dtype,
kv_quant_mode=get_kv_quant_mode(self.attn.kv_cache_dtype),
spec = self.attn.attn_backend.customize_spec(
AttentionSpec(
block_size=self.block_size,
num_kv_heads=self.num_kv_heads,
head_size=self.head_size,
dtype=self.attn.kv_cache_torch_dtype,
kv_quant_mode=get_kv_quant_mode(self.attn.kv_cache_dtype),
)
)
supported = get_supported_kv_cache_layouts([self.attn.attn_backend])
layout = resolve_kv_cache_layout(
@@ -272,6 +274,7 @@ elif current_platform.is_rocm():
"batch_size", [7, 256, 533] if current_platform.is_cuda() else [8]
)
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
@pytest.mark.parametrize("kv_cache_dtype", ["fp8", "nvfp4"])
@pytest.mark.parametrize(
"backend, model_name, model_class, custom_ops",
# Test attention+quant_fp8 fusion with custom and torch impls of QuantFP8
@@ -293,6 +296,7 @@ def test_attention_quant_pattern(
head_size: int,
batch_size: int,
dtype: torch.dtype,
kv_cache_dtype: str,
custom_ops: str,
model_name: str,
model_class: type[AttentionQuantPatternModel],
@@ -306,6 +310,16 @@ def test_attention_quant_pattern(
):
# This also captures the FP4 case
pytest.skip("FlashInfer attn fusion requires Blackwell and flashinfer")
if kv_cache_dtype == "nvfp4" and backend != AttentionBackendEnum.FLASHINFER:
pytest.skip("NVFP4 KV cache is only supported by FlashInfer")
# trtllm-gen has no NVFP4-KV kernel with NVFP4 output, so with an NVFP4 KV
# cache the FP4 output fusion must be rejected: the standalone
# scaled_fp4_quant stays in the graph and the result is unchanged. FP8 output
# fusion with NVFP4 KV is supported.
expect_fusion = not (
kv_cache_dtype == "nvfp4" and model_class.quant_key is kNvfp4Dynamic
)
custom_ops_list = custom_ops.split(",") if custom_ops else []
@@ -332,7 +346,7 @@ def test_attention_quant_pattern(
mode=CompilationMode.VLLM_COMPILE,
custom_ops=custom_ops_list,
),
cache_config=CacheConfig(cache_dtype="fp8"),
cache_config=CacheConfig(cache_dtype=kv_cache_dtype),
attention_config=AttentionConfig(backend=backend),
)
@@ -410,7 +424,7 @@ def test_attention_quant_pattern(
result_fused = compiled_fused(q, k, v)
if backend == AttentionBackendEnum.FLASHINFER:
if backend == AttentionBackendEnum.FLASHINFER and expect_fusion:
# With the Flashinfer backend after the 1st round of the forward
# pass, output quant scale should be loaded into the attn layer's
# _o_scale_float, the 2nd round should reuse the loaded
@@ -423,6 +437,8 @@ def test_attention_quant_pattern(
torch.testing.assert_close(
result_unfused, result_fused_2, atol=1e-2, rtol=1e-2
)
elif not expect_fusion:
assert compiled_fused.attn._o_scale_float is None
# Check attn fusion support
quant_key: QuantKey = model_class.quant_key
@@ -430,9 +446,14 @@ def test_attention_quant_pattern(
layer.impl.fused_output_quant_supported(quant_key)
for key, layer in vllm_config.compilation_config.static_forward_context.items()
]
assert sum(attn_fusion_supported) == len(attn_fusion_supported), (
"All layers should support attention fusion"
)
if expect_fusion:
assert sum(attn_fusion_supported) == len(attn_fusion_supported), (
"All layers should support attention fusion"
)
else:
assert sum(attn_fusion_supported) == 0, (
"No layer should support attention fusion"
)
# Check quantization ops in the graph before and after fusion
quant_op = (
@@ -441,9 +462,19 @@ def test_attention_quant_pattern(
else QUANT_OPS[quant_key]
)
# Note: for fp8, fully_replaced=False because query quant ops remain in graph.
# Only output quant ops are fused into attention.
test_backend.check_before_ops([quant_op], fully_replaced=quant_key is kNvfp4Dynamic)
if expect_fusion:
# Note: for fp8, fully_replaced=False because query quant ops remain in
# graph. Only output quant ops are fused into attention.
test_backend.check_before_ops(
[quant_op], fully_replaced=quant_key is kNvfp4Dynamic
)
else:
# Fusion rejected: the standalone output quant must survive untouched.
num_pre = test_backend.op_count(quant_op, before=True)
assert num_pre > 0, f"Op {quant_op} not found in pre-pass graph"
assert test_backend.op_count(quant_op) == num_pre, (
"Output quant op should remain when fusion is rejected"
)
# access the underlying `AttnQuantFusionPass` on the `LazyInitPass`
assert attn_pass.pass_.matched_count == sum(attn_fusion_supported)
@@ -459,9 +490,14 @@ def test_attention_quant_pattern(
assert attn_nodes_pre[0].kwargs.get("output_scale") is None, (
"Attention should not have output_scale before fusion"
)
assert attn_nodes_post[0].kwargs.get("output_scale") is not None, (
"Attention should have output_scale after fusion"
)
if expect_fusion:
assert attn_nodes_post[0].kwargs.get("output_scale") is not None, (
"Attention should have output_scale after fusion"
)
else:
assert attn_nodes_post[0].kwargs.get("output_scale") is None, (
"Attention should not have output_scale when fusion is rejected"
)
assert attn_nodes_pre[0].kwargs.get("output_block_scale") is None, (
"Attention should not have output_block_scale before fusion"
@@ -477,9 +513,9 @@ def test_attention_quant_pattern(
"The kv_cache_dummy_dep should be consistent before and after fusion"
)
if quant_key.dtype == FP8_DTYPE:
if quant_key.dtype == FP8_DTYPE or not expect_fusion:
assert attn_nodes_post[0].kwargs.get("output_block_scale") is None, (
"Attention should not have output_block_scale after FP8 fusion"
"Attention should not have output_block_scale unless FP4 fusion happened"
)
elif quant_key.dtype == FP4_DTYPE:
assert attn_nodes_post[0].kwargs.get("output_block_scale") is not None, (
vllm/v1/attention/backends/flashinfer.py
+12 -4
View File
@@ -1884,6 +1884,14 @@ class FlashInferImpl(AttentionImpl):
return self.cache_config.get_resolved_kv_cache_layout()
def fused_output_quant_supported(self, quant_key: QuantKey):
if quant_key == kNvfp4Dynamic and self.is_kvcache_nvfp4:
logger.warning_once(
"fuse_attn_quant: not registering the NVFP4-output attention "
"fusion with an NVFP4 KV cache (no such trtllm-gen kernel). "
"Layers whose o_proj takes NVFP4 input keep the unfused path "
"(FP8 -> BF16 -> o_proj quant); FP8-output fusion is unaffected."
)
return False
# XQA does not support FP8/NVFP4 output, so require trtllm-gen
# (SM100+) here. Without that we cannot fuse the output quant.
return (
@@ -2214,7 +2222,7 @@ class FlashInferImpl(AttentionImpl):
if needs_fp8_out_prefill:
output[
num_decode_tokens : num_decode_tokens + num_prefill_tokens
].copy_(out_prefill.to(output.dtype))
].copy_(out_prefill)
else:
assert isinstance(attn_metadata.prefill, TRTLLMPrefill)
# prefill_query may be non-contiguous or have degenerate strides
@@ -2322,7 +2330,7 @@ class FlashInferImpl(AttentionImpl):
if needs_fp8_out:
output[
num_decode_tokens : num_decode_tokens + num_prefill_tokens
].copy_(out[:num_prefill_tokens].to(output.dtype))
].copy_(out[:num_prefill_tokens])
if num_decode_tokens > 0:
decode_query = query[:num_decode_tokens]
@@ -2397,7 +2405,7 @@ class FlashInferImpl(AttentionImpl):
)
if needs_fp8_out:
output[:num_decode_tokens].copy_(out_decode.to(output.dtype))
output[:num_decode_tokens].copy_(out_decode)
else:
assert isinstance(attn_metadata.decode, FlashInferTrtllmAPIDecode)
# decode_query may be non-contiguous or have degenerate strides
@@ -2559,7 +2567,7 @@ class FlashInferImpl(AttentionImpl):
get_dcp_group(),
)
elif needs_fp8_out:
output[:num_decode_tokens].copy_(out.to(output.dtype))
output[:num_decode_tokens].copy_(out)
return output_padded
def do_kv_cache_update(