Skip to content

Commit 7d235cf

Browse files
authored
[TRTLLM-10030][chore] promote SampleState to TypeVar + typing fixes (NVIDIA#11281)
Signed-off-by: ixlmar <206748156+ixlmar@users.noreply.github.com>
1 parent eae480b commit 7d235cf

File tree

4 files changed

+122
-78
lines changed

4 files changed

+122
-78
lines changed

tensorrt_llm/_torch/pyexecutor/llm_request.py

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,14 @@
11
from copy import copy, deepcopy
22
from dataclasses import dataclass
3-
from typing import Any, Dict, List, Optional, Union
3+
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
44

55
import torch
66

77
import tensorrt_llm.bindings
8+
9+
if TYPE_CHECKING:
10+
from tensorrt_llm._torch.pyexecutor.sampler import Strategy
11+
812
from tensorrt_llm._torch.shared_tensor import SharedTensorContainer
913
from tensorrt_llm.bindings import executor as tllm_executor
1014
from tensorrt_llm.executor.result import TokenLogprobs
@@ -583,6 +587,8 @@ def __init__(
583587
additional_outputs=additional_outputs)
584588
self.child_requests = []
585589

590+
self._py_sampling_strategy: "Strategy | None" = None
591+
586592
self._py_embedding_bias_1d: Optional[torch.Tensor] = None
587593
if hasattr(self, 'embedding_bias') and self.embedding_bias is not None:
588594
# Pre-squeeze to 1D if needed (remove batch dimension)

0 commit comments

Comments
 (0)