Fix Mamba2Metadata dropping has_mamba_track_mask (#26877)

This commit is contained in:
Ke Bao
2026-06-01 22:03:59 +08:00
committed by GitHub
parent fd16e05252
commit 693adabff7
2 changed files with 2 additions and 1 deletions
@@ -176,6 +176,7 @@ class Mamba2Metadata(ForwardMetadata):
track_ssm_h_dst=forward_metadata.track_ssm_h_dst,
track_ssm_final_src=forward_metadata.track_ssm_final_src,
track_ssm_final_dst=forward_metadata.track_ssm_final_dst,
has_mamba_track_mask=forward_metadata.has_mamba_track_mask,
num_decodes=len(seq_lens),
num_prefills=0,
num_prefill_tokens=0,
@@ -249,6 +250,7 @@ class Mamba2Metadata(ForwardMetadata):
track_ssm_h_dst=forward_metadata.track_ssm_h_dst,
track_ssm_final_src=forward_metadata.track_ssm_final_src,
track_ssm_final_dst=forward_metadata.track_ssm_final_dst,
has_mamba_track_mask=forward_metadata.has_mamba_track_mask,
num_prefills=num_prefills,
num_prefill_tokens=num_prefill_tokens,
num_decodes=num_decodes,