diff --git a/python/sglang/srt/layers/attention/mamba/mamba2_metadata.py b/python/sglang/srt/layers/attention/mamba/mamba2_metadata.py index 044cb8b14..32a82dbc3 100644 --- a/python/sglang/srt/layers/attention/mamba/mamba2_metadata.py +++ b/python/sglang/srt/layers/attention/mamba/mamba2_metadata.py @@ -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, diff --git a/test/registered/disaggregation/test_disaggregation_hybrid_attention.py b/test/registered/disaggregation/test_disaggregation_hybrid_attention.py index cee19e520..a993f8b58 100644 --- a/test/registered/disaggregation/test_disaggregation_hybrid_attention.py +++ b/test/registered/disaggregation/test_disaggregation_hybrid_attention.py @@ -316,7 +316,6 @@ class TestDisaggregationHybridAttentionMamba(PDDisaggregationServerBase): self.assertGreater(metrics["score"], 0.87) -@unittest.skipIf(is_in_ci(), "Temporarily disable the flaky test.") class TestDisaggregationHybridAttentionMambaExtraBuffer(PDDisaggregationServerBase): @classmethod def setUpClass(cls):