Repository navigation
[DeepEP] Intra-node DeepEP - #40
Conversation
WalkthroughAdds a TileScale-based intra-node DeepEP implementation (EPBuffer, intranode layout/dispatch/combine kernels, utilities, tests, and docs), expands TileLang warp/LD/ST/sync intrinsics and CUDA codegen/templates, introduces host↔device pinned-mapped tensor support and bindings, and removes an older DeepEP benchmark doc/script. Changes
Sequence Diagram(s)sequenceDiagram
participant User
participant EPBuffer
participant Allocator
participant DispatchKernel
participant CombineKernel
participant NVLinkPeers
User->>EPBuffer: init(group, num_nvl_bytes, ...)
EPBuffer->>Allocator: pre-allocate symmetric buffers & counters
User->>EPBuffer: get_dispatch_layout(topk_idx)
EPBuffer->>DispatchKernel: launch layout kernel
DispatchKernel-->>EPBuffer: num_tokens_per_rank, num_tokens_per_expert, is_token_in_rank
User->>EPBuffer: dispatch(x, topk_idx, topk_weights)
EPBuffer->>DispatchKernel: notify_dispatch / dispatch_kernel (or cached path)
DispatchKernel->>NVLinkPeers: put_warp into channel buffers
DispatchKernel-->>EPBuffer: recv_x, recv_topk_idx, recv_topk_weights, handle
User->>User: run local experts on recv_x
User->>EPBuffer: combine(x, handle, topk_weights)
EPBuffer->>CombineKernel: cached_notify_combine / combine_kernel
CombineKernel->>NVLinkPeers: sender pushes contributions
CombineKernel-->>EPBuffer: recv_combined_x, recv_topk_weights
EPBuffer-->>User: combined results
Estimated code review effort🎯 5 (Critical) | ⏱️ ~120 minutes Possibly related PRs
Suggested reviewers
Pre-merge checks and finishing touches❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
|
👋 Hi! Thank you for contributing to the TileLang project. Please remember to run We appreciate you taking this step! Our team will review your contribution, and we look forward to your awesome work! 🚀 |
c359569 to
1412ae7
Compare
c425966 to
4aa257e
Compare
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 13
Note
Due to the large number of review comments, Critical severity comments were prioritized as inline comments.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (3)
src/op/remote_copy.cc (1)
385-404: Input count issues for PutOp and GetOp already flagged.The registrations for PutOp and GetOp at lines 385-388 and 390-393 have incorrect
set_num_inputs(6)which should be 7, as previously noted in the constructor reviews. StOp and LdOp registrations are correct.src/op/sync.h (2)
154-159: SEqualReduce should includeneed_fencefor correctness.The
SEqualReducemethod compareslocal_bar_addr,offset,local_bar, andlocal_indices, but omitsneed_fence. Two barriers with different fence requirements would incorrectly compare as equal.🔎 Suggested fix
bool SEqualReduce(const BarrierBlocksOpNode *other, SEqualReducer equal) const { return equal(local_bar_addr, other->local_bar_addr) && equal(offset, other->offset) && equal(local_bar, other->local_bar) && - equal(local_indices, other->local_indices); + equal(local_indices, other->local_indices) && + equal(need_fence, other->need_fence); }
161-166: SHashReduce should includeneed_fencefor consistency.Similar to SEqualReduce, the hash should include all fields that affect the operator's behavior.
🔎 Suggested fix
void SHashReduce(SHashReducer hash_reduce) const { hash_reduce(local_bar_addr); hash_reduce(offset); hash_reduce(local_bar); hash_reduce(local_indices); + hash_reduce(need_fence); }
🟠 Major comments (9)
examples/distributed/primitives/example_remote_st.py-14-29 (1)
14-29: Remove unusednum_rankparameter or rename the local variable.The function parameter
num_rankon line 14 is shadowed by the local variablenum_rankon line 23, making the parameter unused. Based on the logic, it appearsnum_ranksis obtained dynamically viaT.get_num_ranks(), so the parameter should be removed.🔎 Proposed fix
-def kernel_(M, num_rank, block_M, threads): +def kernel_(M, block_M, threads): @T.prim_func def main( dst: T.Tensor((M), "float32"), src: T.Tensor((M), "float32"), ): with T.Kernel(T.ceildiv(M, block_M), threads=threads) as (bx): rank = T.alloc_local([1], "uint64") - num_rank = T.alloc_local([1], "uint64") + num_ranks = T.alloc_local([1], "uint64") rank[0] = T.get_rank() - num_rank[0] = T.get_num_ranks() + num_ranks[0] = T.get_num_ranks() tx = T.get_thread_binding() T.st(dst[bx * block_M + tx], src[bx * block_M + tx], dst_pe=1 - rank[0]) return mainThen update the call site:
- kernel = tilelang.compile(kernel_(M, num_ranks, BLOCK_M, threads)) + kernel = tilelang.compile(kernel_(M, BLOCK_M, threads))src/tl_templates/cuda/reduce.h-285-294 (1)
285-294: Use__activemask()to match the safer pattern already in this file.The hardcoded mask
0xffffffffat line 287 assumes all 32 threads in the warp are active. On Volta, if threads in the mask are not executing the shuffle instruction, threads still synchronize with them, which can lead to deadlock. The codebase already uses a safer approach at line 74 with__activemask(), which returns a 32-bit unsigned int mask of all currently active threads in the calling warp.For consistency and robustness, replace the hardcoded mask with
__activemask():constexpr uint32_t mask = __activemask();This ensures correct behavior if
warp_reduce()is called in code with potential thread divergence.src/op/remote_copy.h-128-132 (1)
128-132: Fix typos and add missing field to reflection/equality checks (same as PutOpNode).This has identical issues to PutOpNode (Lines 36-40):
- Field name typo:
enable_aggresive_vectorize→enable_aggressive_vectorize- Comment typo: "effctive" → "effective"
- Missing from
RegisterReflection()(lines 143-156)- Missing from
SEqualReduce()(lines 158-169) andSHashReduce()(lines 171-184)🔎 Proposed fixes (parallel to PutOpNode)
Apply the same fixes as suggested for PutOpNode, updating GetOpNode's reflection, SEqualReduce, and SHashReduce methods.
src/op/remote_copy.h-36-40 (1)
36-40: Fix typos and add missing field to reflection/equality checks.Multiple issues with the new field:
- Field name typo:
enable_aggresive_vectorizeshould beenable_aggressive_vectorize- Comment typo: "effctive" should be "effective"
- The field is not exposed in
RegisterReflection()(lines 51-63)- The field is not included in
SEqualReduce()(lines 66-77) orSHashReduce()(lines 79-92)Missing the field from equality/hashing means two PutOp instances that differ only in
enable_aggressive_vectorizewill be considered equal, which is incorrect.🔎 Proposed fixes
Fix the field name and comment:
- bool enable_aggresive_vectorize; ///< Whether to enable aggressive - ///< vectorization, only effctive for + bool enable_aggressive_vectorize; ///< Whether to enable aggressive + ///< vectorization, only effective for ///< warp-scopeAdd to reflection:
static void RegisterReflection() { namespace refl = tvm::ffi::reflection; refl::ObjectDef<PutOpNode>() .def_ro("src_addr", &PutOpNode::src_addr) .def_ro("dst_addr", &PutOpNode::dst_addr) .def_ro("copy_size", &PutOpNode::copy_size) .def_ro("dst_pe", &PutOpNode::dst_pe) .def_ro("unroll_factor", &PutOpNode::unroll_factor) .def_ro("src_buffer", &PutOpNode::src_buffer) .def_ro("dst_buffer", &PutOpNode::dst_buffer) .def_ro("src_indices", &PutOpNode::src_indices) .def_ro("dst_indices", &PutOpNode::dst_indices) - .def_ro("scope", &PutOpNode::scope); + .def_ro("scope", &PutOpNode::scope) + .def_ro("enable_aggressive_vectorize", &PutOpNode::enable_aggressive_vectorize); }Add to SEqualReduce:
bool SEqualReduce(const PutOpNode *other, SEqualReducer equal) const { return equal(src_addr, other->src_addr) && equal(dst_addr, other->dst_addr) && equal(src_offset, other->src_offset) && equal(dst_offset, other->dst_offset) && equal(copy_size, other->copy_size) && equal(dst_pe, other->dst_pe) && equal(unroll_factor, other->unroll_factor) && equal(src_buffer, other->src_buffer) && equal(dst_buffer, other->dst_buffer) && equal(src_indices, other->src_indices) && - equal(dst_indices, other->dst_indices) && scope == other->scope; + equal(dst_indices, other->dst_indices) && scope == other->scope && + enable_aggressive_vectorize == other->enable_aggressive_vectorize; }Add to SHashReduce:
void SHashReduce(SHashReducer hash_reduce) const { hash_reduce(src_addr); hash_reduce(dst_addr); hash_reduce(src_offset); hash_reduce(dst_offset); hash_reduce(copy_size); hash_reduce(dst_pe); hash_reduce(unroll_factor); hash_reduce(src_buffer); hash_reduce(dst_buffer); hash_reduce(src_indices); hash_reduce(dst_indices); hash_reduce(scope); + hash_reduce(enable_aggressive_vectorize); }Committable suggestion skipped: line range outside the PR's diff.
src/tl_templates/cuda/sync.h-187-243 (1)
187-243: Type mismatch: templated wait functions only supportintdespite generic signature.The
wait_eq,wait_ne,wait_ge,wait_le,wait_gt, andwait_lttemplates accept any typeT, but the underlyingld_acquireandld_volatile_globalfunctions (lines 59-69) only operate onint*and returnint. IfTis notint, thereinterpret_cast<T*>(ptr)followed by passing it to these load functions will cause a type mismatch.For example,
wait_ne<uint64_t, uint64_t>would pass auint64_t*told_volatile_global(const int*), which is incorrect.Consider either:
- Constraining
Ttointviastatic_assert- Implementing type-specific overloads of
ld_volatile_globalandld_acquirefor supported types🔎 Proposed fix to constrain T to int
template <typename P, typename T> TL_DEVICE void wait_ne(P ptr, T val) { static_assert(std::is_same_v<P, uint64_t> || std::is_pointer_v<P>, "P must be a pointer or uint64_t"); + static_assert(std::is_same_v<T, int>, "T must be int for wait_ne"); T *flag_ptr = reinterpret_cast<T *>(ptr); // Spin-loop #pragma unroll 1 while (ld_volatile_global(flag_ptr) == val) ; }Apply similar changes to
wait_ge,wait_le,wait_gt,wait_lt, andwait_eq.Committable suggestion skipped: line range outside the PR's diff.
tilelang/utils/ts_ext/ts_ext_bindings.cpp-25-27 (1)
25-27: Missing dtype conversion wrapper for Python compatibility.The
_create_tensorbinding (lines 17-23) uses a lambda to convert Python dtype viatorch::python::detail::py_object_to_dtype, butcreate_host_device_tensoris bound directly. This may cause issues when calling from Python withtorch.float32instead of a rawc10::ScalarType.🔎 Proposed fix to add dtype conversion
- m.def("create_host_device_tensor", &create_host_device_tensor, - "Create host/device shared pinned-mapped tensor (shape + dtype)"); + m.def( + "create_host_device_tensor", + [](const std::vector<int64_t> &shape, const py::object &dtype) { + return create_host_device_tensor( + shape, torch::python::detail::py_object_to_dtype(dtype)); + }, + py::arg("shape"), py::arg("dtype"), + "Create host/device shared pinned-mapped tensor (shape + dtype)");tilelang/utils/ts_ext/tensor.cpp-88-113 (1)
88-113: Memory leak: no custom deleter to free mapped host memory.The tensors created via
torch::from_blobdon't have a custom deleter. When these tensors are destroyed, the memory allocated bycudaHostAllocwon't be freed, causing a memory leak. You need to provide a deleter that callscudaFreeHost(host_ptr).🔎 Proposed fix
std::pair<torch::Tensor, torch::Tensor> create_host_device_tensor(const std::vector<int64_t> &shape, c10::ScalarType dtype) { size_t elem_size = at::elementSize(dtype); int64_t numel = 1; for (int64_t s : shape) numel *= s; size_t bytes = numel * elem_size; void *host_ptr = nullptr; CUDA_CHECK(cudaHostAlloc(&host_ptr, bytes, cudaHostAllocMapped)); void *device_ptr = nullptr; CUDA_CHECK(cudaHostGetDevicePointer(&device_ptr, host_ptr, 0)); + // Custom deleter to free the pinned memory when the tensor is destroyed + auto deleter = [](void *ptr) { + cudaError_t err = cudaFreeHost(ptr); + if (err != cudaSuccess) { + std::fprintf(stderr, "cudaFreeHost failed: %s\n", cudaGetErrorString(err)); + } + }; + auto host_tensor = torch::from_blob( - host_ptr, shape, torch::TensorOptions().dtype(dtype).device(torch::kCPU)); + host_ptr, shape, deleter, torch::TensorOptions().dtype(dtype).device(torch::kCPU)); auto device_tensor = torch::from_blob( - device_ptr, shape, + device_ptr, shape, [](void*){}, // no-op deleter; host_tensor owns the memory torch::TensorOptions().dtype(dtype).device(torch::kCUDA)); return std::make_pair(host_tensor, device_tensor); }src/op/sync.cc-159-161 (1)
159-161: Potential out-of-bounds access inrelation_strarray.The
relation_strarray has 6 elements (indices 0-5), but there's no bounds check onrelationbefore accessingrelation_str[relation]. Ifrelationis out of range, this will cause undefined behavior.🔎 Suggested fix
// Map relation as int to literal_strings const char *relation_str[] = {"eq", "ne", "ge", "le", "gt", "lt"}; + ICHECK(relation >= 0 && relation < 6) + << "relation must be in range [0, 5], got: " << relation; ss << "tl::wait_" << relation_str[relation];examples/distributed/primitives/test_st_options.py-46-48 (1)
46-48: Remove commented-out code.Line 46 contains a commented-out test that is immediately followed by the same test uncommented. This appears to be leftover debugging code.
Proposed cleanup
- # test_st_options("gpu", "weak", False) test_st_options("gpu", "weak", False) test_st_options("gpu", "weak", True)
🟡 Minor comments (11)
examples/distributed/primitives/example_remote_st.py-50-51 (1)
50-51: Shape argument should be a tuple, not an integer.
(M)is not a tuple—it's justMwrapped in parentheses. For a single-element tuple, use(M,). Based ontilelang/utils/tensor.py, the shape parameter expectstuple[int, ...].🔎 Proposed fix
- src = tilelang.tensor((M), torch.float32, allocator=allocator).normal_() - dst = tilelang.tensor((M), torch.float32, allocator=allocator) + src = tilelang.tensor((M,), torch.float32, allocator=allocator).normal_() + dst = tilelang.tensor((M,), torch.float32, allocator=allocator)testing/python/language/test_tilelang_language_vote.py-44-44 (1)
44-44: Incorrect assertion due to operator precedence.The assertion
assert '__any_sync' and '__all_sync' in kernel.get_kernel_source()does not check for both strings. Due to Python's operator precedence, this evaluates asassert '__any_sync' and ('__all_sync' in kernel.get_kernel_source()), where'__any_sync'is always truthy, so only'__all_sync'is actually verified.🔎 Proposed fix
- assert '__any_sync' and '__all_sync' in kernel.get_kernel_source() + source = kernel.get_kernel_source() + assert '__any_sync' in source and '__all_sync' in sourcetilelang/language/reduce.py-278-291 (1)
278-291: Minor docstring parameter name mismatch.The docstring references
xin the Args section (line 286), but the actual parameter is namedvalue. This should be consistent with the other warp_reduce functions which correctly documentvalue.Also, there's a double-space typo in "a register" (line 282) that appears in all five docstrings.
🔎 Proposed fix
def warp_reduce_sum(value: tir.PrimExpr): """Perform warp reduction sum on a register value. This function reduces a value across all threads in a warp using shuffle operations. - Each thread provides a register `value`, and after the reduction, all threads + Each thread provides a register `value`, and after the reduction, all threads will have the sum of all values across the warp. Args: - x (tir.PrimExpr): The input register value to reduce + value (tir.PrimExpr): The input register value to reduce Returns: tir.PrimExpr: The reduced sum value (same on all threads in the warp) """ return tir.call_intrin(value.dtype, tir.op.Op.get("tl.warp_reduce_sum"), value)tilelang/language/builtin.py-765-768 (1)
765-768: Error message incorrectly mentions "release" as a valid semantic forld.The assertion correctly excludes "release" (since release semantics are for stores, not loads), but the error message erroneously lists it as valid.
🔎 Proposed fix
assert sem in [ "weak", "volatile", "acquire", "relaxed" - ], "Semantic must be one of 'weak', 'volatile', 'acquire', 'release', or 'relaxed'." + ], "Semantic must be one of 'weak', 'volatile', 'acquire', or 'relaxed'."tilelang/distributed/testing/test_create_mapped_tensor.py-15-17 (1)
15-17: Add synchronization before checking host tensor after GPU modification.The
random_()operation is asynchronous. Without explicit synchronization, the host tensor comparison may read stale data before the GPU operation completes, causing intermittent test failures.🔎 Proposed fix
# test modification device_tensor.random_() + torch.cuda.synchronize() # Ensure GPU operation completes before reading host assert torch.equal(host_tensor, device_tensor.cpu()), f"{host_tensor=}, {device_tensor=}"tilelang/utils/ts_ext/tensor.cpp-93-95 (1)
93-95: Missing overflow check fornumelcalculation.Unlike
tensor_from_ptrwhich usessafe_mul_int64to detect overflow, this function multiplies shape dimensions without checking for integer overflow, which could lead to incorrect memory allocation sizes.🔎 Proposed fix
size_t elem_size = at::elementSize(dtype); int64_t numel = 1; - for (int64_t s : shape) - numel *= s; + for (int64_t s : shape) { + numel = safe_mul_int64(numel, s); + }examples/distributed/deepseek_deepep/intranode/get_dispatch_layout.py-17-60 (1)
17-60: Return type annotation and docstring don't match actual return.The type annotation specifies 4 return values including
torch.Tensor | Nonefornum_tokens_per_rdma_rank, and the docstring (lines 29-31) documents it, but line 60 only returns 3 values. Either add the fourth return value (asNonefor intranode) or fix the annotations.🔎 Proposed fix (Option A: Fix annotation to match implementation)
def get_dispatch_layout( topk_idx: torch.Tensor, num_experts: int, - num_ranks: int) -> Tuple[torch.Tensor, torch.Tensor | None, torch.Tensor, torch.Tensor]: + num_ranks: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Calculate the layout required for later communication. Arguments: topk_idx: `[num_tokens, num_topk]`, dtype must be `torch.int64`, the expert indices selected by each token, `-1` means no selections. num_experts: the number of experts. num_ranks: the number of ranks. Returns: num_tokens_per_rank: `[num_ranks]` with `torch.int`, the number of tokens to be sent to each rank. - num_tokens_per_rdma_rank: `[num_rdma_ranks]` with `torch.int`, the number of tokens to be sent to each RDMA - rank (with the same GPU index), return `None` for intranode settings. num_tokens_per_expert: `[num_experts]` with `torch.int`, the number of tokens to be sent to each expert. is_token_in_rank: `[num_tokens, num_ranks]` with `torch.bool`, whether a token be sent to a rank. """examples/distributed/deepseek_deepep/utils.py-220-248 (1)
220-248: Inline C++ extension with busy-wait loop.The
wait_for_counters_readyfunction uses a tightwhile (true)spin-loop without any backoff or timeout. This could cause CPU spinning issues in production. Consider:
- Adding a timeout parameter
- Using a yield/sleep in the loop
- Documenting that this is for testing only
🔎 Suggested improvement with timeout
std::tuple<int, std::vector<int>> wait_for_counters_ready( torch::Tensor& moe_recv_counter, torch::Tensor& moe_recv_expert_counter, int64_t timeout_ms = -1) { // -1 means infinite volatile int *counter_ptr = moe_recv_counter.data_ptr<int>(); volatile int *expert_ptr = moe_recv_expert_counter.data_ptr<int>(); const int num_local_experts = moe_recv_expert_counter.size(0); auto start = std::chrono::steady_clock::now(); while (true) { bool ready = counter_ptr[0] >= 0; for (int i = 0; i < num_local_experts && ready; ++i) ready &= expert_ptr[i] >= 0; if (ready) break; if (timeout_ms >= 0) { auto elapsed = std::chrono::duration_cast<std::chrono::milliseconds>( std::chrono::steady_clock::now() - start).count(); if (elapsed > timeout_ms) { throw std::runtime_error("Timeout waiting for counters"); } } std::this_thread::yield(); // Allow other threads to run } // ... rest of function }src/op/sync.h-147-152 (1)
147-152: Reflection should includeneed_fencefield.The
RegisterReflection()method exposeslocal_bar_addr,offset,local_bar, andlocal_indices, but the newneed_fencefield is not included. This could cause issues with serialization or debugging.🔎 Suggested fix
static void RegisterReflection() { namespace refl = tvm::ffi::reflection; refl::ObjectDef<BarrierBlocksOpNode>() .def_ro("local_bar_addr", &BarrierBlocksOpNode::local_bar_addr) .def_ro("offset", &BarrierBlocksOpNode::offset) .def_ro("local_bar", &BarrierBlocksOpNode::local_bar) - .def_ro("local_indices", &BarrierBlocksOpNode::local_indices); + .def_ro("local_indices", &BarrierBlocksOpNode::local_indices) + .def_ro("need_fence", &BarrierBlocksOpNode::need_fence); }examples/distributed/deepseek_deepep/intranode/combine.py-422-422 (1)
422-422: Hardcoded dtype 'bfloat16' ignores input tensor dtype.The
combine_kernelis called withdtype='bfloat16'hardcoded, butx.dtypecould potentially be different. Consider passingstr(x.dtype).replace('torch.', '')or similar to match the input.🔎 Suggested fix
kernel = combine_kernel( num_ranks, config.num_max_nvl_chunked_send_tokens, config.num_max_nvl_chunked_recv_tokens, hidden, num_topk, config.num_sms, - dtype='bfloat16') + dtype='bfloat16' if x.dtype == torch.bfloat16 else str(x.dtype).split('.')[-1])examples/distributed/deepseek_deepep/intranode/dispatch.py-858-861 (1)
858-861: UnconditionalNonedereference at line 860 whenhandleis provided withoutnum_tokens_per_rank.Line 860 accesses
num_tokens_per_rank.shape[0]unconditionally, but the function signature declaresnum_tokens_per_rank: Optional[torch.Tensor] = None. The validation at lines 852-854 only enforces requirements whenhandle is None. In cached mode (whenhandleis provided), a caller can passnum_tokens_per_rank=Nonewithout triggering the assertion, leading to anAttributeErrorat line 860.The issue is reachable:
buffer.py'sdispatchmethod (line 248-269) does not validatenum_tokens_per_rankin cached mode before forwarding it tointranode_dispatch.Add validation to ensure
num_tokens_per_rankis notNonewhen used, or adjust the parameter requirement based on the execution mode.
🧹 Nitpick comments (44)
src/tl_templates/cuda/debug.h (2)
261-261: Correct the misleading comment.The comment refers to this as a "Specialization" but
debug_print_msgis a standalone function, not a template specialization.🔎 Proposed fix
-// Specialization for msg-only debug print +// Device-side utility for message-only debug print __device__ void debug_print_msg(const char *msg) {
261-266: Consider adding a forward declaration for consistency.Unlike
debug_print_var(line 11) anddebug_print_buffer_value(lines 125-127),debug_print_msglacks a forward declaration at the top of the file. Adding one would maintain consistency with the established pattern.🔎 Suggested addition near line 11
// Template declaration for device-side debug printing (variable only) template <typename T> __device__ void debug_print_var(const char *msg, T var); + +// Declaration for device-side debug printing (message only) +__device__ void debug_print_msg(const char *msg);tilelang/language/print.py (1)
152-152: Documentation updates are consistent with the changes.The docstring and error message have been appropriately updated to reflect that
Noneis now a supported type for theobjparameter.Note: Ruff suggests extracting the long error message (TRY003), but this is a minor style preference and the current implementation is clear and readable.
Also applies to: 226-227
examples/distributed/primitives/example_remote_st.py (1)
37-37: Prefix unusedrankvariable with underscore.The
rankvariable is unpacked but never used. Use_or_rankto indicate it's intentionally unused.🔎 Proposed fix
- rank, num_ranks, group = init_dist(local_rank, num_local_ranks) + _rank, num_ranks, group = init_dist(local_rank, num_local_ranks)src/tl_templates/cuda/reduce.h (2)
283-284: Consider expanding the section comment.The comment "TileScale extra" is minimal. Consider adding a brief description of what these utilities do, e.g., "Warp-level reduction primitives for register values."
308-314: Consider addingwarp_reduce_bitxorfor completeness.
BitXorOpis defined at lines 37-41 but there's no correspondingwarp_reduce_bitxorwrapper, while bitwise AND and OR are provided. If bitwise XOR reduction is not needed, this is fine; otherwise, consider adding it for API completeness.src/transform/common/loop_vectorization_utils.h (2)
389-389: Unusedis_loadparameter.The
is_loadparameter is declared but never referenced in the function body. Either this is dead code from an earlier design, or differentiated handling for loads vs. stores was intended but not implemented.Suggested fix: Remove unused parameter or add a TODO
If no differentiated behavior is needed:
- PrimExpr MutateTlLdStExpr_(const CallNode *op, bool is_load) { + PrimExpr MutateTlLdStExpr_(const CallNode *op) {Or add a TODO if differentiated handling is planned:
PrimExpr MutateTlLdStExpr_(const CallNode *op, bool is_load) { // Structure: call_extern("tl::ld<...>", address_of(BufferLoad), value, ...) // or: call_extern("tl::st<...>", address_of(BufferLoad), value, ...) + // TODO: Use is_load for differentiated ld/st handling if needed ICHECK(op->args.size() >= 3) << "tl::ld/st expects at least 3 arguments";
458-468: Consider fallback handling forelem_bitsextraction failure.The default
elem_bits = 16assumption could produce incorrect vectorization widths for non-16-bit types (e.g.,int8orfloat32) if the buffer dtype extraction fails. While this path should rarely be hit, a warning or assertion might help catch unexpected cases.Optional: Add defensive logging
// Try to get element dtype from source buffer auto addr_call_check = new_addr.as<CallNode>(); if (addr_call_check && addr_call_check->op.same_as(builtin::address_of())) { auto buffer_load = addr_call_check->args[0].as<BufferLoadNode>(); if (buffer_load) { elem_bits = buffer_load->buffer->dtype.bits(); } + } else { + LOG(WARNING) << "Could not determine element dtype for tl::ld/st vectorization, defaulting to 16 bits"; }src/transform/vectorize_loop.cc (2)
702-705: Unusedis_loadparameter.Same issue as in
loop_vectorization_utils.h— theis_loadparameter is declared but never used in the function body.
702-834: Significant code duplication withloop_vectorization_utils.h.
MutateTlLdStExpr_is implemented in both this file (~130 lines) andloop_vectorization_utils.hwith nearly identical logic. This duplication:
- Increases maintenance burden
- Risks divergence (as evidenced by the
vec_value_slicebug in this file but not the other)Consider extracting the shared vectorization logic to a common utility function or having
TLVectorizerreuse code from theVectorizerinloop_vectorization_utils.h.testing/python/language/test_tilelang_language_warp_reduce.py (1)
46-46: Remove debug print statement.This
🔎 Proposed fix
def test_warp_reduce_max(): a = torch.randn((32,), dtype=torch.float32, device='cuda') kernel = get_kernel("max", 'float32') - print(kernel.get_kernel_source()) ref = torch.full_like(a, a.max()) kernel(a) torch.testing.assert_close(a, ref)examples/distributed/primitives/example_put_warp.py (1)
71-76: Consider relaxing tolerance for bfloat16.The tolerance
atol=1e-6, rtol=1e-6is quite strict for bfloat16, which has only ~3-4 significant decimal digits. While this should work for a pure copy operation (bit-exact), if you encounter flaky tests, consider increasing toatol=1e-2, rtol=1e-2or usingtorch.equal()for exact comparison.tilelang/language/builtin.py (1)
514-514: PEP 484 prohibits implicitOptional.As flagged by static analysis,
barrier_id: int | PrimExpr = Noneshould explicitly useOptionalwhen the default isNone.🔎 Proposed fix
-def sync_threads(barrier_id: int | PrimExpr = None, arrive_count: int = None): +def sync_threads(barrier_id: int | PrimExpr | None = None, arrive_count: int | None = None):examples/distributed/sp_ag_attention_intra_node.py (1)
13-20: Consider renamingbarrier_all_blocks_sys_kernelfor consistency.The internal implementation now uses
T.barrier_blocks, but the function is still namedbarrier_all_blocks_sys_kernel. Consider renaming tobarrier_blocks_kernelor similar to match the API change.🔎 Proposed rename
@tilelang.jit -def barrier_all_blocks_sys_kernel(num_local_rank,): +def barrier_blocks_kernel(num_local_rank,): @T.prim_func def main(barrier: T.Tensor((num_local_rank), "int32"),): with T.Kernel(1, threads=32): T.barrier_blocks(barrier) return mainAlso update the call site at line 612:
def barrier_all_on_stream(barrier: torch.Tensor, stream: torch.cuda.Stream, world_size: int): - barrier_all_blocks_sys_func = barrier_all_blocks_sys_kernel(world_size) + barrier_all_blocks_sys_func = barrier_blocks_kernel(world_size) barrier_all_blocks_sys_func(barrier, stream=stream.cuda_stream)tilelang/jit/kernel.py (1)
407-415: Use explicitOptional[int]type annotation.The static analysis correctly flags that PEP 484 prohibits implicit
Optional. While functionally equivalent, explicit typing improves code clarity and IDE support.🔎 Proposed fix
+ from typing import Optional + def initialize( self, allocator: BaseAllocator, - stream: int = None, + stream: Optional[int] = None, ):examples/distributed/deepseek_deepep/intranode/__init__.py (1)
1-3: Remove unnecessarynoqadirectives.The static analysis indicates F401 is not enabled, making these
# noqa: F401comments redundant. They can be safely removed to reduce noise.🔎 Proposed fix
-from .get_dispatch_layout import get_dispatch_layout # noqa: F401 -from .dispatch import intranode_dispatch # noqa: F401 -from .combine import intranode_combine # noqa: F401 +from .get_dispatch_layout import get_dispatch_layout +from .dispatch import intranode_dispatch +from .combine import intranode_combinetilelang/distributed/testing/test_create_mapped_tensor.py (1)
4-19: Consider converting to pytest format for CI integration.The test runs as a standalone script but would integrate better with CI if converted to a pytest function. This is optional if manual execution is intentional.
🔎 Example pytest structure
import pytest import torch from tilelang.distributed.utils import create_mapped_tensor def test_create_mapped_tensor(): shape = (1024, 1024) dtype = torch.float32 host_tensor, device_tensor = create_mapped_tensor(shape, dtype) assert device_tensor.device.type == "cuda" assert device_tensor.shape == shape assert device_tensor.dtype == dtype assert torch.equal(host_tensor, device_tensor.cpu()) device_tensor.random_() torch.cuda.synchronize() assert torch.equal(host_tensor, device_tensor.cpu())tilelang/language/distributed/common.py (3)
27-27: Typo: "aggresive" should be "aggressive".The parameter name
enable_aggresive_vectorizehas a spelling error. This affects the public API surface.🔎 Proposed fix
- enable_aggresive_vectorize: bool = False): + enable_aggressive_vectorize: bool = False):Apply the same fix to
get_warp(line 55) and update docstrings (lines 42-44, 70-72) accordingly.
135-162: Type annotationPrimExpr | None = -1is misleading.The type hint suggests
peercan bePrimExprorNone, but the default value is-1(an integer). Consider usingPrimExpr | int = -1orPrimExpr | IntImm = -1to match the actual usage.🔎 Proposed fix
-def wait_ne(ptr: PrimExpr, expected: PrimExpr, peer: PrimExpr | None = -1): +def wait_ne(ptr: PrimExpr, expected: PrimExpr, peer: PrimExpr | int = -1):Apply the same fix to
wait_ge,wait_le,wait_gt, andwait_lt.
116-132:BinaryRelation.EQis defined but unused.
wait_equses a separate"tl.wait_eq"intrinsic instead of usingBinaryRelation.EQ.valuewith"tl.wait". The TODO comment on line 127 acknowledges this semantic difference. Consider removingEQ = 0from the enum if it's not intended for use, or unify the implementation.examples/distributed/deepseek_deepep/intranode/get_dispatch_layout.py (3)
73-73: Missing validation for even expert-to-rank distribution.
experts_per_rank = num_experts // num_ranksassumesnum_expertsis evenly divisible bynum_ranks. If not, the kernel logic (e.g.,rank_idx = expert_idx // experts_per_rank) may produce incorrect results. Consider adding an assertion.🔎 Proposed fix
Add validation in the wrapper function:
# Check inputs assert topk_idx.dtype == torch.int64, "topk_idx must be of dtype torch.int64" assert topk_idx.ndim == 2, "topk_idx must be a 2D tensor" assert topk_idx.is_contiguous(), "topk_idx must be a contiguous tensor" assert num_experts > 0, "num_experts must be greater than 0" + assert num_experts % num_ranks == 0, "num_experts must be evenly divisible by num_ranks"
98-99: Potential data truncation: int64 to int32.
topk_idxisint64butexpert_idxis allocated asint32. If expert indices exceedINT32_MAX, this will truncate. While unlikely for typical MoE setups, consider using consistent types.
104-108: Variable namesumshadows Python built-in.While this works in TileLang's DSL context, using
sumas a variable name can cause confusion. Consider renaming tototaloracc.examples/distributed/deepseek_deepep/intranode/combine.py (6)
4-14: Consider consolidating global state modifications.The
sys.path.appendmanipulation andtilelang.disable_cache()with environment variable modification at import time can cause issues in multi-module scenarios. Consider documenting whytilelang.disable_cache()is needed or making it conditional.
51-51: Magic number forlast_headinitialization.The heuristic value
2**25(33.5M) as "a large number" is undocumented. Consider defining a named constant with documentation explaining the expected upper bound for token counts.🔎 Suggested improvement
+# Maximum expected tokens per channel - used as sentinel value +MAX_EXPECTED_TOKENS = 2**25 + ... - last_head = T.alloc_var('int32', init=2**25) # a heuristic large number + last_head = T.alloc_var('int32', init=MAX_EXPECTED_TOKENS)
115-123: Assertion for hidden dimension alignment.The assertion
assert hidden % 8 == 0on line 123 will fail at JIT compile time rather than at runtime, which is good. However, this constraint should be documented in the publicintranode_combinefunction's docstring as well.
273-274: Magic number for INT32_MAX.
2**31 - 1is used as INT32_MAX. Consider using a named constant for clarity.🔎 Suggested improvement
+INT32_MAX = 2**31 - 1 + ... - min_head = 2**31 - 1 # int32 max + min_head = INT32_MAX ... - if min_head != 2**31 - 1 and min_head > last_head: + if min_head != INT32_MAX and min_head > last_head:
307-315: Spin-wait withT.loop_continue()may cause excessive spinning.The
T.Whileloop withT.warp_any(condvar <= expected_head and expected_head >= 0)followed by immediateT.loop_continue()creates a tight spin-loop. The comment "can we simplify this?" suggests this is a known concern. Consider if backoff or yielding could help reduce contention.
382-443: Theintranode_combinefunction lacks input validation documentation.While
assert handle is not Noneexists, the function unpacks multiple tuples fromhandleandsymm_bufferswithout validation. Consider adding type hints or assertions to catch mismatched tuple structures early.🔎 Suggested improvement
def intranode_combine(rank: int, allocator, symm_buffers, x, config, handle, topk_weights, comm_stream=None): assert handle is not None + assert len(handle) == 6, f"Expected handle with 6 elements, got {len(handle)}" + assert len(symm_buffers) == 11, f"Expected symm_buffers with 11 elements, got {len(symm_buffers)}" rank_prefix_matrix, channel_prefix_matrix, recv_channel_prefix_matrix, recv_src_idx, _, send_head = handleexamples/distributed/deepseek_deepep/intranode/test_intranode.py (4)
30-33: Clean up unused noqa directive and improve exception handling.Per static analysis: the
# noqa: F403is unused (F403 is for star imports, not applicable here). Also,TRY003suggests moving the long message into the exception class.🔎 Suggested fix
try: - import deep_ep # noqa: F403 + import deep_ep except ModuleNotFoundError: - raise ModuleNotFoundError("Please install DeepEP to run this test.") from None + raise ModuleNotFoundError( + "DeepEP is required for this test. Please install it first." + ) from None
40-41: Prefix unused variablerank_idxwith underscore.Per static analysis,
rank_idxis unpacked but never used. Prefix with_to indicate intentional discard.🔎 Suggested fix
- x, topk_idx, topk_weights, rank_idx = gen_inputs(num_tokens, hidden, num_topk, num_experts, + x, topk_idx, topk_weights, _rank_idx = gen_inputs(num_tokens, hidden, num_topk, num_experts, num_ranks)
62-63: Prefix unused variableeventwith underscore.Per static analysis,
eventis unpacked but never used.🔎 Suggested fix
- ref_recv_x, ref_recv_topk_idx, ref_recv_topk_weights, ref_num_recv_tokens_per_expert_list, ref_handle, event = \ + ref_recv_x, ref_recv_topk_idx, ref_recv_topk_weights, ref_num_recv_tokens_per_expert_list, ref_handle, _event = \ deepep_buffer.dispatch(x, None, ref_num_tokens_per_rank, None, ref_is_token_in_rank, ref_num_tokens_per_expert, topk_idx, topk_weights, expert_alignment)
190-204: Bandwidth calculation may be misleading.The bandwidth is calculated using only
recv_x.numel() * 2bytes, but the actual data transferred includestopk_idx,topk_weights, andsrc_idx. This gives an underestimate of actual bandwidth utilization.examples/distributed/deepseek_deepep/utils.py (2)
14-14: Globalnum_smsvariable may cause confusion.There's a module-level
num_sms: int = 20that shadows theConfig.num_smsfield. This could lead to confusion when reading the config maps (lines 46-58, 76-88) which reference this global.🔎 Suggested clarification
-num_sms: int = 20 +# Default number of SMs used for high-throughput kernels +# This global is referenced in Config.get_dispatch_config and get_combine_config +DEFAULT_NUM_SMS: int = 20Then update the config maps to use
DEFAULT_NUM_SMSexplicitly for clarity.
151-152:topk_idxfromtorch.topkwill never contain -1.The generated
topk_idxusestorch.topkwhich always returns valid indices. The subsequentrank_idx.masked_fill_(topk_idx == -1, -1)on line 155 is a no-op sincetopk_idxwill never have -1 values. This is fine for testing purposes but the docstring implies -1 is possible.src/op/sync.h (1)
56-59: Consider documenting therelationenum values.The
relationfield is anintbut its valid values (0-5 mapping to eq/ne/ge/le/gt/lt) are only documented in the .cc file. Consider adding an enum or documentation here.🔎 Suggested improvement
+ /// Relation enum values: + /// 0 = eq, 1 = ne, 2 = ge, 3 = le, 4 = gt, 5 = lt int relation; ///< The relation to compare against.Or better, define an enum:
enum class WaitRelation : int { kEq = 0, ///< Equal kNe = 1, ///< Not equal kGe = 2, ///< Greater or equal kLe = 3, ///< Less or equal kGt = 4, ///< Greater than kLt = 5, ///< Less than };examples/distributed/deepseek_deepep/buffer.py (3)
7-10: Relative imports may fail depending on working directory.The imports
from utils import Configandfrom intranode import ...assume the script is run from thedeepseek_deepepdirectory. Unlikedispatch.pywhich usessys.path.append, this file lacks a similar path adjustment.Consider using explicit relative imports or adding a similar path setup:
🔎 Suggested fix
+import os +import sys +sys.path.append(os.path.dirname(__file__)) + import tilelang -from utils import Config +from .utils import Config from tilelang.distributed.utils import create_mapped_tensor -from intranode import get_dispatch_layout, intranode_dispatch, intranode_combine +from .intranode import get_dispatch_layout, intranode_dispatch, intranode_combineOr use a similar
sys.path.appendpattern as indispatch.py.
160-162: Dead code: internode counter allocation is unreachable.The
_moe_recv_rdma_counterallocation on lines 160-162 will never execute because_pre_alloc_symm_buffers_internode()raisesNotImplementedErrorfornum_ranks > 8. Either remove this dead code or defer it until internode support is implemented.🔎 Suggested fix
def _prepare_counters(self): self._moe_recv_counter, self._moe_recv_counter_mapped = create_mapped_tensor([1], torch.int32) self._moe_recv_expert_counter, self._moe_recv_expert_counter_mapped = create_mapped_tensor( [self.num_local_experts], torch.int32) - - if self.num_ranks > 8: # internode - self._moe_recv_rdma_counter, self._moe_recv_rdma_counter_mapped = create_mapped_tensor( - [1], torch.int32) + # TODO: Add internode counter allocation when internode support is implemented
248-270: Unused unpacked variables from handle.The variables
rank_prefix_matrix,channel_prefix_matrix,recv_channel_prefix_matrix,recv_src_idx, andsend_headare unpacked but never used sincehandleis passed directly tointranode_dispatch. Per the static analysis hint, prefix them with underscores.🔎 Suggested fix
if handle is not None: assert topk_idx is None and topk_weights is None - rank_prefix_matrix, channel_prefix_matrix, recv_channel_prefix_matrix, recv_src_idx, is_token_in_rank, send_head = handle + _rank_prefix_matrix, _channel_prefix_matrix, _recv_channel_prefix_matrix, _recv_src_idx, is_token_in_rank, _send_head = handle recv_x = intranode_dispatch(Alternatively, if only
is_token_in_rankis needed, you could use indexed access:is_token_in_rank = handle[4]examples/distributed/deepseek_deepep/deepep.md (2)
1-9: Fix markdown formatting issues.
- Line 3: Use a proper markdown link instead of a bare URL.
- Lines 7-8: Use
[ ](with space) for unchecked items for proper rendering.🔎 Suggested fix
# DeepEP -To install and compare with the original DeepEP implementation, please refer to https://github.com/deepseek-ai/DeepEP. +To install and compare with the original DeepEP implementation, please refer to [DeepEP](https://github.com/deepseek-ai/DeepEP). ## TODO - [x] Intranode Normal Mode -- [] Internode Normal Mode -- [] Low-latency Mode +- [ ] Internode Normal Mode +- [ ] Low-latency Mode
146-149: Add language specifier to fenced code block.The code block should specify a language for proper syntax highlighting.
🔎 Suggested fix
-``` +```bash TILELANG_USE_DISTRIBUTED=1 python intranode/test_intranode.py \ --num_ranks 8 --num_tokens 4096 --hidden 7168 --num_topk 8 --num_experts 32 [--cached]</details> </blockquote></details> <details> <summary>examples/distributed/deepseek_deepep/intranode/dispatch.py (2)</summary><blockquote> `13-13`: **Remove unused `noqa` directive.** The `# noqa: F403` comment is unnecessary since there's no star import triggering F403. <details> <summary>🔎 Suggested fix</summary> ```diff -from utils import Config, ep_ext # noqa: F403 +from utils import Config, ep_ext
569-579: Consider removing unused kernel arguments.The
num_topkandnum_expertsparameters are passed but never used in the cached dispatch kernel (since it doesn't handle topk data). You can either:
- Remove them if API consistency isn't required
- Prefix with underscore to indicate intentional non-use
🔎 Suggested fix (option 2)
def cached_dispatch_kernel( num_ranks, num_tokens, num_max_send_tokens, # config.num_max_nvl_chunked_send_tokens num_recv_buffer_tokens, # config.num_max_nvl_chunked_recv_tokens hidden, - num_topk, - num_experts, + num_topk, # unused in cached path, kept for API consistency + num_experts, # unused in cached path, kept for API consistency num_sms, dtype: str = 'bfloat16', ):examples/distributed/primitives/test_st_options.py (1)
20-25: Consider using consistent assertion style across test files.Line 24 uses
x.equal(...)while the companion filetest_ld_options.pyline 27 usestorch.equal(x, y). Both work, but consistency aids readability.Suggested change for consistency
- assert x.equal(torch.arange(32, device="cuda")) + assert torch.equal(x, torch.arange(32, device="cuda"))
| src = tilelang.tensor(shape=(M), dtype=torch.bfloat16, allocator=allocator).normal_() | ||
| dst = tilelang.tensor(shape=(M), dtype=torch.bfloat16, allocator=allocator) |
There was a problem hiding this comment.
shape=(M) should be shape=(M,) for a 1D tuple.
shape=(M) evaluates to just the integer M, not a tuple. To create a single-element tuple, use shape=(M,) with a trailing comma.
🔎 Proposed fix
- src = tilelang.tensor(shape=(M), dtype=torch.bfloat16, allocator=allocator).normal_()
- dst = tilelang.tensor(shape=(M), dtype=torch.bfloat16, allocator=allocator)
+ src = tilelang.tensor(shape=(M,), dtype=torch.bfloat16, allocator=allocator).normal_()
+ dst = tilelang.tensor(shape=(M,), dtype=torch.bfloat16, allocator=allocator)🤖 Prompt for AI Agents
In examples/distributed/primitives/example_put_warp.py around lines 58 to 59,
the tensor shapes are specified as shape=(M) which evaluates to the integer M
rather than a 1-tuple; update both tensor shape arguments to use a
single-element tuple syntax shape=(M,) so the tensors are created as 1D arrays
with length M.
| if __name__ == "__main__": | ||
| # from DeepEP all ld instructions | ||
|
|
||
| # ld.acquire.sys.global.s32 / u64 | ||
| test_ld_options(scope="sys", sem="acquire", na=False, nc=False) | ||
|
|
||
| # ld.acquire.gpu.global.s32 | ||
| test_ld_options(scope="gpu", sem="acquire", na=False, nc=False) | ||
|
|
||
| # ld.acquire.cta.s32 | ||
| test_ld_options(scope="cta", sem="acquire", na=False, nc=False) | ||
|
|
||
| # ld.relaxed.gpu.global.L1::no_allocate.b8/b16/b32/b64 | ||
| test_ld_options(scope="gpu", sem="relaxed", na=True, nc=False) | ||
|
|
||
| # ld.volatile.global.s32/f32/s64/u64 | ||
| test_ld_options(scope="gpu", sem="volatile", na=False, nc=False) | ||
|
|
||
| # ld.global.nc.L1::no_allocate.L2::256B (or ld.volatile.global when DISABLE_AGGRESSIVE_PTX_INSTRS) | ||
| test_ld_options(scope="gpu", sem="weak", na=True, nc=True) |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Check CUDA availability handling in the test file
cat -n examples/distributed/primitives/test_ld_options.py
# Check for ld usage patterns in examples
rg -n "\.ld\(" examples/ -A2 -B2 2>/dev/null | head -100
# Look for DeepEP-related files and ld patterns
rg -n "T\.ld\|\.ld\(" . -t py -A2 -B2 2>/dev/null | head -150Repository: tile-ai/tilescale
Length of output: 12361
Add CUDA availability check and verify test pattern coverage.
The test lacks a CUDA availability check and will crash on systems without CUDA. Additionally, several tested parameter combinations (scope="sys", scope="gpu"+sem="relaxed", scope="gpu"+sem="weak") are not found in actual DeepEP usage in combine.py and dispatch.py. Only scope="cta"+sem="acquire" and sem="volatile" patterns are used in production code. Recommend adding torch.cuda.is_available() guard and validating whether all tested combinations reflect real usage requirements.
🤖 Prompt for AI Agents
In examples/distributed/primitives/test_ld_options.py around lines 31-50, the
test immediately runs GPU/DeepEP load-instruction checks and will crash on
machines without CUDA; also several tested (scope,sem) combinations are not used
in production (combine.py/dispatch.py). Add a torch.cuda.is_available() guard at
the top of the main block to skip GPU-specific tests on non-CUDA hosts; then
restrict or document the invoked test_ld_options calls to only the patterns used
in production (e.g., scope="cta", sem="acquire"/"volatile"), or explicitly keep
additional cases if you intend to test broader coverage and mark them as
optional/skip when not present in DeepEP usage. Ensure the guard exits or prints
a skip message and only runs GPU tests when CUDA is available.
| template <int N, int UNROLL_FACTOR, bool enable_aggresive_vectorize = false, | ||
| typename dtype_t> | ||
| TL_DEVICE void cp_warp(uint64_t dst_addr_uint64, | ||
| dtype_t const *const src_addr) { | ||
| dtype_t *dst_addr = reinterpret_cast<dtype_t *>(dst_addr_uint64); | ||
|
|
||
| int lane_id; | ||
| asm("mov.s32 %0, %laneid;" : "=r"(lane_id)); | ||
| constexpr int kLoopStride = 32 * (UNROLL_FACTOR); | ||
| typename std::remove_reference<decltype(LD_FUNC((src_addr) + 0))>::type | ||
| unrolled_values[(UNROLL_FACTOR)]; | ||
| auto __src = (src_addr); | ||
| auto __dst = (dst_addr); | ||
| for (int __i = (lane_id); __i < ((N) / kLoopStride) * kLoopStride; | ||
| __i += kLoopStride) { | ||
| _Pragma("unroll") for (int __j = 0; __j < (UNROLL_FACTOR); ++__j) | ||
| unrolled_values[__j] = LD_FUNC(__src + __i + __j * 32); | ||
| _Pragma("unroll") for (int __j = 0; __j < (UNROLL_FACTOR); ++__j) | ||
| ST_FUNC(__dst + __i + __j * 32, unrolled_values[__j]); | ||
| if constexpr (enable_aggresive_vectorize) { | ||
| int4 *__restrict__ dst_addr_int4 = (int4 *)dst_addr; | ||
| const int4 *__restrict__ src_addr_int4 = (const int4 *)src_addr; | ||
| constexpr int N_int4 = sizeof(dtype_t) * N / 16; | ||
| cp_warp_impl<N_int4, UNROLL_FACTOR>(dst_addr_int4, src_addr_int4); | ||
| } else { | ||
| cp_warp_impl<N, UNROLL_FACTOR>(dst_addr, src_addr); | ||
| } | ||
| for (int __i = ((N) / kLoopStride) * kLoopStride + (lane_id); __i < (N); | ||
| __i += 32) | ||
| ST_FUNC(__dst + __i, LD_FUNC(__src + __i)); | ||
| } |
There was a problem hiding this comment.
Same alignment and size validation issues as above.
This overload has the same critical issues as the previous cp_warp variant (lines 221-241):
- Missing alignment verification for 16-byte boundaries
- Missing size validation for
N * sizeof(dtype_t) % 16 == 0 - Typo:
enable_aggresive_vectorizeshould beenable_aggressive_vectorize
🔎 Proposed fix
template <int N, int UNROLL_FACTOR, bool enable_aggresive_vectorize = false,
typename dtype_t>
TL_DEVICE void cp_warp(uint64_t dst_addr_uint64,
dtype_t const *const src_addr) {
dtype_t *dst_addr = reinterpret_cast<dtype_t *>(dst_addr_uint64);
if constexpr (enable_aggresive_vectorize) {
+ static_assert((N * sizeof(dtype_t)) % 16 == 0,
+ "Aggressive vectorization requires N*sizeof(dtype_t) to be a multiple of 16");
int4 *__restrict__ dst_addr_int4 = (int4 *)dst_addr;
const int4 *__restrict__ src_addr_int4 = (const int4 *)src_addr;
constexpr int N_int4 = sizeof(dtype_t) * N / 16;
cp_warp_impl<N_int4, UNROLL_FACTOR>(dst_addr_int4, src_addr_int4);
} else {
cp_warp_impl<N, UNROLL_FACTOR>(dst_addr, src_addr);
}
}📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| template <int N, int UNROLL_FACTOR, bool enable_aggresive_vectorize = false, | |
| typename dtype_t> | |
| TL_DEVICE void cp_warp(uint64_t dst_addr_uint64, | |
| dtype_t const *const src_addr) { | |
| dtype_t *dst_addr = reinterpret_cast<dtype_t *>(dst_addr_uint64); | |
| int lane_id; | |
| asm("mov.s32 %0, %laneid;" : "=r"(lane_id)); | |
| constexpr int kLoopStride = 32 * (UNROLL_FACTOR); | |
| typename std::remove_reference<decltype(LD_FUNC((src_addr) + 0))>::type | |
| unrolled_values[(UNROLL_FACTOR)]; | |
| auto __src = (src_addr); | |
| auto __dst = (dst_addr); | |
| for (int __i = (lane_id); __i < ((N) / kLoopStride) * kLoopStride; | |
| __i += kLoopStride) { | |
| _Pragma("unroll") for (int __j = 0; __j < (UNROLL_FACTOR); ++__j) | |
| unrolled_values[__j] = LD_FUNC(__src + __i + __j * 32); | |
| _Pragma("unroll") for (int __j = 0; __j < (UNROLL_FACTOR); ++__j) | |
| ST_FUNC(__dst + __i + __j * 32, unrolled_values[__j]); | |
| if constexpr (enable_aggresive_vectorize) { | |
| int4 *__restrict__ dst_addr_int4 = (int4 *)dst_addr; | |
| const int4 *__restrict__ src_addr_int4 = (const int4 *)src_addr; | |
| constexpr int N_int4 = sizeof(dtype_t) * N / 16; | |
| cp_warp_impl<N_int4, UNROLL_FACTOR>(dst_addr_int4, src_addr_int4); | |
| } else { | |
| cp_warp_impl<N, UNROLL_FACTOR>(dst_addr, src_addr); | |
| } | |
| for (int __i = ((N) / kLoopStride) * kLoopStride + (lane_id); __i < (N); | |
| __i += 32) | |
| ST_FUNC(__dst + __i, LD_FUNC(__src + __i)); | |
| } | |
| template <int N, int UNROLL_FACTOR, bool enable_aggresive_vectorize = false, | |
| typename dtype_t> | |
| TL_DEVICE void cp_warp(uint64_t dst_addr_uint64, | |
| dtype_t const *const src_addr) { | |
| dtype_t *dst_addr = reinterpret_cast<dtype_t *>(dst_addr_uint64); | |
| if constexpr (enable_aggresive_vectorize) { | |
| static_assert((N * sizeof(dtype_t)) % 16 == 0, | |
| "Aggressive vectorization requires N*sizeof(dtype_t) to be a multiple of 16"); | |
| int4 *__restrict__ dst_addr_int4 = (int4 *)dst_addr; | |
| const int4 *__restrict__ src_addr_int4 = (const int4 *)src_addr; | |
| constexpr int N_int4 = sizeof(dtype_t) * N / 16; | |
| cp_warp_impl<N_int4, UNROLL_FACTOR>(dst_addr_int4, src_addr_int4); | |
| } else { | |
| cp_warp_impl<N, UNROLL_FACTOR>(dst_addr, src_addr); | |
| } | |
| } |
🤖 Prompt for AI Agents
In src/tl_templates/cuda/copy.h around lines 243-256, this cp_warp overload
needs the same fixes as the previous variant: rename the template parameter and
identifier from enable_aggresive_vectorize to enable_aggressive_vectorize, add
alignment validation (ensure dst_addr_uint64 and src_addr are 16-byte aligned
before casting to int4), and add a size validation that (N * sizeof(dtype_t)) %
16 == 0 before computing N_int4 and calling the int4 path; implement these
checks as compile-time static_asserts where possible (use constexpr expressions
for dtype_t/ N) and fall back to a runtime guard/log or early return if
alignment/size cannot be guaranteed at compile time, then only perform the int4
reinterpret_cast and call cp_warp_impl<N_int4,...> when checks pass.
| template <int N, int UNROLL_FACTOR, bool enable_aggresive_vectorize = false, | ||
| typename dtype_t> | ||
| TL_DEVICE void cp_warp(dtype_t *const dst_addr, uint64_t src_addr_uint64) { | ||
| const dtype_t *src_addr = reinterpret_cast<const dtype_t *>(src_addr_uint64); | ||
|
|
||
| int lane_id; | ||
| asm("mov.s32 %0, %laneid;" : "=r"(lane_id)); | ||
| constexpr int kLoopStride = 32 * (UNROLL_FACTOR); | ||
| typename std::remove_reference<decltype(LD_FUNC((src_addr) + 0))>::type | ||
| unrolled_values[(UNROLL_FACTOR)]; | ||
| auto __src = (src_addr); | ||
| auto __dst = (dst_addr); | ||
| for (int __i = (lane_id); __i < ((N) / kLoopStride) * kLoopStride; | ||
| __i += kLoopStride) { | ||
| _Pragma("unroll") for (int __j = 0; __j < (UNROLL_FACTOR); ++__j) | ||
| unrolled_values[__j] = LD_FUNC(__src + __i + __j * 32); | ||
| _Pragma("unroll") for (int __j = 0; __j < (UNROLL_FACTOR); ++__j) | ||
| ST_FUNC(__dst + __i + __j * 32, unrolled_values[__j]); | ||
| if constexpr (enable_aggresive_vectorize) { | ||
| int4 *__restrict__ dst_addr_int4 = (int4 *)dst_addr; | ||
| const int4 *__restrict__ src_addr_int4 = (const int4 *)src_addr; | ||
| constexpr int N_int4 = sizeof(dtype_t) * N / 16; | ||
| cp_warp_impl<N_int4, UNROLL_FACTOR>(dst_addr_int4, src_addr_int4); | ||
| } else { | ||
| cp_warp_impl<N, UNROLL_FACTOR>(dst_addr, src_addr); | ||
| } | ||
| for (int __i = ((N) / kLoopStride) * kLoopStride + (lane_id); __i < (N); | ||
| __i += 32) | ||
| ST_FUNC(__dst + __i, LD_FUNC(__src + __i)); | ||
| } |
There was a problem hiding this comment.
Same alignment and size validation issues as above.
This overload also has the same critical issues:
- Missing alignment verification for 16-byte boundaries
- Missing size validation for
N * sizeof(dtype_t) % 16 == 0 - Typo:
enable_aggresive_vectorizeshould beenable_aggressive_vectorize
🔎 Proposed fix
template <int N, int UNROLL_FACTOR, bool enable_aggresive_vectorize = false,
typename dtype_t>
TL_DEVICE void cp_warp(dtype_t *const dst_addr, uint64_t src_addr_uint64) {
const dtype_t *src_addr = reinterpret_cast<const dtype_t *>(src_addr_uint64);
if constexpr (enable_aggresive_vectorize) {
+ static_assert((N * sizeof(dtype_t)) % 16 == 0,
+ "Aggressive vectorization requires N*sizeof(dtype_t) to be a multiple of 16");
int4 *__restrict__ dst_addr_int4 = (int4 *)dst_addr;
const int4 *__restrict__ src_addr_int4 = (const int4 *)src_addr;
constexpr int N_int4 = sizeof(dtype_t) * N / 16;
cp_warp_impl<N_int4, UNROLL_FACTOR>(dst_addr_int4, src_addr_int4);
} else {
cp_warp_impl<N, UNROLL_FACTOR>(dst_addr, src_addr);
}
}🤖 Prompt for AI Agents
In src/tl_templates/cuda/copy.h around lines 258-270, fix three issues: rename
the template parameter and all occurrences from enable_aggresive_vectorize to
enable_aggressive_vectorize; add a compile-time check (static_assert) that
(sizeof(dtype_t) * N) % 16 == 0 so vectorized int4 casts are only used when
total bytes are 16-byte aligned; and add runtime alignment validation before
doing the int4 reinterpret_casts (verify both dst_addr and src_addr_uint64 are
16-byte aligned) and handle failure (assert/log and fallback to non-vectorized
path). Ensure the runtime check occurs only when enable_aggressive_vectorize is
true (use if constexpr) and keep the non-vectorized branch unchanged.
| // Reinterpret the value to vector type (e.g., int4 for 8xbf16) | ||
| PrimExpr vec_value = Call(vec_dtype, builtin::reinterpret(), {new_value}); | ||
| PrimExpr vec_value_slice = vec_value.as<CallNode>()->args[0]; | ||
|
|
||
| // Build new args with base addresses and reinterpreted value | ||
| Array<PrimExpr> new_args; | ||
| new_args.push_back(func_name); | ||
| new_args.push_back(new_addr); | ||
| new_args.push_back(vec_value_slice); | ||
| // Copy remaining args (sem, scope, etc.) | ||
| for (size_t i = 3; i < op->args.size(); ++i) { | ||
| new_args.push_back(this->VisitExpr(op->args[i])); | ||
| } | ||
|
|
||
| // Return the vectorized call | ||
| return Call(op->dtype, op->op, new_args); |
There was a problem hiding this comment.
Bug: vec_value_slice extracts the original value, discarding the reinterpret.
Line 798 extracts args[0] from the just-created reinterpret call, which is new_value (the original un-reinterpreted value). This defeats the purpose of the reinterpret operation.
Compare with the correct implementation in loop_vectorization_utils.h (line 485-491) which uses vec_value directly.
Proposed fix
// Reinterpret the value to vector type (e.g., int4 for 8xbf16)
PrimExpr vec_value = Call(vec_dtype, builtin::reinterpret(), {new_value});
- PrimExpr vec_value_slice = vec_value.as<CallNode>()->args[0];
// Build new args with base addresses and reinterpreted value
Array<PrimExpr> new_args;
new_args.push_back(func_name);
new_args.push_back(new_addr);
- new_args.push_back(vec_value_slice);
+ new_args.push_back(vec_value);
// Copy remaining args (sem, scope, etc.)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| // Reinterpret the value to vector type (e.g., int4 for 8xbf16) | |
| PrimExpr vec_value = Call(vec_dtype, builtin::reinterpret(), {new_value}); | |
| PrimExpr vec_value_slice = vec_value.as<CallNode>()->args[0]; | |
| // Build new args with base addresses and reinterpreted value | |
| Array<PrimExpr> new_args; | |
| new_args.push_back(func_name); | |
| new_args.push_back(new_addr); | |
| new_args.push_back(vec_value_slice); | |
| // Copy remaining args (sem, scope, etc.) | |
| for (size_t i = 3; i < op->args.size(); ++i) { | |
| new_args.push_back(this->VisitExpr(op->args[i])); | |
| } | |
| // Return the vectorized call | |
| return Call(op->dtype, op->op, new_args); | |
| // Reinterpret the value to vector type (e.g., int4 for 8xbf16) | |
| PrimExpr vec_value = Call(vec_dtype, builtin::reinterpret(), {new_value}); | |
| // Build new args with base addresses and reinterpreted value | |
| Array<PrimExpr> new_args; | |
| new_args.push_back(func_name); | |
| new_args.push_back(new_addr); | |
| new_args.push_back(vec_value); | |
| // Copy remaining args (sem, scope, etc.) | |
| for (size_t i = 3; i < op->args.size(); ++i) { | |
| new_args.push_back(this->VisitExpr(op->args[i])); | |
| } | |
| // Return the vectorized call | |
| return Call(op->dtype, op->op, new_args); |
🤖 Prompt for AI Agents
In src/transform/vectorize_loop.cc around lines 796 to 811, the code creates a
reinterpret Call (vec_value) but then pulls args[0] from that Call
(vec_value_slice), which is the original un-reinterpreted new_value; replace
that extraction and use the reinterpret Call itself when building new_args.
Concretely, remove the vec_value_slice extraction and push vec_value into
new_args (i.e., use the reinterpret call result directly) so the vectorized call
receives the reinterpreted value.
There was a problem hiding this comment.
Actionable comments posted: 7
🧹 Nitpick comments (14)
examples/distributed/deepseek_deepep/intranode/test_intranode.py (1)
7-10: Test requires 8 GPUs by default — consider adding CI-friendly parameters.The default
--num_ranks=8requires 8 GPUs. For CI environments with fewer GPUs, consider adding explicit parameters like--num_ranks 2to enable testing on smaller setups.@tilelang.testing.requires_cuda def test_intranode(monkeypatch): - monkeypatch.setattr("sys.argv", ["example_intranode.py"]) # optionally add testing params here + monkeypatch.setattr("sys.argv", ["example_intranode.py", "--num_ranks", "2", "--num_experts", "4"]) # 2 GPUs minimum example_intranode.main()examples/distributed/deepseek_deepep/intranode/example_intranode.py (3)
30-33: Address static analysis hints: remove unused noqa and consider exception message.The
noqa: F403directive is unused (line 31 only importsdeep_ep, not*). Also, static analysis flags TRY003 for the long message.🔎 Proposed fix
try: - import deep_ep # noqa: F403 + import deep_ep except ModuleNotFoundError: - raise ModuleNotFoundError("Please install DeepEP to run this test.") from None + raise ModuleNotFoundError("DeepEP required: pip install deep_ep") from None
40-41: Prefix unused variable with underscore.
rank_idxis unpacked but never used. Per static analysis (RUF059), prefix it with an underscore.- x, topk_idx, topk_weights, rank_idx = gen_inputs(num_tokens, hidden, num_topk, num_experts, + x, topk_idx, topk_weights, _rank_idx = gen_inputs(num_tokens, hidden, num_topk, num_experts, num_ranks)
62-63: Prefix unused variable with underscore.
eventis unpacked but never used. Per static analysis (RUF059), prefix it with an underscore.ref_recv_x, ref_recv_topk_idx, ref_recv_topk_weights, ref_num_recv_tokens_per_expert_list, ref_handle, event = \ + ref_recv_x, ref_recv_topk_idx, ref_recv_topk_weights, ref_num_recv_tokens_per_expert_list, ref_handle, _event = \ deepep_buffer.dispatch(x, None, ref_num_tokens_per_rank, None, ref_is_token_in_rank, ref_num_tokens_per_expert, topk_idx, topk_weights, expert_alignment)examples/distributed/deepseek_deepep/deepep_utils.py (3)
14-14: Module-levelnum_smsmay cause confusion withEPBuffer.num_sms.There's a module-level
num_sms = 20here and a class attributeEPBuffer.num_sms = 20inbuffer.py. Changes to one won't affect the other, which could lead to subtle bugs if someone expects them to be synchronized.Consider using a single source of truth, e.g., always referencing
EPBuffer.num_smsor consolidating into a shared constant.
175-177: Minor: Simplify fill operation.
x[:, :].fill_(-1)can be simplified tox.fill_(-1).- x[:, :].fill_(-1) + x.fill_(-1)
231-237: Busy-wait loop may cause high CPU usage.The
while (true)loop polling counters without any backoff or yield could pin a CPU core at 100% while waiting. Consider adding a brief pause or yield if wait times can be significant.// Optional: Add brief pause to reduce CPU spinning #include <thread> // Inside the loop, after the ready check: if (!ready) std::this_thread::yield();examples/distributed/deepseek_deepep/buffer.py (2)
248-270: Address unused unpacked variables from handle.Per static analysis (RUF059), the unpacked handle components at line 250 are never used. Consider using underscores or removing the unpacking since the handle is passed directly to
intranode_dispatch.🔎 Proposed fix
if handle is not None: assert topk_idx is None and topk_weights is None - rank_prefix_matrix, channel_prefix_matrix, recv_channel_prefix_matrix, recv_src_idx, is_token_in_rank, send_head = handle + # Handle structure: (rank_prefix_matrix, channel_prefix_matrix, recv_channel_prefix_matrix, recv_src_idx, is_token_in_rank, send_head) recv_x = intranode_dispatch(
281-302: LGTM!The combine method is a clean wrapper with appropriate documentation. The TODO for bias support is noted.
Would you like me to open an issue to track the bias support TODO mentioned in line 282?
examples/distributed/deepseek_deepep/deepep.md (2)
3-8: Fix markdown formatting issues.The bare URL (line 3) should be a proper link, and the TODO checkboxes (lines 7-8) need spaces to render correctly.
🔎 Proposed fix
-To install and compare with the original DeepEP implementation, please refer to https://github.com/deepseek-ai/DeepEP. +To install and compare with the original DeepEP implementation, please refer to [DeepEP](https://github.com/deepseek-ai/DeepEP). ## TODO - [x] Intranode Normal Mode -- [] Internode Normal Mode -- [] Low-latency Mode +- [ ] Internode Normal Mode +- [ ] Low-latency Mode
146-149: Add language specifier to code block.Per markdown lint (MD040), fenced code blocks should have a language specified.
-``` +```bash CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 TILELANG_USE_DISTRIBUTED=1 python intranode/example_intranode.py \ --num_ranks 8 --num_tokens 4096 --hidden 7168 --num_topk 8 --num_experts 32 [--cached]</blockquote></details> <details> <summary>examples/distributed/deepseek_deepep/intranode/dispatch.py (3)</summary><blockquote> `13-13`: **Remove unused `noqa` directive.** The `noqa: F403` directive is not needed as the wildcard import is not triggering F403. <details> <summary>🔎 Proposed fix</summary> ```diff -from deepep_utils import Config, ep_ext # noqa: F403 +from deepep_utils import Config, ep_extBased on static analysis hints.
15-15: Remove or uncomment the disabled cache line.If
tilelang.disable_cache()is no longer needed, remove it. Otherwise, uncomment it or add a comment explaining why it's disabled.
92-93: Consider the TODO for optimization.The comment suggests returning
per_rank_bufferdirectly asrank_prefix_matrixto avoid the copy operation. This would improve performance slightly.Do you want me to help implement this optimization?
📜 Review details
Configuration used: defaults
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (7)
examples/distributed/deepseek_deepep/buffer.py(1 hunks)examples/distributed/deepseek_deepep/deepep.md(1 hunks)examples/distributed/deepseek_deepep/deepep_utils.py(1 hunks)examples/distributed/deepseek_deepep/intranode/dispatch.py(1 hunks)examples/distributed/deepseek_deepep/intranode/example_intranode.py(1 hunks)examples/distributed/deepseek_deepep/intranode/test_intranode.py(1 hunks)testing/python/language/test_tilelang_language_ldst_options.py(1 hunks)
🧰 Additional context used
🧬 Code graph analysis (4)
examples/distributed/deepseek_deepep/intranode/test_intranode.py (1)
examples/distributed/deepseek_deepep/intranode/example_intranode.py (2)
test_intranode(19-204)main(239-243)
examples/distributed/deepseek_deepep/intranode/example_intranode.py (3)
tilelang/distributed/utils.py (1)
init_dist(40-62)examples/distributed/deepseek_deepep/buffer.py (4)
EPBuffer(13-302)get_dispatch_layout(192-207)dispatch(209-279)combine(281-302)examples/distributed/deepseek_deepep/deepep_utils.py (2)
gen_inputs(129-158)ep_bench(180-217)
testing/python/language/test_tilelang_language_ldst_options.py (1)
tilelang/jit/__init__.py (1)
jit(233-306)
examples/distributed/deepseek_deepep/deepep_utils.py (2)
examples/distributed/deepseek_deepep/buffer.py (1)
num_channels(175-182)tilelang/env.py (1)
get(175-178)
🪛 markdownlint-cli2 (0.18.1)
examples/distributed/deepseek_deepep/deepep.md
3-3: Bare URL used
(MD034, no-bare-urls)
146-146: Fenced code blocks should have a language specified
(MD040, fenced-code-language)
🪛 Ruff (0.14.8)
examples/distributed/deepseek_deepep/intranode/example_intranode.py
31-31: Unused noqa directive (unused: F403)
Remove unused noqa directive
(RUF100)
33-33: Avoid specifying long messages outside the exception class
(TRY003)
40-40: Unpacked variable rank_idx is never used
Prefix it with an underscore or any other dummy variable pattern
(RUF059)
62-62: Unpacked variable event is never used
Prefix it with an underscore or any other dummy variable pattern
(RUF059)
examples/distributed/deepseek_deepep/buffer.py
250-250: Unpacked variable rank_prefix_matrix is never used
Prefix it with an underscore or any other dummy variable pattern
(RUF059)
250-250: Unpacked variable channel_prefix_matrix is never used
Prefix it with an underscore or any other dummy variable pattern
(RUF059)
250-250: Unpacked variable recv_channel_prefix_matrix is never used
Prefix it with an underscore or any other dummy variable pattern
(RUF059)
250-250: Unpacked variable recv_src_idx is never used
Prefix it with an underscore or any other dummy variable pattern
(RUF059)
250-250: Unpacked variable send_head is never used
Prefix it with an underscore or any other dummy variable pattern
(RUF059)
examples/distributed/deepseek_deepep/intranode/dispatch.py
13-13: Unused noqa directive (unused: F403)
Remove unused noqa directive
(RUF100)
508-508: Shebang should contain python, pytest, or uv run
(EXE003)
508-508: Shebang should be at the beginning of the file
(EXE005)
575-575: Unused function argument: num_topk
(ARG001)
576-576: Unused function argument: num_experts
(ARG001)
788-788: Shebang should contain python, pytest, or uv run
(EXE003)
788-788: Shebang should be at the beginning of the file
(EXE005)
🔇 Additional comments (22)
examples/distributed/deepseek_deepep/intranode/example_intranode.py (4)
1-16: Module setup looks reasonable for an example file.The
sys.pathmanipulation is a common pattern for examples that need to import sibling modules. The environment variable setup for NCCL is appropriate for reducing log noise during testing.
188-204: Benchmarking report section looks correct.The bandwidth calculation (
bytes / ms / 1e6) correctly converts to GB/s. The reporting is clear and properly gated to rank 0.
207-222: LGTM!The
runfunction properly initializes the distributed context, executes the test, and cleans up the process group.
225-247: LGTM!Argument parsing and multiprocessing spawn are correctly implemented. Default values align with the PR's benchmark configuration.
examples/distributed/deepseek_deepep/deepep_utils.py (6)
17-90: Config dataclass is well-structured.The factory methods with hardcoded config maps are appropriate for known configurations. The TODO comments for auto-tuning are noted.
93-115: LGTM!Environment setup for NVSHMEM/RDMA is well-documented with inline comments explaining each configuration option.
118-125: LGTM!Simple bias unpacking utility handles both single tensor and tuple cases correctly.
150-158: Note on-1handling in generated inputs.Line 155 masks
rank_idxwheretopk_idx == -1, buttorch.topknever produces-1values. This appears to be defensive code for future cases wheretopk_idxmight have explicit "no selection" markers. The logic is correct.
180-217: LGTM!The benchmarking function follows good practices: L2 cache flush, warmup iterations, skipping the first timed iteration to avoid cold-start effects, and proper CUDA event-based timing.
254-254:-march=nativemay cause portability issues.If the extension is built on one machine and run on another with different CPU features, this could cause illegal instruction errors. For distributed examples this is likely fine, but worth noting.
examples/distributed/deepseek_deepep/buffer.py (5)
1-27: LGTM!Class definition is clean with appropriate docstrings and class-level attributes.
28-78: LGTM!The constructor properly validates inputs, sets up the distributed allocator, and ensures all ranks are synchronized before returning.
86-149: LGTM!Symmetric buffer pre-allocation is well-organized with clear groupings for barrier signals, per-rank/expert metadata, channel offsets, and payload buffers.
154-162: LGTM!Counter preparation correctly handles intranode vs internode cases with conditional RDMA counter allocation.
164-190: LGTM!The properties provide clean access to configuration. Note that
set_num_smsaffects all instances since it modifies the class attribute.examples/distributed/deepseek_deepep/deepep.md (1)
31-202: Documentation is comprehensive and well-organized.The documentation thoroughly covers the intranode API, data structures, kernel responsibilities, execution flow, and usage examples. The notes on limitations (ranks ≤ 8, hidden % 8 == 0 for vectorization) are particularly helpful.
examples/distributed/deepseek_deepep/intranode/test_intranode.py (1)
4-4: Bare import works in pytest context but may fail in spawned child processes.The
import example_intranodeworks when pytest discovers the test, as it adds the test directory tosys.path. However, ifexample_intranode.pyitself spawns child processes (viatorch.multiprocessing.spawnor similar), those child processes won't inherit the test directory fromsys.pathand may fail to import sibling modules. Consider using explicitsys.pathmanipulation inexample_intranode.pyif child process module imports are needed.examples/distributed/deepseek_deepep/intranode/dispatch.py (5)
126-192: LGTM on notify_dispatch wrapper.The function properly initializes the kernel, clears counters, invokes the kernel, and waits for completion before returning the dispatch metadata.
196-217: LGTM on cached_notify_dispatch_kernel.The cached version correctly simplifies the dispatch notification by only clearing the channel buffers while relying on cached routing metadata.
220-245: LGTM on cached_notify_dispatch wrapper.The function properly wraps the cached kernel with appropriate initialization and invocation.
247-562: LGTM on dispatch_kernel implementation.The kernel correctly implements a producer-consumer pattern with:
- Proper backpressure via
T.wait_ge(line 356-359)- Acquire/release memory ordering for synchronization (lines 488, 442, 723)
- Correct topk_idx transformation to local expert indices (lines 408-417)
- Appropriate use of warp-level primitives for data movement
650-651: Verify the difference in division method.Line 650 uses
T.ceildiv(num_tokens, num_channels)while the non-cached dispatch_kernel at line 342 usesT.truncdiv(num_tokens + num_channels - 1, num_channels)with a comment about a TVM bug workaround.Ensure that using
T.ceildivhere doesn't trigger the same TVM bug, or if the bug has been fixed, consider updating both implementations for consistency.
| This example implements DeepEP’s intra‑node (NVLink) dispatch/combine using TileScale kernels. | ||
| z | ||
| The intra‑node path lives under `intranode/` and provides a minimal public API that mirrors DeepEP’s behavior for NVLink‑connected ranks. |
There was a problem hiding this comment.
Remove stray character "z".
Line 27 contains a lone "z" which appears to be an accidental keystroke.
This example implements DeepEP's intra‑node (NVLink) dispatch/combine using TileScale kernels.
-z
+
The intra‑node path lives under `intranode/` and provides a minimal public API that mirrors DeepEP's behavior for NVLink‑connected ranks.📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| This example implements DeepEP’s intra‑node (NVLink) dispatch/combine using TileScale kernels. | |
| z | |
| The intra‑node path lives under `intranode/` and provides a minimal public API that mirrors DeepEP’s behavior for NVLink‑connected ranks. | |
| This example implements DeepEP's intra‑node (NVLink) dispatch/combine using TileScale kernels. | |
| The intra‑node path lives under `intranode/` and provides a minimal public API that mirrors DeepEP's behavior for NVLink‑connected ranks. |
🤖 Prompt for AI Agents
In examples/distributed/deepseek_deepep/deepep.md around lines 26 to 28 there is
a stray character "z" on line 27; remove that lone "z" so the paragraph reads
smoothly and does not contain the accidental keystroke, then save the markdown
file.
| return recv_x, recv_topk_idx, recv_topk_weights, num_recv_tokens_per_expert_list, handle | ||
| else: | ||
| kernel = cached_dispatch_kernel(num_ranks, num_tokens, | ||
| config.num_max_nvl_chunked_send_tokens, | ||
| config.num_max_nvl_chunked_recv_tokens, hidden, num_topk, | ||
| num_experts, config.num_sms, 'bfloat16') | ||
| kernel.initialize(allocator=allocator, stream=comm_stream.cuda_stream) | ||
| kernel( | ||
| rank, | ||
| recv_x, | ||
| recv_src_idx, | ||
| recv_channel_prefix_matrix, | ||
| send_head, | ||
| x, | ||
| is_token_in_rank, | ||
| rank_prefix_matrix, | ||
| channel_prefix_matrix, | ||
| channel_start_offset, | ||
| channel_end_offset, | ||
| channel_head_idx, | ||
| channel_tail_idx, | ||
| channel_x_buffers, | ||
| channel_src_idx_buffers, | ||
| stream=comm_stream.cuda_stream, | ||
| skip_tensor_validation=True) # reduce runtime overhead | ||
| return recv_x |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Find the intranode_dispatch function definition and its calls
rg -nP --type=py 'intranode_dispatch\s*\(' -C 3Repository: tile-ai/tilescale
Length of output: 2253
🏁 Script executed:
#!/bin/bash
# Read the intranode_dispatch function definition and return statements
sed -n '831,972p' examples/distributed/deepseek_deepep/intranode/dispatch.py | head -200Repository: tile-ai/tilescale
Length of output: 5864
🏁 Script executed:
#!/bin/bash
# Check the context around the two intranode_dispatch calls in buffer.py
echo "=== CALL 1 (Line 251) ==="
sed -n '240,275p' examples/distributed/deepseek_deepep/buffer.py
echo -e "\n=== CALL 2 (Line 273) ==="
sed -n '265,310p' examples/distributed/deepseek_deepep/buffer.pyRepository: tile-ai/tilescale
Length of output: 4281
Return type inconsistency in intranode_dispatch function—mitigate by returning consistent tuple structure.
The function returns a 5-tuple (recv_x, recv_topk_idx, recv_topk_weights, num_recv_tokens_per_expert_list, handle) in non-cached mode (handle=None) but only recv_x in cached mode (handle is not None). While the current caller in buffer.py properly handles both cases via explicit if-else branches, this design is error-prone for any future direct callers of intranode_dispatch.
Standardize the return type to always return a 5-tuple, with None values for unused fields in cached mode (e.g., (recv_x, None, None, None, None)), or wrap the function with documented type stubs to clearly indicate the return behavior based on the handle parameter.
🤖 Prompt for AI Agents
In examples/distributed/deepseek_deepep/intranode/dispatch.py around lines 947
to 972, the function returns a 5-tuple in the non-cached branch but only a
single value in the cached branch; change the cached branch to return a
consistent 5-tuple (e.g., (recv_x, None, None, None, None)) so callers always
receive the same structure, or alternatively add clear type annotations and
docstring stating the conditional return shape—implement the former for safety
by returning the unused fields as None in the cached path.
| kernel = get_st_kernel(scope, sem, na) | ||
| x = torch.randint(0, 100, (32,), device="cuda", dtype=torch.int32) | ||
| kernel(x) | ||
| assert x.equal(torch.arange(32, device="cuda")) |
There was a problem hiding this comment.
Use consistent assertion method.
Line 38 uses torch.equal(x, y) while this line uses x.equal(torch.arange(...)). Both work, but consistency improves readability.
🔎 Suggested fix for consistency
- assert x.equal(torch.arange(32, device="cuda"))
+ assert torch.equal(x, torch.arange(32, device="cuda", dtype=torch.int32))📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| assert x.equal(torch.arange(32, device="cuda")) | |
| assert torch.equal(x, torch.arange(32, device="cuda", dtype=torch.int32)) |
🤖 Prompt for AI Agents
In testing/python/language/test_tilelang_language_ldst_options.py at line 66,
the assertion uses the instance method x.equal(torch.arange(...)) which is
inconsistent with line 38 that uses the functional style torch.equal(x, y);
update line 66 to use the same functional form by calling torch.equal(x,
torch.arange(32, device="cuda")) inside the assert so both assertions follow the
same style and improve readability.
There was a problem hiding this comment.
Actionable comments posted: 8
♻️ Duplicate comments (6)
src/transform/vectorize_loop.cc (1)
796-806: Critical: vec_value_lvalue extracts the original value, discarding the reinterpret.Line 800 extracts
args[0]from the just-createdreinterpretcall (line 797), which retrievesnew_value—the original un-reinterpreted value. This defeats the purpose of the reinterpret operation.Compare with the correct implementation in
src/transform/common/loop_vectorization_utils.h(lines 485-491), which usesvec_valuedirectly.🔎 Proposed fix
// Reinterpret the value to vector type (e.g., int4 for 8xbf16) PrimExpr vec_value = Call(vec_dtype, builtin::reinterpret(), {new_value}); - - // A trick to get the lvalue of the vectorized value - PrimExpr vec_value_lvalue = vec_value.as<CallNode>()->args[0]; // Build new args with base addresses and reinterpreted value Array<PrimExpr> new_args; new_args.push_back(func_name); new_args.push_back(new_addr); - new_args.push_back(vec_value_lvalue); + new_args.push_back(vec_value);testing/python/language/test_tilelang_language_ldst_options.py (1)
66-66: Use consistent assertion method and explicit dtype specification.This issue was previously identified: Line 38 uses
torch.equal(x, y)while this line uses the instance methodx.equal(...). Additionally,torch.arange(32, device="cuda")should explicitly specifydtype=torch.int32to match the tensor dtype defined at line 64.🔎 Suggested fix for consistency and type safety
- assert x.equal(torch.arange(32, device="cuda")) + assert torch.equal(x, torch.arange(32, device="cuda", dtype=torch.int32))src/tl_templates/cuda/copy.h (2)
243-270: Critical: Same alignment and size validation issues apply to these overloads.Both overloads have the same critical issues as lines 221-241:
- Missing 16-byte alignment validation for
int4*casts- Missing size validation that
N * sizeof(dtype_t) % 16 == 0These overloads are particularly vulnerable since they accept
uint64_taddresses that might not be properly aligned.🔎 Proposed fix for both overloads
For lines 243-256:
template <int N, int UNROLL_FACTOR, bool enable_aggressive_vectorize = false, typename dtype_t> TL_DEVICE void cp_warp(uint64_t dst_addr_uint64, dtype_t const *const src_addr) { dtype_t *dst_addr = reinterpret_cast<dtype_t *>(dst_addr_uint64); if constexpr (enable_aggressive_vectorize) { + static_assert((N * sizeof(dtype_t)) % 16 == 0, + "Aggressive vectorization requires N*sizeof(dtype_t) to be a multiple of 16"); + if (dst_addr_uint64 % 16 != 0 || (uintptr_t)src_addr % 16 != 0) { + cp_warp_impl<N, UNROLL_FACTOR>(dst_addr, src_addr); + return; + } int4 *__restrict__ dst_addr_int4 = (int4 *)dst_addr; const int4 *__restrict__ src_addr_int4 = (const int4 *)src_addr; constexpr int N_int4 = sizeof(dtype_t) * N / 16; cp_warp_impl<N_int4, UNROLL_FACTOR>(dst_addr_int4, src_addr_int4); } else { cp_warp_impl<N, UNROLL_FACTOR>(dst_addr, src_addr); } }For lines 258-270:
template <int N, int UNROLL_FACTOR, bool enable_aggressive_vectorize = false, typename dtype_t> TL_DEVICE void cp_warp(dtype_t *const dst_addr, uint64_t src_addr_uint64) { const dtype_t *src_addr = reinterpret_cast<const dtype_t *>(src_addr_uint64); if constexpr (enable_aggressive_vectorize) { + static_assert((N * sizeof(dtype_t)) % 16 == 0, + "Aggressive vectorization requires N*sizeof(dtype_t) to be a multiple of 16"); + if ((uintptr_t)dst_addr % 16 != 0 || src_addr_uint64 % 16 != 0) { + cp_warp_impl<N, UNROLL_FACTOR>(dst_addr, src_addr); + return; + } int4 *__restrict__ dst_addr_int4 = (int4 *)dst_addr; const int4 *__restrict__ src_addr_int4 = (const int4 *)src_addr; constexpr int N_int4 = sizeof(dtype_t) * N / 16; cp_warp_impl<N_int4, UNROLL_FACTOR>(dst_addr_int4, src_addr_int4); } else { cp_warp_impl<N, UNROLL_FACTOR>(dst_addr, src_addr); } }
221-241: Critical: Alignment and size validation still missing for aggressive vectorization.The past review identified that casting to
int4*without validation can cause misaligned memory access crashes or silent data truncation. While the typo has been fixed, the core safety issues remain unresolved:
- Missing alignment checks: Both
src_addranddst_addrmust be 16-byte aligned before casting toint4*- Missing size validation:
N * sizeof(dtype_t)must be a multiple of 16, or the integer divisionN_int4 = sizeof(dtype_t) * N / 16will silently truncateCompare with
nvshmem_cp_threadgroup(lines 276-369), which performs runtime alignment checks before vectorized copies.🔎 Proposed fix with validation
template <int N, int UNROLL_FACTOR, bool enable_aggressive_vectorize = false, typename dtype_t> TL_DEVICE void cp_warp(dtype_t const *const dst_addr, dtype_t const *const src_addr) { if constexpr (enable_aggressive_vectorize) { + static_assert((N * sizeof(dtype_t)) % 16 == 0, + "Aggressive vectorization requires N*sizeof(dtype_t) to be a multiple of 16"); + // Runtime alignment check (can use assert in debug builds) + if ((uintptr_t)dst_addr % 16 != 0 || (uintptr_t)src_addr % 16 != 0) { + // Fallback to non-vectorized path + cp_warp_impl<N, UNROLL_FACTOR>(dst_addr, src_addr); + return; + } int4 *__restrict__ dst_addr_int4 = (int4 *)dst_addr; const int4 *__restrict__ src_addr_int4 = (const int4 *)src_addr; constexpr int N_int4 = sizeof(dtype_t) * N / 16; cp_warp_impl<N_int4, UNROLL_FACTOR>(dst_addr_int4, src_addr_int4); } else { cp_warp_impl<N, UNROLL_FACTOR>(dst_addr, src_addr); } }examples/distributed/deepseek_deepep/intranode/dispatch.py (2)
899-902: Critical:recv_src_idxis still incorrectly overwritten in cached mode.As noted in the previous review,
recv_src_idxis extracted from the handle at line 854 and used at line 899 to determinenum_recv_tokens. However, it is unconditionally overwritten at line 902 with a new empty tensor, discarding the cached value.🔎 Proposed fix
recv_x = torch.empty((num_recv_tokens, hidden), dtype=x.dtype, device='cuda') - recv_src_idx = torch.empty((num_recv_tokens,), dtype=torch.int32, device='cuda') + if handle is None: + recv_src_idx = torch.empty((num_recv_tokens,), dtype=torch.int32, device='cuda') + # else: recv_src_idx already unpacked from handle at line 854 if handle is None: recv_topk_idx = torch.empty((num_recv_tokens, num_topk), dtype=torch.int64, device='cuda')
945-970: Return type inconsistency remains unaddressed.The function returns a 5-tuple in non-cached mode (line 945) but only
recv_xin cached mode (line 970). This inconsistency is error-prone for callers.🔎 Proposed fix for consistent return type
else: kernel = cached_dispatch_kernel(num_ranks, num_tokens, config.num_max_nvl_chunked_send_tokens, config.num_max_nvl_chunked_recv_tokens, hidden, config.num_sms, 'bfloat16') kernel.initialize(allocator=allocator, stream=comm_stream.cuda_stream) kernel( rank, recv_x, recv_src_idx, recv_channel_prefix_matrix, send_head, x, is_token_in_rank, rank_prefix_matrix, channel_prefix_matrix, channel_start_offset, channel_end_offset, channel_head_idx, channel_tail_idx, channel_x_buffers, channel_src_idx_buffers, stream=comm_stream.cuda_stream, skip_tensor_validation=True) # reduce runtime overhead - return recv_x + return recv_x, None, None, None, None # Consistent 5-tuple structure
🧹 Nitpick comments (11)
src/op/remote_copy.h (2)
51-92: Includeenable_aggressive_vectorizein equality, hashing, and reflection.The
enable_aggressive_vectorizefield affects code generation behavior (it controls the template parameter passed totl::cp_warpin theLowermethod), but it is not included in:
RegisterReflection()(lines 51-64)SEqualReduce()(lines 66-77)SHashReduce()(lines 79-92)This omission means two
PutOpnodes with different vectorization settings would be considered structurally equal, potentially causing incorrect caching, memoization, or optimization decisions.🔎 Proposed fix
static void RegisterReflection() { namespace refl = tvm::ffi::reflection; refl::ObjectDef<PutOpNode>() .def_ro("src_addr", &PutOpNode::src_addr) .def_ro("dst_addr", &PutOpNode::dst_addr) .def_ro("copy_size", &PutOpNode::copy_size) .def_ro("dst_pe", &PutOpNode::dst_pe) .def_ro("unroll_factor", &PutOpNode::unroll_factor) .def_ro("src_buffer", &PutOpNode::src_buffer) .def_ro("dst_buffer", &PutOpNode::dst_buffer) .def_ro("src_indices", &PutOpNode::src_indices) .def_ro("dst_indices", &PutOpNode::dst_indices) - .def_ro("scope", &PutOpNode::scope); + .def_ro("scope", &PutOpNode::scope) + .def_ro("enable_aggressive_vectorize", &PutOpNode::enable_aggressive_vectorize); } bool SEqualReduce(const PutOpNode *other, SEqualReducer equal) const { return equal(src_addr, other->src_addr) && equal(dst_addr, other->dst_addr) && equal(src_offset, other->src_offset) && equal(dst_offset, other->dst_offset) && equal(copy_size, other->copy_size) && equal(dst_pe, other->dst_pe) && equal(unroll_factor, other->unroll_factor) && equal(src_buffer, other->src_buffer) && equal(dst_buffer, other->dst_buffer) && equal(src_indices, other->src_indices) && - equal(dst_indices, other->dst_indices) && scope == other->scope; + equal(dst_indices, other->dst_indices) && scope == other->scope && + enable_aggressive_vectorize == other->enable_aggressive_vectorize; } void SHashReduce(SHashReducer hash_reduce) const { hash_reduce(src_addr); hash_reduce(dst_addr); hash_reduce(src_offset); hash_reduce(dst_offset); hash_reduce(copy_size); hash_reduce(dst_pe); hash_reduce(unroll_factor); hash_reduce(src_buffer); hash_reduce(dst_buffer); hash_reduce(src_indices); hash_reduce(dst_indices); hash_reduce(scope); + hash_reduce(enable_aggressive_vectorize); }
143-184: Includeenable_aggressive_vectorizein equality, hashing, and reflection.The same issue exists for
GetOpNodeas forPutOpNode: theenable_aggressive_vectorizefield affects code generation but is not included in reflection, equality, or hashing operations.🔎 Proposed fix
static void RegisterReflection() { namespace refl = tvm::ffi::reflection; refl::ObjectDef<GetOpNode>() .def_ro("src_addr", &GetOpNode::src_addr) .def_ro("dst_addr", &GetOpNode::dst_addr) .def_ro("copy_size", &GetOpNode::copy_size) .def_ro("src_pe", &GetOpNode::src_pe) .def_ro("unroll_factor", &GetOpNode::unroll_factor) .def_ro("src_buffer", &GetOpNode::src_buffer) .def_ro("dst_buffer", &GetOpNode::dst_buffer) .def_ro("src_indices", &GetOpNode::src_indices) .def_ro("dst_indices", &GetOpNode::dst_indices) - .def_ro("scope", &GetOpNode::scope); + .def_ro("scope", &GetOpNode::scope) + .def_ro("enable_aggressive_vectorize", &GetOpNode::enable_aggressive_vectorize); } bool SEqualReduce(const GetOpNode *other, SEqualReducer equal) const { return equal(src_addr, other->src_addr) && equal(dst_addr, other->dst_addr) && equal(src_offset, other->src_offset) && equal(dst_offset, other->dst_offset) && equal(copy_size, other->copy_size) && equal(src_pe, other->src_pe) && equal(unroll_factor, other->unroll_factor) && equal(src_buffer, other->src_buffer) && equal(dst_buffer, other->dst_buffer) && equal(src_indices, other->src_indices) && - equal(dst_indices, other->dst_indices) && scope == other->scope; + equal(dst_indices, other->dst_indices) && scope == other->scope && + enable_aggressive_vectorize == other->enable_aggressive_vectorize; } void SHashReduce(SHashReducer hash_reduce) const { hash_reduce(src_addr); hash_reduce(dst_addr); hash_reduce(src_offset); hash_reduce(dst_offset); hash_reduce(copy_size); hash_reduce(src_pe); hash_reduce(unroll_factor); hash_reduce(src_buffer); hash_reduce(dst_buffer); hash_reduce(src_indices); hash_reduce(dst_indices); hash_reduce(scope); + hash_reduce(enable_aggressive_vectorize); }tilelang/language/distributed/common.py (1)
116-122: Consider adding docstring for BinaryRelation enum.The
BinaryRelationenum is used by the wait APIs but lacks a docstring explaining its purpose and relationship to the wait functions.🔎 Suggested docstring
class BinaryRelation(Enum): + """Binary comparison relations for wait operations. + + Used to specify the comparison type in wait_* synchronization primitives. + """ EQ = 0 NE = 1 GE = 2 LE = 3 GT = 4 LT = 5src/tl_templates/cuda/ldst.h (2)
220-234: Consider adding runtime pointer alignment checks.The inline assembly for vectorized operations (especially 16-byte) typically requires proper alignment. Consider adding runtime assertions or compile-time checks to ensure
ptrmeets alignment requirements for the operation size.🔎 Example alignment check
template <Semantic semantic, Scope scope, bool na, typename P, typename T> TL_DEVICE void st(P ptr, T value) { static_assert(sizeof(T) == 2 || sizeof(T) == 4 || sizeof(T) == 8 || sizeof(T) == 16, "tl::st: T must be 2, 4, 8, or 16 bytes"); static_assert(std::is_pointer_v<P> || std::is_same_v<P, uint64_t>, "tl::st: P must be a pointer or uint64_t"); + + // Alignment check for vectorized stores + if constexpr (sizeof(T) >= 8) { + assert((reinterpret_cast<uintptr_t>(ptr) % sizeof(T)) == 0 && + "tl::st: pointer must be aligned to type size"); + } T *ptr_ = reinterpret_cast<T *>(ptr); StImpl<semantic, scope, na>::execute(ptr_, value); }
253-253: Track TODO for additional ld variant.The TODO mentions support for
ld.global.nc.L1::no_allocate.L2::256B. Consider opening an issue to track this if it's needed for performance optimization.Would you like me to open a tracking issue for this feature?
testing/python/language/test_tilelang_language_ldst_options.py (1)
64-64: Consider using zeros instead of randint for minor efficiency.The random tensor initialization is unnecessary since the kernel overwrites all values by storing thread indices. Using
torch.zeros(32, device="cuda", dtype=torch.int32)would be slightly more efficient.🔎 Optional optimization
- x = torch.randint(0, 100, (32,), device="cuda", dtype=torch.int32) + x = torch.zeros(32, device="cuda", dtype=torch.int32)examples/distributed/deepseek_deepep/intranode/combine.py (2)
51-51: Consider using a named constant for the magic number.The value
2**25is used as a heuristic large number. Consider extracting this to a named constant for better readability and maintainability.🔎 Proposed fix
+MAX_EXPECTED_HEAD = 2**25 # Heuristic large value for initial head + @tilelang.jit(pass_configs={"tl.disable_tma_lower": True, "tl.disable_warp_specialized": True}) def cached_notify_combine_kernel(num_ranks, num_sms): ... - last_head = T.alloc_var('int32', init=2**25) # a heuristic large number + last_head = T.alloc_var('int32', init=MAX_EXPECTED_HEAD)
248-258: Clarify the infinite loop pattern.The condition
lane_id < num_ranksis invariant within the loop sincelane_idis derived from thread binding and never changes. This effectively creates an infinite loop for eligible lanes that exits viaT.loop_break()when all warps retire. Consider adding a brief comment explaining this pattern for maintainability.examples/distributed/deepseek_deepep/intranode/dispatch.py (3)
13-13: Remove unusednoqadirective.The
# noqa: F403comment is unnecessary since there's noimport *statement on this line.🔎 Proposed fix
-from deepep_utils import Config, ep_ext # noqa: F403 +from deepep_utils import Config, ep_extBased on static analysis hints.
83-88: Consider usingT.alloc_varinstead of single-element array.Using
T.alloc_local([1], 'int32')for a single scalar is less idiomatic thanT.alloc_var('int32').🔎 Proposed fix
if tx < num_local_experts: - sum = T.alloc_local([1], 'int32') - sum[0] = 0 + sum_val = T.alloc_var('int32') + sum_val = 0 for i in T.serial(0, num_ranks): - sum[0] += per_expert_buffer[i, tx] - sum[0] = T.ceildiv(sum[0], expert_alignment) * expert_alignment # align up - moe_recv_expert_counter_mapped[tx] = sum[0] + sum_val += per_expert_buffer[i, tx] + sum_val = T.ceildiv(sum_val, expert_alignment) * expert_alignment # align up + moe_recv_expert_counter_mapped[tx] = sum_val
509-516: Document the magic numbers inT.put_warpcall.The values
-1and5fordst_peandunroll_factorrespectively are used without explanation. Consider adding a comment or using named constants for clarity.🔎 Proposed fix
+ LOCAL_PE = -1 # -1 indicates local memory copy + UNROLL_FACTOR = 5 T.put_warp( T.address_of(channel_x_buffers[responsible_channel, responsible_rank, token_idx_in_buffer, 0]), T.address_of(recv_x[total_offset + chunk_idx, 0]), hidden, - -1, - 5, + LOCAL_PE, + UNROLL_FACTOR, enable_aggressive_vectorize=True)
📜 Review details
Configuration used: defaults
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (10)
examples/distributed/deepseek_deepep/intranode/combine.py(1 hunks)examples/distributed/deepseek_deepep/intranode/dispatch.py(1 hunks)src/op/remote_copy.cc(4 hunks)src/op/remote_copy.h(3 hunks)src/tl_templates/cuda/copy.h(5 hunks)src/tl_templates/cuda/ldst.h(1 hunks)src/transform/vectorize_loop.cc(2 hunks)testing/python/language/test_tilelang_language_ldst_options.py(1 hunks)tilelang/distributed/utils.py(3 hunks)tilelang/language/distributed/common.py(6 hunks)
🧰 Additional context used
🧬 Code graph analysis (6)
tilelang/distributed/utils.py (1)
tilelang/utils/ts_ext/tensor.cpp (2)
create_host_device_tensor(89-113)create_host_device_tensor(90-91)
testing/python/language/test_tilelang_language_ldst_options.py (2)
tilelang/jit/__init__.py (1)
jit(233-306)src/tl_templates/cuda/ldst.h (2)
ld(238-251)st(221-234)
src/op/remote_copy.h (1)
src/op/remote_copy.cc (10)
is_distributed(86-89)is_distributed(86-86)is_distributed(196-199)is_distributed(196-196)is_distributed(264-267)is_distributed(264-264)is_distributed(333-336)is_distributed(333-333)StOp(248-262)LdOp(316-331)
tilelang/language/distributed/common.py (1)
tilelang/language/tir/op.py (2)
address_of(464-480)call_intrin(120-145)
src/transform/vectorize_loop.cc (2)
tilelang/language/tir/op.py (3)
call_extern(173-195)address_of(464-480)reinterpret(1899-1918)src/transform/common/loop_vectorization_utils.h (1)
MutateTlLdStExpr_(389-521)
src/tl_templates/cuda/ldst.h (3)
src/tl_templates/cuda/nvrtc_std.h (1)
std(47-111)tilelang/language/builtin.py (2)
st(777-808)ld(741-774)src/tl_templates/cuda/debug.h (1)
uint64_t(268-272)
🪛 Ruff (0.14.8)
examples/distributed/deepseek_deepep/intranode/combine.py
238-238: Shebang should contain python, pytest, or uv run
(EXE003)
238-238: Shebang should be at the beginning of the file
(EXE005)
examples/distributed/deepseek_deepep/intranode/dispatch.py
13-13: Unused noqa directive (unused: F403)
Remove unused noqa directive
(RUF100)
508-508: Shebang should contain python, pytest, or uv run
(EXE003)
508-508: Shebang should be at the beginning of the file
(EXE005)
786-786: Shebang should contain python, pytest, or uv run
(EXE003)
786-786: Shebang should be at the beginning of the file
(EXE005)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: build-test-nvidia
🔇 Additional comments (22)
tilelang/distributed/utils.py (3)
22-22: LGTM!The import of
create_host_device_tensorfromtilescale_extis correct and aligns with the C++ extension implementation.
404-405: LGTM!The function correctly delegates to
create_host_device_tensorand the return type annotationtuple[torch.Tensor, torch.Tensor]accurately reflects the underlying C++ implementation that returns a pair of tensors (host and device).
231-231: Parameter order change is complete—all call sites have been updated correctly.The function signature was changed from
perf_fn(fn, rep, warmup)toperf_fn(fn, warmup, rep). All 27+ call sites across benchmark and example files use the correct parameter order, passing warmup iterations before repetitions. The implementation confirms this ordering is correct: warmup iterations are executed first without timing, then rep iterations are timed.src/op/remote_copy.cc (3)
81-81: Typo fix confirmed.The previous typo
enable_aggresive_vectorizehas been correctly fixed toenable_aggressive_vectorize, consistent with the updated field name.
86-89: LGTM! Consistent distributed operation detection.The
is_distributed()implementation correctly identifies distributed operations by checking if the PE parameter is not the sentinel value -1. The updatedLowermethod properly uses this to determine whether to perform remote address remapping.Also applies to: 105-118
385-404: LGTM! Input counts are now consistent.The TIR registrations correctly match the number of arguments accessed in each operator's constructor:
PutOp: 7 inputs (args[0-6])GetOp: 7 inputs (args[0-6])StOp: 6 inputs (args[0-5])LdOp: 7 inputs (args[0-6])The reflection registrations are properly initialized for all operator nodes.
src/op/remote_copy.h (2)
205-260: LGTM! Complete and consistent StOp implementation.The
StOpNodeandStOpclasses are well-structured with:
- All fields properly exposed in reflection
- Complete structural equality and hashing implementations
- Proper TVM object macros and type declarations
262-320: LGTM! Complete and consistent LdOp implementation.The
LdOpNodeandLdOpclasses are well-structured with:
- All fields properly exposed in reflection
- Complete structural equality and hashing implementations
- Proper TVM object macros and type declarations
src/transform/vectorize_loop.cc (1)
459-470: Verify vectorization logic handles non-vectorizable tl::ld/tl::st safely.The new routing to
MutateTlLdStExpr_fortl::ldandtl::stcalls looks appropriate. Ensure that edge cases (insufficient arguments, non-BufferLoad operands, etc.) are handled correctly in the helper to avoid scalarization fallback loops.tilelang/language/distributed/common.py (2)
22-47: LGTM: enable_aggressive_vectorize parameter added correctly.The new
enable_aggressive_vectorizeparameter is properly threaded through to the intrinsic call and documented.
125-162: Clarify semantic difference between wait_eq and other wait functions.Line 127 flags that
wait_eqhas different semantics compared to the other wait functions.wait_equses thetl.wait_eqintrinsic while others usetl.waitwith aBinaryRelationvalue. This inconsistency could confuse users.Consider:
- Documenting the semantic difference in the function docstrings
- Unifying the implementation if the semantics are intended to be the same
- Resolving the TODO comment
Can you clarify the intended semantic difference between
wait_eqand the other wait functions? If there's no fundamental difference, consider refactoringwait_eqto use the sametl.waitintrinsic pattern:def wait_eq(ptr: PrimExpr, expected: PrimExpr, peer: PrimExpr | None = -1): """Wait until *ptr == expected""" return tir.call_intrin("handle", tir.op.Op.get("tl.wait"), BinaryRelation.EQ.value, address_of(ptr), expected, peer)src/tl_templates/cuda/ldst.h (2)
47-59: Verify bfloat16 reinterpret_cast is safe across compilers.The bfloat16 handling uses
reinterpret_cast<uint16_t *>(&value)to extract bits for the assembly constraint. While this works for CUDA's__nv_bfloat16andcutlass::bfloat16_t, ensure this is safe and well-defined for both types.Consider adding a static_assert to verify alignment and layout assumptions if needed for safety across different bf16 implementations.
84-89: Assembly operand order now correct.The 16-byte store assembly format
[%0], {%1, %2, %3, %4}with operand list"l"(ptr), "r"(value.x/y/z/w)is now correct. The past review comment appears to have been addressed.src/tl_templates/cuda/copy.h (4)
114-124: LGTM!The new int16_t and uint16_t specializations correctly follow the established pattern for non-cached global loads, with appropriate PTX instructions (.s16/.u16) and constraint letters.
168-176: LGTM!The new int16_t and uint16_t specializations for non-allocating global stores are implemented correctly with appropriate PTX instructions and constraint letters.
199-219: LGTM!Refactoring the original implementation to
cp_warp_implas an internal helper is a clean approach to support the new vectorization wrappers while preserving the existing logic.
97-97: The verification of this change cannot be completed because git history does not show a transition from "ld.volatile.global" to "ld.nc.global". The code currently uses "ld.nc.global" in cp_warp_impl, which is appropriate for warp-level streaming memory copies that don't require cross-SM synchronization semantics. Confirm whether this change was made in an earlier state or if the review comment is based on a different version of the code.examples/distributed/deepseek_deepep/intranode/combine.py (2)
1-16: LGTM on imports and setup.The imports and environment setup are appropriate for the intranode combine operation. The
tilelang.disable_cache()call and NCCL debug silencing are reasonable for this distributed kernel development context.
330-353: Reduction logic looks correct.The vectorized reduction pattern correctly accumulates values from all topk ranks into the output tensor. The use of
T.vectorized(8)aligns with thehidden % 8 == 0assertion.examples/distributed/deepseek_deepep/intranode/dispatch.py (3)
194-244: LGTM on cached notify dispatch.The cached path correctly performs only buffer clearing and barrier synchronization, avoiding the full prefix computation. The same
comm_streamnull check issue applies here as noted in combine.py.
569-577: Past review issue addressed: unused arguments removed.The
num_topkandnum_expertsparameters have been removed fromcached_dispatch_kernelas suggested in the previous review.
856-859: PotentialAttributeErrorwhen accessing shape ofNonetensors.When
handle is not None,num_tokens_per_expertandtopk_idxcould still beNone(as they'reOptional), but lines 857 and 859 attempt to access.shape[0]and.shape[1]respectively. While the ternary guards against this fornum_expertsandnum_topk, consider adding explicit assertions or restructuring to avoid confusion.
| kernel = cached_notify_combine_kernel(num_ranks, num_sms) | ||
| kernel.initialize(allocator=allocator, stream=comm_stream.cuda_stream) | ||
|
|
||
| kernel( | ||
| send_head, | ||
| channel_head_idx, | ||
| channel_tail_idx, | ||
| barrier_signal, | ||
| stream=comm_stream.cuda_stream, | ||
| skip_tensor_validation=True) # reduce runtime overhead |
There was a problem hiding this comment.
Missing null check for comm_stream parameter.
If comm_stream is None, accessing comm_stream.cuda_stream at lines 87 and 94 will raise an AttributeError. Consider adding a guard or defaulting to the current CUDA stream.
🔎 Proposed fix
def cached_notify_combine(
num_ranks,
num_sms,
##### symm buffers #####
send_head: torch.Tensor,
channel_head_idx: torch.Tensor,
channel_tail_idx: torch.Tensor,
barrier_signal: torch.Tensor,
allocator,
comm_stream=None):
kernel = cached_notify_combine_kernel(num_ranks, num_sms)
- kernel.initialize(allocator=allocator, stream=comm_stream.cuda_stream)
+ stream = comm_stream.cuda_stream if comm_stream is not None else torch.cuda.current_stream().cuda_stream
+ kernel.initialize(allocator=allocator, stream=stream)
kernel(
send_head,
channel_head_idx,
channel_tail_idx,
barrier_signal,
- stream=comm_stream.cuda_stream,
+ stream=stream,
skip_tensor_validation=True) # reduce runtime overhead📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| kernel = cached_notify_combine_kernel(num_ranks, num_sms) | |
| kernel.initialize(allocator=allocator, stream=comm_stream.cuda_stream) | |
| kernel( | |
| send_head, | |
| channel_head_idx, | |
| channel_tail_idx, | |
| barrier_signal, | |
| stream=comm_stream.cuda_stream, | |
| skip_tensor_validation=True) # reduce runtime overhead | |
| kernel = cached_notify_combine_kernel(num_ranks, num_sms) | |
| stream = comm_stream.cuda_stream if comm_stream is not None else torch.cuda.current_stream().cuda_stream | |
| kernel.initialize(allocator=allocator, stream=stream) | |
| kernel( | |
| send_head, | |
| channel_head_idx, | |
| channel_tail_idx, | |
| barrier_signal, | |
| stream=stream, | |
| skip_tensor_validation=True) # reduce runtime overhead |
🤖 Prompt for AI Agents
In examples/distributed/deepseek_deepep/intranode/combine.py around lines 86 to
95, calls access comm_stream.cuda_stream without checking comm_stream; if
comm_stream is None this raises AttributeError. Fix by guarding uses: set a
local variable (e.g., cuda_stream) to comm_stream.cuda_stream if comm_stream is
not None, otherwise obtain a safe default like
torch.cuda.current_stream().cuda_stream or the appropriate default stream object
used elsewhere in the repo, then pass that local cuda_stream into
kernel.initialize(...) and kernel(...); ensure any type expectations match
existing code.
| kernel = combine_kernel( | ||
| num_ranks, | ||
| config.num_max_nvl_chunked_send_tokens, | ||
| config.num_max_nvl_chunked_recv_tokens, | ||
| hidden, | ||
| num_topk, | ||
| config.num_sms, | ||
| dtype='bfloat16') |
There was a problem hiding this comment.
Hardcoded dtype may cause type mismatch.
The kernel is always initialized with dtype='bfloat16' regardless of the actual dtype of input tensor x. This could cause incorrect behavior if x uses a different dtype (e.g., float16, float32).
🔎 Proposed fix
kernel = combine_kernel(
num_ranks,
config.num_max_nvl_chunked_send_tokens,
config.num_max_nvl_chunked_recv_tokens,
hidden,
num_topk,
config.num_sms,
- dtype='bfloat16')
+ dtype=str(x.dtype).split('.')[-1]) # e.g., 'torch.bfloat16' -> 'bfloat16'📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| kernel = combine_kernel( | |
| num_ranks, | |
| config.num_max_nvl_chunked_send_tokens, | |
| config.num_max_nvl_chunked_recv_tokens, | |
| hidden, | |
| num_topk, | |
| config.num_sms, | |
| dtype='bfloat16') | |
| kernel = combine_kernel( | |
| num_ranks, | |
| config.num_max_nvl_chunked_send_tokens, | |
| config.num_max_nvl_chunked_recv_tokens, | |
| hidden, | |
| num_topk, | |
| config.num_sms, | |
| dtype=str(x.dtype).split('.')[-1]) # e.g., 'torch.bfloat16' -> 'bfloat16' |
🤖 Prompt for AI Agents
In examples/distributed/deepseek_deepep/intranode/combine.py around lines 415 to
422, the kernel is created with a hardcoded dtype='bfloat16' which can mismatch
the input tensor x; change the call to pass the actual dtype of x (e.g., use
x.dtype or map x.dtype to the expected kernel string) so the kernel is
constructed with the same data type as the input, and ensure any necessary
conversion/formatting (like converting torch dtypes to kernel string names) is
applied before passing the dtype.
| stream=comm_stream.cuda_stream, | ||
| skip_tensor_validation=True) # reduce runtime overhead | ||
| compute_stream = torch.cuda.current_stream() | ||
| compute_stream.wait_stream(comm_stream) | ||
| return recv_x, recv_topk_weights |
There was a problem hiding this comment.
Missing null check for comm_stream.
Similar to cached_notify_combine, accessing comm_stream.cuda_stream without a null check will fail if comm_stream is None.
🤖 Prompt for AI Agents
In examples/distributed/deepseek_deepep/intranode/combine.py around lines 439 to
443, there is no null check before accessing comm_stream.cuda_stream which will
raise if comm_stream is None; update the code to branch on comm_stream (e.g. set
comm_cuda_stream = comm_stream.cuda_stream if comm_stream is not None else
torch.cuda.current_stream()), pass comm_cuda_stream (or None handled
appropriately) to the operator call, and only call
compute_stream.wait_stream(comm_cuda_stream) when comm_stream is not None so we
avoid dereferencing a null comm_stream.
| kernel = dispatch_kernel(num_ranks, config.num_max_nvl_chunked_send_tokens, | ||
| config.num_max_nvl_chunked_recv_tokens, hidden, num_topk, | ||
| num_experts, config.num_sms, 'bfloat16') |
There was a problem hiding this comment.
Hardcoded dtype may cause type mismatch.
Similar to combine.py, the kernel is initialized with hardcoded 'bfloat16' regardless of the actual dtype of input tensor x. This should be derived from x.dtype to support other precisions.
🔎 Proposed fix
+ dtype_str = str(x.dtype).split('.')[-1] # e.g., 'torch.bfloat16' -> 'bfloat16'
if handle is None:
kernel = dispatch_kernel(num_ranks, config.num_max_nvl_chunked_send_tokens,
config.num_max_nvl_chunked_recv_tokens, hidden, num_topk,
- num_experts, config.num_sms, 'bfloat16')
+ num_experts, config.num_sms, dtype_str)Apply similar change at line 950 for cached_dispatch_kernel.
Committable suggestion skipped: line range outside the PR's diff.
🤖 Prompt for AI Agents
In examples/distributed/deepseek_deepep/intranode/dispatch.py around lines
915-917, the dispatch kernel is created with a hardcoded dtype string 'bfloat16'
which can mismatch the actual input tensor dtype; change the call to derive the
dtype from the input tensor (e.g., use x.dtype or x.dtype.name converted to the
expected kernel string) and pass that value instead of the literal 'bfloat16';
also make the same change for the cached_dispatch_kernel instantiation around
line 950 so both kernels use the actual x.dtype.
| LdOp::LdOp(Array<PrimExpr> args, BufferMap vmap) { | ||
| ObjectPtr<LdOpNode> node = make_object<LdOpNode>(); | ||
| node->src = args[0]; | ||
| ICHECK(node->src.as<CallNode>()) << "src must be a call node"; | ||
| ICHECK(node->src.as<CallNode>()->op.same_as(builtin::address_of())) | ||
| << "src must be address_of op"; | ||
|
|
||
| node->value = args[1]; | ||
| node->sem = args[2].as<IntImm>().value()->value; | ||
| node->scope = args[3].as<IntImm>().value()->value; | ||
| node->na = args[4].as<IntImm>().value()->value; | ||
| node->nc = args[5].as<IntImm>().value()->value; | ||
| node->src_pe = args[6]; | ||
| data_ = std::move(node); | ||
| (void)vmap; | ||
| } |
There was a problem hiding this comment.
Add bounds checking for array indices.
The Lower method at lines 351-352 has the same issue as StOpNode::Lower: it indexes into sem_str and scope_str arrays using unvalidated sem and scope values, causing potential undefined behavior.
🔎 Proposed fix: Add bounds checking
LdOp::LdOp(Array<PrimExpr> args, BufferMap vmap) {
ObjectPtr<LdOpNode> node = make_object<LdOpNode>();
node->src = args[0];
ICHECK(node->src.as<CallNode>()) << "src must be a call node";
ICHECK(node->src.as<CallNode>()->op.same_as(builtin::address_of()))
<< "src must be address_of op";
node->value = args[1];
node->sem = args[2].as<IntImm>().value()->value;
node->scope = args[3].as<IntImm>().value()->value;
node->na = args[4].as<IntImm>().value()->value;
node->nc = args[5].as<IntImm>().value()->value;
node->src_pe = args[6];
+ ICHECK(node->sem >= 0 && node->sem < 5) << "sem must be in range [0, 4], got " << node->sem;
+ ICHECK(node->scope >= 0 && node->scope < 3) << "scope must be in range [0, 2], got " << node->scope;
data_ = std::move(node);
(void)vmap;
}Also applies to: 338-371
🤖 Prompt for AI Agents
In src/op/remote_copy.cc around lines 316-371, LdOp constructor and the Lower
method index into sem_str and scope_str using unvalidated sem and scope values
which can cause undefined behavior; add explicit bounds checks before any array
indexing (e.g., verify 0 <= sem < sem_str.size() and 0 <= scope <
scope_str.size()), and if out of range emit an ICHECK or clear error message and
handle the error path (return/throw) rather than performing the index; apply the
same validation in the corresponding code paths mentioned (lines ~338-371) so
all accesses to sem_str and scope_str are guarded.
| bool enable_aggressive_vectorize; ///< Whether to enable aggressive | ||
| ///< vectorization, only effctive for | ||
| ///< warp-scope |
There was a problem hiding this comment.
Fix documentation typo.
The comment has a typo: "effctive" should be "effective".
🔎 Proposed fix
- bool enable_aggressive_vectorize; ///< Whether to enable aggressive
- ///< vectorization, only effctive for
- ///< warp-scope
+ bool enable_aggressive_vectorize; ///< Whether to enable aggressive
+ ///< vectorization, only effective for
+ ///< warp-scope📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| bool enable_aggressive_vectorize; ///< Whether to enable aggressive | |
| ///< vectorization, only effctive for | |
| ///< warp-scope | |
| bool enable_aggressive_vectorize; ///< Whether to enable aggressive | |
| ///< vectorization, only effective for | |
| ///< warp-scope |
🤖 Prompt for AI Agents
In src/op/remote_copy.h around lines 36 to 38, the inline comment for
enable_aggressive_vectorize contains a typo ("effctive"); update the comment to
read "effective" so it becomes: "Whether to enable aggressive vectorization,
only effective for warp-scope". Keep punctuation and alignment consistent with
surrounding comment style.
| bool enable_aggressive_vectorize; ///< Whether to enable aggressive | ||
| ///< vectorization, only effctive for | ||
| ///< warp-scope |
There was a problem hiding this comment.
Fix documentation typo.
Same typo as in PutOpNode: "effctive" should be "effective".
🔎 Proposed fix
- bool enable_aggressive_vectorize; ///< Whether to enable aggressive
- ///< vectorization, only effctive for
- ///< warp-scope
+ bool enable_aggressive_vectorize; ///< Whether to enable aggressive
+ ///< vectorization, only effective for
+ ///< warp-scope📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| bool enable_aggressive_vectorize; ///< Whether to enable aggressive | |
| ///< vectorization, only effctive for | |
| ///< warp-scope | |
| bool enable_aggressive_vectorize; ///< Whether to enable aggressive | |
| ///< vectorization, only effective for | |
| ///< warp-scope |
🤖 Prompt for AI Agents
In src/op/remote_copy.h around lines 128 to 130, the documentation comment
contains a typo: change "effctive" to "effective" in the comment describing
enable_aggressive_vectorize so the sentence reads "only effective for
warp-scope".
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (2)
examples/distributed/deepseek_deepep/buffer.py (2)
104-119: Consider using consistent shape syntax.The tensor shape definitions mix tuple syntax (lines 90-102) and list syntax (lines 104-146). While both may work, using a consistent style improves readability.
🔎 Suggested fix for consistent tuple syntax
- channel_start_offset = tilelang.tensor([self.num_channels, self.num_ranks], + channel_start_offset = tilelang.tensor((self.num_channels, self.num_ranks), dtype=torch.int32, device='cuda', allocator=self._allocator) - channel_end_offset = tilelang.tensor([self.num_channels, self.num_ranks], + channel_end_offset = tilelang.tensor((self.num_channels, self.num_ranks), dtype=torch.int32, device='cuda', allocator=self._allocator) - channel_head_idx = tilelang.tensor([self.num_channels, self.num_ranks], + channel_head_idx = tilelang.tensor((self.num_channels, self.num_ranks), dtype=torch.int32, device='cuda', allocator=self._allocator) - channel_tail_idx = tilelang.tensor([self.num_channels, self.num_ranks], + channel_tail_idx = tilelang.tensor((self.num_channels, self.num_ranks), dtype=torch.int32, device='cuda', allocator=self._allocator)Apply similar changes to lines 121-146.
250-272: Prefix unused unpacked variables with underscores.Per static analysis, the unpacked variables
rank_prefix_matrix,channel_prefix_matrix,recv_channel_prefix_matrix,recv_src_idx, andsend_headare never used in this branch. Prefix them with underscores to indicate they are intentionally unused.🔎 Suggested fix
if handle is not None: assert topk_idx is None and topk_weights is None - rank_prefix_matrix, channel_prefix_matrix, recv_channel_prefix_matrix, recv_src_idx, is_token_in_rank, send_head = handle + _rank_prefix_matrix, _channel_prefix_matrix, _recv_channel_prefix_matrix, _recv_src_idx, is_token_in_rank, _send_head = handle
📜 Review details
Configuration used: defaults
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
examples/distributed/deepseek_deepep/buffer.py(1 hunks)examples/distributed/deepseek_deepep/intranode/test_intranode.py(1 hunks)
🧰 Additional context used
🧬 Code graph analysis (1)
examples/distributed/deepseek_deepep/intranode/test_intranode.py (1)
examples/distributed/deepseek_deepep/intranode/example_intranode.py (1)
test_intranode(19-204)
🪛 Ruff (0.14.8)
examples/distributed/deepseek_deepep/buffer.py
252-252: Unpacked variable rank_prefix_matrix is never used
Prefix it with an underscore or any other dummy variable pattern
(RUF059)
252-252: Unpacked variable channel_prefix_matrix is never used
Prefix it with an underscore or any other dummy variable pattern
(RUF059)
252-252: Unpacked variable recv_channel_prefix_matrix is never used
Prefix it with an underscore or any other dummy variable pattern
(RUF059)
252-252: Unpacked variable recv_src_idx is never used
Prefix it with an underscore or any other dummy variable pattern
(RUF059)
252-252: Unpacked variable send_head is never used
Prefix it with an underscore or any other dummy variable pattern
(RUF059)
🔇 Additional comments (8)
examples/distributed/deepseek_deepep/intranode/test_intranode.py (3)
1-4: LGTM!The imports are appropriate for this test module. The bare import of
example_intranodeworks correctly when tests are run from within theintranodedirectory or with proper pytest path configuration.
6-9: LGTM!The test function is properly decorated to require CUDA and uses monkeypatch correctly to configure command-line arguments before running the example. This provides good integration test coverage for the intranode dispatch/combine workflow.
12-13: LGTM!Standard test entry point that allows the file to be executed directly.
examples/distributed/deepseek_deepep/buffer.py (5)
1-13: LGTM!Module imports are well-organized with standard library, third-party, and local imports properly grouped.
30-79: LGTM!The
__init__method is well-structured with appropriate validation assertions and a clear initialization sequence. The class properly sets up distributed communication resources, including the allocator, symmetric buffers, and counters.
166-192: LGTM!The static method and properties are well-designed. The comment on line 184 clearly explains the SM allocation strategy (1 SM for send + 1 SM for recv per channel).
194-209: LGTM!Clean delegation to the intranode layout computation with comprehensive documentation.
283-304: LGTM!Clean delegation to the intranode combine implementation with comprehensive documentation. The TODO for bias support is noted.
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (1)
tilelang/distributed/install_deepep.sh (1)
22-28: Add error handling after DeepEP installation.The script doesn't check if the DeepEP installation succeeded before attempting validation, which could lead to confusing error messages.
🔎 Proposed fix to add error checking
# Install DeepEP cd 3rdparty/DeepEP -pip install . --no-build-isolation +if ! pip install . --no-build-isolation; then + echo "Failed to install DeepEP" + exit 1 +fi cd - # Validate -python -c "import deep_ep; print(deep_ep.__version__)" +if ! python -c "import deep_ep; print(deep_ep.__version__)"; then + echo "DeepEP installation validation failed" + exit 1 +fi echo "DeepEP is installed successfully. ✅"
📜 Review details
Configuration used: defaults
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (4)
.github/workflows/ci.yml(1 hunks).gitmodules(1 hunks)pyproject.toml(1 hunks)tilelang/distributed/install_deepep.sh(1 hunks)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (4)
- GitHub Check: Build wheels for Python 3.8 on ubuntu-24.04-arm with CUDA-12.8
- GitHub Check: Build wheels for Python 3.8 on macos-latest with Metal
- GitHub Check: Build wheels for Python 3.8 on ubuntu-latest with CUDA-12.1
- GitHub Check: format-check
🔇 Additional comments (1)
.gitmodules (1)
10-12: LGTM!The DeepEP submodule addition is correctly formatted and consistent with existing submodule entries.
| "torch>=2.7; platform_system == 'Darwin'", | ||
| "tqdm>=4.62.3", | ||
| "typing-extensions>=4.10.0", | ||
| "nvidia-nvshmem-cu12", |
There was a problem hiding this comment.
Add version constraint to nvidia-nvshmem-cu12 dependency.
The nvidia-nvshmem-cu12 dependency has no version constraint, which could lead to compatibility issues if breaking changes are introduced in future releases. Consider pinning to a stable version range such as "nvidia-nvshmem-cu12>=3.4" or "nvidia-nvshmem-cu12~=3.4.5" to ensure reproducible builds and controlled upgrades.
🤖 Prompt for AI Agents
In pyproject.toml around line 34, the dependency "nvidia-nvshmem-cu12" is listed
without any version constraint; update that entry to pin or constrain the
package (for example use a minimum version like "nvidia-nvshmem-cu12>=3.4" or a
compatible release spec like "nvidia-nvshmem-cu12~=3.4.5") to ensure
reproducible builds and avoid accidental breaking changes.
| export NVSHMEM_DIR=$(python -c "import site; print(site.getsitepackages()[0])")/nvidia/nvshmem | ||
| echo "NVSHMEM_DIR is set to $NVSHMEM_DIR" | ||
| ln -sf $NVSHMEM_DIR/lib/libnvshmem_host.so.3 $NVSHMEM_DIR/lib/libnvshmem_host.so |
There was a problem hiding this comment.
🛠️ Refactor suggestion | 🟠 Major
Improve robustness of NVSHMEM path setup and symlink creation.
The current implementation has potential issues:
- Line 17:
site.getsitepackages()[0]may not reliably return the correct site-packages directory in all environments (e.g., virtual environments, conda). - Line 19: The
-fflag inln -sfforces symlink creation even if the source file doesn't exist, masking potential errors.
🔎 Proposed fix for more robust path handling
-# Fix a bug of NVSHMEM path
-export NVSHMEM_DIR=$(python -c "import site; print(site.getsitepackages()[0])")/nvidia/nvshmem
-echo "NVSHMEM_DIR is set to $NVSHMEM_DIR"
-ln -sf $NVSHMEM_DIR/lib/libnvshmem_host.so.3 $NVSHMEM_DIR/lib/libnvshmem_host.so
+# Fix a bug of NVSHMEM path
+export NVSHMEM_DIR=$(python -c "import nvidia.nvshmem; import os; print(os.path.dirname(os.path.dirname(nvidia.nvshmem.__file__)))")
+echo "NVSHMEM_DIR is set to $NVSHMEM_DIR"
+
+# Verify source file exists before creating symlink
+if [ ! -f "$NVSHMEM_DIR/lib/libnvshmem_host.so.3" ]; then
+ echo "Error: libnvshmem_host.so.3 not found at $NVSHMEM_DIR/lib/"
+ exit 1
+fi
+ln -s $NVSHMEM_DIR/lib/libnvshmem_host.so.3 $NVSHMEM_DIR/lib/libnvshmem_host.soCommittable suggestion skipped: line range outside the PR's diff.
| cd 3rdparty/DeepEP | ||
| python setup.py install | ||
| cd - |
There was a problem hiding this comment.
🛠️ Refactor suggestion | 🟠 Major
Replace deprecated setup.py install with modern installation method.
The python setup.py install command is deprecated and discouraged by the Python packaging community. Modern best practice is to use pip install.
🔎 Proposed fix using pip
# Install DeepEP
cd 3rdparty/DeepEP
-python setup.py install
+pip install . --no-build-isolation
cd -📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| cd 3rdparty/DeepEP | |
| python setup.py install | |
| cd - | |
| # Install DeepEP | |
| cd 3rdparty/DeepEP | |
| pip install . --no-build-isolation | |
| cd - |
🤖 Prompt for AI Agents
In tilelang/distributed/install_deepep.sh around lines 22 to 24, replace the
deprecated python setup.py install with a pip-based install; change the command
to use the current interpreter (e.g., python -m pip install .) while still
cd'ing into 3rdparty/DeepEP and returning to the previous directory, so the
script uses modern packaging tools and installs via pip.
There was a problem hiding this comment.
Actionable comments posted: 0
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
.github/workflows/ci.yml (1)
86-86: Inconsistent hash computation breaks cache reliability.The
build-test-nvidiajob usescat requirements-test.txtwhileformat-check(line 29) usessha256sum requirements-test.txt | awk '{print $1}'. This causes:
- REQS_HASH in build-test-nvidia contains the entire file content (multi-line), not a hash
${REQS_HASH:0:8}extracts the first 8 characters of the file content, not a hash- Cache keys become unreliable and inconsistent between jobs
- Changes to later parts of requirements-test.txt won't trigger venv rebuild
🔎 Proposed fix to align hash computation with format-check
- REQS_HASH=$(cat requirements-test.txt 2>/dev/null || true) + REQS_HASH=$(sha256sum requirements-test.txt 2>/dev/null | awk '{print $1}' || echo "no_requirements")
♻️ Duplicate comments (1)
.github/workflows/ci.yml (1)
83-102: DeepEP installation still missing in build-test-nvidia job.This is the same issue previously flagged: the
build-test-nvidiajob installs requirements and flash_attn but does not install DeepEP, while theformat-checkjob does (line 44). The distributed tests at lines 116-123 and 142-149 will likely fail without DeepEP.🔎 Proposed fix to add DeepEP installation
[[ -f requirements-test.txt ]] && \ PIP_NO_BUILD_ISOLATION=1 pip install -r requirements-test.txt --no-user # flash attention usually requires no isolation build pip install flash_attn==2.5.8 --no-user --no-build-isolation + bash tilelang/distributed/install_deepep.sh # Install DeepEP for testing purpose touch "$MARKER" fi
🧹 Nitpick comments (1)
.github/workflows/ci.yml (1)
10-10: Optional: Suppress actionlint warnings for custom runner labels.The static analysis tool flags
nvidiaandhopperas unknown runner labels. These are valid custom labels for self-hosted runners, but you can optionally suppress the warnings by creating anactionlint.yamlconfig file at the repository root.📝 Example actionlint.yaml configuration
Create
.github/actionlint.yaml:self-hosted-runner: labels: - nvidia - hopper
📜 Review details
Configuration used: defaults
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
.github/workflows/ci.yml(2 hunks)
🧰 Additional context used
🪛 actionlint (1.7.9)
.github/workflows/ci.yml
10-10: label "nvidia" is unknown. available labels are "windows-latest", "windows-latest-8-cores", "windows-2025", "windows-2022", "windows-11-arm", "ubuntu-slim", "ubuntu-latest", "ubuntu-latest-4-cores", "ubuntu-latest-8-cores", "ubuntu-latest-16-cores", "ubuntu-24.04", "ubuntu-24.04-arm", "ubuntu-22.04", "ubuntu-22.04-arm", "macos-latest", "macos-latest-xl", "macos-latest-xlarge", "macos-latest-large", "macos-26-xlarge", "macos-26", "macos-15-intel", "macos-15-xlarge", "macos-15-large", "macos-15", "macos-14-xl", "macos-14-xlarge", "macos-14-large", "macos-14", "macos-13-xl", "macos-13-xlarge", "macos-13-large", "macos-13", "self-hosted", "x64", "arm", "arm64", "linux", "macos", "windows". if it is a custom label for self-hosted runner, set list of labels in actionlint.yaml config file
(runner-label)
10-10: label "hopper" is unknown. available labels are "windows-latest", "windows-latest-8-cores", "windows-2025", "windows-2022", "windows-11-arm", "ubuntu-slim", "ubuntu-latest", "ubuntu-latest-4-cores", "ubuntu-latest-8-cores", "ubuntu-latest-16-cores", "ubuntu-24.04", "ubuntu-24.04-arm", "ubuntu-22.04", "ubuntu-22.04-arm", "macos-latest", "macos-latest-xl", "macos-latest-xlarge", "macos-latest-large", "macos-26-xlarge", "macos-26", "macos-15-intel", "macos-15-xlarge", "macos-15-large", "macos-15", "macos-14-xl", "macos-14-xlarge", "macos-14-large", "macos-14", "macos-13-xl", "macos-13-xlarge", "macos-13-large", "macos-13", "self-hosted", "x64", "arm", "arm64", "linux", "macos", "windows". if it is a custom label for self-hosted runner, set list of labels in actionlint.yaml config file
(runner-label)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (4)
- GitHub Check: Build wheels for Python 3.8 on ubuntu-latest with CUDA-12.1
- GitHub Check: Build wheels for Python 3.8 on ubuntu-24.04-arm with CUDA-12.8
- GitHub Check: Build wheels for Python 3.8 on macos-latest with Metal
- GitHub Check: format-check
84ae34e to
ed9c6eb
Compare
There was a problem hiding this comment.
Actionable comments posted: 0
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
.github/workflows/ci.yml (1)
82-102: Verify whether the venv caching strategy should account for DeepEP submodule updates.The venv cache is invalidated only when
requirements-test.txtchanges, but DeepEP is tracked as a git submodule (in.gitmodules). When the DeepEP submodule is updated to a new commit,install_deepep.shwill not run on subsequent CI runs if the venv cache exists, potentially leaving stale DeepEP installations. Consider invalidating the cache when the DeepEP submodule commit changes—for example, by including the DeepEP submodule commit hash in the cache key, or running the DeepEP installation outside the venv recreation conditional.
🧹 Nitpick comments (1)
.github/workflows/ci.yml (1)
10-10: Consider adding custom runner labels to actionlint.yaml to suppress warnings.The custom self-hosted runner labels "nvidia" and "hopper" are valid but trigger static analysis warnings. To suppress these false positives, create an
actionlint.yamlconfiguration file in the repository root with the custom labels defined.📝 Example actionlint.yaml configuration
Create
.github/actionlint.yaml:self-hosted-runner: labels: - nvidia - hopper
📜 Review details
Configuration used: defaults
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
.github/workflows/ci.yml(2 hunks)
🧰 Additional context used
🪛 actionlint (1.7.9)
.github/workflows/ci.yml
10-10: label "nvidia" is unknown. available labels are "windows-latest", "windows-latest-8-cores", "windows-2025", "windows-2022", "windows-11-arm", "ubuntu-slim", "ubuntu-latest", "ubuntu-latest-4-cores", "ubuntu-latest-8-cores", "ubuntu-latest-16-cores", "ubuntu-24.04", "ubuntu-24.04-arm", "ubuntu-22.04", "ubuntu-22.04-arm", "macos-latest", "macos-latest-xl", "macos-latest-xlarge", "macos-latest-large", "macos-26-xlarge", "macos-26", "macos-15-intel", "macos-15-xlarge", "macos-15-large", "macos-15", "macos-14-xl", "macos-14-xlarge", "macos-14-large", "macos-14", "macos-13-xl", "macos-13-xlarge", "macos-13-large", "macos-13", "self-hosted", "x64", "arm", "arm64", "linux", "macos", "windows". if it is a custom label for self-hosted runner, set list of labels in actionlint.yaml config file
(runner-label)
10-10: label "hopper" is unknown. available labels are "windows-latest", "windows-latest-8-cores", "windows-2025", "windows-2022", "windows-11-arm", "ubuntu-slim", "ubuntu-latest", "ubuntu-latest-4-cores", "ubuntu-latest-8-cores", "ubuntu-latest-16-cores", "ubuntu-24.04", "ubuntu-24.04-arm", "ubuntu-22.04", "ubuntu-22.04-arm", "macos-latest", "macos-latest-xl", "macos-latest-xlarge", "macos-latest-large", "macos-26-xlarge", "macos-26", "macos-15-intel", "macos-15-xlarge", "macos-15-large", "macos-15", "macos-14-xl", "macos-14-xlarge", "macos-14-large", "macos-14", "macos-13-xl", "macos-13-xlarge", "macos-13-large", "macos-13", "self-hosted", "x64", "arm", "arm64", "linux", "macos", "windows". if it is a custom label for self-hosted runner, set list of labels in actionlint.yaml config file
(runner-label)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (4)
- GitHub Check: build-test-nvidia
- GitHub Check: Build wheels for Python 3.8 on macos-latest with Metal
- GitHub Check: Build wheels for Python 3.8 on ubuntu-24.04-arm with CUDA-12.8
- GitHub Check: Build wheels for Python 3.8 on ubuntu-latest with CUDA-12.1
ed9c6eb to
dbcee5e
Compare
There was a problem hiding this comment.
Actionable comments posted: 0
♻️ Duplicate comments (1)
.github/workflows/ci.yml (1)
97-97: DeepEP installation is now included.This addresses the previous review feedback. The DeepEP installation script is correctly placed after dependency installation and before project installation.
🧹 Nitpick comments (1)
.github/workflows/ci.yml (1)
85-97: Remove dead code:REQS_HASHandMARKERare computed but never used.Lines 85-86 compute
REQS_HASHandMARKERvariables, but the marker-based venv reuse logic has been removed. These are now dead code and should be cleaned up for clarity.🔎 Proposed cleanup
- name: Ensure venv (local & persistent) run: | set -e - REQS_HASH=$(cat requirements-test.txt 2>/dev/null || true) - MARKER="${{ runner.tool_cache }}/.venv_marker_${{ env.PYTHON_VERSION }}_${REQS_HASH:0:8}" # NOTE(wt): We disable the venv reuse for now to allow installing DeepEP - # echo "venv stale or missing – recreating" rm -rf "${{ runner.tool_cache }}/${{ env.VENV_DIR }}" python -m venv "${{ runner.tool_cache }}/${{ env.VENV_DIR }}" source "${{ runner.tool_cache }}/${{ env.VENV_DIR }}/bin/activate"Additionally, always recreating the venv will slow down CI builds. Once DeepEP can be installed incrementally (or cached), consider restoring the marker-based caching for faster builds.
📜 Review details
Configuration used: defaults
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
.github/workflows/ci.yml(2 hunks)
🧰 Additional context used
🪛 actionlint (1.7.9)
.github/workflows/ci.yml
10-10: label "nvidia" is unknown. available labels are "windows-latest", "windows-latest-8-cores", "windows-2025", "windows-2022", "windows-11-arm", "ubuntu-slim", "ubuntu-latest", "ubuntu-latest-4-cores", "ubuntu-latest-8-cores", "ubuntu-latest-16-cores", "ubuntu-24.04", "ubuntu-24.04-arm", "ubuntu-22.04", "ubuntu-22.04-arm", "macos-latest", "macos-latest-xl", "macos-latest-xlarge", "macos-latest-large", "macos-26-xlarge", "macos-26", "macos-15-intel", "macos-15-xlarge", "macos-15-large", "macos-15", "macos-14-xl", "macos-14-xlarge", "macos-14-large", "macos-14", "macos-13-xl", "macos-13-xlarge", "macos-13-large", "macos-13", "self-hosted", "x64", "arm", "arm64", "linux", "macos", "windows". if it is a custom label for self-hosted runner, set list of labels in actionlint.yaml config file
(runner-label)
10-10: label "hopper" is unknown. available labels are "windows-latest", "windows-latest-8-cores", "windows-2025", "windows-2022", "windows-11-arm", "ubuntu-slim", "ubuntu-latest", "ubuntu-latest-4-cores", "ubuntu-latest-8-cores", "ubuntu-latest-16-cores", "ubuntu-24.04", "ubuntu-24.04-arm", "ubuntu-22.04", "ubuntu-22.04-arm", "macos-latest", "macos-latest-xl", "macos-latest-xlarge", "macos-latest-large", "macos-26-xlarge", "macos-26", "macos-15-intel", "macos-15-xlarge", "macos-15-large", "macos-15", "macos-14-xl", "macos-14-xlarge", "macos-14-large", "macos-14", "macos-13-xl", "macos-13-xlarge", "macos-13-large", "macos-13", "self-hosted", "x64", "arm", "arm64", "linux", "macos", "windows". if it is a custom label for self-hosted runner, set list of labels in actionlint.yaml config file
(runner-label)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (4)
- GitHub Check: Build wheels for Python 3.8 on ubuntu-24.04-arm with CUDA-12.8
- GitHub Check: Build wheels for Python 3.8 on ubuntu-latest with CUDA-12.1
- GitHub Check: Build wheels for Python 3.8 on macos-latest with Metal
- GitHub Check: format-check
🔇 Additional comments (1)
.github/workflows/ci.yml (1)
10-10: Custom self-hosted runner labels are correctly configured.The static analysis tool flags
nvidiaandhopperas unknown labels, but these are valid custom labels for self-hosted runners with NVIDIA Hopper GPUs (H100s per the PR benchmarks). This is standard practice for self-hosted runners. Consider adding anactionlint.yamlconfig to silence these warnings if they become noisy.
dbcee5e to
db6ae71
Compare
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (1)
.github/workflows/ci.yml (1)
85-86: Consider removing unused MARKER variable computation.The
MARKERvariable is computed but never used since the venv reuse logic was disabled. This is dead code that can be removed for clarity.🔎 Proposed cleanup
- REQS_HASH=$(cat requirements-test.txt 2>/dev/null || true) - MARKER="${{ runner.tool_cache }}/.venv_marker_${{ env.PYTHON_VERSION }}_${REQS_HASH:0:8}" # NOTE(wt): We disable the venv reuse for now to allow installing DeepEP
📜 Review details
Configuration used: defaults
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
.github/workflows/ci.yml3rdparty/DeepEP
✅ Files skipped from review due to trivial changes (1)
- 3rdparty/DeepEP
🧰 Additional context used
🪛 actionlint (1.7.9)
.github/workflows/ci.yml
10-10: label "nvidia" is unknown. available labels are "windows-latest", "windows-latest-8-cores", "windows-2025", "windows-2022", "windows-11-arm", "ubuntu-slim", "ubuntu-latest", "ubuntu-latest-4-cores", "ubuntu-latest-8-cores", "ubuntu-latest-16-cores", "ubuntu-24.04", "ubuntu-24.04-arm", "ubuntu-22.04", "ubuntu-22.04-arm", "macos-latest", "macos-latest-xl", "macos-latest-xlarge", "macos-latest-large", "macos-26-xlarge", "macos-26", "macos-15-intel", "macos-15-xlarge", "macos-15-large", "macos-15", "macos-14-xl", "macos-14-xlarge", "macos-14-large", "macos-14", "macos-13-xl", "macos-13-xlarge", "macos-13-large", "macos-13", "self-hosted", "x64", "arm", "arm64", "linux", "macos", "windows". if it is a custom label for self-hosted runner, set list of labels in actionlint.yaml config file
(runner-label)
10-10: label "hopper" is unknown. available labels are "windows-latest", "windows-latest-8-cores", "windows-2025", "windows-2022", "windows-11-arm", "ubuntu-slim", "ubuntu-latest", "ubuntu-latest-4-cores", "ubuntu-latest-8-cores", "ubuntu-latest-16-cores", "ubuntu-24.04", "ubuntu-24.04-arm", "ubuntu-22.04", "ubuntu-22.04-arm", "macos-latest", "macos-latest-xl", "macos-latest-xlarge", "macos-latest-large", "macos-26-xlarge", "macos-26", "macos-15-intel", "macos-15-xlarge", "macos-15-large", "macos-15", "macos-14-xl", "macos-14-xlarge", "macos-14-large", "macos-14", "macos-13-xl", "macos-13-xlarge", "macos-13-large", "macos-13", "self-hosted", "x64", "arm", "arm64", "linux", "macos", "windows". if it is a custom label for self-hosted runner, set list of labels in actionlint.yaml config file
(runner-label)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (4)
- GitHub Check: Build wheels for Python 3.8 on ubuntu-latest with CUDA-12.1
- GitHub Check: Build wheels for Python 3.8 on ubuntu-24.04-arm with CUDA-12.8
- GitHub Check: Build wheels for Python 3.8 on macos-latest with Metal
- GitHub Check: format-check
🔇 Additional comments (3)
.github/workflows/ci.yml (3)
10-10: LGTM: Hopper label added for format-check job.The addition of the "hopper" label is appropriate for self-hosted runners targeting Hopper architecture. The static analysis warnings about unknown labels are expected for custom runner labels.
87-97: Unconditional venv recreation correctly enables DeepEP installation.The change from conditional reuse to unconditional recreation ensures a clean environment for DeepEP installation. This addresses the previous review concern about missing DeepEP in the build-test-nvidia job.
26-45: AI summary claims DeepEP installed in both jobs, but format-check lacks it.The AI summary states that DeepEP installation is added in "both workflows," but only the
build-test-nvidiajob includes the DeepEP installation script (line 97). Theformat-checkjob does not have this installation step, which is likely correct since format checking doesn't require DeepEP functionality.
eef8b0d to
0a04886
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
📜 Review details
Configuration used: defaults
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
.github/workflows/ci.yml3rdparty/DeepEP
🚧 Files skipped from review as they are similar to previous changes (1)
- 3rdparty/DeepEP
🧰 Additional context used
🪛 actionlint (1.7.9)
.github/workflows/ci.yml
10-10: label "nvidia" is unknown. available labels are "windows-latest", "windows-latest-8-cores", "windows-2025", "windows-2022", "windows-11-arm", "ubuntu-slim", "ubuntu-latest", "ubuntu-latest-4-cores", "ubuntu-latest-8-cores", "ubuntu-latest-16-cores", "ubuntu-24.04", "ubuntu-24.04-arm", "ubuntu-22.04", "ubuntu-22.04-arm", "macos-latest", "macos-latest-xl", "macos-latest-xlarge", "macos-latest-large", "macos-26-xlarge", "macos-26", "macos-15-intel", "macos-15-xlarge", "macos-15-large", "macos-15", "macos-14-xl", "macos-14-xlarge", "macos-14-large", "macos-14", "macos-13-xl", "macos-13-xlarge", "macos-13-large", "macos-13", "self-hosted", "x64", "arm", "arm64", "linux", "macos", "windows". if it is a custom label for self-hosted runner, set list of labels in actionlint.yaml config file
(runner-label)
10-10: label "hopper" is unknown. available labels are "windows-latest", "windows-latest-8-cores", "windows-2025", "windows-2022", "windows-11-arm", "ubuntu-slim", "ubuntu-latest", "ubuntu-latest-4-cores", "ubuntu-latest-8-cores", "ubuntu-latest-16-cores", "ubuntu-24.04", "ubuntu-24.04-arm", "ubuntu-22.04", "ubuntu-22.04-arm", "macos-latest", "macos-latest-xl", "macos-latest-xlarge", "macos-latest-large", "macos-26-xlarge", "macos-26", "macos-15-intel", "macos-15-xlarge", "macos-15-large", "macos-15", "macos-14-xl", "macos-14-xlarge", "macos-14-large", "macos-14", "macos-13-xl", "macos-13-xlarge", "macos-13-large", "macos-13", "self-hosted", "x64", "arm", "arm64", "linux", "macos", "windows". if it is a custom label for self-hosted runner, set list of labels in actionlint.yaml config file
(runner-label)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (3)
- GitHub Check: Build wheels for Python 3.8 on ubuntu-24.04-arm with CUDA-12.8
- GitHub Check: Build wheels for Python 3.8 on macos-latest with Metal
- GitHub Check: Build wheels for Python 3.8 on ubuntu-latest with CUDA-12.1
🔇 Additional comments (2)
.github/workflows/ci.yml (2)
10-10: LGTM: Consistent runner configuration.Adding "hopper" aligns the format-check job with build-test-nvidia's runner configuration. The actionlint warnings about unknown labels are false positives—these are valid custom labels for self-hosted runners.
102-102: LGTM: DeepEP installation added as required.This addresses the past review comment. Installing DeepEP after the wheel installation is the correct placement, ensuring distributed tests have the necessary dependencies.
0a04886 to
70c32e2
Compare
There was a problem hiding this comment.
Actionable comments posted: 0
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
.github/workflows/ci.yml (1)
82-96: Remove dead code and restore venv reuse consistency.Several issues with the current venv handling:
- Dead code: Lines 85-86 compute
REQS_HASHandMARKERvariables that are never used. Remove these lines.- Inconsistency: The
format-checkjob (lines 26-45) uses conditional venv reuse with hash-based markers, whilebuild-test-nvidiaunconditionally recreates the venv every run. This inconsistency makes the workflow harder to maintain.- Performance: Unconditionally recreating the venv on every run significantly slows CI execution on self-hosted runners where venv caching is beneficial.
- Unclear reasoning: The comment states venv reuse is disabled "to allow installing DeepEP", but DeepEP is installed after venv creation at line 102. If DeepEP truly requires a fresh environment each time, please document the specific reason why (e.g., state pollution, version conflicts).
Consider restoring the venv reuse logic to match
format-check, or if there's a genuine technical reason for forcing recreation, document it clearly and consider including DeepEP's version in the hash computation instead.🔎 Proposed fix to remove dead code and restore venv reuse
- name: Ensure venv (local & persistent) run: | set -e - REQS_HASH=$(cat requirements-test.txt 2>/dev/null || true) - MARKER="${{ runner.tool_cache }}/.venv_marker_${{ env.PYTHON_VERSION }}_${REQS_HASH:0:8}" - # NOTE(wt): We disable the venv reuse for now to allow installing DeepEP - # echo "venv stale or missing – recreating" - rm -rf "${{ runner.tool_cache }}/${{ env.VENV_DIR }}" - python -m venv "${{ runner.tool_cache }}/${{ env.VENV_DIR }}" - source "${{ runner.tool_cache }}/${{ env.VENV_DIR }}/bin/activate" - python -m pip install --upgrade pip --no-user - [[ -f requirements-test.txt ]] && \ - PIP_NO_BUILD_ISOLATION=1 pip install -r requirements-test.txt --no-user - # flash attention usually requires no isolation build - pip install flash_attn==2.5.8 --no-user --no-build-isolation + REQS_HASH=$(sha256sum requirements-test.txt 2>/dev/null | awk '{print $1}' || echo "no_requirements") + MARKER="${{ runner.tool_cache }}/.venv_marker_${{ env.PYTHON_VERSION }}_${REQS_HASH:0:8}" + + if [[ -f "$MARKER" ]] && [[ -f "${{ runner.tool_cache }}/${{ env.VENV_DIR }}/bin/activate" ]]; then + echo "venv exists and hash matches – reuse it" + else + echo "venv stale or missing – recreating" + rm -rf "${{ runner.tool_cache }}/${{ env.VENV_DIR }}" "$MARKER" + python -m venv "${{ runner.tool_cache }}/${{ env.VENV_DIR }}" + # shellcheck source=/dev/null + source "${{ runner.tool_cache }}/${{ env.VENV_DIR }}/bin/activate" + python -m pip install --upgrade pip --no-user + [[ -f requirements-test.txt ]] && \ + PIP_NO_BUILD_ISOLATION=1 pip install -r requirements-test.txt --no-user + pip install flash_attn==2.5.8 --no-user --no-build-isolation + touch "$MARKER" + fiThis matches the approach used in
format-checkand eliminates the dead code.
📜 Review details
Configuration used: defaults
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
.github/workflows/ci.yml3rdparty/DeepEPtilelang/distributed/install_deepep.sh
✅ Files skipped from review due to trivial changes (1)
- 3rdparty/DeepEP
🚧 Files skipped from review as they are similar to previous changes (1)
- tilelang/distributed/install_deepep.sh
🧰 Additional context used
🪛 actionlint (1.7.9)
.github/workflows/ci.yml
10-10: label "nvidia" is unknown. available labels are "windows-latest", "windows-latest-8-cores", "windows-2025", "windows-2022", "windows-11-arm", "ubuntu-slim", "ubuntu-latest", "ubuntu-latest-4-cores", "ubuntu-latest-8-cores", "ubuntu-latest-16-cores", "ubuntu-24.04", "ubuntu-24.04-arm", "ubuntu-22.04", "ubuntu-22.04-arm", "macos-latest", "macos-latest-xl", "macos-latest-xlarge", "macos-latest-large", "macos-26-xlarge", "macos-26", "macos-15-intel", "macos-15-xlarge", "macos-15-large", "macos-15", "macos-14-xl", "macos-14-xlarge", "macos-14-large", "macos-14", "macos-13-xl", "macos-13-xlarge", "macos-13-large", "macos-13", "self-hosted", "x64", "arm", "arm64", "linux", "macos", "windows". if it is a custom label for self-hosted runner, set list of labels in actionlint.yaml config file
(runner-label)
10-10: label "hopper" is unknown. available labels are "windows-latest", "windows-latest-8-cores", "windows-2025", "windows-2022", "windows-11-arm", "ubuntu-slim", "ubuntu-latest", "ubuntu-latest-4-cores", "ubuntu-latest-8-cores", "ubuntu-latest-16-cores", "ubuntu-24.04", "ubuntu-24.04-arm", "ubuntu-22.04", "ubuntu-22.04-arm", "macos-latest", "macos-latest-xl", "macos-latest-xlarge", "macos-latest-large", "macos-26-xlarge", "macos-26", "macos-15-intel", "macos-15-xlarge", "macos-15-large", "macos-15", "macos-14-xl", "macos-14-xlarge", "macos-14-large", "macos-14", "macos-13-xl", "macos-13-xlarge", "macos-13-large", "macos-13", "self-hosted", "x64", "arm", "arm64", "linux", "macos", "windows". if it is a custom label for self-hosted runner, set list of labels in actionlint.yaml config file
(runner-label)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (4)
- GitHub Check: build-test-nvidia
- GitHub Check: Build wheels for Python 3.8 on macos-latest with Metal
- GitHub Check: Build wheels for Python 3.8 on ubuntu-latest with CUDA-12.1
- GitHub Check: Build wheels for Python 3.8 on ubuntu-24.04-arm with CUDA-12.8
…nd streamline parser.py for better readability
70c32e2 to
c52e4e8
Compare
e16fd2f to
5bbd6dd
Compare
There was a problem hiding this comment.
Actionable comments posted: 7
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
examples/distributed/primitives/example_put_warp.py (1)
71-71: Adjust tolerance thresholds for bfloat16 precision.With the dtype change to bfloat16, the current tolerances (
atol=1e-6, rtol=1e-6) are unrealistically tight. bfloat16 has approximately 3 decimal digits of precision, so tolerances around1e-3are more appropriate to avoid false failures.🔎 Proposed fix
- if torch.allclose(dst_torch, dst, atol=1e-6, rtol=1e-6): + if torch.allclose(dst_torch, dst, atol=1e-3, rtol=1e-3):
♻️ Duplicate comments (22)
testing/python/language/test_tilelang_language_vote.py (1)
44-44: Fix incorrect assertion logic.This issue was previously identified: the assertion doesn't check both conditions due to operator precedence. It evaluates as
assert ('__any_sync') and ('__all_sync' in ...), where the non-empty string'__any_sync'is always truthy, so only'__all_sync'presence is verified.🔎 Proposed fix
- assert '__any_sync' and '__all_sync' in kernel.get_kernel_source() + kernel_src = kernel.get_kernel_source() + assert '__any_sync' in kernel_src and '__all_sync' in kernel_srctesting/python/language/test_tilelang_language_ldst_options.py (1)
66-66: Add dtype specification to prevent type mismatch.The assertion compares an
int32tensor (x) withtorch.arange(32, device="cuda"), which defaults toint64. This type mismatch may cause the assertion to fail or produce incorrect comparisons.🔎 Proposed fix
- assert x.equal(torch.arange(32, device="cuda")) + assert torch.equal(x, torch.arange(32, device="cuda", dtype=torch.int32))
Note: The inconsistent assertion style (instance method vs. functional) was already flagged in a previous review comment.
examples/distributed/primitives/example_put_warp.py (1)
58-59: Critical:shape=(M)evaluates to integer, not a 1D tuple.This issue was already flagged in previous reviews. The
shape=(M)syntax evaluates to the integerMrather than a tuple. Useshape=(M,)with a trailing comma.src/transform/vectorize_loop.cc (1)
796-806: Bug persists:vec_value_lvaluediscards the reinterpret operation.This issue was already flagged in a previous review. The code creates
vec_valueas a reinterpret call but then extractsargs[0]which isnew_value(the original un-reinterpreted value). Compare with the correct implementation inloop_vectorization_utils.h(line 491) which usesvec_valuedirectly.Proposed fix
// Reinterpret the value to vector type (e.g., int4 for 8xbf16) PrimExpr vec_value = Call(vec_dtype, builtin::reinterpret(), {new_value}); - // A trick to get the lvalue of the vectorized value - PrimExpr vec_value_lvalue = vec_value.as<CallNode>()->args[0]; - // Build new args with base addresses and reinterpreted value Array<PrimExpr> new_args; new_args.push_back(func_name); new_args.push_back(new_addr); - new_args.push_back(vec_value_lvalue); + new_args.push_back(vec_value); // Copy remaining args (sem, scope, etc.)src/tl_templates/cuda/copy.h (3)
258-270: Critical: Missing alignment and size validation (same as above).This overload also has the same critical validation issues:
- No alignment verification for 16-byte boundaries
- No size validation for
N * sizeof(dtype_t) % 16 == 0🔎 Proposed fix with validation
template <int N, int UNROLL_FACTOR, bool enable_aggressive_vectorize = false, typename dtype_t> TL_DEVICE void cp_warp(dtype_t *const dst_addr, uint64_t src_addr_uint64) { const dtype_t *src_addr = reinterpret_cast<const dtype_t *>(src_addr_uint64); if constexpr (enable_aggressive_vectorize) { + static_assert((N * sizeof(dtype_t)) % 16 == 0, + "Aggressive vectorization requires N*sizeof(dtype_t) to be a multiple of 16"); + if ((reinterpret_cast<uintptr_t>(dst_addr) % 16 != 0) || + (src_addr_uint64 % 16 != 0)) { + cp_warp_impl<N, UNROLL_FACTOR>(dst_addr, src_addr); + return; + } int4 *__restrict__ dst_addr_int4 = (int4 *)dst_addr; const int4 *__restrict__ src_addr_int4 = (const int4 *)src_addr; constexpr int N_int4 = sizeof(dtype_t) * N / 16; cp_warp_impl<N_int4, UNROLL_FACTOR>(dst_addr_int4, src_addr_int4); } else { cp_warp_impl<N, UNROLL_FACTOR>(dst_addr, src_addr); } }
221-241: Critical: Missing alignment and size validation for aggressive vectorization.While documentation now describes the requirements, the code still lacks validation. The
enable_aggressive_vectorizepath casts pointers toint4*without verifying:
- Both source and destination are 16-byte aligned
N * sizeof(dtype_t)is a multiple of 16This can cause:
- Misaligned memory access crashes if pointers aren't 16-byte aligned
- Silent data truncation if the size isn't a multiple of 16 (integer division
N_int4 = sizeof(dtype_t) * N / 16truncates)Compare with
nvshmem_cp_threadgroup(lines 287-302), which performs runtime alignment checks before vectorized copies.🔎 Proposed fix with validation
template <int N, int UNROLL_FACTOR, bool enable_aggressive_vectorize = false, typename dtype_t> TL_DEVICE void cp_warp(dtype_t const *const dst_addr, dtype_t const *const src_addr) { if constexpr (enable_aggressive_vectorize) { + static_assert((N * sizeof(dtype_t)) % 16 == 0, + "Aggressive vectorization requires N*sizeof(dtype_t) to be a multiple of 16"); + // Runtime alignment check (caller must ensure 16-byte alignment) + if ((reinterpret_cast<uintptr_t>(dst_addr) % 16 != 0) || + (reinterpret_cast<uintptr_t>(src_addr) % 16 != 0)) { + // Fallback to non-vectorized path + cp_warp_impl<N, UNROLL_FACTOR>(dst_addr, src_addr); + return; + } int4 *__restrict__ dst_addr_int4 = (int4 *)dst_addr; const int4 *__restrict__ src_addr_int4 = (const int4 *)src_addr; constexpr int N_int4 = sizeof(dtype_t) * N / 16; cp_warp_impl<N_int4, UNROLL_FACTOR>(dst_addr_int4, src_addr_int4); } else { cp_warp_impl<N, UNROLL_FACTOR>(dst_addr, src_addr); } }
243-256: Critical: Missing alignment and size validation (same as above).This overload has the same critical validation issues as lines 221-241:
- No alignment verification for 16-byte boundaries
- No size validation for
N * sizeof(dtype_t) % 16 == 0🔎 Proposed fix with validation
template <int N, int UNROLL_FACTOR, bool enable_aggressive_vectorize = false, typename dtype_t> TL_DEVICE void cp_warp(uint64_t dst_addr_uint64, dtype_t const *const src_addr) { dtype_t *dst_addr = reinterpret_cast<dtype_t *>(dst_addr_uint64); if constexpr (enable_aggressive_vectorize) { + static_assert((N * sizeof(dtype_t)) % 16 == 0, + "Aggressive vectorization requires N*sizeof(dtype_t) to be a multiple of 16"); + if ((dst_addr_uint64 % 16 != 0) || + (reinterpret_cast<uintptr_t>(src_addr) % 16 != 0)) { + cp_warp_impl<N, UNROLL_FACTOR>(dst_addr, src_addr); + return; + } int4 *__restrict__ dst_addr_int4 = (int4 *)dst_addr; const int4 *__restrict__ src_addr_int4 = (const int4 *)src_addr; constexpr int N_int4 = sizeof(dtype_t) * N / 16; cp_warp_impl<N_int4, UNROLL_FACTOR>(dst_addr_int4, src_addr_int4); } else { cp_warp_impl<N, UNROLL_FACTOR>(dst_addr, src_addr); } }examples/distributed/deepseek_deepep/buffer.py (1)
121-127: Hardcodedbfloat16dtype may cause type mismatch.The
channel_x_buffersallocation usesdtype=torch.bfloat16, but the dispatch and combine operations may receive tensors of different dtypes. Consider parameterizing the dtype via an__init__argument to match the expected input type.src/tl_templates/cuda/sync.h (1)
153-154: Fix typo: "optinal" should be "optional".🔎 Proposed fix
-// Sync blocks at a system-level barrier with an optinal fence +// Sync blocks at a system-level barrier with an optional fencesrc/op/remote_copy.h (1)
36-38: Fix typo: "effctive" should be "effective".The inline comments for
enable_aggressive_vectorizecontain the typo "effctive" (appears in bothPutOpNodeandGetOpNode).🔎 Proposed fix
bool enable_aggressive_vectorize; ///< Whether to enable aggressive - ///< vectorization, only effctive for + ///< vectorization, only effective for ///< warp-scopeAlso applies to: 128-130
src/op/sync.h (1)
134-134: Missingneed_fencein reflection, equality, and hash.The new
need_fencefield is not exposed inRegisterReflection()and is omitted fromSEqualReduceandSHashReduce. This means:
- The field won't be accessible via reflection APIs
- Two barrier ops with different fence requirements would incorrectly be considered equal
🔎 Proposed fix
static void RegisterReflection() { namespace refl = tvm::ffi::reflection; refl::ObjectDef<BarrierBlocksOpNode>() .def_ro("local_bar_addr", &BarrierBlocksOpNode::local_bar_addr) .def_ro("offset", &BarrierBlocksOpNode::offset) .def_ro("local_bar", &BarrierBlocksOpNode::local_bar) - .def_ro("local_indices", &BarrierBlocksOpNode::local_indices); + .def_ro("local_indices", &BarrierBlocksOpNode::local_indices) + .def_ro("need_fence", &BarrierBlocksOpNode::need_fence); } bool SEqualReduce(const BarrierBlocksOpNode *other, SEqualReducer equal) const { return equal(local_bar_addr, other->local_bar_addr) && equal(offset, other->offset) && equal(local_bar, other->local_bar) && - equal(local_indices, other->local_indices); + equal(local_indices, other->local_indices) && + need_fence == other->need_fence; } void SHashReduce(SHashReducer hash_reduce) const { hash_reduce(local_bar_addr); hash_reduce(offset); hash_reduce(local_bar); hash_reduce(local_indices); + hash_reduce(need_fence); }Also applies to: 145-151, 154-159, 161-166
src/op/sync.cc (2)
194-197: Fix arity mismatch: declare 2 inputs, not 1.The
BarrierBlocksOpconstructor accesses bothargs[0](local_bar_addr) andargs[1](need_fence), but.set_num_inputs(1)declares only one input. This causes out-of-bounds access.🔎 Proposed fix
TIR_REGISTER_TL_OP(BarrierBlocksOp, barrier_blocks) - .set_num_inputs(1) + .set_num_inputs(2) .set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kOpaque));
138-146: Add bounds checking forrelationindex.The
WaitOpconstructor extractsrelationfromargs[0]without validation. TheLowermethod indexesrelation_str[relation], which has 6 elements (valid indices 0-5). Invalid values cause undefined behavior.🔎 Proposed fix
WaitOp::WaitOp(Array<PrimExpr> args, BufferMap vmap) { ObjectPtr<WaitOpNode> node = make_object<WaitOpNode>(); node->relation = args[0].as<IntImmNode>()->value; + ICHECK(node->relation >= 0 && node->relation < 6) + << "relation must be in range [0, 5], got " << node->relation; node->addr = args[1]; node->expected = args[2]; node->peer = args[3]; data_ = std::move(node); (void)vmap; }tilelang/language/builtin.py (1)
765-768: Docstring incorrectly lists "release" for load operations.The docstring mentions
'release'as a valid semantic, but the assertion correctly excludes it (release semantics don't apply to loads). Update the docstring to match the assertion.🔎 Proposed fix
assert sem in [ "weak", "volatile", "acquire", "relaxed" - ], "Semantic must be one of 'weak', 'volatile', 'acquire', 'release', or 'relaxed'." + ], "Semantic must be one of 'weak', 'volatile', 'acquire', or 'relaxed'."examples/distributed/deepseek_deepep/deepep.md (3)
3-3: Wrap bare URL in markdown link syntax.The bare URL should be formatted as a markdown link for better readability and accessibility.
🔎 Proposed fix
-To install and compare with the original DeepEP implementation, please refer to https://github.com/deepseek-ai/DeepEP. +To install and compare with the original DeepEP implementation, please refer to [DeepEP](https://github.com/deepseek-ai/DeepEP).
26-28: Remove stray character "z".Line 27 contains a lone "z" which appears to be an accidental keystroke.
🔎 Proposed fix
This example implements DeepEP's intra‑node (NVLink) dispatch/combine using TileScale kernels. -z + The intra‑node path lives under `intranode/` and provides a minimal public API that mirrors DeepEP's behavior for NVLink‑connected ranks.
146-149: Add language identifier to code block.The shell command block should specify
bashorshellas the language identifier for syntax highlighting.🔎 Proposed fix
-``` +```bash CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 TILELANG_USE_DISTRIBUTED=1 python intranode/example_intranode.py \ --num_ranks 8 --num_tokens 4096 --hidden 7168 --num_topk 8 --num_experts 32 [--cached] ```src/target/codegen_cuda.cc (1)
2361-2367:tl::loop_continueis still emitted as an expression instead of a statement
loop_breakcorrectly emits a statement usingthis->PrintIndent()andthis->stream << "break;\n";, butloop_continuestill writes"continue"to the expression streamos. This remains inconsistent and can produce malformed code whentl::loop_continue()appears in statement position. It should be emitted likebreak(indented, with a terminating semicolon and newline) on the statement stream.examples/distributed/deepseek_deepep/intranode/combine.py (1)
98-110: Avoid hardcodingdtype='bfloat16'forcombine_kernel
combine_kernelis parameterized bydtype, butintranode_combinealways instantiates it with'bfloat16'regardless ofx.dtype. This will break or silently misbehave whenxisfloat16/float32/other types.Prefer deriving the dtype string from the input tensor, e.g.:
dtype_str = str(x.dtype).split('.')[-1] # 'torch.bfloat16' -> 'bfloat16' kernel = combine_kernel( num_ranks, config.num_max_nvl_chunked_send_tokens, config.num_max_nvl_chunked_recv_tokens, hidden, num_topk, config.num_sms, dtype=dtype_str, )and ensure the TileLang side supports that dtype.
Also applies to: 415-423
examples/distributed/deepseek_deepep/intranode/dispatch.py (2)
125-187: Handlecomm_streambeingNonein notify paths
notify_dispatchandcached_notify_dispatchboth unconditionally usecomm_stream.cuda_streamforkernel.initialize(...)and kernel launch, yet their signatures defaultcomm_stream=None. If callers don’t provide a stream (or explicitly passNone), you’ll getAttributeError.Either:
- require
comm_stream(assert non-None and drop the default), or- default to the current CUDA stream, e.g.:
stream_obj = comm_stream or torch.cuda.current_stream() cuda_stream = stream_obj.cuda_stream kernel.initialize(allocator=allocator, stream=cuda_stream) kernel(..., stream=cuda_stream, skip_tensor_validation=True)Apply the same pattern wherever
comm_stream.cuda_streamis used.Also applies to: 220-244
829-970: Multiple correctness and API issues inintranode_dispatchWithin
intranode_dispatchthere are several problems that should be addressed together:
- Shape access in cached mode can raise
AttributeErrornum_experts = num_tokens_per_expert.shape[0] if handle is None else 0 num_ranks = num_tokens_per_rank.shape[0] num_topk = topk_idx.shape[1] if handle is None else 0In cached mode (
handle is not None),num_tokens_per_rank,num_tokens_per_expert, andtopk_idxare allowed to beNone, butnum_ranksandnum_experts/num_topkstill dereference.shape. Compute those from handle-derived tensors instead (e.g.,rank_prefix_matrix.shape[0]) and gate any uses ofnum_tokens_per_*/topk_*onhandle is None.
- Hardcoded dtype
'bfloat16'for both kernelsBoth branches construct kernels with a literal
'bfloat16':kernel = dispatch_kernel(..., 'bfloat16') ... kernel = cached_dispatch_kernel(..., 'bfloat16')This ignores
x.dtypeand will misbehave for non-bfloat16 inputs. Derive a dtype string fromx.dtype(as in the earlier review suggestion) and pass that instead.
- Inconsistent return type between cached and non-cached branches
- Non-cached branch returns:
return recv_x, recv_topk_idx, recv_topk_weights, num_recv_tokens_per_expert_list, handle- Cached branch returns only:
return recv_xThis makes the function hard to use safely. Prefer always returning the same 5-tuple, e.g.,
(recv_x, recv_topk_idx, recv_topk_weights, num_recv_tokens_per_expert_list, handle), withNoneplaceholders for fields that are not meaningful in cached mode.
comm_streamandinitializeusage
- Non-cached branch calls
kernel.initialize(allocator=allocator)with no stream, then launches the kernel oncomm_stream.cuda_stream. Sinceinit_tableis now async and stream-aware, consider passingcomm_stream.cuda_streamtoinitializeas well to ensure the metadata copy is ordered with respect to subsequent kernels.- Cached branch again assumes
comm_streamis non-None (stream=comm_stream.cuda_streamandinitialize(..., stream=comm_stream.cuda_stream)), but the function signature allowscomm_stream=None.These should be cleaned up with the same pattern as in the notify helpers: either require
comm_stream, or default it totorch.cuda.current_stream()and pass that into bothinitializeand the kernel launch.examples/distributed/deepseek_deepep/intranode/get_dispatch_layout.py (1)
16-34: Fix mismatch between docstring/signature and actual return valuesThe docstring and type hint declare four return values, including
num_tokens_per_rdma_rank(possiblyNoneintranode), but the implementation returns only three:return num_tokens_per_rank, num_tokens_per_expert, is_token_in_rankEither:
- update the signature and docstring to describe three return values, or
- return a 4-tuple, e.g.:
return num_tokens_per_rank, None, num_tokens_per_expert, is_token_in_rankto match the documented contract.
Also applies to: 60-60
🧹 Nitpick comments (19)
tilelang/utils/sparse.py (1)
15-15: Version-isolated cache prevents library conflicts.The versioned cache directory structure correctly ensures that compiled sparse compressor libraries don't clash across TileLang releases.
As versions accumulate, old cached libraries will remain on disk. Consider documenting a cache cleanup strategy or adding a utility to remove stale version directories in future work.
.github/workflows/ci.yml (3)
85-86: Remove unused REQS_HASH and MARKER variables.These variables are computed but never used since venv reuse is intentionally disabled (line 87). Remove lines 85-86 to eliminate dead code.
🔎 Proposed fix to remove dead code
- name: Ensure venv (local & persistent) run: | set -e - REQS_HASH=$(cat requirements-test.txt 2>/dev/null || true) - MARKER="${{ runner.tool_cache }}/.venv_marker_${{ env.PYTHON_VERSION }}_${REQS_HASH:0:8}" # NOTE(wt): We disable the venv reuse for now to allow installing DeepEP # echo "venv stale or missing – recreating" rm -rf "${{ runner.tool_cache }}/${{ env.VENV_DIR }}"
121-121: Track the temporary test exclusions.The grep filter excludes
sinkandvs_sparsetests as a temporary measure. Ensure there's a tracking issue or plan to re-enable these tests, and consider referencing it in the comment.Would you like me to help create a tracking issue for re-enabling these temporarily disabled tests?
147-147: Track the temporary test exclusions.Similar to line 121, this excludes
tilelibrary_gemmandf16f16f16tests temporarily. Ensure these exclusions are tracked with a plan to re-enable or remove the tests.Would you like me to help create a tracking issue that covers all temporarily disabled tests (lines 121 and 147)?
examples/distributed/primitives/example_remote_st.py (3)
14-14: Remove unusednum_rankparameter.The
num_rankparameter is never used in the function body. The kernel obtains rank information dynamically viaT.get_rank()andT.get_num_ranks()at runtime (lines 24-25), making this parameter unnecessary.🔎 Proposed fix
-def kernel_(M, num_rank, block_M, threads): +def kernel_(M, block_M, threads):Also update the call site on line 45:
- kernel = tilelang.compile(kernel_(M, num_ranks, BLOCK_M, threads)) + kernel = tilelang.compile(kernel_(M, BLOCK_M, threads))
37-37: Prefix unusedrankvariable with underscore.The
rankvariable frominit_distis unpacked but never used. The function useslocal_rankinstead throughout. Consider prefixing it with an underscore to indicate it's intentionally unused.🔎 Proposed fix
- rank, num_ranks, group = init_dist(local_rank, num_local_ranks) + _rank, num_ranks, group = init_dist(local_rank, num_local_ranks)
68-68: Consider using a shorter error message.The error message is specified inline, which can make the exception less reusable. For a simple example like this, the current approach is acceptable, but consider using a brief message or referencing the printed output above.
🔎 Proposed fix
- raise ValueError("Test failed") + raise ValueError("Validation failed")testing/python/language/test_tilelang_language_ldst_options.py (1)
86-87: Add comments for consistency.Lines 86-87 test the "weak" semantic but lack explanatory comments, unlike the test cases above. Adding comments would improve consistency and documentation.
🔎 Suggested addition
+ # st.weak.gpu.global.b* _test_st_options("gpu", "weak", False) + # st.weak.gpu.global.L1::no_allocate.b* _test_st_options("gpu", "weak", True)tilelang/distributed/testing/test_create_mapped_tensor.py (1)
4-19: Consider converting to a proper test framework.The test is currently a standalone script. Converting it to pytest or unittest would provide:
- Automatic test discovery
- Better error reporting and test isolation
- Easier CI/CD integration
- Consistent test patterns across the codebase
🔎 Example pytest conversion
import pytest import torch from tilelang.distributed.utils import create_mapped_tensor @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") def test_create_mapped_tensor(): shape = (1024, 1024) dtype = torch.float32 host_tensor, device_tensor = create_mapped_tensor(shape, dtype) # test meta-data assert device_tensor.device.type == "cuda" assert device_tensor.shape == shape assert device_tensor.dtype == dtype assert torch.equal(host_tensor, device_tensor.cpu()) # test modification device_tensor.random_() torch.cuda.synchronize() assert torch.equal(host_tensor, device_tensor.cpu())tilelang/language/print.py (1)
226-227: Optional: Consider extracting the error message to improve maintainability.The error message correctly reflects the updated supported types. However, as a style improvement, you could define the message as a module-level constant or create a custom exception class to address the Ruff TRY003 hint.
🔎 Optional refactor to address TRY003
+_UNSUPPORTED_TYPE_MSG = "Unexpected type: {}. Supported types are tir.Buffer, tir.PrimExpr and None" + ... else: # Unsupported object type. - raise ValueError( - f"Unexpected type: {type(obj)}. Supported types are tir.Buffer, tir.PrimExpr and None") + raise ValueError(_UNSUPPORTED_TYPE_MSG.format(type(obj)))Based on static analysis (Ruff TRY003).
tilelang/utils/ts_ext/__init__.py (1)
9-18: LGTM!The new
create_host_device_tensoralias is correctly exposed from the C extension and properly added to__all__.Optionally, you could sort
__all__alphabetically to satisfy the Ruff RUF022 lint (isort-style sorting), but this is purely a style preference.src/transform/vectorize_loop.cc (1)
702-836: Consider consolidating withloop_vectorization_utils.h.The
MutateTlLdStExpr_implementation here is nearly identical (~130 lines) to the one inloop_vectorization_utils.h. This duplication:
- Increases maintenance burden
- Risks divergent behavior (as shown by the bug above)
Consider extracting common logic into a shared utility or having
TLVectorizerreuse the implementation fromloop_vectorization_utils.h.tilelang/language/__init__.py (1)
66-70: Consider removing unusednoqadirectives.Per Ruff (RUF100), the
# noqa: F401directives on these lines are unnecessary. However, if consistency with the rest of the file is preferred, this is a minor style choice.🔎 Proposed fix
- warp_reduce_sum, # noqa: F401 - warp_reduce_max, # noqa: F401 - warp_reduce_min, # noqa: F401 - warp_reduce_bitand, # noqa: F401 - warp_reduce_bitor, # noqa: F401 + warp_reduce_sum, + warp_reduce_max, + warp_reduce_min, + warp_reduce_bitand, + warp_reduce_bitor,testing/python/language/test_tilelang_language_warp_reduce.py (1)
43-49: Remove debug print statement.The
print(kernel.get_kernel_source())on line 46 appears to be a debugging artifact. Consider removing it to keep test output clean.🔎 Proposed fix
def test_warp_reduce_max(): a = torch.randn((32,), dtype=torch.float32, device='cuda') kernel = get_kernel("max", 'float32') - print(kernel.get_kernel_source()) ref = torch.full_like(a, a.max()) kernel(a) torch.testing.assert_close(a, ref)examples/distributed/deepseek_deepep/buffer.py (1)
252-252: Prefix unused unpacked variables with underscore.In cached dispatch mode, only
is_token_in_rankfrom the handle is used. Prefix the unused variables (rank_prefix_matrix,channel_prefix_matrix,recv_channel_prefix_matrix,recv_src_idx,send_head) with underscores to signal intent.🔎 Proposed fix
- rank_prefix_matrix, channel_prefix_matrix, recv_channel_prefix_matrix, recv_src_idx, is_token_in_rank, send_head = handle + _rank_prefix_matrix, _channel_prefix_matrix, _recv_channel_prefix_matrix, _recv_src_idx, is_token_in_rank, _send_head = handleexamples/distributed/deepseek_deepep/intranode/example_intranode.py (1)
40-41: Prefix unused unpacked variables with underscore.The
rank_idx(line 40) andevent(line 62) variables are unpacked but never used. Prefix them with underscores to signal they are intentionally unused.🔎 Proposed fix
- x, topk_idx, topk_weights, rank_idx = gen_inputs(num_tokens, hidden, num_topk, num_experts, + x, topk_idx, topk_weights, _rank_idx = gen_inputs(num_tokens, hidden, num_topk, num_experts, num_ranks)- ref_recv_x, ref_recv_topk_idx, ref_recv_topk_weights, ref_num_recv_tokens_per_expert_list, ref_handle, event = \ + ref_recv_x, ref_recv_topk_idx, ref_recv_topk_weights, ref_num_recv_tokens_per_expert_list, ref_handle, _event = \ deepep_buffer.dispatch(x, None, ref_num_tokens_per_rank, None, ref_is_token_in_rank, ref_num_tokens_per_expert, topk_idx, topk_weights, expert_alignment)Also applies to: 62-63
src/tl_templates/cuda/sync.h (1)
187-193: Inconsistent template signature forwait_eq.
wait_equsesvoid *ptrwhile the otherwait_*functions usetemplate<typename P, typename T>withPconstrained to pointer oruint64_t. Consider unifying the signatures for consistency unless there's a specific reason forwait_eqto differ.🔎 Proposed fix (if unification is desired)
-template <typename T> TL_DEVICE void wait_eq(void *ptr, T val) { - T *flag_ptr = reinterpret_cast<T *>(ptr); +template <typename P, typename T> TL_DEVICE void wait_eq(P ptr, T val) { + static_assert(std::is_same_v<P, uint64_t> || std::is_pointer_v<P>, + "P must be a pointer or uint64_t"); + T *flag_ptr = reinterpret_cast<T *>(ptr); // Spin-loop #pragma unroll 1 while (ld_acquire(flag_ptr) != val) ; }Also applies to: 195-243
tilelang/language/builtin.py (1)
514-514: Consider explicitOptionalannotation for clarity.The parameters
barrier_idandarrive_countdefault toNonebut lack explicitOptional[...]type hints. While this works at runtime, PEP 484 recommends being explicit.🔎 Proposed refactor
-def sync_threads(barrier_id: int | PrimExpr = None, arrive_count: int = None): +def sync_threads(barrier_id: int | PrimExpr | None = None, arrive_count: int | None = None):tilelang/jit/adapter/wrapper.py (1)
46-66: Stream-awareinit_tableis fine; ensure callers pass the execution streamSwitching to
cudaMemcpyToSymbolAsync(..., stream)with an explicitcudaStream_tparameter is reasonable and enables overlap. Correctness now depends on callers consistently passing the same stream that subsequent kernels use when readingmeta_data; otherwise, there’s a risk of usingmeta_databefore the async copy completes on a different stream.
📜 Review details
Configuration used: defaults
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (57)
.github/workflows/ci.yml.gitmodules3rdparty/DeepEPbenchmark/distributed/deepep/deepep.mdbenchmark/distributed/deepep/intranode/get_dispatch_layout.pyexamples/distributed/deepseek_deepep/buffer.pyexamples/distributed/deepseek_deepep/deepep.mdexamples/distributed/deepseek_deepep/deepep_utils.pyexamples/distributed/deepseek_deepep/intranode/combine.pyexamples/distributed/deepseek_deepep/intranode/dispatch.pyexamples/distributed/deepseek_deepep/intranode/example_intranode.pyexamples/distributed/deepseek_deepep/intranode/get_dispatch_layout.pyexamples/distributed/deepseek_deepep/intranode/test_intranode.pyexamples/distributed/primitives/example_put_warp.pyexamples/distributed/primitives/example_remote_st.pyexamples/distributed/sp_ag_attention_intra_node.pypyproject.tomlsrc/op/builtin.ccsrc/op/builtin.hsrc/op/remote_copy.ccsrc/op/remote_copy.hsrc/op/sync.ccsrc/op/sync.hsrc/target/codegen_cuda.ccsrc/tl_templates/cuda/common.hsrc/tl_templates/cuda/copy.hsrc/tl_templates/cuda/debug.hsrc/tl_templates/cuda/ldst.hsrc/tl_templates/cuda/reduce.hsrc/tl_templates/cuda/sync.hsrc/transform/common/loop_vectorization_utils.hsrc/transform/loop_vectorize.ccsrc/transform/storage_access.ccsrc/transform/vectorize_loop.cctesting/python/language/test_tilelang_language_elect.pytesting/python/language/test_tilelang_language_ldst_options.pytesting/python/language/test_tilelang_language_vote.pytesting/python/language/test_tilelang_language_warp_reduce.pytilelang/distributed/install_deepep.shtilelang/distributed/testing/sync/test_barrierall_sys.pytilelang/distributed/testing/test_create_mapped_tensor.pytilelang/distributed/utils.pytilelang/jit/adapter/wrapper.pytilelang/jit/kernel.pytilelang/language/__init__.pytilelang/language/builtin.pytilelang/language/distributed/common.pytilelang/language/print.pytilelang/language/reduce.pytilelang/utils/sparse.pytilelang/utils/ts_ext/__init__.pytilelang/utils/ts_ext/exception.htilelang/utils/ts_ext/ipc_ops.cpptilelang/utils/ts_ext/setup.pytilelang/utils/ts_ext/tensor.cpptilelang/utils/ts_ext/ts_ext_bindings.cpptilelang/utils/ts_ext/ts_ext_ops.h
💤 Files with no reviewable changes (2)
- benchmark/distributed/deepep/deepep.md
- benchmark/distributed/deepep/intranode/get_dispatch_layout.py
✅ Files skipped from review due to trivial changes (1)
- 3rdparty/DeepEP
🚧 Files skipped from review as they are similar to previous changes (10)
- pyproject.toml
- tilelang/utils/ts_ext/setup.py
- tilelang/language/reduce.py
- tilelang/utils/ts_ext/ipc_ops.cpp
- tilelang/distributed/testing/sync/test_barrierall_sys.py
- examples/distributed/sp_ag_attention_intra_node.py
- .gitmodules
- tilelang/utils/ts_ext/ts_ext_ops.h
- tilelang/distributed/install_deepep.sh
- tilelang/distributed/utils.py
🧰 Additional context used
🪛 actionlint (1.7.9)
.github/workflows/ci.yml
10-10: label "nvidia" is unknown. available labels are "windows-latest", "windows-latest-8-cores", "windows-2025", "windows-2022", "windows-11-arm", "ubuntu-slim", "ubuntu-latest", "ubuntu-latest-4-cores", "ubuntu-latest-8-cores", "ubuntu-latest-16-cores", "ubuntu-24.04", "ubuntu-24.04-arm", "ubuntu-22.04", "ubuntu-22.04-arm", "macos-latest", "macos-latest-xl", "macos-latest-xlarge", "macos-latest-large", "macos-26-xlarge", "macos-26", "macos-15-intel", "macos-15-xlarge", "macos-15-large", "macos-15", "macos-14-xl", "macos-14-xlarge", "macos-14-large", "macos-14", "macos-13-xl", "macos-13-xlarge", "macos-13-large", "macos-13", "self-hosted", "x64", "arm", "arm64", "linux", "macos", "windows". if it is a custom label for self-hosted runner, set list of labels in actionlint.yaml config file
(runner-label)
10-10: label "hopper" is unknown. available labels are "windows-latest", "windows-latest-8-cores", "windows-2025", "windows-2022", "windows-11-arm", "ubuntu-slim", "ubuntu-latest", "ubuntu-latest-4-cores", "ubuntu-latest-8-cores", "ubuntu-latest-16-cores", "ubuntu-24.04", "ubuntu-24.04-arm", "ubuntu-22.04", "ubuntu-22.04-arm", "macos-latest", "macos-latest-xl", "macos-latest-xlarge", "macos-latest-large", "macos-26-xlarge", "macos-26", "macos-15-intel", "macos-15-xlarge", "macos-15-large", "macos-15", "macos-14-xl", "macos-14-xlarge", "macos-14-large", "macos-14", "macos-13-xl", "macos-13-xlarge", "macos-13-large", "macos-13", "self-hosted", "x64", "arm", "arm64", "linux", "macos", "windows". if it is a custom label for self-hosted runner, set list of labels in actionlint.yaml config file
(runner-label)
🪛 markdownlint-cli2 (0.18.1)
examples/distributed/deepseek_deepep/deepep.md
3-3: Bare URL used
(MD034, no-bare-urls)
146-146: Fenced code blocks should have a language specified
(MD040, fenced-code-language)
🪛 Ruff (0.14.10)
examples/distributed/deepseek_deepep/intranode/example_intranode.py
31-31: Unused noqa directive (unused: F403)
Remove unused noqa directive
(RUF100)
33-33: Avoid specifying long messages outside the exception class
(TRY003)
40-40: Unpacked variable rank_idx is never used
Prefix it with an underscore or any other dummy variable pattern
(RUF059)
62-62: Unpacked variable event is never used
Prefix it with an underscore or any other dummy variable pattern
(RUF059)
tilelang/language/__init__.py
66-66: Unused noqa directive (non-enabled: F401)
Remove unused noqa directive
(RUF100)
67-67: Unused noqa directive (non-enabled: F401)
Remove unused noqa directive
(RUF100)
68-68: Unused noqa directive (non-enabled: F401)
Remove unused noqa directive
(RUF100)
69-69: Unused noqa directive (non-enabled: F401)
Remove unused noqa directive
(RUF100)
70-70: Unused noqa directive (non-enabled: F401)
Remove unused noqa directive
(RUF100)
tilelang/language/print.py
226-227: Avoid specifying long messages outside the exception class
(TRY003)
examples/distributed/deepseek_deepep/buffer.py
252-252: Unpacked variable rank_prefix_matrix is never used
Prefix it with an underscore or any other dummy variable pattern
(RUF059)
252-252: Unpacked variable channel_prefix_matrix is never used
Prefix it with an underscore or any other dummy variable pattern
(RUF059)
252-252: Unpacked variable recv_channel_prefix_matrix is never used
Prefix it with an underscore or any other dummy variable pattern
(RUF059)
252-252: Unpacked variable recv_src_idx is never used
Prefix it with an underscore or any other dummy variable pattern
(RUF059)
252-252: Unpacked variable send_head is never used
Prefix it with an underscore or any other dummy variable pattern
(RUF059)
examples/distributed/deepseek_deepep/intranode/combine.py
238-238: Shebang should contain python, pytest, or uv run
(EXE003)
238-238: Shebang should be at the beginning of the file
(EXE005)
examples/distributed/deepseek_deepep/intranode/dispatch.py
13-13: Unused noqa directive (unused: F403)
Remove unused noqa directive
(RUF100)
508-508: Shebang should contain python, pytest, or uv run
(EXE003)
508-508: Shebang should be at the beginning of the file
(EXE005)
786-786: Shebang should contain python, pytest, or uv run
(EXE003)
786-786: Shebang should be at the beginning of the file
(EXE005)
examples/distributed/primitives/example_remote_st.py
14-14: Unused function argument: num_rank
(ARG001)
37-37: Unpacked variable rank is never used
Prefix it with an underscore or any other dummy variable pattern
(RUF059)
68-68: Avoid specifying long messages outside the exception class
(TRY003)
tilelang/language/builtin.py
514-514: PEP 484 prohibits implicit Optional
Convert to Optional[T]
(RUF013)
tilelang/jit/kernel.py
410-410: PEP 484 prohibits implicit Optional
Convert to Optional[T]
(RUF013)
tilelang/utils/ts_ext/__init__.py
11-18: __all__ is not sorted
Apply an isort-style sorting to __all__
(RUF022)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (3)
- GitHub Check: Build wheels for Python 3.8 on ubuntu-24.04-arm with CUDA-12.8
- GitHub Check: Build wheels for Python 3.8 on macos-latest with Metal
- GitHub Check: Build wheels for Python 3.8 on ubuntu-latest with CUDA-12.1
🔇 Additional comments (51)
testing/python/language/test_tilelang_language_vote.py (1)
8-37: LGTM!The kernel correctly tests the three fundamental cases for warp vote primitives:
- All threads true (value=1)
- All threads false (value=0)
- Mixed values (value=tx%2)
The expected output
[1, 1, 0, 0, 1, 0]correctly reflects the semantics ofwarp_anyandwarp_all.tilelang/utils/sparse.py (1)
9-11: LGTM! Clear versioning strategy.The import and comment clearly document the intent to isolate caches by version, preventing conflicts when multiple TileLang versions coexist.
src/transform/storage_access.cc (1)
291-294: LGTM! Correctly isolates condition accesses from body summary.This change aligns
WhileNodehandling with theIfThenElseNodepattern (lines 251-256). By settingallow_append_ = truebefore visiting the condition and clearingcurr_stmt_.accessafterward, condition-time buffer accesses (e.g., termination flag reads) are properly excluded from the loop body's access summary. This prevents contamination during scope accumulation and is essential for correct storage access analysis..github/workflows/ci.yml (1)
102-102: LGTM! DeepEP installation added.This addresses the previous review comment about missing DeepEP installation in the build-test-nvidia job.
examples/distributed/primitives/example_remote_st.py (1)
73-81: LGTM!The CLI setup is clean and follows standard PyTorch distributed patterns. The default of 2 processes aligns well with the 2-rank assertion in the main function.
testing/python/language/test_tilelang_language_ldst_options.py (4)
7-18: LGTM!The LD kernel builder correctly parameterizes scope, sem, na, and nc options for testing various PTX load instruction variants.
21-30: LGTM!The ST kernel builder correctly parameterizes scope, sem, and na options for testing various PTX store instruction variants.
33-38: LGTM!The LD test helper correctly validates that the load operation preserves data by asserting the loaded output matches the input.
41-60: LGTM!Comprehensive test coverage for LD instructions across different scopes (sys, gpu, cta) and memory semantics (acquire, relaxed, volatile, weak). The inline comments documenting PTX instructions are helpful.
examples/distributed/primitives/example_put_warp.py (1)
18-19: No issues found. The dtype change from float32 to bfloat16 is valid and properly supported throughout the codebase. The tensor dialect is intended to hold core tensor creation and manipulation ops, with the aim that they make sense for any tensor element type, and bfloat16 is a standard, widely-used dtype in this project as evidenced by its extensive usage in quantization, GEMM, attention, and distributed operations.tilelang/distributed/testing/test_create_mapped_tensor.py (1)
1-2: No PyTorch version compatibility concern for these imports. PyTorch's standard CUDA and tensor operations are stable across versions. Thecreate_mapped_tensorimport is from the tilelang library (not PyTorch), so any version compatibility needs would apply to that library, not PyTorch's base functionality.Likely an incorrect or invalid review comment.
tilelang/language/print.py (2)
136-141: LGTM!The new
print_msgfunction is well-structured and follows the existing patterns in the file. The@macrodecorator is appropriate, and the implementation correctly routes to thedebug_print_msgexternal call.
144-144: Verify backward compatibility of signature change.The addition of a default value for
objis logically correct and enables printing standalone messages. The implementation properly routesNonetoprint_msg. All existing call sites across the codebase continue to work as expected with the signature change.tilelang/utils/ts_ext/tensor.cpp (2)
11-11: LGTM!Appropriate include for the new
CUDA_CHECKmacro used increate_host_device_tensor.
38-46: LGTM!Good additions: the
"long int"alias improves compatibility, and quoting the dtype string in the error message aids debugging.tilelang/utils/ts_ext/exception.h (1)
1-36: LGTM!Well-structured exception handling:
TSExceptionproperly extendsstd::exceptionwithnoexceptonwhat().- Macros use the
do-while(0)idiom correctly for safe macro expansion.#ifndefguards prevent redefinition conflicts.CUDA_CHECKcorrectly evaluates the command once before checking the result.tilelang/utils/ts_ext/ts_ext_bindings.cpp (1)
25-27: The binding correctly relies on PyTorch's automatic dtype conversion.The file includes
<torch/extension.h>(line 4), which provides PyTorch's pybind11 type casters including automatic conversion forat::ScalarType(equivalent toc10::ScalarType). Whencreate_host_device_tensoris bound directly to the C++ function acceptingc10::ScalarType, pybind11 automatically converts Pythontorch.dtypeobjects without requiring explicit conversion code. The binding is correct and will function as intended.Likely an incorrect or invalid review comment.
src/transform/loop_vectorize.cc (1)
141-180: LGTM! Vectorization planning fortl::ld/tl::stis well-structured.The pattern detection correctly validates the source address via
address_of(BufferLoad)and appropriately updates the vector size based on buffer indices. The optional destination check (lines 166-169) is reasonable since not alltl::stcalls will have a BufferLoad pattern for the value.src/transform/common/loop_vectorization_utils.h (2)
389-521: LGTM! Vectorization logic fortl::ld/tl::stis well-implemented.The implementation correctly:
- Extracts and visits indices to identify Ramp patterns
- Computes vector lanes from both source and destination
- Determines appropriate vector types (128-bit → int4, 64-bit → int2, 32-bit → int)
- Creates reinterpreted values and uses
vec_valuedirectly (line 491)The fallback to scalarization for unsupported bit widths (line 479) is appropriate.
571-583: LGTM! Routing logic correctly delegates to vectorization handler.The pattern detection and routing to
MutateTlLdStExpr_is consistent with the planning logic inloop_vectorize.cc.src/transform/vectorize_loop.cc (1)
459-470: LGTM for correctness; routing logic is valid.The pattern detection and routing to
MutateTlLdStExpr_is correct.src/tl_templates/cuda/copy.h (4)
114-124: LGTM!The new
int16_tanduint16_tspecializations forld_nc_globalfollow the correct pattern and use appropriate PTX instructions (.s16and.u16) with proper constraint handling.
168-176: LGTM!The new
int16_tanduint16_tspecializations forst_na_globalare correctly implemented with appropriate PTX instructions and constraints.
200-201: LGTM!Clean refactoring to rename the implementation function, allowing new wrapper functions to provide enhanced dispatch logic.
97-97: No change from .volatile to .nc occurred. The LD_NC_FUNC macro and ld_nc_global() function have always used non-coherent cache semantics (.nc), which is intentionally different from the separate ld_volatile_global() function in sync.h. This design properly separates volatile semantics (for synchronization) from non-coherent cache semantics (for efficient data loading), with each serving its appropriate purpose.Likely an incorrect or invalid review comment.
src/tl_templates/cuda/reduce.h (1)
283-314: Well-implemented warp-wide reduction utilities.The butterfly reduction pattern using
__shfl_xor_syncwith offsets 16, 8, 4, 2, 1 correctly reduces across all 32 lanes. The type-safe wrappers for common operations (sum, max, min, bitand, bitor) provide a clean API.examples/distributed/deepseek_deepep/intranode/test_intranode.py (1)
1-13: Test structure follows project conventions.The test appropriately uses
@tilelang.testing.requires_cudato guard execution and relies onexample_intranode.main()to raise on failures. Consider adding explicit assertions or at least a comment clarifying the expected behavior if the test should validate specific outcomes beyond "no exception thrown."testing/python/language/test_tilelang_language_warp_reduce.py (1)
35-41: Good test coverage for warp reduction operations.The tests properly validate each reduction operation against NumPy/torch references. The use of
torch.testing.assert_closeprovides appropriate tolerance for floating-point comparisons.testing/python/language/test_tilelang_language_elect.py (1)
1-30: Well-structured test forelect_one_syncintrinsic.The test correctly validates both codegen (checking for
cute::elect_one_syncin source) and runtime behavior (exactly one thread increments the counter). The SM90+ compute version requirement is appropriate for this Hopper-specific feature.src/op/builtin.cc (1)
318-360: New builtin intrinsics properly registered.The warp-level intrinsics are registered with appropriate arities and effect kinds. Warp reductions correctly use
kOpaque(warp shuffles have synchronization side effects), whilewarp_anyandwarp_allare markedkPure(consistent with vote intrinsics being deterministic given their inputs).src/op/builtin.h (1)
548-599: Header declarations consistent with implementation.The new intrinsic declarations are properly documented and match the corresponding registrations in
builtin.cc. The doxygen comments provide sufficient context for each operation.src/tl_templates/cuda/common.h (1)
298-300: No breaking API changes require migration. The current__sync_thread_partialfunction signature uses runtime parameters, and all call sites in the codebase already invoke it correctly with runtime arguments. No template-style calls exist.Likely an incorrect or invalid review comment.
src/op/remote_copy.cc (1)
86-89: LGTM!The
is_distributed()methods consistently check whether the PE (processing element) is a non-distributed sentinel value (-1). This pattern is correct and aligns acrossPutOpNode,GetOpNode,StOpNode, andLdOpNode.Also applies to: 196-199, 264-267, 333-336
src/op/sync.cc (1)
83-110: LGTM!The
BarrierBlocksOpNode::Lowermethod correctly implements optional fencing withif constexpr (need_fence)and properly generates the barrier call with distributed address calculations.examples/distributed/deepseek_deepep/deepep_utils.py (1)
17-90: LGTM!The
Configdataclass with validation and static factory methods is well-structured. The__post_init__validation ensuresnum_smsis even, and the config mappings for different rank counts are comprehensive.src/tl_templates/cuda/sync.h (1)
156-185: LGTM!The
barrier_blocksimplementation with optional fencing is well-structured. The use ofFINISHED_SUM_TAGfor signaling and the warp-level synchronization with__all_synccorrectly implements the distributed barrier pattern.src/op/remote_copy.h (1)
205-253: LGTM!The new
StOpNodeandLdOpNodedeclarations are well-structured with proper reflection registration, equality/hash methods, and type declarations. They follow the established patterns in the codebase.Also applies to: 262-313
src/op/sync.h (1)
54-105: LGTM!The
WaitOpNodeandWaitOpdeclarations are well-structured with complete reflection registration, equality comparisons, and hash reduction for all fields (addr,expected,peer,relation).tilelang/language/builtin.py (4)
9-9: LGTM!The
Literalimport is correctly added to support the new type-safeldandstfunction signatures.
615-633: LGTM!The distinction between
barrier_blocks(with system-level fence) andsync_blocks(without fence) is clearly documented and correctly implemented through the sharedtl.barrier_blocksintrinsic with different fence parameters.
777-808: LGTM!The
stfunction correctly excludesacquiresemantics (which don't apply to stores) and properly validates scope, semantic, and other parameters before delegating to the intrinsic.
811-849: LGTM!The warp-level intrinsics (
elect_one_sync,sync_warp,loop_continue,warp_any,warp_all) are correctly implemented with appropriate return types and clear documentation. The mask default of-1(0xFFFFFFFF) for vote operations is standard.src/tl_templates/cuda/ldst.h (6)
1-28: LGTM!The enum definitions, type traits, and forward declarations are well-structured. The
always_false_vguard prevents multiple definition issues, and theis_bfloat16trait correctly handles both CUDA native and CUTLASS bfloat16 types.
30-41: LGTM!The fallback templates with
static_assertensure that unsupported load/store configurations are caught at compile-time with clear error messages, preventing silent failures.
43-91: LGTM!The
TL_ST_IMPLmacro comprehensively handles 2/4/8/16-byte stores with appropriate PTX constraints (h,f,r,d,l) and special bfloat16 handling via bit-casting. The 16-byte vector store correctly uses.v4.s32format.
93-145: LGTM!The
TL_LD_IMPLmacro mirrors the store implementation with appropriate output constraints and handles all size/type combinations correctly, including bfloat16 bit-casting and 16-byte vector loads.
147-215: LGTM!The specialization registrations systematically cover all necessary semantic/scope/cache-policy combinations for stores (WEAK/VOLATILE/RELAXED/RELEASE) and loads (nc/WEAK/VOLATILE/RELAXED/ACQUIRE), providing comprehensive support for the documented configurations.
217-255: LGTM!The public
tl::standtl::ldinterfaces provide comprehensive compile-time validation (size, pointer type, semantics) and cleanly delegate to the specialized implementations. The TODO note aboutL2::256Bsupport is informational.src/tl_templates/cuda/debug.h (1)
261-266:debug_print_msghelper looks correct and consistentThe message-only device printf follows the same pattern as existing debug helpers; format string and arguments are aligned, and limiting it to msg + BlockIdx/ThreadIdx is appropriate for lightweight debugging.
src/target/codegen_cuda.cc (1)
810-862: Dynamicstorage_syncarguments are handled correctlyUsing
PrintExprforbarrier_idandthread_countinPrintStorageSynclets you support non-constant barrier parameters while preserving prior behavior for the 1-arg case. The emission oftl::__sync_thread_partial(...)looks consistent with the new sync API.tilelang/language/distributed/common.py (1)
22-47: Distributed put/get and wait helpers look consistent with call sites
- Widening
dst_pe/src_petoPrimExpr | IntImm | Nonewith default-1and threadingenable_aggressive_vectorizethrough to the intrinsic matches howput_warp/get_warpare now used in the DeepEP kernels (e.g., enabling aggressive vectorization for warp copies).- The block variants correctly fix the mode to
"block"and pass a dummy unroll factor (0) plus a finalTrueflag as per the updated intrinsic signature.- The
BinaryRelationenum andwait_*helpers encode relation op +address_of(ptr)cleanly and match usage patterns likeT.wait_ge(channel_head_idx[...], expected, peer=...)andT.wait_ne(channel_start_offset[...], 0).No functional issues stand out; the API surface is coherent with the new intranode dispatch/combine code.
Also applies to: 50-76, 78-113, 116-162
| // Wait for counters to be ready | ||
| while (true) { | ||
| bool ready = counter_ptr[0] >= 0; | ||
| for (int i = 0; i < num_local_experts and ready; ++i) |
There was a problem hiding this comment.
🛠️ Refactor suggestion | 🟠 Major
Use && instead of and in C++ code.
Line 233 uses the Python keyword and in C++ code. While some compilers accept and as an alternative token, && is the standard C++ operator and ensures better portability.
🔎 Proposed fix
- for (int i = 0; i < num_local_experts and ready; ++i)
+ for (int i = 0; i < num_local_experts && ready; ++i)
ready &= expert_ptr[i] >= 0;📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| for (int i = 0; i < num_local_experts and ready; ++i) | |
| for (int i = 0; i < num_local_experts && ready; ++i) |
🤖 Prompt for AI Agents
In examples/distributed/deepseek_deepep/deepep_utils.py around line 233, the
loop uses the Python logical operator "and" in C++ code; replace "and" with the
standard C++ logical AND operator "&&" to ensure correct syntax and portability
(i.e., change the condition to use "&&" between num_local_experts and ready).
| def cached_notify_combine( | ||
| num_ranks, | ||
| num_sms, | ||
| ##### symm buffers ##### | ||
| send_head: torch.Tensor, | ||
| channel_head_idx: torch.Tensor, | ||
| channel_tail_idx: torch.Tensor, | ||
| barrier_signal: torch.Tensor, | ||
| allocator, | ||
| comm_stream=None): | ||
| kernel = cached_notify_combine_kernel(num_ranks, num_sms) | ||
| kernel.initialize(allocator=allocator, stream=comm_stream.cuda_stream) | ||
|
|
||
| kernel( | ||
| send_head, | ||
| channel_head_idx, | ||
| channel_tail_idx, | ||
| barrier_signal, | ||
| stream=comm_stream.cuda_stream, | ||
| skip_tensor_validation=True) # reduce runtime overhead |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
git ls-files examples/distributed/deepseek_deepep/intranode/ 2>/dev/null || echo "Looking for the file..."Repository: tile-ai/tilescale
Length of output: 380
🏁 Script executed:
find / -path "*/examples/distributed/deepseek_deepep/intranode/combine.py" 2>/dev/null | head -1Repository: tile-ai/tilescale
Length of output: 137
🏁 Script executed:
git log --oneline -1 2>/dev/null && pwdRepository: tile-ai/tilescale
Length of output: 99
🏁 Script executed:
cat -n examples/distributed/deepseek_deepep/intranode/combine.py 2>/dev/null | head -100Repository: tile-ai/tilescale
Length of output: 4605
🏁 Script executed:
cat -n examples/distributed/deepseek_deepep/intranode/combine.py | sed -n '380,445p'Repository: tile-ai/tilescale
Length of output: 2577
🏁 Script executed:
grep -n "wait_stream" examples/distributed/deepseek_deepep/intranode/combine.pyRepository: tile-ai/tilescale
Length of output: 108
Guard comm_stream or make it required
Both cached_notify_combine and intranode_combine dereference comm_stream.cuda_stream (lines 87, 94, 423, 439) and later call compute_stream.wait_stream(comm_stream) (line 442) while the function signatures default comm_stream=None. If callers rely on the default, this will raise AttributeError or pass None into wait_stream.
Either:
- require a non-None
comm_stream(assert and drop the default), or - derive a safe default, e.g.:
stream_obj = comm_stream or torch.cuda.current_stream()
cuda_stream = stream_obj.cuda_streamand only call compute_stream.wait_stream(stream_obj) when comm_stream is not None.
Additionally, combine_kernel at line 422 hardcodes dtype='bfloat16' while recv_x at line 412 uses x.dtype. Derive the dtype string from x.dtype instead to ensure consistency.
🤖 Prompt for AI Agents
In examples/distributed/deepseek_deepep/intranode/combine.py around lines 76-95
(and related call sites around 412, 422, 423, 439, 442), callers currently
default comm_stream=None but the code dereferences comm_stream.cuda_stream and
later passes comm_stream into wait_stream; change this by either making
comm_stream a required non-None parameter (remove default and assert) or by
guarding/deriving a safe default: set stream_obj = comm_stream or
torch.cuda.current_stream(), use stream_obj.cuda_stream wherever cuda_stream is
needed, and only call compute_stream.wait_stream(stream_obj) when comm_stream
was provided or when appropriate; additionally, replace the hardcoded
dtype='bfloat16' in combine_kernel (around line 422) with a dtype string derived
from the tensor x.dtype used for recv_x (use x.dtype to construct the matching
kernel dtype string) so the kernel dtype matches the input tensor.
| StOp::StOp(Array<PrimExpr> args, BufferMap vmap) { | ||
| ObjectPtr<StOpNode> node = make_object<StOpNode>(); | ||
| node->dst = args[0]; | ||
| ICHECK(node->dst.as<CallNode>()) << "dst must be a call node"; | ||
| ICHECK(node->dst.as<CallNode>()->op.same_as(builtin::address_of())) | ||
| << "dst must be address_of op"; | ||
|
|
||
| node->value = args[1]; | ||
| node->sem = args[2].as<IntImm>().value()->value; | ||
| node->scope = args[3].as<IntImm>().value()->value; | ||
| node->na = args[4].as<IntImm>().value()->value; | ||
| node->dst_pe = args[5]; | ||
| data_ = std::move(node); | ||
| (void)vmap; | ||
| } |
There was a problem hiding this comment.
Add bounds checking for sem and scope indices.
The StOp and LdOp constructors extract sem and scope from arguments without validation. The Lower methods index sem_str[sem] (5 elements, valid indices 0-4) and scope_str[scope] (3 elements, valid indices 0-2) without bounds checks, causing undefined behavior if invalid values are passed.
🔎 Proposed fix
Add validation in both constructors:
StOp constructor (after line 258):
node->sem = args[2].as<IntImm>().value()->value;
node->scope = args[3].as<IntImm>().value()->value;
node->na = args[4].as<IntImm>().value()->value;
+ ICHECK(node->sem >= 0 && node->sem < 5) << "sem must be in range [0, 4], got " << node->sem;
+ ICHECK(node->scope >= 0 && node->scope < 3) << "scope must be in range [0, 2], got " << node->scope;
node->dst_pe = args[5];LdOp constructor (after line 327):
node->sem = args[2].as<IntImm>().value()->value;
node->scope = args[3].as<IntImm>().value()->value;
node->na = args[4].as<IntImm>().value()->value;
node->nc = args[5].as<IntImm>().value()->value;
+ ICHECK(node->sem >= 0 && node->sem < 5) << "sem must be in range [0, 4], got " << node->sem;
+ ICHECK(node->scope >= 0 && node->scope < 3) << "scope must be in range [0, 2], got " << node->scope;
node->src_pe = args[6];Also applies to: 316-331
🤖 Prompt for AI Agents
In src/op/remote_copy.cc around lines 248-262 (StOp constructor) and also around
lines 316-331 (LdOp constructor), add explicit bounds validation for the
extracted sem and scope indices: after assigning node->sem and node->scope,
assert that 0 <= sem <= 4 and 0 <= scope <= 2 (the valid ranges for sem_str and
scope_str respectively) using the project’s ICHECK/CHECK macro with a clear
error message, so invalid values fail fast instead of causing out-of-bounds
access in Lower.
| if __name__ == "__main__": | ||
| shape = (1024, 1024) | ||
| dtype = torch.float32 | ||
| host_tensor, device_tensor = create_mapped_tensor(shape, dtype) | ||
|
|
||
| # test meta-data | ||
| assert device_tensor.device.type == "cuda" | ||
| assert device_tensor.shape == shape, f"{device_tensor.shape=}" | ||
| assert device_tensor.dtype == dtype, f"{device_tensor.dtype=}" | ||
| assert torch.equal(host_tensor, device_tensor.cpu()), f"{host_tensor=}, {device_tensor=}" |
There was a problem hiding this comment.
Add CUDA availability check.
The test creates CUDA tensors without verifying that CUDA is available, which will cause cryptic failures on non-CUDA systems.
🔎 Proposed fix
if __name__ == "__main__":
+ if not torch.cuda.is_available():
+ print("CUDA not available, skipping test.")
+ exit(0)
+
shape = (1024, 1024)
dtype = torch.float32🤖 Prompt for AI Agents
In tilelang/distributed/testing/test_create_mapped_tensor.py around lines 4 to
13, the test unconditionally creates CUDA tensors; add a CUDA availability check
before creating tensors — e.g., import pytest and if not
torch.cuda.is_available(): pytest.skip("CUDA not available, skipping
CUDA-specific test") — so the test cleanly skips on non-CUDA systems
(alternatively guard with pytest.mark.skipif at the top of the test).
| # test modification | ||
| device_tensor.random_() | ||
| assert torch.equal(host_tensor, device_tensor.cpu()), f"{host_tensor=}, {device_tensor=}" |
There was a problem hiding this comment.
Test doesn't validate mapped memory behavior correctly.
The test uses device_tensor.cpu() to verify equality, which creates a copy from device to CPU. This doesn't test the key property of mapped memory: that host_tensor should reflect changes to device_tensor WITHOUT an explicit copy operation.
Additionally, there's no explicit synchronization after the random_() call. While PyTorch typically synchronizes implicitly, it's better to be explicit when testing low-level memory behavior.
🔎 Proposed fix
# test modification
device_tensor.random_()
+ torch.cuda.synchronize()
- assert torch.equal(host_tensor, device_tensor.cpu()), f"{host_tensor=}, {device_tensor=}"
+ # Verify mapped memory: host_tensor should reflect device changes directly
+ assert torch.equal(host_tensor, device_tensor.cpu()), \
+ f"Mapped memory test failed: {host_tensor=}, {device_tensor.cpu()=}"
+ # Additional check: verify host_tensor can be accessed without .cpu() call
+ # (In true mapped memory, host_tensor and device_tensor share the same underlying buffer)Note: The test should ideally verify that modifications are visible on the host side without needing to call .cpu(). Consider checking the data pointer addresses or using a more direct validation approach that confirms the memory mapping.
Committable suggestion skipped: line range outside the PR's diff.
🤖 Prompt for AI Agents
In tilelang/distributed/testing/test_create_mapped_tensor.py around lines 15-17,
the test incorrectly validates mapped memory by comparing host_tensor to
device_tensor.cpu() (which copies data) and lacks explicit synchronization;
change the test to (1) record the host tensor's data pointer and/or create a
direct NumPy/view into host memory to confirm it's the same backing buffer, (2)
perform device_tensor.random_(), then call the appropriate device
synchronization (e.g., torch.cuda.synchronize() or
torch.device(...).synchronize()) to flush device writes, and (3) assert that
host_tensor now reflects the changes without calling .cpu() (or assert data_ptrs
are equal and element-wise equality between host_tensor and the expected
values), ensuring the test verifies true mapped-memory behavior.
| def initialize( | ||
| self, | ||
| allocator: BaseAllocator, | ||
| stream: int = None, | ||
| ): | ||
| assert allocator.initialized(), "Allocator is not initialized" | ||
| result = self.adapter.lib.init_table( | ||
| ctypes.c_void_p(allocator.table.data_ptr()), allocator.table_size) | ||
| ctypes.c_void_p(allocator.table.data_ptr()), allocator.table_size, | ||
| ctypes.c_void_p(stream) if stream is not None else ctypes.c_void_p(0)) | ||
| if result != 0: |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Read the specific file and lines mentioned in the review
if [ -f "tilelang/jit/kernel.py" ]; then
sed -n '407,420p' tilelang/jit/kernel.py
else
echo "File not found at expected location"
find . -name "kernel.py" -type f 2>/dev/null | head -5
fiRepository: tile-ai/tilescale
Length of output: 673
🏁 Script executed:
#!/bin/bash
# Search for initialize() calls
echo "=== initialize() calls ==="
rg 'initialize\(' --type py -A 2 -B 2 | head -50
echo ""
echo "=== dispatch/combine references ==="
rg 'dispatch|combine' --type py tilelang/jit/ | grep -E '(def|stream)' | head -20Repository: tile-ai/tilescale
Length of output: 4027
🏁 Script executed:
#!/bin/bash
# Check the combine.py file more closely for stream usage
echo "=== Combine path context ==="
grep -n "comm_stream\|cuda_stream\|initialize" examples/distributed/deepseek_deepep/intranode/combine.py | head -20
echo ""
echo "=== Check if there are dispatch patterns with custom streams ==="
rg 'dispatch.*stream|stream.*dispatch' --type py -B 2 -A 2 | head -40
echo ""
echo "=== Check run_once and other kernel execution methods ==="
sed -n '420,450p' tilelang/jit/kernel.pyRepository: tile-ai/tilescale
Length of output: 1525
🏁 Script executed:
#!/bin/bash
# Check Python version requirements and project config
echo "=== Python version target ==="
for f in pyproject.toml setup.py setup.cfg; do
if [ -f "$f" ]; then
echo "Found: $f"
grep -i "python" "$f" | head -10
fi
done
echo ""
echo "=== Check if Optional is already imported in kernel.py ==="
head -30 tilelang/jit/kernel.py | grep -E "^import|^from"Repository: tile-ai/tilescale
Length of output: 1030
🏁 Script executed:
#!/bin/bash
# Final verification: check if there are other kernel execution methods that accept streams
echo "=== Kernel execution methods ==="
rg 'def (run|call|execute|__call__)' tilelang/jit/kernel.py -A 5 | head -50Repository: tile-ai/tilescale
Length of output: 528
🏁 Script executed:
#!/bin/bash
# Check if __call__ also accepts stream parameter
echo "=== __call__ method full signature ==="
sed -n '350,380p' tilelang/jit/kernel.py
echo ""
echo "=== Check what the kernel actually does with init_table ==="
rg 'init_table|get_last_error' tilelang/jit/kernel.py -B 2 -A 2Repository: tile-ai/tilescale
Length of output: 1641
🏁 Script executed:
#!/bin/bash
# Check if initialize has return type annotation
echo "=== initialize method full signature and surrounding code ==="
sed -n '405,425p' tilelang/jit/kernel.py
echo ""
echo "=== Verify typing imports and what's available ==="
head -20 tilelang/jit/kernel.pyRepository: tile-ai/tilescale
Length of output: 1774
Type the stream parameter as optional and ensure consistent usage
The initialize method has two issues:
- Typing
Thestreamparameter should be explicitly typed as optional. Currently annotated asstream: int = None, which is semantically incorrect:
from typing import Optional
def initialize(
self,
allocator: BaseAllocator,
stream: Optional[int] = None,
) -> None:This satisfies type checkers, supports Python 3.8+, and clarifies that None is a valid value.
- CUDA Stream Semantics
Whenstream is None, the code passesctypes.c_void_p(0), routinginit_tableto the default stream. However, callers in the dispatch/combine paths (e.g.,examples/distributed/deepseek_deepep/intranode/combine.py) explicitly pass custom streams viastream=comm_stream.cuda_stream. Ifinit_tableruns on the default stream (0) while subsequent kernels run oncomm_stream, there is no guaranteed ordering between them unless the default stream synchronization semantics apply. To maintain proper stream ordering, callers should consistently pass the same stream toinitialize()that they use for subsequent kernel launches, or rely on explicit stream synchronization barriers.
🧰 Tools
🪛 Ruff (0.14.10)
410-410: PEP 484 prohibits implicit Optional
Convert to Optional[T]
(RUF013)
🤖 Prompt for AI Agents
In tilelang/jit/kernel.py around lines 407 to 416, change the initialize
signature to explicitly type stream as Optional[int] and add a return type
(e.g., -> None) so type checkers know None is allowed; ensure the function
continues to pass a null pointer when stream is None but document/require
callers to pass the same CUDA stream they use for subsequent kernel launches (or
synchronize explicitly) so init_table and later kernels run on the same stream
to preserve ordering.
| std::pair<torch::Tensor, torch::Tensor> | ||
| create_host_device_tensor(const std::vector<int64_t> &shape, | ||
| c10::ScalarType dtype) { | ||
| size_t elem_size = at::elementSize(dtype); | ||
| int64_t numel = 1; | ||
| for (int64_t s : shape) | ||
| numel *= s; | ||
|
|
||
| size_t bytes = numel * elem_size; | ||
|
|
||
| void *host_ptr = nullptr; | ||
| CUDA_CHECK(cudaHostAlloc(&host_ptr, bytes, cudaHostAllocMapped)); | ||
|
|
||
| void *device_ptr = nullptr; | ||
| CUDA_CHECK(cudaHostGetDevicePointer(&device_ptr, host_ptr, 0)); | ||
|
|
||
| auto host_tensor = torch::from_blob( | ||
| host_ptr, shape, torch::TensorOptions().dtype(dtype).device(torch::kCPU)); | ||
|
|
||
| auto device_tensor = torch::from_blob( | ||
| device_ptr, shape, | ||
| torch::TensorOptions().dtype(dtype).device(torch::kCUDA)); | ||
|
|
||
| return std::make_pair(host_tensor, device_tensor); | ||
| } No newline at end of file |
There was a problem hiding this comment.
Memory leak: pinned memory is never freed.
torch::from_blob does not take ownership of the memory. When both tensors are destroyed, cudaHostFree is never called, causing a memory leak. You need to provide a custom deleter to free the pinned memory.
Additionally:
- Missing overflow checks (lines 93-95): Unlike
tensor_from_ptrwhich usessafe_mul_int64, this function doesn't check for overflow when computingnumel. - Missing negative dimension check: Shape dimensions aren't validated.
- Zero-size edge case (line 100): Behavior of
cudaHostAllocwithbytes=0should be handled.
🔎 Proposed fix with custom deleter and validation
std::pair<torch::Tensor, torch::Tensor>
create_host_device_tensor(const std::vector<int64_t> &shape,
c10::ScalarType dtype) {
size_t elem_size = at::elementSize(dtype);
int64_t numel = 1;
- for (int64_t s : shape)
- numel *= s;
+ for (int64_t s : shape) {
+ if (s < 0)
+ throw std::runtime_error("Negative dimension in shape");
+ numel = safe_mul_int64(numel, s);
+ }
size_t bytes = numel * elem_size;
+ if (bytes == 0) {
+ auto opts_cpu = torch::TensorOptions().dtype(dtype).device(torch::kCPU);
+ auto opts_cuda = torch::TensorOptions().dtype(dtype).device(torch::kCUDA);
+ return std::make_pair(torch::empty(shape, opts_cpu),
+ torch::empty(shape, opts_cuda));
+ }
+
void *host_ptr = nullptr;
CUDA_CHECK(cudaHostAlloc(&host_ptr, bytes, cudaHostAllocMapped));
void *device_ptr = nullptr;
CUDA_CHECK(cudaHostGetDevicePointer(&device_ptr, host_ptr, 0));
+ auto deleter = [](void *ptr) {
+ cudaError_t err = cudaHostFree(ptr);
+ if (err != cudaSuccess) {
+ std::fprintf(stderr, "cudaHostFree failed: %s\n",
+ cudaGetErrorString(err));
+ }
+ };
+
auto host_tensor = torch::from_blob(
- host_ptr, shape, torch::TensorOptions().dtype(dtype).device(torch::kCPU));
+ host_ptr, shape, deleter,
+ torch::TensorOptions().dtype(dtype).device(torch::kCPU));
auto device_tensor = torch::from_blob(
device_ptr, shape,
torch::TensorOptions().dtype(dtype).device(torch::kCUDA));
return std::make_pair(host_tensor, device_tensor);
}Note: Only the host tensor should have the deleter since both tensors share the same underlying pinned memory allocation.
Committable suggestion skipped: line range outside the PR's diff.
There was a problem hiding this comment.
Actionable comments posted: 4
♻️ Duplicate comments (26)
src/transform/vectorize_loop.cc (1)
796-813: Bug: Extracts original value, discarding the reinterpret.Line 800 extracts
args[0]from thereinterpretcall created on Line 797, which retrievesnew_value(the original un-reinterpreted value). This defeats the purpose of the reinterpret operation.Compare with the correct implementation in
src/transform/common/loop_vectorization_utils.h(lines 485-491), which usesvec_valuedirectly.🔎 Proposed fix
// Reinterpret the value to vector type (e.g., int4 for 8xbf16) PrimExpr vec_value = Call(vec_dtype, builtin::reinterpret(), {new_value}); - // A trick to get the lvalue of the vectorized value - PrimExpr vec_value_lvalue = vec_value.as<CallNode>()->args[0]; - // Build new args with base addresses and reinterpreted value Array<PrimExpr> new_args; new_args.push_back(func_name); new_args.push_back(new_addr); - new_args.push_back(vec_value_lvalue); + new_args.push_back(vec_value); // Copy remaining args (sem, scope, etc.).github/workflows/ci.yml (1)
85-86: Remove unused MARKER computation.Lines 85-86 compute a
MARKERvariable that is never used because venv reuse is intentionally disabled (per line 87). This dead code should be removed for clarity.🔎 Proposed fix to remove dead code
- name: Ensure venv (local & persistent) run: | set -e - REQS_HASH=$(cat requirements-test.txt 2>/dev/null || true) - MARKER="${{ runner.tool_cache }}/.venv_marker_${{ env.PYTHON_VERSION }}_${REQS_HASH:0:8}" # NOTE(wt): We disable the venv reuse for now to allow installing DeepEP # echo "venv stale or missing – recreating" rm -rf "${{ runner.tool_cache }}/${{ env.VENV_DIR }}"tilelang/jit/kernel.py (1)
407-418: The typing issue from the previous review remains unresolved.As noted in the prior review, the
streamparameter should be typed asOptional[int]rather thanint = Noneto satisfy PEP 484 and type checkers. You'll need to addOptionalto the imports on line 2.The stream-passing logic itself (lines 414-415) is correctly implemented—it conditionally converts to a pointer or passes null—but the signature still needs the type annotation fix.
🔧 Quick fix
Update the import and signature:
-from typing import Any, Callable, Literal +from typing import Any, Callable, Literal, Optionaldef initialize( self, allocator: BaseAllocator, - stream: int = None, -): + stream: Optional[int] = None, +) -> None:src/tl_templates/cuda/copy.h (3)
258-270: Critical: Same alignment and size validation issues as above (unresolved from past review).This overload also has the same critical issues:
- Missing alignment verification for 16-byte boundaries
- Missing size validation for
N * sizeof(dtype_t) % 16 == 0- The
reinterpret_castat line 261 does not guarantee alignment🔎 Proposed fix
template <int N, int UNROLL_FACTOR, bool enable_aggressive_vectorize = false, typename dtype_t> TL_DEVICE void cp_warp(dtype_t *const dst_addr, uint64_t src_addr_uint64) { const dtype_t *src_addr = reinterpret_cast<const dtype_t *>(src_addr_uint64); if constexpr (enable_aggressive_vectorize) { + static_assert((N * sizeof(dtype_t)) % 16 == 0, + "Aggressive vectorization requires N*sizeof(dtype_t) to be a multiple of 16"); + if ((uintptr_t)dst_addr % 16 != 0 || (uintptr_t)src_addr % 16 != 0) { + cp_warp_impl<N, UNROLL_FACTOR>(dst_addr, src_addr); + return; + } int4 *__restrict__ dst_addr_int4 = (int4 *)dst_addr; const int4 *__restrict__ src_addr_int4 = (const int4 *)src_addr; constexpr int N_int4 = sizeof(dtype_t) * N / 16; cp_warp_impl<N_int4, UNROLL_FACTOR>(dst_addr_int4, src_addr_int4); } else { cp_warp_impl<N, UNROLL_FACTOR>(dst_addr, src_addr); } }
221-241: Critical: Missing alignment and size validation for aggressive vectorization (unresolved from past review).The
enable_aggressive_vectorizepath casts pointers toint4*without verifying:
- Both source and destination are 16-byte aligned
N * sizeof(dtype_t)is a multiple of 16This can lead to:
- Misaligned memory access crashes if pointers aren't 16-byte aligned
- Silent data truncation if the size isn't a multiple of 16 (the integer division
N_int4 = sizeof(dtype_t) * N / 16will truncate)Compare with
nvshmem_cp_threadgroup(lines 287-302), which performs runtime alignment checks before attempting vectorized copies.🔎 Proposed fix with validation
template <int N, int UNROLL_FACTOR, bool enable_aggressive_vectorize = false, typename dtype_t> TL_DEVICE void cp_warp(dtype_t const *const dst_addr, dtype_t const *const src_addr) { if constexpr (enable_aggressive_vectorize) { + static_assert((N * sizeof(dtype_t)) % 16 == 0, + "Aggressive vectorization requires N*sizeof(dtype_t) to be a multiple of 16"); + // Runtime alignment check + if ((uintptr_t)dst_addr % 16 != 0 || (uintptr_t)src_addr % 16 != 0) { + // Fall back to non-vectorized path if alignment requirements not met + cp_warp_impl<N, UNROLL_FACTOR>(dst_addr, src_addr); + return; + } int4 *__restrict__ dst_addr_int4 = (int4 *)dst_addr; const int4 *__restrict__ src_addr_int4 = (const int4 *)src_addr; constexpr int N_int4 = sizeof(dtype_t) * N / 16; cp_warp_impl<N_int4, UNROLL_FACTOR>(dst_addr_int4, src_addr_int4); } else { cp_warp_impl<N, UNROLL_FACTOR>(dst_addr, src_addr); } }
243-256: Critical: Same alignment and size validation issues as above (unresolved from past review).This overload has the same critical issues as the previous
cp_warpvariant (lines 221-241):
- Missing alignment verification for 16-byte boundaries
- Missing size validation for
N * sizeof(dtype_t) % 16 == 0- The
reinterpret_castat line 247 does not guarantee alignment🔎 Proposed fix
template <int N, int UNROLL_FACTOR, bool enable_aggressive_vectorize = false, typename dtype_t> TL_DEVICE void cp_warp(uint64_t dst_addr_uint64, dtype_t const *const src_addr) { dtype_t *dst_addr = reinterpret_cast<dtype_t *>(dst_addr_uint64); if constexpr (enable_aggressive_vectorize) { + static_assert((N * sizeof(dtype_t)) % 16 == 0, + "Aggressive vectorization requires N*sizeof(dtype_t) to be a multiple of 16"); + if ((uintptr_t)dst_addr % 16 != 0 || (uintptr_t)src_addr % 16 != 0) { + cp_warp_impl<N, UNROLL_FACTOR>(dst_addr, src_addr); + return; + } int4 *__restrict__ dst_addr_int4 = (int4 *)dst_addr; const int4 *__restrict__ src_addr_int4 = (const int4 *)src_addr; constexpr int N_int4 = sizeof(dtype_t) * N / 16; cp_warp_impl<N_int4, UNROLL_FACTOR>(dst_addr_int4, src_addr_int4); } else { cp_warp_impl<N, UNROLL_FACTOR>(dst_addr, src_addr); } }examples/distributed/deepseek_deepep/deepep.md (3)
3-3: Wrap bare URL in markdown link syntax.The bare URL should be formatted as a markdown link for better readability.
26-28: Remove stray character "z".Line 27 contains a lone "z" which appears to be an accidental keystroke.
146-149: Add language identifier to code block.The shell command block should specify
bashas the language identifier.tilelang/language/reduce.py (1)
278-291: Docstring parameter name mismatch.The docstring references
xbut the actual parameter is namedvalue.src/target/codegen_cuda.cc (1)
2365-2366: Inconsistent handling:loop_continuevsloop_break.
loop_continueoutputs toos(expression stream) whileloop_break(lines 1522-1524) usesthis->PrintIndent()andthis->stream(statement stream). Sincecontinueis a statement, it should follow the same pattern asbreak.examples/distributed/deepseek_deepep/intranode/get_dispatch_layout.py (1)
17-60: Return value count mismatch with docstring and type annotation.The docstring describes 4 return values including
num_tokens_per_rdma_rank, and the type annotation specifiesTuple[torch.Tensor, torch.Tensor | None, torch.Tensor, torch.Tensor], but the function only returns 3 values. Either update the docstring/type annotation to reflect 3 values, or addNoneas the second return value.examples/distributed/deepseek_deepep/buffer.py (1)
121-127: Consider parameterizing the channel buffer dtype.
channel_x_buffersis allocated withdtype=torch.bfloat16, but dispatch and combine operations may receive tensors of different dtypes. Consider adding adtypeparameter to__init__to avoid implicit type conversions.🔎 Proposed fix
def __init__(self, group: dist.ProcessGroup, num_nvl_bytes: int, num_topk: int, num_experts: int, hidden: int, + dtype: torch.dtype = torch.bfloat16, dispatch_cfg: Optional[Config] = None, combine_cfg: Optional[Config] = None): ... + self.dtype = dtype ... def _pre_alloc_symm_buffers_intranode(self): ... channel_x_buffers = tilelang.tensor([ self.num_channels, self.num_ranks, self.dispatch_cfg.num_max_nvl_chunked_recv_tokens, self.hidden ], - dtype=torch.bfloat16, + dtype=self.dtype, device='cuda', allocator=self._allocator)src/tl_templates/cuda/sync.h (1)
153-154: Fix typo in comment.Line 153 contains "optinal" which should be "optional".
🔎 Proposed fix
-// Sync blocks at a system-level barrier with an optinal fence +// Sync blocks at a system-level barrier with an optional fenceexamples/distributed/deepseek_deepep/deepep_utils.py (1)
220-248: Use&&instead ofandin C++ code.Line 233 uses the Python keyword
andin C++ code. While some compilers acceptandas an alternative token,&&is the standard C++ operator and ensures better portability.🔎 Proposed fix
- for (int i = 0; i < num_local_experts and ready; ++i) + for (int i = 0; i < num_local_experts && ready; ++i) ready &= expert_ptr[i] >= 0;src/op/remote_copy.cc (2)
269-302: Add bounds checking forsemandscopearray indices.The
Lowermethod indexessem_str[sem]at line 282 andscope_str[scope]at line 282 without validating thatsemis within [0, 4] andscopeis within [0, 2]. Invalid values cause undefined behavior.🔎 Proposed fix
Add validation in the
StOpconstructor (after line 258):node->sem = args[2].as<IntImm>().value()->value; node->scope = args[3].as<IntImm>().value()->value; node->na = args[4].as<IntImm>().value()->value; + ICHECK(node->sem >= 0 && node->sem < 5) << "sem must be in range [0, 4], got " << node->sem; + ICHECK(node->scope >= 0 && node->scope < 3) << "scope must be in range [0, 2], got " << node->scope; node->dst_pe = args[5];
338-371: Add bounds checking forsemandscopearray indices.The
Lowermethod indexessem_str[sem]at line 351 andscope_str[scope]at line 351 without validating bounds. This causes undefined behavior ifsem >= 5orscope >= 3.🔎 Proposed fix
Add validation in the
LdOpconstructor (after line 327):node->sem = args[2].as<IntImm>().value()->value; node->scope = args[3].as<IntImm>().value()->value; node->na = args[4].as<IntImm>().value()->value; node->nc = args[5].as<IntImm>().value()->value; + ICHECK(node->sem >= 0 && node->sem < 5) << "sem must be in range [0, 4], got " << node->sem; + ICHECK(node->scope >= 0 && node->scope < 3) << "scope must be in range [0, 2], got " << node->scope; node->src_pe = args[6];src/op/sync.cc (1)
153-180: Add bounds checking forrelationarray index.The
relation_strarray has 6 elements (indices 0-5), butrelationis used at line 161 without validation. Invalid values cause undefined behavior.🔎 Proposed fix
Add validation in the
WaitOpconstructor (after line 140):ObjectPtr<WaitOpNode> node = make_object<WaitOpNode>(); node->relation = args[0].as<IntImmNode>()->value; + ICHECK(node->relation >= 0 && node->relation < 6) + << "relation must be in range [0, 5], got " << node->relation; node->addr = args[1];examples/distributed/deepseek_deepep/intranode/combine.py (3)
76-95: Guardcomm_streamaccess or make it required.Both lines 87 and 94 access
comm_stream.cuda_stream, but the parameter defaults toNoneat line 85. If callers rely on the default, this will raiseAttributeError.🔎 Proposed fix
def cached_notify_combine( num_ranks, num_sms, ##### symm buffers ##### send_head: torch.Tensor, channel_head_idx: torch.Tensor, channel_tail_idx: torch.Tensor, barrier_signal: torch.Tensor, allocator, comm_stream=None): + if comm_stream is None: + comm_stream = torch.cuda.current_stream() kernel = cached_notify_combine_kernel(num_ranks, num_sms) kernel.initialize(allocator=allocator, stream=comm_stream.cuda_stream) kernel( send_head, channel_head_idx, channel_tail_idx, barrier_signal, stream=comm_stream.cuda_stream, skip_tensor_validation=True)
415-422: Derive kernel dtype from input tensor.The kernel is initialized with hardcoded
dtype='bfloat16'at line 422, butrecv_xat line 412 usesx.dtype. This mismatch can cause incorrect behavior ifxuses a different dtype.🔎 Proposed fix
kernel = combine_kernel( num_ranks, config.num_max_nvl_chunked_send_tokens, config.num_max_nvl_chunked_recv_tokens, hidden, num_topk, config.num_sms, - dtype='bfloat16') + dtype=str(x.dtype).split('.')[-1])
382-443: Guardcomm_streamaccess.Lines 423 and 439 access
comm_stream.cuda_stream, and line 442 callscompute_stream.wait_stream(comm_stream), butcomm_streamdefaults toNoneat line 389.🔎 Proposed fix
topk_weights, comm_stream=None): + if comm_stream is None: + comm_stream = torch.cuda.current_stream() assert handle is not Nonetilelang/language/builtin.py (1)
765-768: Docstring incorrectly lists "release" as valid for loads.The docstring at line 768 mentions
"'weak', 'volatile', 'acquire', 'release', or 'relaxed'"but the assertion at line 767 correctly excludes"release"for loads (release semantics don't apply to load operations). Update the docstring to match the assertion.🔎 Proposed fix
- ], "Semantic must be one of 'weak', 'volatile', 'acquire', 'release', or 'relaxed'." + ], "Semantic must be one of 'weak', 'volatile', 'acquire', or 'relaxed'."examples/distributed/deepseek_deepep/intranode/dispatch.py (2)
901-902: Critical: recv_src_idx incorrectly overwritten in cached mode.In cached mode (when
handle is not None),recv_src_idxis unpacked from the handle at line 854 and used at line 899 to determinenum_recv_tokens. However, at line 902, a new empty tensor is unconditionally allocated, overwriting the cached one. This causes the kernel to receive an emptyrecv_src_idxin cached mode instead of reusing the existing one from the handle.🔎 Proposed fix
recv_x = torch.empty((num_recv_tokens, hidden), dtype=x.dtype, device='cuda') - recv_src_idx = torch.empty((num_recv_tokens,), dtype=torch.int32, device='cuda') + if handle is None: + recv_src_idx = torch.empty((num_recv_tokens,), dtype=torch.int32, device='cuda') + # else: recv_src_idx already unpacked from handle at line 854 if handle is None: recv_topk_idx = torch.empty((num_recv_tokens, num_topk), dtype=torch.int64, device='cuda')
915-917: Derive dtype from input tensor instead of hardcoding.The kernel is initialized with hardcoded
'bfloat16'regardless of the actual dtype of input tensorx. This prevents support for other precisions (e.g., float16, float32). The dtype should be derived fromx.dtype.🔎 Proposed fix
+ dtype_str = str(x.dtype).split('.')[-1] # e.g., 'torch.bfloat16' -> 'bfloat16' if handle is None: kernel = dispatch_kernel(num_ranks, config.num_max_nvl_chunked_send_tokens, config.num_max_nvl_chunked_recv_tokens, hidden, num_topk, - num_experts, config.num_sms, 'bfloat16') + num_experts, config.num_sms, dtype_str)Apply the same change at line 950 for
cached_dispatch_kernel.src/op/remote_copy.h (2)
36-38: Fix documentation typo.The comment contains a typo: "effctive" should be "effective".
🔎 Proposed fix
bool enable_aggressive_vectorize; ///< Whether to enable aggressive - ///< vectorization, only effctive for + ///< vectorization, only effective for ///< warp-scope
128-130: Fix documentation typo.Same typo as in
PutOpNode: "effctive" should be "effective".🔎 Proposed fix
bool enable_aggressive_vectorize; ///< Whether to enable aggressive - ///< vectorization, only effctive for + ///< vectorization, only effective for ///< warp-scope
🧹 Nitpick comments (10)
tilelang/language/print.py (2)
136-142: Consider enhancing the docstring for consistency.The implementation is correct and follows the established patterns in this file. However, the docstring could be expanded to match the detail level of other functions (e.g.,
print_varat lines 14-24), including a Parameters section and Returns section.🔎 Suggested docstring enhancement
@macro def print_msg(msg: str) -> tir.PrimExpr: """ Prints a message for debugging purposes. + + Parameters: + msg (str): The message to be printed. + + Returns: + tir.PrimExpr: The TIR expression for the debug print operation. """ tir.call_extern("handle", "debug_print_msg", msg)
224-227: The error message is correct; consider the linter suggestion as a low-priority style improvement.The error message accurately reflects the supported types including
None. The Ruff linter suggests extracting long messages into a custom exception class for better maintainability, but this is a minor style preference rather than a functional issue.Based on static analysis hints (Ruff TRY003).
.github/workflows/ci.yml (1)
121-121: Verify temporary test disablement is tracked.The test filtering (lines 121, 147) temporarily disables several tests with comments indicating this is a workaround. Ensure there are corresponding issues or TODO items to track re-enabling these tests once the underlying problems (timeouts, etc.) are resolved.
If you'd like, I can help create GitHub issues to track the re-enablement of these disabled tests:
- Examples:
sink,vs_sparse- Tests:
tilelibrary_gemm,jit_gemm_ctypesAlso applies to: 147-147
tilelang/utils/ts_ext/exception.h (1)
7-17: LGTM – Exception class correctly implements std::exception interface.The class properly inherits from
std::exception, stores the formatted message, and returns it viawhat(). The design is sound and follows standard C++ exception patterns.Optional: More efficient string formatting
The constructor builds the message with multiple string concatenations, creating temporary objects. Since this is only called in error paths, the performance impact is negligible, but you could consider using
std::ostringstreamor C++20'sstd::formatfor cleaner code:#include <sstream> TSException(const char *name, const char *file, int line, const std::string &error) { std::ostringstream oss; oss << "Failed: " << name << " error " << file << ":" << line << " '" << error << "'"; message = oss.str(); }tilelang/utils/ts_ext/ipc_ops.cpp (1)
78-92: Migration to TS_HOST_ASSERT is correct; consider RAII for IPC handles.The assertions properly validate:
- Size consistency (line 78)
- Optional value presence (line 84)
- Handle size correctness (line 86)
However, there's a pre-existing resource management concern: if
cudaIpcOpenMemHandlesucceeds for some iterations but a subsequent assertion or CUDA call fails, the successfully opened IPC handles won't be closed viacudaIpcCloseMemHandle. With exception-based error handling, this leak becomes more likely.Consider using RAII for IPC handle cleanup
Wrap opened IPC handles in a RAII guard to ensure cleanup on early exit:
// Helper to close IPC handle on scope exit struct IpcHandleGuard { void* ptr; bool active; IpcHandleGuard(void* p) : ptr(p), active(true) {} ~IpcHandleGuard() { if (active && ptr) cudaIpcCloseMemHandle(ptr); } void release() { active = false; } }; std::vector<std::unique_ptr<IpcHandleGuard>> guards; for (int i = 0, offset = rdma_rank * num; i < num; ++i) { // ... validation ... if (offset + i != rank) { std::memcpy(ipc_handles[i].reserved, s.data(), CUDA_IPC_HANDLE_SIZE); CUDA_CHECK(cudaIpcOpenMemHandle(&buffer_ptrs[i], ipc_handles[i], cudaIpcMemLazyEnablePeerAccess)); guards.push_back(std::make_unique<IpcHandleGuard>(buffer_ptrs[i])); } } // Success - release guards for (auto& g : guards) g->release();examples/distributed/sp_ag_attention_intra_node.py (1)
12-20: Consider renaming the kernel function for consistency.The kernel function is named
barrier_all_blocks_sys_kernelbut now usesT.barrier_blocks. While functionally correct, consider renaming tobarrier_blocks_kernelfor consistency with the updated API.🔎 Proposed fix
@tilelang.jit -def barrier_all_blocks_sys_kernel(num_local_rank,): +def barrier_blocks_kernel(num_local_rank,): @T.prim_func def main(barrier: T.Tensor((num_local_rank), "int32"),): with T.Kernel(1, threads=32): T.barrier_blocks(barrier) return mainNote: This would also require updating the call site at line 612:
def barrier_all_on_stream(barrier: torch.Tensor, stream: torch.cuda.Stream, world_size: int): - barrier_all_blocks_sys_func = barrier_all_blocks_sys_kernel(world_size) + barrier_all_blocks_sys_func = barrier_blocks_kernel(world_size) barrier_all_blocks_sys_func(barrier, stream=stream.cuda_stream)tilelang/language/reduce.py (1)
278-355: Minor docstring inconsistencies across warp reduction functions.Several docstrings have double spaces before "register" (e.g., "provides a register"). This appears in
warp_reduce_sum,warp_reduce_max,warp_reduce_min,warp_reduce_bitand, andwarp_reduce_bitor.🔎 Example fix for warp_reduce_sum
This function reduces a value across all threads in a warp using shuffle operations. - Each thread provides a register `value`, and after the reduction, all threads + Each thread provides a register `value`, and after the reduction, all threads will have the sum of all values across the warp.tilelang/distributed/utils.py (1)
231-243: Consider simplifying the event synchronization pattern.The current pattern calls
start_event.wait()andstop_event.wait()followed bysynchronize(). The.wait()methods make the current CUDA stream wait for the events (not the CPU). While functionally correct, the more common pattern is:stop_event.record() stop_event.synchronize() # CPU waits for the eventThe current approach works but the explicit
wait()calls beforesynchronize()are redundant sincesynchronize()already ensures completion.tilelang/language/builtin.py (1)
514-514: Fix implicit Optional in type hints.The type annotations use implicit
Optional(e.g.,int | PrimExpr = None), which is prohibited by PEP 484. Explicitly wrap nullable types withOptional.🔎 Proposed fix
-def sync_threads(barrier_id: int | PrimExpr = None, arrive_count: int = None): +def sync_threads(barrier_id: Optional[int | PrimExpr] = None, arrive_count: Optional[int] = None):Based on static analysis hints.
examples/distributed/deepseek_deepep/intranode/dispatch.py (1)
13-13: Remove unusednoqadirective.The
# noqa: F403directive is unnecessary since this is an explicit import, not a wildcard import. The F403 warning only applies tofrom module import *statements.🔎 Proposed fix
-from deepep_utils import Config, ep_ext # noqa: F403 +from deepep_utils import Config, ep_extBased on static analysis hints.
📜 Review details
Configuration used: defaults
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (56)
.github/workflows/ci.yml.gitmodules3rdparty/DeepEPbenchmark/distributed/deepep/deepep.mdbenchmark/distributed/deepep/intranode/get_dispatch_layout.pyexamples/distributed/deepseek_deepep/buffer.pyexamples/distributed/deepseek_deepep/deepep.mdexamples/distributed/deepseek_deepep/deepep_utils.pyexamples/distributed/deepseek_deepep/intranode/combine.pyexamples/distributed/deepseek_deepep/intranode/dispatch.pyexamples/distributed/deepseek_deepep/intranode/example_intranode.pyexamples/distributed/deepseek_deepep/intranode/get_dispatch_layout.pyexamples/distributed/deepseek_deepep/intranode/test_intranode.pyexamples/distributed/primitives/example_put_warp.pyexamples/distributed/primitives/example_remote_st.pyexamples/distributed/sp_ag_attention_intra_node.pypyproject.tomlsrc/op/builtin.ccsrc/op/builtin.hsrc/op/remote_copy.ccsrc/op/remote_copy.hsrc/op/sync.ccsrc/op/sync.hsrc/target/codegen_cuda.ccsrc/tl_templates/cuda/common.hsrc/tl_templates/cuda/copy.hsrc/tl_templates/cuda/debug.hsrc/tl_templates/cuda/ldst.hsrc/tl_templates/cuda/reduce.hsrc/tl_templates/cuda/sync.hsrc/transform/common/loop_vectorization_utils.hsrc/transform/loop_vectorize.ccsrc/transform/storage_access.ccsrc/transform/vectorize_loop.cctesting/python/language/test_tilelang_language_elect.pytesting/python/language/test_tilelang_language_ldst_options.pytesting/python/language/test_tilelang_language_vote.pytesting/python/language/test_tilelang_language_warp_reduce.pytilelang/distributed/install_deepep.shtilelang/distributed/testing/sync/test_barrierall_sys.pytilelang/distributed/testing/test_create_mapped_tensor.pytilelang/distributed/utils.pytilelang/jit/adapter/wrapper.pytilelang/jit/kernel.pytilelang/language/__init__.pytilelang/language/builtin.pytilelang/language/distributed/common.pytilelang/language/print.pytilelang/language/reduce.pytilelang/utils/ts_ext/__init__.pytilelang/utils/ts_ext/exception.htilelang/utils/ts_ext/ipc_ops.cpptilelang/utils/ts_ext/setup.pytilelang/utils/ts_ext/tensor.cpptilelang/utils/ts_ext/ts_ext_bindings.cpptilelang/utils/ts_ext/ts_ext_ops.h
💤 Files with no reviewable changes (2)
- benchmark/distributed/deepep/deepep.md
- benchmark/distributed/deepep/intranode/get_dispatch_layout.py
✅ Files skipped from review due to trivial changes (1)
- 3rdparty/DeepEP
🚧 Files skipped from review as they are similar to previous changes (14)
- src/tl_templates/cuda/common.h
- testing/python/language/test_tilelang_language_vote.py
- tilelang/distributed/install_deepep.sh
- tilelang/utils/ts_ext/tensor.cpp
- examples/distributed/primitives/example_put_warp.py
- src/tl_templates/cuda/debug.h
- testing/python/language/test_tilelang_language_ldst_options.py
- testing/python/language/test_tilelang_language_warp_reduce.py
- src/op/sync.h
- tilelang/utils/ts_ext/ts_ext_bindings.cpp
- examples/distributed/deepseek_deepep/intranode/test_intranode.py
- tilelang/distributed/testing/test_create_mapped_tensor.py
- pyproject.toml
- tilelang/distributed/testing/sync/test_barrierall_sys.py
🧰 Additional context used
🧬 Code graph analysis (26)
src/tl_templates/cuda/reduce.h (3)
src/op/reduce.cc (1)
ReduceOp(24-33)src/op/reduce.h (1)
ReduceOp(146-151)tilelang/language/reduce.py (5)
warp_reduce_sum(278-291)warp_reduce_max(294-307)warp_reduce_min(310-323)warp_reduce_bitand(326-339)warp_reduce_bitor(342-355)
src/transform/loop_vectorize.cc (2)
src/transform/loop_vectorize_dynamic.cc (16)
node(80-85)node(80-80)node(92-96)node(92-92)node(98-112)node(98-98)node(114-120)node(114-114)node(122-125)node(122-122)node(127-135)node(127-127)node(263-266)node(263-263)node(283-286)node(283-283)tilelang/language/tir/op.py (1)
address_of(464-480)
tilelang/utils/ts_ext/ts_ext_ops.h (1)
tilelang/utils/ts_ext/tensor.cpp (2)
create_host_device_tensor(89-113)create_host_device_tensor(90-91)
tilelang/distributed/utils.py (1)
tilelang/utils/ts_ext/tensor.cpp (2)
create_host_device_tensor(89-113)create_host_device_tensor(90-91)
src/transform/vectorize_loop.cc (2)
tilelang/language/tir/op.py (3)
call_extern(173-195)address_of(464-480)reinterpret(1899-1918)src/transform/common/loop_vectorization_utils.h (1)
MutateTlLdStExpr_(389-521)
src/transform/common/loop_vectorization_utils.h (1)
tilelang/language/tir/op.py (3)
address_of(464-480)reinterpret(1899-1918)call_extern(173-195)
examples/distributed/sp_ag_attention_intra_node.py (1)
tilelang/language/builtin.py (1)
barrier_blocks(615-623)
tilelang/language/__init__.py (1)
tilelang/language/reduce.py (5)
warp_reduce_sum(278-291)warp_reduce_max(294-307)warp_reduce_min(310-323)warp_reduce_bitand(326-339)warp_reduce_bitor(342-355)
examples/distributed/primitives/example_remote_st.py (5)
tilelang/distributed/utils.py (1)
init_dist(40-62)tilelang/env.py (1)
disable_cache(286-287)tilelang/language/distributed/common.py (2)
get_rank(10-13)get_num_ranks(16-19)tilelang/utils/allocator.py (1)
get_allocator(236-248)tilelang/utils/tensor.py (1)
tensor(43-56)
tilelang/language/print.py (1)
tilelang/language/tir/op.py (1)
call_extern(173-195)
examples/distributed/deepseek_deepep/intranode/get_dispatch_layout.py (1)
examples/distributed/deepseek_deepep/buffer.py (1)
get_dispatch_layout(194-209)
src/op/builtin.h (2)
tilelang/language/reduce.py (5)
warp_reduce_sum(278-291)warp_reduce_max(294-307)warp_reduce_min(310-323)warp_reduce_bitand(326-339)warp_reduce_bitor(342-355)tilelang/language/builtin.py (5)
elect_one_sync(811-813)sync_warp(816-818)loop_continue(821-823)warp_any(826-836)warp_all(839-849)
src/op/builtin.cc (2)
tilelang/language/reduce.py (5)
warp_reduce_sum(278-291)warp_reduce_max(294-307)warp_reduce_min(310-323)warp_reduce_bitand(326-339)warp_reduce_bitor(342-355)tilelang/language/builtin.py (5)
elect_one_sync(811-813)sync_warp(816-818)loop_continue(821-823)warp_any(826-836)warp_all(839-849)
tilelang/language/distributed/common.py (2)
src/op/remote_copy.h (1)
tvm(16-323)tilelang/language/tir/op.py (2)
address_of(464-480)call_intrin(120-145)
testing/python/language/test_tilelang_language_elect.py (2)
tilelang/language/builtin.py (1)
elect_one_sync(811-813)tilelang/testing/__init__.py (1)
requires_cuda_compute_version(38-102)
tilelang/jit/kernel.py (1)
tilelang/utils/allocator.py (3)
initialized(162-163)table(224-225)table_size(228-229)
src/tl_templates/cuda/sync.h (2)
tilelang/language/builtin.py (1)
barrier_blocks(615-623)tilelang/language/distributed/common.py (3)
wait_eq(125-132)wait_ne(135-138)wait_ge(141-144)
examples/distributed/deepseek_deepep/intranode/combine.py (11)
src/tl_templates/cuda/reduce.h (1)
T(208-280)tilelang/jit/__init__.py (1)
jit(233-306)tilelang/jit/adapter/wrapper.py (2)
prim_func(700-710)prim_func(1150-1160)tilelang/language/proxy.py (1)
Tensor(251-252)tilelang/language/builtin.py (6)
sync_blocks(626-633)barrier_blocks(615-623)ld(741-774)elect_one_sync(811-813)sync_warp(816-818)st(777-808)tilelang/language/ast/ir.py (2)
init(361-369)handle(1467-1497)src/tl_templates/cuda/ldst.h (2)
ld(238-251)st(221-234)tilelang/language/tir/op.py (3)
tvm_warp_shuffle(550-571)if_then_else(2907-2937)address_of(464-480)tilelang/jit/kernel.py (1)
initialize(407-418)tilelang/language/distributed/common.py (1)
wait_ge(141-144)tilelang/language/allocate.py (2)
alloc_shared(24-39)alloc_local(42-53)
examples/distributed/deepseek_deepep/intranode/dispatch.py (3)
tilelang/language/builtin.py (4)
sync_blocks(626-633)st(777-808)barrier_blocks(615-623)ld(741-774)tilelang/language/copy.py (1)
copy(11-87)tilelang/language/distributed/common.py (2)
wait_ge(141-144)wait_ne(135-138)
src/op/sync.cc (2)
src/op/sync.h (4)
BarrierBlocksOp(180-186)Clone(143-152)WaitOp(100-105)is_distributed(61-79)src/op/remote_copy.h (4)
is_distributed(40-64)is_distributed(132-156)is_distributed(214-234)is_distributed(272-293)
src/op/remote_copy.h (1)
src/op/remote_copy.cc (18)
is_distributed(86-89)is_distributed(86-86)is_distributed(196-199)is_distributed(196-196)is_distributed(264-267)is_distributed(264-264)is_distributed(333-336)is_distributed(333-333)Lower(91-122)Lower(91-91)Lower(201-234)Lower(201-201)Lower(269-302)Lower(269-269)Lower(338-371)Lower(338-338)StOp(248-262)LdOp(316-331)
examples/distributed/deepseek_deepep/deepep_utils.py (1)
examples/distributed/deepseek_deepep/buffer.py (1)
num_channels(177-184)
examples/distributed/deepseek_deepep/intranode/example_intranode.py (5)
tilelang/distributed/utils.py (1)
init_dist(40-62)examples/distributed/deepseek_deepep/buffer.py (4)
EPBuffer(15-304)get_dispatch_layout(194-209)dispatch(211-281)combine(283-304)examples/distributed/deepseek_deepep/deepep_utils.py (1)
gen_inputs(129-158)examples/distributed/deepseek_deepep/intranode/test_intranode.py (1)
test_intranode(7-9)testing/python/language/test_tilelang_language_ldst_options.py (2)
main(11-16)main(25-28)
tilelang/language/reduce.py (1)
tilelang/language/tir/op.py (1)
call_intrin(120-145)
tilelang/language/builtin.py (2)
tilelang/language/tir/op.py (2)
call_intrin(120-145)address_of(464-480)src/tl_templates/cuda/ldst.h (2)
ld(238-251)st(221-234)
tilelang/utils/ts_ext/__init__.py (1)
tilelang/utils/ts_ext/tensor.cpp (2)
create_host_device_tensor(89-113)create_host_device_tensor(90-91)
🪛 actionlint (1.7.9)
.github/workflows/ci.yml
10-10: label "nvidia" is unknown. available labels are "windows-latest", "windows-latest-8-cores", "windows-2025", "windows-2022", "windows-11-arm", "ubuntu-slim", "ubuntu-latest", "ubuntu-latest-4-cores", "ubuntu-latest-8-cores", "ubuntu-latest-16-cores", "ubuntu-24.04", "ubuntu-24.04-arm", "ubuntu-22.04", "ubuntu-22.04-arm", "macos-latest", "macos-latest-xl", "macos-latest-xlarge", "macos-latest-large", "macos-26-xlarge", "macos-26", "macos-15-intel", "macos-15-xlarge", "macos-15-large", "macos-15", "macos-14-xl", "macos-14-xlarge", "macos-14-large", "macos-14", "macos-13-xl", "macos-13-xlarge", "macos-13-large", "macos-13", "self-hosted", "x64", "arm", "arm64", "linux", "macos", "windows". if it is a custom label for self-hosted runner, set list of labels in actionlint.yaml config file
(runner-label)
10-10: label "hopper" is unknown. available labels are "windows-latest", "windows-latest-8-cores", "windows-2025", "windows-2022", "windows-11-arm", "ubuntu-slim", "ubuntu-latest", "ubuntu-latest-4-cores", "ubuntu-latest-8-cores", "ubuntu-latest-16-cores", "ubuntu-24.04", "ubuntu-24.04-arm", "ubuntu-22.04", "ubuntu-22.04-arm", "macos-latest", "macos-latest-xl", "macos-latest-xlarge", "macos-latest-large", "macos-26-xlarge", "macos-26", "macos-15-intel", "macos-15-xlarge", "macos-15-large", "macos-15", "macos-14-xl", "macos-14-xlarge", "macos-14-large", "macos-14", "macos-13-xl", "macos-13-xlarge", "macos-13-large", "macos-13", "self-hosted", "x64", "arm", "arm64", "linux", "macos", "windows". if it is a custom label for self-hosted runner, set list of labels in actionlint.yaml config file
(runner-label)
🪛 markdownlint-cli2 (0.18.1)
examples/distributed/deepseek_deepep/deepep.md
3-3: Bare URL used
(MD034, no-bare-urls)
146-146: Fenced code blocks should have a language specified
(MD040, fenced-code-language)
🪛 Ruff (0.14.10)
tilelang/language/__init__.py
66-66: Unused noqa directive (non-enabled: F401)
Remove unused noqa directive
(RUF100)
67-67: Unused noqa directive (non-enabled: F401)
Remove unused noqa directive
(RUF100)
68-68: Unused noqa directive (non-enabled: F401)
Remove unused noqa directive
(RUF100)
69-69: Unused noqa directive (non-enabled: F401)
Remove unused noqa directive
(RUF100)
70-70: Unused noqa directive (non-enabled: F401)
Remove unused noqa directive
(RUF100)
examples/distributed/primitives/example_remote_st.py
14-14: Unused function argument: num_rank
(ARG001)
37-37: Unpacked variable rank is never used
Prefix it with an underscore or any other dummy variable pattern
(RUF059)
68-68: Avoid specifying long messages outside the exception class
(TRY003)
tilelang/language/print.py
226-227: Avoid specifying long messages outside the exception class
(TRY003)
examples/distributed/deepseek_deepep/buffer.py
252-252: Unpacked variable rank_prefix_matrix is never used
Prefix it with an underscore or any other dummy variable pattern
(RUF059)
252-252: Unpacked variable channel_prefix_matrix is never used
Prefix it with an underscore or any other dummy variable pattern
(RUF059)
252-252: Unpacked variable recv_channel_prefix_matrix is never used
Prefix it with an underscore or any other dummy variable pattern
(RUF059)
252-252: Unpacked variable recv_src_idx is never used
Prefix it with an underscore or any other dummy variable pattern
(RUF059)
252-252: Unpacked variable send_head is never used
Prefix it with an underscore or any other dummy variable pattern
(RUF059)
tilelang/jit/kernel.py
410-410: PEP 484 prohibits implicit Optional
Convert to Optional[T]
(RUF013)
examples/distributed/deepseek_deepep/intranode/combine.py
238-238: Shebang should contain python, pytest, or uv run
(EXE003)
238-238: Shebang should be at the beginning of the file
(EXE005)
examples/distributed/deepseek_deepep/intranode/dispatch.py
13-13: Unused noqa directive (unused: F403)
Remove unused noqa directive
(RUF100)
508-508: Shebang should contain python, pytest, or uv run
(EXE003)
508-508: Shebang should be at the beginning of the file
(EXE005)
786-786: Shebang should contain python, pytest, or uv run
(EXE003)
786-786: Shebang should be at the beginning of the file
(EXE005)
examples/distributed/deepseek_deepep/intranode/example_intranode.py
31-31: Unused noqa directive (unused: F403)
Remove unused noqa directive
(RUF100)
33-33: Avoid specifying long messages outside the exception class
(TRY003)
40-40: Unpacked variable rank_idx is never used
Prefix it with an underscore or any other dummy variable pattern
(RUF059)
62-62: Unpacked variable event is never used
Prefix it with an underscore or any other dummy variable pattern
(RUF059)
tilelang/language/builtin.py
514-514: PEP 484 prohibits implicit Optional
Convert to Optional[T]
(RUF013)
tilelang/utils/ts_ext/__init__.py
11-18: __all__ is not sorted
Apply an isort-style sorting to __all__
(RUF022)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (3)
- GitHub Check: Build wheels for Python 3.8 on ubuntu-latest with CUDA-12.1
- GitHub Check: Build wheels for Python 3.8 on macos-latest with Metal
- GitHub Check: Build wheels for Python 3.8 on ubuntu-24.04-arm with CUDA-12.8
| - [] Internode Normal Mode | ||
| - [] Low-latency Mode |
There was a problem hiding this comment.
Fix inconsistent TODO checkbox syntax.
The unchecked items use [] instead of [ ] (with a space), which may not render correctly in some markdown renderers.
- [x] Intranode Normal Mode
-- [] Internode Normal Mode
-- [] Low-latency Mode
+- [ ] Internode Normal Mode
+- [ ] Low-latency Mode📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| - [] Internode Normal Mode | |
| - [] Low-latency Mode | |
| - [ ] Internode Normal Mode | |
| - [ ] Low-latency Mode |
🤖 Prompt for AI Agents
In examples/distributed/deepseek_deepep/deepep.md around lines 7 to 8, the
unchecked checklist items use `[]` instead of the correct `[ ]` syntax; update
both lines to use `[ ]` (i.e., change "- [] Internode Normal Mode" to "- [ ]
Internode Normal Mode" and "- [] Low-latency Mode" to "- [ ] Low-latency Mode")
so they render as unchecked checkboxes in Markdown.
| def kernel_(M, num_rank, block_M, threads): | ||
|
|
||
| @T.prim_func | ||
| def main( | ||
| dst: T.Tensor((M), "float32"), | ||
| src: T.Tensor((M), "float32"), | ||
| ): | ||
| with T.Kernel(T.ceildiv(M, block_M), threads=threads) as (bx): | ||
| rank = T.alloc_local([1], "uint64") | ||
| num_rank = T.alloc_local([1], "uint64") | ||
| rank[0] = T.get_rank() | ||
| num_rank[0] = T.get_num_ranks() | ||
| tx = T.get_thread_binding() | ||
| T.st(dst[bx * block_M + tx], src[bx * block_M + tx], dst_pe=1 - rank[0]) | ||
|
|
||
| return main |
There was a problem hiding this comment.
Remove unused parameter and resolve variable shadowing.
The num_rank parameter is never used in the function body. Additionally, line 23 creates a local variable with the same name, causing shadowing. Since num_rank is fetched dynamically via T.get_num_ranks() at line 25, the parameter is unnecessary.
🔎 Proposed fix
-def kernel_(M, num_rank, block_M, threads):
+def kernel_(M, block_M, threads):And update the caller on line 45:
- kernel = tilelang.compile(kernel_(M, num_ranks, BLOCK_M, threads))
+ kernel = tilelang.compile(kernel_(M, BLOCK_M, threads))Committable suggestion skipped: line range outside the PR's diff.
🧰 Tools
🪛 Ruff (0.14.10)
14-14: Unused function argument: num_rank
(ARG001)
🤖 Prompt for AI Agents
In examples/distributed/primitives/example_remote_st.py around lines 14 to 29,
the kernel_ function declares a num_rank parameter that is never used and is
shadowed by a local variable assigned from T.get_num_ranks(); remove the unused
num_rank parameter from kernel_’s signature and any corresponding argument
passed by the caller (update the call site around line 45), and keep the local
num_rank allocation (or rename it if you prefer) so there is no shadowing or
unused parameter left.
| BLOCK_M = threads = 128 | ||
| assert num_local_ranks == 2, "this example only supports 2 ranks copying to each other" | ||
|
|
||
| rank, num_ranks, group = init_dist(local_rank, num_local_ranks) |
There was a problem hiding this comment.
Remove unused variable from unpacking.
The rank variable is unpacked but never used in this function. Only num_ranks and group are used.
🔎 Proposed fix
- rank, num_ranks, group = init_dist(local_rank, num_local_ranks)
+ _, num_ranks, group = init_dist(local_rank, num_local_ranks)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| rank, num_ranks, group = init_dist(local_rank, num_local_ranks) | |
| _, num_ranks, group = init_dist(local_rank, num_local_ranks) |
🧰 Tools
🪛 Ruff (0.14.10)
37-37: Unpacked variable rank is never used
Prefix it with an underscore or any other dummy variable pattern
(RUF059)
🤖 Prompt for AI Agents
In examples/distributed/primitives/example_remote_st.py around line 37, the code
unpacks three values as "rank, num_ranks, group = init_dist(local_rank,
num_local_ranks)" but never uses rank; remove the unused variable from the
unpacking so only the needed values are assigned (e.g., "num_ranks, group =
init_dist(...)" or adjust init_dist to return a tuple you can unpack into two
variables), and update any related code/comments to reflect the change.
| std::pair<torch::Tensor, torch::Tensor> | ||
| create_host_device_tensor(const std::vector<int64_t> &shape, | ||
| c10::ScalarType dtype); |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
rg -n -C10 'torch::from_blob|cudaHostFree' tilelang/utils/ts_ext/tensor.cppRepository: tile-ai/tilescale
Length of output: 687
Add custom deleter to manage pinned memory lifecycle in tensor.cpp.
The create_host_device_tensor implementation allocates pinned memory via cudaHostAlloc but creates tensors with torch::from_blob without custom deleters. This causes memory leaks since cudaHostAlloc'd memory must be explicitly freed via cudaHostFree. Provide a custom deleter to the from_blob calls to properly release the memory when tensors are destroyed.
🤖 Prompt for AI Agents
In tilelang/utils/ts_ext/ts_ext_ops.h around lines 16 to 18, the header declares
create_host_device_tensor whose implementation in tensor.cpp allocates pinned
host memory with cudaHostAlloc but constructs host tensors via torch::from_blob
without a custom deleter, causing pinned memory leaks; update the from_blob
calls to supply a custom deleter lambda that calls cudaHostFree on the allocated
host pointer (checking for nullptr and handling CUDA errors as appropriate) so
the pinned memory is released when the host tensor is destroyed, and ensure the
deleter signature matches from_blob's expected deleter type so lifetime is
correctly managed.
Intra-node DeepEP
Correctness
Extra
Benchmark Results
Measured on: 8xH100 on NVL, 10 channels, 8 ranks, 32 experts, 7168 hidden, 4096 tokens, normal mode.
Summary by CodeRabbit
New Features
Documentation
Tests
Chores
✏️ Tip: You can customize this high-level summary in your review settings.