deepseek-ai / deepseek-ai/FlashMLA
[Portability][MSVC] Replace __int128_t in shared-memory load/store paths 3
- Dominant language
- C++
- Stars
- 12.9k
- Forks
- 1.2k
- Avg merge
- 4h 20m
- Merged PRs (30d)
- 2
Description
### 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(sK_nope_base + smem_offset) =
+ *reinterpret_cast(&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?
Contributor guide
No contributing guide indexed for this repository
Research direction
Start with csrc/kerutils/include/kerutils/device/sm80/helpers.cuh and csrc/sm90/decode/sparse_fp8/splitkv_mla.cuh, then inspect whether CuTe provides a suitable portable uint128_t transport type. Compare the existing shared-memory helpers and the sparse-FP8 and RoPE stores against the tested uint4 approach, and verify that the selected representation compiles with NVCC and MSVC while preserving the 16-byte transfers.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp
- Domain
- build-system, performance
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 55/100