TensorRT-LLMs/tensorrt_llm/_torch/pyexecutor
liji-nv dca6397d1e
feat: Introduce UB allocator for pytorch flow (#3257)
* Instead of allocating UserBuffers at beginning of runtime, UB buffers
  are now managed with global allocator. The allocator will dynamically
assign free UB buffer or allocate new buffer for torch tensor. It makes
userbuffers easier to use.

* In common usecase, the Userbuffers will be allocated correctly during
  warm up stage. There is no dynamic allocation during inference.

* UB fusion pattern is rewroten using the new UB Allocator. It contains
  following passes:

1. Fuse Quant with allreduce, replace with UB impl, and insert a
   copy_to_userbuffers. Currently the normal allreduce still does not
   support FP8 quant. So this need to be done in UB pass
2. Convert all supported allreduce with UB and insert copy_to_userbuffers.
3. Fuse op before ar with the copy_to_userbuffers. So the op directly
   writes to the userbuffer
4. Remove userbuffers finalize if the output is connect to another UB
   allreduce.

Signed-off-by: Jin Li <59594262+liji-nv@users.noreply.github.com>
2025-04-08 18:39:49 +08:00
..
__init__.py Update TensorRT-LLM (#2755) 2025-02-11 03:01:00 +00:00
_util.py feat: no-cache attention in PyTorch workflow (#3085) 2025-04-05 01:54:32 +08:00
config.py feat: Optionally split MoE inputs into chunks to reduce GPU memory usage (#3104) 2025-04-01 16:07:02 +08:00
cuda_graph_runner.py Refactor imports inside tensorrt_llm._torch. (#3015) 2025-03-26 11:01:07 +08:00
decoder.py fix the py_decoding_iter update in decoder. (#3297) 2025-04-07 11:18:33 +08:00
distributed.py Only gather responses on rank 0 (#3040) 2025-03-24 21:54:51 -07:00
guided_decoder.py Update (#2978) 2025-03-23 16:39:35 +08:00
kv_cache_transceiver.py Update TensorRT-LLM (#2849) 2025-03-04 18:44:00 +08:00
layerwise_nvtx_marker.py Update TensorRT-LLM (#2849) 2025-03-04 18:44:00 +08:00
llm_request.py fix: fix the acceptance rate of pytorch workflow in trtllm-bench (#3240) 2025-04-03 15:12:24 +08:00
model_engine.py feat: Introduce UB allocator for pytorch flow (#3257) 2025-04-08 18:39:49 +08:00
py_executor_creator.py fix: Fix an error related to dummy request when MTP is used (#3146) 2025-04-03 11:08:12 +08:00
py_executor.py feat: Add option to run disaggregated serving without ctx servers,… (#3243) 2025-04-07 21:56:03 -04:00
resource_manager.py feat: Support PeftCacheManager in Torch (#3186) 2025-04-04 12:38:08 +08:00
scheduler.py Update TensorRT-LLM (#2936) 2025-03-18 21:25:19 +08:00