275 lines
8.2 KiB
Python
275 lines
8.2 KiB
Python
import logging
|
|
from typing import Optional, Tuple
|
|
|
|
import torch
|
|
|
|
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
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_flashinfer_comm = None
|
|
_workspace_manager = None
|
|
|
|
if is_flashinfer_available():
|
|
try:
|
|
import flashinfer.comm as 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 "
|
|
"implementation"
|
|
)
|
|
|
|
|
|
class FlashInferWorkspaceManager:
|
|
def __init__(self):
|
|
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(
|
|
self,
|
|
world_size: int,
|
|
rank: int,
|
|
max_token_num: int,
|
|
hidden_dim: int,
|
|
dtype: torch.dtype,
|
|
use_oneshot: Optional[bool] = None,
|
|
):
|
|
"""Initialize workspace"""
|
|
if _flashinfer_comm is None:
|
|
logger.warning(
|
|
"FlashInfer comm not available, skipping workspace " "initialization"
|
|
)
|
|
return
|
|
|
|
self.cleanup()
|
|
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}, 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.workspace is not None:
|
|
try:
|
|
self.workspace.destroy()
|
|
except Exception as e:
|
|
logger.warning(f"Failed to cleanup FlashInfer workspace: {e}")
|
|
finally:
|
|
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,
|
|
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:
|
|
return False
|
|
|
|
world_size = get_tensor_model_parallel_world_size()
|
|
if world_size <= 1:
|
|
return False
|
|
|
|
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,
|
|
dtype=dtype,
|
|
use_oneshot=use_oneshot,
|
|
)
|
|
|
|
return _workspace_manager.initialized
|
|
|
|
|
|
def fake_flashinfer_allreduce_residual_rmsnorm(
|
|
input_tensor: torch.Tensor,
|
|
residual: torch.Tensor,
|
|
weight: torch.Tensor,
|
|
eps: float = 1e-6,
|
|
max_token_num: int = 16384,
|
|
use_oneshot: Optional[bool] = None,
|
|
trigger_completion_at_end: bool = False,
|
|
fp32_acc: bool = False,
|
|
) -> Tuple[torch.Tensor, torch.Tensor]:
|
|
residual_out = torch.empty_like(residual)
|
|
norm_out = torch.empty_like(input_tensor)
|
|
return norm_out, residual_out
|
|
|
|
|
|
@register_custom_op(
|
|
mutates_args=["input_tensor", "residual", "weight"],
|
|
fake_impl=fake_flashinfer_allreduce_residual_rmsnorm,
|
|
)
|
|
def flashinfer_allreduce_residual_rmsnorm(
|
|
input_tensor: torch.Tensor,
|
|
residual: torch.Tensor,
|
|
weight: torch.Tensor,
|
|
eps: float = 1e-6,
|
|
max_token_num: int = 2048,
|
|
use_oneshot: Optional[bool] = None,
|
|
trigger_completion_at_end: bool = False,
|
|
fp32_acc: bool = False,
|
|
) -> Tuple[torch.Tensor, torch.Tensor]:
|
|
"""
|
|
Use FlashInfer's fused allreduce + residual + RMS norm operation
|
|
|
|
Args:
|
|
input_tensor: Input tensor that needs allreduce
|
|
residual: Residual tensor
|
|
weight: RMS norm weight
|
|
eps: RMS norm epsilon
|
|
max_token_num: Maximum token number
|
|
use_oneshot: Whether to use oneshot mode
|
|
trigger_completion_at_end: Whether to trigger completion at end
|
|
fp32_acc: Whether to use fp32 precision
|
|
|
|
Returns:
|
|
Tuple[torch.Tensor, torch.Tensor]: (norm_output, residual_output)
|
|
"""
|
|
if not is_flashinfer_available() or _flashinfer_comm is None:
|
|
logger.debug(
|
|
"FlashInfer not available, falling back to standard " "implementation"
|
|
)
|
|
return None, None
|
|
|
|
world_size = get_tensor_model_parallel_world_size()
|
|
if world_size <= 1:
|
|
logger.debug("Single GPU, no need for allreduce fusion")
|
|
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],
|
|
dtype=input_tensor.dtype,
|
|
token_num=input_tensor.shape[0],
|
|
use_oneshot=use_oneshot,
|
|
):
|
|
logger.debug("FlashInfer workspace not available")
|
|
return None, None
|
|
|
|
residual_out = torch.empty_like(residual)
|
|
norm_out = torch.empty_like(input_tensor)
|
|
|
|
_flashinfer_comm.allreduce_fusion(
|
|
input=input_tensor,
|
|
workspace=_workspace_manager.workspace,
|
|
pattern=_flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNorm,
|
|
launch_with_pdl=True,
|
|
residual_out=residual_out,
|
|
norm_out=norm_out,
|
|
residual_in=residual,
|
|
rms_gamma=weight,
|
|
rms_eps=eps,
|
|
use_oneshot=use_oneshot,
|
|
fp32_acc=fp32_acc,
|
|
)
|
|
|
|
return norm_out, residual_out
|
|
|
|
|
|
def cleanup_flashinfer_workspace():
|
|
global _workspace_manager
|
|
if _workspace_manager is not None:
|
|
_workspace_manager.cleanup()
|