refactor context parallel state (#17213)

Co-authored-by: Shunkang <182541032+Shunkangz@users.noreply.github.co>

test/registered/8-gpu-models/test_deepseek_v32_cp_single_node.py

@@ -18,6 +18,7 @@ BASE_ARGS = [
 DP_ARGS = [
     "--tp=8",
     "--dp=2",
+    "--attn-cp-size=4",
     "--enable-dp-attention",
 ]
 
@@ -43,6 +44,7 @@ CP_IN_SEQ_SPLIT_ARGS = [
 CP_ROUND_ROBIN_ARGS = [
     "--enable-nsa-prefill-context-parallel",
     "--nsa-prefill-cp-mode=round-robin-split",
+    "--attn-cp-size=8",
 ]