A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 6, 00:45 UTC
[ROCm][Test] Fix ROCM_AITER_UNIFIED_ATTN attn+quant fusion test (#37640)
Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com>
vllmellm committed
Mar 25, 2026 at 13:06 UTC
42e95479761a055d57c246a9bbb5121b9605951d
1 file changed
+7
-1
tests/compile/passes/test_fusion_attn.py
+7
-1
@@ -53,6 +53,7 @@ class AttentionQuantPatternModel(torch.nn.Module):
53
kv_cache_dtype: torch.dtype,
54
device: torch.device,
55
vllm_config: VllmConfig,
56
+ block_size: int,
57
**kwargs,
58
):
59
super().__init__()
@@ -74,7 +75,7 @@ class AttentionQuantPatternModel(torch.nn.Module):
75
self.attn._k_scale = self.attn._k_scale.to(device)
76
self.attn._v_scale = self.attn._v_scale.to(device)
77
77
- self.block_size = 16
78
+ self.block_size = block_size
79
80
# Initialize attn MetadataBuilder
81
self.builder = self.attn.attn_backend.get_builder_cls()(
@@ -299,6 +300,9 @@ def test_attention_quant_pattern(
300
torch.set_default_dtype(dtype)
301
torch.manual_seed(42)
302
303
+ backend_cls = backend.get_class()
304
+ block_size = backend_cls.get_preferred_block_size(16)
305
+
306
model_config = ModelConfig(
307
model=model_name,
308
max_model_len=2048,
@@ -342,6 +346,7 @@ def test_attention_quant_pattern(
346
kv_cache_dtype=FP8_DTYPE,
347
device=device,
348
vllm_config=vllm_config_unfused,
349
+ block_size=block_size,
350
)
351
model_unfused = model_unfused.to(device)
352
result_unfused_0 = model_unfused(q, k, v) # noqa: F841 HACK: See #131044
@@ -370,6 +375,7 @@ def test_attention_quant_pattern(
375
device=device,
376
vllm_config=vllm_config,
377
w=model_unfused.w,
378
+ block_size=block_size,
379
)
380
model_fused = model_fused.to(device)
381