[feat] hma connector supports GPU buffer MR for GPUDirct RDMA#981
[feat] hma connector supports GPU buffer MR for GPUDirct RDMA#981relat-ivity wants to merge 7 commits into
Conversation
relat-ivity
requested review from
Infinite666,
harrisonyhq,
mag1c-h,
qyh111 and
ygwpz
as code owners
relat-ivity
requested review from
FangRun2,
Tarrei,
Wwwzff and
flesher0813
as code owners
ygwpz
left a comment
ygwpz
left a comment
There was a problem hiding this comment.
Follow-up Review - PR #981
Previous Concerns: All Addressed ✅
All 7 concerns from the previous review have been successfully addressed:
- ✅ Data integrity assertion (Line 167-169): Added
assert len(buffer_addrs) == len(buffer_sizes) - ✅ Code clarity (Line 29-30): Added comment explaining GPU buffer purpose
- ✅ Debug logging (Line 119-125): Added logging for GPU buffer registration
- ✅ Redundant int() conversion (Line 171): Simplified key creation
- ✅ Type safety (Line 112-116): Added validation for non-empty lists
- ✅ Overflow protection (Line 46): Changed to safe sum calculation
- ✅ Warning for missing layouts (Line 159-163): Added logging instead of silent skip
New Observations (Minor - L3-L5)
Three minor suggestions for code quality improvements. See inline comments below.
Summary
The new commits thoroughly address all previous concerns with appropriate safeguards, logging, and validation. The implementation is now more robust and maintainable. Good work!
The three minor suggestions below are optional improvements that could be addressed in a follow-up PR if desired.
| tensor_size = math.prod([t.shape[i] for i in size_dims]) * t.element_size() | ||
| # GPU buffer sizes for GPUDirect RDMA registration in store. | ||
| # Total buffer size = number of blocks (shape[0]) × bytes per block stride. | ||
| buffer_sizes.append(int(t.shape[0]) * block_stride) |
There was a problem hiding this comment.
🔴 Critical: Potential integer overflow. When t.shape[0] is very large (e.g., millions of blocks) and block_stride is also large, this multiplication could overflow in Python's int conversion. Consider using int(t.shape[0]) * int(block_stride) or add bounds checking to ensure the product doesn't exceed expected limits for GPUDirect RDMA registration.
| ), "KV cache buffer addresses and sizes must have the same length." | ||
| for addr, size in zip(buffer_addrs, buffer_sizes): | ||
| key = (addr, size) | ||
| if key in gpu_kv_buffer_set: |
There was a problem hiding this comment.
💡 Suggestion: For better readability, use addr and size variables directly instead of key[0] and key[1] in the append statements below:
gpu_kv_buffer_set.add((addr, size))
gpu_kv_buffer_addrs.append(addr)
gpu_kv_buffer_sizes.append(size)2 participants
Purpose
Enable the HMA connector to provide GPU KV buffer address and size metadata to UCM stores, following the existing ucm_connector registration logic in PR #958.
Modifications
hma_connector.py: Collects GPU buffer addresses and sizes of the vLLM KV cache, and passes them to UCM store for GDR pre-registration.Test