From 9511a3f8eec6992d8834ad85af855683bd74ba29 Mon Sep 17 00:00:00 2001
From: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com>
Date: Wed, 25 Feb 2026 21:01:10 -0500
Subject: [PATCH] [Bugfix] Fix AttributeError in SMControlContextManager
 (#35338)

Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
---
 vllm/v1/worker/gpu_ubatch_wrapper.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/vllm/v1/worker/gpu_ubatch_wrapper.py b/vllm/v1/worker/gpu_ubatch_wrapper.py
index 45ba1bef9..754f2981c 100644
--- a/vllm/v1/worker/gpu_ubatch_wrapper.py
+++ b/vllm/v1/worker/gpu_ubatch_wrapper.py
@@ -74,7 +74,7 @@ class SMControlContextManager:
             "SM control is currently only supported on CUDA"
         )
 
-        total_sms = num_compute_units(torch.cuda.current_device().index)
+        total_sms = num_compute_units(torch.cuda.current_device())
 
         assert comm_sms < total_sms
         self.total_sms = total_sms