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
[Bugfix] Use a separate FlashInfer workspace buffer for trtllm-gen (#25520)
Benjamin Chislett committed
Sep 23, 2025 at 20:19 UTC
1983609239caaab24ab1ed2bfa2aa92e8c76c1b1
1 file changed
+12
-2
vllm/v1/attention/backends/flashinfer.py
+12
-2
@@ -48,6 +48,16 @@ FP4_DTYPE = torch.uint8
48
49
logger = init_logger(__name__)
50
51
+trtllm_gen_workspace_buffer = None
52
+
53
+
54
+def _get_trtllm_gen_workspace_buffer():
55
+ global trtllm_gen_workspace_buffer
56
+ if trtllm_gen_workspace_buffer is None:
57
+ trtllm_gen_workspace_buffer = torch.zeros(
58
+ FLASHINFER_WORKSPACE_BUFFER_SIZE, dtype=torch.uint8, device='cuda')
59
+ return trtllm_gen_workspace_buffer
60
+
61
62
@triton.jit
63
def _trtllm_prefill_attn_kvfp8_dequant(
@@ -862,7 +872,7 @@ class FlashInferImpl(AttentionImpl):
872
else:
873
# prefill_query may be non-contiguous
874
prefill_query = prefill_query.contiguous()
865
- workspace_buffer = prefill_wrapper._float_workspace_buffer
875
+ workspace_buffer = _get_trtllm_gen_workspace_buffer()
876
block_tables_prefill = attn_metadata.block_table_tensor[
877
num_decode_tokens:]
878
seq_lens_prefill = attn_metadata.seq_lens[num_decode_tokens:]
@@ -943,7 +953,7 @@ class FlashInferImpl(AttentionImpl):
953
else:
954
# decode_query may be non-contiguous
955
decode_query = decode_query.contiguous()
946
- workspace_buffer = decode_wrapper._float_workspace_buffer
956
+ workspace_buffer = _get_trtllm_gen_workspace_buffer()
957
block_tables_decode = attn_metadata.\
958
block_table_tensor[:num_decode_tokens]
959
seq_lens_decode = attn_metadata.seq_lens[:num_decode_tokens]