remove assertions that cause synchronizations

Signed-off-by: Chih-Chieh-Yang <[email protected]>

Author: cyang49

vllm/model_executor/layers/mamba/ops/ssd_chunk_scan.py

@@ -511,9 +511,6 @@ def _chunk_scan_fwd(
             # with initial states, we need to take care of how
             # seq_idx crosses the boundaries
             assert batch == 1, "chunk scan only supports initial states with batch 1"
-            # The following assertion causes D2H copy - removing for now
-            # assert initial_states.shape == (seq_idx[0].max() + 1, nheads,
-            #                                 headdim, dstate)
 
             if initial_states.shape[0] == 1:
                 # no in this case no point to use initial states

vllm/model_executor/layers/mamba/ops/ssd_state_passing.py

@@ -150,9 +150,6 @@ def _state_passing_fwd(
             #   are used for continuous batching. In which case we
             #   require seq_idx to be provided
             assert seq_idx is not None, ""
-            # The following assertion causes D2H copy - removing for now
-            # assert initial_states.shape == (seq_idx.max().item() + 1, nheads,
-            #                                 dim)
         else:
             # - this is the regular batching case, where initial
             #   states are used are for each example of the batch.