Fix Mamba2Metadata dropping has_mamba_track_mask (#26877)
This commit is contained in:
@@ -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,
|
||||
|
||||
@@ -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):
|
||||
|
||||
Reference in New Issue
Block a user