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 5, 13:30 UTC
[Perf][GDN] Align TMA usage with upstream FLA (#38981)
Signed-off-by: Artem Perevedentsev <aperevedents@nvidia.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Artem Perevedentsev committed
Apr 4, 2026 at 19:38 UTC
99e5539a6701e77b0586e7d5a1fdddf19a1d42f2
1 file changed
+7
-3
vllm/model_executor/layers/fla/ops/utils.py
+7
-3
@@ -154,9 +154,13 @@ is_nvidia_hopper = is_nvidia and (
154
)
155
use_cuda_graph = is_nvidia and os.environ.get("FLA_USE_CUDA_GRAPH", "0") == "1"
156
is_gather_supported = hasattr(triton.language, "gather")
157
-is_tma_supported = (is_nvidia and torch.cuda.get_device_capability(0)[0] >= 9) and (
158
- hasattr(triton.language, "_experimental_make_tensor_descriptor")
159
- or hasattr(triton.language, "make_tensor_descriptor")
157
+is_tma_supported = (
158
+ is_nvidia_hopper
159
+ and os.getenv("FLA_USE_TMA", "0") == "1"
160
+ and (
161
+ hasattr(triton.language, "_experimental_make_tensor_descriptor")
162
+ or hasattr(triton.language, "make_tensor_descriptor")
163
+ )
164
)
165
166