[Portability][MSVC] Replace __int128_t in shared-memory load/store paths 3

Background

I encountered this issue while building the vllm-project/FlashMLA fork through SystemPanic/vllm-windows for a deployment serving DeepSeek-V4-Flash-0731 on 4 x H200 NVL GPUs.

Problem

Two FlashMLA CUDA source files use __int128_t for 128-bit shared-memory transport:

csrc/kerutils/include/kerutils/device/sm80/helpers.cuh
csrc/sm90/decode/sparse_fp8/splitkv_mla.cuh

The helper file uses __int128_t with ld.shared.b128 and st.shared.b128, including the float4 load/store helpers.

The sparse-FP8 decode kernel uses it for 16-byte shared-memory stores:

*(__int128_t*)(sK_nope_base + smem_offset) =
    *(__int128_t*)&cur_bf16x8;

__int128_t is available with common GCC and Clang host toolchains, but MSVC does not provide this type. Consequently, these CUDA sources cannot be compiled unchanged when NVCC uses MSVC as its host compiler.

I understand that native Windows may not be part of the supported FlashMLA build matrix. However, these uses only transport 128 bits without performing integer arithmetic, so replacing the compiler-specific type with a portable 16-byte transport type appears to be a cross-platform source-portability improvement.

Tested workaround

The working downstream build replaces __int128_t with CUDA's uint4.

For the shared-memory helpers, the tested implementation uses four 32-bit PTX operands:

__device__ __forceinline__ void st_shared(void* ptr, uint4 val) {
    uint32_t addr = cute::cast_smem_ptr_to_uint(ptr);
    asm volatile(
        "st.shared.v4.u32 [%0], {%1, %2, %3, %4};"
        :
        : "r"(addr), "r"(val.x), "r"(val.y), "r"(val.z), "r"(val.w)
    );
}

__device__ __forceinline__ uint4 ld_shared_u4(void* ptr) {
    uint32_t addr = cute::cast_smem_ptr_to_uint(ptr);
    uint4 val;
    asm volatile(
        "ld.shared.v4.u32 {%0, %1, %2, %3}, [%4];"
        : "=r"(val.x), "=r"(val.y), "=r"(val.z), "=r"(val.w)
        : "r"(addr)
    );
    return val;
}

The float4 helpers reinterpret the same 16 bytes through uint4.

For the sparse-FP8 stores, the tested change is conceptually:

- *reinterpret_cast<__int128_t*>(sK_nope_base + smem_offset) =
-     *reinterpret_cast<__int128_t*>(&cur_bf16x8);
+ *reinterpret_cast<uint4*>(sK_nope_base + smem_offset) =
+     *reinterpret_cast<uint4*>(&cur_bf16x8);

The equivalent RoPE store receives the same change.

This allowed the FlashMLA kernels to compile successfully with NVCC and MSVC.

Semantics

These values are used only for bit transport:

• No signed integer arithmetic is performed.
• Both representations are 16 bytes.
• uint4 has the required 16-byte alignment.
• Byte ordering and total shared-memory transaction width are preserved.
• st.shared.v4.u32 and ld.shared.v4.u32 transfer the same total 128 bits as the existing .b128 operations.

The register representation differs, however: the tested workaround uses four 32-bit operands instead of the original 128-bit PTX operand.

CuTe also appears to provide a portable uint128_t transport type. If appropriate for these call sites, using the existing CuTe type may preserve the current .b128 PTX while removing the dependency on the host compiler's __int128_t support.

Environment

• Windows Server 2022
• Visual Studio Build Tools with MSVC
• CUDA/NVCC 13.2
• Python 3.12
• PyTorch 2.11.0+cu130
• SystemPanic/vllm-windows 0.25-based deployment
• DeepSeek-V4-Flash-0731
• 4 x H200 NVL GPUs, SM90a

Possible resolutions

1. Use an existing portable CuTe 128-bit transport type if it supports these PTX operands and toolchains.
2. Use uint4 with the tested .v4.u32 load/store implementation.
3. Keep the existing path for GCC/Clang and select a uint4 implementation under MSVC if preserving the current Linux PTX is preferred.

Would one of these portable representations be suitable for the canonical FlashMLA implementation?