[Bugfix] speedup nvfp4 kv for FMHA (#55031)
Signed-off-by: Shiyang Chen <shiychen@nvidia.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
sychen52 committed
Sep 8, 2026 at 10:25 UTC
1b2c591cd0c3bb5a85ac7f3d6cbaa2fa7df6bc7d
2 files changed
+67
-23
tests/compile/passes/test_fusion_attn.py
+55
-19
@@ -119,12 +119,14 @@ class AttentionQuantPatternModel(torch.nn.Module):
119
max_blocks = (max(batch_spec.seq_lens) + self.block_size - 1) // self.block_size
120
num_blocks = batch_size * max_blocks
121
122
- spec = AttentionSpec(
123
- block_size=self.block_size,
124
- num_kv_heads=self.num_kv_heads,
125
- head_size=self.head_size,
126
- dtype=self.attn.kv_cache_torch_dtype,
127
- kv_quant_mode=get_kv_quant_mode(self.attn.kv_cache_dtype),
122
+ spec = self.attn.attn_backend.customize_spec(
123
+ AttentionSpec(
124
+ block_size=self.block_size,
125
+ num_kv_heads=self.num_kv_heads,
126
+ head_size=self.head_size,
127
+ dtype=self.attn.kv_cache_torch_dtype,
128
+ kv_quant_mode=get_kv_quant_mode(self.attn.kv_cache_dtype),
129
+ )
130
)
131
supported = get_supported_kv_cache_layouts([self.attn.attn_backend])
132
layout = resolve_kv_cache_layout(
@@ -272,6 +274,7 @@ elif current_platform.is_rocm():
274
"batch_size", [7, 256, 533] if current_platform.is_cuda() else [8]
275
)
276
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
277
+@pytest.mark.parametrize("kv_cache_dtype", ["fp8", "nvfp4"])
278
@pytest.mark.parametrize(
279
"backend, model_name, model_class, custom_ops",
280
# Test attention+quant_fp8 fusion with custom and torch impls of QuantFP8
@@ -293,6 +296,7 @@ def test_attention_quant_pattern(
296
head_size: int,
297
batch_size: int,
298
dtype: torch.dtype,
299
+ kv_cache_dtype: str,
300
custom_ops: str,
301
model_name: str,
302
model_class: type[AttentionQuantPatternModel],
@@ -306,6 +310,16 @@ def test_attention_quant_pattern(
310
):
311
# This also captures the FP4 case
312
pytest.skip("FlashInfer attn fusion requires Blackwell and flashinfer")
313
+ if kv_cache_dtype == "nvfp4" and backend != AttentionBackendEnum.FLASHINFER:
314
+ pytest.skip("NVFP4 KV cache is only supported by FlashInfer")
315
+
316
+ # trtllm-gen has no NVFP4-KV kernel with NVFP4 output, so with an NVFP4 KV
317
+ # cache the FP4 output fusion must be rejected: the standalone
318
+ # scaled_fp4_quant stays in the graph and the result is unchanged. FP8 output
319
+ # fusion with NVFP4 KV is supported.
320
+ expect_fusion = not (
321
+ kv_cache_dtype == "nvfp4" and model_class.quant_key is kNvfp4Dynamic
322
+ )
323
324
custom_ops_list = custom_ops.split(",") if custom_ops else []
325
@@ -332,7 +346,7 @@ def test_attention_quant_pattern(
346
mode=CompilationMode.VLLM_COMPILE,
347
custom_ops=custom_ops_list,
348
),
335
- cache_config=CacheConfig(cache_dtype="fp8"),
349
+ cache_config=CacheConfig(cache_dtype=kv_cache_dtype),
350
attention_config=AttentionConfig(backend=backend),
351
)
352
@@ -410,7 +424,7 @@ def test_attention_quant_pattern(
424
425
result_fused = compiled_fused(q, k, v)
426
413
- if backend == AttentionBackendEnum.FLASHINFER:
427
+ if backend == AttentionBackendEnum.FLASHINFER and expect_fusion:
428
# With the Flashinfer backend after the 1st round of the forward
429
# pass, output quant scale should be loaded into the attn layer's
430
# _o_scale_float, the 2nd round should reuse the loaded
@@ -423,6 +437,8 @@ def test_attention_quant_pattern(
437
torch.testing.assert_close(
438
result_unfused, result_fused_2, atol=1e-2, rtol=1e-2
439
)
440
+ elif not expect_fusion:
441
+ assert compiled_fused.attn._o_scale_float is None
442
443
# Check attn fusion support
444
quant_key: QuantKey = model_class.quant_key
@@ -430,9 +446,14 @@ def test_attention_quant_pattern(
446
layer.impl.fused_output_quant_supported(quant_key)
447
for key, layer in vllm_config.compilation_config.static_forward_context.items()
448
]
433
- assert sum(attn_fusion_supported) == len(attn_fusion_supported), (
434
- "All layers should support attention fusion"
435
- )
449
+ if expect_fusion:
450
+ assert sum(attn_fusion_supported) == len(attn_fusion_supported), (
451
+ "All layers should support attention fusion"
452
+ )
453
+ else:
454
+ assert sum(attn_fusion_supported) == 0, (
455
+ "No layer should support attention fusion"
456
+ )
457
458
# Check quantization ops in the graph before and after fusion
459
quant_op = (
@@ -441,9 +462,19 @@ def test_attention_quant_pattern(
462
else QUANT_OPS[quant_key]
463
)
464
444
- # Note: for fp8, fully_replaced=False because query quant ops remain in graph.
445
- # Only output quant ops are fused into attention.
446
- test_backend.check_before_ops([quant_op], fully_replaced=quant_key is kNvfp4Dynamic)
465
+ if expect_fusion:
466
+ # Note: for fp8, fully_replaced=False because query quant ops remain in
467
+ # graph. Only output quant ops are fused into attention.
468
+ test_backend.check_before_ops(
469
+ [quant_op], fully_replaced=quant_key is kNvfp4Dynamic
470
+ )
471
+ else:
472
+ # Fusion rejected: the standalone output quant must survive untouched.
473
+ num_pre = test_backend.op_count(quant_op, before=True)
474
+ assert num_pre > 0, f"Op {quant_op} not found in pre-pass graph"
475
+ assert test_backend.op_count(quant_op) == num_pre, (
476
+ "Output quant op should remain when fusion is rejected"
477
+ )
478
479
# access the underlying `AttnQuantFusionPass` on the `LazyInitPass`
480
assert attn_pass.pass_.matched_count == sum(attn_fusion_supported)
@@ -459,9 +490,14 @@ def test_attention_quant_pattern(
490
assert attn_nodes_pre[0].kwargs.get("output_scale") is None, (
491
"Attention should not have output_scale before fusion"
492
)
462
- assert attn_nodes_post[0].kwargs.get("output_scale") is not None, (
463
- "Attention should have output_scale after fusion"
464
- )
493
+ if expect_fusion:
494
+ assert attn_nodes_post[0].kwargs.get("output_scale") is not None, (
495
+ "Attention should have output_scale after fusion"
496
+ )
497
+ else:
498
+ assert attn_nodes_post[0].kwargs.get("output_scale") is None, (
499
+ "Attention should not have output_scale when fusion is rejected"
500
+ )
501
502
assert attn_nodes_pre[0].kwargs.get("output_block_scale") is None, (
503
"Attention should not have output_block_scale before fusion"
@@ -477,9 +513,9 @@ def test_attention_quant_pattern(
513
"The kv_cache_dummy_dep should be consistent before and after fusion"
514
)
515
480
- if quant_key.dtype == FP8_DTYPE:
516
+ if quant_key.dtype == FP8_DTYPE or not expect_fusion:
517
assert attn_nodes_post[0].kwargs.get("output_block_scale") is None, (
482
- "Attention should not have output_block_scale after FP8 fusion"
518
+ "Attention should not have output_block_scale unless FP4 fusion happened"
519
)
520
elif quant_key.dtype == FP4_DTYPE:
521
assert attn_nodes_post[0].kwargs.get("output_block_scale") is not None, (
vllm/v1/attention/backends/flashinfer.py
+12
-4
@@ -1884,6 +1884,14 @@ class FlashInferImpl(AttentionImpl):
1884
return self.cache_config.get_resolved_kv_cache_layout()
1885
1886
def fused_output_quant_supported(self, quant_key: QuantKey):
1887
+ if quant_key == kNvfp4Dynamic and self.is_kvcache_nvfp4:
1888
+ logger.warning_once(
1889
+ "fuse_attn_quant: not registering the NVFP4-output attention "
1890
+ "fusion with an NVFP4 KV cache (no such trtllm-gen kernel). "
1891
+ "Layers whose o_proj takes NVFP4 input keep the unfused path "
1892
+ "(FP8 -> BF16 -> o_proj quant); FP8-output fusion is unaffected."
1893
+ )
1894
+ return False
1895
# XQA does not support FP8/NVFP4 output, so require trtllm-gen
1896
# (SM100+) here. Without that we cannot fuse the output quant.
1897
return (
@@ -2214,7 +2222,7 @@ class FlashInferImpl(AttentionImpl):
2222
if needs_fp8_out_prefill:
2223
output[
2224
num_decode_tokens : num_decode_tokens + num_prefill_tokens
2217
- ].copy_(out_prefill.to(output.dtype))
2225
+ ].copy_(out_prefill)
2226
else:
2227
assert isinstance(attn_metadata.prefill, TRTLLMPrefill)
2228
# prefill_query may be non-contiguous or have degenerate strides
@@ -2322,7 +2330,7 @@ class FlashInferImpl(AttentionImpl):
2330
if needs_fp8_out:
2331
output[
2332
num_decode_tokens : num_decode_tokens + num_prefill_tokens
2325
- ].copy_(out[:num_prefill_tokens].to(output.dtype))
2333
+ ].copy_(out[:num_prefill_tokens])
2334
2335
if num_decode_tokens > 0:
2336
decode_query = query[:num_decode_tokens]
@@ -2397,7 +2405,7 @@ class FlashInferImpl(AttentionImpl):
2405
)
2406
2407
if needs_fp8_out:
2400
- output[:num_decode_tokens].copy_(out_decode.to(output.dtype))
2408
+ output[:num_decode_tokens].copy_(out_decode)
2409
else:
2410
assert isinstance(attn_metadata.decode, FlashInferTrtllmAPIDecode)
2411
# decode_query may be non-contiguous or have degenerate strides
@@ -2559,7 +2567,7 @@ class FlashInferImpl(AttentionImpl):
2567
get_dcp_group(),
2568
)
2569
elif needs_fp8_out:
2562
- output[:num_decode_tokens].copy_(out.to(output.dtype))
2570
+ output[:num_decode_tokens].copy_(out)
2571
return output_padded
2572
2573
def do_kv_cache_update(