Fix int32 truncation of 64-bit ssd_row_addrs in unrolled forward path

[EddyLXJ]

Summary:
X-link: https://github.com/facebookresearch/FBGEMM/pull/2674

CONTEXT: KVZCH SSD TBE forward kernel deterministically page-faults on AMD MI350X with HSA Memory access faults at addresses showing a bimodal distribution: low (e.g., 0x4989000, 0x670ac000) AND sign-extended high (e.g., 0xfffff5295000, 0xffffe62d4000). Reproduced in 7+ MAST jobs.

ROOT CAUSE: In the ROCm-only unrolled outer loop in embedding_forward_split_kernel_template.cu, the per-thread cache-index staging array was hardcoded as int32_t at lines 330 and 422. For SSD path, locs_or_addrs_type is int64_t (the array entries are raw 64-bit row pointers from ssd_row_addrs). SHFL_SYNC returns the full int64_t value, but writing it to int32_t silently truncates the upper 32 bits. On reload, sign-extension of the low 32 bits reconstructs a corrupted pointer:

• Low-32 with bit 31 = 0 -> small positive value (matches low fault addresses)
• Low-32 with bit 31 = 1 -> sign-extended 0xFFFF... (matches high fault addresses)
  The dereference at load_and_accumulate / load_weights macros (lines 78, 173) reads through this corrupted pointer -> HSA fault.

The non-unrolled tail loop at line 482 already uses {{ locs_or_addrs_type }} correctly. CUDA path uses the non-unrolled path so never hits this bug.

WHAT:
Replace the hardcoded int32_t with the jinja {{ locs_or_addrs_type }} template variable at both sites (line 330 staging array and line 422 reload), matching the type used elsewhere in the same template. Add an explanatory comment at the staging array site.

Differential Revision: D104181076

[meta-codesync]

@EddyLXJ has exported this pull request. If you are a Meta employee, you can view the originating Diff in D104181076.