diff --git a/benchmarks/flashnext_prefill/DECODE_PARITY_20261015.md b/benchmarks/flashnext_prefill/DECODE_PARITY_20261015.md index 968d92d6f..938e89b63 100644 --- a/benchmarks/flashnext_prefill/DECODE_PARITY_20261015.md +++ b/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 diff --git a/benchmarks/flashnext_prefill/hc_sequence_rows.py b/benchmarks/flashnext_prefill/hc_sequence_rows.py index 4b565f60f..9e6a931b3 100644 --- a/benchmarks/flashnext_prefill/hc_sequence_rows.py +++ b/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( diff --git a/benchmarks/flashnext_prefill/hc_sequence_runtime.py b/benchmarks/flashnext_prefill/hc_sequence_runtime.py index eec067054..1e04365e3 100644 --- a/benchmarks/flashnext_prefill/hc_sequence_runtime.py +++ b/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( diff --git a/tests/models/qwen4_exp/test_hc_prefill_peer.py b/tests/models/qwen4_exp/test_hc_prefill_peer.py index 46b4e0f32..7909c20c9 100644 --- a/tests/models/qwen4_exp/test_hc_prefill_peer.py +++ b/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 diff --git a/vllm/models/qwen4_exp/nvidia/model.py b/vllm/models/qwen4_exp/nvidia/model.py index 6e6007eeb..7eb4a70c0 100644 --- a/vllm/models/qwen4_exp/nvidia/model.py +++ b/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 diff --git a/vllm/models/qwen4_exp/nvidia/sm70_hc_prefill_peer.py b/vllm/models/qwen4_exp/nvidia/sm70_hc_prefill_peer.py index fb763092d..15edf6a33 100644 --- a/vllm/models/qwen4_exp/nvidia/sm70_hc_prefill_peer.py +++ b/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]