[Perf] Admit HC row ownership with the qualified MTP GEMM policy

Communicate every HC prefill partial in FP32 and preserve the established process GEMM policy so MTP packed decode remains admitted. Qualify existing true-reduction state/graph cases without relaxing the failed FP16 drift limit; full-model resource and verifier gates remain pending.

Assisted-by: Codex
Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com>

benchmarks/flashnext_prefill/DECODE_PARITY_20261015.md

@@ -100,6 +100,33 @@ the requested allocator operation; models without a host-QSA consumer do
 not apply it. CPU admission/hash tests pass; GPU graph/cache-pressure and
 matched model/headroom/performance qualification are pending.
 
+## HC prefill admission under MTP
+
+Both allocator384 and512 candidates fail actual prefill while still leaving
+physical bytes free. The cached fallback requests144/160MiB full residual
+temporaries. Inspection finds that the HC row-owner admission additionally
+requires FP16 reduced-precision reductions to be disabled. The established
+MTP decode policy retains them, so requesting HC sequence prefill alone did
+not select the measured low-memory path. The earlier whole-model result is
+not evidence that row ownership executed; its stage tests explicitly set
+the different reduction policy.
+
+The repair preserves the existing process GEMM policy and MTP packed routes.
+HC prefill partials communicate in FP32 for attention's single contribution
+as well as MoE's separate contributions, then retain the existing FP16
+combine/projection interface. Admission may therefore preserve either
+qualified reduction setting; FP16 accumulation and high-level captured
+prefill remain excluded. A route log distinguishes actual execution from
+the requested flag. Other model/shape/topology guards remain in place.
+
+The true-reduction FP16 transport screen fails the original48-step factor
+allowance at8192 rows and is retained as failed. FP32 transport passes the
+unchanged independent-reference and48-step allowances at7344/8192 on all
+four ranks. Model-adapter padding, changed-input capture and compilation
+checks pass4096/4097/7344/7345/8192 with the actual MTP reduction setting.
+CPU admission/owner tests pass43 cases. Whole-model quality, selected route,
+memory and verifier curves remain required before promoting this repair.
+
 ## Rejected screens
 
 Warp-aggregated histograms regress random inputs. Warp-local emission

benchmarks/flashnext_prefill/hc_sequence_rows.py

@@ -142,6 +142,10 @@ def main():
     )
     parser.add_argument("--tokens", type=int, nargs="+", default=[7344, 8976, 11424])
     parser.add_argument("--drift-steps", type=int, default=0)
+    parser.add_argument("--reduced-precision", action="store_true")
+    parser.add_argument(
+        "--drift-communication", choices=("fp16", "fp32"), default="fp16"
+    )
     args = parser.parse_args()
     rank = int(os.environ["LOCAL_RANK"])
     torch.accelerator.set_device_index(rank)
@@ -149,7 +153,9 @@ def main():
     assert dist.get_world_size() == 4
     nccl = PyNcclCommunicator(dist.group.WORLD, rank)
     assert not nccl.disabled
-    torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False
+    torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = (
+        args.reduced_precision
+    )
     report = {"cases": [], "paired_max_row_l2_limit": 0.001}
     passed = True
     try:
@@ -288,7 +294,12 @@ def main():
                         ch, cb, ci = control(
                             hidden=ch, injection=ci, partial=(noise + cb * 0.025).half()
                         )
-                        sh, sb, si = candidate(
+                        drift_candidate = (
+                            fp32_candidate
+                            if args.drift_communication == "fp32"
+                            else candidate
+                        )
+                        sh, sb, si = drift_candidate(
                             hidden=sh, injection=si, partial=(noise + sb * 0.025).half()
                         )
                         rh, rb, ri = fp32_reference(

benchmarks/flashnext_prefill/hc_sequence_runtime.py

@@ -194,13 +194,16 @@ def main():
     parser = argparse.ArgumentParser(description=__doc__)
     parser.add_argument("--output", type=Path, required=True)
     parser.add_argument("--peer", action="store_true")
+    parser.add_argument("--reduced-precision", action="store_true")
     parser.add_argument("--tokens", type=int, nargs="+", default=[4096, 7344, 11424])
     args = parser.parse_args()
     rank = int(os.environ["LOCAL_RANK"])
     torch.accelerator.set_device_index(rank)
     torch.set_grad_enabled(False)
     dist.init_process_group("gloo")
-    torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False
+    torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = (
+        args.reduced_precision
+    )
     group = Group(rank)
     pp = NS(is_first_rank=True, is_last_rank=True)
     cfg = NS(

tests/models/qwen4_exp/test_hc_prefill_peer.py

@@ -10,6 +10,69 @@ from vllm.models.qwen4_exp.nvidia import sm70_hc_prefill_peer as peer
 from vllm.runtime_resources import release_runtime_resources, runtime_resources_for
 
 
+@pytest.mark.parametrize("reduced", [False, True])
+@pytest.mark.parametrize(
+    "rows,accepted",
+    [(4095, False), (4096, True), (4097, True), (16384, True), (16385, False)],
+)
+def test_prefill_admission_preserves_mtp_precision_policy(
+    monkeypatch, reduced, rows, accepted
+):
+    from vllm.config import CUDAGraphMode
+    from vllm.models.qwen4_exp.nvidia import model
+
+    monkeypatch.setattr(
+        torch.backends.cuda.matmul, "allow_fp16_reduced_precision_reduction", reduced
+    )
+    monkeypatch.setattr(torch.backends.cuda.matmul, "allow_fp16_accumulation", False)
+    monkeypatch.setattr(model, "is_forward_context_available", lambda: True)
+    context = NS(cudagraph_runtime_mode=CUDAGraphMode.NONE)
+    monkeypatch.setattr(model, "get_forward_context", lambda: context)
+    owner = NS(_hc_prefill_sequence_qualified=True)
+    inputs = torch.empty(rows, dtype=torch.long)
+    selected = model.Qwen4ExpForCausalLM._use_hc_prefill_sequence_rows(
+        owner, inputs, None, None, {}
+    )
+    assert selected is accepted
+    assert torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction is reduced
+    context.cudagraph_runtime_mode = CUDAGraphMode.FULL
+    assert not model.Qwen4ExpForCausalLM._use_hc_prefill_sequence_rows(
+        owner, inputs, None, None, {}
+    )
+    context.cudagraph_runtime_mode = CUDAGraphMode.NONE
+    monkeypatch.setattr(torch.backends.cuda.matmul, "allow_fp16_accumulation", True)
+    assert not model.Qwen4ExpForCausalLM._use_hc_prefill_sequence_rows(
+        owner, inputs, None, None, {}
+    )
+
+
+def test_single_attention_partial_uses_fp32_transport(monkeypatch):
+    from vllm import distributed
+    from vllm.models.qwen4_exp.nvidia.ops import hc
+
+    seen = []
+
+    def reduce_scatter(value, dim):
+        seen.append(value)
+        assert value.dtype == torch.float32 and dim == 0
+        return value[:1024]
+
+    monkeypatch.setattr(
+        distributed, "get_tp_group", lambda: NS(reduce_scatter=reduce_scatter)
+    )
+    monkeypatch.setattr(peer, "current_runtime_resources", lambda: {})
+    monkeypatch.setattr(
+        hc, "hc_combine_norm", lambda hidden, payload, *args: (payload, hidden)
+    )
+    payload = torch.arange(4096 * 4).reshape(4096, 4).half()
+    hidden = torch.zeros(1024, 16, dtype=torch.half)
+    actual, _ = peer._combine_norm(
+        hidden, payload, torch.zeros(1024, 4), torch.zeros(16), 1e-6, False, True, True
+    )
+    assert len(seen) == 1 and torch.equal(seen[0], payload.float())
+    assert actual.dtype == payload.dtype and torch.equal(actual, payload[:1024])
+
+
 @pytest.mark.parametrize("separate", [False, True])
 def test_peer_fallback_keeps_unrounded_moe_sum(monkeypatch, separate):
     from vllm import distributed

vllm/models/qwen4_exp/nvidia/model.py

@@ -1359,6 +1359,10 @@ class Qwen4ExpForCausalLM(
             input_ids, inputs_embeds, intermediate_tensors, kwargs
         )
         if sequence_parallel:
+            logger.info_once(
+                "Using qualified TP4 HC prefill row owners with FP32 partial "
+                "communication; decode keeps its existing GEMM policy."
+            )
             if self._hc_prefill_sequence_model is None:
                 config = copy.copy(self.vllm_config)
                 config.compilation_config = copy.copy(config.compilation_config)
@@ -1398,7 +1402,6 @@ class Qwen4ExpForCausalLM(
             and intermediate_tensors is None
             and kwargs.get("deepstack_input_embeds") is None
             and (inputs_embeds is None or inputs_embeds.dtype == torch.float16)
-            and not torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction
             and not torch.backends.cuda.matmul.allow_fp16_accumulation
             and is_forward_context_available()
             and get_forward_context().cudagraph_runtime_mode == CUDAGraphMode.NONE

vllm/models/qwen4_exp/nvidia/sm70_hc_prefill_peer.py

@@ -313,7 +313,9 @@ def _combine_norm(
             hidden, first, second, injection, weight, eps, normalize
         )
     if partial:
-        source = first if second is None else first.float() + second.float()
+        # Attention can supply one partial, while MoE supplies two. Both
+        # use the qualified FP32 transport before the original FP16 combine.
+        source = first.float() if second is None else first.float() + second.float()
         reduced = group.reduce_scatter(source, dim=0).to(first.dtype)
     else:
         reduced = first.reshape(4, -1, first.shape[-1])[group.rank_in_group]