#003

causal_conv1d

bf16 sglang · prefill · sglang.srt.layers.attention.mamba.causal_conv1d_triton · importance 2.2%

Reference Implementation

reference.py

import torch
import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module):

    def __init__(self, activation: str | None='silu') -> None:
        super().__init__()
        self.activation = activation

    def forward(self, x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor, initial_state: torch.Tensor) -> torch.Tensor:
        dtype = x.dtype
        seq = x.unsqueeze(0).to(weight.dtype)
        state = initial_state.to(weight.dtype)
        width = weight.shape[1]
        seq_with_state = torch.cat([state[:, :, -(width - 1):], seq], dim=-1)
        out = F.conv1d(seq_with_state, weight.unsqueeze(1), bias, groups=weight.shape[0])
        out = out[:, :, -x.shape[-1]:]
        if self.activation in ('silu', 'swish'):
            out = F.silu(out)
        elif self.activation is not None:
            raise ValueError(f'unsupported activation: {self.activation}')
        return out.squeeze(0).to(dtype)

Shapes

T_SOL hardware:

#	token_count	dim	width	T_SOL(XPU-A)	T_Prod	S
0	1	4096	4	0.02 us	7.60 us	0.3%
1	128	4096	4	0.41 us	35.70 us	1.1%
2	256	4096	4	0.80 us	77.10 us	1.0%
3	514	4096	4	1.60 us	100.60 us	1.6%
4	1024	4096	4	3.18 us	272.40 us	1.2%
5	1	8192	4	0.03 us	7.70 us	0.4%
6	4195	8192	4	25.96 us	1.33 ms	1.9%
7	11027	8192	4	68.20 us	3.50 ms	2.0%
8	14807	8192	4	91.57 us	4.66 ms	2.0%

Input Generation

input.py

import torch

def _make_inputs(token_count: int, dim: int, width: int=4) -> dict[str, torch.Tensor]:
    x = torch.randn(dim, token_count, dtype=torch.bfloat16, device='cuda') * 0.02
    weight = torch.randn(dim, width, dtype=torch.bfloat16, device='cuda') * 0.02
    bias = torch.randn(dim, dtype=torch.bfloat16, device='cuda') * 0.02
    initial_state = torch.randn(1, dim, width - 1, dtype=torch.bfloat16, device='cuda') * 0.02
    return {'x': x, 'weight': weight, 'bias': bias, 'initial_state': initial_state}