fix: guard SM90 fused logp label indexing

Signed-off-by: Dnoob <dxpouo@gmail.com>

Author: Dnoob

csrc/cuda/fused_logp_sm90.cu

@@ -2,6 +2,7 @@
 // Copyright (c) 2026 RL-Kernel Contributors
 
 #include "../utils/tma_utils.cuh"
+#include <cstdint>
 #include <math_constants.h>
 #include <torch/extension.h>
 #include <tuple>
@@ -140,8 +141,13 @@ __global__ void fused_logp_online_tma_kernel(
                 logsum_out[row_idx] = log_sum;
             }
             int label_idx = labels[row_idx];
-            float label_val = __bfloat162float(logits_gmem[row_idx * vocab_size + label_idx]);
-            output_logp[row_idx] = label_val - row_max - log_sum;
+            if (label_idx >= 0 && label_idx < vocab_size) {
+                int64_t label_offset = static_cast<int64_t>(row_idx) * vocab_size + label_idx;
+                float label_val = __bfloat162float(logits_gmem[label_offset]);
+                output_logp[row_idx] = label_val - row_max - log_sum;
+            } else {
+                output_logp[row_idx] = 0.0f;
+            }
         }
     }
 }

tests/test_fused_logp_backward.py

@@ -459,6 +459,18 @@ def test_partial_last_tile_matches_reference(self):
             rtol=1e-5,
         )
 
+    def test_invalid_labels_return_zero(self):
+        logits, token_ids = _make_inputs(4, 520, dtype=torch.bfloat16, seed=5)
+        labels = token_ids.to(torch.int32)
+        labels[0] = -1
+        labels[1] = logits.size(1)
+
+        logp, _, _ = _C.fused_logp_sm90_with_lse(logits, labels)
+        assert torch.equal(logp[:2], torch.zeros_like(logp[:2]))
+
+        ref = selected_logprobs_reference(logits, token_ids)
+        assert torch.allclose(logp[2:], ref[2:], atol=1e-3, rtol=1e-5)
+
     def test_unaligned_row_stride_uses_generic_fallback(self, monkeypatch):
         from rl_engine.kernels.ops.cuda.loss.logp import FusedLogpSM90Op