From 235d5b314fc7a7799521362aa9d1f2a09634d229 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sun, 17 May 2026 08:27:47 +0000 Subject: [PATCH] fix: fallback token indices allocation with verify+rebuild --- vllm/nvfp4_cutedsl.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/vllm/nvfp4_cutedsl.py b/vllm/nvfp4_cutedsl.py index 6a83d668..e662e930 100644 --- a/vllm/nvfp4_cutedsl.py +++ b/vllm/nvfp4_cutedsl.py @@ -130,6 +130,16 @@ class CuTeDSLMoERunner: self._token_indices = torch.arange( self.max_num_tokens, dtype=torch.int32 ).unsqueeze(1).expand(-1, self.top_k).contiguous().view(-1).to(self.device) + # Verify the tensor was transferred correctly (CuTeDSL JIT can corrupt GPU state) + _verify = self._token_indices[:8].cpu().tolist() + if _verify != [0, 0, 1, 1, 2, 2, 3, 3]: + # Fallback: allocate directly on GPU + self._token_indices = torch.zeros( + self.max_num_tokens * self.top_k, dtype=torch.int32, device=self.device + ) + for i in range(self.max_num_tokens): + self._token_indices[i * self.top_k] = i + self._token_indices[i * self.top_k + 1] = i self._expert_id_range = torch.arange( self.num_experts, dtype=torch.int32 ).to(self.device)