Port TRT-LLM communication kernels to flashinfer #2102
Conversation
WalkthroughAdds a throughput‑optimized Mixture‑of‑Experts all‑to‑all backend: new CUDA dispatch/combine/sanitize kernels and headers, C++ FFI entry points and env helpers, Python/JIT bindings and workspace management, metainfo helpers, and extensive single‑GPU and MPI multi‑rank tests and docs. Changes
Sequence Diagram(s)sequenceDiagram
participant Py as Python API (MoeAlltoAll)
participant JIT as JIT Module
participant FFI as C++ FFI
participant GPU as GPU Kernels
participant Meta as Metainfo / Flags
Note over Py,JIT: ensure JIT module built/loaded
Py->>JIT: ensure module
Py->>FFI: moe_a2a_initialize(workspace, ep_rank, ep_size, max_tokens)
FFI->>Meta: init metainfo, counters, flags
Py->>FFI: moe_a2a_dispatch(routing, payloads, workspace, metainfo, ...)
FFI->>GPU: launch prepare_dispatch & moeA2ADispatchKernel
GPU->>Meta: write recv buffers, update counters & completion flags
FFI-->>Py: return recv tensors + combine_payload_offset
Py->>FFI: moe_a2a_combine(payload, local_num_tokens, workspace, metainfo, ..., offset)
FFI->>GPU: launch prepare_combine & moeA2ACombineKernel
GPU->>Meta: read recv buffers, accumulate TOP_K contributions, write outputs
FFI-->>Py: return combined outputs
Py->>FFI: moe_a2a_sanitize_expert_ids(expert_ids, workspace, metainfo, ep_rank, invalid_id)
FFI->>GPU: launch sanitize kernel
FFI-->>Py: sanitization complete
Estimated code review effort🎯 5 (Critical) | ⏱️ ~120 minutes Areas needing extra attention:
Suggested reviewers
Poem
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 |
djns99
force-pushed
the
djns99/update-trtllm-kernels
branch
5 times, most recently
from
710a388 to
bd82a2b
Compare
| #define check_timeout(s) false | ||
| #else | ||
| // 300 * 2000 MHz - should be high enough on any GPU but will prevent a hang | ||
| #define check_timeout(s) ((clock64() - (s)) > (300ll * 2000ll * 1000ll * 1000ll)) |
There was a problem hiding this comment.
I have manually added this, can I get someone to sanity check my logic here
djns99
force-pushed
the
djns99/update-trtllm-kernels
branch
from
f766bfe to
8cdf8d8
Compare
| ) | ||
|
|
||
|
|
||
| def moe_a2a_get_workspace_size_per_rank( |
There was a problem hiding this comment.
This API is also an addition of mine, some extra eyes on this would be helpful
There was a problem hiding this comment.
It looks like TRT-LLM just added (in the last 24 hours) something similar which looks slightly more user friendly: NVIDIA/TensorRT-LLM@8b5eded
Maybe in a future PR that could be integrated to flashinfer
There was a problem hiding this comment.
Yeah I am in touch with the author, I was going to port across as part of this, it should be a small change
There was a problem hiding this comment.
BTW, I am still seeing names like "MNNVL". This term is not precise as the kernel is also supposed to work for the nvlink domain within a single node like DGX B200. Could you modify it to be sth like "trtllm" "nvlink_one_sided" to reduce discrepancy?
There was a problem hiding this comment.
Currently this uses MnnvlMemory.initialize() in the MoeAllToAll class and uses that to allocate the workspace, does that not imply that this is still MNNVL only?
I agree we can rename, but I wonder if it would be more consistent with flashinfer to keep the mnnvl name?
That said, better to fix the name now while no-one is using it
There was a problem hiding this comment.
Also nvlink_one_sided is really a technical detail that I don't think we should be exposing to end users. Instead should I call it nvlink_moe_alltoall or something similar. Though I guess the dispatch is not specific to MOE
djns99
requested review from
aleozlx,
bkryu,
cyx-6,
djmmoss,
jiahanc,
joker-eph,
nvmbreughe,
wenscarl,
yongwww and
yzh119
as code owners
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 5
🧹 Nitpick comments (18)
csrc/nv_internal/cpp/common/envUtils.cpp (1)
357-357: Consider cachinggetEnvEplbForceGdrcopylike other bool env helpers
getEnvEplbForceGdrcopycallsgetBoolEnv(and thusstd::getenv) on every invocation, while most other helpers in this file cache the value in astaticlocal. Functionally this is fine, but for consistency and to avoid repeated env lookups in hot paths you might want to align it:-bool getEnvEplbForceGdrcopy() { return getBoolEnv("TRTLLM_EPLB_FORCE_GDRCOPY"); } +bool getEnvEplbForceGdrcopy() { + static bool const forceGdrcopy = getBoolEnv("TRTLLM_EPLB_FORCE_GDRCOPY"); + return forceGdrcopy; +}Not critical, but it would match the rest of the env-utils style.
csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu (2)
428-431: Consider documenting why the acquire fence is commented out.The
fence.acquire.sysat line 430 is commented out after the dispatch wait loop. While the combine kernel (line 735) does have an acquire fence, having it commented here without explanation could cause confusion for future maintainers. If this is intentional (relying on the combine kernel's fence), a brief comment explaining the design decision would help.} - // asm volatile("fence.acquire.sys;"); + // NOTE: Acquire fence intentionally omitted here; combine kernel provides + // the acquire semantics before reading peer data. #endif
596-609: Generic fallback is unreachable code.The generic fallback reduction loop (lines 599-608) can never be reached because the
SWITCH_TOP_Kmacro (lines 53-78) only allowsTOP_Kvalues of 1, 2, 4, or 8, and all these cases have explicit handling above. Consider removing this dead code or adding astatic_assertto document the constraint.} else if constexpr (TOP_K == 1) { // nothing to do - } else { - // Generic fallback: accumulate all into acc[0] - T* a0 = reinterpret_cast<T*>(&acc[0]); -#pragma unroll - for (int k = 1; k < TOP_K; ++k) { - T* ak = reinterpret_cast<T*>(&acc[k]); -#pragma unroll - for (int j = 0; j < elems_per_vec; ++j) { - a0[j] += ak[j]; - } - } + } else { + static_assert(TOP_K == 1 || TOP_K == 2 || TOP_K == 4 || TOP_K == 8, + "Only TOP_K values 1, 2, 4, 8 are supported"); }scripts/task_test_multi_node_comm_kernels.sh (1)
9-13: Disabling cache cleanup may cause stale import issues.The cache cleanup commands are commented out. If module refactoring occurs between test runs, stale
.pycfiles could cause import errors or unexpected behavior. Consider re-enabling these commands or documenting why they're disabled.tests/comm/test_trtllm_moe_alltoall.py (4)
1-2: Copyright year should be updated to 2025.The license header shows 2024 but this is a new file created in 2025.
-Copyright (c) 2024 by FlashInfer team. +Copyright (c) 2025 by FlashInfer team.
112-112: Potential issue with payload size calculation.
x[0].numel()gets the number of elements in the first row, but ifinput_tensorsis a list of 2D tensors, this calculates size per token correctly. However, the variable namepayload_size_per_tokenand the indexingx[0]could be clearer.- payload_size_per_token = sum([x[0].numel() * x.itemsize for x in input_tensors]) + payload_size_per_token = sum([x.shape[-1] * x.element_size() for x in input_tensors])
207-236: CUDA streams created but not explicitly cleaned up.The
cuda_streams_all_rankslist creates CUDA streams that are not explicitly destroyed. While Python's garbage collector will eventually clean them up, for test reliability consider using a context manager or explicit cleanup.
411-411: Minor typo in comment.Extra slash at end of comment.
- # For each expert selected for this token/ + # For each expert selected for this tokentests/comm/test_mnnvl_moe_alltoall.py (4)
37-46: Consider usingraisewithout exception name per Python best practices.The explicit
raise eis redundant; bareraisepreserves the traceback better.def safe_run(func, *args, **kwargs): comm = MPI.COMM_WORLD try: func(*args, **kwargs) except MPIExit as e: - raise e + raise except Exception as e: traceback.print_exc() comm.allgather(True) - raise e + raise
49-51: Test fixture should yield for proper cleanup semantics.Even though no cleanup is needed, the fixture pattern should include
yieldfor consistency.@pytest.fixture(autouse=True) def setup_test(): torch.manual_seed(0x1234) + yield
571-576: Blind exception catch may mask real initialization errors.Catching bare
Exceptionwhen checking MNNVL support could hide legitimate configuration issues. Consider catching specific exception types or at least logging the exception.try: MnnvlMemory.initialize() if not MnnvlMemory.supports_mnnvl(): pytest.skip("MNNVL not supported on this system") - except Exception: + except (RuntimeError, pynvml.NVMLError) as e: + # Log exception for debugging if needed pytest.skip("MNNVL not supported on this system")
709-712: Unused variableexpert_id_payload_indexas flagged by static analysis.The unpacked variable is never used. Either prefix with underscore or remove from unpacking.
- payloads, expert_id_payload_index = make_bfloat16_payloads( + payloads, _expert_id_payload_index = make_bfloat16_payloads( local_num_tokens, hidden_size, top_k, rank, token_selected_experts )flashinfer/comm/trtllm_moe_alltoall.py (5)
8-8: TODO comment should be addressed or tracked.The
# TODO Reviewcomment at the top suggests this module needs review. Consider removing after review or converting to a tracked issue.Would you like me to open an issue to track any remaining review items?
351-351: Mutable class attribute should useClassVarannotation.Per static analysis and Python best practices, mutable class attributes should be annotated with
typing.ClassVar.+from typing import ClassVar + class MoeAlltoAll: ... # Single shared workspace across the process - _WORKSPACE_CACHE: dict[tuple[int, int, int, int], dict] = {} + _WORKSPACE_CACHE: ClassVar[dict[tuple[int, int, int, int], dict]] = {}
456-463: Assertions for validation could use proper exceptions in production.Using
assertfor validation is acceptable for debug builds but these checks may be skipped in optimized Python (python -O). Consider using explicitif/raisefor critical invariants.
610-610: Inefficient way to get element size.Creating an empty tensor just to get element size is wasteful. Use
torch.finfoortorch.iinfoor a lookup table instead.- element_size = torch.tensor([], dtype=dtype).element_size() + # More efficient: use dtype itemsize directly + element_size = torch.empty(0, dtype=dtype).element_size()Or better, consider caching element sizes or using:
element_size = torch.finfo(dtype).bits // 8 if dtype.is_floating_point else torch.iinfo(dtype).bits // 8
621-628:__all__is not sorted as noted by static analysis.Consider sorting for consistency, though this is a minor issue.
__all__ = [ "MoeAlltoAll", "moe_a2a_initialize", + "moe_a2a_combine", "moe_a2a_dispatch", - "moe_a2a_combine", + "moe_a2a_get_workspace_size_per_rank", "moe_a2a_sanitize_expert_ids", - "moe_a2a_get_workspace_size_per_rank", ]csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h (1)
78-120: Well-documented struct with clear field descriptions.The
MoeA2ADispatchParamsstruct has excellent inline documentation explaining each field's purpose and dimensions. The TODO on line 90-91 about renamingmax_tokens_per_ranktoruntime_max_tokens_per_rankshould be tracked.Would you like me to open an issue to track the TODO about renaming
max_tokens_per_rank?
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (16)
csrc/nv_internal/cpp/common/envUtils.cpp(2 hunks)csrc/nv_internal/tensorrt_llm/common/envUtils.h(2 hunks)csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu(1 hunks)csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h(1 hunks)csrc/nv_internal/tensorrt_llm/thop/moeAlltoAllMeta.h(1 hunks)csrc/trtllm_moe_a2a.cu(1 hunks)docs/api/comm.rst(1 hunks)flashinfer/aot.py(1 hunks)flashinfer/comm/__init__.py(1 hunks)flashinfer/comm/trtllm_moe_alltoall.py(1 hunks)flashinfer/jit/__init__.py(1 hunks)flashinfer/jit/comm.py(1 hunks)scripts/task_test_multi_node_comm_kernels.sh(1 hunks)tests/comm/test_mnnvl_memory.py(1 hunks)tests/comm/test_mnnvl_moe_alltoall.py(1 hunks)tests/comm/test_trtllm_moe_alltoall.py(1 hunks)
🧰 Additional context used
🧬 Code graph analysis (11)
flashinfer/jit/__init__.py (1)
flashinfer/jit/comm.py (1)
gen_mnnvl_a2a_module(83-109)
csrc/nv_internal/tensorrt_llm/common/envUtils.h (1)
csrc/nv_internal/cpp/common/envUtils.cpp (8)
getEnvKVCacheTimeOutputPath(275-278)getEnvKVCacheTimeOutputPath(275-275)getEnvMoeA2AOneBlockPerToken(326-333)getEnvMoeA2AOneBlockPerToken(326-326)getEnvMoeA2ADispatchBlockSize(347-350)getEnvMoeA2ADispatchBlockSize(347-347)getEnvMoeA2ACombineBlockSize(352-355)getEnvMoeA2ACombineBlockSize(352-352)
tests/comm/test_mnnvl_memory.py (1)
flashinfer/comm/mapping.py (1)
local_rank(391-392)
csrc/nv_internal/tensorrt_llm/thop/moeAlltoAllMeta.h (1)
csrc/trtllm_moe_a2a.cu (2)
getMoeA2AMetaInfoIndexPairs(395-407)getMoeA2AMetaInfoIndexPairs(395-395)
csrc/trtllm_moe_a2a.cu (1)
csrc/nv_internal/cpp/common/envUtils.cpp (2)
getEnvMoeA2AOneBlockPerToken(326-333)getEnvMoeA2AOneBlockPerToken(326-326)
flashinfer/aot.py (1)
flashinfer/jit/comm.py (1)
gen_mnnvl_a2a_module(83-109)
flashinfer/jit/comm.py (1)
flashinfer/jit/core.py (2)
JitSpec(213-312)gen_jit_spec(315-381)
csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu (1)
csrc/nv_internal/cpp/common/envUtils.cpp (4)
getEnvMoeA2ADispatchBlockSize(347-350)getEnvMoeA2ADispatchBlockSize(347-347)getEnvMoeA2ACombineBlockSize(352-355)getEnvMoeA2ACombineBlockSize(352-352)
flashinfer/comm/trtllm_moe_alltoall.py (4)
flashinfer/comm/mnnvl.py (5)
MnnvlMemory(232-551)MnnvlConfig(224-229)as_torch_strided_tensor(264-273)initialize(276-285)set_comm_from_config(288-293)flashinfer/comm/mapping.py (2)
Mapping(21-475)moe_ep_rank(349-350)flashinfer/jit/comm.py (1)
gen_mnnvl_a2a_module(83-109)include/flashinfer/trtllm/fused_moe/runner.h (1)
num_experts(263-263)
tests/comm/test_mnnvl_moe_alltoall.py (2)
flashinfer/comm/trtllm_moe_alltoall.py (4)
MoeAlltoAll(336-618)dispatch(484-541)get_combine_payload_tensor_in_workspace(585-618)combine(543-583)flashinfer/comm/mnnvl.py (3)
MnnvlMemory(232-551)initialize(276-285)supports_mnnvl(545-551)
csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h (3)
csrc/nv_internal/tensorrt_llm/common/envUtils.h (1)
tensorrt_llm(23-104)csrc/nv_internal/tensorrt_llm/thop/moeAlltoAllMeta.h (1)
mnnvl_throughput(25-58)csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu (8)
moe_a2a_dispatch_launch(445-506)moe_a2a_dispatch_launch(445-445)moe_a2a_prepare_dispatch_launch(436-439)moe_a2a_prepare_dispatch_launch(436-436)moe_a2a_combine_launch(792-842)moe_a2a_combine_launch(792-792)moe_a2a_sanitize_expert_ids_launch(864-872)moe_a2a_sanitize_expert_ids_launch(864-866)
🪛 Clang (14.0.6)
csrc/nv_internal/tensorrt_llm/thop/moeAlltoAllMeta.h
[error] 19-19: 'array' file not found
(clang-diagnostic-error)
csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h
[error] 18-18: 'cuda_bf16.h' file not found
(clang-diagnostic-error)
🪛 Ruff (0.14.6)
flashinfer/comm/trtllm_moe_alltoall.py
351-351: Mutable class attributes should be annotated with typing.ClassVar
(RUF012)
443-443: Avoid specifying long messages outside the exception class
(TRY003)
445-445: Avoid specifying long messages outside the exception class
(TRY003)
606-608: Avoid specifying long messages outside the exception class
(TRY003)
621-628: __all__ is not sorted
Apply an isort-style sorting to __all__
(RUF022)
tests/comm/test_mnnvl_moe_alltoall.py
34-34: Avoid specifying long messages outside the exception class
(TRY003)
42-42: Use raise without specifying exception name
Remove exception name
(TRY201)
46-46: Use raise without specifying exception name
Remove exception name
(TRY201)
575-575: Do not catch blind exception: Exception
(BLE001)
676-676: Do not catch blind exception: Exception
(BLE001)
710-710: Unpacked variable expert_id_payload_index is never used
Prefix it with an underscore or any other dummy variable pattern
(RUF059)
🔇 Additional comments (38)
csrc/nv_internal/cpp/common/envUtils.cpp (2)
275-277: KV cache time output path now cached by reference – behavior looks consistentReturning
std::string const&to a function-local static is safe here and matches the pattern used by other helpers in this file (single evaluation of the env var, cheap repeated access). The only behavior change is that changes toTRTLLM_KVCACHE_TIME_OUTPUT_PATHafter the first call are no longer observed, which is usually fine for env-based config.If you don’t rely on mutating the env mid-process (including in tests), this change looks good to me.
326-355: Based on my analysis with the available information, I can now rewrite the review comment:Verify MOE A2A environment variable prefixes (
TLLM_vsTRTLLM_) against project documentation and call sites; fix misleading comment insanitizeBlockSizeto reflect "round up" instead of "nearest"The new helpers follow the file's consistent patterns (single-read statics, sane defaults), but need attention in two areas:
Env var prefixes for MoE A2A knobs
- These functions read
TLLM_MOE_A2A_ONE_BLOCK_PER_TOKEN,TLLM_MOE_A2A_DISPATCH_BLOCK_SIZE, andTLLM_MOE_A2A_COMBINE_BLOCK_SIZE- Review the project documentation, call sites, and surrounding code to confirm whether the
TLLM_prefix matches intended usage; most other envs in this file appear to useTRTLLM_prefix and silently using the wrong names would be problematic.
sanitizeBlockSizecomment accuracy
- The implementation rounds up to a multiple of 32:
block = (block + 31) / 32 * 32;(e.g., 33 → 64, not 32)- The comment currently states "Round to nearest multiple of 32 (warp size)", which is misleading.
Fix the comment:
- // Round to nearest multiple of 32 (warp size)
- // Round up to the next multiple of 32 (warp size)
The extra `if (block == 0) block = 256;` check after clamping is redundant but harmless.csrc/nv_internal/tensorrt_llm/common/envUtils.h (1)
95-102: LGTM! New MoE A2A environment variable accessors are properly declared.The three new accessor functions (
getEnvMoeA2AOneBlockPerToken,getEnvMoeA2ADispatchBlockSize,getEnvMoeA2ACombineBlockSize) are well-documented with default behaviors and align with their implementations inenvUtils.cpp.csrc/nv_internal/tensorrt_llm/thop/moeAlltoAllMeta.h (3)
19-22: Static analysis false positive - standard headers are valid.The Clang error about
<array>not being found is a false positive. This is a standard C++11 header that should be available in any modern C++ environment. The includes are correct.
28-43: LGTM! Well-structured metadata index enum.The
MoeA2AMetaInfoIndexenum provides clear, sequential indexing for metadata fields withNUM_METAINFO_FIELDS = 9correctly representing the count of actual data fields (0-8). TheMoeA2ADataOffsetstype alias correctly uses this count for the array size.
45-58: LGTM! Useful name-to-index mapping function.The inline
getMoeA2AMetaInfoIndexPairs()function provides a clean way to expose metadata field names and their corresponding indices, which is consumed by the TVM FFI interface incsrc/trtllm_moe_a2a.cu.csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu (4)
114-116: Timeout calculation looks reasonable.The timeout of
300ll * 2000ll * 1000ll * 1000llcycles (~600 billion) translates to approximately 300 seconds at 2 GHz, which provides a generous upper bound to prevent infinite hangs while allowing ample time for legitimate synchronization delays. The calculation avoids overflow by usinglong longliterals.
272-276: LGTM! Single-threaded flag increment is safe.The
flag_valincrement occurs only whenidx == 0, ensuring single-threaded access. Since this kernel runs sequentially in the stream before the dispatch kernel, there's no race condition.
844-872: LGTM! Sanitize kernel implementation is correct.The kernel correctly identifies invalid tokens (where
token_idx >= recv_counters[source_rank]) and sets their expert IDs toinvalid_id. Each thread operates on disjoint memory locations, avoiding any race conditions.
315-344: Based on my investigation, I cannot access the repository or find public documentation for thekMaxRanksconstant definition. The repository clone failed, and web searches returned no results for this internal NVIDIA TensorRT-LLM code.However, the core concern raised in the review comment remains valid and cannot be conclusively verified without access to:
- The header file containing
kMaxRanksdefinition- The actual value of
kMaxRanks- Runtime validation constraints on
ep_size(ensemble parallel size)The potential undefined behavior is legitimate: if
target_rankcan be 64 or greater, the bit shift1ULL << target_rankon auint64_twould indeed cause undefined behavior in C++.
Verify
kMaxRanksdoes not exceed 64 to avoid undefined behavior.The
already_copiedbitmask usesuint64_twith bit operations1ULL << target_rank. Iftarget_rankcan be 64 or greater, this causes undefined behavior (shifting by >= width of type). The code validatesparams.ep_size <= kMaxRanksat line 448, so ensurekMaxRanksis defined as ≤ 64 in the header.flashinfer/aot.py (1)
515-522: LGTM! MNNVL A2A module integration follows existing patterns.The new
gen_mnnvl_a2a_moduleis correctly imported within theadd_commblock and added under thehas_sm100condition, consistent with the existinggen_trtllm_comm_moduleandgen_trtllm_mnnvl_comm_moduleplacement.tests/comm/test_mnnvl_memory.py (1)
125-125: LGTM! Correct device selection for multi-node scenarios.Using
self.local_rankinstead ofself.rankis the correct fix for multi-node setups where the global rank may exceed the local GPU count. This aligns with the setup fixture (line 51) and theMapping.local_rankproperty shown in the relevant snippet.scripts/task_test_multi_node_comm_kernels.sh (1)
17-19: LGTM - new MoE A2A test coverage added.The addition of
test_mnnvl_moe_alltoall.pyaligns with the new MoE A2A functionality introduced in this PR.flashinfer/jit/__init__.py (1)
80-80: LGTM!The re-export follows the established pattern for module generators in this file.
flashinfer/comm/__init__.py (1)
42-52: LGTM!The new MoE A2A exports follow the established import patterns and properly expose the public API surface for the throughput backend.
csrc/trtllm_moe_a2a.cu (6)
53-99: LGTM - offset calculations and workspace sizing are well-structured.The alignment logic using cache-line boundaries (128 bytes) is appropriate for GPU memory access patterns. The offset calculations correctly account for metadata fields and payload regions.
101-123: LGTM - initialization logic is correct.Input validation is comprehensive. The
cudaStreamSynchronizeat line 119 ensures the memset completes before returning, which is necessary for correctness.
132-262: LGTM - dispatch implementation is solid.The function has comprehensive input validation, proper payload descriptor setup, workspace size verification, and appropriate CUDA error checking after kernel launch.
264-277: LGTM!The dtype conversion correctly handles the common floating-point types used in MoE computations. The unreachable return after the throw silences compiler warnings.
279-361: LGTM - combine implementation is well-validated.The function properly validates all inputs including the optional
payloadInWorkspaceoptimization for zero-copy scenarios. Error handling is comprehensive.
393-416: LGTM!The metainfo accessor correctly wraps the C++ pairs for Python accessibility, and all FFI exports are properly registered.
flashinfer/jit/comm.py (1)
83-109: I'll verify whether architecture-specific NVCC flags are required for thegen_mnnvl_a2a_modulefunction. Let me search the codebase for relevant information.
<function_calls>
#!/bin/bashFind the moeAlltoAllKernels.cu file
find . -name "moeAlltoAllKernels.cu" -type f 2>/dev/null
</function_calls>
<function_calls>
#!/bin/bashCheck if the file exists and search for architecture-specific patterns
if [ -f "csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu" ]; then
echo "=== Checking for SM guards and architecture macros ==="
rg -n 'sm|SM[0-9]|CUDA_ARCH|CUDA_SM' "csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu" || echo "No explicit SM guards found"echo ""
echo "=== Checking for tensor core/MMA operations ==="
rg -n 'mma|wmma|tensorop|__mma|__hmma' "csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu" || echo "No tensor core intrinsics found"
else
echo "File not found at expected path"
fi
</function_calls><function_calls>
#!/bin/bashLook at the comm.py file to understand nvcc_flags definition and usage
echo "=== Current comm.py around nvcc_flags definition ==="
rg -n 'nvcc_flags' flashinfer/jit/comm.py -B3 -A3echo ""
echo "=== All functions in comm.py to see pattern ==="
rg -n '^def gen_' flashinfer/jit/comm.py
</function_calls><function_calls>
#!/bin/bashLet's see the full context of comm.py to understand nvcc_flags and module patterns
wc -l flashinfer/jit/comm.py
</function_calls><function_calls>
#!/bin/bashRead the entire comm.py file to understand structure and nvcc_flags
cat -n flashinfer/jit/comm.py
</function_calls><function_calls>
#!/bin/bashCheck what nvcc_flags is and where it comes from
rg -n 'nvcc_flags\s*=' flashinfer/jit/comm.py -B2 -A2
echo ""
echo "=== Check for imports in comm.py ==="
head -30 flashinfer/jit/comm.py
</function_calls><function_calls>
#!/bin/bashCheck the jit module structure
ls -la flashinfer/jit/
</function_calls><function_calls>
#!/bin/bashCheck jit_env for nvcc_flags
rg -n 'nvcc_flags' flashinfer/jit/ -A2 -B2
</function_calls>tests/comm/test_trtllm_moe_alltoall.py (6)
25-29: Session-scoped fixture withtorch.manual_seedmay not reset seed between tests.The fixture sets the seed once at session start but individual tests may want reproducible seeds. Consider if this is intentional or if test-level seeding is needed.
62-69: Good defensive check for SM resources.This helper appropriately skips tests when insufficient SMs are available for parallel kernel execution. The check prevents hangs on systems with limited GPU resources.
72-84: LGTM!The
make_payloadhelper correctly distinguishes between integer and floating-point types for random tensor generation.
138-141: Sorting approach for validation is reasonable but fragile.Sorting both input and output tensors to compare them works for this test case but relies on unique values. If there are duplicate values, the sort order could differ. Consider documenting this assumption or using a more robust comparison.
388-429: Reference implementation for fake_moe looks correct.The fake_moe function properly handles expert parallelism filtering and accumulation. The tree reduction comment on line 423 correctly explains why results are summed after collection.
530-536: Relatively loose tolerance for numerical comparison.Using
atol=1.5e-2andrtol=1.5e-2is quite loose for bf16/fp16. This may mask precision issues. Verify this tolerance is intentional given the accumulation order differences mentioned elsewhere.tests/comm/test_mnnvl_moe_alltoall.py (3)
293-293: Direct modification of class variable_WORKSPACEis concerning.Setting
MoeAlltoAll._WORKSPACE = Nonedirectly before instantiation suggests test isolation concerns. This should be documented or handled via a proper reset method.Consider whether
_reset_workspace()method fromMoeAlltoAllshould be used instead, or if this pattern is intentional for test setup.
800-813: Good documentation of tolerance rationale.The comment on line 809 clearly explains why a 99% match threshold is used instead of exact comparison due to bf16 accumulation order differences. This is helpful for future maintainers.
836-838: Helpful run instructions in docstring.The comment showing how to run with
mpirunis useful for developers unfamiliar with MPI testing.flashinfer/comm/trtllm_moe_alltoall.py (3)
353-383: Workspace caching strategy looks correct.The caching by
(workspace_size_per_rank, ep_rank, ep_size, max_num_tokens)tuple allows reusing workspaces across instances with compatible configurations. This addresses the past review comment about supporting different shaped communicators.
470-482:_reset_workspacemethod deletes from class cache without thread safety.If multiple threads could access this class simultaneously, the
deloperation on_WORKSPACE_CACHEcould cause issues. Document that this method is not thread-safe.def _reset_workspace(self): - """Reset the workspace to free up its state. This is mainly used for testing. Use this with caution. This object is no longer usable after this.""" + """Reset the workspace to free up its state. + + Warning: This method is not thread-safe and is mainly used for testing. + This object is no longer usable after calling this method. + """
505-508: Good use of state machine pattern for dispatch/combine sequencing.The phase checking prevents calling dispatch twice without combine and ensures proper operation ordering. This is a clean design.
csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h (4)
17-19: Static analysis reports missingcuda_bf16.h- this is a false positive.The
cuda_bf16.handcuda_fp16.hheaders are provided by the CUDA toolkit and will be available during compilation with nvcc. Static analysis tools without CUDA environment cannot find these headers.
23-27: Configuration constants are well-documented and reasonable.The limits (256 experts, 8 top-k, 8 payloads, 64 ranks) provide good flexibility while keeping fixed-size arrays manageable. Consider whether these should be configurable at runtime if larger deployments are anticipated.
173-179: Function declarations are clean and match the implementation.The kernel launch function declarations align with the implementations shown in the relevant code snippets from
moeAlltoAllKernels.cu.
148-148: Unable to verify include configuration due to repository access failure.The repository clone failed, preventing me from examining the file's include structure, verifying whether
nvinfer1::DataTypeis actually used, or confirming if the necessary headers are already present. Manual verification is required to confirm:
- Whether
csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.hcurrently includes NvInfer headers- Whether
nvinfer1::DataTypeis actually declared in the file or included transitively- Whether the code compiles successfully without the suggested include
csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu
Show resolved
Hide resolved
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (3)
flashinfer/comm/trtllm_moe_alltoall.py (3)
379-379: Annotate mutable class attribute withClassVar.Per Python best practices, mutable class attributes should be annotated with
ClassVarto make clear they are shared across instances.+from typing import ClassVar + class MoeAlltoAll: ... - _WORKSPACE_CACHE: dict[tuple[int, int, int, int], dict] = {} + _WORKSPACE_CACHE: ClassVar[dict[tuple[int, int, int, int], dict]] = {}
638-638: Consider usingtorch.finfoortorch.iinfofor element size.Creating an empty tensor just to get element size has minor overhead. Consider using dtype introspection directly.
- element_size = torch.tensor([], dtype=dtype).element_size() + element_size = torch._utils._element_size(dtype)Alternatively, keep the current approach if you prefer avoiding private APIs.
649-656: Consider addingmoe_a2a_wrap_payload_tensor_in_workspaceto__all__.This function is used in tests and appears to be part of the public API. Also consider sorting
__all__for consistency.__all__ = [ "MoeAlltoAll", + "moe_a2a_combine", "moe_a2a_initialize", "moe_a2a_dispatch", - "moe_a2a_combine", + "moe_a2a_get_workspace_size_per_rank", "moe_a2a_sanitize_expert_ids", - "moe_a2a_get_workspace_size_per_rank", + "moe_a2a_wrap_payload_tensor_in_workspace", ]
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (6)
csrc/trtllm_moe_alltoall.cu(1 hunks)flashinfer/aot.py(1 hunks)flashinfer/comm/trtllm_moe_alltoall.py(1 hunks)flashinfer/jit/__init__.py(1 hunks)flashinfer/jit/comm.py(1 hunks)tests/comm/test_trtllm_moe_alltoall.py(1 hunks)
🚧 Files skipped from review as they are similar to previous changes (2)
- flashinfer/aot.py
- flashinfer/jit/comm.py
🧰 Additional context used
🧬 Code graph analysis (2)
csrc/trtllm_moe_alltoall.cu (2)
csrc/tvm_ffi_utils.h (3)
Tensor(282-284)get_current_stream(266-270)encode_dlpack_dtype(29-31)flashinfer/comm/trtllm_moe_alltoall.py (6)
moe_a2a_get_workspace_size_per_rank(175-198)moe_a2a_get_workspace_size_per_rank(350-361)moe_a2a_initialize(41-47)moe_a2a_initialize(210-218)moe_a2a_dispatch(53-93)moe_a2a_dispatch(251-309)
tests/comm/test_trtllm_moe_alltoall.py (7)
flashinfer/comm/mapping.py (1)
Mapping(21-475)tests/test_helpers/test_helpers.py (1)
get_device_properties(10-11)include/flashinfer/trtllm/fused_moe/runner.h (1)
num_experts(263-263)flashinfer/comm/trtllm_moe_alltoall.py (13)
moe_a2a_get_workspace_size_per_rank(175-198)moe_a2a_get_workspace_size_per_rank(350-361)MoeAlltoAll(364-646)dispatch(512-569)get_combine_payload_tensor_in_workspace(613-646)combine(571-611)moe_a2a_initialize(41-47)moe_a2a_initialize(210-218)moe_a2a_dispatch(53-93)moe_a2a_dispatch(251-309)moe_a2a_sanitize_expert_ids(146-155)moe_a2a_sanitize_expert_ids(338-347)moe_a2a_wrap_payload_tensor_in_workspace(221-248)flashinfer/fused_moe/utils.py (1)
_(157-163)csrc/xqa/mha.cu (1)
any(157-157)tests/comm/test_mnnvl_moe_alltoall.py (1)
fake_moe(121-181)
🪛 Ruff (0.14.6)
flashinfer/comm/trtllm_moe_alltoall.py
379-379: Mutable class attributes should be annotated with typing.ClassVar
(RUF012)
471-471: Avoid specifying long messages outside the exception class
(TRY003)
473-473: Avoid specifying long messages outside the exception class
(TRY003)
634-636: Avoid specifying long messages outside the exception class
(TRY003)
649-656: __all__ is not sorted
Apply an isort-style sorting to __all__
(RUF022)
🔇 Additional comments (16)
flashinfer/jit/__init__.py (1)
80-80: LGTM!The new import follows the established pattern for re-exporting JIT module generators from the comm submodule.
csrc/trtllm_moe_alltoall.cu (5)
53-88: LGTM!The offset calculation logic correctly aligns data structures to cache-line boundaries where needed, following a clear sequential layout pattern.
101-123: LGTM!The initialization correctly validates inputs, zeros the workspace region, and returns metadata offsets. The stream synchronization ensures the workspace is properly initialized before returning.
278-360: LGTM!The combine operation has thorough input validation, properly handles the workspace-backed payload case, and includes appropriate error checking after kernel launch.
362-394: LGTM!The sanitization operation correctly validates inputs and launches the kernel with proper error checking.
396-419: LGTM!The metainfo index pairs helper provides a clean mechanism to expose C++ constants to Python, and all required functions are properly exported.
tests/comm/test_trtllm_moe_alltoall.py (6)
74-86: LGTM!The payload generator correctly handles both integer and floating-point dtypes for test data generation.
93-161: LGTM!Comprehensive single-GPU test covering multiple payload dtypes, dispatch/combine workflow, and workspace-backed tensor operations.
164-240: LGTM!The helper correctly simulates multi-rank dispatch on a single GPU using separate CUDA streams, with proper synchronization.
302-344: LGTM!The multi-rank test correctly validates token routing across simulated ranks with proper verification of payload delivery.
390-431: LGTM!The reference MoE implementation provides a deterministic baseline for verifying combine correctness, with appropriate handling of expert-parallel scenarios.
434-551: LGTM!Comprehensive combine test covering multiple dtypes, workspace configurations, and ranks with appropriate numerical tolerances for reduced-precision arithmetic.
flashinfer/comm/trtllm_moe_alltoall.py (4)
32-207: LGTM!The JIT module getter follows the established pattern with proper caching and custom op registration.
221-248: LGTM!The function correctly creates a workspace-backed tensor view with properly documented parameters.
470-473: LGTM!The validation logic is appropriate and the exception messages are concise.
498-510: LGTM!The reset method appropriately handles workspace cleanup for testing scenarios, with clear documentation about post-call state.
djns99
force-pushed
the
djns99/update-trtllm-kernels
branch
from
6e9bed5 to
a51b1ea
Compare
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (6)
csrc/trtllm_moe_alltoall.cu (1)
263-276: Consider extending dtype support for future flexibility.The
toNvDataTypefunction currently supports half, bfloat16, and float32. Consider documenting supported types or adding int8/fp8 support if those are common in MoE workloads.tests/comm/test_mnnvl_moe_alltoall.py (3)
711-712: Unused variable is intentional; consider underscore prefix per Ruff hint.The
expert_id_payload_indexis returned by the helper but not used in this test. Consider renaming to_expert_id_payload_indexto signal intentional discard.- payloads, expert_id_payload_index = make_bfloat16_payloads( + payloads, _expert_id_payload_index = make_bfloat16_payloads(
293-294: Setting class attribute_WORKSPACE = Nonemay conflict with class-level cache.Assigning
MoeAlltoAll._WORKSPACE = Noneresets a non-existent instance attribute. The class uses_WORKSPACE_CACHEfor caching. This assignment has no effect but is misleading.Consider removing this line or using
MoeAlltoAll._WORKSPACE_CACHE.clear()if the intent is to reset the cache:- MoeAlltoAll._WORKSPACE = None + MoeAlltoAll._WORKSPACE_CACHE.clear()
742-742: Same issue:_WORKSPACE = Noneassignment is ineffective.This line also sets a non-existent attribute. Consider removing or using
_WORKSPACE_CACHE.clear().- MoeAlltoAll._WORKSPACE = Noneflashinfer/comm/trtllm_moe_alltoall.py (2)
375-377: Annotate mutable class attribute withClassVarper Ruff hint.The
_WORKSPACE_CACHEis a mutable class-level attribute that should be annotated withClassVarto make the intent clear.+from typing import ClassVar + class MoeAlltoAll: ... - _WORKSPACE_CACHE: dict[tuple[int, int, int, int], dict] = {} + _WORKSPACE_CACHE: ClassVar[dict[tuple[int, int, int, int], dict]] = {}
496-508: Document that_reset_workspaceinvalidates the instance.The docstring mentions this but it's critical: after calling
_reset_workspace, the object is unusable. Consider adding a stronger warning or raising an exception on subsequent method calls.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (4)
csrc/trtllm_moe_alltoall.cu(1 hunks)flashinfer/comm/trtllm_moe_alltoall.py(1 hunks)tests/comm/test_mnnvl_moe_alltoall.py(1 hunks)tests/comm/test_trtllm_moe_alltoall.py(1 hunks)
🧰 Additional context used
🧬 Code graph analysis (3)
csrc/trtllm_moe_alltoall.cu (2)
csrc/tvm_ffi_utils.h (3)
Tensor(282-284)get_current_stream(266-270)encode_dlpack_dtype(29-31)flashinfer/comm/trtllm_moe_alltoall.py (11)
moe_a2a_get_workspace_size_per_rank(173-196)moe_a2a_get_workspace_size_per_rank(348-359)moe_a2a_initialize(39-45)moe_a2a_initialize(208-216)moe_a2a_dispatch(51-91)moe_a2a_dispatch(249-307)moe_a2a_combine(97-138)moe_a2a_combine(310-333)moe_a2a_sanitize_expert_ids(144-153)moe_a2a_sanitize_expert_ids(336-345)moe_a2a_get_metainfo_index_pairs(159-167)
flashinfer/comm/trtllm_moe_alltoall.py (3)
flashinfer/comm/mnnvl.py (5)
MnnvlMemory(232-551)MnnvlConfig(224-229)as_torch_strided_tensor(264-273)initialize(276-285)set_comm_from_config(288-293)flashinfer/comm/mapping.py (2)
Mapping(21-475)moe_ep_rank(349-350)flashinfer/jit/comm.py (1)
gen_mnnvl_moe_alltoall_module(83-109)
tests/comm/test_mnnvl_moe_alltoall.py (3)
flashinfer/comm/trtllm_moe_alltoall.py (4)
MoeAlltoAll(362-644)dispatch(510-567)get_combine_payload_tensor_in_workspace(611-644)combine(569-609)flashinfer/comm/mapping.py (1)
Mapping(21-475)flashinfer/comm/mnnvl.py (3)
MnnvlMemory(232-551)initialize(276-285)supports_mnnvl(545-551)
🪛 Ruff (0.14.6)
flashinfer/comm/trtllm_moe_alltoall.py
377-377: Mutable class attributes should be annotated with typing.ClassVar
(RUF012)
469-469: Avoid specifying long messages outside the exception class
(TRY003)
471-471: Avoid specifying long messages outside the exception class
(TRY003)
632-634: Avoid specifying long messages outside the exception class
(TRY003)
tests/comm/test_mnnvl_moe_alltoall.py
34-34: Avoid specifying long messages outside the exception class
(TRY003)
576-576: Do not catch blind exception: Exception
(BLE001)
677-677: Do not catch blind exception: Exception
(BLE001)
711-711: Unpacked variable expert_id_payload_index is never used
Prefix it with an underscore or any other dummy variable pattern
(RUF059)
🔇 Additional comments (44)
csrc/trtllm_moe_alltoall.cu (8)
1-51: LGTM! Well-structured header and utility functions.The license header, includes, and utility functions are appropriately organized. The
alignOffsetfunction correctly implements cache-line alignment using bitwise operations.
53-88: LGTM! Offset calculation is well-structured.The
calculateOffsetsfunction properly calculates memory offsets for various MoE A2A data structures with appropriate alignment for cache-line boundaries.
90-99: LGTM! Workspace size calculation.The workspace size calculation correctly accounts for metadata, payload, and combine regions with proper alignment.
101-123: LGTM! Initialize operation with proper validation.Good input validation for workspace dimensions, rank bounds, and proper error checking for CUDA operations. The synchronization before returning metainfo is appropriate.
125-261: LGTM! Dispatch operation is well-implemented.The dispatch function has comprehensive input validation, proper payload descriptor handling, and correct workspace pointer arithmetic. Error checking after kernel launch is appropriate.
278-360: LGTM! Combine operation with proper validation.The combine function correctly validates payload dimensions, workspace pointer alignment, and handles the
payloadInWorkspaceflag appropriately. Error checking after kernel launch is proper.
362-394: LGTM! Sanitize operation is correctly implemented.Proper input validation and error checking for the sanitize expert IDs kernel.
396-419: LGTM! Metainfo export and FFI registration.The metainfo index pairs function and TVM FFI exports are correctly implemented, providing clean Python interoperability.
tests/comm/test_trtllm_moe_alltoall.py (12)
25-29: LGTM! Docstring has been corrected.The fixture docstring now accurately describes that it sets the torch seed for deterministic tests.
32-60: Good test parameter coverage.The test parameters cover a good range of configurations (small, medium, large) for both single-GPU and multi-rank scenarios, with various dtypes and payload configurations.
63-72: Good resource-aware skip logic.The SM count check appropriately skips tests when hardware resources are insufficient, preventing false failures on less capable GPUs.
74-86: LGTM! Payload generation helper.The
make_payloadfunction correctly handles both integer and floating-point dtypes with appropriate random value generation.
89-162: Comprehensive single-GPU test with proper verification.The test covers dispatch and combine flows with multiple dtypes, validates output via sorting and exact comparison, and tests the workspace-backed combine path.
164-240: LGTM! Multi-rank dispatch helper is well-structured.The helper properly manages workspaces, initializes per-rank metadata, uses separate CUDA streams for parallel execution, and synchronizes appropriately.
243-259: LGTM! Sanitize helper function.Simple and correct delegation to the underlying sanitize function for each rank.
262-299: LGTM! Combine helper with parallel execution.The combine helper correctly uses separate streams per rank and synchronizes before returning results.
302-345: LGTM! Multi-rank test with proper verification.Good verification logic that filters non-zero tensors and compares sorted outputs against the reference filtered by expert assignment.
347-388: LGTM! Sanitize test with comprehensive verification.The test properly clones tensors before sanitization to enable before/after comparison and correctly verifies the sanitization logic.
390-431: LGTM! Reference MoE implementation for verification.The
fake_moefunction provides a clear reference implementation for verifying the distributed MoE behavior, with proper EP-rank filtering logic.
434-555: Good end-to-end combine test with tolerance handling.The test covers the full dispatch-process-combine cycle with both in-workspace and external payload paths. The tolerance values for bf16 are reasonable.
tests/comm/test_mnnvl_moe_alltoall.py (11)
27-46: MPI error handling utilities are well-designed.The
MPIExitexception,check_any_rank_failed, andsafe_runpattern provide robust MPI coordination for test failures across ranks, ensuring clean error propagation.
49-52: LGTM! Test fixture for deterministic seeding.
55-88: LGTM! Helper functions for expert routing and token generation.
compute_target_rank_idcorrectly implements contiguous expert partitioning, andgenerate_token_selected_expertsproperly generates random expert assignments.
91-119: LGTM! Expert weight creation with reproducible seeding.Using
ep_rank * 1000 + ias a seed ensures reproducibility across runs while differentiating experts per rank.
122-182: LGTM! Comprehensive fake MoE reference implementation.The function correctly handles both EP-rank and global modes with proper local expert ID conversion.
185-258: LGTM! Payload creation helpers.Both NV FP4 and BFloat16 payload creators are well-structured with appropriate rank-specific patterns for verification.
261-383: LGTM! Single-rank dispatch worker function.Comprehensive workspace setup, dispatch execution, and metadata extraction for MPI-based testing.
386-556: LGTM! Thorough dispatch verification.The
verify_dispatchfunction provides exhaustive validation of shapes, dtypes, counters, routing, and payload content. This is excellent for catching regressions.
572-577: Bare Exception catch is intentional for MNNVL availability check.The broad exception catch here is acceptable since it's used to detect MNNVL support availability across various failure modes (driver issues, missing hardware, etc.).
673-678: Bare Exception catch is acceptable for capability detection.Same as above - this is intentional for gracefully skipping tests on systems without MNNVL support.
657-818: LGTM! Full dispatch+combine cycle test.The test properly verifies the complete MoE A2A workflow with appropriate tolerance for bf16 accumulation order differences. The 99% match threshold is reasonable given the expected numerical variations.
flashinfer/comm/trtllm_moe_alltoall.py (13)
1-19: LGTM! Module header and imports are well-organized.Clean module docstring and appropriate imports for the MoE A2A functionality.
21-28: LGTM! State dataclass is appropriate.The
_A2AStatedataclass cleanly tracks the dispatch/combine lifecycle with appropriate fields.
30-205: LGTM! JIT module initialization with custom op registration.The lazy module loading with
@functools.cacheand custom op registration is well-structured. The returnedSimpleNamespaceprovides a clean API surface.
208-246: LGTM! Public wrapper functions are clean delegations.The top-level
moe_a2a_*functions provide clean interfaces to the JIT module, with appropriate docstrings where needed.
249-307: LGTM! Dispatch wrapper with tensor wrapping.The dispatch function correctly wraps the raw offsets into workspace-backed tensors for each payload.
310-359: LGTM! Combine, sanitize, and workspace size wrappers.Clean delegation to the underlying JIT module.
379-409: LGTM! Workspace caching with proper key management.The
get_workspaceclassmethod correctly caches workspaces by configuration tuple, preventing redundant allocations.
411-432: LGTM! Lazy metainfo constant initialization.The
_init_constantsmethod properly strips prefixes for a cleaner Python API.
434-494: LGTM! Constructor with proper validation and MNNVL configuration.Good input validation for
top_kandnum_experts, with optional MnnvlConfig support as discussed in past reviews.
510-567: LGTM! Dispatch method with proper state management.Good state assertions, lifecycle tracking, and optional sanitization flow.
569-609: LGTM! Combine method with state reset.Proper state validation and reset after combine completes, enabling the next dispatch/combine cycle.
611-645: LGTM! Workspace-backed tensor accessor.The
get_combine_payload_tensor_in_workspacemethod correctly computes slice bounds and validates state.
647-654: LGTM! Clean__all__export list.Explicitly defines the public API surface.
|
/bot run |
djns99
force-pushed
the
djns99/update-trtllm-kernels
branch
from
13f86a2 to
f99d8ba
Compare
|
/bot help |
FlashInfer CI BotAvailable Commands:
Authorization:Only whitelisted users can trigger CI. Contact a maintainer for access. How It Works:
Note: Any whitelisted user can trigger CI for any PR, not just their own. |
|
/bot run |
There was a problem hiding this comment.
Actionable comments posted: 2
♻️ Duplicate comments (2)
csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu (1)
748-786: Block size inconsistency between prepare and combine launch functions.
moe_a2a_prepare_combine_launchuses hardcodedkBlockSize = 256, whilemoe_a2a_combine_launch(line 800) usesgetEnvMoeA2ACombineBlockSize(). IfTLLM_MOE_A2A_COMBINE_BLOCK_SIZEis set to a different value, the behavior may be inconsistent.However, reviewing the actual usage: the prepare kernel's block size only affects the grid/block launch configuration for the copy operation, not the actual combine logic. The copy uses
ThreadingPolicywhich adapts to any block size. So while technically inconsistent, the functional impact is limited to performance characteristics of the prepare phase.Consider aligning for consistency:
void moe_a2a_prepare_combine_launch(MoeA2ACombineParams const& params) { - constexpr int kBlockSize = 256; + int const kBlockSize = tensorrt_llm::common::getEnvMoeA2ACombineBlockSize(); - constexpr int kWarpsPerBlock = kBlockSize / 32; + int const kWarpsPerBlock = kBlockSize / 32;flashinfer/comm/__init__.py (1)
42-56: Missing export formoe_a2a_get_metainfo_index_pairs.The function
moe_a2a_get_metainfo_index_pairsis defined intrtllm_moe_alltoall.pyand referenced in the__all__list there, but it's not exported from this module. If this is intended to be part of the public API (as suggested by docs/api/comm.rst), add the export:from .trtllm_moe_alltoall import ( moe_a2a_wrap_payload_tensor_in_workspace as moe_a2a_wrap_payload_tensor_in_workspace, ) +from .trtllm_moe_alltoall import ( + moe_a2a_get_metainfo_index_pairs as moe_a2a_get_metainfo_index_pairs, +)
🧹 Nitpick comments (14)
csrc/nv_internal/cpp/common/envUtils.cpp (1)
335-345: Unreachable code on line 343.After line 339 ensures
block >= 256and line 342 rounds to nearest multiple of 32,blockcan never be 0. The checkif (block == 0) block = 256;on line 343 is unreachable.Consider removing the dead code:
static int sanitizeBlockSize(std::optional<int32_t> const& val) { // Default 256 when not set or invalid int block = val.value_or(256); // Clamp to sane CUDA bounds and warp multiples if (block <= 0) block = 256; if (block > 1024) block = 1024; // Round to nearest multiple of 32 (warp size) block = (block + 31) / 32 * 32; - if (block == 0) block = 256; return block; }csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h (1)
173-179: Comment describes parameters in different order than declaration.The inline comments describe parameters in order:
expert_ids,recv_counters,ep_size,max_tokens_per_rank,top_k,invalid_id. However, the actual declaration placesinvalid_idbeforeep_size. Consider aligning the comment with the actual signature for clarity.// Sanitize expert IDs for invalid tokens // expert_ids: [ep_size, max_tokens_per_rank, top_k] (int32) // recv_counters: [ep_size] (int32), number of valid tokens per source // invalid_id: value to fill for invalid tokens' expert ids +// Parameters: expert_ids, recv_counters, invalid_id, ep_size, max_tokens_per_rank, top_k, stream void moe_a2a_sanitize_expert_ids_launch(int32_t* expert_ids, int32_t const* recv_counters, int32_t invalid_id, int ep_size, int max_tokens_per_rank, int top_k, cudaStream_t stream);csrc/trtllm_moe_alltoall.cu (1)
95-117: Consider removing the stream synchronization or documenting why it's needed.Line 113-114 performs a
cudaStreamSynchronizeaftercudaMemsetAsync. If this is for correctness (ensuring the memset completes before returning the metainfo tensor), it could potentially be relaxed since subsequent kernel launches on the same stream will see the zeroed memory. However, if the metainfo tensor is accessed on the CPU immediately after this call, the sync is necessary.If the sync is intentional for a specific reason, consider adding a brief comment explaining why.
flashinfer/comm/trtllm_moe_alltoall.py (5)
386-389: Annotate mutable class attribute withClassVar.Per PEP 526 and Ruff RUF012, mutable class attributes should be annotated with
typing.ClassVarto indicate they are class-level, not instance-level.+from typing import ClassVar + class MoeAlltoAll: ... # Single shared workspace across the process - _WORKSPACE_CACHE: dict[tuple[int, int, int, int], dict] = {} + _WORKSPACE_CACHE: ClassVar[dict[tuple[int, int, int, int], dict]] = {}Based on static analysis hints.
487-496: Use explicitOptionalor| Nonefor optional parameters.PEP 484 prohibits implicit
Optional(using= Nonewithout the type hint includingNone). Update the type hints for clarity.def __init__( self, mapping: Mapping, max_num_tokens: int, top_k: int, num_experts: int, - workspace_size_per_rank: int = None, - hidden_size: int = None, + workspace_size_per_rank: int | None = None, + hidden_size: int | None = None, mnnvl_config: Optional[MnnvlConfig] = None, ):Based on static analysis hints.
698-708: Avoid creating empty tensor just to get element size.Creating an empty tensor to get
element_size()is inefficient. Usetorch.finfoortorch.iinfoinstead, or use thedtype.itemsizeattribute directly (available in PyTorch).- element_size = torch.tensor([], dtype=dtype).element_size() + element_size = torch.empty(0, dtype=dtype).element_size() # or use dtype-specific size lookupOr more efficiently:
element_size = torch.empty(0, dtype=dtype).element_size()Note:
torch.dtypeobjects don't have a directitemsizeattribute in PyTorch, but usingtorch.empty(0, dtype=dtype).element_size()avoids the list literal.
229-236: Bounds check uses assert which can be disabled.The assertion at line 231-233 will be skipped if Python is run with
-Oflag. For production code that handles user input, consider using a proper exception.- assert slice_end <= workspace.numel(), ( - f"slice_end {slice_end} exceeds workspace size {workspace.numel()}" - ) + if slice_end > workspace.numel(): + raise ValueError( + f"slice_end {slice_end} exceeds workspace size {workspace.numel()}" + )
560-572: Consider using a more graceful cleanup pattern for_reset_workspace.The
_reset_workspacemethod deletesself._WORKSPACEbefore accessingself.workspace_size_per_rank, which could cause issues if the method is called incorrectly. Also, settingphase = "deleted"is good for safety, but the object attributes likeworkspaceandmetainfowill still reference stale data.Consider either nullifying those references or documenting more explicitly that the object is unusable after this call.
csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu (1)
111-116: Timeout value assumes GPU clock speed.The timeout calculation
300ll * 2000ll * 1000ll * 1000llassumes a 2000 MHz clock. While this is a reasonable worst-case assumption and 300 seconds is generous, the comment could be more precise about the calculation (e.g., "300 seconds * 2000 MHz").tests/comm/test_mnnvl_moe_alltoall.py (5)
185-223: Address the TODO comment.The TODO on line 206 indicates that FP8 (
.to(torch.float8_e4m3fn)) testing is currently disabled due to test failures.Do you want me to open an issue to track enabling FP8 testing, or would you like help investigating why the test fails with FP8?
572-577: Consider more specific exception handling.While catching bare
Exceptionforpytest.skipis a common pattern, catching a more specific exception (e.g.,RuntimeError,ImportError) would be more precise if you know what errors indicate MNNVL unavailability.If you can identify the specific exceptions raised when MNNVL is unavailable, apply this pattern:
try: MnnvlMemory.initialize() if not MnnvlMemory.supports_mnnvl(): pytest.skip("MNNVL not supported on this system") - except Exception: + except (RuntimeError, ImportError, OSError): pytest.skip("MNNVL not supported on this system")
673-678: Consider more specific exception handling.Same as the earlier comment: if specific exception types indicate MNNVL unavailability, catching those explicitly would be more precise.
802-814: Consider documenting the tolerance rationale more clearly.The match rate threshold and tolerance values (2e-2) are chosen to handle bfloat16 accumulation order differences. While the comment on lines 810-811 explains this, consider adding a brief inline comment near the tolerance values themselves.
# Verify against reference num_matches = ( - torch.isclose(combined_output, reference_output, atol=2e-2, rtol=2e-2) + torch.isclose(combined_output, reference_output, atol=2e-2, rtol=2e-2) # Tolerant to bf16 accumulation order .sum() .item() ) match_rate = num_matches / combined_output.numel() - match_threshold = 0.99 + match_threshold = 0.99 # Allow 1% mismatches due to non-deterministic reduction order
816-816: Remove or uncomment the assertion.The commented-out assertion on line 816 should either be removed (if the match rate check above is sufficient) or uncommented (if stricter validation is needed).
If the match rate check (lines 802-814) is the intended validation approach, remove the commented line:
- # torch.testing.assert_close(combined_output, reference_output, rtol=6e-2, atol=6e-2) - check_any_rank_failed()
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (17)
csrc/nv_internal/cpp/common/envUtils.cpp(2 hunks)csrc/nv_internal/tensorrt_llm/common/envUtils.h(2 hunks)csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu(1 hunks)csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h(1 hunks)csrc/nv_internal/tensorrt_llm/thop/moeAlltoAllMeta.h(1 hunks)csrc/trtllm_moe_alltoall.cu(1 hunks)docs/api/comm.rst(1 hunks)flashinfer/aot.py(1 hunks)flashinfer/comm/__init__.py(1 hunks)flashinfer/comm/trtllm_moe_alltoall.py(1 hunks)flashinfer/jit/__init__.py(1 hunks)flashinfer/jit/comm.py(1 hunks)scripts/task_test_multi_node_comm_kernels.sh(1 hunks)scripts/task_test_single_node_comm_kernels.sh(1 hunks)tests/comm/test_mnnvl_memory.py(1 hunks)tests/comm/test_mnnvl_moe_alltoall.py(1 hunks)tests/comm/test_trtllm_moe_alltoall.py(1 hunks)
🚧 Files skipped from review as they are similar to previous changes (6)
- docs/api/comm.rst
- scripts/task_test_multi_node_comm_kernels.sh
- scripts/task_test_single_node_comm_kernels.sh
- csrc/nv_internal/tensorrt_llm/common/envUtils.h
- flashinfer/jit/init.py
- flashinfer/aot.py
🧰 Additional context used
🧬 Code graph analysis (5)
csrc/nv_internal/cpp/common/envUtils.cpp (1)
include/flashinfer/trtllm/common.h (1)
getBoolEnv(195-198)
tests/comm/test_mnnvl_memory.py (1)
flashinfer/comm/mapping.py (1)
local_rank(391-392)
flashinfer/comm/__init__.py (1)
flashinfer/comm/trtllm_moe_alltoall.py (11)
MoeAlltoAll(373-708)moe_a2a_combine(97-138)moe_a2a_combine(301-324)moe_a2a_dispatch(51-91)moe_a2a_dispatch(240-298)moe_a2a_initialize(39-45)moe_a2a_initialize(199-207)moe_a2a_get_workspace_size_per_rank(339-370)moe_a2a_sanitize_expert_ids(144-153)moe_a2a_sanitize_expert_ids(327-336)moe_a2a_wrap_payload_tensor_in_workspace(210-237)
csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h (2)
csrc/nv_internal/tensorrt_llm/common/envUtils.h (1)
tensorrt_llm(23-104)csrc/nv_internal/tensorrt_llm/thop/moeAlltoAllMeta.h (1)
mnnvl_throughput(25-58)
csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu (1)
csrc/nv_internal/cpp/common/envUtils.cpp (4)
getEnvMoeA2ADispatchBlockSize(347-350)getEnvMoeA2ADispatchBlockSize(347-347)getEnvMoeA2ACombineBlockSize(352-355)getEnvMoeA2ACombineBlockSize(352-352)
🪛 Clang (14.0.6)
csrc/nv_internal/tensorrt_llm/thop/moeAlltoAllMeta.h
[error] 19-19: 'array' file not found
(clang-diagnostic-error)
csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h
[error] 18-18: 'cuda_bf16.h' file not found
(clang-diagnostic-error)
🪛 Ruff (0.14.7)
flashinfer/comm/trtllm_moe_alltoall.py
388-388: Mutable class attributes should be annotated with typing.ClassVar
(RUF012)
493-493: PEP 484 prohibits implicit Optional
Convert to T | None
(RUF013)
494-494: PEP 484 prohibits implicit Optional
Convert to T | None
(RUF013)
533-533: Avoid specifying long messages outside the exception class
(TRY003)
535-535: Avoid specifying long messages outside the exception class
(TRY003)
696-698: Avoid specifying long messages outside the exception class
(TRY003)
tests/comm/test_mnnvl_moe_alltoall.py
34-34: Avoid specifying long messages outside the exception class
(TRY003)
576-576: Do not catch blind exception: Exception
(BLE001)
677-677: Do not catch blind exception: Exception
(BLE001)
711-711: Unpacked variable expert_id_payload_index is never used
Prefix it with an underscore or any other dummy variable pattern
(RUF059)
🔇 Additional comments (50)
flashinfer/jit/comm.py (1)
83-109: LGTM!The new
gen_mnnvl_moe_alltoall_module()function follows the established pattern of other JIT module generators in this file. The source file paths and include paths are correctly constructed using thePathoperator.csrc/nv_internal/cpp/common/envUtils.cpp (2)
326-333: LGTM!The
getEnvMoeA2AOneBlockPerToken()function correctly implements the "default true, false only if explicitly set to 0" semantics with proper static caching.
347-357: LGTM!The block size getters and
getEnvEplbForceGdrcopy()follow established patterns in this file with proper static initialization for thread-safe caching.csrc/nv_internal/tensorrt_llm/thop/moeAlltoAllMeta.h (2)
28-41: LGTM!The enum values are sequentially numbered and
NUM_METAINFO_FIELDS = 9correctly matches the count of preceding fields.
45-58: LGTM!The helper function correctly maps each enum value to its string name for Python/C++ interop. The inline function returning by value is appropriate for this infrequently-called utility.
csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h (4)
17-19: LGTM!The CUDA type headers are correctly included. The static analysis error about
cuda_bf16.hnot found is a false positive—this is a standard CUDA Toolkit header that will be available when compiling with nvcc.
23-27: LGTM!The configuration constants define sensible upper bounds for the MoE A2A kernels. Fixed-size arrays with these bounds are appropriate for device-side kernel parameter structures.
78-120: LGTM!The
MoeA2ADispatchParamsstruct comprehensively captures all parameters needed for the dispatch phase including threading policy, EP configuration, token routing, payload descriptors, and synchronization primitives.
146-148: Verify TensorRT header inclusion fornvinfer1::DataType.The
dtypefield usesnvinfer1::DataTypefrom the TensorRT SDK. Ensure thatnvinfer.hor an equivalent TensorRT header providing this type is included either directly in this file or transitively through other includes, and that TensorRT headers are properly configured in the project's build system.tests/comm/test_mnnvl_memory.py (1)
125-125: LGTM!The change to use
self.local_rankinstead ofself.rankfor device selection is correct. This aligns with the setup fixture (line 51) and ensures proper device mapping in multi-node configurations whererankmay exceed the local GPU count.csrc/trtllm_moe_alltoall.cu (6)
1-52: LGTM: Header includes and utility definitions.The file structure is clean with appropriate includes for TVM FFI, CUDA, and the internal TensorRT-LLM headers. The anonymous namespace properly encapsulates internal implementation details, and the alignment helper function is correct.
53-93: LGTM: Workspace offset calculation.The
calculateOffsetsfunction correctly computes the memory layout for the A2A workspace with proper alignment padding to 128 bytes at key boundaries. The offsets are computed sequentially and the function is well-structured.
119-255: LGTM: Dispatch operation implementation.The dispatch operation has thorough input validation, correct payload descriptor preparation, and proper workspace size checking. The alignment check at line 183-185 with the improved error message is helpful. The kernel launch error checking is correctly placed after the launch calls.
257-354: LGTM: Combine operation implementation.The combine operation correctly validates inputs, handles the
payloadInWorkspaceflag with a pointer check, allocates the output tensor, and sets up the combine parameters. Error checking after kernel launches is properly implemented.
356-388: LGTM: Sanitize expert IDs operation.The sanitize operation correctly extracts dimensions from the expert_ids tensor and calls the underlying kernel with proper error checking.
390-413: LGTM: Metainfo index pairs and FFI exports.The
getMoeA2AMetaInfoIndexPairsfunction correctly wraps the internal helper and returns a TVM-compatible tuple. All FFI exports are properly registered.flashinfer/comm/trtllm_moe_alltoall.py (5)
711-719:__all__is missingmoe_a2a_get_metainfo_index_pairs.The function
moe_a2a_get_metainfo_index_pairsis defined but not included in__all__. If it's intended to be private, that's fine. If it should be public (for testing/debugging as mentioned in the class docstring), add it to__all__.
1-197: LGTM: Module setup and JIT module wrapper.The module docstring is clear, the
@functools.cachedecorator correctly ensures the JIT module is built once, and the custom op registrations follow the expected pattern. TheSimpleNamespacereturn cleanly bundles the registered functions.
373-462: LGTM: MoeAlltoAll class workspace management.The workspace cache pattern correctly handles multiple communicators with different configurations. The
get_moe_workspace_size_per_rankstatic method provides a sensible default calculation. The initialization flow with MnnvlMemory is well-structured.
574-631: LGTM: Dispatch method with proper state tracking.The dispatch method correctly enforces the state machine (idle → dispatched), validates runtime parameters, and optionally sanitizes expert IDs. State updates are properly sequenced after the kernel call.
633-673: LGTM: Combine method with state reset.The combine method correctly checks state, calls the underlying operation, and resets state for the next round.
tests/comm/test_trtllm_moe_alltoall.py (10)
28-32: LGTM: Fixture docstring corrected.The docstring now accurately describes that this fixture sets up the torch seed for deterministic tests.
66-75: LGTM: SM count guard prevents test hangs.The
check_sufficient_sm_countfunction correctly skips tests that require more SMs than available, preventing potential hangs when simulating multiple ranks on a single GPU.
77-90: Consider usingtorch.randintconsistently.The
make_payloadfunction has a subtle issue: for integer types, it usestorch.iinfo(dtype).minandtorch.iinfo(dtype).maxas bounds, buttorch.randintuses half-open interval[low, high). Fortorch.uint8,maxis 255, so the actual range is[0, 255)which excludes 255. This is likely fine for testing but worth noting.
100-169: LGTM: Single GPU test is comprehensive.The test covers the full dispatch-combine cycle with multiple data types, validates output correctness using sorting to handle shuffling, and tests the in-workspace combine path.
171-247: LGTM: Multi-rank dispatch helper.The helper function correctly sets up per-rank workspaces, initializes metainfo, and executes dispatch in parallel using CUDA streams. The synchronization pattern is correct.
309-351: LGTM: Multi-rank dispatch test.The test validates that tokens are correctly routed to their target ranks and that the payload data is preserved.
354-395: LGTM: Sanitize expert IDs test.The test verifies that invalid expert IDs (tokens not belonging to a rank) are correctly set to the invalid marker value.
397-438: LGTM: Fake MoE reference implementation.The
fake_moefunction provides a clear reference implementation that applies expert-specific scaling and handles expert parallelism correctly.
441-558: LGTM: End-to-end combine test.The test exercises the full dispatch-sanitize-process-combine pipeline and validates against the reference implementation with reasonable tolerances for floating-point accumulation.
561-614: LGTM: Workspace size calculation test.The test verifies consistency between the low-level and high-level workspace size calculation APIs, and validates the relationship between aux data size and total workspace size.
csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu (10)
1-51: LGTM: Header and utility macros.The file has proper copyright header, includes, and well-structured utility macros (
SWITCH_BOOL,SWITCH_TOP_K, etc.) for compile-time specialization.
140-158: LGTM: Threading policy abstractions.The
WarpPolicyandBlockPolicystructs cleanly abstract the threading model differences, allowing the same kernel code to work with both warp-per-token and block-per-token configurations.
160-189: LGTM: Vectorized copy implementation.The vectorized copy correctly selects vector width based on alignment and uses the flashinfer vec_t abstraction for efficient memory operations.
264-277: Non-atomic increment offlag_val_ptrin prepare kernel.Line 275 uses a non-atomic increment:
*flag_val_ptr = *flag_val_ptr + 1. Since this is guarded byidx == 0, it's safe in a single-kernel context. However, if multiple prepare kernels could run concurrently (e.g., from different streams), this would be a race. Assuming the dispatch/combine workflow is strictly sequential per communicator, this is fine.
287-434: LGTM: Dispatch kernel implementation.The dispatch kernel correctly handles per-token routing with TOP_K fanout, uses shared memory to stage routing data, implements proper synchronization barriers, and handles the cross-rank completion signaling with appropriate memory fences. The early-exit path for already-copied ranks (line 321) is a good optimization.
512-635: LGTM: Vectorized combine implementation.The
vectorized_combine_implcorrectly implements the TOP_K reduction with specialized unrolled paths for common values (1, 2, 4, 8) and a generic fallback. The accumulation is done in registers before a single store.
637-669: LGTM: Prepare combine kernel.The kernel correctly copies payload data to receive buffers with proper bounds checking against
recv_counters.
675-746: LGTM: Combine kernel with synchronization.The combine kernel implements proper cross-rank synchronization using flag values, with acquire/release semantics and timeout handling. The
__syncthreads()after the synchronization loop ensures all threads in the block see the updated data.
792-842: LGTM: Combine launch function.The launch function correctly validates parameters, sets up kernel pointers, and uses the macro-based dispatch to select the appropriate kernel specialization.
844-872: LGTM: Sanitize expert IDs kernel and launch.The sanitize kernel correctly identifies invalid tokens (beyond
recv_counters) and sets their expert IDs to the invalid marker.tests/comm/test_mnnvl_moe_alltoall.py (9)
49-52: LGTM!The autouse fixture ensures deterministic test behavior by setting a fixed random seed.
55-67: LGTM!The contiguous expert partitioning logic is correct and well-documented.
70-88: LGTM!The expert selection logic correctly generates random top-k expert assignments per token, handling the empty token case appropriately.
91-119: LGTM!The per-expert seeding strategy ensures reproducible weight initialization across test runs while maintaining diversity between experts.
122-182: LGTM!The MoE emulation correctly handles both EP-aware and global modes, properly filtering experts by rank ownership and accumulating scaled outputs.
226-258: LGTM!The bfloat16 payload structure mirrors the FP4 version appropriately, with rank-specific patterns to aid verification.
386-556: LGTM!The verification function performs comprehensive validation including counter consistency, payload routing correctness, and expert ID sanitization for invalid tokens.
711-711: Static analysis false positive.Line 711's
expert_id_payload_indexis used on line 617, so the Ruff warning (RUF059) is incorrect. No action needed.
294-294: Verify if_WORKSPACEreset is necessary for test isolation.Setting
MoeAlltoAll._WORKSPACE = Nonemanipulates class-level state at line 294. Confirm whether:
- This is intentional cleanup in a teardown/cleanup context
- The workspace is designed as global/shared or if it should be per-instance
- Test isolation/parallelization (e.g., pytest-xdist) requires this reset
If this is cleanup code within a test function or fixture teardown, the concern may be unfounded; however, ensure the reset doesn't interfere with concurrent test execution if parallelization is enabled.
|
[FAILED] Pipeline #39628456: 4/20 passed |
6 participants
📌 Description
This ports the latest MNNVL A2A communication implementation from TRT-LLM
🔍 Related Issues
#2094
🚀 Pull Request Checklist
Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete.
✅ Pre-commit Checks
pre-commitby runningpip install pre-commit(or used your preferred method).pre-commit install.pre-commit run --all-filesand fixed any reported issues.🧪 Tests
unittest, etc.).Reviewer Notes
Summary by CodeRabbit
New Features
Configuration
Documentation
Tests
✏️ Tip: You can customize this high-level summary in your review settings.