/tag-run-ci-label

/tag-and-rerun-ci

/rerun-failed-ci

@Shunkangz Thank you for the contribution, I left some reviews. There will be another round of more thorough review. We can discuss in the slack channel as well. Couple of things I want to call out:

• During development of MLA CP, I found out two notable bugs in current MHA CP impl. 1 is

  sglang/python/sglang/srt/layers/attention/flashattention_backend.py

  Lines 492 to 508 in 5466e9f

  	# MLA/MHA CP: prepare_mlp_sync_batch pads extend tokens up to
  	# lcm(attn_tp_size, attn_cp_size), so cache_seqlens_cp can exceed
  	# seq_lens_cpu.max(). Widen page_table by the pad delta to keep
  	# FA3's causal reads in-bounds; widened columns index KV slot 0
  	# (req_to_token is zero-init) and outputs for padding queries are
  	# discarded downstream.
  	if (
  	self.attn_cp_size > 1
  	and forward_batch.global_num_tokens_cpu is not None
  	and forward_batch.extend_num_tokens is not None
  	and forward_batch.extend_seq_lens_cpu is not None
  	):
  	padded_extend = int(forward_batch.extend_num_tokens)
  	real_extend = int(sum(forward_batch.extend_seq_lens_cpu))
  	pad_delta = padded_extend - real_extend
  	if pad_delta > 0:
  	metadata.max_seq_len_k += pad_delta

  
  and 2 is

  sglang/python/sglang/srt/layers/utils/cp_utils.py

  Lines 471 to 472 in 5466e9f

  	# Derive prefix offset from unpadded CPU tensors. Both `seqs_len` and `extend_lens` are unpadded by the caller
  	# Using the padded `kv_len` here would undercount `prefix_len` by the padding amount and shift the FA causal horizon.

• We need MHA CP zigzag mode test for attn_cp_size == 4 and bs > 1

Thank you for pointing this out. In this PR, I only want to support the bs > 1 with GQA model. For MLA CP part, I left it as original implementation. I believe that the MLA CP should be refactored and aligned with our existing logic such as args, layer communicator and so on.

@Shunkangz Thank you for the contribution, I left some reviews. There will be another round of more thorough review. We can discuss in the slack channel as well. Couple of things I want to call out:

• During development of MLA CP, I found out two notable bugs in current MHA CP impl. 1 is

  sglang/python/sglang/srt/layers/attention/flashattention_backend.py

  Lines 492 to 508 in 5466e9f

  	# MLA/MHA CP: prepare_mlp_sync_batch pads extend tokens up to
  	# lcm(attn_tp_size, attn_cp_size), so cache_seqlens_cp can exceed
  	# seq_lens_cpu.max(). Widen page_table by the pad delta to keep
  	# FA3's causal reads in-bounds; widened columns index KV slot 0
  	# (req_to_token is zero-init) and outputs for padding queries are
  	# discarded downstream.
  	if (
  	self.attn_cp_size > 1
  	and forward_batch.global_num_tokens_cpu is not None
  	and forward_batch.extend_num_tokens is not None
  	and forward_batch.extend_seq_lens_cpu is not None
  	):
  	padded_extend = int(forward_batch.extend_num_tokens)
  	real_extend = int(sum(forward_batch.extend_seq_lens_cpu))
  	pad_delta = padded_extend - real_extend
  	if pad_delta > 0:
  	metadata.max_seq_len_k += pad_delta

  and 2 is

  sglang/python/sglang/srt/layers/utils/cp_utils.py

  Lines 471 to 472 in 5466e9f

  	# Derive prefix offset from unpadded CPU tensors. Both `seqs_len` and `extend_lens` are unpadded by the caller
  	# Using the padded `kv_len` here would undercount `prefix_len` by the padding amount and shift the FA causal horizon.

• We need MHA CP zigzag mode test for attn_cp_size == 4 and bs > 1

Thank you for pointing this out. In this PR, I only want to support the bs > 1 with GQA model. For MLA CP part, I left it as original implementation. I believe that the MLA CP should be refactored and aligned with our existing logic such as args, layer communicator and so on.

Ah sorry I should be clearer here:

1. is an issue about the padding happening in prepare_mlp_sync_batch causing cache_seqlens_cp to go over seq_lens_cpu.max(). This bug occur also for MHA CP and will likely impact this PR.
2. is an issue about padded kv_len messing up metadata computation. I think we agree on this already Support batch size > 1 when enable CP #23269 (comment)

@Shunkangz Thank you for the contribution, I left some reviews. There will be another round of more thorough review. We can discuss in the slack channel as well. Couple of things I want to call out:

• During development of MLA CP, I found out two notable bugs in current MHA CP impl. 1 is

  sglang/python/sglang/srt/layers/attention/flashattention_backend.py

  Lines 492 to 508 in 5466e9f

  	# MLA/MHA CP: prepare_mlp_sync_batch pads extend tokens up to
  	# lcm(attn_tp_size, attn_cp_size), so cache_seqlens_cp can exceed
  	# seq_lens_cpu.max(). Widen page_table by the pad delta to keep
  	# FA3's causal reads in-bounds; widened columns index KV slot 0
  	# (req_to_token is zero-init) and outputs for padding queries are
  	# discarded downstream.
  	if (
  	self.attn_cp_size > 1
  	and forward_batch.global_num_tokens_cpu is not None
  	and forward_batch.extend_num_tokens is not None
  	and forward_batch.extend_seq_lens_cpu is not None
  	):
  	padded_extend = int(forward_batch.extend_num_tokens)
  	real_extend = int(sum(forward_batch.extend_seq_lens_cpu))
  	pad_delta = padded_extend - real_extend
  	if pad_delta > 0:
  	metadata.max_seq_len_k += pad_delta

  and 2 is

  sglang/python/sglang/srt/layers/utils/cp_utils.py

  Lines 471 to 472 in 5466e9f

  	# Derive prefix offset from unpadded CPU tensors. Both `seqs_len` and `extend_lens` are unpadded by the caller
  	# Using the padded `kv_len` here would undercount `prefix_len` by the padding amount and shift the FA causal horizon.

• We need MHA CP zigzag mode test for attn_cp_size == 4 and bs > 1

Thank you for pointing this out. In this PR, I only want to support the bs > 1 with GQA model. For MLA CP part, I left it as original implementation. I believe that the MLA CP should be refactored and aligned with our existing logic such as args, layer communicator and so on.

Ah sorry I should be clearer here:

1. is an issue about the padding happening in prepare_mlp_sync_batch causing cache_seqlens_cp to go over seq_lens_cpu.max(). This bug occur also for MHA CP and will likely impact this PR.
2. is an issue about padded kv_len messing up metadata computation. I think we agree on this already Support batch size > 1 when enable CP #23269 (comment)

For 1, let's discuss in details through slack. For 2, I think that the existing TestContextParallelMetadata already cover this. Can you confirm it?

@Shunkangz Thank you for the contribution, I left some reviews. There will be another round of more thorough review. We can discuss in the slack channel as well. Couple of things I want to call out:

• During development of MLA CP, I found out two notable bugs in current MHA CP impl. 1 is

  sglang/python/sglang/srt/layers/attention/flashattention_backend.py

  Lines 492 to 508 in 5466e9f

  	# MLA/MHA CP: prepare_mlp_sync_batch pads extend tokens up to
  	# lcm(attn_tp_size, attn_cp_size), so cache_seqlens_cp can exceed
  	# seq_lens_cpu.max(). Widen page_table by the pad delta to keep
  	# FA3's causal reads in-bounds; widened columns index KV slot 0
  	# (req_to_token is zero-init) and outputs for padding queries are
  	# discarded downstream.
  	if (
  	self.attn_cp_size > 1
  	and forward_batch.global_num_tokens_cpu is not None
  	and forward_batch.extend_num_tokens is not None
  	and forward_batch.extend_seq_lens_cpu is not None
  	):
  	padded_extend = int(forward_batch.extend_num_tokens)
  	real_extend = int(sum(forward_batch.extend_seq_lens_cpu))
  	pad_delta = padded_extend - real_extend
  	if pad_delta > 0:
  	metadata.max_seq_len_k += pad_delta

  and 2 is

  sglang/python/sglang/srt/layers/utils/cp_utils.py

  Lines 471 to 472 in 5466e9f

  	# Derive prefix offset from unpadded CPU tensors. Both `seqs_len` and `extend_lens` are unpadded by the caller
  	# Using the padded `kv_len` here would undercount `prefix_len` by the padding amount and shift the FA causal horizon.

• We need MHA CP zigzag mode test for attn_cp_size == 4 and bs > 1

Thank you for pointing this out. In this PR, I only want to support the bs > 1 with GQA model. For MLA CP part, I left it as original implementation. I believe that the MLA CP should be refactored and aligned with our existing logic such as args, layer communicator and so on.

Ah sorry I should be clearer here:

1. is an issue about the padding happening in prepare_mlp_sync_batch causing cache_seqlens_cp to go over seq_lens_cpu.max(). This bug occur also for MHA CP and will likely impact this PR.
2. is an issue about padded kv_len messing up metadata computation. I think we agree on this already Support batch size > 1 when enable CP #23269 (comment)

For 1, let's discuss in details through slack. For 2, I think that the existing TestContextParallelMetadata already cover this. Can you confirm it?

Yes let's discuss 1 further in slack tomorrow. For 2, sorry what test is this?

• We need MHA CP zigzag mode test for attn_cp_size == 4 and bs > 1

I mean the TestContextParallelMetadata might already cover this.

/rerun-failed-ci