[FlashInfer] Switch FlashInfer allreduce fusion to unified API (#18341)
This commit is contained in:
Mohammad Miadh Angkad
committed by
GitHub
GitHub
parent
bf36aa4c31
commit
7b0fb43c7a
@@ -2,9 +2,11 @@ import logging
|
||||
from typing import Optional, Tuple
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
|
||||
from sglang.srt.distributed import get_tensor_model_parallel_world_size
|
||||
from sglang.srt.distributed import (
|
||||
get_tensor_model_parallel_rank,
|
||||
get_tensor_model_parallel_world_size,
|
||||
)
|
||||
from sglang.srt.utils import is_flashinfer_available
|
||||
from sglang.srt.utils.custom_op import register_custom_op
|
||||
|
||||
@@ -17,7 +19,15 @@ if is_flashinfer_available():
|
||||
try:
|
||||
import flashinfer.comm as comm
|
||||
|
||||
_flashinfer_comm = comm
|
||||
if hasattr(comm, "allreduce_fusion") and hasattr(
|
||||
comm, "create_allreduce_fusion_workspace"
|
||||
):
|
||||
_flashinfer_comm = comm
|
||||
else:
|
||||
logger.warning(
|
||||
"flashinfer.comm unified allreduce_fusion API is not available, "
|
||||
"falling back to standard implementation"
|
||||
)
|
||||
except ImportError:
|
||||
logger.warning(
|
||||
"flashinfer.comm is not available, falling back to standard "
|
||||
@@ -27,10 +37,12 @@ if is_flashinfer_available():
|
||||
|
||||
class FlashInferWorkspaceManager:
|
||||
def __init__(self):
|
||||
self.workspace_tensor = None
|
||||
self.ipc_handles = None
|
||||
self.workspace = None
|
||||
self.world_size = None
|
||||
self.rank = None
|
||||
self.max_token_num = None
|
||||
self.hidden_dim = None
|
||||
self.dtype = None
|
||||
self.initialized = False
|
||||
|
||||
def initialize(
|
||||
@@ -39,13 +51,10 @@ class FlashInferWorkspaceManager:
|
||||
rank: int,
|
||||
max_token_num: int,
|
||||
hidden_dim: int,
|
||||
group=None,
|
||||
use_fp32_lamport: bool = False,
|
||||
dtype: torch.dtype,
|
||||
use_oneshot: Optional[bool] = None,
|
||||
):
|
||||
"""Initialize workspace"""
|
||||
if self.initialized and self.world_size == world_size:
|
||||
return
|
||||
|
||||
if _flashinfer_comm is None:
|
||||
logger.warning(
|
||||
"FlashInfer comm not available, skipping workspace " "initialization"
|
||||
@@ -53,47 +62,82 @@ class FlashInferWorkspaceManager:
|
||||
return
|
||||
|
||||
self.cleanup()
|
||||
|
||||
self.ipc_handles, self.workspace_tensor = (
|
||||
comm.trtllm_create_ipc_workspace_for_all_reduce_fusion(
|
||||
rank,
|
||||
world_size,
|
||||
max_token_num,
|
||||
hidden_dim,
|
||||
group=group,
|
||||
use_fp32_lamport=use_fp32_lamport,
|
||||
try:
|
||||
self.workspace = _flashinfer_comm.create_allreduce_fusion_workspace(
|
||||
backend="trtllm",
|
||||
world_size=world_size,
|
||||
rank=rank,
|
||||
max_token_num=max_token_num,
|
||||
hidden_dim=hidden_dim,
|
||||
dtype=dtype,
|
||||
force_oneshot_support=bool(use_oneshot),
|
||||
)
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to initialize FlashInfer workspace: {e}")
|
||||
self.workspace = None
|
||||
self.initialized = False
|
||||
return
|
||||
|
||||
self.world_size = world_size
|
||||
self.rank = rank
|
||||
self.max_token_num = max_token_num
|
||||
self.hidden_dim = hidden_dim
|
||||
self.dtype = dtype
|
||||
self.initialized = True
|
||||
|
||||
backend = getattr(self.workspace, "backend", "unknown")
|
||||
logger.info(
|
||||
f"FlashInfer workspace initialized for rank {rank}, "
|
||||
f"world_size {world_size}"
|
||||
f"world_size {world_size}, backend {backend}"
|
||||
)
|
||||
|
||||
def is_buffer_size_sufficient(
|
||||
self,
|
||||
token_num: int,
|
||||
hidden_dim: int,
|
||||
dtype: torch.dtype,
|
||||
use_oneshot: Optional[bool] = None,
|
||||
) -> bool:
|
||||
if not self.initialized or self.workspace is None:
|
||||
return False
|
||||
try:
|
||||
return self.workspace.is_buffer_size_sufficient(
|
||||
tp_size=self.world_size,
|
||||
num_tokens=token_num,
|
||||
hidden_dim=hidden_dim,
|
||||
dtype=dtype,
|
||||
use_oneshot=use_oneshot,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.debug(f"FlashInfer workspace size check failed: {e}")
|
||||
return False
|
||||
|
||||
def cleanup(self):
|
||||
"""Clean up workspace"""
|
||||
if self.initialized and self.ipc_handles is not None:
|
||||
if self.workspace is not None:
|
||||
try:
|
||||
_flashinfer_comm.trtllm_destroy_ipc_workspace_for_all_reduce(
|
||||
self.ipc_handles, group=dist.group.WORLD
|
||||
)
|
||||
self.workspace.destroy()
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to cleanup FlashInfer workspace: {e}")
|
||||
finally:
|
||||
self.workspace_tensor = None
|
||||
self.ipc_handles = None
|
||||
self.workspace = None
|
||||
self.initialized = False
|
||||
self.world_size = None
|
||||
self.rank = None
|
||||
self.max_token_num = None
|
||||
self.hidden_dim = None
|
||||
self.dtype = None
|
||||
|
||||
|
||||
_workspace_manager = FlashInferWorkspaceManager()
|
||||
|
||||
|
||||
def ensure_workspace_initialized(
|
||||
max_token_num: int = 2048, hidden_dim: int = 4096, use_fp32_lamport: bool = False
|
||||
max_token_num: int = 2048,
|
||||
hidden_dim: int = 4096,
|
||||
dtype: torch.dtype = torch.float16,
|
||||
token_num: Optional[int] = None,
|
||||
use_oneshot: Optional[bool] = None,
|
||||
):
|
||||
"""Ensure workspace is initialized"""
|
||||
if not is_flashinfer_available() or _flashinfer_comm is None:
|
||||
@@ -103,18 +147,27 @@ def ensure_workspace_initialized(
|
||||
if world_size <= 1:
|
||||
return False
|
||||
|
||||
rank = dist.get_rank()
|
||||
rank = get_tensor_model_parallel_rank()
|
||||
token_num = token_num or max_token_num
|
||||
|
||||
if (
|
||||
not _workspace_manager.initialized
|
||||
or _workspace_manager.world_size != world_size
|
||||
or _workspace_manager.rank != rank
|
||||
or not _workspace_manager.is_buffer_size_sufficient(
|
||||
token_num=token_num,
|
||||
hidden_dim=hidden_dim,
|
||||
dtype=dtype,
|
||||
use_oneshot=use_oneshot,
|
||||
)
|
||||
):
|
||||
_workspace_manager.initialize(
|
||||
world_size=world_size,
|
||||
rank=rank,
|
||||
max_token_num=max_token_num,
|
||||
hidden_dim=hidden_dim,
|
||||
use_fp32_lamport=use_fp32_lamport,
|
||||
dtype=dtype,
|
||||
use_oneshot=use_oneshot,
|
||||
)
|
||||
|
||||
return _workspace_manager.initialized
|
||||
@@ -177,42 +230,39 @@ def flashinfer_allreduce_residual_rmsnorm(
|
||||
return None, None
|
||||
|
||||
assert input_tensor.shape[0] <= max_token_num
|
||||
if (
|
||||
not input_tensor.is_contiguous()
|
||||
or not residual.is_contiguous()
|
||||
or not weight.is_contiguous()
|
||||
):
|
||||
logger.debug("Non-contiguous tensors, skipping FlashInfer allreduce fusion")
|
||||
return None, None
|
||||
|
||||
if not ensure_workspace_initialized(
|
||||
max_token_num=max_token_num,
|
||||
hidden_dim=input_tensor.shape[-1],
|
||||
use_fp32_lamport=(input_tensor.dtype == torch.float32),
|
||||
dtype=input_tensor.dtype,
|
||||
token_num=input_tensor.shape[0],
|
||||
use_oneshot=use_oneshot,
|
||||
):
|
||||
logger.debug("FlashInfer workspace not available")
|
||||
return None, None
|
||||
|
||||
token_num, hidden_dim = input_tensor.shape
|
||||
|
||||
residual_out = torch.empty_like(residual)
|
||||
norm_out = torch.empty_like(input_tensor)
|
||||
|
||||
_flashinfer_comm.trtllm_allreduce_fusion(
|
||||
allreduce_in=input_tensor,
|
||||
world_size=world_size,
|
||||
world_rank=dist.get_rank(),
|
||||
token_num=token_num,
|
||||
hidden_dim=hidden_dim,
|
||||
workspace_ptrs=_workspace_manager.workspace_tensor,
|
||||
_flashinfer_comm.allreduce_fusion(
|
||||
input=input_tensor,
|
||||
workspace=_workspace_manager.workspace,
|
||||
pattern=_flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNorm,
|
||||
launch_with_pdl=True,
|
||||
use_oneshot=use_oneshot,
|
||||
trigger_completion_at_end=trigger_completion_at_end,
|
||||
fp32_acc=fp32_acc,
|
||||
pattern_code=(_flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNorm),
|
||||
allreduce_out=None,
|
||||
residual_in=residual,
|
||||
residual_out=residual_out,
|
||||
norm_out=norm_out,
|
||||
quant_out=None,
|
||||
scale_out=None,
|
||||
residual_in=residual,
|
||||
rms_gamma=weight,
|
||||
rms_eps=eps,
|
||||
scale_factor=None,
|
||||
layout_code=None,
|
||||
use_oneshot=use_oneshot,
|
||||
fp32_acc=fp32_acc,
|
||||
)
|
||||
|
||||
return norm_out, residual_out
|
||||
|
||||
Reference in New Issue
Block a user