perf: avoid full-vocab gather for greedy sampling

lightllm/common/basemodel/triton_kernel/post_process/greedy_sample.py

@@ -0,0 +1,104 @@
+"""Local greedy statistics for distributed vocabulary shards."""
+
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _greedy_sample_stage1_kernel(
+    logits,
+    partial_max,
+    partial_sum,
+    partial_argmax,
+    stride_row,
+    stride_col,
+    vocab_size: tl.constexpr,
+    num_blocks: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    row = tl.program_id(0)
+    block = tl.program_id(1)
+    offsets = block * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    values = tl.load(
+        logits + row * stride_row + offsets * stride_col,
+        mask=offsets < vocab_size,
+        other=-float("inf"),
+    )
+    values = values.to(tl.float32)
+
+    block_max = tl.max(values, axis=0)
+    block_sum = tl.sum(tl.exp(values - block_max), axis=0)
+    block_argmax = tl.argmax(values, axis=0) + block * BLOCK_SIZE
+    output_offset = row * num_blocks + block
+    tl.store(partial_max + output_offset, block_max)
+    tl.store(partial_sum + output_offset, block_sum)
+    tl.store(partial_argmax + output_offset, block_argmax)
+
+
+@triton.jit
+def _greedy_sample_stage2_stats_kernel(
+    partial_max,
+    partial_sum,
+    partial_argmax,
+    output_stats,
+    num_blocks: tl.constexpr,
+    batch_size: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    row = tl.program_id(0)
+    offsets = tl.arange(0, BLOCK_SIZE)
+    mask = offsets < num_blocks
+    input_offset = row * num_blocks + offsets
+    block_max = tl.load(partial_max + input_offset, mask=mask, other=-float("inf"))
+    block_sum = tl.load(partial_sum + input_offset, mask=mask, other=0.0)
+    block_argmax = tl.load(partial_argmax + input_offset, mask=mask, other=0x7FFFFFFF)
+
+    global_max = tl.max(block_max, axis=0)
+    global_sum = tl.sum(block_sum * tl.exp(block_max - global_max), axis=0)
+    candidate_ids = tl.where(block_max == global_max, block_argmax, 0x7FFFFFFF)
+    global_argmax = tl.min(candidate_ids, axis=0)
+    tl.store(output_stats + row, global_max)
+    tl.store(output_stats + batch_size + row, global_max + tl.log(global_sum))
+    tl.store(output_stats + 2 * batch_size + row, global_argmax)
+
+
+def _launch_stage1(logits: torch.Tensor, scratch: torch.Tensor, block_size: int, num_blocks: int) -> None:
+    batch_size, vocab_size = logits.shape
+    _greedy_sample_stage1_kernel[(batch_size, num_blocks)](
+        logits,
+        scratch[0],
+        scratch[1],
+        scratch[2],
+        logits.stride(0),
+        logits.stride(1),
+        vocab_size=vocab_size,
+        num_blocks=num_blocks,
+        BLOCK_SIZE=block_size,
+        num_warps=8,
+    )
+
+
+@torch.no_grad()
+def greedy_sample_local_stats(logits: torch.Tensor, alloc_func=torch.empty) -> torch.Tensor:
+    """Return local max, logsumexp and argmax rows for distributed greedy sampling."""
+
+    assert logits.ndim == 2 and logits.is_cuda and logits.is_contiguous()
+    batch_size, vocab_size = logits.shape
+    block_size = 4096
+    num_blocks = triton.cdiv(vocab_size, block_size)
+    scratch = alloc_func((3, batch_size, num_blocks), dtype=torch.float32, device=logits.device)
+    output_stats = alloc_func((3, batch_size), dtype=torch.float32, device=logits.device)
+
+    _launch_stage1(logits, scratch, block_size, num_blocks)
+    _greedy_sample_stage2_stats_kernel[(batch_size,)](
+        scratch[0],
+        scratch[1],
+        scratch[2],
+        output_stats,
+        num_blocks=num_blocks,
+        batch_size=batch_size,
+        BLOCK_SIZE=triton.next_power_of_2(num_blocks),
+        num_warps=4,
+    )
+    return output_stats

lightllm/common/basemodel/triton_kernel/post_process/vocab_parallel_greedy.py

@@ -0,0 +1,119 @@
+"""Greedy sampling directly from tensor-parallel vocabulary shards."""
+
+import torch
+import triton
+import triton.language as tl
+
+from lightllm.common.basemodel.triton_kernel.post_process.greedy_sample import (
+    greedy_sample_local_stats,
+)
+from lightllm.common.basemodel.triton_kernel.transpose_convert import (
+    transpose_convert_2d,
+)
+from lightllm.distributed.communication_op import all_gather
+from lightllm.utils.envs_utils import enable_env_vars
+
+
+VOCAB_PARALLEL_GREEDY_ENV = "LIGHTLLM_VOCAB_PARALLEL_GREEDY"
+
+
+def is_vocab_parallel_greedy_enabled() -> bool:
+    return enable_env_vars(VOCAB_PARALLEL_GREEDY_ENV)
+
+
+@triton.jit
+def _combine_vocab_parallel_stats_kernel(
+    gathered_stats,
+    output,
+    token_num,
+    vocab_size: tl.constexpr,
+    tp_world_size: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    token_offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    token_mask = token_offsets < token_num
+    rank_stride = 3 * token_num
+
+    global_max = tl.full((BLOCK_SIZE,), -float("inf"), tl.float32)
+    global_id = tl.full((BLOCK_SIZE,), 0x7FFFFFFF, tl.int32)
+    for rank in tl.static_range(tp_world_size):
+        rank_base = rank * rank_stride
+        local_max = tl.load(
+            gathered_stats + rank_base + token_offsets,
+            mask=token_mask,
+            other=-float("inf"),
+        )
+        local_id = tl.load(
+            gathered_stats + rank_base + 2 * token_num + token_offsets,
+            mask=token_mask,
+            other=0x7FFFFFFF,
+        ).to(tl.int32)
+        local_id += (rank * vocab_size) // tp_world_size
+        wins = (local_max > global_max) | ((local_max == global_max) & (local_id < global_id))
+        global_max = tl.where(wins, local_max, global_max)
+        global_id = tl.where(wins, local_id, global_id)
+
+    global_sum = tl.zeros((BLOCK_SIZE,), tl.float32)
+    for rank in tl.static_range(tp_world_size):
+        rank_base = rank * rank_stride
+        local_lse = tl.load(
+            gathered_stats + rank_base + token_num + token_offsets,
+            mask=token_mask,
+            other=-float("inf"),
+        )
+        global_sum += tl.exp(local_lse - global_max)
+
+    tl.store(output + token_offsets * 2, global_id, mask=token_mask)
+    tl.store(output + token_offsets * 2 + 1, -tl.log(global_sum), mask=token_mask)
+
+
+@torch.no_grad()
+def vocab_parallel_greedy(
+    local_logits: torch.Tensor,
+    *,
+    vocab_size: int,
+    tp_world_size: int,
+    group,
+    alloc_func,
+) -> torch.Tensor:
+    """Return packed ``[token_id, logprob]`` rows without gathering full logits."""
+
+    assert local_logits.ndim == 2 and local_logits.is_cuda and local_logits.is_contiguous()
+    local_vocab_size, token_num = local_logits.shape
+    assert local_vocab_size in {
+        vocab_size // tp_world_size,
+        (vocab_size + tp_world_size - 1) // tp_world_size,
+    }
+
+    transposed_logits = alloc_func(
+        (token_num, local_vocab_size),
+        dtype=local_logits.dtype,
+        device=local_logits.device,
+    )
+    transpose_convert_2d(local_logits, transposed_logits)
+    local_stats = greedy_sample_local_stats(transposed_logits, alloc_func=alloc_func)
+
+    gathered_stats = alloc_func((tp_world_size, 3, token_num), dtype=torch.float32, device=local_logits.device)
+    all_gather(
+        [gathered_stats[rank] for rank in range(tp_world_size)],
+        local_stats,
+        group=group,
+        async_op=False,
+    )
+
+    output = alloc_func((token_num, 2), dtype=torch.float32, device=local_logits.device)
+    _combine_vocab_parallel_stats_kernel[(triton.cdiv(token_num, 256),)](
+        gathered_stats,
+        output,
+        token_num,
+        vocab_size=vocab_size,
+        tp_world_size=tp_world_size,
+        BLOCK_SIZE=256,
+        num_warps=4,
+    )
+    return output
+
+
+def unpack_vocab_parallel_greedy(output: torch.Tensor):
+    assert output.ndim == 2 and output.shape[1] == 2 and output.dtype == torch.float32
+    return output[:, 0].to(torch.int64), output[:, 1]

lightllm/common/basemodel/triton_kernel/transpose_convert.py

@@ -0,0 +1,65 @@
+"""Tiled transpose kernels used by the post-layer logits path."""
+
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _transpose_convert_2d_kernel(
+    input_ptr,
+    output_ptr,
+    rows,
+    cols,
+    input_stride_0,
+    input_stride_1,
+    output_stride_0,
+    output_stride_1,
+    BLOCK_ROWS: tl.constexpr,
+    BLOCK_COLS: tl.constexpr,
+):
+    row_offsets = tl.program_id(0) * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS)
+    col_offsets = tl.program_id(1) * BLOCK_COLS + tl.arange(0, BLOCK_COLS)
+    input_offsets = row_offsets[:, None] * input_stride_0 + col_offsets[None, :] * input_stride_1
+    mask = (row_offsets[:, None] < rows) & (col_offsets[None, :] < cols)
+    values = tl.load(input_ptr + input_offsets, mask=mask)
+
+    output_offsets = col_offsets[:, None] * output_stride_0 + row_offsets[None, :] * output_stride_1
+    tl.store(output_ptr + output_offsets, tl.trans(values), mask=tl.trans(mask))
+
+
+@torch.no_grad()
+def transpose_convert_2d(
+    input_tensor: torch.Tensor,
+    output_tensor: torch.Tensor,
+    *,
+    block_rows: int = 64,
+    block_cols: int = 64,
+    num_warps: int = 8,
+    num_stages: int = 1,
+) -> torch.Tensor:
+    """Transpose a contiguous 2-D CUDA tensor while converting its dtype."""
+
+    assert input_tensor.is_cuda and output_tensor.is_cuda
+    assert input_tensor.device == output_tensor.device
+    assert input_tensor.ndim == 2 and output_tensor.ndim == 2
+    assert output_tensor.shape == (input_tensor.shape[1], input_tensor.shape[0])
+    assert input_tensor.is_contiguous() and output_tensor.is_contiguous()
+
+    rows, cols = input_tensor.shape
+    grid = (triton.cdiv(rows, block_rows), triton.cdiv(cols, block_cols))
+    _transpose_convert_2d_kernel[grid](
+        input_tensor,
+        output_tensor,
+        rows,
+        cols,
+        input_tensor.stride(0),
+        input_tensor.stride(1),
+        output_tensor.stride(0),
+        output_tensor.stride(1),
+        BLOCK_ROWS=block_rows,
+        BLOCK_COLS=block_cols,
+        num_warps=num_warps,
+        num_stages=num_stages,
+    )
+    return output_tensor

lightllm/models/llama/layer_infer/post_layer_infer.py

@@ -7,6 +7,10 @@
 from lightllm.models.llama.layer_weights.pre_and_post_layer_weight import LlamaPreAndPostLayerWeight
 from lightllm.models.llama.infer_struct import LlamaInferStateInfo
 from lightllm.common.basemodel import PostLayerInferTpl
+from lightllm.common.basemodel.triton_kernel.post_process.vocab_parallel_greedy import (
+    is_vocab_parallel_greedy_enabled,
+    vocab_parallel_greedy,
+)
 from lightllm.distributed.communication_op import all_gather
 
 
@@ -80,7 +84,7 @@ def _token_forward(
         if prompt_logics_hiddens is not None:
             prompt_token_num = prompt_logics_hiddens.shape[0]
             infer_state.prompt_logics = self._lm_head_and_gather(
-                prompt_logics_hiddens, prompt_token_num, layer_weight, infer_state
+                prompt_logics_hiddens, prompt_token_num, layer_weight, infer_state, force_full_logits=True
             )
 
         return ans_logics
@@ -91,13 +95,23 @@ def _lm_head_and_gather(
         token_num: int,
         layer_weight: LlamaPreAndPostLayerWeight,
         infer_state: LlamaInferStateInfo,
+        force_full_logits: bool = False,
     ) -> torch.Tensor:
         normed = self._norm(hidden, infer_state, layer_weight)
         normed = normed.permute(1, 0).view(-1, token_num)
         logic_batch = layer_weight.lm_head_weight_(input=normed, alloc_func=self.alloc_tensor)
         normed = None
 
         vocab_size = layer_weight.lm_head_weight_.vocab_size
+        if is_vocab_parallel_greedy_enabled() and not force_full_logits:
+            return vocab_parallel_greedy(
+                logic_batch,
+                vocab_size=vocab_size,
+                tp_world_size=self.tp_world_size_,
+                group=infer_state.dist_group,
+                alloc_func=self.alloc_tensor,
+            )
+
         if self.tp_world_size_ == 1:
             gather_data = logic_batch
         else:

lightllm/server/router/model_infer/mode_backend/base_backend.py

@@ -45,6 +45,10 @@
 from lightllm.server.router.model_infer.mode_backend.overlap_events import OverlapEventManager, OverlapEventPack
 from lightllm.server.router.model_infer.mode_backend.generic_post_process import sample
 from lightllm.common.basemodel.triton_kernel.gather_token_id import scatter_token
+from lightllm.common.basemodel.triton_kernel.post_process.vocab_parallel_greedy import (
+    is_vocab_parallel_greedy_enabled,
+    unpack_vocab_parallel_greedy,
+)
 from lightllm.server.pd_io_struct import PDChunckedTransTaskRet
 from .multi_level_kv_cache import MultiLevelKvCacheModule
 from lightllm.utils.profiler import ProcessProfiler, ProfilerCmd
@@ -383,6 +387,13 @@ def _get_next_token_ranks(self, logits: torch.Tensor, next_token_ids: torch.Tens
         仅 ``--enable_rl`` 时做真实 rank；否则返回 GPU 常量 ``-1``，避免 O(batch * vocab) 比较。
         下游 async_copy 在同样条件下会忽略该返回值。
         """
+        if is_vocab_parallel_greedy_enabled():
+            return g_pin_mem_manager.get_const_gpu_tensor(
+                key="next_token_ranks",
+                shape=next_token_ids.shape,
+                fill_value=1 if self.args.enable_rl else -1,
+                dtype=torch.int32,
+            )
         if not self.args.enable_rl:
             return g_pin_mem_manager.get_const_gpu_tensor(
                 key="next_token_ranks",
@@ -840,10 +851,16 @@ def _trans_req_ids_to_req_objs(self, req_ids: List[int]) -> List[InferReq]:
 
     def _gen_argmax_token_ids(self, model_output: ModelOutput):
         logits = model_output.logits
+        if is_vocab_parallel_greedy_enabled():
+            token_ids, _ = unpack_vocab_parallel_greedy(logits)
+            return token_ids
         return torch.argmax(logits, dim=-1)
 
     def _gen_argmax_token_ids_and_prob(self, model_output: ModelOutput):
         logits = model_output.logits
+        if is_vocab_parallel_greedy_enabled():
+            token_ids, token_logprobs = unpack_vocab_parallel_greedy(logits)
+            return token_ids, torch.exp(token_logprobs)
         probs = torch.softmax(logits, dim=-1)
         max_probs, draft_next_token_ids_gpu = torch.max(probs, dim=-1)
         return draft_next_token_ids_gpu, max_probs
@@ -860,6 +877,8 @@ def _sample_and_scatter_token(
     ):
 
         if mask_func is not None:
+            if is_vocab_parallel_greedy_enabled():
+                raise RuntimeError("LIGHTLLM_VOCAB_PARALLEL_GREEDY does not support constrained logits")
             assert len(run_reqs) == logits.shape[0]
             mask_func(run_reqs, logits)