diff --git a/.gitignore b/.gitignore index d340682..9bbce33 100644 --- a/.gitignore +++ b/.gitignore @@ -72,4 +72,22 @@ ncu_tmp #nsys *.nsys-rep -*.sqlite \ No newline at end of file +*.sqlite + +# demo / CI run artifacts +outputs/ + +# standalone kernel test binaries (tests/standalone/build.sh) +tests/standalone/test_mfma_simple +tests/standalone/test_mfma_pipeline_hazards + +# hipcc -save-temps intermediates +*-hip-amdgcn-amd-amdhsa-*.bc +*-hip-amdgcn-amd-amdhsa-*.hipi +*-hip-amdgcn-amd-amdhsa-*.s +*-hip-amdgcn-amd-amdhsa-*.out +*-hip-amdgcn-amd-amdhsa-*.out.resolution.txt +*-host-x86_64-unknown-linux-gnu.bc +*-host-x86_64-unknown-linux-gnu.hipi +*-host-x86_64-unknown-linux-gnu.s +*-hip-amdgcn-amd-amdhsa.hipfb \ No newline at end of file diff --git a/README.md b/README.md index aed74af..d863239 100644 --- a/README.md +++ b/README.md @@ -24,7 +24,7 @@ export MODEL_PATH=/path/to/gpt-oss-120b # local GPT-OSS 120B weights (or HF re export GPU=0 # target GPU id rm -rf demo/gpt_oss/permanent_output_dir -USE_FP8_ACT=1 HIP_VISIBLE_DEVICES=$GPU \ +HIP_VISIBLE_DEVICES=$GPU \ python3 demo/gpt_oss/demo.py --use-mirage \ --model-path "$MODEL_PATH" \ --prompt "Tell me the history of america" \ diff --git a/demo/gpt_oss/README.md b/demo/gpt_oss/README.md index 6c05a67..edc1f17 100644 --- a/demo/gpt_oss/README.md +++ b/demo/gpt_oss/README.md @@ -18,7 +18,7 @@ All commands below run from the repo root. ```bash rm -rf demo/gpt_oss/permanent_output_dir -USE_FP8_ACT=1 HIP_VISIBLE_DEVICES=$GPU \ +HIP_VISIBLE_DEVICES=$GPU \ python3 demo/gpt_oss/demo.py --use-mirage \ --model-path "$MODEL_PATH" \ --prompt "Tell me the history of america" \ @@ -39,7 +39,7 @@ The input prompt is set with `--prompt` (wrap it in double quotes): ```bash rm -rf demo/gpt_oss/permanent_output_dir -USE_FP8_ACT=1 HIP_VISIBLE_DEVICES=$GPU \ +HIP_VISIBLE_DEVICES=$GPU \ python3 demo/gpt_oss/demo.py --use-mirage \ --model-path "$MODEL_PATH" \ --prompt "Explain how transformers work in simple terms" \ @@ -52,13 +52,15 @@ Notes: change the token count and make `Decode avg` non-comparable. - `--max-seq-length` is the total sequence length (prompt + generated tokens). Keep it at `512` for benchmarking; shorter values give artificially low TPOT. +- The GEMM arithmetic is not selectable: every MI300 kernel runs MXFP4 weights + against FP8 E4M3 activations (E8M0 per-128-element block scales) on + `v_mfma_scale_f32_16x16x128_f8f6f4`, accumulating in f32. ## Required environment | Variable | Value | Why | |----------|-------|-----| | `MIRAGE_HOME` | repo root | Mirage repo root. | -| `USE_FP8_ACT` | `1` | FP8 activations (the benchmark config). | | `HIP_VISIBLE_DEVICES` | GPU id | Target GPU. | ## Key flags diff --git a/demo/gpt_oss/demo.py b/demo/gpt_oss/demo.py index b76e0f9..3aca5b6 100644 --- a/demo/gpt_oss/demo.py +++ b/demo/gpt_oss/demo.py @@ -34,6 +34,96 @@ PADDED_INTERMEDIATE_SIZE = 2944 +# WikiText-2 raw, test split. Blank lines and the "= Section =" headers are +# dropped, the rest joined with "\n\n", then truncated to --ppl-max-tokens. +# Recording the recipe here (rather than a token count alone) is what makes +# numbers from different runs comparable. +PPL_CORPUS_DESC = "wikitext-2-raw-v1/test, non-header non-blank lines, '\\n\\n'-joined" + + +def load_ppl_corpus(tokenizer, corpus: str, max_tokens: int): + """Return up to `max_tokens` token ids for the perplexity corpus. + + `corpus` is either 'wikitext2' or a path to a UTF-8 text file. The file + fallback exists so the measurement runs on a machine with no network. + """ + if corpus == "wikitext2": + from datasets import load_dataset + ds = load_dataset("wikitext", "wikitext-2-raw-v1", split="test") + lines = [ + t.strip() for t in ds["text"] + if t.strip() and not t.strip().startswith("=") + ] + text = "\n\n".join(lines) + else: + with open(corpus, "r", encoding="utf-8") as f: + text = f.read() + ids = tokenizer(text, return_tensors=None, add_special_tokens=False)["input_ids"] + return ids[:max_tokens] + + +def report_perplexity(mode: str, nll_sum: float, n_scored: int, args, + corpus_tokens: int, per_pos=None, top1=None, + targets=None, tokenizer=None, ent=None): + """Print (and optionally dump) a perplexity result. + + per_pos/top1/targets are optional diagnostics: with PPL_DEBUG=1 they are + printed per position, which is how you tell a uniformly-degraded + distribution apart from a handful of catastrophic rows. + """ + if per_pos is not None and os.environ.get("PPL_DEBUG", "0") == "1": + print(f"\n[PPL_DEBUG {mode}] per-position NLL " + f"(pos, target, nll, top1, top1==target)") + for i, nll in enumerate(per_pos): + t = int(targets[i]) if targets is not None else -1 + p = int(top1[i]) if top1 is not None else -1 + print(f" pos={i + 1:4d} tgt={t:6d} nll={nll:8.4f} " + f"top1={p:6d} {'HIT' if p == t else ''}") + if top1 is not None and targets is not None: + import numpy as _np + hits = sum(1 for i in range(len(per_pos)) + if int(top1[i]) == int(targets[i])) + print(f" top-1 accuracy: {hits}/{len(per_pos)} " + f"({100.0 * hits / len(per_pos):.1f}%)") + mean_nll = nll_sum / n_scored + ppl = math.exp(mean_nll) + print(f"\n{'=' * 60}") + print(f"PERPLEXITY ({mode})") + print(f"{'=' * 60}") + print(f" corpus : {args.ppl_corpus} ({PPL_CORPUS_DESC})") + print(f" corpus tokens : {corpus_tokens}") + print(f" scored positions: {n_scored}") + print(f" mean NLL : {mean_nll:.6f}") + print(f" perplexity : {ppl:.4f}") + if ent: + print(f" mean entropy : {sum(ent) / len(ent):.4f} nats" + f" (sharpness; a noisier GEMM raises this)") + if top1 is not None and targets is not None: + hits = sum(1 for i in range(len(top1)) + if int(top1[i]) == int(targets[i])) + print(f" top-1 accuracy : {hits}/{len(top1)} " + f"({100.0 * hits / len(top1):.2f}%)") + print(f"{'=' * 60}") + if args.ppl_out: + os.makedirs(os.path.dirname(args.ppl_out) or ".", exist_ok=True) + with open(args.ppl_out, "w") as f: + json.dump({ + "mode": mode, + "corpus": args.ppl_corpus, + "corpus_desc": PPL_CORPUS_DESC, + "corpus_tokens": corpus_tokens, + "scored_positions": n_scored, + "mean_nll": mean_nll, + "perplexity": ppl, + "per_position_nll": per_pos, + "top1": top1, + "targets": targets, + "entropy": ent, + }, f, indent=2) + print(f"Saved perplexity to {args.ppl_out}") + return ppl + + def grid_for_rmsnorm_linear_layer(size: int): if size % 64 == 0: return size // 64 @@ -257,20 +347,6 @@ def pack_mxfp4_workgroup(blocks: torch.Tensor, scales: torch.Tensor, return packed.contiguous() -def pack_bf16_workgroup(weight_bf16: torch.Tensor, - output_per_wg: int) -> torch.Tensor: - """Pack BF16 weight [out_dim, K] into workgroup layout [n_wgs, OPW*K*2] bytes. - - Each workgroup: OPW rows x K bf16 values = OPW*K*2 bytes (no scales section). - """ - out_dim, K = weight_bf16.shape - assert out_dim % output_per_wg == 0 - n_wgs = out_dim // output_per_wg - # [n_wgs, OPW, K] bf16 -> view as uint8 [n_wgs, OPW*K*2] - return weight_bf16.reshape(n_wgs, output_per_wg, K).contiguous().view( - torch.uint8).reshape(n_wgs, -1) - - def dequant_mxfp4_to_bf16(blocks: torch.Tensor, scales: torch.Tensor, target_out_dim: int = None, target_reduction: int = None) -> torch.Tensor: @@ -360,6 +436,23 @@ def dequant_mxfp4_to_bf16(blocks: torch.Tensor, scales: torch.Tensor, help="Only use first N layers (for memory-constrained testing)") parser.add_argument("--verify", action="store_true", help="Run both PyTorch and Mirage, compare intermediates") + parser.add_argument( + "--ppl-corpus", default="wikitext2", + help=("Corpus for PPL_MODE=1. Either 'wikitext2' (HuggingFace " + "wikitext/wikitext-2-raw-v1, test split) or a path to a UTF-8 " + "text file."), + ) + parser.add_argument( + "--ppl-max-tokens", default=512, type=int, + help=("Number of corpus tokens to score in PPL_MODE. The logits sink " + "is [max_seq_length+1, padded_vocab] float32 (~400KB/position), " + "and the megakernel runs one iteration per token, so this is " + "both the memory and the runtime knob."), + ) + parser.add_argument( + "--ppl-out", default=None, + help="Dump the PPL_MODE result to this JSON path.", + ) args = parser.parse_args() if args.verify: args.use_mirage = True @@ -418,25 +511,79 @@ def dequant_mxfp4_to_bf16(blocks: torch.Tensor, scales: torch.Tensor, if args.max_layers is not None: num_layers = min(num_layers, args.max_layers) print(f"Using {num_layers} layers (out of {config.num_hidden_layers})") + # Truncate the reference too. num_layers only bounds the MPK task + # graph; GptOssModel.forward iterates self.layers unconditionally, so + # without this the Torch path silently keeps running all 36 layers and + # any MPK-vs-Torch comparison under --max-layers is meaningless. + model.model.layers = model.model.layers[:num_layers] total_num_requests = 1 if not args.use_mirage else args.max_num_batched_requests + + # ── Perplexity mode ─────────────────────────────────────────────────── + # Score a fixed corpus instead of generating. The megakernel already does + # teacher forcing during prefill: prepare_next_batch only copies a sampled + # token into tokens[] once `step + 1 >= prompt_length`, so while we are + # still inside the prompt every position conditions on the *reference* + # prefix. Loading the corpus as one long prompt and running prefill-only + # is therefore exactly the teacher-forced pass perplexity needs -- no + # per-step host round trip and no change to the megakernel loop. + ppl_mode = os.environ.get("PPL_MODE", "0") == "1" + ppl_token_ids = None + if ppl_mode: + if args.use_mirage and args.max_num_batched_tokens != 1: + # The LM head task RMSNorms batch_count rows but feeds only row 0 + # to the GEMM, so a multi-token iteration would emit one logit row + # for a batch of positions. Perplexity needs one row per position. + raise ValueError( + "PPL_MODE requires --max-num-batched-tokens 1; the LM head " + f"emits one logit row per iteration (got " + f"{args.max_num_batched_tokens})." + ) + if args.use_mirage and os.environ.get("FUSE_TAIL", "0") == "1": + # The fused tail never dereferences its lm_logits output pointer, + # so it has no logits sink to attach. + raise ValueError("PPL_MODE is incompatible with FUSE_TAIL=1.") + ppl_token_ids = load_ppl_corpus( + tokenizer, args.ppl_corpus, args.ppl_max_tokens + ) + n_ppl = len(ppl_token_ids) + if n_ppl < 2: + raise ValueError( + f"PPL corpus tokenized to {n_ppl} tokens; need at least 2 " + f"to score a single next-token prediction." + ) + # One extra slot so the last scored position has somewhere to land and + # prepare_next_batch's `step + num_tokens + 1 >= max_seq_length` stop + # fires the moment prefill completes -- prefill-only, no decode. + args.max_seq_length = n_ppl + 1 + print(f"[PPL] corpus={args.ppl_corpus} tokens={n_ppl} " + f"max_seq_length={args.max_seq_length}") + tokens = torch.full((total_num_requests, args.max_seq_length), 0, dtype=torch.long, device="cuda") - # Tokenize prompt (apply chat template if available) - text = args.prompt - if hasattr(tokenizer, 'chat_template') and tokenizer.chat_template: - messages = [{"role": "user", "content": text}] - formatted = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) - model_inputs = tokenizer([formatted], return_tensors="pt", add_special_tokens=False).to("cuda") - print(f"Chat template applied: {len(model_inputs.input_ids[0])} tokens") + if ppl_mode: + ids = torch.tensor(ppl_token_ids, dtype=torch.long, device="cuda") + for r in range(total_num_requests): + tokens[r, :n_ppl] = ids + prompt_lengths = torch.full( + (total_num_requests,), n_ppl, dtype=torch.int, device="cuda" + ) else: - model_inputs = tokenizer([text], return_tensors="pt").to("cuda") - for r in range(total_num_requests): - for i in range(model_inputs.input_ids.shape[-1]): - tokens[r, i] = model_inputs.input_ids[0, i] - prompt_lengths = torch.full( - (total_num_requests,), model_inputs.input_ids.shape[-1], - dtype=torch.int, device="cuda" - ) + # Tokenize prompt (apply chat template if available) + text = args.prompt + if hasattr(tokenizer, 'chat_template') and tokenizer.chat_template: + messages = [{"role": "user", "content": text}] + formatted = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) + model_inputs = tokenizer([formatted], return_tensors="pt", add_special_tokens=False).to("cuda") + print(f"Chat template applied: {len(model_inputs.input_ids[0])} tokens") + else: + model_inputs = tokenizer([text], return_tensors="pt").to("cuda") + for r in range(total_num_requests): + for i in range(model_inputs.input_ids.shape[-1]): + tokens[r, i] = model_inputs.input_ids[0, i] + prompt_lengths = torch.full( + (total_num_requests,), model_inputs.input_ids.shape[-1], + dtype=torch.int, device="cuda" + ) # Position embeddings positions = torch.arange(args.max_seq_length).unsqueeze(0).to("cuda") @@ -984,18 +1131,40 @@ def aiter_profiled_layer_forward(self, hidden_states, position_embeddings=None, # required code paths exist in paged_attention_decode_minimal_hd64_mi300.cuh # (chunk-aware partition) and merge_splitkv.cuh (with optional sinks). # - # Runtime selector based on max-seq-length (measured 2026-05-03 on MI350): - # seq <= 1024 -> chunks=8 (mean 3.404→3.470 ms) - # seq >= 2048 -> chunks=16 (mean 3.529→3.766 ms; chunks=8 is 4-13% slower) - # Override with CK_FMHA_NUM_KV_CHUNKS env var. + # Chunks are claimed by xcd_rank inside the fused full-layer gang task + # (gang_full_layer_fused_mi300.cuh: `if (xcd_rank < NUM_KV_CHUNKS)`), so + # NUM_KV_CHUNKS must never exceed the workers available on one XCD -- + # otherwise the chunk barrier never reaches NUM_KV_CHUNKS-1, the merge + # never fires, and the megakernel deadlocks. + # + # Attention time per chunk is proportional to seqlen/NUM_KV_CHUNKS, so + # scale chunks with sequence length to keep decode latency flat. The + # decode kernel already stamps LSE=-inf for chunks that get no KV tiles, + # so over-provisioning chunks at short seqlen is safe (just wasteful). + # + # KV_TILE=64 in paged_attention_decode_minimal_hd64_mi300.cuh; aim for + # >=2 tiles per chunk so a chunk is worth its merge overhead. + _nw, _ = mi.get_configurations_from_gpu(rank) + MAX_KV_CHUNKS = _nw // 8 # workers per XCD (240/8 = 30 on MI350) _env_chunks = os.environ.get("CK_FMHA_NUM_KV_CHUNKS") if _env_chunks is not None: ck_fmha_num_kv_chunks = int(_env_chunks) else: - ck_fmha_num_kv_chunks = 8 if args.max_seq_length <= 1024 else 16 + _kv_tiles = max(1, (args.max_seq_length + 63) // 64) + ck_fmha_num_kv_chunks = max(8, min(MAX_KV_CHUNKS, _kv_tiles // 2)) assert ck_fmha_num_kv_chunks >= 1 use_split_attn_chunks = (ck_fmha_num_kv_chunks > 1) fuse_full_layer = os.environ.get("FUSE_FULL_LAYER", "1") == "1" + if fuse_full_layer and ck_fmha_num_kv_chunks > MAX_KV_CHUNKS: + raise ValueError( + f"CK_FMHA_NUM_KV_CHUNKS={ck_fmha_num_kv_chunks} exceeds the " + f"{MAX_KV_CHUNKS} workers per XCD available to claim chunks in " + f"the fused full-layer gang task; the split-KV merge would " + f"never fire and the kernel would hang." + ) + print(f"[CFG] max_seq_length={args.max_seq_length} " + f"ck_fmha_num_kv_chunks={ck_fmha_num_kv_chunks} " + f"(max {MAX_KV_CHUNKS})") fuse_tail = os.environ.get("FUSE_TAIL", "0") == "1" if args.profiling: @@ -1161,10 +1330,17 @@ def make_tensor(name, dims, torch_dtype=torch.bfloat16): fuse_qkv_attn = True if fuse_qkv_attn: qkv_attn_barrier = make_tensor("qkv_attn_barrier", (16,), torch_dtype=torch.int32) - # Hierarchical barrier for fused W13+W2 kernel [16*E int32]: - # Per expert (1 cache line = 64 bytes): - # [0..7]: xcd_arrive[8], [8]: global_arrive, [9]: w2_done - moe_fused_barrier = make_tensor("moe_fused_barrier", (16 * num_experts,), torch_dtype=torch.int32) + # Hierarchical barrier for fused W13+W2 kernel [160*E int32]. + # Per expert, 10 slots of one 64-byte cache line (16 int32) each: + # [x*16] for x in 0..7: per-XCD release flag (st_wt, bypasses L2) + # [8*16]: global_arrive (atomic, lives in L2) + # [9*16]: reserved + # The one-line-per-slot spacing is required, not padding: an L2-resident + # atomic sharing a line with write-through release flags can write the + # stale L2 copy back over them, reverting releases that already + # happened and deadlocking the W2 workers for that expert. See the + # layout note in gang_moe_fused_mxfp4_mi300.cuh (MOE_BAR_*). + moe_fused_barrier = make_tensor("moe_fused_barrier", (160 * num_experts,), torch_dtype=torch.int32) # W13+SwiGLU fused output: [bs, top_k, padded_intermediate] # (SwiGLU is fused into W13 epilogue — no separate mlp_mid buffer) swiglu_out = make_tensor("swiglu_out", (bs, num_experts_per_tok, PADDED_INTERMEDIATE_SIZE)) @@ -1181,6 +1357,22 @@ def make_tensor(name, dims, torch_dtype=torch.bfloat16): argmax_out = mpk.attach_input(torch_tensor=output_tokens, name="output_token") if fuse_tail: argmax_in = make_tensor("argmax_in", (bs, vocab_size)) + # Perplexity sink: full logit row per scored position. Row r holds the + # distribution over tokens[r], written by the iteration that consumed + # tokens[r-1] (task_register passes runtime_config.step[0] + 1), so + # row 0 is never written and rows 1..n_ppl are. float32, not bf16 -- + # bf16's ~0.4% relative precision is the same order as the GEMM error + # this buffer exists to measure. + ppl_logits = None + if ppl_mode: + ppl_bytes = args.max_seq_length * vocab_size * 4 + print(f"[PPL] logits sink: [{args.max_seq_length}, {vocab_size}] " + f"f32 = {ppl_bytes / 1e9:.2f} GB") + ppl_logits = make_tensor( + "ppl_logits", (args.max_seq_length, vocab_size), + torch_dtype=torch.float32, + ) + ppl_logits_torch = _tensor_refs["ppl_logits"] # Split-K workspace for linear_with_residual on MI300 # Must include done counter space: per XCD we need n_tiles_per_xcd ints @@ -1331,19 +1523,12 @@ def _attach_input_keep(torch_tensor, name): qkv_out_size = w_qkv_shuffled.shape[0] # fused_qkv_dim # Quantize/pack weights for workgroup layout qkv_output_per_wg = 64 # 10 tiles/XCD fits in 30 workers, kvupd fusion needs OPW==head_dim - use_bf16_native = os.environ.get("USE_FP16_ACT") == "1" - if use_bf16_native: - # BF16 native weights: store pre-dequanted BF16 in workgroup layout - w_qkv_packed = pack_bf16_workgroup(w_qkv_shuffled, output_per_wg=qkv_output_per_wg) - w_qkv_mxfp4 = _attach_input_keep( - w_qkv_packed, f"layer_{i}_qkv_bf16") - else: - qkv_blocks, qkv_scales = quantize_bf16_to_mxfp4(w_qkv_shuffled) - w_qkv_packed = pack_mxfp4_workgroup( - qkv_blocks, qkv_scales, output_per_wg=qkv_output_per_wg, - ).squeeze(0) # [n_wgs, wg_bytes] - w_qkv_mxfp4 = _attach_input_keep( - w_qkv_packed, f"layer_{i}_qkv_mxfp4") + qkv_blocks, qkv_scales = quantize_bf16_to_mxfp4(w_qkv_shuffled) + w_qkv_packed = pack_mxfp4_workgroup( + qkv_blocks, qkv_scales, output_per_wg=qkv_output_per_wg, + ).squeeze(0) # [n_wgs, wg_bytes] + w_qkv_mxfp4 = _attach_input_keep( + w_qkv_packed, f"layer_{i}_qkv_mxfp4") # QKV bias: shuffle Q/K/V biases to match interleaved weight layout q_bias = layer.self_attn.q_proj.bias.data.to("cuda") k_bias = layer.self_attn.k_proj.bias.data.to("cuda") @@ -1752,17 +1937,12 @@ def _attach_input_keep(torch_tensor, name): if is_rocm: # Quantize/pack O-proj weight o_output_per_wg = 16 - if use_bf16_native: - w_o_packed = pack_bf16_workgroup(w_o, output_per_wg=o_output_per_wg) - w_o_mxfp4 = _attach_input_keep( - w_o_packed, f"layer_{i}_o_proj_bf16") - else: - o_blocks, o_scales = quantize_bf16_to_mxfp4(w_o) - w_o_packed = pack_mxfp4_workgroup( - o_blocks, o_scales, output_per_wg=o_output_per_wg, - ).squeeze(0) # [n_wgs, wg_bytes] - w_o_mxfp4 = _attach_input_keep( - w_o_packed, f"layer_{i}_o_proj_mxfp4") + o_blocks, o_scales = quantize_bf16_to_mxfp4(w_o) + w_o_packed = pack_mxfp4_workgroup( + o_blocks, o_scales, output_per_wg=o_output_per_wg, + ).squeeze(0) # [n_wgs, wg_bytes] + w_o_mxfp4 = _attach_input_keep( + w_o_packed, f"layer_{i}_o_proj_mxfp4") # === MoE block weight prep (needed before fused path) === post_norm_w_padded = pad_weight_1d( @@ -2024,6 +2204,7 @@ def _attach_input_keep(torch_tensor, name): output_per_wg=lm_head_output_per_wg, output_stride=vocab_size, block_dim=(256, 1, 1), + ppl_logits=ppl_logits, ) mpk.argmax_reduce_layer( input=(argmax_part_value, argmax_part_index), @@ -2055,8 +2236,90 @@ def _attach_input_keep(torch_tensor, name): tokens.size(1) - prompt_lengths[0].item() ) output_len = max(0, min(output_len, tokens.size(1) - prompt_lengths[0].item())) + if ppl_mode: + # Prefill-only: every scored position must condition on the reference + # prefix, and a single generated token would start feeding the model + # its own output. + output_len = 0 + + if ppl_mode and not args.use_mirage: + # Torch reference perplexity on the same slice. One causal forward + # over the whole sequence is teacher forcing by construction. + def _mxfp4_roundtrip(w): + """Push a bf16 weight through the same quantizer MPK uses.""" + b, s = quantize_bf16_to_mxfp4(w.data) + w.data = dequant_mxfp4_to_bf16(b, s)[0].to(w.dtype).reshape( + w.data.shape + ) - if not args.use_mirage: + # MPK quantizes weights the checkpoint stores in bf16 -- the LM head, + # QKV and O-proj -- down to MXFP4, while this reference keeps them in + # bf16. (The MoE experts are natively MXFP4 in both paths, so they are + # not part of the difference.) Comparing the two as-is charges that + # quantization loss to the megakernel. Round-tripping the reference's + # weights through the same quantizer separates "the kernel computes + # something different" from "the kernel was handed coarser weights". + # + # PPL_MXFP4_HEAD=1 head only (the original, narrower control) + # PPL_MXFP4_MATCH=1 head + QKV + O-proj: the matched-precision run + _match = os.environ.get("PPL_MXFP4_MATCH", "0") == "1" + if _match or os.environ.get("PPL_MXFP4_HEAD", "0") == "1": + _mxfp4_roundtrip(model.lm_head.weight) + print("[PPL] Torch LM head round-tripped through MXFP4") + if _match: + n_rt = 0 + for _lyr in model.model.layers: + for _w in (_lyr.self_attn.q_proj, _lyr.self_attn.k_proj, + _lyr.self_attn.v_proj, _lyr.self_attn.o_proj): + _mxfp4_roundtrip(_w.weight) + n_rt += 1 + print(f"[PPL] Torch QKV/O-proj round-tripped through MXFP4 " + f"({n_rt} weights)") + ids = tokens[:1, :n_ppl] + cos_e = position_embeddings[0][:, :n_ppl] + sin_e = position_embeddings[1][:, :n_ppl] + hidden, _ = model.model( + input_ids=ids, position_embeddings=(cos_e, sin_e), step=step, + ) + targets = tokens[0, 1:n_ppl] + # Chunk the LM head: [n, 201088] float32 logits at once is avoidable + # memory pressure and the sum is exact either way. + nll_sum = 0.0 + per_pos, top1, ent = [], [], [] + CH = 64 + for lo in range(0, n_ppl - 1, CH): + hi = min(lo + CH, n_ppl - 1) + chunk_logits = model.lm_head(hidden[0, lo:hi, :]).float() + losses = torch.nn.functional.cross_entropy( + chunk_logits, targets[lo:hi], reduction="none" + ) + nll_sum += losses.sum().item() + per_pos.extend(losses.tolist()) + top1.extend(chunk_logits.argmax(dim=-1).tolist()) + lp = torch.log_softmax(chunk_logits, dim=-1) + ent.extend((-(lp.exp() * lp).sum(dim=-1)).tolist()) + # Raw logit rows for a direct MPK-vs-Torch comparison. Derived metrics + # (NLL, entropy) can only say the distributions differ; the raw vectors + # say *how* -- a scale error, an offset, or unstructured noise are three + # different bugs and they look identical after a softmax. + if os.environ.get("PPL_DUMP_LOGITS"): + rows = [int(x) for x in + os.environ.get("PPL_DUMP_ROWS", "1,2,5,10,50,100").split(",") + if int(x) < n_ppl - 1] + torch.save( + {"rows": rows, + "logits": {r: model.lm_head(hidden[0, r, :]).float().cpu() + for r in rows}, + "hidden": {r: hidden[0, r, :].float().cpu() for r in rows}}, + os.environ["PPL_DUMP_LOGITS"]) + print(f"[PPL] dumped rows {rows} to " + f"{os.environ['PPL_DUMP_LOGITS']}") + report_perplexity( + "torch", nll_sum, n_ppl - 1, args, corpus_tokens=n_ppl, + per_pos=per_pos, top1=top1, targets=targets.tolist(), + tokenizer=tokenizer, ent=ent, + ) + elif not args.use_mirage: prompt_len = prompt_lengths[0].item() decode_limit = prompt_len + output_len for cur_pos in range(prompt_len, decode_limit): @@ -2185,6 +2448,169 @@ def _attach_input_keep(torch_tensor, name): ): _fwd_times[int(_m.group(1))] = float(_m.group(2)) + # The device-side per-iter ring holds FWDPASS_LOG_MAX (8192) samples. + # Longer runs drop the tail, and since per-iter latency grows with + # sequence length, averaging only what survived understates the real + # number. [FWD_PASS_TOTAL] is accumulated over every iteration, so + # prefer it whenever samples were dropped. + _fwd_dropped = 0 + _fwd_total_avg = None + _fwd_total_iters = 0 + _m_tot = re.search( + r"\[FWD_PASS_TOTAL\] iters=(\d+) total_ms=[\d.]+ " + r"avg_ms=([\d.]+) dropped=(\d+)", + _captured, + ) + if _m_tot: + _fwd_total_iters = int(_m_tot.group(1)) + _fwd_total_avg = float(_m_tot.group(2)) + _fwd_dropped = int(_m_tot.group(3)) + + if ppl_mode: + # ppl_logits[r] is the distribution over tokens[0, r], written by + # the iteration that consumed tokens[0, r-1]. Row 0 is never + # written, so scored positions are 1..n_ppl-1. + # + # Slice to config.vocab_size: the buffer is padded to 201216 and + # the pad columns were filled by rows of the zero-padded LM head + # weight. They are not real vocabulary and must not enter the + # softmax denominator. + real_vocab = config.vocab_size + targets = tokens[0, 1:n_ppl] + nll_sum = 0.0 + per_pos, top1, ent = [], [], [] + CH = 64 + for lo in range(1, n_ppl, CH): + hi = min(lo + CH, n_ppl) + chunk = ppl_logits_torch[lo:hi, :real_vocab].float() + losses = torch.nn.functional.cross_entropy( + chunk, targets[lo - 1:hi - 1], reduction="none" + ) + nll_sum += losses.sum().item() + per_pos.extend(losses.tolist()) + top1.extend(chunk.argmax(dim=-1).tolist()) + # Distribution sharpness. Numeric noise in the GEMM flattens + # the softmax, which *lowers* NLL at positions the model gets + # wrong -- so entropy has to be reported alongside perplexity + # or a noisier kernel can look like a better one. + lp = torch.log_softmax(chunk, dim=-1) + ent.extend((-(lp.exp() * lp).sum(dim=-1)).tolist()) + # A row the kernel never touched is all zeros -- uniform over the + # vocabulary, ln(201088) = 12.21 nats. Catching that here beats + # reporting a plausible-looking but meaningless number. + # + # Both of these run chunked. A whole-tensor `== 0.0` on the sink + # allocates an [n, vocab] bool and `.float()` an [n, vocab] f32 -- + # at 32k that is 6 GB and 25 GB on top of the 25 GB sink, i.e. an + # OOM in the diagnostic rather than in the thing being measured. + n_zero = 0 + zero_total = 0 + first_zero_row = -1 + first_zero_cols = None + pad_max = 0.0 + for lo in range(1, n_ppl, CH): + hi = min(lo + CH, n_ppl) + blk = ppl_logits_torch[lo:hi, :real_vocab] + zc = (blk == 0.0) + per_row = zc.sum(dim=1) + zero_total += int(per_row.sum().item()) + n_zero += int((per_row == real_vocab).sum().item()) + if first_zero_row < 0 and bool((per_row > 0).any().item()): + i0 = int((per_row > 0).nonzero()[0].item()) + first_zero_row = lo + i0 + first_zero_cols = zc[i0].nonzero().flatten()[:16].tolist() + if vocab_size > real_vocab: + pad_max = max(pad_max, float( + ppl_logits_torch[lo:hi, real_vocab:].abs().max().item() + )) + if n_zero: + print(f"[PPL] WARNING: {n_zero}/{n_ppl - 1} scored rows are " + f"all-zero -- the logits sink was not written for them.") + # Per-column coverage. An exactly-0.0 logit is possible but + # vanishingly unlikely in float32, so a nonzero count here means + # columns the kernel never wrote -- which reads as logit 0 and + # produces a ~17-nat NLL whenever the target lands on one. + print(f"[PPL] zero columns: total={zero_total} " + f"per-row mean={zero_total / max(1, n_ppl - 1):.1f} " + f"of {real_vocab}") + if first_zero_row >= 0: + print(f"[PPL] first affected row {first_zero_row}: " + f"first 16 zero cols = {first_zero_cols}") + print(f"[PPL] pad-column max |logit| (excluded): {pad_max:.4f}") + + # Self-consistency: the last prefill iteration consumed + # tokens[n_ppl-1], wrote sink row n_ppl, AND -- because + # step + 1 == prompt_length there -- had its argmax copied into + # tokens[0, n_ppl] by prepare_next_batch. If the sink is a + # faithful copy of the values the in-register argmax reduced, + # those two must name the same token. This checks the sink + # against the kernel's own reduction rather than against Torch, + # so it isolates "is the sink right" from "is MXFP4 accurate". + if n_ppl < args.max_seq_length: + sink_top = int( + ppl_logits_torch[n_ppl, :real_vocab].argmax().item() + ) + kernel_top = int(tokens[0, n_ppl].item()) + ok = "OK" if sink_top == kernel_top else "MISMATCH" + print(f"[PPL] sink/argmax self-check: sink_argmax={sink_top} " + f"kernel_token={kernel_top} -> {ok}") + + # Stronger: the same last iteration also left 240 per-worker + # (max, abs_idx) pairs in argmax_part_*. Each worker owns a + # known set of 64-column tiles, so recomputing its max from + # the sink and comparing checks every column of the row, not + # just the single winner above. + pv = _tensor_refs["argmax_part_value"][0].float() + pi = _tensor_refs["argmax_part_index"][0] + wpx = mpk.num_workers // 8 # workers per XCD + nwg = (vocab_size // lm_head_output_per_wg) // 8 + sink_row = ppl_logits_torch[n_ppl] + bad_idx = bad_val = 0 + for p in range(8): + pstart = p * nwg * lm_head_output_per_wg + for r in range(wpx): + cols = torch.cat([ + torch.arange( + pstart + wg * lm_head_output_per_wg, + pstart + (wg + 1) * lm_head_output_per_wg, + device="cuda") + for wg in range(r, nwg, wpx) + ]) + vals = sink_row[cols] + k = int(vals.argmax().item()) + w = p * wpx + r + if int(cols[k].item()) != int(pi[w].item()): + bad_idx += 1 + # argmax_part_value is bf16: 8 mantissa bits, so + # compare at bf16 resolution, not exactly. + elif abs(float(vals[k]) - float(pv[w])) > \ + 0.02 * max(1.0, abs(float(pv[w]))): + bad_val += 1 + print(f"[PPL] sink/per-worker-argmax check over all " + f"{mpk.num_workers} workers: " + f"{bad_idx} index mismatches, {bad_val} value " + f"mismatches -> " + f"{'OK' if bad_idx == 0 and bad_val == 0 else 'MISMATCH'}") + if os.environ.get("PPL_DUMP_LOGITS"): + rows = [int(x) for x in + os.environ.get("PPL_DUMP_ROWS", + "1,2,5,10,50,100").split(",") + if int(x) < n_ppl - 1] + # Sink row r+1 holds the distribution over tokens[r+1], i.e. + # the same position the Torch dump indexes as row r. + torch.save( + {"rows": rows, + "logits": {r: ppl_logits_torch[r + 1, :real_vocab] + .float().cpu() for r in rows}}, + os.environ["PPL_DUMP_LOGITS"]) + print(f"[PPL] dumped rows {rows} to " + f"{os.environ['PPL_DUMP_LOGITS']}") + report_perplexity( + "mpk", nll_sum, n_ppl - 1, args, corpus_tokens=n_ppl, + per_pos=per_pos, top1=top1, targets=targets.tolist(), + tokenizer=tokenizer, ent=ent, + ) + #print("tokens.shape = ", tokens.shape, flush=True) #print("All tokens:", tokens[0].tolist()) #print("Step:", step.tolist()) @@ -2285,6 +2711,15 @@ def _attach_input_keep(torch_tensor, name): and _it - 1 <= total_iterations] print(f" Decode per-iter range: min={min(_decode_samples):.3f}ms " f"max={max(_decode_samples):.3f}ms") + if _fwd_dropped > 0: + print("-" * 80) + print(f" NOTE: device per-iter ring overflowed — {_fwd_dropped} of " + f"{_fwd_total_iters} samples dropped. The prefill/decode " + f"splits above cover only the first " + f"{_fwd_total_iters - _fwd_dropped} iterations and understate " + f"latency (per-iter grows with seq len).") + print(f" All-iteration device average: {_fwd_total_avg:.3f}ms/iter " + f"over {_fwd_total_iters} iters") print("=" * 80) # === Verification: compare Mirage intermediates with PyTorch reference === diff --git a/include/mirage/persistent_kernel/persistent_kernel.cuh b/include/mirage/persistent_kernel/persistent_kernel.cuh index ac4c5e0..97b39ec 100644 --- a/include/mirage/persistent_kernel/persistent_kernel.cuh +++ b/include/mirage/persistent_kernel/persistent_kernel.cuh @@ -25,7 +25,11 @@ #include #include #endif +#include +#include #include +#include +#include #include #include #include @@ -73,6 +77,226 @@ __device__ int g_oproj_inner_iters; // total iterations accumulated __device__ int g_oproj_inner_reset_flag; // CAS flag for per-iter reset #endif +#ifdef MPK_NIL_TRIPWIRE +// Sub-phase breadcrumbs from inside the gang task kernels. +// +// The outer loop's breadcrumb only resolves to "somewhere inside +// _execute_gang_task", which is the whole layer. These let a task kernel name +// the phase it was in. Declared before task_header.cuh so the task kernels can +// see them; the pointer is published by the worker loop (every block stores +// the same value, so the race is benign). +// +// A worker block's slot index is blockIdx.x: worker_id defaults to blockIdx.x +// for every block that runs the worker loop. +__device__ unsigned long long *g_tw_dev; + +__device__ __forceinline__ void + mpk_tw_sub(int sub, unsigned long long aux, int tid) { + if (tid == 0 && g_tw_dev != nullptr) { + unsigned long long *b = + g_tw_dev + MPK_TW_HDR + blockIdx.x * MPK_TW_PER_WORKER; + b[4] = (unsigned long long)sub; + b[5] = aux; + } +} +#define MPK_TW_SUB(sub, aux) mpk_tw_sub((sub), (unsigned long long)(aux), tid) +#else +#define MPK_TW_SUB(sub, aux) ((void)0) +#endif + +// Intra-layer phase breadcrumb for the worker-state dump. +// +// The 40000+ml marker says which layer a stalled worker is in, but a layer +// spans eight phases and several cross-XCD barriers, so it cannot say which +// barrier is holding. MPK_TW_SUB resolves that, but only under +// MPK_NIL_TRIPWIRE, which costs enough to perturb the timing of the very +// race being chased. This writes the same sub-phase code into the existing +// worker-state slot, so MPK_WORKER_STATE alone gets intra-layer resolution: +// a single relaxed store by tid 0, no atomics and no extra buffers. +// +// Encoded as 50000000 + xcd*100000 + phase*1000 + layer%1000, so one int +// carries all three. The XCD matters because every barrier here is either +// per-XCD (the QKV epoch) or cross-XCD (attn_global, MoE W13->W2), and +// telling those apart is exactly what identifies the blocking worker. +__device__ int *g_ws_dev; + +__device__ __forceinline__ void + mpk_ws_phase(int phase, int layer, int xcd, int tid) { + if (tid == 0 && g_ws_dev != nullptr) { + __atomic_store_n(&g_ws_dev[blockIdx.x * 4 + 3], + 50000000 + xcd * 100000 + phase * 1000 + (layer % 1000), + __ATOMIC_RELAXED); + } +} +#ifdef MPK_WORKER_STATE +#define MPK_WS_PHASE(phase, layer, xcd) \ + mpk_ws_phase((phase), (layer), (xcd), tid) +#else +#define MPK_WS_PHASE(phase, layer, xcd) ((void)0) +#endif + +// Barrier watch: what a spinning worker is actually waiting *for*. +// +// MPK_WS_PHASE names the phase, which narrows a stall to a barrier, but not +// to a cause. Every barrier here is a "poll until counter >= expected" loop, +// and the two ways it hangs are indistinguishable from the phase alone: +// either the producer never arrived (observed < expected, and the shortfall +// says how many arrivals are missing), or the waiter computed an expected +// value the producer will never reach (observed >= expected but the load is +// reading a stale cache line, or expected ran ahead by a full epoch). This +// records observed/expected so the dump can tell those apart. +// +// Lives in the second half of the worker-state buffer -- 4 more ints per +// worker at [num_workers*4 + blockIdx.x*4] -- to avoid widening RuntimeConfig +// for a debug-only path. Slots: 0 = barrier id, 1 = last observed value, +// 2 = expected value, 3 = spin iterations. +// +// Cost in a healthy run is two relaxed stores per barrier entry: the observed +// value is refreshed only once every MPK_WS_WAIT_REFRESH spins, and these +// polls clear in far fewer than that, so the loop body itself stays untouched. +__device__ int g_ws_nworkers; +#define MPK_WS_WAIT_REFRESH 4096 + +__device__ __forceinline__ void + mpk_ws_wait_begin(int barrier_id, int expected, int tid) { + if (tid == 0 && g_ws_dev != nullptr) { + int *b = g_ws_dev + g_ws_nworkers * 4 + blockIdx.x * 4; + __atomic_store_n(&b[0], barrier_id, __ATOMIC_RELAXED); + __atomic_store_n(&b[2], expected, __ATOMIC_RELAXED); + __atomic_store_n(&b[3], 0, __ATOMIC_RELAXED); + } +} + +__device__ __forceinline__ void + mpk_ws_wait_tick(int observed, int spins, int tid) { + if (tid == 0 && g_ws_dev != nullptr) { + int *b = g_ws_dev + g_ws_nworkers * 4 + blockIdx.x * 4; + __atomic_store_n(&b[1], observed, __ATOMIC_RELAXED); + __atomic_store_n(&b[3], spins, __ATOMIC_RELAXED); + } +} + +// Third quarter of the buffer: four barrier-specific auxiliary values. +// +// observed/expected says a release is missing but not why. For the MoE +// W13->W2 barrier the discriminating value is the raw arrival counter: the +// release fires on (prev % W13_TILES) == W13_TILES-1, on a counter that is +// never reset and is shared by every layer that activates the expert. If the +// counter is sitting on a multiple of W13_TILES the arrivals all landed and +// the release value is wrong; if it is not, arrivals were lost in some earlier +// layer and the modular boundary has been permanently skewed past. +__device__ __forceinline__ void + mpk_ws_wait_aux(int a0, int a1, int a2, int a3, int tid) { + if (tid == 0 && g_ws_dev != nullptr) { + int *b = g_ws_dev + g_ws_nworkers * 8 + blockIdx.x * 4; + __atomic_store_n(&b[0], a0, __ATOMIC_RELAXED); + __atomic_store_n(&b[1], a1, __ATOMIC_RELAXED); + // b[2] and b[3] belong to the per-wave masks below -- not written here. + (void)a2; + (void)a3; + } +} +#define MPK_WS_WAIT_AUX(a0, a1, a2, a3) \ + mpk_ws_wait_aux((a0), (a1), (a2), (a3), tid) + +// Per-wave exit mark for a *thread-divergent* poll. +// +// Every other tracer here is tid==0 only, which is enough while a barrier is +// entered and left by the whole block together. The MoE W13->W2 poll is not: +// it deliberately has each thread test the release flag on its own with no +// __syncthreads, so tid 0 can clear the poll and run on while other waves of +// the same block are still spinning. From tid 0's slot that is invisible -- +// it reports a straight-line mark past the barrier while the block as a whole +// is still stuck at it, which is exactly how the gx_1 capture looked. +// +// One relaxed store per wave, into the aux quarter's low bits: wave w sets +// bit w when it leaves the poll. A block whose waves all exited reads 0xf. +// Anything less names the waves that never got the release. +__device__ __forceinline__ void mpk_ws_wave_exit(int wave, int tid) { + if ((tid & 63) == 0 && g_ws_dev != nullptr) { + int *b = g_ws_dev + g_ws_nworkers * 8 + blockIdx.x * 4; + __atomic_fetch_or(&b[3], 1 << (wave & 31), __ATOMIC_RELAXED); + } +} +// Each wave clears its OWN bit in both masks. Doing it this way instead of +// "tid 0 zeroes the word, then __syncthreads" matters: adding a __syncthreads +// to a region whose whole point is that it is thread-divergent changes the +// control flow being measured. The ix_1 capture is not trustworthy for that +// reason. Per-wave clears race with nothing, and need no barrier. +__device__ __forceinline__ void mpk_ws_wave_clear(int wave, int tid) { + if ((tid & 63) == 0 && g_ws_dev != nullptr) { + int *b = g_ws_dev + g_ws_nworkers * 8 + blockIdx.x * 4; + __atomic_fetch_and(&b[2], ~(1 << (wave & 31)), __ATOMIC_RELAXED); + __atomic_fetch_and(&b[3], ~(1 << (wave & 31)), __ATOMIC_RELAXED); + } +} +#define MPK_WS_WAVE_EXIT(wave) mpk_ws_wave_exit((wave), tid) +#define MPK_WS_WAVE_CLEAR(wave) mpk_ws_wave_clear((wave), tid) + +// Same per-wave mask, one slot over, for arrival at a __syncthreads. +// +// hx_2 showed the W13->W2 poll mask at 0xf -- every wave cleared the barrier +// -- with the block still parked at mark 8303, i.e. inside the FP8 quant. +// The only thing in that function that can block is its trailing +// __syncthreads, so the question becomes which wave fails to reach it. +__device__ __forceinline__ void mpk_ws_wave_sync(int wave, int tid) { + if ((tid & 63) == 0 && g_ws_dev != nullptr) { + int *b = g_ws_dev + g_ws_nworkers * 8 + blockIdx.x * 4; + __atomic_fetch_or(&b[2], 1 << (wave & 31), __ATOMIC_RELAXED); + } +} +#define MPK_WS_WAVE_SYNC(wave) mpk_ws_wave_sync((wave), tid) + +// Straight-line progress mark, for code that is not a spin loop. +// +// The barrier watch only sees workers parked in a poll. A worker stuck +// *between* polls -- inside MoE compute, say -- registers nowhere, which is +// exactly the state the first captured deadlock left its two blockers in. +// This writes the same slots with spins = -1 so the host can tell a mark from +// a barrier and print the last code the worker got past. +// One store, not four. These fire per MoE tile -- several times per worker +// per layer -- and g_ws_dev is pinned *host* memory, so every store is a PCIe +// write. The four-store version cost 2.3x (6.99 vs 3.01 ms/token), which is +// far past the point where the instrumentation perturbs the race it is meant +// to catch. Packed as code*100000 + aux%100000, written to the spins slot as +// a negative so the host can distinguish it from a poll count. +__device__ __forceinline__ void mpk_ws_mark(int code, int aux, int tid) { + if (tid == 0 && g_ws_dev != nullptr) { + int *b = g_ws_dev + g_ws_nworkers * 4 + blockIdx.x * 4; + __atomic_store_n( + &b[3], -(code * 100000 + (aux % 100000)), __ATOMIC_RELAXED); + } +} +// Compile-time gate for the inline worker-state dump sites below. The runtime +// null check alone still costs a global load and a branch at ~30 sites, several +// of them in the per-layer dispatch loop; measured 2.386 -> 2.321 ms/iter when +// compiled out. Define MPK_WORKER_STATE to get the dump back. +#ifdef MPK_WORKER_STATE +#define MPK_WS_ON(cfg) ((cfg).precomp_dbg_worker_state != nullptr) +#else +#define MPK_WS_ON(cfg) (false) +#endif + +#ifdef MPK_WORKER_STATE +#define MPK_WS_MARK(code, aux) mpk_ws_mark((code), (aux), tid) +#else +#define MPK_WS_MARK(code, aux) ((void)0) +#endif + +#ifdef MPK_WORKER_STATE +#define MPK_WS_WAIT_BEGIN(barrier_id, expected) \ + mpk_ws_wait_begin((barrier_id), (expected), tid) +#define MPK_WS_WAIT_TICK(observed, spins) \ + do { \ + if (((spins) & (MPK_WS_WAIT_REFRESH - 1)) == 0) { \ + mpk_ws_wait_tick((observed), (spins), tid); \ + } \ + } while (0) +#else +#define MPK_WS_WAIT_BEGIN(barrier_id, expected) ((void)0) +#define MPK_WS_WAIT_TICK(observed, spins) ((void)0) +#endif + #if defined(MIRAGE_GRACE_HOPPER) #include "tasks/hopper/task_header.cuh" #elif defined(MIRAGE_GRACE_BLACKWELL) @@ -86,7 +310,7 @@ __device__ int g_oproj_inner_reset_flag; // CAS flag for per-iter reset // HIP/ROCm compatibility macros and type aliases #if defined(__HIP_PLATFORM_AMD__) || defined(MIRAGE_AMD_MI300) #include - // CUDA runtime API -> HIP equivalents +// CUDA runtime API -> HIP equivalents #define cudaMalloc hipMalloc #define cudaFree hipFree #define cudaStreamSynchronize hipStreamSynchronize @@ -102,9 +326,9 @@ __device__ int g_oproj_inner_reset_flag; // CAS flag for per-iter reset hipFuncAttributeMaxDynamicSharedMemorySize #define cudaStreamNonBlocking hipStreamNonBlocking #define cudaEventDisableTiming hipEventDisableTiming - // HIP's hipFuncSetAttribute requires - // const void*; wrap to cast function - // pointers +// HIP's hipFuncSetAttribute requires +// const void*; wrap to cast function +// pointers #define cudaFuncSetAttribute(func, attr, value) \ hipFuncSetAttribute(reinterpret_cast(func), attr, value) #define cudaStreamCreateWithFlags hipStreamCreateWithFlags @@ -843,12 +1067,29 @@ __device__ __forceinline__ int _span_stage_for_task(int task_type, } #endif -// Deferred FWD_PASS log: store per-iteration timing, print at termination +// Deferred FWD_PASS log: store per-iteration timing, print at termination. +// +// The ring is bounded, so runs longer than FWDPASS_LOG_MAX iterations cannot +// keep every sample. They must not silently keep only the *first* window +// either: per-iter latency grows with sequence length, so truncating to the +// first 8k iterations biased every reported average downward at long seq len +// (a 49k run reported the latency of its cheapest 8k iterations). +// +// The ring therefore stride-decimates: when it fills it compacts in place, +// keeping every 2nd sample and doubling the stride. Coverage stays uniform +// across the entire run at any length, which is what makes latency-vs-seqlen +// reconstructible. g_fwdpass_stride is the current decimation factor; the +// totals accumulate over every iteration regardless. #define FWDPASS_LOG_MAX 8192 __device__ int g_fwdpass_count; +__device__ int g_fwdpass_stride = 1; __device__ unsigned long long g_fwdpass_time_ns[FWDPASS_LOG_MAX]; // iter duration __device__ int g_fwdpass_tokens[FWDPASS_LOG_MAX]; // num_active_tokens +// Untruncated aggregates: cover every iteration regardless of ring capacity. +__device__ int g_fwdpass_dropped; +__device__ unsigned long long g_fwdpass_total_ns; +__device__ int g_fwdpass_total_iters; __device__ __forceinline__ void execute_worker(RuntimeConfig config, int assigned_worker_id = -1) { @@ -893,6 +1134,29 @@ __device__ __forceinline__ void execute_worker(RuntimeConfig config, #endif int const worker_id = (assigned_worker_id >= 0) ? assigned_worker_id : blockIdx.x; +#ifdef MPK_NIL_TRIPWIRE + // Hand the tripwire buffer to the task kernels, which take input/output + // pointer arrays rather than the RuntimeConfig and so cannot reach it. + // Every block stores the same value, so the write race is benign. + // MPK_TW_SUB indexes by blockIdx.x, which equals worker_id here (nothing + // passes assigned_worker_id), and MPK_TW_SLOTS covers 256 blocks. + if (threadIdx.x == 0) { + g_tw_dev = config.tripwire; + } + __syncthreads(); +#endif +#ifdef MPK_WORKER_STATE + // Same publication for the worker-state phase breadcrumb (see MPK_WS_PHASE). + // Indexed by blockIdx.x, matching the 4-int-per-worker dump layout. + // Gated like the tripwire above: with MPK_WORKER_STATE off nothing reads + // g_ws_dev, so this store and its __syncthreads are pure overhead. + if (threadIdx.x == 0) { + g_ws_dev = config.precomp_dbg_worker_state; + // Stride to the barrier-watch half of the buffer (see MPK_WS_WAIT_BEGIN). + g_ws_nworkers = config.num_workers; + } + __syncthreads(); +#endif worker_queues[0] = config.worker_queues[worker_id]; worker_queue_ids[0] = worker_id; int num_worker_queues = 1; @@ -1267,7 +1531,7 @@ __device__ __forceinline__ void execute_worker(RuntimeConfig config, if (threadIdx.x == 0) { #ifdef MPK_PRECOMPUTED_DISPATCH // Debug: write current task position (state[0]) before dep check - if (config.precomp_dbg_worker_state != nullptr) { + if (MPK_WS_ON(config)) { int *ws = config.precomp_dbg_worker_state + worker_id * 4; __atomic_store_n(&ws[0], (int)get_task_position_index(task_ids[queue_pos]), @@ -1285,7 +1549,7 @@ __device__ __forceinline__ void execute_worker(RuntimeConfig config, EventCounter actual_counts = 0; #ifdef MPK_PRECOMPUTED_DISPATCH // Debug: write dep_event index (state[1]) - if (config.precomp_dbg_worker_state != nullptr) { + if (MPK_WS_ON(config)) { int *ws = config.precomp_dbg_worker_state + worker_id * 4; __atomic_store_n(&ws[1], (int)event_index, __ATOMIC_RELAXED); } @@ -1319,7 +1583,7 @@ __device__ __forceinline__ void execute_worker(RuntimeConfig config, #ifdef MPK_PRECOMPUTED_DISPATCH // Mark as spinning: ws[3] = -(needed - actual) to distinguish from // "not spinning" - if (config.precomp_dbg_worker_state != nullptr) { + if (MPK_WS_ON(config)) { int *ws = config.precomp_dbg_worker_state + worker_id * 4; __atomic_store_n(&ws[3], -((int)(needed_counts - actual_counts)), @@ -1351,7 +1615,7 @@ __device__ __forceinline__ void execute_worker(RuntimeConfig config, } #ifdef MPK_PRECOMPUTED_DISPATCH // Phase 10 = dep check done, about to execute task - if (config.precomp_dbg_worker_state != nullptr) { + if (MPK_WS_ON(config)) { int *ws = config.precomp_dbg_worker_state + worker_id * 4; __atomic_store_n(&ws[3], 10, __ATOMIC_RELAXED); } @@ -1783,7 +2047,7 @@ __device__ __forceinline__ void execute_worker(RuntimeConfig config, __syncthreads(); #ifdef MPK_PRECOMPUTED_DISPATCH // Phase 11 = past gang __syncthreads, entering tile loop - if (threadIdx.x == 0 && config.precomp_dbg_worker_state != nullptr) { + if (threadIdx.x == 0 && MPK_WS_ON(config)) { int *ws = config.precomp_dbg_worker_state + worker_id * 4; __atomic_store_n(&ws[3], 11, __ATOMIC_RELAXED); } @@ -1822,13 +2086,23 @@ __device__ __forceinline__ void execute_worker(RuntimeConfig config, // Layer 0: task_desc already loaded from precomputed dispatch // buffer with correct per-XCD pointers. Skip the copy. if (ml > 0) { - int ml_in_base = (xcd_id * config.ml_num_layers + ml) * 24; - int ml_out_base = (xcd_id * config.ml_num_layers + ml) * 11; - for (int i = threadIdx.x; i < 24; i += blockDim.x) { + // Widths must match the host-side ML_N_IN / ML_N_OUT that + // built these tables (see the multi-layer scan), or the + // strides disagree and every layer past 0 reads the wrong + // slots. Both sides key off the TaskDesc capacity so the + // LM-head variant's input_ptrs[24..27] / output_ptrs[12] get + // refreshed too. + int ml_in_base = + (xcd_id * config.ml_num_layers + ml) * MAX_INPUTS_PER_TASK; + int ml_out_base = + (xcd_id * config.ml_num_layers + ml) * MAX_OUTPUTS_PER_TASK; + for (int i = threadIdx.x; i < MAX_INPUTS_PER_TASK; + i += blockDim.x) { task_desc->input_ptrs[i] = config.ml_input_table[ml_in_base + i]; } - for (int i = threadIdx.x; i < 11; i += blockDim.x) { + for (int i = threadIdx.x; i < MAX_OUTPUTS_PER_TASK; + i += blockDim.x) { task_desc->output_ptrs[i] = config.ml_output_table[ml_out_base + i]; } @@ -1837,11 +2111,83 @@ __device__ __forceinline__ void execute_worker(RuntimeConfig config, } __syncthreads(); } +#ifdef MPK_NIL_TRIPWIRE + // Breadcrumb: which layer this worker reached, and whether the + // pointer set it is about to hand to the kernel contains a null. + // Written to pinned host memory so it survives an abort. + // Deliberately only plain stores to this worker's own slots: no + // scan over the pointer set (the host-side table validation + // already covers nulls there) and no global atomics on the hot + // path. An earlier version scanned all 35 pointers per layer and + // cost 4.5x -- enough to perturb timing and hide the very race + // it was meant to catch, the same way rocgdb does. + if (threadIdx.x == 0 && config.tripwire != nullptr) { + unsigned long long *b = config.tripwire + MPK_TW_HDR + + worker_id * MPK_TW_PER_WORKER; + b[0] = (unsigned long long)ml; + b[1] = 100; // phase: entering layer execution + b[3] = (unsigned long long)task_desc->input_ptrs[0]; + // Decode iteration: distinguishes "faults on the very first + // pass" from "faults after hundreds of clean ones". The host + // only sees that it died within a second of launch, which at + // ~2.5 ms/iter is anywhere in the first few hundred. + b[6] = (unsigned long long)pc_iter; + } +#endif + + // Publish the deterministic layer counter for this layer. + // + // The fused-layer task derives its three barrier release values + // from this instead of snapshotting "current counter + 1". The + // snapshot was only valid if the reader was ordered before this + // layer's producer, which nothing guaranteed -- that is the race + // f1fa720 worked around by making *every* worker arrive at the + // Phase 2 barrier, at a cost of ~0.27ms/iter. + // + // pc_iter is the decode iteration (1-based) and ml the layer, so + // this counts exactly what the per-XCD epoch counts: one bump per + // layer, never reset, monotonic across the whole run. Every + // worker computes the same value with no read of a shared + // counter, so there is nothing left to race. + // + // _linear_reserved is the free int32 of the n_tile union member + // this task already uses (n_tile_start / n_tile_count are the two + // uint16s below it); it is declared in runtime_header.h and read + // nowhere else in the tree. + // + // task_desc is a reused shared-memory slot, so this store must be + // visible to the whole block before any worker enters the task. + if (threadIdx.x == 0) { + task_desc->task_metadata._linear_reserved = + (int32_t)((pc_iter - 1) * config.ml_num_layers + ml); + } + __syncthreads(); // Execute this layer int my_tiles = 0; + // Publish the layer index into the worker-state phase slot. + // + // ws[3] was set to 11 once at gang entry and then never touched + // again for the whole task. Since all 36 layers run inside this + // one task, every worker reported phase=11 for the entire run + // and a hang dump showed "240 workers at gang entry" no matter + // which layer they were actually stuck in -- the counter could + // not distinguish a worker on layer 0 from one on layer 35. + // Encoding the layer makes the dump localize the stall. + if (threadIdx.x == 0 && MPK_WS_ON(config)) { + int *ws = config.precomp_dbg_worker_state + worker_id * 4; + __atomic_store_n(&ws[3], 40000 + ml, __ATOMIC_RELAXED); + } for (int t = block_xcd_local_rank; t < ml_n_tile_count; t += block_workers_on_xcd) { +#ifdef MPK_NIL_TRIPWIRE + if (threadIdx.x == 0 && config.tripwire != nullptr) { + unsigned long long *b = config.tripwire + MPK_TW_HDR + + worker_id * MPK_TW_PER_WORKER; + b[2] = (unsigned long long)(ml_n_tile_start + t); + b[1] = 200; // phase: inside gang tile execution + } +#endif _execute_gang_task(task_desc, config, ml_n_tile_start + t); my_tiles++; } @@ -2002,8 +2348,7 @@ __device__ __forceinline__ void execute_worker(RuntimeConfig config, int tile_idx = n_tile_start + t; #ifdef MPK_PRECOMPUTED_DISPATCH // Phase 12 = about to execute gang tile - if (threadIdx.x == 0 && - config.precomp_dbg_worker_state != nullptr) { + if (threadIdx.x == 0 && MPK_WS_ON(config)) { int *ws = config.precomp_dbg_worker_state + worker_id * 4; __atomic_store_n(&ws[3], 1200 + tile_idx, __ATOMIC_RELAXED); } @@ -2013,7 +2358,7 @@ __device__ __forceinline__ void execute_worker(RuntimeConfig config, } #ifdef MPK_PRECOMPUTED_DISPATCH // Phase 13 = tile loop done - if (threadIdx.x == 0 && config.precomp_dbg_worker_state != nullptr) { + if (threadIdx.x == 0 && MPK_WS_ON(config)) { int *ws = config.precomp_dbg_worker_state + worker_id * 4; __atomic_store_n(&ws[3], 13, __ATOMIC_RELAXED); } @@ -2064,12 +2409,23 @@ __device__ __forceinline__ void execute_worker(RuntimeConfig config, __threadfence(); } } +#endif +#ifdef MPK_NIL_TRIPWIRE + // Breadcrumb for the non-gang path, so a fault outside the multi-layer + // loop is still attributable to a worker and a task type. + if (threadIdx.x == 0 && config.tripwire != nullptr) { + unsigned long long *b = + config.tripwire + MPK_TW_HDR + worker_id * MPK_TW_PER_WORKER; + b[1] = 300; // phase: non-gang _execute_task + b[2] = (unsigned long long)task_desc->task_type; + b[3] = (unsigned long long)task_desc->input_ptrs[0]; + } #endif _execute_task(task_desc, config); } #ifdef MPK_PRECOMPUTED_DISPATCH // Phase 15 = task execution done, about to syncthreads - if (threadIdx.x == 0 && config.precomp_dbg_worker_state != nullptr) { + if (threadIdx.x == 0 && MPK_WS_ON(config)) { int *ws = config.precomp_dbg_worker_state + worker_id * 4; __atomic_store_n(&ws[3], 15, __ATOMIC_RELAXED); } @@ -2271,7 +2627,7 @@ __device__ __forceinline__ void execute_worker(RuntimeConfig config, #ifdef MPK_PRECOMPUTED_DISPATCH // Debug: increment per-worker tasks_done (state[2]) // Phase 20 = about to signal event - if (config.precomp_dbg_worker_state != nullptr) { + if (MPK_WS_ON(config)) { int *ws = config.precomp_dbg_worker_state + worker_id * 4; __atomic_store_n(&ws[2], __atomic_load_n(&ws[2], __ATOMIC_RELAXED) + 1, @@ -2502,7 +2858,7 @@ __device__ __forceinline__ void execute_worker(RuntimeConfig config, #endif #ifdef MPK_PRECOMPUTED_DISPATCH // Phase 30 = event signaling complete, looping back - if (threadIdx.x == 0 && config.precomp_dbg_worker_state != nullptr) { + if (threadIdx.x == 0 && MPK_WS_ON(config)) { int *ws = config.precomp_dbg_worker_state + worker_id * 4; __atomic_store_n(&ws[3], 30, __ATOMIC_RELAXED); } @@ -2666,10 +3022,43 @@ __device__ __forceinline__ void execute_scheduler(RuntimeConfig config, // Store per-iteration timing (deferred print at termination) if (prev_end_of_graph_clk != 0) { int idx = end_of_graph_count; // 1-based after first iter - if (idx < FWDPASS_LOG_MAX) { - g_fwdpass_time_ns[idx] = iter_end_clk - prev_end_of_graph_clk; - g_fwdpass_tokens[idx] = num_active_tokens; - g_fwdpass_count = end_of_graph_count + 1; + unsigned long long dur = iter_end_clk - prev_end_of_graph_clk; + // Aggregates first: these must cover every iteration, including the + // ones past the end of the ring. + g_fwdpass_total_ns += dur; + g_fwdpass_total_iters++; + // Stride-decimate rather than truncate. + // + // The old code kept iterations 0..8191 and dropped every one after, + // which is precisely backwards for this workload: per-iter latency + // GROWS with sequence length, so the discarded tail is the only part + // that shows the scaling. A 32k run reported the average of its + // cheapest 8k iterations and looked flat no matter how badly the + // tail degraded. + // + // Instead, once the ring fills, halve the resolution and compact: + // keep every 2nd sample, then every 4th, and so on. The ring then + // spans the whole run at uniform stride for any length, and the + // recorded token count lets the host reconstruct latency-vs-seqlen. + if (idx / g_fwdpass_stride < FWDPASS_LOG_MAX) { + if (idx % g_fwdpass_stride == 0) { + g_fwdpass_time_ns[idx / g_fwdpass_stride] = dur; + g_fwdpass_tokens[idx / g_fwdpass_stride] = num_active_tokens; + g_fwdpass_count = idx / g_fwdpass_stride + 1; + } + } else { + // Ring full at this stride: compact in place, doubling the stride. + for (int i = 0; i * 2 + 1 < FWDPASS_LOG_MAX; i++) { + g_fwdpass_time_ns[i] = g_fwdpass_time_ns[i * 2]; + g_fwdpass_tokens[i] = g_fwdpass_tokens[i * 2]; + } + g_fwdpass_stride *= 2; + g_fwdpass_count = FWDPASS_LOG_MAX / 2; + if (idx % g_fwdpass_stride == 0) { + g_fwdpass_time_ns[idx / g_fwdpass_stride] = dur; + g_fwdpass_tokens[idx / g_fwdpass_stride] = num_active_tokens; + g_fwdpass_count = idx / g_fwdpass_stride + 1; + } } } prev_end_of_graph_clk = iter_end_clk; @@ -2688,13 +3077,26 @@ __device__ __forceinline__ void execute_scheduler(RuntimeConfig config, iteration_num, (double)(prep_done_clk - iter_end_clk) / 1000.0); #endif - // Dump deferred FWD_PASS log (after timing, before terminate) + // Dump deferred FWD_PASS log (after timing, before terminate). + // iter is reconstructed from the decimation stride. for (int i = 1; i < g_fwdpass_count && i < FWDPASS_LOG_MAX; i++) { printf("[FWD_PASS] iter=%d time_ms=%.3f num_active_tokens=%d\n", - i, + i * g_fwdpass_stride, (double)g_fwdpass_time_ns[i] / 1000000.0, g_fwdpass_tokens[i]); } + // Untruncated summary. dropped>0 means the per-iter lines above are + // only the first FWDPASS_LOG_MAX iterations and must not be averaged + // as if they were the whole run -- use total_ms/iters instead. + printf("[FWD_PASS_TOTAL] iters=%d total_ms=%.3f avg_ms=%.3f " + "dropped=%d\n", + g_fwdpass_total_iters, + (double)g_fwdpass_total_ns / 1000000.0, + g_fwdpass_total_iters > 0 + ? (double)g_fwdpass_total_ns / 1000000.0 / + (double)g_fwdpass_total_iters + : 0.0, + g_fwdpass_dropped); #ifdef MPK_ENABLE_MOE_SUBPHASE // Raw timestamps: scratch[0]=entry, [1]=before_lds, // [4]=after_compute, [2]=after_barrier @@ -3158,6 +3560,139 @@ static int g_dbg_num_events = 0; static int g_dbg_num_xcds = 0; #endif +#ifdef MPK_NIL_TRIPWIRE +// ── Tripwire for the nil-address GPU memory fault ──────────────────────── +// The fault aborts the process. Device printf output is lost, and the GPU +// coredump ROCm writes carries no wavefront state, so the usual channels +// tell us nothing. This buffer is pinned host memory mapped into the GPU +// address space: kernel writes reach host RAM as they happen and survive +// the abort, and a SIGABRT handler prints them. +// +// Layout: MPK_TW_HDR unused header slots, then MPK_TW_PER_WORKER per worker. +// Per worker w, base = MPK_TW_HDR + w * MPK_TW_PER_WORKER: +// [base+0] last layer this worker entered +// [base+1] outer phase marker (100 layer entry, 200 gang tile, 300 non-gang) +// [base+2] last tile_idx +// [base+3] the pointer this worker was about to dereference +// [base+4] sub-phase inside the fused layer kernel (see MPK_TW_SUB calls: +// 1 entry, 10 QKV, 20 QKV barrier, 30 attn, 40 chunk barrier, +// 50 merge, 60 cross-XCD barrier, 70 O-proj/TopK, 75 TopK wait, +// 80 MoE, 90 kernel exit) +// [base+5] sub-phase aux value (meaning depends on the sub-phase) +// [base+6] decode iteration (pc_iter) +static unsigned long long *g_tw_host = nullptr; +static int g_tw_num_workers = 0; +static bool g_tw_handler_installed = false; + +// Async-signal-safe-ish dump. fprintf from a fatal handler is not strictly +// legal, but the process is aborting anyway and getting the data out matters +// more than standards purity here. +static void tripwire_dump(int sig) { + if (g_tw_host == nullptr) { + _exit(134); + } + fprintf(stderr, "\n===== MPK NIL TRIPWIRE (signal %d) =====\n", sig); + // Per-worker breadcrumbs only; there is no global header counter, since + // maintaining one would mean a device-wide atomic on the hot path. + // What matters is the spread: if one worker is on a different layer from + // the rest, or holds a null ptr, that names the culprit. + int printed = 0; + unsigned long long min_layer = ~0ull, max_layer = 0; + int nulls = 0; + for (int w = 0; w < g_tw_num_workers; w++) { + unsigned long long const *b = + g_tw_host + MPK_TW_HDR + w * MPK_TW_PER_WORKER; + if (b[0] == 0 && b[1] == 0 && b[2] == 0 && b[3] == 0) { + continue; // worker never wrote a breadcrumb + } + fprintf(stderr, + " w%-3d it=%llu layer=%llu phase=%llu sub=%llu aux=%llu " + "tile=%llu ptr=0x%llx%s\n", + w, + b[6], + b[0], + b[1], + b[4], + b[5], + b[2], + b[3], + b[3] == 0 ? " <-- NULL" : ""); + if (b[3] == 0) { + nulls++; + } + if (b[0] < min_layer) { + min_layer = b[0]; + } + if (b[0] > max_layer) { + max_layer = b[0]; + } + printed++; + } + fprintf( + stderr, + "(%d workers left breadcrumbs; layer span %llu..%llu; %d null ptrs)\n", + printed, + printed ? min_layer : 0, + max_layer, + nulls); + fprintf(stderr, "===== END TRIPWIRE =====\n"); + fflush(stderr); + // Restore default disposition and re-raise so the exit status is unchanged. + signal(sig, SIG_DFL); + raise(sig); +} + +// Snapshot the breadcrumbs to disk on a timer. +// +// The signal handler alone is not enough: the ROCm runtime installs its own +// SIGABRT disposition after we arm ours, so on a real memory fault the +// process dies without our handler ever running (observed). A file on disk +// survives regardless of who wins the signal-handler race, and the fault +// aborts within ~1s of launch, so a short interval catches it. +static bool volatile g_tw_snap_stop = false; +static std::thread g_tw_snap_thread; + +static void tripwire_snapshot_loop() { + while (!g_tw_snap_stop) { + FILE *f = fopen("/tmp/mpk_tripwire.txt", "w"); + if (f != nullptr) { + for (int w = 0; w < g_tw_num_workers; w++) { + unsigned long long const *b = + g_tw_host + MPK_TW_HDR + w * MPK_TW_PER_WORKER; + if (b[0] == 0 && b[1] == 0 && b[2] == 0 && b[3] == 0) { + continue; + } + fprintf(f, + "w%-3d it=%llu layer=%llu phase=%llu sub=%llu aux=%llu " + "tile=%llu ptr=0x%llx%s\n", + w, + b[6], + b[0], + b[1], + b[4], + b[5], + b[2], + b[3], + b[3] == 0 ? " <-- NULL" : ""); + } + fclose(f); + } + std::this_thread::sleep_for(std::chrono::milliseconds(100)); + } +} + +static void install_tripwire_handler() { + if (g_tw_handler_installed) { + return; + } + signal(SIGABRT, tripwire_dump); + signal(SIGSEGV, tripwire_dump); + g_tw_snap_thread = std::thread(tripwire_snapshot_loop); + g_tw_snap_thread.detach(); + g_tw_handler_installed = true; +} +#endif + // meta_tensors[0]: seq_length // meta_tensors[1]: tokens // meta_tensors[2]: input_tokens @@ -3344,8 +3879,17 @@ extern "C" void init_persistent_kernel(std::vector meta_tensors, #ifdef MPK_FUSED_LAYER_BATCHING { constexpr int NUM_XCDS_ML = 8; - constexpr int ML_N_IN = 24; - constexpr int ML_N_OUT = 11; + // Must cover every slot the fused layer kernels read, not just the ones + // the plain variant uses. The LM-head variant + // (TASK_GANG_FULL_LAYER_WITH_LMHEAD_FUSED_MI300, FUSE_TAIL=1) reads + // input_ptrs[24..27] and output_ptrs[12] -- see + // gang_full_layer_with_lmhead_fused_mi300.cuh. These were 24 and 11, + // sized for the plain variant, so with FUSE_TAIL=1 the ml loop would + // refresh only the first 24/11 slots and leave the LM-head pointers + // holding layer 0's values for every later layer. Sizing to the TaskDesc + // capacity keeps this correct for any variant. + constexpr int ML_N_IN = MAX_INPUTS_PER_TASK; // 28 + constexpr int ML_N_OUT = MAX_OUTPUTS_PER_TASK; // 13 int n_tasks_pre = (int)all_tasks.size(); int n_events_pre = (int)all_events.size(); @@ -3406,6 +3950,69 @@ extern "C" void init_persistent_kernel(std::vector meta_tensors, (int)td.task_type); } + // Validate the tables before upload. The layer loop dereferences every + // slot the kernel reads, so a null there is a nil-address GPU fault with + // no usable backtrace -- catching it on the host names the exact slot. + // + // Only the slots the task type actually declares are checked. The tables + // are sized to the TaskDesc capacity so the LM-head variant's higher + // slots get refreshed, but the plain variant declares 24 inputs / 11 + // outputs (see task_config in graph.cc), and slots past its count are + // legitimately null -- TaskDesc value-initializes them. Flagging those + // would report 1152/576 "nulls" on a perfectly healthy graph and train + // the reader to ignore this line. + { + int null_in = 0, null_out = 0, unused_in = 0, unused_out = 0; + for (int L = 0; L < ml_layers; L++) { + bool is_lmhead = all_tasks[fused_layer_positions[L]].task_type == + TASK_GANG_FULL_LAYER_WITH_LMHEAD_FUSED_MI300; + int used_in = is_lmhead ? 28 : 24; + int used_out = is_lmhead ? 13 : 11; + for (int xcd = 0; xcd < NUM_XCDS_ML; xcd++) { + int base = (xcd * ml_layers + L) * ML_N_IN; + for (int i = 0; i < ML_N_IN; i++) { + if (h_input_table[base + i] == nullptr) { + if (i >= used_in) { + unused_in++; + continue; + } + if (null_in < 8) { + printf( + "[MPK] ML_NULL_IN layer=%d xcd=%d idx=%d\n", L, xcd, i); + } + null_in++; + } + } + int obase = (xcd * ml_layers + L) * ML_N_OUT; + for (int i = 0; i < ML_N_OUT; i++) { + if (h_output_table[obase + i] == nullptr) { + if (i >= used_out) { + unused_out++; + continue; + } + if (null_out < 8) { + printf( + "[MPK] ML_NULL_OUT layer=%d xcd=%d idx=%d\n", L, xcd, i); + } + null_out++; + } + } + } + } + // Only report when something is actually wrong. A healthy graph has + // null_in == null_out == 0, and printing that on every run trains the + // reader to skip the line that matters. + if (null_in || null_out) { + printf("[MPK] ML table validation: %d null inputs, %d null outputs " + "(%d/%d null in undeclared slots, expected)\n", + null_in, + null_out, + unused_in, + unused_out); + fflush(stdout); + } + } + // === Compact task graph: remove layers 1..35 tasks and inter-layer // events === Collect tasks and events to remove std::set tasks_to_remove; // positions in all_tasks @@ -3790,11 +4397,6 @@ extern "C" void init_persistent_kernel(std::vector meta_tensors, 0, NUM_XCDS_PC * sizeof(int)); - printf("[MPK] XCD templates uploaded: %d entries x %d max_tpw\n", - total_template_entries, - max_tpw); - fflush(stdout); - // Use device memory for iter_ready and terminate (GPU-only polling) { global_runtime_config.precomp_iter_ready = @@ -3821,10 +4423,79 @@ extern "C" void init_persistent_kernel(std::vector meta_tensors, 0, sizeof(unsigned long long)); - // Debug worker state — disabled for performance + // Debug worker state — off by default (the per-task relaxed stores cost + // a little on the hot path), opt in with MPK_WORKER_STATE=1 to make the + // host poll loop dump where every worker is parked when it hangs. + // + // Host-mapped, not device memory: on a hang the kernel never returns, so + // the host must be able to read this while the kernel is still spinning. global_runtime_config.precomp_dbg_worker_state = nullptr; g_dbg_h_worker_state = nullptr; g_dbg_num_workers = num_workers; + { + char const *ws_env = getenv("MPK_WORKER_STATE"); + if (ws_env != nullptr && atoi(ws_env) != 0) { + // Three thirds: [0, num_workers*4) is the per-worker phase state, + // [num_workers*4, num_workers*8) is the barrier watch written by + // MPK_WS_WAIT_BEGIN / MPK_WS_WAIT_TICK, and + // [num_workers*8, num_workers*12) is the per-barrier aux written by + // MPK_WS_WAIT_AUX. + size_t ws_bytes = (size_t)num_workers * 12 * sizeof(int); + int *ws_host = nullptr; + if (hipHostMalloc(reinterpret_cast(&ws_host), + ws_bytes, + hipHostMallocMapped | hipHostMallocNonCoherent) == + hipSuccess && + ws_host != nullptr) { + memset(ws_host, 0, ws_bytes); + int *ws_dev = nullptr; + if (hipHostGetDevicePointer(reinterpret_cast(&ws_dev), + ws_host, + 0) == hipSuccess) { + g_dbg_h_worker_state = ws_host; + global_runtime_config.precomp_dbg_worker_state = ws_dev; + fprintf(stderr, + "[HOST_DBG] worker-state tracing ON (%d workers)\n", + num_workers); + } else { + (void)hipHostFree(ws_host); + } + } + } + } + +#ifdef MPK_NIL_TRIPWIRE + // Tripwire for the nil-address memory fault. Backed by *pinned host* + // memory, not device memory: the fault aborts the process, so anything + // living only in VRAM (and device printf, and the GPU coredump) is + // gone by the time we could read it. Writes from the kernel land in + // host RAM over PCIe as they happen, so whatever the last worker wrote + // before the fault is still there for the SIGABRT handler to print. + (void)hipHostMalloc(reinterpret_cast(&g_tw_host), + MPK_TW_SLOTS * sizeof(unsigned long long), + hipHostMallocMapped | hipHostMallocNonCoherent); + if (g_tw_host != nullptr) { + for (int i = 0; i < MPK_TW_SLOTS; i++) { + g_tw_host[i] = 0; + } + unsigned long long *tw_dev = nullptr; + (void)hipHostGetDevicePointer( + reinterpret_cast(&tw_dev), g_tw_host, 0); + global_runtime_config.tripwire = tw_dev; + g_tw_num_workers = num_workers; + install_tripwire_handler(); + printf("[MPK] nil tripwire armed: host=%p dev=%p slots=%d\n", + (void *)g_tw_host, + (void *)tw_dev, + MPK_TW_SLOTS); + } else { + global_runtime_config.tripwire = nullptr; + printf("[MPK] nil tripwire: hipHostMalloc FAILED, not armed\n"); + } + fflush(stdout); +#else + global_runtime_config.tripwire = nullptr; +#endif // Clear host debug pointers (no longer host-mapped) g_dbg_h_iter_ready = nullptr; @@ -4211,39 +4882,104 @@ extern "C" void launch_persistent_kernel(cudaStream_t default_stream) { global_runtime_config.precomp_terminate) { hipStream_t dbg_stream; (void)hipStreamCreate(&dbg_stream); - for (int dbg_i = 0; dbg_i < 30; dbg_i++) { - std::this_thread::sleep_for(std::chrono::seconds(1)); - unsigned long long ir = 0; - int term = 0; - unsigned long long td = 0; - (void)hipMemcpyAsync(&ir, - global_runtime_config.precomp_iter_ready, - sizeof(unsigned long long), - hipMemcpyDeviceToHost, - dbg_stream); - (void)hipMemcpyAsync(&term, - global_runtime_config.precomp_terminate, - sizeof(int), - hipMemcpyDeviceToHost, - dbg_stream); - if (global_runtime_config.precomp_dbg_tasks_done) { - (void)hipMemcpyAsync(&td, - global_runtime_config.precomp_dbg_tasks_done, + // The device-memory poll runs on its own detached thread. + // + // This is not premature generality: when the kernel hangs, the ROCm + // runtime wedges with it, and hipMemcpyAsync/hipStreamSynchronize below + // block forever *even on a private stream*. Running them inline is what + // made this instrumentation useless on the one case it exists for -- the + // hung 32k run printed "Starting poll loop..." and then not one line + // more, because it never got past the first sync to reach the dump. + // + // The worker-state and event-counter buffers are pinned host memory, so + // the dumps below need no runtime call and keep printing regardless. + // dev_polls staying at 0 is itself the signal that HIP is stuck. + static std::atomic dbg_a_ir{0}; + static std::atomic dbg_a_term{0}; + static std::atomic dbg_a_td{0}; + static std::atomic dbg_a_polls{0}; + static std::atomic dbg_a_stop{false}; + // Detached, and every object it touches has static storage duration -- + // a thread blocked in HIP outlives this scope and must not dangle. + std::thread([dbg_stream]() { + while (!dbg_a_stop.load(std::memory_order_relaxed)) { + unsigned long long ir = 0; + int term = 0; + unsigned long long td = 0; + (void)hipMemcpyAsync(&ir, + global_runtime_config.precomp_iter_ready, sizeof(unsigned long long), hipMemcpyDeviceToHost, dbg_stream); + (void)hipMemcpyAsync(&term, + global_runtime_config.precomp_terminate, + sizeof(int), + hipMemcpyDeviceToHost, + dbg_stream); + if (global_runtime_config.precomp_dbg_tasks_done) { + (void)hipMemcpyAsync(&td, + global_runtime_config.precomp_dbg_tasks_done, + sizeof(unsigned long long), + hipMemcpyDeviceToHost, + dbg_stream); + } + (void)hipStreamSynchronize(dbg_stream); + dbg_a_ir.store(ir, std::memory_order_relaxed); + dbg_a_term.store(term, std::memory_order_relaxed); + dbg_a_td.store(td, std::memory_order_relaxed); + dbg_a_polls.fetch_add(1, std::memory_order_relaxed); + if (term) { + break; + } + std::this_thread::sleep_for(std::chrono::milliseconds(200)); + } + }).detach(); + // A 32k run takes minutes, so the original 30-tick (30s) bound stopped + // watching long before the interesting window: a hang at iteration + // ~20000 produced no dump at all, because the loop had already exited + // into the blocking stream sync below. Poll until the kernel actually + // reports terminate (the `if (term) break` at the bottom), bounded only + // by a generous ceiling so a wedged run still exits the loop and lets + // the enclosing `timeout` reap it. + int const dbg_max_ticks = 900; + for (int dbg_i = 0; dbg_i < dbg_max_ticks; dbg_i++) { + std::this_thread::sleep_for(std::chrono::seconds(1)); + unsigned long long ir = dbg_a_ir.load(std::memory_order_relaxed); + int term = dbg_a_term.load(std::memory_order_relaxed); + unsigned long long td = dbg_a_td.load(std::memory_order_relaxed); + unsigned long long polls = dbg_a_polls.load(std::memory_order_relaxed); + // dev_polls is normally 0 for the whole run -- a healthy 32k run shows + // 0 at every tick. hipMemcpyAsync does not complete while the + // persistent kernel occupies the GPU, even on a private stream, so + // iter_ready/terminate/tasks_done stay 0 and mean nothing until the + // kernel exits. Do NOT read dev_polls=0 as a hang. + // + // The real liveness signal is the per-worker `done` counter in the + // dump below: it lives in pinned host memory and advances even while + // HIP is blocked. Frozen `done` across ticks = actually stuck. + if (dbg_i % 30 == 0 || term) { + fprintf( + stderr, + "[HOST_DBG] t=%ds iter_ready=%llu terminate=%d " + "tasks_done=%llu dev_polls=%llu (0 is normal; watch `done`)\n", + dbg_i + 1, + ir, + term, + td, + polls); } - (void)hipStreamSynchronize(dbg_stream); - fprintf( - stderr, - "[HOST_DBG] t=%ds iter_ready=%llu terminate=%d tasks_done=%llu\n", - dbg_i + 1, - ir, - term, - td); // Dump per-worker state + summary if (g_dbg_h_worker_state && dbg_i >= 2) { int phase_hist[4] = {}; + // Fused-layer phase histogram over ALL workers, not just the 32 + // printed individually. A stall where 239 workers sit on one barrier + // and a single straggler sits elsewhere is invisible in a 32-worker + // sample, and that straggler is the one holding the barrier. + // Indexed by (phase-50000)/100; also track min/max layer to expose + // cross-worker layer skew, which is the real deadlock signature. + std::map fl_phase_hist; + std::map fl_px_hist; // key = phase*10 + xcd + int fl_layer_min = 1 << 30, fl_layer_max = -1; int dep_hist[256] = {}; int phase_gang_stuck = 0; int phase_gang_entry = 0; @@ -4251,6 +4987,37 @@ extern "C" void launch_persistent_kernel(cudaStream_t default_stream) { // 3=w2poll(3001) 4=w2polldone(3002) // 5=w2quant(3003) 6=w2lds(3004) 7=w2mfma(3005) int moe_epilogue = 0; // 3006 + // Liveness first, so the verbose decision below can use it. + // + // Comparing each worker's `done` against the previous tick is the + // only reliable progress test here (see the dev_polls note above). + // Tracked per-worker, not as an aggregate, so one frozen straggler + // among 239 live workers cannot hide. + static std::vector prev_done; + static int stall_ticks = 0; + int frozen = 0; + bool have_prev = ((int)prev_done.size() == g_dbg_num_workers); + { + std::vector cur_done(g_dbg_num_workers); + for (int w = 0; w < g_dbg_num_workers; w++) { + cur_done[w] = __atomic_load_n(&g_dbg_h_worker_state[w * 4 + 2], + __ATOMIC_RELAXED); + } + if (have_prev) { + for (int w = 0; w < g_dbg_num_workers; w++) { + if (cur_done[w] == prev_done[w]) { + frozen++; + } + } + stall_ticks = (frozen == g_dbg_num_workers) ? stall_ticks + 1 : 0; + } + prev_done.swap(cur_done); + } + // Now that the loop runs for the whole kernel rather than 30s, the + // per-worker lines would be ~32 lines/sec of healthy-run noise and + // would bury the failure. Print them only while a stall is building + // or on a periodic checkpoint. + bool verbose = (stall_ticks >= 1) || (dbg_i % 30 == 0); for (int w = 0; w < g_dbg_num_workers; w++) { int *ws = g_dbg_h_worker_state + w * 4; int tpos = __atomic_load_n(&ws[0], __ATOMIC_RELAXED); @@ -4260,20 +5027,96 @@ extern "C" void launch_persistent_kernel(cudaStream_t default_stream) { int sc = phase / 100000; int moe_xcd = (phase / 10000) % 10; int moe_tile = phase % 10000; - if (w < 32 || sc == 3001) { - fprintf(stderr, - " w%d: pos=%d dep=%d done=%d phase=%d (sc=%d xcd=%d " - "tile=%d)\n", - w, - tpos, - dep, - done, - phase, - sc, - moe_xcd, - moe_tile); + if (verbose && (w < 32 || sc == 3001)) { + // 40000+ml is the multi-layer marker written in the ml loop. + // Print it as a layer number rather than a raw phase, since + // "which layer" is the whole point of that encoding. + if (phase >= 50000000 && phase < 51000000) { + // 50000000 + xcd*100000 + phase*1000 + layer%1000. + int fx = (phase - 50000000) / 100000; + int fp = ((phase - 50000000) / 1000) % 100; + int fl = (phase - 50000000) % 1000; + char const *pn = "?"; + switch (fp) { + case 20: + pn = "P2-qkv-epoch"; + break; + case 30: + pn = "P3-attn-chunk"; + break; + case 40: + pn = "P4-chunk-barrier"; + break; + case 50: + pn = "P5-merge"; + break; + case 60: + pn = "P6-attn-xcd-barrier"; + break; + case 70: + pn = "P7-oproj"; + break; + case 75: + pn = "P7b-routing-poll"; + break; + case 80: + pn = "P8-moe"; + break; + case 90: + pn = "P9-layer-done"; + break; + } + fprintf(stderr, + " w%d: pos=%d dep=%d done=%d xcd=%d epoch=%d %s\n", + w, + tpos, + dep, + done, + fx, + fl, + pn); + } else if (phase >= 40000 && phase < 40000 + 1024) { + fprintf(stderr, + " w%d: pos=%d dep=%d done=%d layer=%d\n", + w, + tpos, + dep, + done, + phase - 40000); + } else { + fprintf(stderr, + " w%d: pos=%d dep=%d done=%d phase=%d (sc=%d xcd=%d " + "tile=%d)\n", + w, + tpos, + dep, + done, + phase, + sc, + moe_xcd, + moe_tile); + } } - if (phase >= 1200 && phase < 300100000) { + // Exclude the 40000+ml multi-layer marker: it means "executing + // layer ml", which is normal progress, not a stuck gang tile. + if (phase >= 50000000 && phase < 51000000) { + int fp = ((phase - 50000000) / 1000) % 100; + int fx = (phase - 50000000) / 100000; + int fl = (phase - 50000000) % 1000; + fl_phase_hist[fp]++; + // Per-(phase,XCD) tally: a barrier stall shows as one XCD short + // of its expected arrivals while the others are complete. + fl_px_hist[fp * 10 + fx]++; + if (fl < fl_layer_min) { + fl_layer_min = fl; + } + if (fl > fl_layer_max) { + fl_layer_max = fl; + } + } + if (phase >= 1200 && phase < 300100000 && + !(phase >= 40000 && phase < 40000 + 1024) && + !(phase >= 50000000 && phase < 51000000)) { phase_gang_stuck++; } if (sc == 2000) { @@ -4311,36 +5154,393 @@ extern "C" void launch_persistent_kernel(cudaStream_t default_stream) { phase_hist[3]++; } } - fprintf(stderr, - " phases: spinning=%d dep_done=%d signaling=%d sig_done=%d " - "gang_tile=%d gang_entry=%d", - phase_hist[0], - phase_hist[1], - phase_hist[2], - phase_hist[3], - phase_gang_stuck, - phase_gang_entry); - fprintf(stderr, - "\n moe: w13=%d w13sig=%d w2entry=%d w2poll=%d w2done=%d " - "w2quant=%d w2lds=%d w2mfma=%d w2epi=%d", - moe_sub[0], - moe_sub[1], - moe_sub[2], - moe_sub[3], - moe_sub[4], - moe_sub[5], - moe_sub[6], - moe_sub[7], - moe_epilogue); - if (phase_hist[0] > 0) { - fprintf(stderr, " | stuck_deps:"); - for (int d = 0; d < 256; d++) { - if (dep_hist[d] > 0) { - fprintf(stderr, " ev%d=%d", d, dep_hist[d]); + if (have_prev && (verbose || stall_ticks >= 1)) { + fprintf(stderr, + " progress: %d/%d workers advanced since last tick%s\n", + g_dbg_num_workers - frozen, + g_dbg_num_workers, + stall_ticks >= 2 ? " *** STALLED: no worker has advanced " + "for 3+ ticks ***" + : ""); + } + // Barrier watch. Only meaningful once the run is actually wedged -- + // in a healthy run these slots hold whatever barrier each worker + // last cleared, which is noise. On a stall it is the whole answer: + // observed vs expected separates "the producer never arrived" + // (observed short by N) from "this waiter expects an epoch the + // producer will never reach" (observed >= expected, i.e. the poll + // should have cleared and the load is reading a stale line). + if (stall_ticks >= 2) { + // key = barrier_id, value = (count, min observed, expected, + // max spins) aggregated over workers still spinning. + std::map> bw; + // The watch slots are last-write-wins and are never cleared on + // exit, so a worker that finished its layer long ago still shows + // whatever mark it passed on the way out. Reading them without + // cross-checking the live phase is how the first capture produced + // 56 "blockers" that were all idle at the scheduler (phase == -1) + // while the one worker actually inside the layer went unprinted. + // Only a worker whose phase says it is still executing can be + // holding anything, so gate every watch read on that. + for (int w = 0; w < g_dbg_num_workers; w++) { + int wphase = __atomic_load_n(&g_dbg_h_worker_state[w * 4 + 3], + __ATOMIC_RELAXED); + bool in_task = (wphase >= 0); + int *b = g_dbg_h_worker_state + g_dbg_num_workers * 4 + w * 4; + int bid = __atomic_load_n(&b[0], __ATOMIC_RELAXED); + int obs = __atomic_load_n(&b[1], __ATOMIC_RELAXED); + int exp_ = __atomic_load_n(&b[2], __ATOMIC_RELAXED); + int spins = __atomic_load_n(&b[3], __ATOMIC_RELAXED); + if (!in_task) { + continue; // idle at the scheduler; its watch slot is stale + } + // Whatever this worker's watch slot says, it is genuinely still + // inside the task -- print the live phase too, since that is the + // ground truth the watch slot has to be read against. + { + char const *pn = "?"; + int fx = -1, fp = -1, fl = -1; + if (wphase >= 50000000 && wphase < 51000000) { + fx = (wphase - 50000000) / 100000; + fp = ((wphase - 50000000) / 1000) % 100; + fl = (wphase - 50000000) % 1000; + switch (fp) { + case 20: + pn = "P2-qkv-epoch"; + break; + case 30: + pn = "P3-attn-chunk"; + break; + case 40: + pn = "P4-chunk-barrier"; + break; + case 50: + pn = "P5-merge"; + break; + case 60: + pn = "P6-attn-xcd-barrier"; + break; + case 70: + pn = "P7-oproj"; + break; + case 75: + pn = "P7b-routing-poll"; + break; + case 80: + pn = "P8-moe"; + break; + case 90: + pn = "P9-layer-done"; + break; + } + } + int *a = g_dbg_h_worker_state + g_dbg_num_workers * 8 + w * 4; + int a0 = __atomic_load_n(&a[0], __ATOMIC_RELAXED); + int a1 = __atomic_load_n(&a[1], __ATOMIC_RELAXED); + int a2 = __atomic_load_n(&a[2], __ATOMIC_RELAXED); + int a3 = __atomic_load_n(&a[3], __ATOMIC_RELAXED); + fprintf(stderr, + " >>> w%d IN TASK: phase=%d (%s xcd=%d epoch=%d) " + "watch[bid=%d obs=%d exp=%d spins=%d]\n", + w, + wphase, + pn, + fx, + fl, + bid, + obs, + exp_, + spins); + // Per-wave exit mask for the MoE W13->W2 poll. That poll is + // thread-divergent by construction (each thread tests the + // release flag on its own, no __syncthreads), so waves of one + // block leave it independently. tid 0 leaving is NOT the block + // leaving: gx_1 showed w30 reporting a straight-line mark past + // the barrier (8303, inside the quant's trailing + // __syncthreads) while the block was still short by 2. A mask + // below 0xf names the waves that never saw the release. + if (fp == 80) { + int smask = a2 & 0xf; + fprintf(stderr, + " quant sync mask=0x%x (%s)\n", + smask, + smask == 0xf + ? "all 4 waves reached the quant __syncthreads" + : "waves missing -- note NSUBBLOCKS(96) < " + "blockDim(256), so waves 1..3 do fewer or " + "zero loop trips and legitimately arrive " + "early; only a mask frozen across dumps is " + "evidence"); + int mask = a3 & 0xf; + fprintf(stderr, + " wave exit mask=0x%x (%s)\n", + mask, + mask == 0xf + ? "all 4 waves cleared the W13->W2 poll" + : (mask == 0 + ? "no wave cleared it" + : "<== BLOCK SPLIT ACROSS THE BARRIER: " + "some waves cleared, others still " + "spinning -- the cleared waves are " + "parked in the next __syncthreads")); + } + // spins == 0 means the poll cleared on its very first load + // without ever ticking, so obs/exp and the aux below are + // leftovers from an earlier layer and must not be read as this + // worker's current state. An in-task worker with spins == 0 is + // stuck somewhere *past* its last barrier, not at one -- which + // is exactly what the fx_8 capture showed after the cache-line + // split, and what the aux misreported as a lost fan-out. + if (spins == 0) { + fprintf(stderr, + " (watch slot STALE: last barrier cleared on " + "its first load -- this worker is stuck past it, " + "read the mark, not obs/exp)\n"); + } + // a2/a3 are the per-wave masks now, so the slot-spread + // decoding below no longer applies. Keep the arrival counter, + // which still lives in a0. + if (spins > 0 && bid >= 800 && bid < 900) { + fprintf(stderr, + " MoE arrivals=%d (arrivals%%46=%d) " + "expert_id=%d\n", + a0, + a0 % 46, + a1); + } + if (false) { + // a0 = raw arrival counter. If a0 is an exact multiple of + // W13_TILES every producer arrived and the release fired. + // a2 packs (slots_at_or_past_expected)*1e6 + (max-min) + // across the 8 per-XCD release slots of this expert; a3 is + // the min. All 8 are written by one producer in one loop, so + // spread != 0 means the fan-out writes are being lost. + int n_ok = a2 / 1000000; + int spread = a2 % 1000000; + fprintf(stderr, + " MoE aux: arrivals=%d (arrivals%%46=%d) " + "expert_id=%d slots_ok=%d/8 spread=%d min=%d %s\n", + a0, + a0 % 46, + a1, + n_ok, + spread, + a3, + spread != 0 + ? "<== SLOTS DISAGREE: release fan-out lost" + : (a0 % 46 == 0 + ? "<== all arrived, all slots equal: " + "release VALUE is stale" + : "<== arrivals missing")); + } + } + // spins == -1 is a straight-line mark (MPK_WS_MARK), not a poll: + // the worker is running, and the code says where. Print those + // separately -- on a total stall they are the blockers, since a + // worker that is not in any poll is the one everyone waits on. + if (spins < 0) { + int mcode = (-spins) / 100000; + int maux = (-spins) % 100000; + char const *mn = "?"; + switch (mcode) { + case 8000: + mn = "MoE tile loop"; + break; + case 8100: + mn = "MoE exit: past tile range"; + break; + case 8101: + mn = "MoE exit: W13 padding tile"; + break; + case 8102: + mn = "MoE exit: token out of batch"; + break; + case 8103: + mn = "MoE exit: token not routed"; + break; + case 8200: + mn = "MoE W13 compute"; + break; + case 8201: + mn = "MoE W13 done, at barrier"; + break; + case 8300: + mn = "MoE W2 entry"; + break; + case 8302: + mn = "MoE W2: cleared W13->W2 barrier"; + break; + case 8303: + mn = "MoE W2: FP8 quant"; + break; + case 8304: + mn = "MoE W2: drain HBM loads"; + break; + case 8305: + mn = "MoE W2: MFMA loop"; + break; + case 8306: + mn = "MoE W2: epilogue"; + break; + } + fprintf(stderr, + " w%d NOT POLLING: mark %d (%s) tile=%d" + " <== running, not blocked; everyone else waits on " + "this\n", + w, + mcode, + mn, + maux); + continue; + } + // spins==0 means this worker cleared its last barrier without + // ever ticking; it is not waiting, so it is not the blocker. + if (spins == 0) { + continue; + } + auto it = bw.find(bid); + if (it == bw.end()) { + bw[bid] = {1, obs, exp_, spins, w}; + } else { + it->second[0]++; + if (obs < it->second[1]) { + it->second[1] = obs; + } + if (spins > it->second[3]) { + it->second[3] = spins; + it->second[4] = w; + } + } + } + if (!bw.empty()) { + fprintf(stderr, + " *** BARRIER WATCH (workers still spinning) " + "***\n"); + for (auto const &kv : bw) { + char const *bn = "?"; + int bid = kv.first; + if (bid >= 800 && bid < 900) { + bn = "MoE-W13->W2"; + } else if (bid == 20) { + bn = "P2-qkv-epoch"; + } else if (bid == 60) { + bn = "P6-attn-xcd"; + } else if (bid == 75) { + bn = "P7b-routing"; + } + fprintf(stderr, + " barrier %d (%s%s%d): %lld workers spinning, " + "observed=%lld expected=%lld short_by=%lld " + "max_spins=%lld (w%lld)%s\n", + bid, + bn, + bid >= 800 && bid < 900 ? " expert " : "", + bid >= 800 && bid < 900 ? bid - 800 : 0, + kv.second[0], + kv.second[1], + kv.second[2], + kv.second[2] - kv.second[1], + kv.second[3], + kv.second[4], + kv.second[1] >= kv.second[2] + ? " <== STALE READ: observed >= expected, the " + "poll should have cleared" + : ""); } } } - fprintf(stderr, "\n"); + if (verbose) { + if (!fl_phase_hist.empty()) { + fprintf(stderr, + " fused-layer phases (all %d workers):", + g_dbg_num_workers); + for (auto const &kv : fl_phase_hist) { + char const *pn = "?"; + switch (kv.first) { + case 20: + pn = "P2-qkv-epoch"; + break; + case 30: + pn = "P3-attn-chunk"; + break; + case 40: + pn = "P4-chunk-barrier"; + break; + case 50: + pn = "P5-merge"; + break; + case 60: + pn = "P6-attn-xcd-barrier"; + break; + case 70: + pn = "P7-oproj"; + break; + case 75: + pn = "P7b-routing-poll"; + break; + case 80: + pn = "P8-moe"; + break; + case 90: + pn = "P9-layer-done"; + break; + } + fprintf(stderr, " %s=%d", pn, kv.second); + } + fprintf(stderr, + " epoch=[%d..%d]%s\n", + fl_layer_min, + fl_layer_max, + fl_layer_min != fl_layer_max + ? " *** EPOCH SKEW: workers on different layers ***" + : ""); + // Per-XCD breakdown. Each XCD runs 30 workers, so a phase that + // shows fewer than 30 on some XCD is missing arrivals there -- + // that XCD is the one holding a per-XCD barrier. + for (auto const &kv : fl_phase_hist) { + fprintf(stderr, " phase %d per-xcd:", kv.first); + for (int x = 0; x < 8; x++) { + auto it = fl_px_hist.find(kv.first * 10 + x); + fprintf(stderr, + " x%d=%d", + x, + it == fl_px_hist.end() ? 0 : it->second); + } + fprintf(stderr, "\n"); + } + } + fprintf( + stderr, + " phases: spinning=%d dep_done=%d signaling=%d sig_done=%d " + "gang_tile=%d gang_entry=%d", + phase_hist[0], + phase_hist[1], + phase_hist[2], + phase_hist[3], + phase_gang_stuck, + phase_gang_entry); + fprintf(stderr, + "\n moe: w13=%d w13sig=%d w2entry=%d w2poll=%d w2done=%d " + "w2quant=%d w2lds=%d w2mfma=%d w2epi=%d", + moe_sub[0], + moe_sub[1], + moe_sub[2], + moe_sub[3], + moe_sub[4], + moe_sub[5], + moe_sub[6], + moe_sub[7], + moe_epilogue); + if (phase_hist[0] > 0) { + fprintf(stderr, " | stuck_deps:"); + for (int d = 0; d < 256; d++) { + if (dep_hist[d] > 0) { + fprintf(stderr, " ev%d=%d", d, dep_hist[d]); + } + } + } + fprintf(stderr, "\n"); + } // verbose } // Dump event counters to diagnose which events haven't fired if (g_dbg_h_event_counters && dbg_i == 5) { @@ -4400,7 +5600,11 @@ extern "C" void launch_persistent_kernel(cudaStream_t default_stream) { break; } } - (void)hipStreamDestroy(dbg_stream); + dbg_a_stop.store(true, std::memory_order_relaxed); + // Deliberately leaking dbg_stream: the poll thread is detached and may + // still be blocked inside HIP on it. Destroying the stream out from + // under it would turn a diagnosable hang into a use-after-free. This + // path only runs under MPK_WORKER_STATE debug builds. } #endif (void)cudaStreamSynchronize(global_runtime_config.worker_stream); diff --git a/include/mirage/persistent_kernel/runtime_header.h b/include/mirage/persistent_kernel/runtime_header.h index 623920b..c079501 100644 --- a/include/mirage/persistent_kernel/runtime_header.h +++ b/include/mirage/persistent_kernel/runtime_header.h @@ -94,6 +94,13 @@ typedef unsigned long long int EventCounter; int const MAX_INPUTS_PER_TASK = 28; int const MAX_OUTPUTS_PER_TASK = 13; + +// Nil-address tripwire buffer geometry (see MPK_NIL_TRIPWIRE). +// Per worker: [0] layer [1] outer phase [2] tile [3] input_ptrs[0] +// [4] fused-kernel sub-phase [5] sub-phase aux value +#define MPK_TW_HDR 4 +#define MPK_TW_PER_WORKER 8 +#define MPK_TW_SLOTS (MPK_TW_HDR + 256 * MPK_TW_PER_WORKER) // Increased to 304 to support full CU utilization on AMD MI300X (304 CUs) // and NVIDIA Blackwell (160+ SMs which uses 144 workers) int const MAX_NUM_WORKERS = 304; @@ -299,8 +306,20 @@ static_assert( struct alignas(16) TaskDesc { TaskDesc(FullTaskDesc t) : task_type(t.task_type), variant_id(t.variant_id), - trigger_event(t.trigger_event), dependent_event(t.dependent_event), + trigger_event(t.trigger_event), + dependent_event(t.dependent_event), input_ptrs{}, output_ptrs{}, +#ifdef MPK_ENABLE_TMA + input_tma_desc_ptrs{}, output_tma_desc_ptrs{}, +#endif task_metadata(t.task_metadata) { + // The pointer arrays are value-initialized above, not just the first + // num_inputs/num_outputs entries. The whole TaskDesc is memcpy'd to the + // GPU and, for the fused layer path, copied into a *reused* shared-memory + // slot -- so any slot this constructor skips carried indeterminate host + // stack bytes into device memory, and on the device it aliases whatever + // the previous task left behind. Consumers index by slot, not by count + // (see the ml pointer-table build in persistent_kernel.cuh, which snapshots + // all MAX_* slots), so "past num_outputs" is not the same as "never read". for (int i = 0; i < t.num_inputs; i++) { input_ptrs[i] = t.inputs[i].base_ptr; } @@ -320,7 +339,13 @@ struct alignas(16) TaskDesc { } #endif } - __host__ __device__ TaskDesc() { + __host__ __device__ TaskDesc() + : input_ptrs{}, output_ptrs{} +#ifdef MPK_ENABLE_TMA + , + input_tma_desc_ptrs{}, output_tma_desc_ptrs{} +#endif + { task_metadata.raw_payload = ~0ull; } TaskType task_type; @@ -418,6 +443,10 @@ struct RuntimeConfig { unsigned long long *precomp_dbg_tasks_done; // debug: host-mapped task counter int *precomp_dbg_worker_state; // debug: [num_workers*4] = {task_pos, // dep_event, tasks_done, stuck_count} + // Nil-address fault tripwire: pinned host memory mapped for device access, + // so breadcrumbs survive the abort that a memory fault triggers. Null + // unless built with MPK_NIL_TRIPWIRE. See persistent_kernel.cuh. + unsigned long long *tripwire; // Cross-XCD gang barrier: workers sync before executing gang tasks with // internal barriers unsigned long long diff --git a/include/mirage/persistent_kernel/tasks/mi300/gang_full_layer_fused_mi300.cuh b/include/mirage/persistent_kernel/tasks/mi300/gang_full_layer_fused_mi300.cuh index 99ff616..f19dfc5 100644 --- a/include/mirage/persistent_kernel/tasks/mi300/gang_full_layer_fused_mi300.cuh +++ b/include/mirage/persistent_kernel/tasks/mi300/gang_full_layer_fused_mi300.cuh @@ -92,7 +92,8 @@ __device__ __noinline__ void int oproj_tiles_per_xcd, int moe_total_tiles_per_xcd, int workers_per_xcd, - int tile_idx) { + int tile_idx, + int task_layer_idx) { // (All phases enabled — buffer_inv before Phase 6 poll fixes L2 stale read) // input_ptrs layout: // [0] workspace_f32 [1] residual [2] norm_weight_pre @@ -121,17 +122,9 @@ __device__ __noinline__ void // previous iteration) can persist in vL1 across gang task boundaries. asm volatile("buffer_inv" ::: "memory"); - { - constexpr int LAYER_IDX_SMEM_OFF = - mirage::runtime::MAX_DYNAMIC_SHARED_MEMORY_SIZE - - mirage::runtime::LAYER_IDX_SMEM_OFFSET_FROM_END; - extern __shared__ char _layer_smem_init[]; - if (tid == 0) { - int *p = reinterpret_cast(&_layer_smem_init[LAYER_IDX_SMEM_OFF]); - *p = 0; - } - } - __syncthreads(); + // NOTE: the layer counter that the MoE W13->W2 barrier derives its release + // value from is published further down, once qkv_epoch_expected is known. + // See the LAYER_IDX_SMEM_OFF store just before Phase 1. #ifdef MPK_ENABLE_DEVICE_TASK_TIMING unsigned long long _fused_t0 = __builtin_amdgcn_s_memrealtime(); @@ -155,6 +148,8 @@ __device__ __noinline__ void } #endif + MPK_TW_SUB(1, tile_idx); + int *oproj_counters_base = static_cast(input_ptrs[16]); int *attn_global = oproj_counters_base + FULL_LAYER_ATTN_GLOBAL_COUNTER_SLOT; int *qkv_epoch = oproj_counters_base + FULL_LAYER_QKV_EPOCH_SLOT; @@ -162,25 +157,91 @@ __device__ __noinline__ void int *routing_ready = oproj_counters_base + 10 * 16; int *attn_release = oproj_counters_base + FULL_LAYER_ATTN_XCD_RELEASE_SLOT; - // All threads independently compute expected values (ld_nt is uniform). - // Eliminates shared variables and __syncthreads broadcast. - // attn_expected and qkv_expected are only used by tid==0, computed inline at - // use. - int routing_expected = ld_nt_s32(routing_ready) + 1; - int attn_release_expected = ld_nt_s32(&attn_release[xcd_id * 16]) + 1; + // Barrier release values, derived from the layer counter rather than read. + // + // These used to be snapshots of "current value + 1". That is only correct if + // this worker reads *before* this layer's producer bumps the counter, and + // nothing guaranteed it: all 36 layers run inside a single task (see the ml + // loop in persistent_kernel.cuh) with only a per-block __syncthreads between + // them, so workers skew freely across layer boundaries. A worker that fell a + // full layer behind could read an already-bumped counter and then wait for a + // bump this layer never produces -- an intermittent deadlock, because it + // needs a full-layer skew to happen. + // + // f1fa720 fixed that by forcing *every* worker on the XCD to arrive at the + // Phase 2 barrier, which ordered each read before every producer. It worked, + // and it also cost 2.19 -> 2.46 ms/iter. 2c1071c recovered most of that by + // splitting arrival from waiting (all 30 arrive, only ranks < NUM_KV_CHUNKS + // block), reaching 2.39. The remaining ~0.2ms was *not* recovered by this + // change and is still unattributed -- device timing under MPK_DEVICE_TIMING=1 + // cannot localize it, because the ~147k printfs inflate iterations to ~56ms + // and the skew lands in Phase 6's xcd_barrier, whose median then moves + // opposite to real latency. This change is a correctness fix; treat its + // latency effect as neutral. + // + // The layer counter removes the race at its source. All three counters start + // at 0, bump exactly once per layer, and are never reset, so the value this + // layer drives them to is a pure function of the layer index -- no shared + // read, nothing to order, and no arrival requirement. See the + // _linear_reserved store in the ml loop for how it is published. + // + // qkv_epoch[x] bumped once by the last worker to arrive on XCD x + // attn_release[x] written by the last XCD to reach the attn_global + // barrier, using its *own* expected value -- so producer + // and consumer now agree by construction + // routing_ready[*] read-modify-written once by the single TopK completer + // + // Because they are all the same per-layer count, all three expected values + // are the same number. + int const layer_counter = task_layer_idx; + int const routing_expected = layer_counter + 1; + int const attn_release_expected = layer_counter + 1; + int const qkv_epoch_expected = layer_counter + 1; + + // Only the workers that read this layer's QKV output take part in the epoch + // barrier. The expected values above no longer depend on arrival ordering, + // so the participant set is free to be the set that actually needs the + // barrier. This is the pre-f1fa720 participant set, now safe to use because + // nothing reads a shared counter -- it measured neutral, not faster. + int const qkv_epoch_participants = total_qkv_tiles_per_xcd > NUM_KV_CHUNKS + ? total_qkv_tiles_per_xcd + : NUM_KV_CHUNKS; + + // Publish the layer counter the MoE W13->W2 barrier keys off. + // + // That barrier derives its release value as layer_idx + 1 + // (gang_moe_fused_mxfp4_mi300.cuh), and its d_barrier is monotonic -- never + // reset. So layer_idx must be monotonic across the whole run, not per-layer. + // This slot used to be stored as 0 on every entry, which made release_val + // permanently 1: after the very first layer wrote 1, the barrier was already + // satisfied for every later layer, so W2 workers stopped waiting for their + // own layer's W13 and could read swiglu_out before it was written. The tile + // ordering (all W13 tiles precede all W2 tiles, padded so every worker + // starts on W13) usually hid it, but nothing enforced it. + // + // qkv_epoch_expected is exactly the counter needed: it counts + // (iterations * num_layers + layer), monotonic and never reset. It is now a + // pure function of the layer index rather than a snapshot, so every worker + // -- on this XCD or any other -- computes the same value with no barrier + // required to make it agree. That is strictly stronger than what this slot + // relied on before. + { + constexpr int LAYER_IDX_SMEM_OFF = + mirage::runtime::MAX_DYNAMIC_SHARED_MEMORY_SIZE - + mirage::runtime::LAYER_IDX_SMEM_OFFSET_FROM_END; + extern __shared__ char _layer_smem_init[]; + if (tid == 0) { + *reinterpret_cast(&_layer_smem_init[LAYER_IDX_SMEM_OFF]) = + qkv_epoch_expected; + } + } + __syncthreads(); // ══════════════════════════════════════════════════════════════════ // Phase 1: QKV GEMM // ══════════════════════════════════════════════════════════════════ + MPK_TW_SUB(10, xcd_rank); if (xcd_rank < total_qkv_tiles_per_xcd) { - // Snapshot qkv_expected BEFORE QKV GEMM (which may atomicAdd to epoch). - // Only tid==0 uses it, but read is cheap and avoids shared-var broadcast. - int qkv_expected; - if (tid == 0) { - qkv_expected = - __atomic_load_n(&qkv_epoch[xcd_id * 16], __ATOMIC_RELAXED) + 1; - } - gang_resaddf32_rmsnorm_linear_mxfp4_bias_kvupd_kernel 10 only + // 10 workers ever reached the chunk barrier, the + // `(s_chunk_prev % NUM_KV_CHUNKS) == NUM_KV_CHUNKS-1` merge condition never + // fired, and the megakernel deadlocked. Attention chunks are now served by + // any of the workers_per_xcd (30) workers on this XCD. + MPK_TW_SUB(30, xcd_rank); + MPK_WS_PHASE(30, qkv_epoch_expected, xcd_id); + { if (xcd_rank < NUM_KV_CHUNKS) { int kv_chunk_idx = xcd_rank; using bf16_t = __hip_bfloat16; @@ -294,7 +393,49 @@ __device__ __noinline__ void // ══════════════════════════════════════════════════════════════════ // Phase 4: Chunk barrier — CROC last-chunk-worker runs merge // ══════════════════════════════════════════════════════════════════ + MPK_TW_SUB(40, kv_chunk_idx); + MPK_WS_PHASE(40, qkv_epoch_expected, xcd_id); + // Flush this chunk's o_acc/lse_acc partials before arriving. + // + // The chunk kernel writes its partials with ordinary stores, which may + // still be sitting in this CU's write buffer / vL1 when the atomic + // below retires. atom_add_release_gpu_s32 orders *this thread's* prior + // writes, but tid 0 is not the thread that wrote most of this chunk's + // partials -- the other 255 threads did, and __syncthreads is a + // block-execution barrier, not a memory-visibility one to other CUs. + // Without this, the merging worker (a *different* block, possibly a + // different CU) can read a partially-written o_acc/lse_acc slot. + // + // The result is silent numerical corruption, not a crash: merge weights + // each chunk by exp2(lse - m_global), so a stale/torn lse reads as + // garbage magnitude. A torn lse that lands large makes m_global huge and + // drives every other chunk's weight to zero; one that lands as raw + // uninitialized bits can be NaN, which propagates through the whole + // attention output. That is the seq-len-dependent part: at 512 tokens + // ntiles=32 so 14 of 30 chunks exit early via the empty-chunk path and + // never race, while at 32k ntiles=2048 gives every one of the 30 chunks + // real work to write, so the window is open on all of them every layer. + // + // Scope: this barrier is per-XCD (chunk_barrier[xcd_id * 16]), so the + // chunk workers and the merging worker are always on the same XCD and + // share one 32MB L2. Making the partials visible therefore only requires + // getting them *into* L2 -- not flushing L2 to HBM. + // + // db48239 used threadfence_gpu() here, which is agent scope and lowers to + // `buffer_wbl2 sc1; s_waitcnt vmcnt(0)` (verified in gfx950 ISA). The + // buffer_wbl2 is an L2->HBM writeback of the whole cache, paid by every + // chunk worker on every layer, and it buys nothing a same-XCD consumer + // can observe: see the intra-XCD section of mpk_atoms.cuh, "all CUs + // within an XCD share the same 32MB L2 ... No buffer_wbl2 required". + // Dropping it is worth ~0.16 ms/iter at seq 512. + // + // What remains is the part that is actually load-bearing: s_waitcnt + // vmcnt(0) retires all 256 threads' outstanding stores so they have + // reached L2 before tid 0's release atomic. Workgroup-scope fences emit + // no instruction at all on gfx950, so they cannot substitute -- the + // consumer is a different block. __syncthreads(); + asm volatile("s_waitcnt vmcnt(0)" ::: "memory"); __shared__ int s_chunk_prev; if (tid == 0) { s_chunk_prev = atom_add_release_gpu_s32(&chunk_barrier[xcd_id * 16], 1); @@ -303,10 +444,27 @@ __device__ __noinline__ void if ((s_chunk_prev % NUM_KV_CHUNKS) == NUM_KV_CHUNKS - 1) { // Last chunk worker: run merge (no reset needed — modular check) + // + // Acquire the partials the other chunks released above. buffer_inv + // drops this CU's stale vL1 lines so the merge reads what they + // actually wrote rather than a cached copy from a previous layer -- + // vL1 is per-CU, so this is required even though the producers share + // our L2. + // + // Plain `buffer_inv` (no sc1) invalidates vL1 only. The acquire fence + // that used to precede it was agent scope, which emits `buffer_inv + // sc1` and additionally invalidates L2 -- discarding lines this XCD's + // own chunk workers had just written, and forcing them to be re-read + // from HBM. The producers are all on this XCD (the barrier is + // per-XCD), so their stores are already in our L2 and invalidating it + // is both unnecessary and actively harmful. + asm volatile("buffer_inv" ::: "memory"); // ══════════════════════════════════════════════════════════════════ // Phase 5: Merge (write-through fused) + signal // ══════════════════════════════════════════════════════════════════ + MPK_TW_SUB(50, s_chunk_prev); + MPK_WS_PHASE(50, qkv_epoch_expected, xcd_id); // WRITE_THROUGH=true: merge writes bf16 output directly via st_wt, // eliminating the separate __syncthreads + readback + flush pass. merge_splitkv_ck_fmha<__hip_bfloat16, @@ -366,10 +524,11 @@ __device__ __noinline__ void // Issue buffer_load_lds for O-proj weights BEFORE the barrier poll // so DMA runs in the background during the spin-wait (~5-60us). // ══════════════════════════════════════════════════════════════════ + MPK_TW_SUB(60, attn_release_expected); + MPK_WS_PHASE(60, qkv_epoch_expected, xcd_id); int oproj_topk_tiles_per_xcd = oproj_tiles_per_xcd > router_tile_n ? oproj_tiles_per_xcd : router_tile_n; -#if !defined(USE_BF16_NATIVE_WEIGHTS) && !defined(USE_BF16_ACTIVATIONS) { constexpr int OPROJ_WG_DATA = OPROJ_OUTPUT_PER_WG * (OPROJ_REDUCTION_SIZE / 2); @@ -383,13 +542,10 @@ __device__ __noinline__ void constexpr int OPROJ_SLPT = (OPROJ_N16_SCALE + 255) / 256; constexpr int OPROJ_SCALE_PAD = OPROJ_SLPT * 256 * 16; -#ifdef USE_FP4_ACTIVATIONS - constexpr int OPROJ_FP_TOK = OPROJ_REDUCTION_SIZE / 2; - constexpr int OPROJ_FP_SCL = OPROJ_NUM_B32 * 4; -#else + // Activations are FP8 E4M3: one byte per element, one E8M0 scale byte per + // 32-element block. See _gang_wave_parallel_fp8_quant. constexpr int OPROJ_FP_TOK = OPROJ_REDUCTION_SIZE; constexpr int OPROJ_FP_SCL = OPROJ_NUM_B32; -#endif constexpr int OPROJ_LDS_W_OFF = ((OPROJ_FP_TOK + OPROJ_FP_SCL + 15) / 16) * 16; @@ -445,10 +601,17 @@ __device__ __noinline__ void } } } -#endif - while (ld_nt_s32(&attn_release[xcd_id * 16]) < attn_release_expected) { - __builtin_amdgcn_s_sleep(1); + MPK_WS_WAIT_BEGIN(60, attn_release_expected); + { + int _obs; + int _spins = 0; + while ((_obs = ld_nt_s32(&attn_release[xcd_id * 16])) < + attn_release_expected) { + MPK_WS_WAIT_TICK(_obs, _spins); + _spins++; + __builtin_amdgcn_s_sleep(1); + } } asm volatile("buffer_inv" ::: "memory"); asm volatile("s_waitcnt vmcnt(0)" ::: "memory"); @@ -460,6 +623,8 @@ __device__ __noinline__ void // ══════════════════════════════════════════════════════════════════ // Phase 7: O-proj + RMSNorm + Router + TopK // ══════════════════════════════════════════════════════════════════ + MPK_TW_SUB(70, oproj_topk_tiles_per_xcd); + MPK_WS_PHASE(70, qkv_epoch_expected, xcd_id); { if (xcd_rank < oproj_topk_tiles_per_xcd) { int oproj_tile_idx = xcd_id * oproj_topk_tiles_per_xcd + xcd_rank; @@ -496,12 +661,12 @@ __device__ __noinline__ void oproj_topk_tiles_per_xcd, oproj_tile_idx, routing_ready, -#ifdef MPK_ENABLE_DEVICE_TASK_TIMING - _ts_base -#else - nullptr -#endif - ); + // ts_base: the O-proj/TopK kernel's optional per-sub-op timestamp + // sink. Nothing here reads those slots -- the [FUSED_PHASE] printf + // below derives every number from _fused_t0.._fused_t4, which are + // taken in this function. This used to name a _ts_base that was + // never declared, so MPK_DEVICE_TIMING=1 did not compile. + nullptr); } } @@ -512,9 +677,16 @@ __device__ __noinline__ void // signals routing_ready. Other workgroups poll here until TopK // results are globally visible. // All threads poll independently — eliminates __syncthreads overhead. + MPK_TW_SUB(75, routing_expected); + MPK_WS_PHASE(75, qkv_epoch_expected, xcd_id); { int *my_release = &routing_ready[(1 + xcd_id) * 16]; - while (ld_nt_s32(my_release) < routing_expected) { + MPK_WS_WAIT_BEGIN(75, routing_expected); + int _obs; + int _spins = 0; + while ((_obs = ld_nt_s32(my_release)) < routing_expected) { + MPK_WS_WAIT_TICK(_obs, _spins); + _spins++; __builtin_amdgcn_s_sleep(1); } } @@ -529,6 +701,8 @@ __device__ __noinline__ void // ══════════════════════════════════════════════════════════════════ for (int moe_t = xcd_rank; moe_t < moe_total_tiles_per_xcd; moe_t += workers_per_xcd) { + MPK_TW_SUB(80, moe_t); + MPK_WS_PHASE(80, qkv_epoch_expected, xcd_id); gang_moe_fused_mxfp4_kernel_mi300W2 barrier, and W2. The +// nil-address fault lands somewhere in there, and at that resolution a worker +// that faulted is indistinguishable from one merely parked at the barrier. +// +// aux carries the decoded tile identity, because every address this kernel +// computes is derived from it -- expert_id indexes the weight bases and the +// barrier, and the w13/w2 split decides which pointer set is live: +// [15:0] global_tile [23:16] expert_idx [31:24] expert_id +// [39:32] num_activated_experts [40] is_w2 +#define MOE_TW_AUX() \ + (((unsigned long long)(unsigned short)global_tile) | \ + (((unsigned long long)(unsigned char)expert_idx) << 16) | \ + (((unsigned long long)(unsigned char)expert_id) << 24) | \ + (((unsigned long long)(unsigned char)num_activated_experts) << 32) | \ + (((unsigned long long)(is_w2 ? 1 : 0)) << 40)) +#define MOE_DBG_SUBPHASE(code) MPK_TW_SUB((code), MOE_TW_AUX()) +// Pre-decode marker: expert_id/is_w2 do not exist yet, so pass aux explicitly. +// Distinguishes a fault in the routing-mask read itself from one in the +// compute that follows it. +#define MOE_DBG_ENTRY(code, aux) MPK_TW_SUB((code), (aux)) +#else #define MOE_DBG_SUBPHASE(code) ((void)0) +#define MOE_DBG_ENTRY(code, aux) ((void)0) +#endif namespace kernel { +// Per-expert MoE barrier geometry. One 64-byte line per slot (see the layout +// note at the top of this file): 8 per-XCD release flags then the arrival +// counter, so 9 lines used out of 10 reserved per expert. +// [xcd * MOE_BAR_LINE] per-XCD release flag (st_wt, HBM) +// [MOE_BAR_COUNTER_SLOT * ..] global arrival count (atomic, L2) +constexpr int MOE_BAR_LINE = 16; // int32 per cache line +constexpr int MOE_BAR_COUNTER_SLOT = 8; // line index of the arrival counter +constexpr int MOE_BAR_SLOTS = 10; // lines reserved per expert +constexpr int MOE_BAR_STRIDE = MOE_BAR_SLOTS * MOE_BAR_LINE; // ints per expert + template = total_tiles) { + MPK_WS_MARK(8100, global_tile); // exit: past end of tile range return; } @@ -147,6 +202,7 @@ __device__ __noinline__ void gang_moe_fused_mxfp4_kernel_mi300( phase_tile = global_tile % W13_TILES; // Padding tile: expert_idx beyond activated range → skip if (expert_idx >= num_activated_experts) { + MPK_WS_MARK(8101, global_tile); // exit: W13 padding tile return; } } else { @@ -159,15 +215,26 @@ __device__ __noinline__ void gang_moe_fused_mxfp4_kernel_mi300( int tok_idx = phase_tile / n_wgs; int wg_idx = phase_tile % n_wgs; + // Marker 1001: tile decoded, expert_id read. If num_activated_experts or + // expert_id is out of range here, every pointer built below is wild -- + // this is the marker that separates "bad routing input" from "bad compute". + MOE_DBG_ENTRY( + 1001, + ((unsigned long long)(unsigned short)global_tile) | + (((unsigned long long)(unsigned char)expert_idx) << 16) | + (((unsigned long long)(unsigned char)num_activated_experts) << 32) | + (((unsigned long long)(is_w2 ? 1 : 0)) << 40)); int expert_id = d_mask[expert_idx]; int const *expert_routing = d_routing + expert_id * BATCH_SIZE; if (tok_idx >= BATCH_SIZE) { + MPK_WS_MARK(8102, global_tile); // exit: token out of batch return; } int route_val = expert_routing[tok_idx]; if (route_val == 0) { + MPK_WS_MARK(8103, global_tile); // exit: token not routed here return; } int topk_slot = route_val - 1; @@ -181,6 +248,7 @@ __device__ __noinline__ void gang_moe_fused_mxfp4_kernel_mi300( // ══════════════════════════════════════════════════════════════════════════ if (!is_w2) { MOE_DBG_SUBPHASE(2000); + MPK_WS_MARK(8200, global_tile); // W13 compute // Shared memory layout: FP8 quantized tokens + scales uint8_t *s_tok_fp8 = (uint8_t *)_fused_smem; uint8_t *s_tok_scales = s_tok_fp8 + W13_K; @@ -421,60 +489,93 @@ __device__ __noinline__ void gang_moe_fused_mxfp4_kernel_mi300( unsigned ts_addr = (unsigned)(uintptr_t)(s_tok_scales); asm volatile( // Zero accumulator + // ── Two disjoint operand banks ── + // Bank 0: A v[22:25], A scale v7, B v[8:15], B scale v16 + // Bank 1: A v[26:29], A scale v18, B v[32:39], B scale v19 + // Address scratch v17, accumulator a[0:3]. + // + // Prefetching into the registers the current MFMA reads is a WAR + // race: lgkmcnt tracks when LDS data lands in the VGPR, not when + // the MFMA finished sampling its operands, and a 16x16x128 MFMA + // streams them over the op rather than latching at issue. When + // LDS returns fast the write-back lands mid-MFMA and the op sees + // mixed-iteration operands (~17-22% of launches before banking). + // Ping-pong: while the MFMA consumes bank X, prefetch writes bank + // 1-X, so no register is ever both a live source and an in-flight + // LDS destination. + // + // Verified by tests/standalone/test_mfma_pipeline_hazards.hip. "v_accvgpr_write_b32 a0, 0\n" "v_accvgpr_write_b32 a1, 0\n" "v_accvgpr_write_b32 a2, 0\n" "v_accvgpr_write_b32 a3, 0\n" - // Pre-issue 5 reads for iteration 0 - "ds_read_b128 v[22:25], %[wa]\n" // weight A (16B FP4) - "ds_read_u8 v7, %[wsa]\n" // weight A scale - "ds_read_b128 v[8:11], %[ta]\n" // token B lo (16B FP8) - "ds_read_b128 v[12:15], %[ta] offset:64\n" // token B hi (16B FP8) - "ds_read_u8 v16, %[tsa]\n" // token B scale - "s_mov_b32 s13, 0\n" // loop counter + // Pre-issue 5 reads for iteration 0 into bank 0 + "ds_read_b128 v[22:25], %[wa]\n" + "ds_read_u8 v7, %[wsa]\n" + "ds_read_b128 v[8:11], %[ta]\n" + "ds_read_b128 v[12:15], %[ta] offset:64\n" + "ds_read_u8 v16, %[tsa]\n" + "s_mov_b32 s13, 0\n" - // ── Iterations 0..22: prefetch next, MFMA current ── "PIPELINED_W13_T0_%=:\n" - "s_waitcnt lgkmcnt(0)\n" // iter N reads complete - - // Advance addresses for iter N+1 + // ---- consume bank 0, prefetch into bank 1 ---- + "s_waitcnt lgkmcnt(0)\n" "v_add_u32_e32 %[wa], 64, %[wa]\n" "v_add_u32_e32 %[wsa], 4, %[wsa]\n" "v_add_u32_e32 %[ta], 0x80, %[ta]\n" "s_add_i32 s13, s13, 1\n" - - // Token B scale → v17 FIRST (oldest in queue, completes first) "v_add_u32_e32 v17, s13, %[tsa]\n" - "ds_read_u8 v17, v17\n" // [lgkmcnt +1] oldest - // Prefetch iter N+1 data into SAME regs (MFMA reads old values) - "ds_read_b128 v[22:25], %[wa]\n" // [lgkmcnt +2] - "ds_read_u8 v7, %[wsa]\n" // [lgkmcnt +3] - "ds_read_b128 v[8:11], %[ta]\n" // [lgkmcnt +4] - "ds_read_b128 v[12:15], %[ta] offset:64\n" // [lgkmcnt +5] - - // MFMA from iter N data (32 cycles, reads v[22:25] v[8:15] v7 - // v16) + "ds_read_u8 v19, v17\n" + "ds_read_b128 v[26:29], %[wa]\n" + "ds_read_u8 v18, %[wsa]\n" + "ds_read_b128 v[32:35], %[ta]\n" + "ds_read_b128 v[36:39], %[ta] offset:64\n" "v_mfma_scale_f32_16x16x128_f8f6f4 a[0:3], v[22:25], v[8:15], " "a[0:3], v7, v16 op_sel_hi:[0,0,0] cbsz:4\n" + "s_cmpk_lt_i32 s13, %[iters_m1]\n" + "s_cbranch_scc0 W13_T0_TAIL_B1_%=\n" - // Copy next token scale during MFMA execution. - // lgkmcnt(4): wait for token scale (oldest, issued first), leave - // 4 data reads flying - "s_waitcnt lgkmcnt(4)\n" - "v_mov_b32_e32 v16, v17\n" - + // ---- consume bank 1, prefetch into bank 0 ---- + "s_waitcnt lgkmcnt(0)\n" + "v_add_u32_e32 %[wa], 64, %[wa]\n" + "v_add_u32_e32 %[wsa], 4, %[wsa]\n" + "v_add_u32_e32 %[ta], 0x80, %[ta]\n" + "s_add_i32 s13, s13, 1\n" + "v_add_u32_e32 v17, s13, %[tsa]\n" + "ds_read_u8 v16, v17\n" + "ds_read_b128 v[22:25], %[wa]\n" + "ds_read_u8 v7, %[wsa]\n" + "ds_read_b128 v[8:11], %[ta]\n" + "ds_read_b128 v[12:15], %[ta] offset:64\n" + "v_mfma_scale_f32_16x16x128_f8f6f4 a[0:3], v[26:29], v[32:39], " + "a[0:3], v18, v19 op_sel_hi:[0,0,0] cbsz:4\n" "s_cmpk_lt_i32 s13, %[iters_m1]\n" "s_cbranch_scc1 PIPELINED_W13_T0_%=\n" - // ── Final iteration: no more prefetch needed ── + // ── Final MFMA ── + // Both tails are emitted because exit parity decides which bank + // holds the final operands. Live MFMA_ITERS is 23 (odd), so the + // loop falls out of the bank 1 half with the last operands in + // BANK 0 -- this path. An even count exits via W13_T0_TAIL_B1 + // with them in bank 1. One tail alone would silently use the + // wrong bank for one parity. "s_waitcnt lgkmcnt(0)\n" "v_mfma_scale_f32_16x16x128_f8f6f4 a[0:3], v[22:25], v[8:15], " "a[0:3], v7, v16 op_sel_hi:[0,0,0] cbsz:4\n" + "s_branch W13_T0_ACC_%=\n" - // Read accumulator into output - "s_nop 7\n" - "s_nop 0\n" + "W13_T0_TAIL_B1_%=:\n" + "s_waitcnt lgkmcnt(0)\n" + "v_mfma_scale_f32_16x16x128_f8f6f4 a[0:3], v[26:29], v[32:39], " + "a[0:3], v18, v19 op_sel_hi:[0,0,0] cbsz:4\n" + + "W13_T0_ACC_%=:\n" + // 32 clocks: the scaled MFMA is a 32-cycle op on CDNA4. The old + // "s_nop 7; s_nop 0" was 9 clocks (correct only for a 4-pass + // MFMA) and returned a partially-retired accumulator every time. + "s_nop 15\n" + "s_nop 15\n" "v_accvgpr_read_b32 %[acc0], a0\n" "v_accvgpr_read_b32 %[acc1], a1\n" "v_accvgpr_read_b32 %[acc2], a2\n" @@ -500,10 +601,24 @@ __device__ __noinline__ void gang_moe_fused_mxfp4_kernel_mi300( "v15", "v16", "v17", + "v18", + "v19", "v22", "v23", "v24", "v25", + "v26", + "v27", + "v28", + "v29", + "v32", + "v33", + "v34", + "v35", + "v36", + "v37", + "v38", + "v39", "a0", "a1", "a2", @@ -724,39 +839,94 @@ __device__ __noinline__ void gang_moe_fused_mxfp4_kernel_mi300( unsigned ts_addr = (unsigned)(uintptr_t)(s_tok_scales); asm volatile( + // ── Two disjoint operand banks ── + // Bank 0: A v[22:25], A scale v7, B v[8:15], B scale v16 + // Bank 1: A v[26:29], A scale v18, B v[32:39], B scale v19 + // Address scratch v17, accumulator a[0:3]. + // + // Prefetching into the registers the current MFMA reads is a + // WAR race: lgkmcnt tracks when LDS data lands in the VGPR, not + // when the MFMA finished sampling its operands, and a 16x16x128 + // MFMA streams them over the op rather than latching at issue. + // When LDS returns fast the write-back lands mid-MFMA and the + // op sees mixed-iteration operands (~17-22% of launches before + // banking). Ping-pong: while the MFMA consumes bank X, prefetch + // writes bank 1-X, so no register is ever both a live source + // and an in-flight LDS destination. + // + // Verified by tests/standalone/test_mfma_pipeline_hazards.hip. "v_accvgpr_write_b32 a0, 0\n" "v_accvgpr_write_b32 a1, 0\n" "v_accvgpr_write_b32 a2, 0\n" "v_accvgpr_write_b32 a3, 0\n" + + // Pre-issue 5 reads for iteration 0 into bank 0 "ds_read_b128 v[22:25], %[wa]\n" "ds_read_u8 v7, %[wsa]\n" "ds_read_b128 v[8:11], %[ta]\n" "ds_read_b128 v[12:15], %[ta] offset:64\n" "ds_read_u8 v16, %[tsa]\n" "s_mov_b32 s13, 0\n" + "PIPELINED_W13_T1_%=:\n" + // ---- consume bank 0, prefetch into bank 1 ---- "s_waitcnt lgkmcnt(0)\n" "v_add_u32_e32 %[wa], 64, %[wa]\n" "v_add_u32_e32 %[wsa], 4, %[wsa]\n" "v_add_u32_e32 %[ta], 0x80, %[ta]\n" "s_add_i32 s13, s13, 1\n" "v_add_u32_e32 v17, s13, %[tsa]\n" - "ds_read_u8 v17, v17\n" + "ds_read_u8 v19, v17\n" + "ds_read_b128 v[26:29], %[wa]\n" + "ds_read_u8 v18, %[wsa]\n" + "ds_read_b128 v[32:35], %[ta]\n" + "ds_read_b128 v[36:39], %[ta] offset:64\n" + "v_mfma_scale_f32_16x16x128_f8f6f4 a[0:3], v[22:25], v[8:15], " + "a[0:3], v7, v16 op_sel_hi:[0,0,0] cbsz:4\n" + "s_cmpk_lt_i32 s13, %[iters_m1]\n" + "s_cbranch_scc0 W13_T1_TAIL_B1_%=\n" + + // ---- consume bank 1, prefetch into bank 0 ---- + "s_waitcnt lgkmcnt(0)\n" + "v_add_u32_e32 %[wa], 64, %[wa]\n" + "v_add_u32_e32 %[wsa], 4, %[wsa]\n" + "v_add_u32_e32 %[ta], 0x80, %[ta]\n" + "s_add_i32 s13, s13, 1\n" + "v_add_u32_e32 v17, s13, %[tsa]\n" + "ds_read_u8 v16, v17\n" "ds_read_b128 v[22:25], %[wa]\n" "ds_read_u8 v7, %[wsa]\n" "ds_read_b128 v[8:11], %[ta]\n" "ds_read_b128 v[12:15], %[ta] offset:64\n" - "v_mfma_scale_f32_16x16x128_f8f6f4 a[0:3], v[22:25], v[8:15], " - "a[0:3], v7, v16 op_sel_hi:[0,0,0] cbsz:4\n" - "s_waitcnt lgkmcnt(4)\n" - "v_mov_b32_e32 v16, v17\n" + "v_mfma_scale_f32_16x16x128_f8f6f4 a[0:3], v[26:29], v[32:39], " + "a[0:3], v18, v19 op_sel_hi:[0,0,0] cbsz:4\n" "s_cmpk_lt_i32 s13, %[iters_m1]\n" "s_cbranch_scc1 PIPELINED_W13_T1_%=\n" + + // ── Final MFMA ── + // Both tails are emitted because exit parity decides which bank + // holds the final operands. Live MFMA_ITERS is 23 (odd), so the + // loop falls out of the bank 1 half with the last operands in + // BANK 0 -- this path. An even count exits via W13_T1_TAIL_B1 + // with them in bank 1. One tail alone would silently use the + // wrong bank for one parity. "s_waitcnt lgkmcnt(0)\n" "v_mfma_scale_f32_16x16x128_f8f6f4 a[0:3], v[22:25], v[8:15], " "a[0:3], v7, v16 op_sel_hi:[0,0,0] cbsz:4\n" - "s_nop 7\n" - "s_nop 0\n" + "s_branch W13_T1_ACC_%=\n" + + "W13_T1_TAIL_B1_%=:\n" + "s_waitcnt lgkmcnt(0)\n" + "v_mfma_scale_f32_16x16x128_f8f6f4 a[0:3], v[26:29], v[32:39], " + "a[0:3], v18, v19 op_sel_hi:[0,0,0] cbsz:4\n" + + "W13_T1_ACC_%=:\n" + // 32 clocks: the scaled MFMA is a 32-cycle op on CDNA4. The old + // "s_nop 7; s_nop 0" was 9 clocks (correct only for a 4-pass + // MFMA) and returned a partially-retired accumulator every + // time. + "s_nop 15\n" + "s_nop 15\n" "v_accvgpr_read_b32 %[acc0], a0\n" "v_accvgpr_read_b32 %[acc1], a1\n" "v_accvgpr_read_b32 %[acc2], a2\n" @@ -782,10 +952,24 @@ __device__ __noinline__ void gang_moe_fused_mxfp4_kernel_mi300( "v15", "v16", "v17", + "v18", + "v19", "v22", "v23", "v24", "v25", + "v26", + "v27", + "v28", + "v29", + "v32", + "v33", + "v34", + "v35", + "v36", + "v37", + "v38", + "v39", "a0", "a1", "a2", @@ -1051,11 +1235,24 @@ __device__ __noinline__ void gang_moe_fused_mxfp4_kernel_mi300( // ──────────────────────────────── Uses layer index from shared memory for // monotonically increasing release MOE_DBG_SUBPHASE(2001); - constexpr int HIER_STRIDE = 16; + MPK_WS_MARK(8201, global_tile); // W13 done, arriving at barrier if (tid == 0) { - int base = expert_id * HIER_STRIDE; - // Single global arrival (all W13 tiles increment one counter) - int prev_global = atom_add_release_gpu_s32(&d_barrier[base + 8], 1); + int base = expert_id * MOE_BAR_STRIDE; + // Single global arrival (all W13 tiles increment one counter). + // + // The counter must NOT share a cache line with the release slots below. + // The release fan-out uses st_wt (sc0 sc1), which bypasses L2 and writes + // straight to HBM, while this atomic is an L2-resident read-modify-write + // on the same 64-byte line. When both are in flight on the same line the + // L2 copy -- still holding the *old* release values -- is written back + // over the fresh write-through data, silently reverting slots that were + // already released. The captured deadlock is exactly that: every producer + // had arrived (arrivals % W13_TILES == 0, so the release did fire) yet + // the per-XCD slots of one expert held *different* epochs, which is + // impossible if the eight stores from one producer all survived. + // COUNTER_OFF puts the counter on the next line. + int prev_global = atom_add_release_gpu_s32( + &d_barrier[base + MOE_BAR_COUNTER_SLOT * MOE_BAR_LINE], 1); if ((prev_global % W13_TILES) == W13_TILES - 1) { // Last W13 arrival: write per-XCD release = layer_idx + 1 constexpr int LAYER_IDX_SMEM_OFF = @@ -1065,7 +1262,8 @@ __device__ __noinline__ void gang_moe_fused_mxfp4_kernel_mi300( *reinterpret_cast(&_fused_smem[LAYER_IDX_SMEM_OFF]); int release_val = layer_idx + 1; for (int x = 0; x < 8; x++) { - st_wt_u32((void *)&d_barrier[base + x], (unsigned)release_val); + st_wt_u32((void *)&d_barrier[base + x * MOE_BAR_LINE], + (unsigned)release_val); } asm volatile("s_waitcnt vmcnt(0)" ::: "memory"); } @@ -1089,6 +1287,7 @@ __device__ __noinline__ void gang_moe_fused_mxfp4_kernel_mi300( uint8_t *s_tok_scales = s_tok_fp8 + W2_K; MOE_DBG_SUBPHASE(3000); + MPK_WS_MARK(8300, global_tile); // W2 entry // Weight pointers — depend only on expert_id/wg_idx, available before barrier uint8_t const *expert_weight = W_down + static_cast(expert_id) * W2_EXPERT_BYTES; @@ -1124,8 +1323,7 @@ __device__ __noinline__ void gang_moe_fused_mxfp4_kernel_mi300( MOE_DBG_SUBPHASE(3001); // All threads independently read layer_idx from LDS (uniform value). // Eliminates shared variable and __syncthreads broadcast. - constexpr int HIER_STRIDE = 16; - int base = expert_id * HIER_STRIDE; + int base = expert_id * MOE_BAR_STRIDE; int w2_expected; { constexpr int LAYER_IDX_SMEM_OFF = @@ -1275,17 +1473,62 @@ __device__ __noinline__ void gang_moe_fused_mxfp4_kernel_mi300( // Eliminates tid==0 + __syncthreads — each thread confirms barrier itself. { int expected = w2_expected; - while (ld_nt_s32(&d_barrier[base + xcd_id]) < expected) { + // Barrier id encodes the expert so the dump can tell which of the 4 + // activated experts never got its W13 release (see MPK_WS_WAIT_BEGIN). + MPK_WS_WAIT_BEGIN(800 + expert_idx, expected); + // Each wave clears its own bits in both masks before it starts spinning, + // so what the dump reads describes this poll and not an earlier one. No + // __syncthreads here on purpose -- this poll is deliberately divergent. + MPK_WS_WAVE_CLEAR(warp_id); + int _obs; + int _spins = 0; + while ((_obs = ld_nt_s32(&d_barrier[base + xcd_id * MOE_BAR_LINE])) < + expected) { + MPK_WS_WAIT_TICK(_obs, _spins); + // Refresh the discriminating values on the same cadence as the tick: + // the raw arrival counter (whether it sits on a multiple of W13_TILES + // separates "release fired but was lost" from "arrivals never landed"), + // and how many of the 8 per-XCD slots agree. All 8 are written by one + // producer in one loop, so any spread means releases are being lost. + if ((_spins & (MPK_WS_WAIT_REFRESH - 1)) == 0) { + int _n_ok = 0, _mn = 0x7fffffff, _mx = -0x7fffffff; + for (int _x = 0; _x < 8; _x++) { + int _v = ld_nt_s32(&d_barrier[base + _x * MOE_BAR_LINE]); + if (_v >= expected) { + _n_ok++; + } + if (_v < _mn) { + _mn = _v; + } + if (_v > _mx) { + _mx = _v; + } + } + // a3 is now the per-wave exit mask (MPK_WS_WAVE_EXIT), so fold _mn + // into a2 instead of overwriting it. + MPK_WS_WAIT_AUX( + ld_nt_s32(&d_barrier[base + MOE_BAR_COUNTER_SLOT * MOE_BAR_LINE]), + expert_id, + _n_ok * 1000000 + (_mx - _mn), + -1); + } + _spins++; __builtin_amdgcn_s_sleep(1); } + // This wave's threads all cleared the release. Record it: the poll is + // per-thread with no __syncthreads, so waves leave independently and a + // block can be split across the barrier. + MPK_WS_WAVE_EXIT(warp_id); } MOE_DBG_SUBPHASE(3002); + MPK_WS_MARK(8302, global_tile); // W2: cleared W13->W2 barrier // No buffer_inv needed — NT loads bypass L2 entirely. // FP8 quant of SwiGLU output — writes to LDS[0..W2_K+scales] // buffer_load_lds writes to LDS[W2_OFF..] — no conflict, both in flight. MOE_DBG_SUBPHASE(3003); + MPK_WS_MARK(8303, global_tile); // W2: FP8 quant of SwiGLU output { unsigned short const *w2_input_base = d_swiglu_out + tok_idx * (NUM_TOPK * INTERMEDIATE_SIZE) + @@ -1297,6 +1540,7 @@ __device__ __noinline__ void gang_moe_fused_mxfp4_kernel_mi300( // Drain ALL pending HBM loads: buffer_load_lds (weight) + scale loads // Weight loads were issued before barrier poll, should be done by now. MOE_DBG_SUBPHASE(3004); + MPK_WS_MARK(8304, global_tile); // W2: drain HBM loads asm volatile("s_waitcnt vmcnt(0)" ::: "memory"); asm volatile("s_waitcnt lgkmcnt(0)" ::: "memory"); { @@ -1320,6 +1564,7 @@ __device__ __noinline__ void gang_moe_fused_mxfp4_kernel_mi300( g_subphase_scratch[6] = __builtin_amdgcn_s_memrealtime(); #endif MOE_DBG_SUBPHASE(3005); + MPK_WS_MARK(8305, global_tile); // W2: MFMA loop // LDS-based MFMA loop: weights already in LDS, compiler pipelines ds_reads. // Assembly shows lgkmcnt(7)/lgkmcnt(1) interleaving — much better than // the HBM path's vmcnt(0) stalls before every MFMA group. @@ -1359,7 +1604,7 @@ __device__ __noinline__ void gang_moe_fused_mxfp4_kernel_mi300( &d_w2_bias[expert_id * W2_OUTPUT_SIZE + out_n_base]; asm volatile("global_load_dword %0, %2, off\n" "global_load_dwordx2 %1, %3, off" - : "=v"(pf_rw), "=v"(pf_bias) + : "=&v"(pf_rw), "=&v"(pf_bias) : "v"(rw_ptr), "v"(bias_ptr) : "memory"); } @@ -1381,59 +1626,93 @@ __device__ __noinline__ void gang_moe_fused_mxfp4_kernel_mi300( unsigned w2_ts_addr = (unsigned)(uintptr_t)(s_tok_scales); asm volatile( // Zero accumulator + // ── Two disjoint operand banks ── + // Bank 0: A v[22:25], A scale v7, B v[8:15], B scale v16 + // Bank 1: A v[26:29], A scale v18, B v[32:39], B scale v19 + // Address scratch v17, accumulator a[0:3]. + // + // Prefetching into the registers the current MFMA reads is a WAR + // race: lgkmcnt tracks when LDS data lands in the VGPR, not when + // the MFMA finished sampling its operands, and a 16x16x128 MFMA + // streams them over the op rather than latching at issue. When LDS + // returns fast the write-back lands mid-MFMA and the op sees + // mixed-iteration operands (~17-22% of launches before banking). + // Ping-pong: while the MFMA consumes bank X, prefetch writes bank + // 1-X, so no register is ever both a live source and an in-flight + // LDS destination. + // + // Verified by tests/standalone/test_mfma_pipeline_hazards.hip. "v_accvgpr_write_b32 a0, 0\n" "v_accvgpr_write_b32 a1, 0\n" "v_accvgpr_write_b32 a2, 0\n" "v_accvgpr_write_b32 a3, 0\n" - // Pre-issue 5 reads for iteration 0 - "ds_read_b128 v[22:25], %[wa]\n" // weight A (16B FP4) - "ds_read_u8 v7, %[wsa]\n" // weight A scale - "ds_read_b128 v[8:11], %[ta]\n" // token B lo (16B FP8) - "ds_read_b128 v[12:15], %[ta] offset:64\n" // token B hi (16B FP8) - "ds_read_u8 v16, %[tsa]\n" // token B scale - "s_mov_b32 s13, 0\n" // loop counter + // Pre-issue 5 reads for iteration 0 into bank 0 + "ds_read_b128 v[22:25], %[wa]\n" + "ds_read_u8 v7, %[wsa]\n" + "ds_read_b128 v[8:11], %[ta]\n" + "ds_read_b128 v[12:15], %[ta] offset:64\n" + "ds_read_u8 v16, %[tsa]\n" + "s_mov_b32 s13, 0\n" - // ── Iterations 0..W2_MFMA_ITERS-2: prefetch next, MFMA current ── "PIPELINED_W2_T0_%=:\n" - "s_waitcnt lgkmcnt(0)\n" // iter N reads complete - - // Advance addresses for iter N+1 + // ---- consume bank 0, prefetch into bank 1 ---- + "s_waitcnt lgkmcnt(0)\n" "v_add_u32_e32 %[wa], 64, %[wa]\n" "v_add_u32_e32 %[wsa], 4, %[wsa]\n" "v_add_u32_e32 %[ta], 0x80, %[ta]\n" "s_add_i32 s13, s13, 1\n" - - // Token B scale → v17 FIRST (oldest in queue, completes first) "v_add_u32_e32 v17, s13, %[tsa]\n" - "ds_read_u8 v17, v17\n" // [lgkmcnt +1] oldest - // Prefetch iter N+1 data into SAME regs (MFMA reads old values) - "ds_read_b128 v[22:25], %[wa]\n" // [lgkmcnt +2] - "ds_read_u8 v7, %[wsa]\n" // [lgkmcnt +3] - "ds_read_b128 v[8:11], %[ta]\n" // [lgkmcnt +4] - "ds_read_b128 v[12:15], %[ta] offset:64\n" // [lgkmcnt +5] - - // MFMA from iter N data (32 cycles, reads v[22:25] v[8:15] v7 v16) + "ds_read_u8 v19, v17\n" + "ds_read_b128 v[26:29], %[wa]\n" + "ds_read_u8 v18, %[wsa]\n" + "ds_read_b128 v[32:35], %[ta]\n" + "ds_read_b128 v[36:39], %[ta] offset:64\n" "v_mfma_scale_f32_16x16x128_f8f6f4 a[0:3], v[22:25], v[8:15], " "a[0:3], v7, v16 op_sel_hi:[0,0,0] cbsz:4\n" + "s_cmpk_lt_i32 s13, %[iters_m1]\n" + "s_cbranch_scc0 W2_T0_TAIL_B1_%=\n" - // Copy next token scale during MFMA execution. - // lgkmcnt(4): wait for token scale (oldest, issued first), leave 4 - // data reads flying - "s_waitcnt lgkmcnt(4)\n" - "v_mov_b32_e32 v16, v17\n" - + // ---- consume bank 1, prefetch into bank 0 ---- + "s_waitcnt lgkmcnt(0)\n" + "v_add_u32_e32 %[wa], 64, %[wa]\n" + "v_add_u32_e32 %[wsa], 4, %[wsa]\n" + "v_add_u32_e32 %[ta], 0x80, %[ta]\n" + "s_add_i32 s13, s13, 1\n" + "v_add_u32_e32 v17, s13, %[tsa]\n" + "ds_read_u8 v16, v17\n" + "ds_read_b128 v[22:25], %[wa]\n" + "ds_read_u8 v7, %[wsa]\n" + "ds_read_b128 v[8:11], %[ta]\n" + "ds_read_b128 v[12:15], %[ta] offset:64\n" + "v_mfma_scale_f32_16x16x128_f8f6f4 a[0:3], v[26:29], v[32:39], " + "a[0:3], v18, v19 op_sel_hi:[0,0,0] cbsz:4\n" "s_cmpk_lt_i32 s13, %[iters_m1]\n" "s_cbranch_scc1 PIPELINED_W2_T0_%=\n" - // ── Final iteration: no more prefetch needed ── + // ── Final MFMA ── + // Both tails are emitted because exit parity decides which bank + // holds the final operands. Live MFMA_ITERS is 23 (odd), so the + // loop falls out of the bank 1 half with the last operands in + // BANK 0 -- this path. An even count exits via W2_T0_TAIL_B1 + // with them in bank 1. One tail alone would silently use the + // wrong bank for one parity. "s_waitcnt lgkmcnt(0)\n" "v_mfma_scale_f32_16x16x128_f8f6f4 a[0:3], v[22:25], v[8:15], " "a[0:3], v7, v16 op_sel_hi:[0,0,0] cbsz:4\n" + "s_branch W2_T0_ACC_%=\n" - // Read accumulator into output - "s_nop 7\n" - "s_nop 0\n" + "W2_T0_TAIL_B1_%=:\n" + "s_waitcnt lgkmcnt(0)\n" + "v_mfma_scale_f32_16x16x128_f8f6f4 a[0:3], v[26:29], v[32:39], " + "a[0:3], v18, v19 op_sel_hi:[0,0,0] cbsz:4\n" + + "W2_T0_ACC_%=:\n" + // 32 clocks: the scaled MFMA is a 32-cycle op on CDNA4. The old + // "s_nop 7; s_nop 0" was 9 clocks (correct only for a 4-pass + // MFMA) and returned a partially-retired accumulator every time. + "s_nop 15\n" + "s_nop 15\n" "v_accvgpr_read_b32 %[acc0], a0\n" "v_accvgpr_read_b32 %[acc1], a1\n" "v_accvgpr_read_b32 %[acc2], a2\n" @@ -1459,10 +1738,24 @@ __device__ __noinline__ void gang_moe_fused_mxfp4_kernel_mi300( "v15", "v16", "v17", + "v18", + "v19", "v22", "v23", "v24", "v25", + "v26", + "v27", + "v28", + "v29", + "v32", + "v33", + "v34", + "v35", + "v36", + "v37", + "v38", + "v39", "a0", "a1", "a2", @@ -1470,6 +1763,7 @@ __device__ __noinline__ void gang_moe_fused_mxfp4_kernel_mi300( } MOE_DBG_SUBPHASE(3006); + MPK_WS_MARK(8306, global_tile); // W2: epilogue asm volatile("s_waitcnt vmcnt(0)" ::: "memory"); if (col == 0 && out_n_base < W2_OUTPUT_SIZE) { unsigned bt0 = (pf_bias.x & 0xFFFFu) << 16; @@ -1681,7 +1975,7 @@ __device__ __noinline__ void gang_moe_fused_mxfp4_kernel_mi300( &d_w2_bias[expert_id * W2_OUTPUT_SIZE + out_n_base]; asm volatile("global_load_dword %0, %2, off\n" "global_load_dwordx2 %1, %3, off" - : "=v"(pf_rw), "=v"(pf_bias) + : "=&v"(pf_rw), "=&v"(pf_bias) : "v"(rw_ptr), "v"(bias_ptr) : "memory"); } diff --git a/include/mirage/persistent_kernel/tasks/mi300/gang_moe_linear_mxfp4_mi300.cuh b/include/mirage/persistent_kernel/tasks/mi300/gang_moe_linear_mxfp4_mi300.cuh index 878a24e..7003fc5 100644 --- a/include/mirage/persistent_kernel/tasks/mi300/gang_moe_linear_mxfp4_mi300.cuh +++ b/include/mirage/persistent_kernel/tasks/mi300/gang_moe_linear_mxfp4_mi300.cuh @@ -194,11 +194,18 @@ __device__ __forceinline__ void // Combine amaxes from 4 sub-blocks sharing the same 128-element // super-block. Threads sb, sb+1, sb+2, sb+3 are consecutive lanes in the // same wave. Use __shfl to read each neighbor's amax (4 reads, 3 fmaxf). + // + // The partner index is clamped for the same reason as in the NT variant + // below: NSUBBLOCKS = REDUCTION_SIZE/32 need not be a multiple of 4, so + // the tail super-block would otherwise reduce against lanes whose loop + // condition failed and whose `amax` register was never written. int base_lane = lane_id & ~3; // round down to group of 4 + int const sb_first = sb - sub_idx; + int const n_valid = min(4, (NSUBBLOCKS - 1) - sb_first + 1); float a0 = __shfl(amax, base_lane); - float a1 = __shfl(amax, base_lane + 1); - float a2 = __shfl(amax, base_lane + 2); - float a3 = __shfl(amax, base_lane + 3); + float a1 = __shfl(amax, base_lane + min(1, n_valid - 1)); + float a2 = __shfl(amax, base_lane + min(2, n_valid - 1)); + float a3 = __shfl(amax, base_lane + min(3, n_valid - 1)); float block_amax = fmaxf(fmaxf(a0, a1), fmaxf(a2, a3)); // Compute E8M0 scale @@ -255,15 +262,32 @@ __device__ __forceinline__ void _gang_wave_parallel_fp8_quant_nt( uint32_t const *base_ptr = src32 + base / 2; // 4 wide NT loads (64 bytes = 32 bf16) + // + // The outputs MUST be early-clobber ("=&v"). This is one asm block with + // four separate instructions, so the compiler is free to allocate an + // output register on top of an input it believes is dead after the + // block -- and it does: without the '&' it emits + // global_load_dwordx4 v[4:7], v[4:5], ... + // global_load_dwordx4 v[8:11], v[6:7], ... + // global_load_dwordx4 v[12:15], v[8:9], ... + // global_load_dwordx4 v[16:19], v[10:11],... + // where the first load's destination overwrites the address operands of + // the next three before they issue. Those loads then use whatever the + // returned data happened to be as an address. A wild address that never + // completes leaves the wave parked on the s_waitcnt vmcnt(0) below + // forever, which hangs the __syncthreads at the end of this function and + // through it the whole block -- the captured deadlock is exactly that: + // wave 0 missing from the quant sync mask (0xe) while every wave had + // already cleared the W13->W2 barrier (0xf). uint32_t dw[16]; asm volatile("global_load_dwordx4 %0, %4, off sc0 sc1 nt\n" "global_load_dwordx4 %1, %5, off sc0 sc1 nt\n" "global_load_dwordx4 %2, %6, off sc0 sc1 nt\n" "global_load_dwordx4 %3, %7, off sc0 sc1 nt" - : "=v"(*(i32x4_t *)&dw[0]), - "=v"(*(i32x4_t *)&dw[4]), - "=v"(*(i32x4_t *)&dw[8]), - "=v"(*(i32x4_t *)&dw[12]) + : "=&v"(*(i32x4_t *)&dw[0]), + "=&v"(*(i32x4_t *)&dw[4]), + "=&v"(*(i32x4_t *)&dw[8]), + "=&v"(*(i32x4_t *)&dw[12]) : "v"(base_ptr), "v"(base_ptr + 4), "v"(base_ptr + 8), @@ -283,12 +307,23 @@ __device__ __forceinline__ void _gang_wave_parallel_fp8_quant_nt( amax = fmaxf(amax, fmaxf(fabsf(lo), fabsf(hi))); } - // Combine amaxes from 4 sub-blocks via shuffle + // Combine amaxes across the 4 sub-blocks of this 128-element super-block. + // + // NSUBBLOCKS is REDUCTION_SIZE/32 and is NOT guaranteed to be a multiple + // of 4: for the W2 path REDUCTION_SIZE = INTERMEDIATE_SIZE = 2880, giving + // NSUBBLOCKS = 90. The last super-block therefore has only 2 real + // sub-blocks (sb 88, 89), but the shuffles below still read lanes for + // sb 90 and 91 -- threads whose loop condition failed, so their `amax` + // is an uninitialized register. Clamping the partner index to the last + // valid sub-block makes the reduction read only lanes that ran. int base_lane = lane_id & ~3; + int const sb_first = sb - sub_idx; // first sb of this super-block + int const sb_last = NSUBBLOCKS - 1; // last sb that actually runs + int const n_valid = min(4, sb_last - sb_first + 1); float a0 = __shfl(amax, base_lane); - float a1 = __shfl(amax, base_lane + 1); - float a2 = __shfl(amax, base_lane + 2); - float a3 = __shfl(amax, base_lane + 3); + float a1 = __shfl(amax, base_lane + min(1, n_valid - 1)); + float a2 = __shfl(amax, base_lane + min(2, n_valid - 1)); + float a3 = __shfl(amax, base_lane + min(3, n_valid - 1)); float block_amax = fmaxf(fmaxf(a0, a1), fmaxf(a2, a3)); // Compute E8M0 scale @@ -321,6 +356,11 @@ __device__ __forceinline__ void _gang_wave_parallel_fp8_quant_nt( s_tok_scales[super_blk] = se; } } + // Record that this wave finished its strided share and is entering the + // block-wide sync. If a stall shows a mask below 0xf here, the missing + // wave never got out of the loop above -- almost certainly stuck on the + // vmcnt(0) that drains its NT loads of d_swiglu_out. + MPK_WS_WAVE_SYNC(tid >> 6); __syncthreads(); } @@ -457,15 +497,18 @@ __device__ __forceinline__ void uint32_t const *base_ptr = src32 + base / 2; // 4 wide NT loads (64 bytes = 32 bf16) instead of 16 individual dword loads + // Early-clobber outputs are required here for the same reason as in + // _gang_wave_parallel_fp8_quant_nt: without '&' the allocator puts the + // first load's destination on top of the later loads' address registers. uint32_t dw[16]; asm volatile("global_load_dwordx4 %0, %4, off sc0 sc1 nt\n" "global_load_dwordx4 %1, %5, off sc0 sc1 nt\n" "global_load_dwordx4 %2, %6, off sc0 sc1 nt\n" "global_load_dwordx4 %3, %7, off sc0 sc1 nt" - : "=v"(*(i32x4_t *)&dw[0]), - "=v"(*(i32x4_t *)&dw[4]), - "=v"(*(i32x4_t *)&dw[8]), - "=v"(*(i32x4_t *)&dw[12]) + : "=&v"(*(i32x4_t *)&dw[0]), + "=&v"(*(i32x4_t *)&dw[4]), + "=&v"(*(i32x4_t *)&dw[8]), + "=&v"(*(i32x4_t *)&dw[12]) : "v"(base_ptr), "v"(base_ptr + 4), "v"(base_ptr + 8), diff --git a/include/mirage/persistent_kernel/tasks/mi300/gang_oproj_topk_moe_fused_mi300.cuh b/include/mirage/persistent_kernel/tasks/mi300/gang_oproj_topk_moe_fused_mi300.cuh index 6656441..8086266 100644 --- a/include/mirage/persistent_kernel/tasks/mi300/gang_oproj_topk_moe_fused_mi300.cuh +++ b/include/mirage/persistent_kernel/tasks/mi300/gang_oproj_topk_moe_fused_mi300.cuh @@ -117,6 +117,25 @@ __device__ __attribute__((always_inline)) void __syncthreads(); int expected = s_routing_expected; + // Publish the layer counter for the MoE W13->W2 barrier, which reads it from + // this fixed LDS offset and releases with layer_idx + 1. Its d_barrier is + // monotonic and never reset, so this must advance once per layer for the + // whole run or the barrier stops gating after the first layer. + // + // The routing epoch is that counter: bumped exactly once per layer by the + // TopK completer, never reset, and identical for every worker here because + // it is read before Phase 1 signals it. + { + constexpr int LAYER_IDX_SMEM_OFF = + mirage::runtime::MAX_DYNAMIC_SHARED_MEMORY_SIZE - + mirage::runtime::LAYER_IDX_SMEM_OFFSET_FROM_END; + extern __shared__ char _oproj_moe_smem[]; + if (tid == 0) { + *reinterpret_cast(&_oproj_moe_smem[LAYER_IDX_SMEM_OFF]) = expected; + } + } + __syncthreads(); + // ── Phase 1: O-PROJ + RMSNorm + Router + TopK ── // Need max(oproj_tiles_per_xcd, router_tile_n) workers: O-proj may need // fewer tiles than TopK experts when OUTPUT_PER_WG is large. diff --git a/include/mirage/persistent_kernel/tasks/mi300/gang_rmsnorm_linear_mxfp4_bias_argmax_mi300.cuh b/include/mirage/persistent_kernel/tasks/mi300/gang_rmsnorm_linear_mxfp4_bias_argmax_mi300.cuh index 75747b8..6a9b743 100644 --- a/include/mirage/persistent_kernel/tasks/mi300/gang_rmsnorm_linear_mxfp4_bias_argmax_mi300.cuh +++ b/include/mirage/persistent_kernel/tasks/mi300/gang_rmsnorm_linear_mxfp4_bias_argmax_mi300.cuh @@ -40,6 +40,12 @@ __device__ __noinline__ void gang_rmsnorm_linear_mxfp4_bias_argmax_kernel( void const *bias_ptr, void *argmax_val_ptr, // [num_workers] bf16: per-worker max value void *argmax_idx_ptr, // [num_workers] int64: per-worker absolute index + void *logits_out_ptr, // [max_seq_length, output_stride] f32, or nullptr. + // Perplexity mode only: when non-null the full logit + // row is also written to HBM at row `step`. Null on + // the serving path, which keeps logits in registers + // and pays no HBM traffic. + int step, // row of logits_out_ptr to write (ignored if null) int num_active_tokens, int n_wgs_per_xcd, int workers_per_xcd, @@ -286,6 +292,19 @@ __device__ __noinline__ void gang_rmsnorm_linear_mxfp4_bias_argmax_kernel( thread_max = val; thread_max_abs_idx = (long long)abs_idx; } + // Perplexity mode: also spill the logit to HBM. abs_idx is the + // vocab column this lane owns, so the writes across all workers + // tile the row exactly once -- no atomics, no races. + // + // Stored as f32, not bf16: bf16 carries ~0.4% relative precision, + // which is the same order as the GEMM error this buffer exists to + // measure. Truncating here would fold a comparable one-sided error + // into the very quantity under test. + if (logits_out_ptr != nullptr) { + reinterpret_cast( + logits_out_ptr)[(long long)step * output_stride + abs_idx] = + val; + } } } } diff --git a/include/mirage/persistent_kernel/tasks/mi300/gang_rmsnorm_linear_mxfp4_bias_mi300.cuh b/include/mirage/persistent_kernel/tasks/mi300/gang_rmsnorm_linear_mxfp4_bias_mi300.cuh index 70dcd43..e861f20 100644 --- a/include/mirage/persistent_kernel/tasks/mi300/gang_rmsnorm_linear_mxfp4_bias_mi300.cuh +++ b/include/mirage/persistent_kernel/tasks/mi300/gang_rmsnorm_linear_mxfp4_bias_mi300.cuh @@ -2374,7 +2374,27 @@ __device__ __noinline__ void "v_accvgpr_write_b32 a2, 0\n" "v_accvgpr_write_b32 a3, 0\n" - // Pre-issue 5 reads for iteration 0 + // ── Two disjoint operand banks (see below) ── + // Bank 0: A v[22:25], A scale v7, B v[8:15], B scale v16 + // Bank 1: A v[26:29], A scale v18, B v[32:39], B scale v19 + // Address scratch v17, accumulator a[0:3]. + // + // The loop MUST NOT prefetch into the registers the current MFMA + // reads as sources. lgkmcnt tracks when LDS data lands in the + // VGPR; it says nothing about when the MFMA has finished sampling + // its operands, and a 16x16x128 MFMA streams them over the op's + // duration rather than latching at issue. Overwriting in place is + // therefore a WAR race: when LDS returns fast the write-back lands + // mid-MFMA and the op sees mixed-iteration operands. Measured at + // ~17-22% of launches wrong before banking. + // + // Ping-pong rule: while the MFMA consumes bank X, the prefetch + // writes bank 1-X. No register is ever both a live MFMA source and + // an in-flight LDS destination, so the race cannot occur. + // + // Verified by tests/standalone/test_mfma_pipeline_hazards.hip. + + // Pre-issue 5 reads for iteration 0 into bank 0 "ds_read_b128 v[22:25], %[wa]\n" // weight A (16B FP4) "ds_read_u8 v7, %[wsa]\n" // weight A scale "ds_read_b128 v[8:11], %[ta]\n" // token B lo (16B FP8) @@ -2382,46 +2402,77 @@ __device__ __noinline__ void "ds_read_u8 v16, %[tsa]\n" // token B scale "s_mov_b32 s13, 0\n" // loop counter - // ── Iterations 0..22: prefetch next, MFMA current ── + // ── MFMA_ITERS-1 in-loop MFMAs, alternating banks ── "PIPELINED_QKV_%=:\n" - "s_waitcnt lgkmcnt(0)\n" // iter N reads complete - // Advance addresses for iter N+1 + // ---- consume bank 0, prefetch into bank 1 ---- + "s_waitcnt lgkmcnt(0)\n" // bank 0 operands resident + + // Advance addresses for the next iteration "v_add_u32_e32 %[wa], 64, %[wa]\n" "v_add_u32_e32 %[wsa], 4, %[wsa]\n" "v_add_u32_e32 %[ta], 0x80, %[ta]\n" "s_add_i32 s13, s13, 1\n" - // Token B scale → v17 FIRST (oldest in queue, completes first) "v_add_u32_e32 v17, s13, %[tsa]\n" - "ds_read_u8 v17, v17\n" // [lgkmcnt +1] oldest - // Prefetch iter N+1 data into SAME regs (MFMA reads old values) - "ds_read_b128 v[22:25], %[wa]\n" // [lgkmcnt +2] - "ds_read_u8 v7, %[wsa]\n" // [lgkmcnt +3] - "ds_read_b128 v[8:11], %[ta]\n" // [lgkmcnt +4] - "ds_read_b128 v[12:15], %[ta] offset:64\n" // [lgkmcnt +5] - - // MFMA from iter N data (32 cycles, reads v[22:25] v[8:15] v7 v16) + "ds_read_u8 v19, v17\n" // B scale -> bank 1 + "ds_read_b128 v[26:29], %[wa]\n" // A data -> bank 1 + "ds_read_u8 v18, %[wsa]\n" // A scale -> bank 1 + "ds_read_b128 v[32:35], %[ta]\n" // B lo -> bank 1 + "ds_read_b128 v[36:39], %[ta] offset:64\n" // B hi -> bank 1 + "v_mfma_scale_f32_16x16x128_f8f6f4 a[0:3], v[22:25], v[8:15], " "a[0:3], v7, v16 op_sel_hi:[0,0,0] cbsz:4\n" - // Copy next token scale during MFMA execution. - // lgkmcnt(4): wait for token scale (oldest, issued first), leave 4 - // data reads flying - "s_waitcnt lgkmcnt(4)\n" - "v_mov_b32_e32 v16, v17\n" + "s_cmpk_lt_i32 s13, %[iters_m1]\n" + "s_cbranch_scc0 QKV_TAIL_B1_%=\n" + + // ---- consume bank 1, prefetch into bank 0 ---- + "s_waitcnt lgkmcnt(0)\n" // bank 1 operands resident + + "v_add_u32_e32 %[wa], 64, %[wa]\n" + "v_add_u32_e32 %[wsa], 4, %[wsa]\n" + "v_add_u32_e32 %[ta], 0x80, %[ta]\n" + "s_add_i32 s13, s13, 1\n" + + "v_add_u32_e32 v17, s13, %[tsa]\n" + "ds_read_u8 v16, v17\n" // B scale -> bank 0 + "ds_read_b128 v[22:25], %[wa]\n" // A data -> bank 0 + "ds_read_u8 v7, %[wsa]\n" // A scale -> bank 0 + "ds_read_b128 v[8:11], %[ta]\n" // B lo -> bank 0 + "ds_read_b128 v[12:15], %[ta] offset:64\n" // B hi -> bank 0 + + "v_mfma_scale_f32_16x16x128_f8f6f4 a[0:3], v[26:29], v[32:39], " + "a[0:3], v18, v19 op_sel_hi:[0,0,0] cbsz:4\n" "s_cmpk_lt_i32 s13, %[iters_m1]\n" "s_cbranch_scc1 PIPELINED_QKV_%=\n" - // ── Final iteration: no more prefetch needed ── + // ── Final MFMA ── + // Exit parity matters, so BOTH tails are emitted. In the live + // instantiation REDUCTION_SIZE is 2944, so MFMA_ITERS is 23 (odd) + // and the loop falls out of the bank 1 half with the final + // operands in BANK 0 -- this path. An even MFMA_ITERS exits via + // QKV_TAIL_B1 with the final operands in bank 1. Emitting only one + // tail would silently use the wrong bank for one of the two + // parities and reintroduce mixed-iteration operands. "s_waitcnt lgkmcnt(0)\n" "v_mfma_scale_f32_16x16x128_f8f6f4 a[0:3], v[22:25], v[8:15], " "a[0:3], v7, v16 op_sel_hi:[0,0,0] cbsz:4\n" + "s_branch QKV_ACC_%=\n" - // Read accumulator into output - "s_nop 7\n" - "s_nop 0\n" + "QKV_TAIL_B1_%=:\n" + "s_waitcnt lgkmcnt(0)\n" + "v_mfma_scale_f32_16x16x128_f8f6f4 a[0:3], v[26:29], v[32:39], " + "a[0:3], v18, v19 op_sel_hi:[0,0,0] cbsz:4\n" + + "QKV_ACC_%=:\n" + // 32 clocks before reading the accumulator. The scaled MFMA is a + // 32-cycle op on CDNA4; the previous "s_nop 7; s_nop 0" was 9 + // clocks (the correct wait for a 4-pass MFMA) and returned a + // partially-retired accumulator on every launch. + "s_nop 15\n" + "s_nop 15\n" "v_accvgpr_read_b32 %[acc0], a0\n" "v_accvgpr_read_b32 %[acc1], a1\n" "v_accvgpr_read_b32 %[acc2], a2\n" @@ -2447,10 +2498,24 @@ __device__ __noinline__ void "v15", "v16", "v17", + "v18", + "v19", "v22", "v23", "v24", "v25", + "v26", + "v27", + "v28", + "v29", + "v32", + "v33", + "v34", + "v35", + "v36", + "v37", + "v38", + "v39", "a0", "a1", "a2", diff --git a/python/mirage/mpk/persistent_kernel.py b/python/mirage/mpk/persistent_kernel.py index fb20e79..f4f006b 100644 --- a/python/mirage/mpk/persistent_kernel.py +++ b/python/mirage/mpk/persistent_kernel.py @@ -377,6 +377,19 @@ def get_compile_command( if int(os.environ.get("PRECOMPUTED_DISPATCH", "1")) == 1: flags = flags + ["-DMPK_PRECOMPUTED_DISPATCH"] flags = flags + ["-DMPK_FUSED_LAYER_BATCHING"] + if int(os.environ.get("MPK_NIL_TRIPWIRE", "0")) == 1: + # Breadcrumbs in pinned host memory + SIGABRT dump, to attribute + # the nil-address memory fault. Off by default: the per-layer + # writes cost a little and only matter while chasing that bug. + flags = flags + ["-DMPK_NIL_TRIPWIRE"] + if int(os.environ.get("MPK_WORKER_STATE", "0")) == 1: + # Per-phase worker-state breadcrumbs: which phase/barrier each + # worker is in, dumped on a hang. This is how the fused-layer + # deadlocks were attributed, so keep it reachable -- but off by + # default. The stores go to pinned *host* memory over PCIe from + # ~30 sites in the two hottest task headers, several inside spin + # loops; measured 2.386 -> 2.321 ms/iter when compiled out. + flags = flags + ["-DMPK_WORKER_STATE"] if int(os.environ.get("TRACE_MOE", "0")) == 1: flags = flags + ["-DMPK_TRACE_MOE_DISPATCH"] if int(os.environ.get("EMBED_DEBUG", "0")) == 1: @@ -2500,6 +2513,7 @@ def gang_rmsnorm_linear_mxfp4_bias_argmax_layer( output_per_wg: int, output_stride: int, block_dim: tuple = (256, 1, 1), + ppl_logits: DTensor = None, ): """Fused RMSNorm + MXFP4 Gang Linear + Bias + Argmax (norm-once). @@ -2510,6 +2524,11 @@ def gang_rmsnorm_linear_mxfp4_bias_argmax_layer( total_tiles_per_xcd = workers_per_xcd (each worker enters once). Output: one (bf16 max, int64 abs_idx) per worker. Follow with argmax_reduce_layer(CHUNK_SIZE=0) for final token. + + ppl_logits: optional [max_seq_length, output_stride] float32 sink. + When supplied the kernel ALSO writes the full logit row for the + position being scored, which perplexity needs and argmax discards. + Costs a 393KB HBM write per step, so leave it None for serving. """ assert norm_input.num_dims == 2 assert mxfp4_weight.num_dims == 2 @@ -2528,11 +2547,20 @@ def gang_rmsnorm_linear_mxfp4_bias_argmax_layer( tb_graph.new_input(bias, (1, -1, -1), 1, True) tb_graph.new_input(argmax_part_value, (1, -1, -1), -1, True) tb_graph.new_input(argmax_part_index, (1, -1, -1), -1, True) - self.kn_graph.customized( - [norm_input, norm_weight, norm_output, mxfp4_weight, bias, - argmax_part_value, argmax_part_index], - tb_graph, - ) + tensors = [norm_input, norm_weight, norm_output, mxfp4_weight, bias, + argmax_part_value, argmax_part_index] + if ppl_logits is not None: + # input_map (-1,-1,-1), NOT (1,-1,-1) like the argmax partials. + # A non-negative map partitions that dim across grid_dim.x, and the + # runtime pre-shifts each block's base pointer by + # bid.x * vocab/8. The kernel indexes logits by the *absolute* + # vocab column (the same abs_idx it uses for argmax), so a + # pre-shifted base would double-count the partition offset -- + # blocks 0..3 would write every other 25152-column stripe and + # blocks 4..7 would run off the end of the row into the next one. + tb_graph.new_input(ppl_logits, (-1, -1, -1), -1, True) + tensors.append(ppl_logits) + self.kn_graph.customized(tensors, tb_graph) self.kn_graph.register_task( tb_graph, "gang_rmsnorm_linear_mxfp4_bias_argmax_mi300", [output_stride, output_per_wg, n_wgs_per_xcd, diff --git a/src/kernel/graph.cc b/src/kernel/graph.cc index 47ddf80..d7fc01f 100644 --- a/src/kernel/graph.cc +++ b/src/kernel/graph.cc @@ -654,8 +654,16 @@ void Graph::register_task(char const *task_type, std::vector params) { task_register ->register_gang_rmsnorm_linear_mxfp4_bias_argmax_mi300_task( customized->bgraph, params); - task_config[op] = std::make_tuple( - 5, 2, TASK_GANG_RMSNORM_LINEAR_MXFP4_BIAS_ARGMAX_MI300, variant_id); + // 2 outputs normally (argmax value/index); 3 when a perplexity logits + // sink is attached as a trailing output. Derive rather than hardcode so + // both shapes reach runtime.cc with a consistent input/output split. + int argmax_num_outputs = (int)customized->bgraph.operators.size() - 5; + assert(argmax_num_outputs == 2 || argmax_num_outputs == 3); + task_config[op] = + std::make_tuple(5, + argmax_num_outputs, + TASK_GANG_RMSNORM_LINEAR_MXFP4_BIAS_ARGMAX_MI300, + variant_id); gang_task_tiles_per_xcd[op] = params[3]; // total_tiles_per_xcd } else if (name == "gang_mulsumradd_rmsnorm_linear_mxfp4_bias_mi300") { assert(params.size() == 7 && diff --git a/src/kernel/runtime.cc b/src/kernel/runtime.cc index 33ec3aa..1253583 100644 --- a/src/kernel/runtime.cc +++ b/src/kernel/runtime.cc @@ -349,7 +349,13 @@ void register_mugraph( : desc.stride[d + 1] * output_ops[0]->dtensor.dim[d + 1]; } - task.inputs[task.num_outputs++] = desc; + // Was task.inputs[task.num_outputs++]: wrote the output + // descriptor into the *input* array, indexed by the output + // counter. With num_inputs==2 from the loop above and + // num_outputs==0 here, that clobbered inputs[0] and left + // outputs[0] unset, so TASK_REDUCE got a null output pointer and + // a corrupted first input. + task.outputs[task.num_outputs++] = desc; all_tasks.push_back(task); // Update current task map cur_task_map[bid] = all_tasks.size() - 1; diff --git a/src/kernel/task_register.cc b/src/kernel/task_register.cc index 4e555ad..2a7b8c1 100644 --- a/src/kernel/task_register.cc +++ b/src/kernel/task_register.cc @@ -1187,6 +1187,9 @@ int TaskRegister::register_gang_rmsnorm_linear_mxfp4_bias_mi300_task( // total_tiles_per_xcd = workers_per_xcd (each worker enters once, loops // internally). Inputs: [norm_input, norm_weight, norm_output, mxfp4_weight, // bias] Outputs: [argmax_part_value (bf16), argmax_part_index (int64)] +// plus an OPTIONAL third output [ppl_logits (f32)] used only by perplexity +// mode. When absent the kernel gets a nullptr and skips the HBM logits +// write entirely, so the serving path is byte-for-byte the previous code. int TaskRegister::register_gang_rmsnorm_linear_mxfp4_bias_argmax_mi300_task( threadblock::Graph const &bgraph, std::vector const ¶ms) { assert(params.size() == 5); @@ -1199,9 +1202,10 @@ int TaskRegister::register_gang_rmsnorm_linear_mxfp4_bias_argmax_mi300_task( std::vector input_ops; std::vector output_ops; int num_inputs = 5; - int num_outputs = 2; - - assert(bgraph.operators.size() == (size_t)num_inputs + num_outputs); + // 2 outputs normally, 3 when a perplexity logits sink is attached. + int num_outputs = (int)bgraph.operators.size() - num_inputs; + assert(num_outputs == 2 || num_outputs == 3); + bool emit_logits = (num_outputs == 3); for (auto const &op : bgraph.operators) { assert(op->op_type == mirage::type::TB_INPUT_OP); if (input_ops.size() < (size_t)num_inputs) { @@ -1228,6 +1232,14 @@ int TaskRegister::register_gang_rmsnorm_linear_mxfp4_bias_argmax_mi300_task( code.e(" task_desc->input_ptrs[4],"); // bias code.e(" task_desc->output_ptrs[0],"); // argmax_part_value (bf16) code.e(" task_desc->output_ptrs[1],"); // argmax_part_index (int64) + if (emit_logits) { + code.e(" task_desc->output_ptrs[2],"); // ppl_logits (f32) + } else { + code.e(" nullptr,"); // no logits sink + } + // Row to write in the logits buffer: the position being scored. step[0] is + // the last consumed position, so the token produced here lands at step+1. + code.e(" runtime_config.step[0] + 1,"); code.e(" runtime_config.qo_indptr_buffer[MPK_MAX_NUM_BATCHED_REQUESTS],"); code.e(" $,", n_wgs_per_xcd); code.e(" $,", workers_per_xcd); @@ -2091,7 +2103,12 @@ int TaskRegister::register_gang_full_layer_fused_mi300_task( code.e(" $,", oproj_tiles_per_xcd); code.e(" $,", moe_total_tiles_per_xcd); code.e(" $,", workers_per_xcd); - code.e(" tile_idx);"); + code.e(" tile_idx,"); + // Deterministic layer counter, published by the ml loop into the free int32 + // of the n_tile union member. The task derives its barrier release values + // from this instead of snapshotting a shared counter -- see the + // layer_counter comment in gang_full_layer_fused_mi300.cuh. + code.e(" (int)task_desc->task_metadata._linear_reserved);"); return register_task_variant(TASK_GANG_FULL_LAYER_FUSED_MI300, code.to_string()); } @@ -2211,7 +2228,9 @@ int TaskRegister::register_gang_full_layer_with_lmhead_fused_mi300_task( code.e(" $,", lm_n_wgs_per_xcd); code.e(" $,", lm_output_stride); code.e(" $,", lm_actual_hidden_dim); - code.e(" tile_idx);"); + code.e(" tile_idx,"); + // See the matching comment in the non-LM-head variant above. + code.e(" (int)task_desc->task_metadata._linear_reserved);"); return register_task_variant(TASK_GANG_FULL_LAYER_WITH_LMHEAD_FUSED_MI300, code.to_string()); } diff --git a/tests/ci-tests/run_ci_tests_gpt_oss.sh b/tests/ci-tests/run_ci_tests_gpt_oss.sh index 631fad2..7fd57dc 100755 --- a/tests/ci-tests/run_ci_tests_gpt_oss.sh +++ b/tests/ci-tests/run_ci_tests_gpt_oss.sh @@ -13,7 +13,6 @@ set -euo pipefail ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" export MIRAGE_HOME="${MIRAGE_HOME:-$ROOT}" export HIP_VISIBLE_DEVICES="${HIP_VISIBLE_DEVICES:-0}" -export USE_FP8_ACT="${USE_FP8_ACT:-1}" MODEL_PATH="${MODEL_PATH:-${GPT_OSS_MODEL_PATH:-openai/gpt-oss-120b}}" PROMPT="${GPT_OSS_PROMPT:-Tell me the history of america}" MAX_SEQ_LEN="${GPT_OSS_MAX_SEQ_LEN:-512}" diff --git a/tests/ci-tests/run_ci_tests_gpt_oss_perplexity.sh b/tests/ci-tests/run_ci_tests_gpt_oss_perplexity.sh new file mode 100755 index 0000000..ec2ae47 --- /dev/null +++ b/tests/ci-tests/run_ci_tests_gpt_oss_perplexity.sh @@ -0,0 +1,34 @@ +#!/usr/bin/env bash +# GPT-OSS 120B perplexity: score a WikiText-2 slice with the Torch reference and +# with Mirage (MPK), then compare (tests/ci-tests/test_gpt_oss_perplexity.py). +# +# PPL_MODE=1 loads the corpus as one long prompt and runs prefill only. The +# megakernel does not overwrite tokens[] inside the prompt, so every position +# conditions on the reference prefix -- prefill is teacher forcing. The LM head +# scores one row per iteration, hence --max-num-batched-tokens 1. +# +# Env (override as needed): +# MODEL_PATH local GPT-OSS 120B dir or HF repo id (required) +# HIP_VISIBLE_DEVICES target GPU (default 0) +# PPL_MAX_TOKENS corpus tokens to score (default 512) +set -euo pipefail + +ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" +export MIRAGE_HOME="${MIRAGE_HOME:-$ROOT}" +export HIP_VISIBLE_DEVICES="${HIP_VISIBLE_DEVICES:-0}" +export PPL_MODE=1 +MODEL_PATH="${MODEL_PATH:-${GPT_OSS_MODEL_PATH:-openai/gpt-oss-120b}}" +PPL_MAX_TOKENS="${PPL_MAX_TOKENS:-512}" +OUT_DIR="${GPT_OSS_OUTPUT_DIR:-$ROOT/outputs/gpt_oss}" + +DEMO="$ROOT/demo/gpt_oss/demo.py" +COMMON=(--model-path "$MODEL_PATH" --max-num-batched-tokens 1 + --ppl-max-tokens "$PPL_MAX_TOKENS") + +echo "MIRAGE_HOME=$MIRAGE_HOME HIP_VISIBLE_DEVICES=$HIP_VISIBLE_DEVICES MODEL_PATH=$MODEL_PATH" +echo "Scoring with the Torch reference..." +python3 "$DEMO" "${COMMON[@]}" --ppl-out "$OUT_DIR/torch_ppl.json" +echo "Scoring with Mirage (MPK)..." +python3 "$DEMO" --use-mirage "${COMMON[@]}" --ppl-out "$OUT_DIR/mpk_ppl.json" +echo "Comparing perplexity..." +GPT_OSS_OUTPUT_DIR="$OUT_DIR" pytest -q -s "$ROOT/tests/ci-tests/test_gpt_oss_perplexity.py" diff --git a/tests/ci-tests/run_gpt_oss_ppl_sweep.sh b/tests/ci-tests/run_gpt_oss_ppl_sweep.sh new file mode 100755 index 0000000..268ad10 --- /dev/null +++ b/tests/ci-tests/run_gpt_oss_ppl_sweep.sh @@ -0,0 +1,62 @@ +#!/usr/bin/env bash +# Perplexity vs sequence length for GPT-OSS 120B, 512 -> 32768 tokens. +# +# Scores a WikiText-2 prefix at each length with Mirage (MPK), and with the +# Torch reference wherever the reference can actually run. The reference +# attention in demo/gpt_oss/models/modeling_gpt_oss.py materializes a +# [64, n, n] float32 score matrix (:172), which is 64 GB per copy at 16k and +# 256 GB at 32k -- it OOMs well before MPK does. So Torch anchors the short +# end and the long end is MPK-only. TORCH_MAX_LEN sets that cutoff. +# +# Env: +# MODEL_PATH local GPT-OSS 120B dir or HF repo id +# HIP_VISIBLE_DEVICES target GPU (default 0) +# PPL_LENS lengths to sweep (default "512 1024 2048 4096 8192 16384 32768") +# TORCH_MAX_LEN longest length to also score with Torch (default 4096) +# PPL_SWEEP_OUT results dir (default outputs/gpt_oss/ppl_sweep) +set -uo pipefail # NOT -e: a single length OOMing must not kill the sweep + +ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" +export MIRAGE_HOME="${MIRAGE_HOME:-$ROOT}" +export HIP_VISIBLE_DEVICES="${HIP_VISIBLE_DEVICES:-0}" +export PPL_MODE=1 +MODEL_PATH="${MODEL_PATH:-${GPT_OSS_MODEL_PATH:-openai/gpt-oss-120b}}" +LENS="${PPL_LENS:-512 1024 2048 4096 8192 16384 32768}" +TORCH_MAX_LEN="${TORCH_MAX_LEN:-4096}" +OUT="${PPL_SWEEP_OUT:-$ROOT/outputs/gpt_oss/ppl_sweep}" +mkdir -p "$OUT" + +DEMO="$ROOT/demo/gpt_oss/demo.py" +LOG="$OUT/sweep.log" +: > "$LOG" + +echo "MODEL_PATH=$MODEL_PATH GPU=$HIP_VISIBLE_DEVICES lengths: $LENS" | tee -a "$LOG" +echo "Torch reference up to $TORCH_MAX_LEN tokens (longer OOMs the reference attention)" | tee -a "$LOG" + +run_one() { # $1=len $2=mpk|torch + local n="$1" mode="$2" + local tag="${mode}_${n}" + local extra=() + [ "$mode" = "mpk" ] && extra=(--use-mirage) + echo "=== $mode n=$n ===" | tee -a "$LOG" + timeout 3600 python3 "$DEMO" --model-path "$MODEL_PATH" "${extra[@]}" \ + --max-num-batched-tokens 1 --ppl-max-tokens "$n" \ + --ppl-out "$OUT/$tag.json" > "$OUT/$tag.log" 2>&1 + local rc=$? + if [ $rc -ne 0 ]; then + echo " FAILED rc=$rc (see $OUT/$tag.log)" | tee -a "$LOG" + tail -5 "$OUT/$tag.log" | sed 's/^/ /' | tee -a "$LOG" + else + grep -E "perplexity|mean entropy|top-1 accuracy|zero columns|per-worker-argmax|FWD_PASS_TOTAL" \ + "$OUT/$tag.log" | sed 's/^/ /' | tee -a "$LOG" + fi +} + +for n in $LENS; do + run_one "$n" mpk + if [ "$n" -le "$TORCH_MAX_LEN" ]; then run_one "$n" torch; fi +done + +echo "" | tee -a "$LOG" +python3 "$ROOT/tests/ci-tests/summarize_ppl_sweep.py" "$OUT" 2>&1 | tee -a "$LOG" +echo "Results in $OUT (per-run logs: $OUT/_.log)" | tee -a "$LOG" diff --git a/tests/ci-tests/summarize_ppl_sweep.py b/tests/ci-tests/summarize_ppl_sweep.py new file mode 100755 index 0000000..b672e63 --- /dev/null +++ b/tests/ci-tests/summarize_ppl_sweep.py @@ -0,0 +1,132 @@ +#!/usr/bin/env python3 +"""Tabulate a perplexity-vs-sequence-length sweep produced by +run_gpt_oss_ppl_sweep.sh. Writes summary.md and summary.csv next to the dumps. + +Usage: summarize_ppl_sweep.py + +Reads _.json. Entropy is reported alongside perplexity on purpose: +numeric noise flattens the softmax, which *lowers* NLL at positions the model +gets wrong, so perplexity alone can make a noisier run look like a better one. +""" +import csv +import json +import os +import re +import sys + + +def main(d): + rows = {} + for fn in sorted(os.listdir(d)): + m = re.fullmatch(r"(mpk|torch)_(\d+)\.json", fn) + if not m: + continue + mode, n = m.group(1), int(m.group(2)) + with open(os.path.join(d, fn)) as f: + j = json.load(f) + ent = j.get("entropy") or [] + top1, tgts = j.get("top1") or [], j.get("targets") or [] + acc = (sum(1 for a, b in zip(top1, tgts) if int(a) == int(b)) + / len(top1)) if top1 else None + rows.setdefault(n, {})[mode] = { + "ppl": j.get("perplexity"), + "nll": j.get("mean_nll"), + "scored": j.get("scored_positions"), + "ent": (sum(ent) / len(ent)) if ent else None, + "acc": acc, + "top1": top1, + "per_pos": j.get("per_position_nll") or [], + } + + if not rows: + print(f"No _.json dumps found in {d}") + return 1 + + # Each length scores a longer prefix of the same corpus, so the full-slice + # perplexity moves with what text got included, not just with context + # length -- WikiText-2 gets markedly harder around position 512 (6.3 nats + # vs 3.6 for the first 511), which shows up as rising perplexity on the + # *Torch reference* too. Restricting every run to the positions all runs + # share removes that and leaves only the effect of longer context. + shortest = min(rows) + common = min( + (len(v["per_pos"]) for r in rows.values() for v in r.values() + if v["per_pos"]), + default=0, + ) + for r in rows.values(): + for v in r.values(): + pp = v["per_pos"][:common] + v["ppl_common"] = ( + __import__("math").exp(sum(pp) / len(pp)) if pp else None + ) + + def f(x, spec="8.3f"): + return format(x, spec) if isinstance(x, (int, float)) else " n/a" + + out = [] + out.append(f"| seq len | scored | MPK ppl | Torch ppl | ratio | " + f"MPK ppl@{common} | Torch ppl@{common} | " + f"MPK ent | Torch ent | MPK acc | Torch acc | agree |") + out.append("|--------:|-------:|--------:|----------:|------:|" + "----------:|------------:|" + "--------:|----------:|--------:|----------:|------:|") + csv_rows = [] + for n in sorted(rows): + r = rows[n] + mk, tc = r.get("mpk", {}), r.get("torch", {}) + ratio = (mk.get("ppl") / tc["ppl"] + if mk.get("ppl") and tc.get("ppl") else None) + agree = None + if mk.get("top1") and tc.get("top1"): + k = min(len(mk["top1"]), len(tc["top1"])) + agree = sum(1 for i in range(k) + if int(mk["top1"][i]) == int(tc["top1"][i])) / k + out.append( + f"| {n:7d} | {mk.get('scored') or tc.get('scored') or 0:6d} " + f"| {f(mk.get('ppl'))} | {f(tc.get('ppl'))} | {f(ratio, '6.3f')} " + f"| {f(mk.get('ppl_common'), '10.3f')} " + f"| {f(tc.get('ppl_common'), '12.3f')} " + f"| {f(mk.get('ent'), '7.3f')} | {f(tc.get('ent'), '9.3f')} " + f"| {f(mk.get('acc'), '7.4f')} | {f(tc.get('acc'), '9.4f')} " + f"| {f(agree, '6.4f')} |" + ) + csv_rows.append({ + "seq_len": n, "scored": mk.get("scored") or tc.get("scored"), + "mpk_ppl": mk.get("ppl"), "torch_ppl": tc.get("ppl"), + "ratio": ratio, + "mpk_ppl_common_prefix": mk.get("ppl_common"), + "torch_ppl_common_prefix": tc.get("ppl_common"), + "mpk_entropy": mk.get("ent"), + "torch_entropy": tc.get("ent"), "mpk_acc": mk.get("acc"), + "torch_acc": tc.get("acc"), "top1_agreement": agree, + }) + + table = "\n".join(out) + print("\nPerplexity vs sequence length (WikiText-2 prefix)\n") + print(table) + print("\nppl/ent = perplexity / mean predictive entropy (nats). " + "acc = argmax vs corpus targets.\nagree = MPK and Torch pick the " + "same argmax. Blank Torch cells = the reference\nattention " + "([64,n,n] f32) does not fit at that length.") + print(f"\nppl@{common} is the SAME first {common} positions in every run " + f"-- read this column for the\neffect of context length. The plain " + f"ppl column also moves with which text each\nslice includes " + f"(WikiText-2 gets harder past ~512 tokens), which is why the Torch\n" + f"reference rises too.") + + with open(os.path.join(d, "summary.md"), "w") as fh: + fh.write("# GPT-OSS 120B perplexity vs sequence length\n\n" + "Corpus: WikiText-2 raw test, first N tokens. MPK is " + "prefill-only (teacher forced).\n\n" + table + "\n") + with open(os.path.join(d, "summary.csv"), "w", newline="") as fh: + w = csv.DictWriter(fh, fieldnames=list(csv_rows[0].keys())) + w.writeheader() + w.writerows(csv_rows) + print(f"\nWrote {d}/summary.md and {d}/summary.csv") + return 0 + + +if __name__ == "__main__": + sys.exit(main(sys.argv[1] if len(sys.argv) > 1 + else "outputs/gpt_oss/ppl_sweep")) diff --git a/tests/ci-tests/test_gpt_oss_perplexity.py b/tests/ci-tests/test_gpt_oss_perplexity.py new file mode 100644 index 0000000..7bc4e29 --- /dev/null +++ b/tests/ci-tests/test_gpt_oss_perplexity.py @@ -0,0 +1,140 @@ +"""GPT-OSS 120B perplexity: Mirage (MPK) vs the Torch reference. + +Companion to test_gpt_oss_inference_output.py. That test gates on generated +tokens, which is a thresholded single-prompt proxy: it says "the two paths +agree" or "they don't" and nothing in between. This one gates on a continuous, +corpus-wide number, so a numerical regression that only bends the distribution +(rather than flipping the argmax) still shows up. + +Both dumps come from demo.py with PPL_MODE=1, which loads a corpus as one long +prompt and runs prefill only. Inside the prompt the megakernel never overwrites +tokens[] (persistent_kernel.cuh guards the writeback on step+1 >= prompt_len), +so every position conditions on the reference prefix -- prefill IS teacher +forcing, which is exactly what perplexity needs. + +Produce the inputs first (from repo root): + PPL_MODE=1 python3 demo/gpt_oss/demo.py --model-path \ + --max-num-batched-tokens 1 --ppl-max-tokens 512 \ + --ppl-out outputs/gpt_oss/torch_ppl.json + PPL_MODE=1 python3 demo/gpt_oss/demo.py --model-path --use-mirage \ + --max-num-batched-tokens 1 --ppl-max-tokens 512 \ + --ppl-out outputs/gpt_oss/mpk_ppl.json + +Why a ceiling and not an exact value: MXFP4 decode is not bit-deterministic, so +the MPK number moves run to run. On a 512-token WikiText-2 slice the observed +spread was ~92-100 against a Torch reference of ~36. The defaults below leave +headroom above that spread while still catching a real blowup (a broken logits +sink lands in the hundreds or worse -- an early partition-offset bug in the sink +scored 64.5 against a Torch 8.15 on a 64-token slice). + +Tunables (env): + GPT_OSS_OUTPUT_DIR dir holding the two json dumps (default outputs/gpt_oss) + GPT_OSS_PPL_MAX absolute MPK perplexity ceiling (default 150) + GPT_OSS_PPL_RATIO_MAX max MPK/Torch perplexity ratio (default 3.5) + GPT_OSS_PPL_MIN_SCORED min scored positions to accept a run (default 64) +""" +import json +import math +import os +import pytest + +DEFAULT_OUTPUT_DIR = os.environ.get( + "GPT_OSS_OUTPUT_DIR", os.path.join("outputs", "gpt_oss") +) +TORCH_PPL = os.path.join(DEFAULT_OUTPUT_DIR, "torch_ppl.json") +MPK_PPL = os.path.join(DEFAULT_OUTPUT_DIR, "mpk_ppl.json") + +PPL_MAX = float(os.environ.get("GPT_OSS_PPL_MAX", "150")) +PPL_RATIO_MAX = float(os.environ.get("GPT_OSS_PPL_RATIO_MAX", "3.5")) +MIN_SCORED = int(os.environ.get("GPT_OSS_PPL_MIN_SCORED", "64")) + + +def _load_ppl(path): + if not os.path.exists(path): + pytest.fail( + f"Missing perplexity file: {path}. Run demo/gpt_oss/demo.py with " + f"PPL_MODE=1 --max-num-batched-tokens 1 --ppl-out {path} " + f"(and --use-mirage for the MPK dump) first." + ) + with open(path) as f: + data = json.load(f) + ppl = data.get("perplexity") + if not isinstance(ppl, (int, float)): + pytest.fail(f"'perplexity' missing or not a number in {path}") + if not math.isfinite(ppl): + pytest.fail( + f"Non-finite perplexity ({ppl}) in {path} -- the logits sink " + f"produced inf/nan, which is a kernel bug, not a quality result." + ) + if ppl < 1.0: + pytest.fail( + f"Perplexity {ppl} < 1 in {path}, which is impossible for a " + f"cross-entropy over a real distribution. The scoring slice or the " + f"vocab truncation is wrong." + ) + return ppl, data + + +def _top1_agreement(a, b): + """Fraction of positions where the two runs pick the same argmax token.""" + if not a or not b: + return None + n = min(len(a), len(b)) + return sum(1 for i in range(n) if int(a[i]) == int(b[i])) / n + + +def test_gpt_oss_mpk_perplexity(): + torch_ppl, torch_meta = _load_ppl(TORCH_PPL) + mpk_ppl, mpk_meta = _load_ppl(MPK_PPL) + + # The two numbers are only comparable if they scored the same text. A + # silent corpus/slice mismatch would otherwise read as a quality delta. + for key in ("corpus", "corpus_desc", "corpus_tokens", "scored_positions"): + if torch_meta.get(key) != mpk_meta.get(key): + pytest.fail( + f"Dumps scored different slices: {key} is " + f"{torch_meta.get(key)!r} (torch) vs {mpk_meta.get(key)!r} " + f"(mpk). Re-run both with the same --ppl-corpus and " + f"--ppl-max-tokens." + ) + + n_scored = mpk_meta.get("scored_positions", 0) + if n_scored < MIN_SCORED: + pytest.fail( + f"Only {n_scored} scored positions (require >= {MIN_SCORED}); " + f"too few for the mean NLL to be stable. Raise --ppl-max-tokens." + ) + + ratio = mpk_ppl / torch_ppl + agreement = _top1_agreement(mpk_meta.get("top1"), torch_meta.get("top1")) + + # Reported, not gated. Top-1 agreement is the sharper discriminator of the + # two, but it is also the noisier one across runs of the same build, so + # gating on it would flake. Read it when the perplexity gate fires. + agree_str = f"{agreement:.2%}" if agreement is not None else "n/a" + print( + f"[gpt-oss perplexity] scored {n_scored} positions of " + f"{mpk_meta.get('corpus')}: mpk_ppl={mpk_ppl:.4f} " + f"(ceiling {PPL_MAX}), torch_ppl={torch_ppl:.4f}, " + f"ratio={ratio:.3f} (max {PPL_RATIO_MAX}), " + f"mpk/torch top-1 agreement={agree_str}" + ) + + if mpk_ppl > PPL_MAX: + pytest.fail( + f"MPK perplexity {mpk_ppl:.4f} exceeds ceiling {PPL_MAX} on " + f"{n_scored} positions (torch reference {torch_ppl:.4f}, " + f"ratio {ratio:.3f}, top-1 agreement {agree_str}). Either a " + f"numerical regression reached the LM head or the logits sink is " + f"writing the wrong columns -- check the per-column coverage " + f"diagnostic the demo prints in PPL_MODE." + ) + + if ratio > PPL_RATIO_MAX: + pytest.fail( + f"MPK perplexity {mpk_ppl:.4f} is {ratio:.3f}x the Torch " + f"reference {torch_ppl:.4f} (max {PPL_RATIO_MAX}) over " + f"{n_scored} positions, top-1 agreement {agree_str}. The absolute " + f"ceiling passed, so the corpus may just be hard -- but the two " + f"paths disagree by more than MXFP4 run-to-run drift explains." + ) diff --git a/tests/standalone/build.sh b/tests/standalone/build.sh index 3bd0949..15d721b 100755 --- a/tests/standalone/build.sh +++ b/tests/standalone/build.sh @@ -10,6 +10,19 @@ hipcc -o test_mfma_simple test_mfma_simple.hip \ -O3 \ -std=c++17 +# gfx950-only: the scaled-MFMA hazard regression uses +# v_mfma_scale_f32_16x16x128_f8f6f4, which does not assemble on other targets. +if [ "${GFX_ARCH:-gfx950}" = "gfx950" ]; then + echo "Building scaled-MFMA pipeline hazard regression..." + hipcc -o test_mfma_pipeline_hazards test_mfma_pipeline_hazards.hip \ + -D__HIP_PLATFORM_AMD__ \ + --offload-arch=gfx950 \ + -munsafe-fp-atomics \ + -O3 \ + -std=c++17 +fi + echo "Build complete!" echo "" echo "Run with: ./test_mfma_simple" +echo " ./test_mfma_pipeline_hazards [launches] # gfx950 only" diff --git a/tests/standalone/test_mfma_pipeline_hazards.hip b/tests/standalone/test_mfma_pipeline_hazards.hip new file mode 100644 index 0000000..262b52c --- /dev/null +++ b/tests/standalone/test_mfma_pipeline_hazards.hip @@ -0,0 +1,412 @@ +/* Copyright 2025 CMU + * + * Regression test for the two gfx950 scaled-MFMA pipelining hazards. + * + * These hazards are SILENT: they corrupt roughly one iteration out of a + * 22-iteration accumulation, which on an FP4-quantized model looks like + * quantization noise rather than a bug. Nothing hangs, nothing faults, and + * the model still emits fluent text. The only way to catch them is to + * compare against a schedule that cannot race. + * + * Hazard 1 (WAR on MFMA source registers) + * The pipelined loop issues iteration N+1's ds_read into the SAME VGPRs + * that iteration N's MFMA is reading, and relies on the reads landing + * "late". lgkmcnt tracks when LDS data reaches the VGPR -- it says nothing + * about when the MFMA has finished sampling its sources. When LDS returns + * fast the write-back lands mid-MFMA and the op sees mixed-iteration + * operands. Intermittent: ~17-22% of launches. + * + * Hazard 2 (AccVGPR read before retirement) + * "s_nop 7; s_nop 0" is a 9-clock gap. The FP4xFP8 scaled MFMA is a + * 32-cycle op on CDNA4, so v_accvgpr_read_b32 samples the accumulator + * mid-retirement. Deterministic: fails every launch. The accumulator + * retires progressively, so the signature depends on the gap: at 9 clocks + * a0..a2 are stale with a3 correct, at 10 clocks only a0 is stale. + * + * Measured threshold on gfx950 is exactly 11 clocks -- 10 fails 300/300, + * 11 passes 300/300. Note that LLVM's hazard recognizer emits precisely + * "s_nop 7; s_nop 2" (11 clocks) for intrinsic-based MFMAs elsewhere in + * this codebase, which corroborates the number. Inline asm bypasses that + * recognizer entirely -- that is why the hand-written tail sat 2 clocks + * short. We still pad to 32 rather than 11: the threshold is not + * architecturally documented, and the tail runs once per GEMM. + * + * The fix for hazard 1 is two disjoint operand banks (ping-pong), so no + * register is ever both a live MFMA source and an in-flight LDS destination. + * The fix for hazard 2 is "s_nop 15" x2 == 32 clocks. + * + * NOTE ON LOOP PARITY: the ping-pong loop's exit parity decides which bank + * holds the final operands. An ODD iteration count falls out of the bank-1 + * half (final operands in bank 0); an EVEN count exits via the bank-0 half + * (final operands in bank 1). A single-tail port gets one of the two cases + * backwards and silently reintroduces the bug, so BOTH exit paths are spelled + * out below. Production runs 23 iterations (REDUCTION_SIZE 2944 / + * K_PER_MFMA 128) for QKV, W13 and W2 alike. Rebuild with -DMFMA_ITERS=22 to + * exercise the even case. + * + * Build: ./build.sh (or see the hipcc line in that script) + * Run: ./test_mfma_pipeline_hazards [launches] + * + * Exit code 0 = fixed schedule is bit-exact across all launches. + * Exit code 1 = a hazard is present; the mfma_fixed kernel below is the + * schedule production code should match. + */ + +#include +#include +#include +#include +#include +#include + +#ifndef MFMA_ITERS +// Production runs 23 (REDUCTION_SIZE 2944 / K_PER_MFMA 128). Both parities +// are exercised -- see the ITERS note in the header comment. +#define MFMA_ITERS 23 +#endif + +#define HIP_CHECK(call) \ + do { \ + hipError_t err = call; \ + if (err != hipSuccess) { \ + std::fprintf(stderr, \ + "HIP error at %s:%d: %s\n", \ + __FILE__, \ + __LINE__, \ + hipGetErrorString(err)); \ + std::exit(2); \ + } \ + } while (0) + +// LDS staging shared by every variant: A data 64B/iter, A scale 1B/iter, +// B data 128B/iter, B scale 1B/iter. +#define HAZ_PROLOGUE \ + extern __shared__ uint8_t smem[]; \ + int tid = threadIdx.x; \ + uint8_t *sAd = smem; \ + uint8_t *sAs = sAd + 64 * MFMA_ITERS; \ + uint8_t *sBd = sAs + 64; \ + uint8_t *sBs = sBd + 128 * MFMA_ITERS; \ + for (int i = tid; i < 64 * MFMA_ITERS; i += 256) \ + sAd[i] = Ad[i]; \ + for (int i = tid; i < MFMA_ITERS; i += 256) \ + sAs[i] = As[i]; \ + for (int i = tid; i < 128 * MFMA_ITERS; i += 256) \ + sBd[i] = Bd[i]; \ + for (int i = tid; i < MFMA_ITERS; i += 256) \ + sBs[i] = Bs[i]; \ + __syncthreads(); \ + int grp = (tid & 63) >> 4; \ + unsigned w_addr = (unsigned)(uintptr_t)(sAd) + (unsigned)(grp * 16); \ + unsigned ws_addr = (unsigned)(uintptr_t)(sAs); \ + unsigned t_addr = (unsigned)(uintptr_t)(sBd) + (unsigned)(grp * 16); \ + unsigned ts_addr = (unsigned)(uintptr_t)(sBs); \ + float acc0, acc1, acc2, acc3; + +#define HAZ_EPILOGUE \ + size_t off = ((size_t)blockIdx.x * 256 + tid) * 4; \ + out[off + 0] = acc0; \ + out[off + 1] = acc1; \ + out[off + 2] = acc2; \ + out[off + 3] = acc3; + +#define HAZ_OUTS \ + [a0] "=v"(acc0), [a1] "=v"(acc1), [a2] "=v"(acc2), [a3] "=v"(acc3), \ + [wa] "+v"(w_addr), [wsa] "+v"(ws_addr), [ta] "+v"(t_addr) + +#define HAZ_INS [tsa] "v"(ts_addr), [im1] "n"(MFMA_ITERS - 1) + +// Both operand banks are clobbered in every variant so the register +// allocator cannot reuse them and perturb the schedule under test. +#define HAZ_CLOBBER \ + "memory", "s13", "v7", "v8", "v9", "v10", "v11", "v12", "v13", "v14", "v15", \ + "v16", "v17", "v18", "v19", "v22", "v23", "v24", "v25", "v26", "v27", \ + "v28", "v29", "v32", "v33", "v34", "v35", "v36", "v37", "v38", "v39", \ + "a0", "a1", "a2", "a3" + +// --------------------------------------------------------------------------- +// BROKEN: the historical production schedule. Both hazards present. +// Kept so the test proves it is actually detecting something. +// --------------------------------------------------------------------------- +__global__ void mfma_broken(float *out, + uint8_t const *Ad, + uint8_t const *As, + uint8_t const *Bd, + uint8_t const *Bs) { + HAZ_PROLOGUE + asm volatile( + "v_accvgpr_write_b32 a0, 0\nv_accvgpr_write_b32 a1, 0\n" + "v_accvgpr_write_b32 a2, 0\nv_accvgpr_write_b32 a3, 0\n" + "ds_read_b128 v[22:25], %[wa]\nds_read_u8 v7, %[wsa]\n" + "ds_read_b128 v[8:11], %[ta]\nds_read_b128 v[12:15], %[ta] offset:64\n" + "ds_read_u8 v16, %[tsa]\ns_mov_b32 s13, 0\n" + "HAZBAD%=:\n" + "s_waitcnt lgkmcnt(0)\n" + "v_add_u32_e32 %[wa], 64, %[wa]\nv_add_u32_e32 %[wsa], 1, %[wsa]\n" + "v_add_u32_e32 %[ta], 0x80, %[ta]\ns_add_i32 s13, s13, 1\n" + "v_add_u32_e32 v17, s13, %[tsa]\nds_read_u8 v17, v17\n" + // HAZARD 1: these four reads target the MFMA's own source registers. + "ds_read_b128 v[22:25], %[wa]\nds_read_u8 v7, %[wsa]\n" + "ds_read_b128 v[8:11], %[ta]\nds_read_b128 v[12:15], %[ta] offset:64\n" + "v_mfma_scale_f32_16x16x128_f8f6f4 a[0:3], v[22:25], v[8:15], a[0:3], " + "v7, v16 op_sel_hi:[0,0,0] cbsz:4\n" + "s_waitcnt lgkmcnt(4)\nv_mov_b32_e32 v16, v17\n" + "s_cmpk_lt_i32 s13, %[im1]\ns_cbranch_scc1 HAZBAD%=\n" + "s_waitcnt lgkmcnt(0)\n" + "v_mfma_scale_f32_16x16x128_f8f6f4 a[0:3], v[22:25], v[8:15], a[0:3], " + "v7, v16 op_sel_hi:[0,0,0] cbsz:4\n" + // HAZARD 2: 9 clocks before reading a 32-cycle op's accumulator. + "s_nop 7\ns_nop 0\n" + "v_accvgpr_read_b32 %[a0], a0\nv_accvgpr_read_b32 %[a1], a1\n" + "v_accvgpr_read_b32 %[a2], a2\nv_accvgpr_read_b32 %[a3], a3\n" + : HAZ_OUTS:HAZ_INS + : HAZ_CLOBBER); + HAZ_EPILOGUE +} + +// --------------------------------------------------------------------------- +// FIXED: two disjoint operand banks + a 32-clock retirement gap. +// This is the schedule production code should match. +// --------------------------------------------------------------------------- +__global__ void mfma_fixed(float *out, + uint8_t const *Ad, + uint8_t const *As, + uint8_t const *Bd, + uint8_t const *Bs) { + HAZ_PROLOGUE + asm volatile( + "v_accvgpr_write_b32 a0, 0\nv_accvgpr_write_b32 a1, 0\n" + "v_accvgpr_write_b32 a2, 0\nv_accvgpr_write_b32 a3, 0\n" + // Preload bank 0. + "ds_read_b128 v[22:25], %[wa]\nds_read_u8 v7, %[wsa]\n" + "ds_read_b128 v[8:11], %[ta]\nds_read_b128 v[12:15], %[ta] offset:64\n" + "ds_read_u8 v16, %[tsa]\ns_mov_b32 s13, 0\n" + "HAZFIX%=:\n" + // --- consume bank 0, prefetch into bank 1 (disjoint registers) --- + "s_waitcnt lgkmcnt(0)\n" + "v_add_u32_e32 %[wa], 64, %[wa]\nv_add_u32_e32 %[wsa], 1, %[wsa]\n" + "v_add_u32_e32 %[ta], 0x80, %[ta]\ns_add_i32 s13, s13, 1\n" + "v_add_u32_e32 v17, s13, %[tsa]\nds_read_u8 v19, v17\n" + "ds_read_b128 v[26:29], %[wa]\nds_read_u8 v18, %[wsa]\n" + "ds_read_b128 v[32:35], %[ta]\nds_read_b128 v[36:39], %[ta] offset:64\n" + "v_mfma_scale_f32_16x16x128_f8f6f4 a[0:3], v[22:25], v[8:15], a[0:3], " + "v7, v16 op_sel_hi:[0,0,0] cbsz:4\n" + "s_cmpk_lt_i32 s13, %[im1]\ns_cbranch_scc0 HAZTAIL1%=\n" + // --- consume bank 1, prefetch into bank 0 --- + "s_waitcnt lgkmcnt(0)\n" + "v_add_u32_e32 %[wa], 64, %[wa]\nv_add_u32_e32 %[wsa], 1, %[wsa]\n" + "v_add_u32_e32 %[ta], 0x80, %[ta]\ns_add_i32 s13, s13, 1\n" + "v_add_u32_e32 v17, s13, %[tsa]\nds_read_u8 v16, v17\n" + "ds_read_b128 v[22:25], %[wa]\nds_read_u8 v7, %[wsa]\n" + "ds_read_b128 v[8:11], %[ta]\nds_read_b128 v[12:15], %[ta] offset:64\n" + "v_mfma_scale_f32_16x16x128_f8f6f4 a[0:3], v[26:29], v[32:39], a[0:3], " + "v18, v19 op_sel_hi:[0,0,0] cbsz:4\n" + "s_cmpk_lt_i32 s13, %[im1]\ns_cbranch_scc1 HAZFIX%=\n" + // Exited after a bank-1 iteration: final operands live in bank 0. + "s_waitcnt lgkmcnt(0)\n" + "v_mfma_scale_f32_16x16x128_f8f6f4 a[0:3], v[22:25], v[8:15], a[0:3], " + "v7, v16 op_sel_hi:[0,0,0] cbsz:4\n" + "s_branch HAZDONE%=\n" + "HAZTAIL1%=:\n" + // Exited after a bank-0 iteration: final operands live in bank 1. + "s_waitcnt lgkmcnt(0)\n" + "v_mfma_scale_f32_16x16x128_f8f6f4 a[0:3], v[26:29], v[32:39], a[0:3], " + "v18, v19 op_sel_hi:[0,0,0] cbsz:4\n" + "HAZDONE%=:\n" + // 32 clocks: matches the scaled-MFMA latency on CDNA4. + "s_nop 15\ns_nop 15\n" + "v_accvgpr_read_b32 %[a0], a0\nv_accvgpr_read_b32 %[a1], a1\n" + "v_accvgpr_read_b32 %[a2], a2\nv_accvgpr_read_b32 %[a3], a3\n" + : HAZ_OUTS:HAZ_INS + : HAZ_CLOBBER); + HAZ_EPILOGUE +} + +// --------------------------------------------------------------------------- +// REFERENCE: no overlap whatsoever. Slow, obviously correct, used as truth. +// Every operand is fully resident before the MFMA issues, and the MFMA is +// fully retired before anything touches its registers or accumulator. +// --------------------------------------------------------------------------- +__global__ void mfma_reference(float *out, + uint8_t const *Ad, + uint8_t const *As, + uint8_t const *Bd, + uint8_t const *Bs) { + HAZ_PROLOGUE + asm volatile( + "v_accvgpr_write_b32 a0, 0\nv_accvgpr_write_b32 a1, 0\n" + "v_accvgpr_write_b32 a2, 0\nv_accvgpr_write_b32 a3, 0\n" + "s_mov_b32 s13, 0\n" + "HAZREF%=:\n" + "ds_read_b128 v[22:25], %[wa]\nds_read_u8 v7, %[wsa]\n" + "ds_read_b128 v[8:11], %[ta]\nds_read_b128 v[12:15], %[ta] offset:64\n" + "v_add_u32_e32 v17, s13, %[tsa]\nds_read_u8 v16, v17\n" + "s_waitcnt lgkmcnt(0)\ns_nop 15\ns_nop 15\n" + "v_mfma_scale_f32_16x16x128_f8f6f4 a[0:3], v[22:25], v[8:15], a[0:3], " + "v7, v16 op_sel_hi:[0,0,0] cbsz:4\n" + "s_nop 15\ns_nop 15\ns_nop 15\ns_nop 15\n" + "v_add_u32_e32 %[wa], 64, %[wa]\nv_add_u32_e32 %[wsa], 1, %[wsa]\n" + "v_add_u32_e32 %[ta], 0x80, %[ta]\ns_add_i32 s13, s13, 1\n" + "s_cmpk_lt_i32 s13, %[it]\ns_cbranch_scc1 HAZREF%=\n" + "s_nop 15\ns_nop 15\ns_nop 15\ns_nop 15\n" + "v_accvgpr_read_b32 %[a0], a0\nv_accvgpr_read_b32 %[a1], a1\n" + "v_accvgpr_read_b32 %[a2], a2\nv_accvgpr_read_b32 %[a3], a3\n" + : HAZ_OUTS + : [tsa] "v"(ts_addr), [it] "n"(MFMA_ITERS) + : HAZ_CLOBBER); + HAZ_EPILOGUE +} + +typedef void (*hazard_kernel_t)(float *, + uint8_t const *, + uint8_t const *, + uint8_t const *, + uint8_t const *); + +int main(int argc, char **argv) { + int const launches = (argc > 1) ? std::atoi(argv[1]) : 300; + int const grid = 64; + + size_t const szA = 64 * MFMA_ITERS, szAs = MFMA_ITERS; + size_t const szB = 128 * MFMA_ITERS, szBs = MFMA_ITERS; + std::vector hA(szA), hAs(szAs), hB(szB), hBs(szBs); + + // Deterministic inputs. A is FP4 (e2m1) packed two per byte -- every nibble + // is finite. B is FP8 e4m3, where exponent 0xF with a nonzero mantissa is + // NaN, so the exponent is held in [1,13] to keep every byte finite and + // normal. Random bytes here would decode to NaN and mask the comparison. + unsigned state = 12345u; + auto next = [&]() { + state = state * 1664525u + 1013904223u; + return state >> 16; + }; + for (size_t i = 0; i < szA; i++) { + hA[i] = (uint8_t)(next() & 0xff); + } + for (size_t i = 0; i < szB; i++) { + unsigned r = next(); + hB[i] = (uint8_t)(((r & 1) << 7) | ((1 + ((r >> 1) % 13)) << 3) | + ((r >> 8) & 0x7)); + } + // E8M0 scales clustered near 1.0 (bias 127). + for (size_t i = 0; i < szAs; i++) { + hAs[i] = (uint8_t)(124 + (next() % 7)); + } + for (size_t i = 0; i < szBs; i++) { + hBs[i] = (uint8_t)(124 + (next() % 7)); + } + + uint8_t *dA, *dAs, *dB, *dBs; + float *dOut; + HIP_CHECK(hipMalloc(&dA, szA)); + HIP_CHECK(hipMalloc(&dAs, szAs)); + HIP_CHECK(hipMalloc(&dB, szB)); + HIP_CHECK(hipMalloc(&dBs, szBs)); + HIP_CHECK(hipMalloc(&dOut, (size_t)grid * 1024 * sizeof(float))); + HIP_CHECK(hipMemcpy(dA, hA.data(), szA, hipMemcpyHostToDevice)); + HIP_CHECK(hipMemcpy(dAs, hAs.data(), szAs, hipMemcpyHostToDevice)); + HIP_CHECK(hipMemcpy(dB, hB.data(), szB, hipMemcpyHostToDevice)); + HIP_CHECK(hipMemcpy(dBs, hBs.data(), szBs, hipMemcpyHostToDevice)); + + size_t const smem = 64 * MFMA_ITERS + 64 + 128 * MFMA_ITERS + 64 + 256; + + std::vector host((size_t)grid * 1024); + auto launch = [&](hazard_kernel_t k, int blocks) { + HIP_CHECK(hipMemset(dOut, 0, (size_t)blocks * 1024 * sizeof(float))); + hipLaunchKernelGGL( + k, dim3(blocks), dim3(256), smem, 0, dOut, dA, dAs, dB, dBs); + HIP_CHECK(hipGetLastError()); + HIP_CHECK(hipDeviceSynchronize()); + HIP_CHECK(hipMemcpy(host.data(), + dOut, + (size_t)blocks * 1024 * sizeof(float), + hipMemcpyDeviceToHost)); + }; + + // Ground truth from the non-overlapping reference. + launch(mfma_reference, 1); + std::vector truth(host.begin(), host.begin() + 1024); + for (int i = 0; i < 1024; i++) { + if (!std::isfinite(truth[i])) { + std::printf("FAIL: reference produced a non-finite value at %d\n", i); + return 2; + } + } + + std::printf("gfx950 scaled-MFMA pipeline hazard regression\n"); + std::printf( + " iterations %d, grid %d, launches %d\n", MFMA_ITERS, grid, launches); + std::printf(" reference lane0 accumulator: %.4f\n\n", truth[0]); + + struct Variant { + char const *name; + hazard_kernel_t kernel; + bool must_match; + } const variants[] = { + {"reference (no overlap)", mfma_reference, true}, + {"fixed (banked + 32clk)", mfma_fixed, true}, + {"broken (historical)", mfma_broken, false}, + }; + + int failures = 0; + for (auto const &v : variants) { + long wrong_blocks = 0, total_blocks = 0; + double worst_rel = 0.0; + for (int t = 0; t < launches; t++) { + launch(v.kernel, grid); + for (int b = 0; b < grid; b++) { + total_blocks++; + bool bad = false; + for (int i = 0; i < 1024; i++) { + float got = host[(size_t)b * 1024 + i]; + if (got != truth[i]) { + bad = true; + if (std::fabs(truth[i]) > 1e-6) { + worst_rel = std::max(worst_rel, + std::fabs((double)got - truth[i]) / + std::fabs(truth[i])); + } + } + } + if (bad) { + wrong_blocks++; + } + } + } + std::printf( + " %-28s %7ld/%ld blocks differ", v.name, wrong_blocks, total_blocks); + if (wrong_blocks) { + std::printf(" (worst rel err %.2f%%)", 100.0 * worst_rel); + } + std::printf("\n"); + + if (v.must_match && wrong_blocks != 0) { + failures++; + } + // If the broken variant stops reproducing, this test has gone blind -- + // a toolchain or scheduler change may be hiding the hazard rather than + // fixing it. Report loudly instead of passing quietly. + if (!v.must_match && wrong_blocks == 0) { + std::printf( + " WARNING: the known-broken schedule now matches the " + "reference.\n" + " This test can no longer prove it detects the " + "hazard.\n" + " Inspect the emitted ISA before trusting a pass.\n"); + } + } + + std::printf("\n"); + if (failures) { + std::printf("FAIL: a hazard-free schedule did not reproduce the " + "reference.\n" + " Compare the emitted ISA against the mfma_fixed kernel " + "in this file:\n" + " hipcc --offload-arch=gfx950 -S " + "test_mfma_pipeline_hazards.hip -o -\n"); + return 1; + } + std::printf("PASS: hazard-free schedules are bit-exact across all " + "launches.\n"); + return 0; +}