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 15, 18:30 UTC
[Neuron] trim attention kernel tests to fit trn1.2x instance (#14988)
Signed-off-by: Liangfu Chen <liangfc@amazon.com>
Liangfu Chen committed
Mar 18, 2025 at 00:05 UTC
53a0cf8b95beb0eaf1b374146b0f33e216fe82d0
1 file changed
+1
-1
tests/neuron/1_core/test_prefix_prefill.py
+1
-1
@@ -314,7 +314,7 @@ def get_active_block_tables(block_tables, query_lens, seq_lens, block_size,
314
315
# Test edge cases
316
(1, 128, 16, 1024, 4, 2, 16, False), # large decode batch
317
- (16, 4, 8, 8192, 48, 1, 128, True), # large prefill batch
317
+ (16, 4, 8, 1024, 4, 2, 128, True), # large prefill batch
318
(4, 12, 32, 2048, 16, 1, 32, True), # multi-head attention (MHA)
319
(4, 12, 32, 2048, 16, 16, 32, True), # multi-query attention (MQA)
320
])