Neuron: eager embedding/gather without dtype casts or host scalar operands

Both drain the device stream in torch-neuronx (a cast round-trips through the
host, a scalar operand is a host->device copy): int32 ids index directly and
out-of-range rows are cleared with masked_fill. Tiny cp=2 nextn=5 on device:
9/9 golden tokens.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

deepinfer/model_executor/layers/tilelang_ops/neuron/gather.py

@@ -34,9 +34,9 @@ def gather_tl(D: int, dtype: torch.dtype, index_dtype: torch.dtype, x_stride: in
         if n == 0:
             y.zero_()
             return
-        valid = (index >= 0) & (index < n)
-        safe = torch.where(valid, index, torch.zeros_like(index))
-        rows = _gather_rows(x[:, :D], safe)
-        y.copy_(torch.where(valid.view(m, 1), rows, torch.zeros((), dtype=rows.dtype, device=rows.device)))
+        # No scalar tensor operands: on the neuron device a host scalar is a host->device copy, which drains the
+        # stream.
+        rows = _gather_rows(x[:, :D], index.clamp(0, n - 1))
+        y.copy_(rows.masked_fill(((index < 0) | (index >= n)).view(m, 1), 0))
 
     return _gather

deepinfer/model_executor/layers/vocab_parallel_embedding.py

@@ -28,11 +28,12 @@ class VocabParallelEmbedding(torch.nn.Module):
             # Eagerly the zero-filling gather is ~5 NEFF launches, and its scalar operand is a host->device copy,
             # which drains the stream. One lookup instead; out-of-range ids still read zeros unless the caller's
             # ids are always valid (ids_in_range: DSpark's Markov embedding of sampled draft tokens).
+            # No dtype casts either: eagerly they round-trip through the host (int32 ids index directly).
             n = self.weight.shape[0]
-            rows = torch.nn.functional.embedding(input_.long().clamp(0, n - 1), self.weight)
+            rows = torch.nn.functional.embedding(input_.clamp(0, n - 1), self.weight)
             if getattr(self, "ids_in_range", False):
                 return rows
-            return rows * ((input_ >= 0) & (input_ < n)).unsqueeze(-1).to(rows.dtype)
+            return rows.masked_fill(((input_ < 0) | (input_ >= n)).unsqueeze(-1), 0)
         output = gather(self.weight, input_)
         return output