diff --git a/benchmarks/single_node/fixed_seq_len/minimaxm3_fp4_b300.sh b/benchmarks/single_node/fixed_seq_len/minimaxm3_fp4_b300.sh index f91419edb..ad7ed2059 100755 --- a/benchmarks/single_node/fixed_seq_len/minimaxm3_fp4_b300.sh +++ b/benchmarks/single_node/fixed_seq_len/minimaxm3_fp4_b300.sh @@ -3,7 +3,13 @@ # MiniMax-M3 NVFP4 B300 single-node vLLM recipe. # Same shape as minimaxm3_fp8_b300.sh but uses the nvidia/MiniMax-M3-NVFP4 # checkpoint. MiniMax-M3 modelopt NVFP4 support (vllm-project/vllm PR #46380) is -# baked into the perf container image, so no runtime patch is needed. +# baked into the perf container image. +# +# At runtime the recipe swaps the image's FlashInfer for the first pinned +# nightly containing the upstream SM100 low-M MXFP8 split-K kernel +# (flashinfer-ai/flashinfer#3847), then backports the AutoTuner non-Tensor guard +# fix from flashinfer-ai/flashinfer#3918 and reverts the communication workspace +# changes from flashinfer-ai/flashinfer#3745 for performance isolation. source "$(dirname "$0")/../../benchmark_lib.sh" @@ -19,6 +25,43 @@ check_env_vars \ RANDOM_RANGE_RATIO \ RESULT_FILENAME +# --- FlashInfer nightly + targeted runtime patches -------------------------- +FLASHINFER_VERSION=0.6.15.dev20260710 +FLASHINFER_NIGHTLY_TAG=nightly-v0.6.15-20260710 +FLASHINFER_RELEASE_URL="https://github.com/flashinfer-ai/flashinfer/releases/download/${FLASHINFER_NIGHTLY_TAG}" + +python3 -m pip uninstall -y flashinfer-python flashinfer-cubin flashinfer-jit-cache + +python3 -m pip install \ + "${FLASHINFER_RELEASE_URL}/flashinfer_python-${FLASHINFER_VERSION}-py3-none-any.whl" \ + "${FLASHINFER_RELEASE_URL}/flashinfer_cubin-${FLASHINFER_VERSION}-py3-none-any.whl" \ + "${FLASHINFER_RELEASE_URL}/flashinfer_jit_cache-${FLASHINFER_VERSION}+cu130-cp39-abi3-manylinux_2_28_$(uname -m).whl" \ + || { echo "FlashInfer nightly install failed" >&2; exit 1; } + +# The pinned nightly predates flashinfer-ai/flashinfer#3918. Apply only its +# runtime fix; the upstream test change is intentionally not backported. +AUTOTUNER_PATCH="$(dirname "$0")/patches/flashinfer-autotuner-non-tensor-guard.patch" +if ! command -v patch >/dev/null 2>&1; then + apt-get update -y && apt-get install -y --no-install-recommends patch \ + || { echo "Failed to install patch(1)" >&2; exit 1; } +fi +SITE_PACKAGES=$(dirname "$(python3 -c "import importlib.util; print(importlib.util.find_spec('flashinfer').submodule_search_locations[0])")") \ + || { echo "Could not locate the installed flashinfer package" >&2; exit 1; } +patch --dry-run -p1 -d "${SITE_PACKAGES}" < "${AUTOTUNER_PATCH}" >/dev/null \ + || { echo "FlashInfer AutoTuner non-Tensor guard patch does not apply" >&2; exit 1; } +patch -p1 -d "${SITE_PACKAGES}" < "${AUTOTUNER_PATCH}" \ + || { echo "FlashInfer AutoTuner non-Tensor guard patch failed" >&2; exit 1; } + +# Reverse the runtime communication changes from flashinfer-ai/flashinfer#3745 +# while leaving the 0.6.15 GEMM, MoE, and attention code unchanged. +COMM_REVERT_PATCH="$(dirname "$0")/patches/flashinfer-revert-pr-3745.patch" +patch --dry-run -p1 -d "${SITE_PACKAGES}" < "${COMM_REVERT_PATCH}" >/dev/null \ + || { echo "FlashInfer PR #3745 revert patch does not apply" >&2; exit 1; } +patch -p1 -d "${SITE_PACKAGES}" < "${COMM_REVERT_PATCH}" \ + || { echo "FlashInfer PR #3745 revert patch failed" >&2; exit 1; } + +# ----------------------------------------------------------------------------- + if [[ -n "${MODEL_PATH:-}" ]]; then if [[ ! -d "$MODEL_PATH" || -z "$(ls -A "$MODEL_PATH" 2>/dev/null)" ]]; then hf download "$MODEL" --local-dir "$MODEL_PATH" @@ -39,6 +82,7 @@ SERVER_LOG=/workspace/server.log export VLLM_ENGINE_READY_TIMEOUT_S=3600 export VLLM_FLOAT32_MATMUL_PRECISION=high export VLLM_FLASHINFER_ALLREDUCE_BACKEND=trtllm +export VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS=1800 if [ "${DP_ATTENTION}" = "true" ]; then PARALLEL_ARGS="--tensor-parallel-size=1 --data-parallel-size=$TP --enable-expert-parallel" @@ -57,6 +101,7 @@ start_gpu_monitor set -x vllm serve "$MODEL_PATH" --served-model-name "$MODEL" --host 0.0.0.0 --port $PORT \ $PARALLEL_ARGS \ +--attention_config.indexer_kv_dtype fp8 \ --gpu-memory-utilization 0.95 \ --max-model-len $MAX_MODEL_LEN \ --kv-cache-dtype fp8 \ diff --git a/benchmarks/single_node/fixed_seq_len/patches/flashinfer-autotuner-non-tensor-guard.patch b/benchmarks/single_node/fixed_seq_len/patches/flashinfer-autotuner-non-tensor-guard.patch new file mode 100644 index 000000000..711af8b84 --- /dev/null +++ b/benchmarks/single_node/fixed_seq_len/patches/flashinfer-autotuner-non-tensor-guard.patch @@ -0,0 +1,38 @@ +diff --git a/flashinfer/autotuner/autotuner.py b/flashinfer/autotuner/autotuner.py +index 54b90a0..76443ca 100644 +--- a/flashinfer/autotuner/autotuner.py ++++ b/flashinfer/autotuner/autotuner.py +@@ -2074,14 +2074,19 @@ class AutoTuner: + def _prepare_input_tensors_with_batches( + self, +- inputs: list[torch.Tensor], ++ inputs: list[Any], + tuning_config: TuningConfig, +- ) -> list[list[torch.Tensor]]: ++ ) -> list[list[Any]]: + """Create multiple input copies to flush the L2 cache between profiling iterations.""" + if not tuning_config.use_cold_l2_cache: + return [inputs] + +- one_buffer_bytes = sum(input.numel() * input.element_size() for input in inputs) ++ one_buffer_bytes = sum( ++ input.numel() * input.element_size() ++ if isinstance(input, torch.Tensor) ++ else 0 ++ for input in inputs ++ ) + if one_buffer_bytes <= 0: + logger.debug( + "[Autotuner] No tensor inputs or zero-sized tensors; falling back to single-batch profiling." + ) +@@ -2093,7 +2098,9 @@ class AutoTuner: + + inputs_list = [inputs] + for _ in range(num_buffers - 1): +- inputs_list.append(list(t.clone() for t in inputs)) ++ inputs_list.append( ++ [t.clone() if isinstance(t, torch.Tensor) else t for t in inputs] ++ ) + + logger.debug( + f"[Autotuner] use_cold_l2_cache={tuning_config.use_cold_l2_cache}, use {num_buffers} different tensors for profiling" diff --git a/benchmarks/single_node/fixed_seq_len/patches/flashinfer-revert-pr-3745.patch b/benchmarks/single_node/fixed_seq_len/patches/flashinfer-revert-pr-3745.patch new file mode 100644 index 000000000..3629ea80f --- /dev/null +++ b/benchmarks/single_node/fixed_seq_len/patches/flashinfer-revert-pr-3745.patch @@ -0,0 +1,672 @@ +diff --git a/flashinfer/comm/allreduce.py b/flashinfer/comm/allreduce.py +index 87295b870e25cd534ab1dc35d1782c7e5fc05f25..c9fe5400f1196594770175fe5006d83ec833b5d7 100644 +--- a/flashinfer/comm/allreduce.py ++++ b/flashinfer/comm/allreduce.py +@@ -63,7 +63,6 @@ from flashinfer.trace.templates.comm import allreduce_fusion_trace + + from .trtllm_ar import trtllm_allreduce_fusion + from .trtllm_ar import trtllm_create_ipc_workspace_for_all_reduce_fusion +-from .trtllm_ar import _initialize_allreduce_fusion_protocol + from .trtllm_ar import check_trtllm_allreduce_fusion_workspace_metadata + from .trtllm_ar import trtllm_moe_allreduce_fusion + from .trtllm_ar import trtllm_moe_finalize_allreduce_fusion +@@ -179,59 +178,6 @@ class TRTLLMAllReduceFusionWorkspace(AllReduceFusionWorkspace): + logger.warning("Workspace is insufficient for problem size. %s", e) + return False + +- @flashinfer_api +- def checkpoint_prepare(self) -> None: +- """Detach physical backing; repeated successful calls are no-ops.""" +- if not self.mem_handles or not all( +- isinstance(handle, SymmDeviceMemory) for handle in self.mem_handles +- ): +- raise NotImplementedError( +- "Stable-VA checkpointing is unavailable for workspaces backed " +- "by torch symmetric memory" +- ) +- +- mapped = [handle.mapped for handle in self.mem_handles] +- if not any(mapped): +- return +- if not all(mapped): +- raise RuntimeError("TRT-LLM symmetric-memory handle state is inconsistent") +- +- for handle in self.mem_handles: +- handle._unmap_and_release_handles() +- # Do not return until every rank has released all workspace handles. +- self.mem_handles[0].comm_backend.barrier() +- +- @flashinfer_api +- def checkpoint_restore(self, comm_backend: CommBackend) -> None: +- """Restore physical backing; repeated successful calls are no-ops.""" +- if not self.mem_handles or not all( +- isinstance(handle, SymmDeviceMemory) for handle in self.mem_handles +- ): +- raise NotImplementedError( +- "Stable-VA checkpointing is unavailable for workspaces backed " +- "by torch symmetric memory" +- ) +- +- mapped = [handle.mapped for handle in self.mem_handles] +- if all(mapped): +- return +- if any(mapped): +- raise RuntimeError("TRT-LLM symmetric-memory handle state is inconsistent") +- for handle in self.mem_handles: +- handle._create_and_map_handles(comm_backend) +- +- _initialize_allreduce_fusion_protocol( +- ipc_handles=self.ipc_handles, +- tp_rank=self.rank, +- flag_size=self.metadata["flag_size"], +- lamport_buffer_size=self.metadata["lamport_buffer_size"], +- lamport_comm_size=self.metadata["lamport_comm_size"], +- use_fp32_lamport=self.metadata["use_fp32_lamport"], +- control_flag_ptr=self.metadata["control_flag_ptr"], +- ) +- torch.cuda.synchronize() +- comm_backend.barrier() +- + def destroy(self) -> None: + """Destroy workspace and free resources.""" + if getattr(self, "_destroyed", False): +@@ -730,11 +676,6 @@ def allreduce_fusion( + # Dispatch based on workspace type + if isinstance(workspace, TRTLLMAllReduceFusionWorkspace): + # TensorRT-LLM backend implementation +- if any( +- isinstance(handle, SymmDeviceMemory) and not handle.mapped +- for handle in workspace.mem_handles +- ): +- raise RuntimeError("TRT-LLM symmetric-memory handles are not attached") + + # ---- MOE Reduction pattern ---- + if pattern == AllReduceFusionPattern.kMoEReductionARResidualRMSNorm: +diff --git a/flashinfer/comm/mnnvl.py b/flashinfer/comm/mnnvl.py +index e127b813ced8475550637fc8d4295df5f5bd8f38..5dec41b6f5477f79488748f181231f041c222259 100644 +--- a/flashinfer/comm/mnnvl.py ++++ b/flashinfer/comm/mnnvl.py +@@ -1004,7 +1004,6 @@ class SymmDeviceMemory: + self.signal_pad_offset = 0 + self.allocation_size = 0 + self.comm_backend = comm_backend_for_handle_transfer or MPIBackend() +- self._enable_multicast = enable_multicast + + # CUDA memory handles and pointers + self.mc_ptr = 0 # CUdeviceptr mMcPtr +@@ -1016,24 +1015,22 @@ class SymmDeviceMemory: + self.uc_handles: List[ + int + ] = [] # std::vector mUcHandles +- self._mapped = False + + # Signal pad constants + self.SIGNAL_PAD_ALIGNMENT = 16 + self.SIGNAL_PAD_SIZE = SIGNAL_PAD_SIZE + + # Check if device supports multicasting +- if self._enable_multicast: +- multicast_supported = checkCudaErrors( +- cuda.cuDeviceGetAttribute( +- cuda.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_MULTICAST_SUPPORTED, +- device_idx, +- ) ++ multicast_supported = checkCudaErrors( ++ cuda.cuDeviceGetAttribute( ++ cuda.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_MULTICAST_SUPPORTED, ++ device_idx, ++ ) ++ ) ++ if multicast_supported == 0: ++ raise RuntimeError( ++ "[SymmDeviceMemory] Device does not support multicasting." + ) +- if multicast_supported == 0: +- raise RuntimeError( +- "[SymmDeviceMemory] Device does not support multicasting." +- ) + + # Calculate signal pad offset with alignment (matching C++ exactly) + self.signal_pad_offset = round_up(buf_size, self.SIGNAL_PAD_ALIGNMENT) +@@ -1044,8 +1041,16 @@ class SymmDeviceMemory: + f"Signal pad offset: {self.signal_pad_offset}" + ) + +- self._exchanger: Optional[HandleExchanger] = None +- self._create_and_map_handles(self.comm_backend) ++ # Create handle exchanger ++ if is_mnnvl_fabric_supported(device_idx): ++ self._exchanger: HandleExchanger = FabricHandleExchanger( ++ self.comm_backend, self.group_rank, self.group_size ++ ) ++ else: ++ self._exchanger = PosixFDHandleExchanger( ++ self.comm_backend, self.group_rank, self.group_size ++ ) ++ self._alloc_mn_mcast_mem(buf_size, enable_multicast) + + if allocate_signal_pads: + # Initialize signal pads +@@ -1063,7 +1068,7 @@ class SymmDeviceMemory: + def __del__(self): + """Destructor - cleanup allocated memory""" + +- if hasattr(self, "_exchanger") and self._exchanger is not None: ++ if hasattr(self, "_exchanger"): + self._exchanger.close() + + # Skip cleanup during Python finalization to avoid segfaults +@@ -1121,8 +1126,6 @@ class SymmDeviceMemory: + checkCudaErrors(cuda.cuMemRelease(self.mc_handle)) + except Exception as e: + logger.warning("Destructor: Failed to release MC handle: %s", e) +- elif hasattr(self, "mc_ptr") and self.mc_ptr: +- checkCudaErrors(cuda.cuMemAddressFree(self.mc_ptr, self.allocation_size)) + + def get_signal_pad_ptrs_host(self) -> List[int]: + """Get the raw array of signal pad pointers to all ranks (including self)""" +@@ -1172,65 +1175,20 @@ class SymmDeviceMemory: + """Get the usable buffer size (excluding signal pad)""" + return self.allocation_size - self.SIGNAL_PAD_SIZE + +- @property +- def mapped(self) -> bool: +- return self._mapped +- +- def _create_and_map_handles(self, comm: CommBackend) -> None: +- """Create physical backing and map it at the reserved addresses.""" +- # Create handle exchanger +- if is_mnnvl_fabric_supported(self.device_idx): +- self._exchanger = FabricHandleExchanger( +- comm, self.group_rank, self.group_size +- ) +- else: +- self._exchanger = PosixFDHandleExchanger( +- comm, self.group_rank, self.group_size +- ) +- ++ def _alloc_mn_mcast_mem(self, buf_size: int, enable_multicast: bool): ++ """Allocate multi-node multicast memory using MNNVL""" + self._verify_cuda_context() + + # Compute allocation size and get allocation properties +- allocation_prop, mc_prop = self._get_allocation_prop(self.buf_size) ++ allocation_prop, mc_prop = self._get_allocation_prop(buf_size) + + # Allocate, exchange, and map unicast buffers + self._allocate_unicast_buffers(allocation_prop) + + # Setup multicast object, exchange handles, map and bind memory +- if self._enable_multicast: ++ if enable_multicast: + self._setup_multicast(mc_prop) + +- self.comm_backend = comm +- self._mapped = True +- +- def _unmap_and_release_handles(self) -> None: +- """Unmap and release physical backing while retaining reserved addresses.""" +- # Drain local work, then align ranks before changing shared mappings. +- cuda.cuCtxSynchronize() +- self.comm_backend.barrier() +- +- if self._enable_multicast: +- checkCudaErrors( +- cuda.cuMulticastUnbind( +- self.mc_handle, self.device_idx, 0, self.allocation_size +- ) +- ) +- checkCudaErrors(cuda.cuMemUnmap(self.mc_ptr, self.allocation_size)) +- +- for ptr in self.uc_ptrs: +- checkCudaErrors(cuda.cuMemUnmap(ptr, self.allocation_size)) +- +- if self._enable_multicast: +- checkCudaErrors(cuda.cuMemRelease(self.mc_handle)) +- self.mc_handle = 0 +- for handle in self.uc_handles: +- checkCudaErrors(cuda.cuMemRelease(handle)) +- +- self._exchanger.close() +- self.uc_handles = [0] * self.group_size +- self._exchanger = None +- self._mapped = False +- + def _verify_cuda_context(self): + """Verify CUDA context is set to the correct device.""" + try: +@@ -1268,23 +1226,20 @@ class SymmDeviceMemory: + buf_size + self.SIGNAL_PAD_SIZE, alloc_granularity + ) + +- self._mc_granularity = alloc_granularity +- mc_prop = None +- if self._enable_multicast: +- # Set up multicast properties +- mc_prop = cuda.CUmulticastObjectProp() +- mc_prop.numDevices = self.group_size +- mc_prop.size = self.allocation_size +- mc_prop.handleTypes = self._exchanger.handle_type +- +- # Get multicast granularity and adjust allocation size +- self._mc_granularity = checkCudaErrors( +- cuda.cuMulticastGetGranularity( +- mc_prop, +- cuda.CUmulticastGranularity_flags.CU_MULTICAST_GRANULARITY_RECOMMENDED, +- ) ++ # Set up multicast properties ++ mc_prop = cuda.CUmulticastObjectProp() ++ mc_prop.numDevices = self.group_size ++ mc_prop.size = self.allocation_size ++ mc_prop.handleTypes = self._exchanger.handle_type ++ ++ # Get multicast granularity and adjust allocation size ++ self._mc_granularity = checkCudaErrors( ++ cuda.cuMulticastGetGranularity( ++ mc_prop, ++ cuda.CUmulticastGranularity_flags.CU_MULTICAST_GRANULARITY_RECOMMENDED, + ) +- self.allocation_size = round_up(self.allocation_size, self._mc_granularity) ++ ) ++ self.allocation_size = round_up(self.allocation_size, self._mc_granularity) + + return allocation_prop, mc_prop + +@@ -1320,24 +1275,21 @@ class SymmDeviceMemory: + self._exchanger.handle_type, + ) + ) +- self._exchanger.cleanup(all_shareable_uc_handles[p]) +- self._exchanger.cleanup(local_shareable_uc_handle) ++ self._exchanger.cleanup(all_shareable_uc_handles[p]) + + # Reserve address space for UC pointers +- if not self.uc_ptrs: +- self.uc_ptrs = [0] * self.group_size +- total_uc_size = self.allocation_size * self.group_size +- self.total_uc_size = total_uc_size +- uc_base_ptr = checkCudaErrors( +- cuda.cuMemAddressReserve(total_uc_size, self._mc_granularity, 0, 0) +- ) +- self.uc_base_ptr = uc_base_ptr +- for i in range(self.group_size): +- offset = self.allocation_size * i +- self.uc_ptrs[i] = int(uc_base_ptr) + offset ++ self.uc_ptrs = [0] * self.group_size ++ total_uc_size = self.allocation_size * self.group_size ++ self.total_uc_size = total_uc_size ++ uc_base_ptr = checkCudaErrors( ++ cuda.cuMemAddressReserve(total_uc_size, self._mc_granularity, 0, 0) ++ ) ++ self.uc_base_ptr = uc_base_ptr + + # Map UC memory + for i in range(self.group_size): ++ offset = self.allocation_size * i ++ self.uc_ptrs[i] = int(uc_base_ptr) + offset + checkCudaErrors( + cuda.cuMemMap( + self.uc_ptrs[i], self.allocation_size, 0, self.uc_handles[i], 0 +@@ -1347,7 +1299,7 @@ class SymmDeviceMemory: + # Set memory access permissions for UC + access_desc = self._get_mem_access_desc() + checkCudaErrors( +- cuda.cuMemSetAccess(self.uc_base_ptr, self.total_uc_size, [access_desc], 1) ++ cuda.cuMemSetAccess(uc_base_ptr, total_uc_size, [access_desc], 1) + ) + + def _setup_multicast(self, mc_prop): +@@ -1377,18 +1329,15 @@ class SymmDeviceMemory: + self._exchanger.handle_type, + ) + ) +- self._exchanger.cleanup(shareable_mc_handle) ++ self._exchanger.cleanup(shareable_mc_handle) + + # Add device to multicast + checkCudaErrors(cuda.cuMulticastAddDevice(self.mc_handle, self.device_idx)) + + # Reserve and map MC pointer +- if not self.mc_ptr: +- self.mc_ptr = checkCudaErrors( +- cuda.cuMemAddressReserve( +- self.allocation_size, self._mc_granularity, 0, 0 +- ) +- ) ++ self.mc_ptr = checkCudaErrors( ++ cuda.cuMemAddressReserve(self.allocation_size, self._mc_granularity, 0, 0) ++ ) + checkCudaErrors( + cuda.cuMemMap(self.mc_ptr, self.allocation_size, 0, self.mc_handle, 0) + ) +diff --git a/flashinfer/comm/trtllm_ar.py b/flashinfer/comm/trtllm_ar.py +index d8a3666a6c0d68a4f8c8531570755a6497f0b02b..77fba4049ab56a5445650cc25b9ed1c127cd804e 100644 +--- a/flashinfer/comm/trtllm_ar.py ++++ b/flashinfer/comm/trtllm_ar.py +@@ -16,7 +16,7 @@ limitations under the License. + + import functools + import logging +-from ctypes import c_void_p, cast, create_string_buffer ++from ctypes import c_void_p, cast + from types import SimpleNamespace + from typing import List, Optional, Tuple, Union + from typing_extensions import deprecated +@@ -431,7 +431,7 @@ OneShotMaxToken = 128 + MAX_ALL_REDUCE_BLOCKS = 24 + LamportTokenNumThreshold = 16 + +-_symm_workspace_refs: dict[int, list[object]] = {} ++_symm_workspace_refs: dict[int, list[torch.Tensor]] = {} + + + @deprecated( +@@ -555,34 +555,6 @@ BarrierFlagCount = 256 + MAX_COMM_SIZE = 2147483647 & ~((1 << 21) - 1) # MAX_INT32 rounded down to 2MB + + +-def _initialize_allreduce_fusion_protocol( +- ipc_handles: List[List[int]], +- tp_rank: int, +- flag_size: int, +- lamport_buffer_size: int, +- lamport_comm_size: int, +- use_fp32_lamport: bool, +- control_flag_ptr: int, +-) -> None: +- cudart.cudaMemset(c_void_p(ipc_handles[1][tp_rank]), 0, flag_size) +- +- lamport_dtype = torch.float32 if use_fp32_lamport else torch.float16 +- aligned_size = round_up(lamport_buffer_size, 16) +- trtllm_lamport_initialize( +- ipc_handles[2][tp_rank], +- aligned_size // (4 if use_fp32_lamport else 2), +- lamport_dtype, +- ) +- +- cudart.cudaMemset(c_void_p(control_flag_ptr), 0, 5 * 4) +- lamport_comm_size_bytes = create_string_buffer( +- lamport_comm_size.to_bytes(4, byteorder="little"), 4 +- ) +- cudart.cudaMemcpy( +- c_void_p(control_flag_ptr + 3 * 4), cast(lamport_comm_size_bytes, c_void_p), 4 +- ) +- +- + @deprecated( + "use the unified API allreduce.py instead. It will internally call trtllm_create_ipc_workspace_for_all_reduce_fusion." + ) +@@ -620,8 +592,7 @@ def trtllm_create_ipc_workspace_for_all_reduce_fusion( + use_fp32_lamport, buffer_size, flag_size, lamport_comm_size, lamport_buffer_size + - If create_metadata=True: and use_symm_dev_mem=True: (ipc_handles, workspace_tensor, mem_handles,metadata) + where metadata contains: tp_rank, tp_size, max_token_num, hidden_dim, +- use_fp32_lamport, buffer_size, flag_size, lamport_comm_size, +- lamport_buffer_size, control_flag_ptr ++ use_fp32_lamport, buffer_size, flag_size, lamport_comm_size, lamport_buffer_size + and mem_handles is a list of SymmDeviceMemory objects. + + Note: The optional parameters make the API clunky at this time. This will be refactored in the future, at the cost of backward compatibility, where the default behavior will be +@@ -668,7 +639,12 @@ def trtllm_create_ipc_workspace_for_all_reduce_fusion( + lamport_buffer_size = lamport_comm_size * 3 + + device = torch.device(f"cuda:{torch.cuda.current_device()}") +- symm_refs: list[object] = [] ++ group_name = ( ++ group.group_name ++ if group is not None ++ else torch.distributed.group.WORLD.group_name ++ ) ++ symm_refs: list[torch.Tensor] = [] + + # we should init 3 buffers for all reduce fusion: + # [buffer_size, flag_size, lamport_buffer_size] +@@ -683,35 +659,16 @@ def trtllm_create_ipc_workspace_for_all_reduce_fusion( + ]: + aligned_size = round_up(size, 16) + +- if use_symm_dev_mem: +- assert comm_backend is not None +- handle = SymmDeviceMemory( +- buf_size=aligned_size, +- group_size=tp_size, +- group_rank=tp_rank, +- device_idx=device.index, +- comm_backend_for_handle_transfer=comm_backend, +- enable_multicast=False, +- allocate_signal_pads=False, +- ) +- ptrs = handle.get_buffer_ptrs_host() +- symm_refs.append(handle) +- mem_handles.append(handle) +- else: +- group_name = ( +- group.group_name +- if group is not None +- else torch.distributed.group.WORLD.group_name +- ) +- ptrs, tensor, handle = _alloc_symm_buffer_bytes( +- aligned_size, +- tp_size, +- dtype, +- device, +- group_name, +- ) +- symm_refs.append((tensor, handle)) ++ ptrs, tensor, handle = _alloc_symm_buffer_bytes( ++ aligned_size, ++ tp_size, ++ dtype, ++ device, ++ group_name, ++ ) ++ symm_refs.append((tensor, handle)) + ipc_handles.append(ptrs) ++ mem_handles.append(handle) + + logger.debug( + "rank %s allocated ipc_handles: %s", +@@ -721,9 +678,6 @@ def trtllm_create_ipc_workspace_for_all_reduce_fusion( + + _symm_workspace_refs[id(ipc_handles)] = symm_refs + +- if use_symm_dev_mem: +- cudart.cudaMemset(c_void_p(ipc_handles[1][tp_rank]), 0, flag_size) +- + # Initialize lamport buffer + aligned_lamport_buffer_size = round_up(lamport_buffer_size, 16) + if use_fp32_lamport: +@@ -774,7 +728,6 @@ def trtllm_create_ipc_workspace_for_all_reduce_fusion( + ) + + if use_symm_dev_mem: +- torch.cuda.synchronize() + comm_backend.barrier() # must sync after create_workspace + else: + dist.barrier(group=group) +@@ -792,7 +745,6 @@ def trtllm_create_ipc_workspace_for_all_reduce_fusion( + "lamport_buffer_size": lamport_buffer_size, + } + if use_symm_dev_mem: +- metadata["control_flag_ptr"] = flag_ptr.value + return ipc_handles, workspace_tensor, mem_handles, metadata + else: + return ipc_handles, workspace_tensor, metadata +diff --git a/flashinfer/comm/trtllm_mnnvl_ar.py b/flashinfer/comm/trtllm_mnnvl_ar.py +index 0a76914d20a33a81227adf3c8c30eaf6280a98fd..bc48915c7b2419d004848143ff868a5a43e98e87 100644 +--- a/flashinfer/comm/trtllm_mnnvl_ar.py ++++ b/flashinfer/comm/trtllm_mnnvl_ar.py +@@ -14,14 +14,15 @@ import torch + from typing_extensions import deprecated + + from flashinfer.comm.mapping import Mapping +-from flashinfer.api_logging import flashinfer_api ++from flashinfer.comm.mnnvl import TorchDistBackend + + from ..jit import gen_trtllm_mnnvl_comm_module + from ..utils import register_custom_op + from ..fp4_quantization import _compute_swizzled_layout_sf_size +-from .mnnvl import CommBackend, McastGPUBuffer, MPIBackend, SymmDeviceMemory ++from .mnnvl import CommBackend, MPIBackend + from .trtllm_ar import QuantizationSFLayout + from .workspace_base import AllReduceFusionWorkspace ++from .torch_symmetric_memory import _alloc_symm_buffer_bytes + + + def mpi_barrier(): +@@ -146,16 +147,26 @@ class MNNVLAllReduceFusionWorkspace(AllReduceFusionWorkspace): + # support base_gpu_id != 0 scenarios where the actual CUDA device + # index differs from the TP rank / local_rank. + device = torch.device("cuda", torch.cuda.current_device()) +- self.handle = McastGPUBuffer( +- buf_size=requested_workspace_size, +- group_size=mapping.tp_size, +- group_rank=mapping.tp_rank, +- device=device, +- comm_backend_for_handle_transfer=comm_backend, ++ if isinstance(comm_backend, TorchDistBackend): ++ group = ( ++ comm_backend._group ++ if comm_backend._group is not None ++ else torch.distributed.group.WORLD ++ ) ++ group_name = group.group_name ++ else: ++ group_name = torch.distributed.group.WORLD.group_name ++ self.ptrs, self.tensor, self.handle = _alloc_symm_buffer_bytes( ++ requested_workspace_size, ++ mapping.tp_size, ++ torch.float32, ++ device, ++ group_name, + ) +- self.ptrs = self.handle.mcast_device_memory.get_buffer_ptrs_host() + +- allocated_size = self.handle.buf_size ++ # handle.buffer_size is the usable data size. torch symmetric memory ++ # allocator places signal_pad on top of it, not carved from within. ++ allocated_size = self.handle.buffer_size + # We want the buffer size to be aligned to 16B which is the granularity for buffer management. + self.buffer_size_bytes = ( + math.floor(allocated_size / self.NUM_LAMPORT_BUFFERS) // 16 * 16 +@@ -167,7 +178,9 @@ class MNNVLAllReduceFusionWorkspace(AllReduceFusionWorkspace): + f"[MNNVL Allreduce] Actual allocated size: {allocated_size} bytes, Actual buffer size per lamport buffer: {self.buffer_size_bytes} bytes, total workspace: {self.workspace_size_bytes} bytes." + ) + +- self.handle.lamport_initialize(self.rank, torch.float32) ++ # lamport initialize tensor to negative zero. ++ self.tensor.fill_(-0.0) ++ # Wait until the initialization is done + torch.cuda.synchronize() + comm_backend.barrier() + +@@ -182,9 +195,9 @@ class MNNVLAllReduceFusionWorkspace(AllReduceFusionWorkspace): + device=torch.device("cuda", torch.cuda.current_device()), + ) + +- self.uc_ptrs_dev = self.handle.get_buffer_ptrs_dev() +- self.uc_ptr_local = self.handle.get_unicast_ptr(self.rank) +- self.mc_ptr = self.handle.get_multicast_ptr() ++ self.uc_ptrs_dev = self.handle.buffer_ptrs_dev ++ self.uc_ptr_local = self.handle.buffer_ptrs[self.rank] ++ self.mc_ptr = self.handle.multicast_ptr + + @functools.cache + def is_buffer_size_sufficient( +@@ -239,53 +252,6 @@ class MNNVLAllReduceFusionWorkspace(AllReduceFusionWorkspace): + def backend(self) -> str: + return "mnnvl" + +- def _require_handles_attached(self) -> None: +- memory = getattr(self.handle, "mcast_device_memory", None) +- if isinstance(memory, SymmDeviceMemory) and not memory.mapped: +- raise RuntimeError("MNNVL handles are not attached") +- +- def _initialize_protocol(self) -> None: +- self.handle.lamport_initialize(self.rank, torch.float32) +- num_bytes_to_clear = [0] * 4 +- self.buffer_flags.copy_( +- torch.tensor( +- [0, 2, self.buffer_size_bytes, 0, *num_bytes_to_clear, 0], +- dtype=torch.uint32, +- device=self.buffer_flags.device, +- ) +- ) +- torch.cuda.synchronize() +- +- @flashinfer_api +- def checkpoint_prepare(self) -> None: +- """Detach physical backing; repeated successful calls are no-ops.""" +- memory = getattr(self.handle, "mcast_device_memory", None) +- if not isinstance(memory, SymmDeviceMemory): +- raise NotImplementedError( +- "Stable-VA checkpointing is unavailable for workspaces backed " +- "by torch symmetric memory" +- ) +- if not memory.mapped: +- return +- memory._unmap_and_release_handles() +- # Do not return until every rank has released its workspace handles. +- memory.comm_backend.barrier() +- +- @flashinfer_api +- def checkpoint_restore(self, comm_backend: CommBackend) -> None: +- """Restore physical backing; repeated successful calls are no-ops.""" +- memory = getattr(self.handle, "mcast_device_memory", None) +- if not isinstance(memory, SymmDeviceMemory): +- raise NotImplementedError( +- "Stable-VA checkpointing is unavailable for workspaces backed " +- "by torch symmetric memory" +- ) +- if memory.mapped: +- return +- memory._create_and_map_handles(comm_backend) +- self._initialize_protocol() +- comm_backend.barrier() +- + def destroy(self) -> None: + """Destroy workspace and free resources.""" + if getattr(self, "_destroyed", False): +@@ -295,6 +261,7 @@ class MNNVLAllReduceFusionWorkspace(AllReduceFusionWorkspace): + del self.uc_ptrs_dev + del self.uc_ptr_local + del self.mc_ptr ++ del self.tensor + del self.handle + del self.ptrs + self._destroyed = True +@@ -466,8 +433,6 @@ def trtllm_mnnvl_allreduce( + f"The output tensor must be 2D, got {len(output.shape)}D. The shape is {output.shape}." + ) + +- workspace._require_handles_attached() +- + module = get_trtllm_mnnvl_comm_module() + + if strategy == MNNVLAllreduceFusionStrategy.AUTO: +@@ -568,8 +533,6 @@ def trtllm_mnnvl_fused_allreduce_add_rmsnorm( + f"The residual output tensor must be 2D, got {len(residual_out.shape)}D. The shape is {residual_out.shape}." + ) + +- workspace._require_handles_attached() +- + module = get_trtllm_mnnvl_comm_module() + + if strategy == MNNVLAllreduceFusionStrategy.AUTO: +@@ -809,8 +772,6 @@ def trtllm_mnnvl_fused_allreduce_add_rmsnorm_quant( + else: + raise ValueError(f"Unsupported MNNVL quant_type: {quant_type}") + +- workspace._require_handles_attached() +- + if strategy == MNNVLAllreduceFusionStrategy.AUTO: + strategy = MNNVLAllreduceFusionStrategy.select_strategy( + workspace.tp_size, token_num, hidden_dim, input.dtype +@@ -884,7 +845,7 @@ def get_allreduce_mnnvl_workspace( + + Returns: + Tuple containing: +- - MNNVLAllReduceFusionWorkspace: The CUDA VMM-backed workspace object ++ - MNNVLAllReduceFusionWorkspace: The workspace object backed by torch symmetric memory + - torch.Tensor: Buffer flags tensor tracking state + - int: Maximum number of elements that can fit in buffer + """ diff --git a/configs/nvidia-master.yaml b/configs/nvidia-master.yaml index 6be8e006f..3f3fc6c4c 100644 --- a/configs/nvidia-master.yaml +++ b/configs/nvidia-master.yaml @@ -13559,7 +13559,7 @@ minimaxm3-fp8-b300-vllm: # weights are pre-staged read-only at /scratch/models/MiniMax-M3-NVFP4 (added to # the STAGED_MODELS allow-list in launch_b300-nv.sh). minimaxm3-fp4-b300-vllm: - image: vllm/vllm-openai:nightly-93d8f834dd8acf33eb0e2a75b2711b628cb6e226 + image: vllm/vllm-openai:nightly-2afa3f7e950264bb179d030c23a1ed1f46558fd9 model: nvidia/MiniMax-M3-NVFP4 model-prefix: minimaxm3 runner: b300 diff --git a/perf-changelog.yaml b/perf-changelog.yaml index be28208ab..926f01e31 100644 --- a/perf-changelog.yaml +++ b/perf-changelog.yaml @@ -4716,3 +4716,14 @@ - "Clean the export envs" - "Enable two batch overlap" pr-link: https://github.com/SemiAnalysisAI/InferenceX/pull/2093 + +- config-keys: + - minimaxm3-fp4-b300-vllm + description: + - "Update the vLLM image to nightly-2afa3f7e950264bb179d030c23a1ed1f46558fd9" + - "Install FlashInfer 0.6.15.dev20260710" + - "Use the upstream SM100 low-M MXFP8 split-K implementation included in FlashInfer 0.6.15" + - "Apply the FlashInfer AutoTuner non-Tensor guard fix from PR #3918" + - "Revert the FlashInfer PR #3745 runtime communication workspace changes via a local patch" + - "Use the explicit FlashInfer TRT-LLM all-reduce backend override" + pr-link: https://github.com/SemiAnalysisAI/InferenceX/pull/2124