Skip to content

perf(mamba): prefill text prompts in multi-token chunks #2184

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
chaofengw-nv wants to merge 1 commit into
base: main
Choose a base branch
from
+574 −13
Open

perf(mamba): prefill text prompts in multi-token chunks #2184

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
181 changes: 178 additions & 3 deletions families/mamba/model.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,7 @@
_transpose_2d,
)
from . import graph_ops
from . import scan_rows
from .parallel import normalize_parallel_config


Expand Down Expand Up @@ -202,9 +203,15 @@ def build_engine(
verbose: bool = False,
debug_layer_outputs: bool = False,
parallel_config=None,
prefill_rows: int | None = None,
) -> bytes:
"""Build TRT engine for Mamba SSM.

``prefill_rows`` builds the prefill engine of a split bundle instead:
``token_id`` is ``[S]`` with one profile ``1..prefill_rows``, every layer
runs its selective scan over the ``S`` rows, the logits come from the
last row and the present states are those after the last row.

max_cache_length is accepted for the shared build signature but is not used
by Mamba (SSM state is constant size regardless of sequence length).

Expand All @@ -219,6 +226,10 @@ def build_engine(
present_ssm_0..N: float32 [1, d_inner * state_size]
"""
parallel = normalize_parallel_config(parallel_config)
if prefill_rows is not None and (parallel.enabled or debug_layer_outputs):
raise NotImplementedError("Mamba prefill engines are single-GPU without debug outputs")
if prefill_rows is not None and not 1 <= prefill_rows <= scan_rows.MAX_ROWS:
raise ValueError(f"Mamba prefill_rows must be in 1..{scan_rows.MAX_ROWS}")
if parallel.enabled:
parallel.validate()
from .tp_builder import build_mamba_tp_engine
Expand Down Expand Up @@ -257,7 +268,14 @@ def build_engine(
# -----------------------------------------------------------
# Inputs
# -----------------------------------------------------------
token_id = network.add_input("token_id", trt.int32, (1,))
rows = 1 if prefill_rows is None else None
token_id = network.add_input("token_id", trt.int32, (1 if rows else -1,))
if prefill_rows is not None:
profile = builder.create_optimization_profile()
profile.set_shape(
"token_id", (1,), (min(PREFILL_OPT_TOKENS, prefill_rows),), (prefill_rows,)
)
trt_config.add_optimization_profile(profile)

conv_state_inputs = []
ssm_state_inputs = []
Expand Down Expand Up @@ -323,6 +341,7 @@ def build_engine(
conv_kernel=conv_kernel,
dt_rank=dt_rank,
dtype=work_np_dtype,
sequence_length=rows,
)

hidden_state = result["hidden"]
Expand All @@ -332,6 +351,10 @@ def build_engine(
if debug_layer_outputs:
_mark_debug_output(network, hidden_state, f"debug_hidden_{layer_idx}")

if rows is None:
# Only the last row can affect the next token.
hidden_state = scan_rows.last_row(network, hidden_state, hidden)

# -----------------------------------------------------------
# Final norm
# -----------------------------------------------------------
Expand Down Expand Up @@ -418,8 +441,9 @@ def _add_mamba_layer(
conv_kernel: int,
dt_rank: int,
dtype: np.dtype = np.float32,
sequence_length: int | None = 1,
) -> dict[str, trt.ITensor]:
"""Add one Mamba SSM layer.
"""Add one Mamba SSM layer over one row, or ``S`` rows when ``sequence_length`` is None.

Returns: {hidden, present_conv, present_ssm}
"""
Expand All @@ -437,6 +461,38 @@ def _add_mamba_layer(
z = graph_ops.add_matmul_rhs_constant(
network, normed, hidden_size, d_inner, weights[f"{prefix}.w_in_z"], dtype=dtype
)
if sequence_length is None:
y, present_conv, present_ssm = _add_mamba_mixer_rows(
network,
x,
weights=weights,
prefix=prefix,
conv_state_in=conv_state_in,
ssm_state_in=ssm_state_in,
d_inner=d_inner,
state_size=state_size,
conv_kernel=conv_kernel,
dt_rank=dt_rank,
dtype=dtype,
)
if y.dtype != hidden.dtype:
y = network.add_cast(y, hidden.dtype).get_output(0)
z_activated = graph_ops.add_activation(network, z, "silu", dtype=dtype)
gated = network.add_elementwise(y, z_activated, trt.ElementWiseOperation.PROD)
out = graph_ops.add_matmul_rhs_constant(
network,
gated.get_output(0),
d_inner,
hidden_size,
weights[f"{prefix}.w_out"],
dtype=dtype,
)
residual = network.add_elementwise(hidden, out, trt.ElementWiseOperation.SUM)
return {
"hidden": residual.get_output(0),
"present_conv": present_conv,
"present_ssm": present_ssm,
}

# ===== 3. Conv1d step =====
# conv_state_in: [d_inner, conv_kernel]
Expand Down Expand Up @@ -632,6 +688,113 @@ def _add_mamba_layer(
}


def _add_mamba_mixer_rows(
network: trt.INetworkDefinition,
x: trt.ITensor,
*,
weights: WeightDict,
prefix: str,
conv_state_in: trt.ITensor,
ssm_state_in: trt.ITensor,
d_inner: int,
state_size: int,
conv_kernel: int,
dt_rank: int,
dtype: np.dtype,
) -> tuple[trt.ITensor, trt.ITensor, trt.ITensor]:
"""Conv and selective scan over ``x [S, d_inner]``: ``(y [S, d_inner] FP32, conv, ssm)``.

Same arithmetic as the one-row layer applied token by token; the scan
combines the rows as described in :mod:`.scan_rows`.
"""
prod_op = trt.ElementWiseOperation.PROD

def fp32(tensor):
return (
tensor
if tensor.dtype == trt.float32
else network.add_cast(tensor, trt.float32).get_output(0)
)

conv_w = graph_ops.add_constant(
network, (d_inner, conv_kernel), weights[f"{prefix}.conv1d_weight"], dtype=dtype
)
conv_columns, present_conv = scan_rows.add_short_conv_rows(
network,
columns=scan_rows.shuffle(network, x, transpose=(1, 0)),
conv_state=conv_state_in,
weight=conv_w,
channels=d_inner,
kernel=conv_kernel,
)
conv_out = graph_ops.add_bias_sum(
network,
scan_rows.shuffle(network, conv_columns, transpose=(1, 0)),
d_inner,
weights[f"{prefix}.conv1d_bias"],
dtype=dtype,
)
activated = graph_ops.add_activation(network, conv_out, "silu", dtype=dtype) # [S, d]

dt_in = graph_ops.add_matmul_rhs_constant(
network, activated, d_inner, dt_rank, weights[f"{prefix}.w_dt_in"], dtype=dtype
)
b = graph_ops.add_matmul_rhs_constant(
network, activated, d_inner, state_size, weights[f"{prefix}.w_B"], dtype=dtype
)
c = graph_ops.add_matmul_rhs_constant(
network, activated, d_inner, state_size, weights[f"{prefix}.w_C"], dtype=dtype
)
dt = graph_ops.add_matmul_rhs_constant(
network, dt_in, dt_rank, d_inner, weights[f"{prefix}.w_dt_out"], dtype=dtype
)
dt = fp32(
graph_ops.add_bias_sum(network, dt, d_inner, weights[f"{prefix}.dt_proj_bias"], dtype=dtype)
)
one = graph_ops.add_constant(network, (1, 1), np.array([1.0], dtype=np.float32))
dt = network.add_unary(
network.add_elementwise(
network.add_unary(dt, trt.UnaryOperation.EXP).get_output(0),
one,
trt.ElementWiseOperation.SUM,
).get_output(0),
trt.UnaryOperation.LOG,
).get_output(0) # softplus, [S, d]

dt_3d = scan_rows.shuffle(network, dt, reshape=(-1, d_inner, 1))
a_const = graph_ops.add_constant(
network, (1, d_inner, state_size), weights[f"{prefix}.A"].reshape(1, d_inner, state_size)
)
decay = network.add_unary(
network.add_elementwise(dt_3d, a_const, prod_op).get_output(0), trt.UnaryOperation.EXP
).get_output(0) # exp(dt A): [S, d, n]
x_3d = scan_rows.shuffle(network, fp32(activated), reshape=(-1, d_inner, 1))
b_3d = scan_rows.shuffle(network, fp32(b), reshape=(-1, 1, state_size))
update = network.add_elementwise(
network.add_elementwise(dt_3d, b_3d, prod_op).get_output(0), x_3d, prod_op
).get_output(0) # dt B x: [S, d, n]
states, present_ssm = scan_rows.add_selective_scan_rows(
network, decay=decay, update=update, state=ssm_state_in, width=(d_inner, state_size)
)
c_3d = scan_rows.shuffle(network, fp32(c), reshape=(-1, 1, state_size))
y = network.add_reduce(
network.add_elementwise(states, c_3d, prod_op).get_output(0),
trt.ReduceOperation.SUM,
1 << 2,
keep_dims=False,
).get_output(0) # [S, d]
d_const = graph_ops.add_constant(network, (1, d_inner), weights[f"{prefix}.D"])
y = network.add_elementwise(
y,
network.add_elementwise(d_const, fp32(activated), prod_op).get_output(0),
trt.ElementWiseOperation.SUM,
).get_output(0)
return y, present_conv, present_ssm


# Optimum rows of the prefill engine.
PREFILL_OPT_TOKENS = 32

_BUNDLE_FILES = (
"tokenizer.json",
"tokenizer_config.json",
Expand Down Expand Up @@ -757,7 +920,19 @@ def build(request: "BuildRequest", writer: "BundleWriter") -> None:
parallel_config=parallel,
)
writer.add_bytes("engine.plan", plan)
layout = "single"
# A second engine prefills the prompt in chunks; decode keeps the one-row engine.
writer.add_bytes(
"prefill.plan",
model.build_engine(
config,
weights,
max_sequence_length,
precision=precision,
verbose=bool(request.verbose),
prefill_rows=scan_rows.MAX_ROWS,
),
)
layout = "split"
writer.add_json(
"runtime.json",
_runtime_config(
Expand Down
52 changes: 45 additions & 7 deletions families/mamba/runtime/pipeline.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,11 +29,21 @@ RecurrentPipeline::RecurrentPipeline(std::unique_ptr<ITrtModule> decoder,
std::unique_ptr<MambaInferenceState> state,
RecurrentGenConfig config, cudaStream_t stream,
const char* name, std::shared_ptr<ITokenizer> tokenizer,
std::string model_id_str)
: decoder_(std::move(decoder)), state_(std::move(state)), config_(config), stream_(stream),
name_(name), tokenizer_(std::move(tokenizer)), model_id_(std::move(model_id_str)) {
std::string model_id_str, std::unique_ptr<ITrtModule> prefill)
: decoder_(std::move(decoder)), state_(std::move(state)), prefill_(std::move(prefill)),
config_(config), stream_(stream), name_(name), tokenizer_(std::move(tokenizer)),
model_id_(std::move(model_id_str)) {
if (!decoder_ || !decoder_->ok())
throw std::runtime_error(std::string(name_) + ": invalid decoder module");
if (prefill_) {
if (!prefill_->ok())
throw std::runtime_error(std::string(name_) + ": invalid prefill module");
const auto rows = prefill_->input_profile_shape("token_id", prefill_->profile_idx(),
ProfileShapeSelector::kMax);
if (rows.size() != 1 || rows[0] < 1)
throw std::runtime_error(std::string(name_) + ": invalid prefill token profile");
prefill_rows_ = static_cast<int32_t>(rows[0]);
}
}

static std::vector<int32_t> encode_prompt(const ITokenizer& tokenizer,
Expand Down Expand Up @@ -103,10 +113,12 @@ std::vector<int32_t> RecurrentPipeline::generate_from_ids(const std::vector<int3

// ── Prefill phase ──
auto t_prefill_start = SteadyClock::now();
for (std::size_t i = 0; i + 1 < input_ids.size(); ++i)
run_step(input_ids[i], logits);

run_step(input_ids.back(), logits);
if (prefill_) {
run_prefill_chunks(input_ids, logits);
} else {
for (const auto token : input_ids)
run_step(token, logits);
}
auto t_prefill_end = SteadyClock::now();

// ── Decode phase ──
Expand Down Expand Up @@ -182,6 +194,32 @@ void RecurrentPipeline::report_timing(SteadyClock::time_point t_prefill_start,
}
}

void RecurrentPipeline::run_prefill_chunks(const std::vector<int32_t>& input_ids,
std::vector<float>& logits) {
// The prefill engine reads and writes the same recurrent state buffers.
state_->bind_to(*prefill_);
const auto total = static_cast<int32_t>(input_ids.size());
for (int32_t start = 0; start < total; start += prefill_rows_) {
const int32_t rows = std::min(prefill_rows_, total - start);
TensorMap inputs;
inputs["token_id"] =
Tensor{const_cast<int32_t*>(input_ids.data() + start), {rows}, DType::kInt32};
state_->prepare_step(inputs, rows);
prefill_->forward_async(inputs);
state_->advance(rows);
}
prefill_->sync();
const auto* device_logits = prefill_->device_ptr("logits");
if (device_logits == nullptr)
throw std::runtime_error(std::string(name_) + ": prefill engine has no logits output");
logits.resize(static_cast<std::size_t>(config_.vocab_size));
const cudaError_t copied = cudaMemcpy(logits.data(), device_logits,
logits.size() * sizeof(float), cudaMemcpyDeviceToHost);
if (copied != cudaSuccess)
throw std::runtime_error(std::string(name_) +
": failed to copy prefill logits: " + cudaGetErrorString(copied));
}

void RecurrentPipeline::run_step(int32_t token_id, std::vector<float>& logits) {
auto t0 = SteadyClock::now();

Expand Down
8 changes: 7 additions & 1 deletion families/mamba/runtime/pipeline.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,7 @@ class RecurrentPipeline final : public ITextGeneration {
std::unique_ptr<MambaInferenceState> state, RecurrentGenConfig config,
cudaStream_t stream, const char* name,
std::shared_ptr<ITokenizer> tokenizer = nullptr,
std::string model_id_str = "");
std::string model_id_str = "", std::unique_ptr<ITrtModule> prefill = nullptr);

TextResult generate(const std::string& prompt, const TextGenerationConfig& cfg = {}) override;
int32_t default_max_new_tokens() const override { return 128; }
Expand All @@ -51,6 +51,9 @@ class RecurrentPipeline final : public ITextGeneration {
private:
std::unique_ptr<ITrtModule> decoder_;
std::unique_ptr<MambaInferenceState> state_;
// Multi-row prefill engine of a split bundle; decode stays on decoder_.
std::unique_ptr<ITrtModule> prefill_;
int32_t prefill_rows_{0};
RecurrentGenConfig config_;
cudaStream_t stream_;
const char* name_;
Expand All @@ -62,6 +65,9 @@ class RecurrentPipeline final : public ITextGeneration {
const MambaSamplingParams& params);

void run_step(int32_t token_id, std::vector<float>& logits);
// Feed the prompt to prefill_ in chunks of up to prefill_rows_ tokens and
// return the logits of the last token.
void run_prefill_chunks(const std::vector<int32_t>& input_ids, std::vector<float>& logits);

using SteadyClock = std::chrono::steady_clock;
void report_timing(SteadyClock::time_point t_prefill_start,
Expand Down
18 changes: 16 additions & 2 deletions families/mamba/runtime/plugin.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -77,7 +77,8 @@ ITask* create(const FamilyContext& context) {
if (hidden_size <= 0 || num_layers <= 0 || vocab_size <= 0 || inner <= 0 || state_size <= 0 ||
conv_kernel <= 1 || tp_size <= 0 || max_cache_length <= 0 || num_attention_heads <= 0 ||
num_key_value_heads <= 0 || head_dim <= 0 || expand <= 0 ||
(layout != "single" && layout != "dual_profile") ||
(layout != "single" && layout != "dual_profile" && layout != "split") ||
(layout == "split" && tp_size != 1) ||
tp_mode != (tp_size > 1 ? "tensor_parallel" : "single") ||
(precision != "fp16" && precision != "bf16" && precision != "fp32")) {
throw std::runtime_error("mamba runtime.json contains invalid geometry");
Expand All @@ -97,6 +98,18 @@ ITask* create(const FamilyContext& context) {
throw std::runtime_error("mamba failed to load decoder");
decoder->set_timing_label("mamba decoder");
const cudaStream_t stream = decoder->stream();
std::unique_ptr<ITrtModule> prefill;
if (layout == "split") {
// The prefill engine shares the decoder's stream so state copies stay ordered.
ModuleCreateOptions prefill_options = options;
prefill_options.stream = stream;
const auto& prefill_plan = require_section(context.reader, "prefill.plan");
prefill = context.backend.create_module(prefill_plan.data(), prefill_plan.size(),
prefill_options);
if (prefill == nullptr || !prefill->ok())
throw std::runtime_error("mamba failed to load the prefill engine");
prefill->set_timing_label("mamba prefill");
}

std::vector<MambaRecurrentState::TensorSpec> specs = {
{"conv_state", {inner, conv_kernel}, "present_conv"},
Expand All @@ -114,7 +127,8 @@ ITask* create(const FamilyContext& context) {
generation.chat_template_format =
mamba_detect_chat_template_format(chat_template(context.reader));
return new RecurrentPipeline(std::move(decoder), std::move(state), std::move(generation),
stream, "Mamba", create_tokenizer(context.reader), std::string{});
stream, "Mamba", create_tokenizer(context.reader), std::string{},
std::move(prefill));
}

} // namespace trtmc::mamba
Expand Down
Loading
Loading