Batch-invariant weighted SwiGLU reads invalid memory at large input offsets

[arcusbuilds]

Describe the bug

_weighted_silu_mul_bounded_kernel uses 32-bit offsets for its input loads. The address wraps when an input row starts at element 2^31; the read then fails with an illegal memory access. vllm_fused_moe passes this kernel a bound computed as int64, with no size guard.

@NVIDIA/mcore-oncall

Steps/Code to reproduce bug

Requires CUDA, Triton, and about 6.5 GiB of free GPU memory:

import torch
from megatron.core.inference.moe.batch_invariant import weighted_silu_mul_bounded

rows, half_n = 262145, 4096
y = torch.zeros((rows, 2 * half_n), device="cuda", dtype=torch.bfloat16)
weights = torch.ones(rows, device="cuda", dtype=torch.float32)
y[-1] = 1
bound = torch.tensor(rows * half_n, device="cuda", dtype=torch.int64)
expected = (
    torch.nn.functional.silu(y[-1, :half_n].float()) * y[-1, half_n:].float()
).bfloat16()

out = weighted_silu_mul_bounded(y, weights, bound)
torch.cuda.synchronize()
assert torch.equal(out[-1], expected)

Expected behavior

out[-1] equals expected. On main, the kernel reports CUDA error: an illegal memory access was encountered.

Additional context

The unweighted SwiGLU kernel casts its row index to int64 before multiplying by the row width. The same cast on xoffset fixes this case without changing the activation arithmetic.