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 14, 05:30 UTC
[FlashInfer] Cache hyper params in metadata builder (#23732)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Woosuk Kwon committed
Aug 27, 2025 at 03:45 UTC
11eddf02f0234f79435d747f2d3dce117ab39aa1
1 file changed
+15
-15
vllm/v1/attention/backends/flashinfer.py
+15
-15
@@ -214,6 +214,10 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
214
# TODO: discard this for trtllm-gen backend
215
self.global_hyperparameters = infer_global_hyperparameters(
216
get_per_layer_parameters(vllm_config, layer_names, FlashInferImpl))
217
+ self.sm_scale = self.global_hyperparameters.sm_scale
218
+ self.window_left = self.global_hyperparameters.window_left
219
+ self.logits_soft_cap = self.global_hyperparameters.logits_soft_cap
220
+ self.has_sinks = self.global_hyperparameters.has_sinks
221
222
# Preparing persistent buffers (device-side)
223
self.paged_kv_indptr = torch.zeros(max_num_reqs + 1,
@@ -381,8 +385,6 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
385
)
386
387
# Check if any layer uses sinks (requires TRTLLM attention)
384
- has_sinks = self.global_hyperparameters.has_sinks
385
-
388
prefill_use_trtllm = use_trtllm_attention(self.num_qo_heads,
389
self.num_kv_heads,
390
num_prefill_tokens,
@@ -390,7 +392,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
392
self.cache_dtype,
393
self.q_data_type,
394
is_prefill=True,
393
- has_sinks=has_sinks)
395
+ has_sinks=self.has_sinks)
396
decode_use_trtllm = use_trtllm_attention(self.num_qo_heads,
397
self.num_kv_heads,
398
num_decode_tokens,
@@ -398,7 +400,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
400
self.cache_dtype,
401
self.q_data_type,
402
is_prefill=False,
401
- has_sinks=has_sinks)
403
+ has_sinks=self.has_sinks)
404
405
attn_metadata = FlashInferMetadata(
406
num_actual_tokens=num_actual_tokens,
@@ -433,9 +435,9 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
435
self.head_dim,
436
self.page_size,
437
causal=True,
436
- sm_scale=self.global_hyperparameters.sm_scale,
437
- window_left=self.global_hyperparameters.window_left,
438
- logits_soft_cap=self.global_hyperparameters.logits_soft_cap,
438
+ sm_scale=self.sm_scale,
439
+ window_left=self.window_left,
440
+ logits_soft_cap=self.logits_soft_cap,
441
q_data_type=self.q_data_type,
442
kv_data_type=self.kv_cache_dtype,
443
)
@@ -472,10 +474,9 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
474
self.head_dim,
475
self.page_size,
476
causal=True,
475
- sm_scale=self.global_hyperparameters.sm_scale,
476
- window_left=self.global_hyperparameters.window_left,
477
- logits_soft_cap=self.global_hyperparameters.
478
- logits_soft_cap,
477
+ sm_scale=self.sm_scale,
478
+ window_left=self.window_left,
479
+ logits_soft_cap=self.logits_soft_cap,
480
q_data_type=self.q_data_type,
481
kv_data_type=self.kv_cache_dtype,
482
)
@@ -525,10 +526,9 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
526
self.page_size,
527
# Disable flashinfer's pos encoding and use vllm's rope.
528
pos_encoding_mode="NONE",
528
- sm_scale=self.global_hyperparameters.sm_scale,
529
- window_left=self.global_hyperparameters.window_left,
530
- logits_soft_cap=self.global_hyperparameters.
531
- logits_soft_cap,
529
+ sm_scale=self.sm_scale,
530
+ window_left=self.window_left,
531
+ logits_soft_cap=self.logits_soft_cap,
532
q_data_type=self.q_data_type,
533
kv_data_type=self.kv_cache_dtype,
534
)