Optimize the Performance of a LM

class NDArray: # Tensor Implementation
    def _init(self, other):
        self._shape = other._shape
        self._strides = other._strides
        self._offset = other._offset
        self._device = other._device
        self._handle = other._handle

class Linear(Module):
    def __init__(
        self, in_features, out_features, bias=True, device=None, dtype="float32"
    ):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.weight = Parameter(
            init.kaiming_uniform(in_features, out_features, device=device, dtype=dtype))
        if bias:
            self.bias = Parameter(
                init.kaiming_uniform(out_features, 1, device=device, dtype=dtype)
                .reshape((1, out_features)))

    def forward(self, X: Tensor) -> Tensor:
        Z = X @ self.weight
        if not self.bias is None:
            Z = Z + self.bias.broadcast_to(Z.shape)
        return Z

class TransformerLM(nn.Module):
    def __init__(
        self,
        vocab_size: int,
        context_length: int,
        d_model: int,
        num_layers: int,
        num_heads: int,
        d_ff: int,
        attn_pdrop: float,
        residual_pdrop: float,
        **kwargs
    ):
        super().__init__()
        self.embed = nn.Embedding(vocab_size, d_model)
        self.pos_embed = nn.Parameter(torch.zeros(context_length, d_model))
        self.drop = nn.Dropout(residual_pdrop)
        self.layers = nn.Sequential(*[TransformerBlock(d_model, num_heads, d_ff, attn_pdrop, residual_pdrop) for _ in range(num_layers)])
        self.ln_f = RMSNorm(d_model)
        self.head = nn.Linear(d_model, vocab_size, bias=False)
        self.context_length = context_length
    
    def forward(self, x):
        x = self.embed(x)
        x = x + self.pos_embed[:x.size(1), :]
        x = self.drop(x)
        x = self.layers(x)
        x = self.ln_f(x)
        x = self.head(x)
        return x

{   
  	"vocab_size": 10000,
    "context_length": 256,
    "d_model": 512,
    "num_layers": 4,
    "num_heads": 16,
    "d_ff": 2048,    
}

def run_benchmark(model: nn.Module, input_data: torch.Tensor, device: torch.device,
                  warmup_steps: int, timed_steps: int, pass_type: str):
    ...
    # --- Warm-up Phase ---
    print("Running warm-up steps...")
    with context:
        for _ in range(warmup_steps):
            outputs = model(input_data)
            if pass_type == 'forward_backward':
                loss = outputs.sum() # Dummy loss
                optimizer.zero_grad(set_to_none=True) # More efficient
                loss.backward()
                optimizer.step()
            synchronize(device) # Synchronize after each warm-up step, beacuse cuda is calld asycnly

    # --- Timing Phase ---
    print(f"Running {timed_steps} timed steps...")
    total_elapsed_time = 0.0
    with context:
        for i in range(timed_steps):
            synchronize(device) # Ensure previous work is done (esp. CUDA async)
            start_time = timeit.default_timer()

            outputs = model(input_data)

            if pass_type == 'forward_backward':
                loss = outputs.sum() # Dummy loss
                optimizer.zero_grad(set_to_none=True)
                loss.backward()
                optimizer.step()

            synchronize(device) # Ensure operation completes before stopping timer
            end_time = timeit.default_timer()
            total_elapsed_time += (end_time - start_time)

    avg_time_per_step = total_elapsed_time / timed_steps
    print("--- Benchmark Complete ---")
    return avg_time_per_step

def run_profile(model: nn.Module, input_data: torch.Tensor, device: torch.device,
                profile_steps: int, pass_type: str, profile_memory: bool, with_stack: bool,
                output_dir: str = "."):
    ...

    # Define the step function to be profiled
    def profile_step():
        ...
        
    profile_step()
    synchronize(device)
    print("Starting profiler context...")

    if profile_memory and device.type == 'cuda':
        torch.cuda.memory._record_memory_history(max_entries=10000)

    with profile(activities=activities,
                 record_shapes=True,
                 profile_memory=profile_memory,
                 with_stack=with_stack,
                 schedule=torch.profiler.schedule(wait=0, warmup=0, active=1, repeat=profile_steps),
                 experimental_config=torch._C._profiler._ExperimentalConfig(verbose=True)) as prof: # with_stack MUST be True for export_stacks
        for _ in range(profile_steps):
            profile_step()
            # Need to sync *inside* the step for CUDA kernel timings in profiler results
            # but *outside* the record_functions if you want to profile the sync itself.
            # Syncing after the main work is usually best for seeing kernel times.
            synchronize(device)
            prof.step() # Signal the profiler that one step is done

    print("--- Profiling Complete ---")
		...

       with autocast(device_type=device.type, dtype=torch.float16, enabled=(device.type == 'cuda')):
            outputs = model(dummy_data)
            loss = criterion(outputs, dummy_labels)

def softmax(x, dim):
    max_vals = torch.max(x, dim=dim, keepdim=True).values
    shifted_x = x - max_vals
    x_exp = torch.exp(shifted_x)
    x_exp_sum = torch.sum(x_exp, dim=dim, keepdim=True)
    return x_exp / x_exp_sum

@triton.jit
def softmax_kernel(
    output_ptr,
    input_ptr,  
    input_row_stride,  
    output_row_stride, 
    n_cols,           
    BLOCK_SIZE_N: tl.constexpr,
):
    row_idx = tl.program_id(axis=0)
    row_start_ptr_input = input_ptr + (row_idx * input_row_stride)
    row_start_ptr_output = output_ptr + (row_idx * output_row_stride)

    col_offsets = tl.arange(0, BLOCK_SIZE_N)

    input_mask = col_offsets < n_cols

    row_data = tl.load(row_start_ptr_input + col_offsets, mask=input_mask, other=-float('inf'))
    row_max = tl.max(row_data, axis=0)

    numerator = row_data - row_max
    numerator = tl.exp(numerator)
    denominator = tl.sum(numerator, axis=0)
    softmax_output = numerator / denominator
    tl.store(row_start_ptr_output + col_offsets, softmax_output, mask=input_mask)


# --- Wrapper function to launch the kernel ---
def softmax(x: torch.Tensor):
    if not x.is_cuda:
        print("输入张量必须在 CUDA 设备上。正在移至 CUDA...")
        x = x.to('cuda')
    x = x.contiguous()
    n_rows, n_cols = x.shape
    output = torch.empty_like(x)
    BLOCK_SIZE_N = triton.next_power_of_2(n_cols)
    if BLOCK_SIZE_N > 16384 and n_cols <= 16384 : 
        BLOCK_SIZE_N = triton.next_power_of_2(n_cols) if n_cols > 0 else 1
    elif n_cols == 0:
        return output
    grid = (n_rows,) 
    softmax_kernel[grid](
        output,
        x,
        x.stride(0),
        output.stride(0),
        n_cols,
        BLOCK_SIZE_N=BLOCK_SIZE_N,
    )
    return output

const eventSource = new EventSource('/api/stream');

eventSource.onmessage = (event) => {
    const data = JSON.parse(event.data);
    document.getElementById('output').innerHTML += data.token;
};

eventSource.addEventListener('end', () => {
    eventSource.close();
});

from fastapi import Request
import openai
import asyncio

@app.get("/api/stream")
async def stream_response(request: Request):
    async def event_generator():
        response = await openai.ChatCompletion.acreate(
            model="gpt-4",
            messages=[{"role": "user", "content": "写一个递归的快速排序函数"}],
            stream=True
        )
        async for chunk in response:
            token = chunk.choices[0].delta.get("content", "")
            yield f'data: {{"token": "{token}"}}\n\n'
        yield "event: end\ndata: {}\n\n"
    
    return EventResponse(event_generator())