Implementation:Hpcaitech ColossalAI Stream Chat Patch

Knowledge Sources	Hpcaitech_ColossalAI
Domains	Natural Language Processing, Text Generation, Streaming Inference
Last Updated	2026-02-09 00:00 GMT

Overview

Streaming chat utility that provides token-by-token text generation for Colossal-LLaMA models.

Description

This module implements streaming chat functionality for Colossal-LLaMA language models. It provides three main functions: get_prompt_template for formatting conversation history into a prompt string, streaming_chat for generating responses in a streaming fashion with configurable sampling parameters, and stream_generate for low-level autoregressive token generation with support for incremental decoding via past key values. The implementation is adapted from the ChatGLM3 streaming generation approach and supports left-padded tokenization.

Usage

Use this module when deploying Colossal-LLaMA models for interactive chat applications that require real-time streaming of generated text tokens. It is particularly useful for building chat interfaces where users need to see responses as they are generated rather than waiting for the complete output.

Code Reference

Source Location

• Repository: Hpcaitech_ColossalAI
• File: applications/Colossal-LLaMA/colossal_llama/utils/stream_chat_patch.py
• Lines: 1-253

Signature

def get_prompt_template(
    input_query: str,
    history: List[Dict] = None,
    roles: list = ["", "Human", "Assistant"],
) -> str:

@torch.inference_mode()
def streaming_chat(
    model: Any,
    tokenizer: PreTrainedTokenizer,
    input_query: str,
    history: List[Dict] = None,
    roles: list = ["", "Human", "Assistant"],
    past_key_values: Tuple[Tuple[torch.FloatTensor, Any], Any] = None,
    temperature: float = 0.8,
    top_p: float = 0.95,
    top_k: int = 50,
    do_sample: bool = True,
    length_penalty: float = 1.2,
    max_new_tokens: int = 512,
    logits_processor: LogitsProcessorList = None,
    return_past_key_values: bool = False,
    **kwargs,
):

@torch.inference_mode()
def stream_generate(
    model: Any,
    input_ids: torch.Tensor,
    generation_config: Optional[GenerationConfig] = None,
    logits_processor: Optional[LogitsProcessorList] = None,
    stopping_criteria: Optional[StoppingCriteriaList] = None,
    prefix_allowed_tokens_fn: Optional[Callable[[int, torch.Tensor], List[int]]] = None,
    return_past_key_values: bool = False,
    **kwargs,
):

Import

from colossal_llama.utils.stream_chat_patch import get_prompt_template, streaming_chat, stream_generate

I/O Contract

Inputs (streaming_chat)

Name	Type	Required	Description
model	Any	Yes	The language model to generate responses
tokenizer	PreTrainedTokenizer	Yes	Tokenizer compatible with the model (must use left padding)
input_query	str	Yes	The current user input to respond to
history	List[Dict]	No	List of past conversations with 'role' and 'message' keys
roles	list	No	Roles in the conversation, defaults to ["", "Human", "Assistant"]
past_key_values	Tuple	No	Past key values for incremental decoding
temperature	float	No	Temperature for token sampling, defaults to 0.8
top_p	float	No	Nucleus sampling probability threshold, defaults to 0.95
top_k	int	No	Top-K filtering threshold, defaults to 50
do_sample	bool	No	Whether to sample responses, defaults to True
length_penalty	float	No	Penalty for response length, defaults to 1.2
max_new_tokens	int	No	Maximum number of new tokens to generate, defaults to 512
logits_processor	LogitsProcessorList	No	Custom logits processors
return_past_key_values	bool	No	Whether to return past key values, defaults to False

Outputs

Name	Type	Description
response	str	The generated text response (yielded incrementally)
history	List[Dict]	Updated conversation history
past_key_values	Tuple (optional)	Updated past key values if return_past_key_values is True

Usage Examples

from colossal_llama.utils.stream_chat_patch import streaming_chat, get_prompt_template

# Stream responses from a Colossal-LLaMA model
for response, history in streaming_chat(
    model=model,
    tokenizer=tokenizer,
    input_query="What is ColossalAI?",
    history=[],
    temperature=0.8,
    max_new_tokens=512,
):
    print(response, end="", flush=True)

# Generate a prompt template
prompt = get_prompt_template(
    input_query="Hello, how are you?",
    history=[{"role": "Human", "message": "Hi"}, {"role": "Assistant", "message": "Hello!"}],
)

Related Pages

• Environment:Hpcaitech_ColossalAI_CUDA_GPU_Environment