From cda967a18c18231e3a019017042654554e8178d3 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Wed, 29 Jul 2026 04:19:19 -0400 Subject: [PATCH 01/85] Add Qwen3.6-35B-A3B text frontend --- docs/nexn2_usage.md | 13 +- docs/qwen36_moe_usage.md | 162 +++++++++++ flash_rt/api.py | 12 +- flash_rt/frontends/torch/_nexn2_rtx_decode.py | 24 +- .../torch/_nexn2_rtx_nvfp4_weights.py | 3 +- flash_rt/frontends/torch/nexn2_rtx.py | 21 +- flash_rt/frontends/torch/qwen36_moe_rtx.py | 255 ++++++++++++++++++ flash_rt/hardware/__init__.py | 7 +- tests/test_qwen36_moe_gpu.py | 69 +++++ tests/test_qwen36_moe_smoke.py | 235 ++++++++++++++++ 10 files changed, 783 insertions(+), 18 deletions(-) create mode 100644 docs/qwen36_moe_usage.md create mode 100644 flash_rt/frontends/torch/qwen36_moe_rtx.py create mode 100644 tests/test_qwen36_moe_gpu.py create mode 100644 tests/test_qwen36_moe_smoke.py diff --git a/docs/nexn2_usage.md b/docs/nexn2_usage.md index 671f5928..dd2905aa 100644 --- a/docs/nexn2_usage.md +++ b/docs/nexn2_usage.md @@ -99,7 +99,10 @@ Nexn2TorchFrontendRtx( Methods: `set_prompt(text)`, `infer() -> (1, S, vocab)`, `generate(max_new_tokens) -> list[int]`, `tokenizer`, `latency_records` (list[float], per `infer()`). -Env knobs: `FLASHRT_NEXN2_PREFILL_CHUNK` (chunked-prefill block size, default 8192; 0 disables), `FLASHRT_NEXN2_GRAPH_CACHE_MAX` (decode CUDA-graph LRU cap, default 256). +Env knobs: `FLASHRT_QWEN35MOE_PREFILL_CHUNK` (chunked-prefill block size, +default 8192; 0 disables) and `FLASHRT_QWEN35MOE_GRAPH_CACHE_MAX` (decode +CUDA-graph LRU cap, default 256). The older `FLASHRT_NEXN2_PREFILL_CHUNK` and +`FLASHRT_NEXN2_GRAPH_CACHE_MAX` names remain compatible aliases. ## 5. Performance @@ -185,8 +188,8 @@ logits that seed decode are stable: activations no longer bound it; the residual limit on a 32 GB card is the bf16 KV cache (~5.4 GB at 256k over the 10 full-attn layers) alongside the ~22 GB of weights. `generate()` auto-chunks; tune the block via - `FLASHRT_NEXN2_PREFILL_CHUNK` (default 8192; lower trades a little throughput - for headroom). The raw `infer()` path returns all-position logits and is a - separate single-pass validation tool, capped at ~4k by the `(S, 248320)` - logit tensor — use `generate()` for long context. + `FLASHRT_QWEN35MOE_PREFILL_CHUNK` (default 8192; lower trades a little + throughput for headroom). The raw `infer()` path returns all-position logits + and is a separate single-pass validation tool, capped at ~4k by the + `(S, 248320)` logit tensor — use `generate()` for long context. * Text LLM only — not wired into the `load_model` / VLA `predict()` API. diff --git a/docs/qwen36_moe_usage.md b/docs/qwen36_moe_usage.md new file mode 100644 index 00000000..937c3905 --- /dev/null +++ b/docs/qwen36_moe_usage.md @@ -0,0 +1,162 @@ +# Qwen3.6-35B-A3B text inference + +FlashRT runs the language backbone from the official +`Qwen/Qwen3.6-35B-A3B` BF16 checkpoint on an RTX 5090. The checkpoint uses the +same `qwen3_5_moe` text architecture as Nex-N2-mini, so both models share the +same weight loader, prefill, attention, MoE, recurrent-state, and CUDA Graph +decode implementation. + +This entry is text-only. It does not load the vision tower and it validates but +does not execute the checkpoint's MTP head. Image/video input and speculative +decode are not part of this interface. + +## Requirements + +| | | +|---|---| +| Checkpoint | `Qwen/Qwen3.6-35B-A3B` BF16 safetensors | +| Hardware | RTX 5090 / SM120 | +| GPU memory | 32 GB | +| Framework | PyTorch | +| Runtime quantization | NVFP4 | +| Build flag | `-DFLASHRT_ENABLE_QWEN35MOE=ON` | + +Configure and build the gated `qwen3_5_moe` kernels: + +```bash +cmake -S . -B build -DFLASHRT_ENABLE_QWEN35MOE=ON +cmake --build build -j +pip install -e ".[torch]" +``` + +## Usage + +```python +from flash_rt.frontends.torch.qwen36_moe_rtx import ( + Qwen36MoeTextFrontendRtx, +) + +frontend = Qwen36MoeTextFrontendRtx( + "/models/Qwen3.6-35B-A3B", + device="cuda:0", + max_seq=4096, + kernelized=True, + quant_scope="experts", +) +frontend.set_prompt("Explain why deterministic reductions matter.") +token_ids = frontend.generate(max_new_tokens=64) +print(frontend.tokenizer.decode(token_ids)) +``` + +`set_prompt()` accepts already-rendered text. For chat requests, render the +checkpoint's own template first: + +```python +messages = [{"role": "user", "content": "Write a CUDA reduction checklist."}] +prompt = frontend.tokenizer.apply_chat_template( + messages, + tokenize=False, + add_generation_prompt=True, +) +frontend.set_prompt(prompt) +``` + +The direct frontend is intentional. `flash_rt.load_model()` wraps VLA models +with a `predict(images, ...)` API and therefore redirects +`config="qwen36_moe"` to the class above. + +## Checkpoint contract + +Before allocating GPU memory, the frontend checks: + +- the exact 40-layer `qwen3_5_moe` text geometry; +- the 30 linear-attention / 10 full-attention schedule; +- all 693 text-backbone tensors consumed by the shared pipeline; +- all 19 official MTP tensors; +- every safetensors shard referenced by the index. + +Extra vision tensors are allowed and ignored by this text-only entry. The +validation can be run without loading model weights: + +```bash +PYTHONPATH=. python - <<'PY' +from flash_rt.frontends.torch.qwen36_moe_rtx import ( + validate_qwen36_moe_checkpoint, +) +print(validate_qwen36_moe_checkpoint("/models/Qwen3.6-35B-A3B")) +PY +``` + +## Runtime controls + +The shared architecture uses: + +- `FLASHRT_QWEN35MOE_PREFILL_CHUNK` — chunked-prefill block size, default + `8192`; `0` disables chunking. +- `FLASHRT_QWEN35MOE_GRAPH_CACHE_MAX` — decode CUDA Graph LRU capacity, + default `256`. + +The older `FLASHRT_NEXN2_PREFILL_CHUNK` and +`FLASHRT_NEXN2_GRAPH_CACHE_MAX` names remain compatible aliases. + +## Validation + +The repository smoke test is checkpoint-independent: + +```bash +PYTHONPATH=. PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 \ + pytest -q -p no:cacheprovider tests/test_qwen36_moe_smoke.py +``` + +Set `FLASHRT_QWEN36_MOE_CKPT_DIR` to include the official checkpoint contract +test. Performance and precision numbers must be measured on Qwen3.6 weights; +Nex-N2-mini measurements are not interchangeable even though the compute +pipeline is shared. + +The optional real-model test loads the checkpoint, checks finite logits, and +compares cold and warm CUDA Graph generation with an official Transformers +BF16 greedy-token fixture: + +```bash +FLASHRT_QWEN36_MOE_CKPT_DIR=/models/Qwen3.6-35B-A3B \ +PYTHONPATH=. PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 \ +pytest -q -p no:cacheprovider tests/test_qwen36_moe_gpu.py +``` + +First-light data on an RTX 5090, PyTorch 2.9.1, CUDA 12.8 runtime, and the +official BF16 checkpoint: + +| Measurement | Result | +|---|---:| +| Runtime weight load | 47.96 s | +| Resident allocated memory after load | 21.44 GiB | +| Peak allocated memory during load | 22.94 GiB | +| First 21-token prefill, including warmup | 230.95 ms | +| Subsequent 20–45-token prefill | 28.99–35.12 ms | +| 64-token prompt, 32-token eager decode | 48.14 tok/s | +| 64-token prompt, 32-token warm CUDA Graph decode | 195.49 tok/s | + +The eager, first-capture, and warm-graph runs produced the same 32 token IDs. +These numbers are a first-light correctness run, not a context-length sweep. + +Four chat prompts from 12 to 45 tokens were also compared with the official +Transformers BF16 implementation: + +| Precision check | Result | +|---|---:| +| Last-token logit cosine, minimum | 0.95635 | +| Last-token logit cosine, mean | 0.96455 | +| First-token argmax matches | 4 / 4 | +| Greedy generation matches | 64 / 64 tokens | + +The logit cosine is lower than the Nex-N2-mini measurement, but the tested +greedy sequences were token-exact for 16 generated tokens on all four prompts. + +## Limitations + +- Text only; the vision tower is not loaded. +- Greedy decode only on the kernelized path. +- The MTP tensors are validated but not loaded, so speculative decode is not + enabled. +- Only the BF16 source checkpoint with runtime NVFP4 conversion is supported. +- SM120 only. diff --git a/flash_rt/api.py b/flash_rt/api.py index 4e4a64d7..4497d568 100644 --- a/flash_rt/api.py +++ b/flash_rt/api.py @@ -456,11 +456,12 @@ def load_model(checkpoint, framework="torch", num_views=2, autotune=3, "Supported: pi0, pi05, llm, mllm") elif config not in ("pi05", "groot", "groot_n17", "pi0", "pi0fast", "motus", "wan22_ti2v_5b", "cosmos3_video", - "cosmos3_edge", "nexn2"): + "cosmos3_edge", "nexn2", "qwen36_moe"): raise ValueError( f"Unknown config: {config}. " f"Supported: pi05, groot, groot_n17, pi0, pi0fast, motus, " - f"wan22_ti2v_5b, cosmos3_video, cosmos3_edge, nexn2") + f"wan22_ti2v_5b, cosmos3_video, cosmos3_edge, nexn2, " + f"qwen36_moe") if framework not in ("torch", "jax", "jetson_pi"): raise ValueError( f"Unknown framework: {framework}. Supported: torch, jax, jetson_pi") @@ -538,6 +539,13 @@ def load_model(checkpoint, framework="torch", num_views=2, autotune=3, " from flash_rt.frontends.torch.nexn2_rtx import " "Nexn2TorchFrontendRtx\n" "See docs/nexn2_usage.md.") + if config == "qwen36_moe": + raise NotImplementedError( + "config='qwen36_moe' is a text LLM and is not served through " + "load_model's VLA wrapper. Construct it directly:\n" + " from flash_rt.frontends.torch.qwen36_moe_rtx import " + "Qwen36MoeTextFrontendRtx\n" + "See docs/qwen36_moe_usage.md.") from flash_rt.hardware import detect_arch, resolve_pipeline_class arch = detect_arch() if hardware == "auto" else hardware diff --git a/flash_rt/frontends/torch/_nexn2_rtx_decode.py b/flash_rt/frontends/torch/_nexn2_rtx_decode.py index 1a2cbaf7..7ff6f8f0 100644 --- a/flash_rt/frontends/torch/_nexn2_rtx_decode.py +++ b/flash_rt/frontends/torch/_nexn2_rtx_decode.py @@ -34,6 +34,13 @@ from flash_rt.frontends.torch._nexn2_rtx_nvfp4_weights import _sf_swz_bytes from flash_rt.hardware.rtx.attn_backend_nexn2 import RtxFlashAttnBackendNexn2 + +def _qwen35moe_env(name: str, default: str) -> str: + generic = f"FLASHRT_QWEN35MOE_{name}" + legacy = f"FLASHRT_NEXN2_{name}" + return os.environ.get(generic, os.environ.get(legacy, default)) + + # Prompt length at/above which the batched prefill wins over the per-token loop # (batched has a fixed forward overhead; below this the loop's lower latency # wins). See Nexn2DecodeState.batched_prefill. @@ -222,7 +229,7 @@ def __init__(self, handles, max_seq, device): # used pos once over the cap; 0/negative disables the bound (legacy). self._graphs = collections.OrderedDict() self.graph_cache_max = int( - os.environ.get('FLASHRT_NEXN2_GRAPH_CACHE_MAX', '256')) + _qwen35moe_env("GRAPH_CACHE_MAX", "256")) self._snap_lin = [torch.empty_like(t) for t in self.lin_state] self._snap_conv = [torch.empty_like(t) for t in self.lin_conv_state] # Pre-allocated KV snapshot rows (one [.,1,.] slice each) reused every @@ -265,7 +272,7 @@ def __init__(self, handles, max_seq, device): # the ~16k single-pass ceiling on a 32 GB card. A multiple of the WY # chunk (64). 0 disables (always single-pass). self.prefill_chunk = int( - os.environ.get('FLASHRT_NEXN2_PREFILL_CHUNK', '8192')) + _qwen35moe_env("PREFILL_CHUNK", "8192")) def reset(self): for s in self.lin_state: @@ -465,7 +472,18 @@ def _moe_layer_decode(h, ld, state, fvk, device): inter.data_ptr(), dn_p.data_ptr(), dn_s.data_ptr(), dn_a.data_ptr(), idx.data_ptr(), d_dn.data_ptr(), TOPK, n_dn, INTER, INTER, dn_p[0].numel(), dn_s[0].numel(), s) - out = (tw_row.float() @ d_dn.float()).unsqueeze(0) # (1, n_dn) weighted sum + # Fixed-order weighted sum. The generic torch matmul may choose a + # reduction whose accumulation order changes between launches, which can + # flip a later greedy decision when two logits are nearly tied. + if 'decode_topk_rows' not in ld: + ld['decode_topk_rows'] = torch.arange( + TOPK, dtype=torch.int32, device=device) + out = torch.empty(n_dn, dtype=torch.float32, device=device) + fvk.moe_weighted_sum_sm120_bf16( + d_dn.data_ptr(), ld['decode_topk_rows'].data_ptr(), + tw_row.data_ptr(), out.data_ptr(), + 1, TOPK, n_dn, n_dn, s) + out = out.unsqueeze(0) if fused_rs: # already projected above sg, su = sg_f, su_f diff --git a/flash_rt/frontends/torch/_nexn2_rtx_nvfp4_weights.py b/flash_rt/frontends/torch/_nexn2_rtx_nvfp4_weights.py index b68e7122..985162aa 100644 --- a/flash_rt/frontends/torch/_nexn2_rtx_nvfp4_weights.py +++ b/flash_rt/frontends/torch/_nexn2_rtx_nvfp4_weights.py @@ -57,7 +57,8 @@ def _open_shards(ckpt_dir: str): idx_path = os.path.join(ckpt_dir, 'model.safetensors.index.json') if not os.path.isfile(idx_path): raise RuntimeError( - f'Nex-N2 ckpt missing model.safetensors.index.json: {ckpt_dir!r}') + 'qwen3_5_moe checkpoint missing ' + f'model.safetensors.index.json: {ckpt_dir!r}') wmap = json.load(open(idx_path))['weight_map'] handles_d = {} for shard in set(wmap.values()): diff --git a/flash_rt/frontends/torch/nexn2_rtx.py b/flash_rt/frontends/torch/nexn2_rtx.py index 91cd9448..b1791552 100644 --- a/flash_rt/frontends/torch/nexn2_rtx.py +++ b/flash_rt/frontends/torch/nexn2_rtx.py @@ -34,22 +34,25 @@ ) -def _require_kernels(fvk) -> None: +def _require_kernels( + fvk, *, model_label: str = "Nex-N2", + usage_doc: str = "docs/nexn2_usage.md") -> None: """Raise a clear RuntimeError if the gated qwen3_5_moe kernels or the FA2 module are missing (build was not configured with -DFLASHRT_ENABLE_QWEN35MOE=ON, or flash_rt_fa2 is absent).""" missing = [s for s in _REQUIRED_FVK if not hasattr(fvk, s)] if missing: raise RuntimeError( - "Nex-N2 kernelized path needs the qwen3_5_moe SM120 kernels, which " + f"{model_label} kernelized path needs the qwen3_5_moe SM120 " + "kernels, which " "are absent from flash_rt_kernels (missing: " f"{', '.join(missing)}). Rebuild on an SM120 toolchain with " - "-DFLASHRT_ENABLE_QWEN35MOE=ON. See docs/nexn2_usage.md.") + f"-DFLASHRT_ENABLE_QWEN35MOE=ON. See {usage_doc}.") try: from flash_rt import flash_rt_fa2 as _fa2 except Exception as e: # pragma: no cover raise RuntimeError( - "Nex-N2 full attention needs the vendored FA2 module " + f"{model_label} full attention needs the vendored FA2 module " "(flash_rt_fa2), which failed to import. Build with FA2 enabled " "(ENABLE_FA2, auto-on for SM120).") from e fa2_missing = [s for s in ('fwd_bf16', 'fwd_bf16_causal') @@ -64,6 +67,9 @@ def _require_kernels(fvk) -> None: class Nexn2TorchFrontendRtx: """Nex-N2-mini inference frontend (PyTorch + RTX SM120).""" + _MODEL_LABEL = "Nex-N2" + _USAGE_DOC = "docs/nexn2_usage.md" + def __init__(self, checkpoint_path: str, *, device: str = 'cuda:0', max_seq: int = 2048, @@ -148,7 +154,12 @@ def _build_kernelized_nvfp4(self) -> None: extract_weights_nexn2_nvfp4, ) - _require_kernels(fvk) # fail fast before loading the 35B ckpt + # Fail before loading the 35B checkpoint. + _require_kernels( + fvk, + model_label=self._MODEL_LABEL, + usage_doc=self._USAGE_DOC, + ) self._tokenizer = AutoTokenizer.from_pretrained(self.checkpoint_path) self._fvk = fvk diff --git a/flash_rt/frontends/torch/qwen36_moe_rtx.py b/flash_rt/frontends/torch/qwen36_moe_rtx.py new file mode 100644 index 00000000..df98f705 --- /dev/null +++ b/flash_rt/frontends/torch/qwen36_moe_rtx.py @@ -0,0 +1,255 @@ +"""Qwen3.6-35B-A3B text inference on RTX SM120. + +The language backbone is the same ``qwen3_5_moe`` architecture used by +Nex-N2-mini, so this frontend reuses that implementation. The official +Qwen3.6 checkpoint also contains a vision tower and an MTP head; this entry +validates those weights but intentionally exposes only text prefill and greedy +decode. Vision and speculative decoding are separate integration surfaces. +""" + +from __future__ import annotations + +import json +import os +from typing import Any + +from flash_rt.frontends.torch.nexn2_rtx import Nexn2TorchFrontendRtx + + +_EXPECTED_LAYER_TYPES = tuple( + "full_attention" if (i + 1) % 4 == 0 else "linear_attention" + for i in range(40) +) + +_EXPECTED_TEXT_CONFIG = { + "model_type": "qwen3_5_moe_text", + "num_hidden_layers": 40, + "hidden_size": 2048, + "vocab_size": 248320, + "num_attention_heads": 16, + "num_key_value_heads": 2, + "head_dim": 256, + "num_experts": 256, + "num_experts_per_tok": 8, + "linear_num_key_heads": 16, + "linear_num_value_heads": 32, + "linear_key_head_dim": 128, + "linear_value_head_dim": 128, + "full_attention_interval": 4, + "mtp_num_hidden_layers": 1, +} + +_MTP_KEYS = { + "mtp.fc.weight", + "mtp.norm.weight", + "mtp.pre_fc_norm_embedding.weight", + "mtp.pre_fc_norm_hidden.weight", + "mtp.layers.0.input_layernorm.weight", + "mtp.layers.0.post_attention_layernorm.weight", + "mtp.layers.0.self_attn.q_proj.weight", + "mtp.layers.0.self_attn.k_proj.weight", + "mtp.layers.0.self_attn.v_proj.weight", + "mtp.layers.0.self_attn.o_proj.weight", + "mtp.layers.0.self_attn.q_norm.weight", + "mtp.layers.0.self_attn.k_norm.weight", + "mtp.layers.0.mlp.gate.weight", + "mtp.layers.0.mlp.experts.gate_up_proj", + "mtp.layers.0.mlp.experts.down_proj", + "mtp.layers.0.mlp.shared_expert.gate_proj.weight", + "mtp.layers.0.mlp.shared_expert.up_proj.weight", + "mtp.layers.0.mlp.shared_expert.down_proj.weight", + "mtp.layers.0.mlp.shared_expert_gate.weight", +} + + +def _required_text_keys(layer_types: tuple[str, ...]) -> set[str]: + keys = { + "lm_head.weight", + "model.language_model.embed_tokens.weight", + "model.language_model.norm.weight", + } + mlp_suffixes = ( + "mlp.gate.weight", + "mlp.experts.gate_up_proj", + "mlp.experts.down_proj", + "mlp.shared_expert.gate_proj.weight", + "mlp.shared_expert.up_proj.weight", + "mlp.shared_expert.down_proj.weight", + "mlp.shared_expert_gate.weight", + ) + linear_suffixes = ( + "linear_attn.in_proj_qkv.weight", + "linear_attn.in_proj_z.weight", + "linear_attn.in_proj_a.weight", + "linear_attn.in_proj_b.weight", + "linear_attn.conv1d.weight", + "linear_attn.A_log", + "linear_attn.dt_bias", + "linear_attn.norm.weight", + "linear_attn.out_proj.weight", + ) + full_suffixes = ( + "self_attn.q_proj.weight", + "self_attn.k_proj.weight", + "self_attn.v_proj.weight", + "self_attn.o_proj.weight", + "self_attn.q_norm.weight", + "self_attn.k_norm.weight", + ) + for i, layer_type in enumerate(layer_types): + prefix = f"model.language_model.layers.{i}." + keys.add(prefix + "input_layernorm.weight") + keys.add(prefix + "post_attention_layernorm.weight") + keys.update(prefix + suffix for suffix in mlp_suffixes) + suffixes = ( + full_suffixes + if layer_type == "full_attention" + else linear_suffixes + ) + keys.update(prefix + suffix for suffix in suffixes) + return keys + + +def validate_qwen36_moe_checkpoint(checkpoint_path: str) -> dict[str, Any]: + """Validate the official BF16 checkpoint before allocating GPU weights.""" + checkpoint_path = os.path.abspath(os.fspath(checkpoint_path)) + config_path = os.path.join(checkpoint_path, "config.json") + index_path = os.path.join( + checkpoint_path, "model.safetensors.index.json") + + for path in (config_path, index_path): + if not os.path.isfile(path): + raise FileNotFoundError( + f"Qwen3.6-35B-A3B checkpoint is missing {path!r}") + + with open(config_path, "r", encoding="utf-8") as f: + config = json.load(f) + if config.get("model_type") != "qwen3_5_moe": + raise ValueError( + "Qwen3.6-35B-A3B text frontend requires " + "model_type='qwen3_5_moe'; " + f"got {config.get('model_type')!r} in {config_path}") + + text_config = config.get("text_config") + if not isinstance(text_config, dict): + raise ValueError(f"missing text_config object in {config_path}") + + mismatches = [] + for name, expected in _EXPECTED_TEXT_CONFIG.items(): + actual = text_config.get(name) + if actual != expected: + mismatches.append(f"{name}={actual!r} (expected {expected!r})") + layer_types = tuple(text_config.get("layer_types") or ()) + if layer_types != _EXPECTED_LAYER_TYPES: + mismatches.append( + "layer_types does not match the 30-linear/10-full attention " + "qwen3_5_moe schedule") + if mismatches: + raise ValueError( + "checkpoint is not compatible with the Qwen3.6-35B-A3B " + "SM120 text pipeline: " + "; ".join(mismatches)) + + with open(index_path, "r", encoding="utf-8") as f: + index = json.load(f) + weight_map = index.get("weight_map") + if not isinstance(weight_map, dict): + raise ValueError(f"missing weight_map object in {index_path}") + + required_text = _required_text_keys(layer_types) + missing_text = sorted(required_text.difference(weight_map)) + if missing_text: + preview = ", ".join(missing_text[:8]) + if len(missing_text) > 8: + preview += f", ... ({len(missing_text)} missing)" + raise ValueError( + "Qwen3.6-35B-A3B checkpoint is missing text backbone tensors: " + + preview) + + missing_mtp = sorted(_MTP_KEYS.difference(weight_map)) + if missing_mtp: + preview = ", ".join(missing_mtp[:8]) + if len(missing_mtp) > 8: + preview += f", ... ({len(missing_mtp)} missing)" + raise ValueError( + "Qwen3.6-35B-A3B checkpoint is missing its MTP tensor group: " + + preview) + + shards = sorted(set(weight_map.values())) + bad_shards = [] + for shard in shards: + path = os.path.join(checkpoint_path, shard) + if not os.path.isfile(path) or os.path.getsize(path) == 0: + bad_shards.append(shard) + if bad_shards: + preview = ", ".join(bad_shards[:8]) + if len(bad_shards) > 8: + preview += f", ... ({len(bad_shards)} missing or empty)" + raise FileNotFoundError( + "Qwen3.6-35B-A3B checkpoint has missing or empty shards: " + + preview) + + return { + "checkpoint_path": checkpoint_path, + "text_tensor_count": len(required_text), + "mtp_tensor_count": len(_MTP_KEYS), + "vision_tensor_count": sum( + ".visual." in name for name in weight_map), + "tensor_count": len(weight_map), + "shard_count": len(shards), + } + + +class Qwen36MoeTextFrontendRtx(Nexn2TorchFrontendRtx): + """Qwen3.6-35B-A3B text-only frontend for RTX SM120.""" + + _MODEL_LABEL = "Qwen3.6-35B-A3B text" + _USAGE_DOC = "docs/qwen36_moe_usage.md" + + def __init__(self, checkpoint_path: str, *, + device: str = "cuda:0", + max_seq: int = 2048, + quant: str = "nvfp4", + kernelized: bool = False, + quant_scope: str = "experts") -> None: + if quant != "nvfp4": + raise NotImplementedError( + f"quant={quant!r} is not implemented; only 'nvfp4' is " + "supported") + contract = validate_qwen36_moe_checkpoint(checkpoint_path) + super().__init__( + checkpoint_path, + device=device, + max_seq=max_seq, + quant=quant, + kernelized=kernelized, + quant_scope=quant_scope, + ) + self._checkpoint_contract = contract + + def generate(self, max_new_tokens: int, *, do_sample: bool = False): + """Generate tokens with the shared qwen3_5_moe CUDA Graph path.""" + if not self._kernelized: + return super().generate( + max_new_tokens=max_new_tokens, do_sample=do_sample) + if self._prompt_ids is None: + raise ValueError("call set_prompt(...) before generate()") + if do_sample: + raise NotImplementedError( + "the Qwen3.6-35B-A3B kernelized path supports greedy " + "decoding only") + + from flash_rt.frontends.torch._nexn2_rtx_decode import ( + Nexn2DecodeState, + generate_greedy_graph, + ) + + if self._decode_state is None: + self._decode_state = Nexn2DecodeState( + self._weights, self._user_max_seq, self.device) + return generate_greedy_graph( + self._decode_state, + self._prompt_ids, + max_new_tokens, + self._fvk, + self.device, + ) diff --git a/flash_rt/hardware/__init__.py b/flash_rt/hardware/__init__.py index 4ee54985..3b12332c 100644 --- a/flash_rt/hardware/__init__.py +++ b/flash_rt/hardware/__init__.py @@ -156,10 +156,13 @@ def detect_arch() -> str: # (-DFLASHRT_ENABLE_QWEN35MOE=ON). Registered here for discoverability / # resolve_pipeline_class, but the frontend exposes an LLM surface # (infer()->logits, generate_greedy) rather than the VLA predict(images) - # API, so it is used via direct instantiation of Nexn2TorchFrontendRtx - # (see docs/nexn2_usage.md) rather than load_model's VLAModel wrapper. + # API, so these are used via direct frontend construction rather than + # load_model's VLAModel wrapper. ("nexn2", "torch", "rtx_sm120"): ("flash_rt.frontends.torch.nexn2_rtx", "Nexn2TorchFrontendRtx"), + ("qwen36_moe", "torch", "rtx_sm120"): + ("flash_rt.frontends.torch.qwen36_moe_rtx", + "Qwen36MoeTextFrontendRtx"), # ── Pi0-FAST ── (SM120 runtime fork inside pipeline, no AttentionBackend protocol.) ("pi0fast", "torch", "thor"): diff --git a/tests/test_qwen36_moe_gpu.py b/tests/test_qwen36_moe_gpu.py new file mode 100644 index 00000000..c16be2d3 --- /dev/null +++ b/tests/test_qwen36_moe_gpu.py @@ -0,0 +1,69 @@ +"""Optional first-light test for Qwen3.6-35B-A3B on SM120. + +Run with: + + FLASHRT_QWEN36_MOE_CKPT_DIR=/models/Qwen3.6-35B-A3B \ + PYTHONPATH=. pytest -q tests/test_qwen36_moe_gpu.py +""" + +from __future__ import annotations + +import os + +import pytest + + +CKPT = os.environ.get("FLASHRT_QWEN36_MOE_CKPT_DIR") + +pytestmark = pytest.mark.skipif( + not CKPT, + reason="set FLASHRT_QWEN36_MOE_CKPT_DIR for the real-model test", +) + + +def test_qwen36_moe_first_light_matches_hf_golden(): + import torch + + if not torch.cuda.is_available(): + pytest.skip("CUDA is unavailable") + if torch.cuda.get_device_capability() != (12, 0): + pytest.skip("the kernelized qwen3_5_moe path requires SM120") + + from flash_rt.frontends.torch.qwen36_moe_rtx import ( + Qwen36MoeTextFrontendRtx, + ) + + frontend = Qwen36MoeTextFrontendRtx( + CKPT, + device="cuda:0", + max_seq=128, + kernelized=True, + quant_scope="experts", + ) + messages = [{ + "role": "user", + "content": "Write a Python function that merges two sorted lists.", + }] + prompt = frontend.tokenizer.apply_chat_template( + messages, + tokenize=False, + add_generation_prompt=True, + ) + frontend.set_prompt(prompt) + + with torch.no_grad(): + logits = frontend.infer() + assert logits.shape == (1, 20, 248320) + assert torch.isfinite(logits).all() + + # Official Transformers BF16 greedy output for the prompt above. + golden = [ + 8160, 579, 264, 7047, 1817, 25, 271, 16, + 13, 220, 2972, 15771, 2598, 279, 2570, 5952, + ] + with torch.no_grad(): + generations = [ + frontend.generate(max_new_tokens=len(golden)) + for _ in range(8) + ] + assert generations == [golden] * len(generations) diff --git a/tests/test_qwen36_moe_smoke.py b/tests/test_qwen36_moe_smoke.py new file mode 100644 index 00000000..af24bbe1 --- /dev/null +++ b/tests/test_qwen36_moe_smoke.py @@ -0,0 +1,235 @@ +"""Smoke tests for the Qwen3.6-35B-A3B text frontend.""" + +from __future__ import annotations + +import json +import os + +import pytest + + +def _config(): + layer_types = [ + "full_attention" if (i + 1) % 4 == 0 else "linear_attention" + for i in range(40) + ] + return { + "model_type": "qwen3_5_moe", + "text_config": { + "model_type": "qwen3_5_moe_text", + "num_hidden_layers": 40, + "hidden_size": 2048, + "vocab_size": 248320, + "num_attention_heads": 16, + "num_key_value_heads": 2, + "head_dim": 256, + "num_experts": 256, + "num_experts_per_tok": 8, + "linear_num_key_heads": 16, + "linear_num_value_heads": 32, + "linear_key_head_dim": 128, + "linear_value_head_dim": 128, + "full_attention_interval": 4, + "mtp_num_hidden_layers": 1, + "layer_types": layer_types, + }, + } + + +def _checkpoint(tmp_path): + from flash_rt.frontends.torch.qwen36_moe_rtx import ( + _MTP_KEYS, + _required_text_keys, + ) + + config = _config() + layer_types = tuple(config["text_config"]["layer_types"]) + keys = _required_text_keys(layer_types) | _MTP_KEYS + shard = "model-00001-of-00001.safetensors" + (tmp_path / shard).write_bytes(b"checkpoint") + (tmp_path / "config.json").write_text( + json.dumps(config), encoding="utf-8") + (tmp_path / "model.safetensors.index.json").write_text( + json.dumps({"weight_map": {key: shard for key in keys}}), + encoding="utf-8", + ) + return tmp_path + + +def test_frontend_is_a_thin_qwen_entry(): + from flash_rt.frontends.torch.nexn2_rtx import Nexn2TorchFrontendRtx + from flash_rt.frontends.torch.qwen36_moe_rtx import ( + Qwen36MoeTextFrontendRtx, + ) + + assert issubclass(Qwen36MoeTextFrontendRtx, Nexn2TorchFrontendRtx) + assert Qwen36MoeTextFrontendRtx._MODEL_LABEL == ( + "Qwen3.6-35B-A3B text") + + +def test_registry_resolves_qwen36_moe(): + from flash_rt.hardware import _PIPELINE_MAP, resolve_pipeline_class + + assert _PIPELINE_MAP[("qwen36_moe", "torch", "rtx_sm120")] == ( + "flash_rt.frontends.torch.qwen36_moe_rtx", + "Qwen36MoeTextFrontendRtx", + ) + cls = resolve_pipeline_class("qwen36_moe", "torch", "rtx_sm120") + assert cls.__name__ == "Qwen36MoeTextFrontendRtx" + + +def test_load_model_redirects_to_text_frontend(): + import flash_rt + + with pytest.raises(NotImplementedError) as exc: + flash_rt.load_model("/nonexistent", config="qwen36_moe") + message = str(exc.value) + assert "Qwen36MoeTextFrontendRtx" in message + assert "text LLM" in message + + +def test_constructor_rejects_quant_before_checkpoint_access(): + from flash_rt.frontends.torch.qwen36_moe_rtx import ( + Qwen36MoeTextFrontendRtx, + ) + + with pytest.raises(NotImplementedError, match="only 'nvfp4'"): + Qwen36MoeTextFrontendRtx("/nonexistent", quant="fp8") + + +def test_checkpoint_contract_accepts_complete_layout(tmp_path): + from flash_rt.frontends.torch.qwen36_moe_rtx import ( + validate_qwen36_moe_checkpoint, + ) + + result = validate_qwen36_moe_checkpoint(str(_checkpoint(tmp_path))) + assert result["text_tensor_count"] == 693 + assert result["mtp_tensor_count"] == 19 + assert result["tensor_count"] == 712 + assert result["shard_count"] == 1 + + +def test_checkpoint_contract_rejects_wrong_geometry(tmp_path): + from flash_rt.frontends.torch.qwen36_moe_rtx import ( + validate_qwen36_moe_checkpoint, + ) + + checkpoint = _checkpoint(tmp_path) + config_path = checkpoint / "config.json" + config = json.loads(config_path.read_text(encoding="utf-8")) + config["text_config"]["num_experts"] = 128 + config_path.write_text(json.dumps(config), encoding="utf-8") + + with pytest.raises(ValueError, match="num_experts=128"): + validate_qwen36_moe_checkpoint(str(checkpoint)) + + +def test_checkpoint_contract_rejects_missing_text_tensor(tmp_path): + from flash_rt.frontends.torch.qwen36_moe_rtx import ( + validate_qwen36_moe_checkpoint, + ) + + checkpoint = _checkpoint(tmp_path) + index_path = checkpoint / "model.safetensors.index.json" + index = json.loads(index_path.read_text(encoding="utf-8")) + del index["weight_map"]["lm_head.weight"] + index_path.write_text(json.dumps(index), encoding="utf-8") + + with pytest.raises(ValueError, match="lm_head.weight"): + validate_qwen36_moe_checkpoint(str(checkpoint)) + + +def test_checkpoint_contract_rejects_partial_mtp(tmp_path): + from flash_rt.frontends.torch.qwen36_moe_rtx import ( + validate_qwen36_moe_checkpoint, + ) + + checkpoint = _checkpoint(tmp_path) + index_path = checkpoint / "model.safetensors.index.json" + index = json.loads(index_path.read_text(encoding="utf-8")) + del index["weight_map"]["mtp.fc.weight"] + index_path.write_text(json.dumps(index), encoding="utf-8") + + with pytest.raises(ValueError, match="MTP tensor group"): + validate_qwen36_moe_checkpoint(str(checkpoint)) + + +def test_checkpoint_contract_rejects_missing_shard(tmp_path): + from flash_rt.frontends.torch.qwen36_moe_rtx import ( + validate_qwen36_moe_checkpoint, + ) + + checkpoint = _checkpoint(tmp_path) + (checkpoint / "model-00001-of-00001.safetensors").unlink() + + with pytest.raises(FileNotFoundError, match="missing or empty shards"): + validate_qwen36_moe_checkpoint(str(checkpoint)) + + +def test_generic_env_names_precede_legacy_aliases(monkeypatch): + from flash_rt.frontends.torch._nexn2_rtx_decode import _qwen35moe_env + + monkeypatch.setenv("FLASHRT_NEXN2_PREFILL_CHUNK", "4096") + assert _qwen35moe_env("PREFILL_CHUNK", "8192") == "4096" + monkeypatch.setenv("FLASHRT_QWEN35MOE_PREFILL_CHUNK", "2048") + assert _qwen35moe_env("PREFILL_CHUNK", "8192") == "2048" + + +def test_kernelized_generate_uses_shared_graph_path(monkeypatch): + from flash_rt.frontends.torch import _nexn2_rtx_decode as decode + from flash_rt.frontends.torch.qwen36_moe_rtx import ( + Qwen36MoeTextFrontendRtx, + ) + + calls = {} + + class FakeState: + def __init__(self, weights, max_seq, device): + calls["state"] = (weights, max_seq, device) + + def fake_generate(state, prompt_ids, count, fvk, device): + calls["generate"] = (state, prompt_ids, count, fvk, device) + return [7] * count + + monkeypatch.setattr(decode, "Nexn2DecodeState", FakeState) + monkeypatch.setattr(decode, "generate_greedy_graph", fake_generate) + + frontend = Qwen36MoeTextFrontendRtx.__new__( + Qwen36MoeTextFrontendRtx) + frontend._kernelized = True + frontend._prompt_ids = object() + frontend._decode_state = None + frontend._weights = object() + frontend._user_max_seq = 128 + frontend.device = "cuda:0" + frontend._fvk = object() + + assert frontend.generate(3) == [7, 7, 7] + assert calls["state"] == ( + frontend._weights, frontend._user_max_seq, frontend.device) + assert calls["generate"][1:] == ( + frontend._prompt_ids, 3, frontend._fvk, frontend.device) + with pytest.raises(NotImplementedError, match="greedy"): + frontend.generate(1, do_sample=True) + + +@pytest.mark.skipif( + not os.environ.get("FLASHRT_QWEN36_MOE_CKPT_DIR"), + reason="set FLASHRT_QWEN36_MOE_CKPT_DIR for checkpoint validation", +) +def test_real_checkpoint_contract(): + from flash_rt.frontends.torch.qwen36_moe_rtx import ( + validate_qwen36_moe_checkpoint, + ) + + result = validate_qwen36_moe_checkpoint( + os.environ["FLASHRT_QWEN36_MOE_CKPT_DIR"]) + assert result == { + "checkpoint_path": os.path.abspath( + os.environ["FLASHRT_QWEN36_MOE_CKPT_DIR"]), + "text_tensor_count": 693, + "mtp_tensor_count": 19, + "vision_tensor_count": 333, + "tensor_count": 1045, + "shard_count": 26, + } From 1db9b463164342b160cf16b55bacf912115d6e7c Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Wed, 29 Jul 2026 11:37:32 -0400 Subject: [PATCH 02/85] Harden Qwen3.6 runtime contracts --- docs/qwen36_moe_usage.md | 31 ++- flash_rt/frontends/torch/qwen36_moe_rtx.py | 212 ++++++++++++++------- tests/test_qwen36_moe_gpu.py | 87 ++++++++- tests/test_qwen36_moe_smoke.py | 92 ++++++++- 4 files changed, 335 insertions(+), 87 deletions(-) diff --git a/docs/qwen36_moe_usage.md b/docs/qwen36_moe_usage.md index 937c3905..d2f31235 100644 --- a/docs/qwen36_moe_usage.md +++ b/docs/qwen36_moe_usage.md @@ -19,12 +19,14 @@ decode are not part of this interface. | GPU memory | 32 GB | | Framework | PyTorch | | Runtime quantization | NVFP4 | -| Build flag | `-DFLASHRT_ENABLE_QWEN35MOE=ON` | +| Build flags | `-DGPU_ARCH=120 -DFLASHRT_ENABLE_QWEN35MOE=ON` | Configure and build the gated `qwen3_5_moe` kernels: ```bash -cmake -S . -B build -DFLASHRT_ENABLE_QWEN35MOE=ON +cmake -S . -B build \ + -DGPU_ARCH=120 \ + -DFLASHRT_ENABLE_QWEN35MOE=ON cmake --build build -j pip install -e ".[torch]" ``` @@ -40,7 +42,6 @@ frontend = Qwen36MoeTextFrontendRtx( "/models/Qwen3.6-35B-A3B", device="cuda:0", max_seq=4096, - kernelized=True, quant_scope="experts", ) frontend.set_prompt("Explain why deterministic reductions matter.") @@ -65,14 +66,22 @@ The direct frontend is intentional. `flash_rt.load_model()` wraps VLA models with a `predict(images, ...)` API and therefore redirects `config="qwen36_moe"` to the class above. +The frontend defaults to `kernelized=True`. Setting `kernelized=False` is +rejected because it would select the parent Transformers reference path, which +loads the complete multimodal BF16 model rather than this text-only NVFP4 +runtime. + ## Checkpoint contract Before allocating GPU memory, the frontend checks: -- the exact 40-layer `qwen3_5_moe` text geometry; +- the exact 40-layer `qwen3_5_moe` text geometry, including MoE widths, + convolution width, gated attention, normalization epsilon, and RoPE + parameters; - the 30 linear-attention / 10 full-attention schedule; -- all 693 text-backbone tensors consumed by the shared pipeline; -- all 19 official MTP tensors; +- the exact shapes of all 693 text-backbone tensors consumed by the shared + pipeline; +- the exact shapes of all 19 official MTP tensors; - every safetensors shard referenced by the index. Extra vision tensors are allowed and ignored by this text-only entry. The @@ -113,9 +122,10 @@ test. Performance and precision numbers must be measured on Qwen3.6 weights; Nex-N2-mini measurements are not interchangeable even though the compute pipeline is shared. -The optional real-model test loads the checkpoint, checks finite logits, and -compares cold and warm CUDA Graph generation with an official Transformers -BF16 greedy-token fixture: +The optional GPU suite checks the shared weighted-sum reducer against the +former Torch reduction, repeats it eagerly and through CUDA Graph replay, then +loads the checkpoint, checks finite logits, and compares cold and warm CUDA +Graph generation with an official Transformers BF16 greedy-token fixture: ```bash FLASHRT_QWEN36_MOE_CKPT_DIR=/models/Qwen3.6-35B-A3B \ @@ -155,7 +165,8 @@ greedy sequences were token-exact for 16 generated tokens on all four prompts. ## Limitations - Text only; the vision tower is not loaded. -- Greedy decode only on the kernelized path. +- The kernelized runtime NVFP4 path is required. +- Greedy decode only. - The MTP tensors are validated but not loaded, so speculative decode is not enabled. - Only the BF16 source checkpoint with runtime NVFP4 conversion is supported. diff --git a/flash_rt/frontends/torch/qwen36_moe_rtx.py b/flash_rt/frontends/torch/qwen36_moe_rtx.py index df98f705..cf4b3e85 100644 --- a/flash_rt/frontends/torch/qwen36_moe_rtx.py +++ b/flash_rt/frontends/torch/qwen36_moe_rtx.py @@ -11,6 +11,7 @@ import json import os +from contextlib import ExitStack from typing import Any from flash_rt.frontends.torch.nexn2_rtx import Nexn2TorchFrontendRtx @@ -29,85 +30,133 @@ "num_attention_heads": 16, "num_key_value_heads": 2, "head_dim": 256, + "attention_bias": False, + "attn_output_gate": True, + "hidden_act": "silu", "num_experts": 256, "num_experts_per_tok": 8, + "moe_intermediate_size": 512, + "shared_expert_intermediate_size": 512, "linear_num_key_heads": 16, "linear_num_value_heads": 32, "linear_key_head_dim": 128, "linear_value_head_dim": 128, + "linear_conv_kernel_dim": 4, + "mamba_ssm_dtype": "float32", "full_attention_interval": 4, + "partial_rotary_factor": 0.25, + "rms_norm_eps": 1e-6, "mtp_num_hidden_layers": 1, + "mtp_use_dedicated_embeddings": False, + "tie_word_embeddings": False, } -_MTP_KEYS = { - "mtp.fc.weight", - "mtp.norm.weight", - "mtp.pre_fc_norm_embedding.weight", - "mtp.pre_fc_norm_hidden.weight", - "mtp.layers.0.input_layernorm.weight", - "mtp.layers.0.post_attention_layernorm.weight", - "mtp.layers.0.self_attn.q_proj.weight", - "mtp.layers.0.self_attn.k_proj.weight", - "mtp.layers.0.self_attn.v_proj.weight", - "mtp.layers.0.self_attn.o_proj.weight", - "mtp.layers.0.self_attn.q_norm.weight", - "mtp.layers.0.self_attn.k_norm.weight", - "mtp.layers.0.mlp.gate.weight", - "mtp.layers.0.mlp.experts.gate_up_proj", - "mtp.layers.0.mlp.experts.down_proj", - "mtp.layers.0.mlp.shared_expert.gate_proj.weight", - "mtp.layers.0.mlp.shared_expert.up_proj.weight", - "mtp.layers.0.mlp.shared_expert.down_proj.weight", - "mtp.layers.0.mlp.shared_expert_gate.weight", +_EXPECTED_ROPE_PARAMETERS = { + "rope_type": "default", + "rope_theta": 10000000, + "partial_rotary_factor": 0.25, + "mrope_interleaved": True, + "mrope_section": [11, 11, 10], } -def _required_text_keys(layer_types: tuple[str, ...]) -> set[str]: - keys = { - "lm_head.weight", - "model.language_model.embed_tokens.weight", - "model.language_model.norm.weight", +def _expected_mlp_shapes(prefix: str) -> dict[str, tuple[int, ...]]: + return { + prefix + "mlp.gate.weight": (256, 2048), + prefix + "mlp.experts.gate_up_proj": (256, 1024, 2048), + prefix + "mlp.experts.down_proj": (256, 2048, 512), + prefix + "mlp.shared_expert.gate_proj.weight": (512, 2048), + prefix + "mlp.shared_expert.up_proj.weight": (512, 2048), + prefix + "mlp.shared_expert.down_proj.weight": (2048, 512), + prefix + "mlp.shared_expert_gate.weight": (1, 2048), + } + + +def _expected_attention_shapes( + prefix: str, layer_type: str) -> dict[str, tuple[int, ...]]: + if layer_type == "full_attention": + return { + prefix + "self_attn.q_proj.weight": (8192, 2048), + prefix + "self_attn.k_proj.weight": (512, 2048), + prefix + "self_attn.v_proj.weight": (512, 2048), + prefix + "self_attn.o_proj.weight": (2048, 4096), + prefix + "self_attn.q_norm.weight": (256,), + prefix + "self_attn.k_norm.weight": (256,), + } + return { + prefix + "linear_attn.in_proj_qkv.weight": (8192, 2048), + prefix + "linear_attn.in_proj_z.weight": (4096, 2048), + prefix + "linear_attn.in_proj_a.weight": (32, 2048), + prefix + "linear_attn.in_proj_b.weight": (32, 2048), + prefix + "linear_attn.conv1d.weight": (8192, 1, 4), + prefix + "linear_attn.A_log": (32,), + prefix + "linear_attn.dt_bias": (32,), + prefix + "linear_attn.norm.weight": (128,), + prefix + "linear_attn.out_proj.weight": (2048, 4096), + } + + +def _expected_text_shapes( + layer_types: tuple[str, ...]) -> dict[str, tuple[int, ...]]: + shapes = { + "lm_head.weight": (248320, 2048), + "model.language_model.embed_tokens.weight": (248320, 2048), + "model.language_model.norm.weight": (2048,), } - mlp_suffixes = ( - "mlp.gate.weight", - "mlp.experts.gate_up_proj", - "mlp.experts.down_proj", - "mlp.shared_expert.gate_proj.weight", - "mlp.shared_expert.up_proj.weight", - "mlp.shared_expert.down_proj.weight", - "mlp.shared_expert_gate.weight", - ) - linear_suffixes = ( - "linear_attn.in_proj_qkv.weight", - "linear_attn.in_proj_z.weight", - "linear_attn.in_proj_a.weight", - "linear_attn.in_proj_b.weight", - "linear_attn.conv1d.weight", - "linear_attn.A_log", - "linear_attn.dt_bias", - "linear_attn.norm.weight", - "linear_attn.out_proj.weight", - ) - full_suffixes = ( - "self_attn.q_proj.weight", - "self_attn.k_proj.weight", - "self_attn.v_proj.weight", - "self_attn.o_proj.weight", - "self_attn.q_norm.weight", - "self_attn.k_norm.weight", - ) for i, layer_type in enumerate(layer_types): prefix = f"model.language_model.layers.{i}." - keys.add(prefix + "input_layernorm.weight") - keys.add(prefix + "post_attention_layernorm.weight") - keys.update(prefix + suffix for suffix in mlp_suffixes) - suffixes = ( - full_suffixes - if layer_type == "full_attention" - else linear_suffixes - ) - keys.update(prefix + suffix for suffix in suffixes) - return keys + shapes[prefix + "input_layernorm.weight"] = (2048,) + shapes[prefix + "post_attention_layernorm.weight"] = (2048,) + shapes.update(_expected_mlp_shapes(prefix)) + shapes.update(_expected_attention_shapes(prefix, layer_type)) + return shapes + + +def _expected_mtp_shapes() -> dict[str, tuple[int, ...]]: + prefix = "mtp.layers.0." + shapes = { + "mtp.fc.weight": (2048, 4096), + "mtp.norm.weight": (2048,), + "mtp.pre_fc_norm_embedding.weight": (2048,), + "mtp.pre_fc_norm_hidden.weight": (2048,), + prefix + "input_layernorm.weight": (2048,), + prefix + "post_attention_layernorm.weight": (2048,), + } + shapes.update(_expected_attention_shapes(prefix, "full_attention")) + shapes.update(_expected_mlp_shapes(prefix)) + return shapes + + +_MTP_KEYS = set(_expected_mtp_shapes()) + + +def _required_text_keys(layer_types: tuple[str, ...]) -> set[str]: + return set(_expected_text_shapes(layer_types)) + + +def _read_tensor_shapes( + checkpoint_path: str, + weight_map: dict[str, str], + tensor_names: set[str], +) -> dict[str, tuple[int, ...]]: + from safetensors import safe_open + + shapes = {} + with ExitStack() as stack: + readers = { + shard: stack.enter_context( + safe_open( + os.path.join(checkpoint_path, shard), + framework="pt", + device="cpu", + ) + ) + for shard in set(weight_map[name] for name in tensor_names) + } + for name in tensor_names: + shapes[name] = tuple( + readers[weight_map[name]].get_slice(name).get_shape()) + return shapes def validate_qwen36_moe_checkpoint(checkpoint_path: str) -> dict[str, Any]: @@ -139,6 +188,17 @@ def validate_qwen36_moe_checkpoint(checkpoint_path: str) -> dict[str, Any]: actual = text_config.get(name) if actual != expected: mismatches.append(f"{name}={actual!r} (expected {expected!r})") + rope_parameters = text_config.get("rope_parameters") + if not isinstance(rope_parameters, dict): + mismatches.append( + "rope_parameters is missing or is not an object") + else: + for name, expected in _EXPECTED_ROPE_PARAMETERS.items(): + actual = rope_parameters.get(name) + if actual != expected: + mismatches.append( + f"rope_parameters.{name}={actual!r} " + f"(expected {expected!r})") layer_types = tuple(text_config.get("layer_types") or ()) if layer_types != _EXPECTED_LAYER_TYPES: mismatches.append( @@ -188,6 +248,23 @@ def validate_qwen36_moe_checkpoint(checkpoint_path: str) -> dict[str, Any]: "Qwen3.6-35B-A3B checkpoint has missing or empty shards: " + preview) + expected_shapes = _expected_text_shapes(layer_types) + expected_shapes.update(_expected_mtp_shapes()) + actual_shapes = _read_tensor_shapes( + checkpoint_path, weight_map, set(expected_shapes)) + shape_mismatches = [ + f"{name}={actual_shapes[name]!r} (expected {expected!r})" + for name, expected in sorted(expected_shapes.items()) + if actual_shapes[name] != expected + ] + if shape_mismatches: + preview = "; ".join(shape_mismatches[:8]) + if len(shape_mismatches) > 8: + preview += f"; ... ({len(shape_mismatches)} mismatched)" + raise ValueError( + "Qwen3.6-35B-A3B checkpoint tensor shape mismatches: " + + preview) + return { "checkpoint_path": checkpoint_path, "text_tensor_count": len(required_text), @@ -209,12 +286,16 @@ def __init__(self, checkpoint_path: str, *, device: str = "cuda:0", max_seq: int = 2048, quant: str = "nvfp4", - kernelized: bool = False, + kernelized: bool = True, quant_scope: str = "experts") -> None: if quant != "nvfp4": raise NotImplementedError( f"quant={quant!r} is not implemented; only 'nvfp4' is " "supported") + if not kernelized: + raise NotImplementedError( + "Qwen3.6-35B-A3B text only supports kernelized=True with " + "runtime NVFP4 conversion") contract = validate_qwen36_moe_checkpoint(checkpoint_path) super().__init__( checkpoint_path, @@ -228,9 +309,6 @@ def __init__(self, checkpoint_path: str, *, def generate(self, max_new_tokens: int, *, do_sample: bool = False): """Generate tokens with the shared qwen3_5_moe CUDA Graph path.""" - if not self._kernelized: - return super().generate( - max_new_tokens=max_new_tokens, do_sample=do_sample) if self._prompt_ids is None: raise ValueError("call set_prompt(...) before generate()") if do_sample: diff --git a/tests/test_qwen36_moe_gpu.py b/tests/test_qwen36_moe_gpu.py index c16be2d3..76c3e4ef 100644 --- a/tests/test_qwen36_moe_gpu.py +++ b/tests/test_qwen36_moe_gpu.py @@ -8,6 +8,7 @@ from __future__ import annotations +import importlib import os import pytest @@ -15,19 +16,93 @@ CKPT = os.environ.get("FLASHRT_QWEN36_MOE_CKPT_DIR") -pytestmark = pytest.mark.skipif( - not CKPT, - reason="set FLASHRT_QWEN36_MOE_CKPT_DIR for the real-model test", -) - -def test_qwen36_moe_first_light_matches_hf_golden(): +def _require_sm120(): import torch if not torch.cuda.is_available(): pytest.skip("CUDA is unavailable") if torch.cuda.get_device_capability() != (12, 0): pytest.skip("the kernelized qwen3_5_moe path requires SM120") + try: + kernels = importlib.import_module("flash_rt.flash_rt_kernels") + except ImportError as exc: + pytest.skip(f"FlashRT CUDA extension is unavailable: {exc}") + if not hasattr(kernels, "moe_weighted_sum_sm120_bf16"): + pytest.skip("FlashRT was built without qwen3_5_moe kernels") + return torch, kernels + + +def _weighted_sum(kernels, d_dn, rows, weights, out): + import torch + + status = kernels.moe_weighted_sum_sm120_bf16( + d_dn.data_ptr(), + rows.data_ptr(), + weights.data_ptr(), + out.data_ptr(), + 1, + weights.numel(), + d_dn.shape[1], + d_dn.shape[1], + torch.cuda.current_stream().cuda_stream, + ) + assert status == 0 + + +def test_qwen35moe_weighted_sum_is_graph_deterministic(): + torch, kernels = _require_sm120() + + generator = torch.Generator(device="cuda").manual_seed(7) + d_dn = torch.randn( + 8, 2048, dtype=torch.bfloat16, device="cuda", + generator=generator) + rows = torch.arange(8, dtype=torch.int32, device="cuda") + weights = torch.softmax( + torch.randn(8, dtype=torch.float32, device="cuda", + generator=generator), + dim=0, + ).contiguous() + reference = weights @ d_dn.float() + out = torch.empty(2048, dtype=torch.float32, device="cuda") + + _weighted_sum(kernels, d_dn, rows, weights, out) + torch.cuda.synchronize() + assert torch.isfinite(out).all() + assert (out - reference).abs().max().item() <= 5e-7 + + eager_results = [] + for _ in range(20): + _weighted_sum(kernels, d_dn, rows, weights, out) + eager_results.append(out.clone()) + torch.cuda.synchronize() + assert all(torch.equal(eager_results[0], result) + for result in eager_results[1:]) + + graph_out = torch.empty_like(out) + _weighted_sum(kernels, d_dn, rows, weights, graph_out) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + _weighted_sum(kernels, d_dn, rows, weights, graph_out) + + graph_results = [] + for _ in range(20): + graph.replay() + graph_results.append(graph_out.clone()) + torch.cuda.synchronize() + assert torch.isfinite(graph_results[0]).all() + assert all(torch.equal(graph_results[0], result) + for result in graph_results[1:]) + assert torch.equal(graph_results[0], eager_results[0]) + + +@pytest.mark.skipif( + not CKPT, + reason="set FLASHRT_QWEN36_MOE_CKPT_DIR for the real-model test", +) +def test_qwen36_moe_first_light_matches_hf_golden(): + torch, _ = _require_sm120() from flash_rt.frontends.torch.qwen36_moe_rtx import ( Qwen36MoeTextFrontendRtx, diff --git a/tests/test_qwen36_moe_smoke.py b/tests/test_qwen36_moe_smoke.py index af24bbe1..191dc6c7 100644 --- a/tests/test_qwen36_moe_smoke.py +++ b/tests/test_qwen36_moe_smoke.py @@ -2,6 +2,7 @@ from __future__ import annotations +import inspect import json import os @@ -23,14 +24,32 @@ def _config(): "num_attention_heads": 16, "num_key_value_heads": 2, "head_dim": 256, + "attention_bias": False, + "attn_output_gate": True, + "hidden_act": "silu", "num_experts": 256, "num_experts_per_tok": 8, + "moe_intermediate_size": 512, + "shared_expert_intermediate_size": 512, "linear_num_key_heads": 16, "linear_num_value_heads": 32, "linear_key_head_dim": 128, "linear_value_head_dim": 128, + "linear_conv_kernel_dim": 4, + "mamba_ssm_dtype": "float32", "full_attention_interval": 4, + "partial_rotary_factor": 0.25, + "rms_norm_eps": 1e-6, "mtp_num_hidden_layers": 1, + "mtp_use_dedicated_embeddings": False, + "tie_word_embeddings": False, + "rope_parameters": { + "rope_type": "default", + "rope_theta": 10000000, + "partial_rotary_factor": 0.25, + "mrope_interleaved": True, + "mrope_section": [11, 11, 10], + }, "layer_types": layer_types, }, } @@ -56,6 +75,22 @@ def _checkpoint(tmp_path): return tmp_path +def _mock_checkpoint_shapes(monkeypatch, overrides=None): + from flash_rt.frontends.torch import qwen36_moe_rtx + + shapes = qwen36_moe_rtx._expected_text_shapes( + qwen36_moe_rtx._EXPECTED_LAYER_TYPES) + shapes.update(qwen36_moe_rtx._expected_mtp_shapes()) + shapes.update(overrides or {}) + monkeypatch.setattr( + qwen36_moe_rtx, + "_read_tensor_shapes", + lambda checkpoint_path, weight_map, tensor_names: { + name: shapes[name] for name in tensor_names + }, + ) + + def test_frontend_is_a_thin_qwen_entry(): from flash_rt.frontends.torch.nexn2_rtx import Nexn2TorchFrontendRtx from flash_rt.frontends.torch.qwen36_moe_rtx import ( @@ -65,6 +100,8 @@ def test_frontend_is_a_thin_qwen_entry(): assert issubclass(Qwen36MoeTextFrontendRtx, Nexn2TorchFrontendRtx) assert Qwen36MoeTextFrontendRtx._MODEL_LABEL == ( "Qwen3.6-35B-A3B text") + assert inspect.signature( + Qwen36MoeTextFrontendRtx).parameters["kernelized"].default is True def test_registry_resolves_qwen36_moe(): @@ -97,11 +134,21 @@ def test_constructor_rejects_quant_before_checkpoint_access(): Qwen36MoeTextFrontendRtx("/nonexistent", quant="fp8") -def test_checkpoint_contract_accepts_complete_layout(tmp_path): +def test_constructor_rejects_reference_path_before_checkpoint_access(): + from flash_rt.frontends.torch.qwen36_moe_rtx import ( + Qwen36MoeTextFrontendRtx, + ) + + with pytest.raises(NotImplementedError, match="kernelized=True"): + Qwen36MoeTextFrontendRtx("/nonexistent", kernelized=False) + + +def test_checkpoint_contract_accepts_complete_layout(tmp_path, monkeypatch): from flash_rt.frontends.torch.qwen36_moe_rtx import ( validate_qwen36_moe_checkpoint, ) + _mock_checkpoint_shapes(monkeypatch) result = validate_qwen36_moe_checkpoint(str(_checkpoint(tmp_path))) assert result["text_tensor_count"] == 693 assert result["mtp_tensor_count"] == 19 @@ -109,7 +156,28 @@ def test_checkpoint_contract_accepts_complete_layout(tmp_path): assert result["shard_count"] == 1 -def test_checkpoint_contract_rejects_wrong_geometry(tmp_path): +@pytest.mark.parametrize( + ("path", "invalid", "message"), + [ + (("moe_intermediate_size",), 1024, "moe_intermediate_size=1024"), + ( + ("shared_expert_intermediate_size",), + 1024, + "shared_expert_intermediate_size=1024", + ), + (("linear_conv_kernel_dim",), 3, "linear_conv_kernel_dim=3"), + (("partial_rotary_factor",), 0.5, "partial_rotary_factor=0.5"), + (("rms_norm_eps",), 1e-5, "rms_norm_eps=1e-05"), + (("attn_output_gate",), False, "attn_output_gate=False"), + ( + ("rope_parameters", "rope_theta"), + 10000, + "rope_parameters.rope_theta=10000", + ), + ], +) +def test_checkpoint_contract_rejects_wrong_geometry( + tmp_path, path, invalid, message): from flash_rt.frontends.torch.qwen36_moe_rtx import ( validate_qwen36_moe_checkpoint, ) @@ -117,10 +185,13 @@ def test_checkpoint_contract_rejects_wrong_geometry(tmp_path): checkpoint = _checkpoint(tmp_path) config_path = checkpoint / "config.json" config = json.loads(config_path.read_text(encoding="utf-8")) - config["text_config"]["num_experts"] = 128 + target = config["text_config"] + for name in path[:-1]: + target = target[name] + target[path[-1]] = invalid config_path.write_text(json.dumps(config), encoding="utf-8") - with pytest.raises(ValueError, match="num_experts=128"): + with pytest.raises(ValueError, match=message): validate_qwen36_moe_checkpoint(str(checkpoint)) @@ -166,6 +237,19 @@ def test_checkpoint_contract_rejects_missing_shard(tmp_path): validate_qwen36_moe_checkpoint(str(checkpoint)) +def test_checkpoint_contract_rejects_wrong_tensor_shape( + tmp_path, monkeypatch): + from flash_rt.frontends.torch.qwen36_moe_rtx import ( + validate_qwen36_moe_checkpoint, + ) + + name = "model.language_model.layers.0.linear_attn.conv1d.weight" + _mock_checkpoint_shapes(monkeypatch, {name: (8192, 1, 3)}) + + with pytest.raises(ValueError, match=r"conv1d\.weight=.*8192, 1, 3"): + validate_qwen36_moe_checkpoint(str(_checkpoint(tmp_path))) + + def test_generic_env_names_precede_legacy_aliases(monkeypatch): from flash_rt.frontends.torch._nexn2_rtx_decode import _qwen35moe_env From c92dc19e01aa3545ffa4ba4123174c4ecd4fcdc5 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 30 Jul 2026 04:50:06 -0400 Subject: [PATCH 03/85] Add Qwen3.6-MoE edge checkpoint utilities Development tooling for a memory-constrained Qwen3.6-35B-A3B runtime: - probe.py projects per-category checkpoint sizes and expert-cache quotas for a given memory budget, and scores sampled experts under INT8/INT4. - quantize_experts.py writes routed experts as fixed-size INT8 or INT4 blocks. INT4 follows the sign-magnitude / UE4M3-per-16 contract and optionally applies the orthonormal H16 transform. - route_trace.py records router selections and simulates a bounded per-layer LRU to size the expert cache. The router trace hook in the shared decode path is eager-only and stays disabled unless a caller opts in, so CUDA Graph capture is unaffected. --- flash_rt/frontends/torch/_nexn2_rtx_decode.py | 8 + qwen36_moe_edge/README.md | 64 ++++ qwen36_moe_edge/__init__.py | 1 + qwen36_moe_edge/probe.py | 324 ++++++++++++++++++ qwen36_moe_edge/quantize_experts.py | 225 ++++++++++++ qwen36_moe_edge/route_trace.py | 122 +++++++ tests/test_qwen36_moe_edge_quant.py | 98 ++++++ 7 files changed, 842 insertions(+) create mode 100644 qwen36_moe_edge/README.md create mode 100644 qwen36_moe_edge/__init__.py create mode 100644 qwen36_moe_edge/probe.py create mode 100644 qwen36_moe_edge/quantize_experts.py create mode 100644 qwen36_moe_edge/route_trace.py create mode 100644 tests/test_qwen36_moe_edge_quant.py diff --git a/flash_rt/frontends/torch/_nexn2_rtx_decode.py b/flash_rt/frontends/torch/_nexn2_rtx_decode.py index 7ff6f8f0..2704a8e1 100644 --- a/flash_rt/frontends/torch/_nexn2_rtx_decode.py +++ b/flash_rt/frontends/torch/_nexn2_rtx_decode.py @@ -273,6 +273,10 @@ def __init__(self, handles, max_seq, device): # chunk (64). 0 disables (always single-pass). self.prefill_chunk = int( _qwen35moe_env("PREFILL_CHUNK", "8192")) + # Optional eager-only routing trace used to size edge expert caches. + # Keep this disabled during CUDA Graph capture. + self.router_trace = None + self._active_layer = -1 def reset(self): for s in self.lin_state: @@ -443,6 +447,9 @@ def _moe_layer_decode(h, ld, state, fvk, device): fvk.moe_router_topk_sm120_bf16(lr.data_ptr(), idx.data_ptr(), topv.data_ptr(), lr.numel(), TOPK, s) tw_row = F.softmax(topv, -1) # (TOPK,) device + if state.router_trace is not None: + state.router_trace[state._active_layer].append( + tuple(int(v) for v in idx.cpu().tolist())) if 'experts_gate_up_alpha_dev' not in ld: # cache once/layer ld['experts_gate_up_alpha_dev'] = \ @@ -528,6 +535,7 @@ def decode_step(state, token_id, pos, fvk, device): h = res + attn res = h n = _rms_fvk(h, ld['post_norm_w_t'], fvk, device, state.eps) + state._active_layer = L h = res + _moe_layer_decode(n, ld, state, fvk, device) h = _rms_fvk(h, p['final_norm_w_t'], fvk, device, state.eps) diff --git a/qwen36_moe_edge/README.md b/qwen36_moe_edge/README.md new file mode 100644 index 00000000..c3f973fa --- /dev/null +++ b/qwen36_moe_edge/README.md @@ -0,0 +1,64 @@ +# Qwen3.6-MoE edge experiments + +This directory contains checkpoint-independent development utilities for a +memory-constrained Qwen3.6-35B-A3B runtime. It is not a production frontend. + +The intended runtime layout follows the MiniMax-M3 Spark prototype: + +- non-routed weights remain resident in a mixed-precision format; +- each routed expert is stored as one fixed-size block; +- a bounded per-layer LRU holds hot expert blocks; +- misses are read from local storage into reusable staging buffers. + +Inspect projected checkpoint sizes and sampled expert quality: + +```bash +PYTHONPATH=. python qwen36_moe_edge/probe.py \ + --checkpoint /models/Qwen3.6-35B-A3B \ + --mode memory \ + --group-size 16 + +PYTHONPATH=. python qwen36_moe_edge/probe.py \ + --checkpoint /models/Qwen3.6-35B-A3B \ + --mode quality \ + --group-size 16 +``` + +Generate fixed-size routed-expert blocks for a layer range: + +```bash +PYTHONPATH=. python qwen36_moe_edge/quantize_experts.py \ + --checkpoint /models/Qwen3.6-35B-A3B \ + --output /models/Qwen3.6-35B-A3B-INT8E \ + --format int8 \ + --layers 0:40 + +PYTHONPATH=. python qwen36_moe_edge/quantize_experts.py \ + --checkpoint /models/Qwen3.6-35B-A3B \ + --output /models/Qwen3.6-35B-A3B-INT4E-RHT16 \ + --format int4-rht \ + --group-size 16 \ + --layers 0:40 +``` + +INT8 uses symmetric per-output-channel FP16 scales. INT4 follows the Thor +Pi0.5 numerical contract: sign-magnitude values, one UE4M3 scale per 16 K +values, and two values per byte with the low nibble first. `int4-rht` applies +the same orthonormal H16/4 transform to every K block that the runtime applies +to activations. Scale bytes in these edge block files are linear; a Thor +loader must convert them to the SM1xx SFB tile-interleaved layout before +calling the native block-scaled MMA kernels. + +An SM120 machine can collect real router selections for cache sizing: + +```bash +PYTHONPATH=. python qwen36_moe_edge/route_trace.py \ + --checkpoint /models/Qwen3.6-35B-A3B \ + --prompt "Explain edge mixture-of-experts inference. " \ + --prompt-tokens 32 \ + --new-tokens 64 \ + --output qwen36_moe_route_trace.json +``` + +Tracing deliberately uses eager per-token prefill. It must not be enabled +during CUDA Graph capture. diff --git a/qwen36_moe_edge/__init__.py b/qwen36_moe_edge/__init__.py new file mode 100644 index 00000000..b423fee7 --- /dev/null +++ b/qwen36_moe_edge/__init__.py @@ -0,0 +1 @@ +"""Development utilities for memory-constrained Qwen3.6-MoE inference.""" diff --git a/qwen36_moe_edge/probe.py b/qwen36_moe_edge/probe.py new file mode 100644 index 00000000..3a3850f3 --- /dev/null +++ b/qwen36_moe_edge/probe.py @@ -0,0 +1,324 @@ +#!/usr/bin/env python3 +"""Memory and quantization probes for the Qwen3.6-MoE edge path.""" + +from __future__ import annotations + +import argparse +import json +import statistics +from pathlib import Path + +import torch +import torch.nn.functional as F +from safetensors import safe_open + +from qwen36_moe_edge.quantize_experts import ( + HIDDEN, + INTERMEDIATE, + NUM_EXPERTS, + NUM_LAYERS, + CheckpointReader, + _hadamard16, + _int4_weight, + _int8_weight, + _layout, +) + + +def _numel(shape: tuple[int, ...]) -> int: + result = 1 + for value in shape: + result *= value + return result + + +def _category(name: str) -> str: + if ".mlp.experts." in name: + return "routed_experts" + if name in ( + "lm_head.weight", + "model.language_model.embed_tokens.weight", + ): + return "embed_lm_head" + if ".linear_attn." in name and not name.endswith("norm.weight"): + return "gdn_weights" + dense_markers = ( + ".self_attn.q_proj.weight", + ".self_attn.k_proj.weight", + ".self_attn.v_proj.weight", + ".self_attn.o_proj.weight", + ".linear_attn.out_proj.weight", + ".mlp.shared_expert.gate_proj.weight", + ".mlp.shared_expert.up_proj.weight", + ".mlp.shared_expert.down_proj.weight", + ) + if any(marker in name for marker in dense_markers): + return "other_dense" + return "norm_router_misc" + + +def _quantized_bytes( + shape: tuple[int, ...], *, bits: int, group_size: int +) -> int: + elements = _numel(shape) + if len(shape) < 2: + return elements * 2 + rows = _numel(shape[:-1]) + columns = shape[-1] + if bits == 8: + return elements + rows * 2 + return (elements + 1) // 2 + rows * ( + (columns + group_size - 1) // group_size) + + +def memory_probe( + checkpoint: Path, + group_size: int, + budget_gib: float, + runtime_reserve_gib: float) -> None: + with (checkpoint / "model.safetensors.index.json").open( + encoding="utf-8") as f: + weight_map = json.load(f)["weight_map"] + readers = { + shard: safe_open( + checkpoint / shard, framework="pt", device="cpu") + for shard in set(weight_map.values()) + } + categories: dict[str, list[tuple[int, tuple[int, ...]]]] = {} + for name, shard in weight_map.items(): + if ".visual." in name or name.startswith("mtp."): + continue + shape = tuple(readers[shard].get_slice(name).get_shape()) + categories.setdefault(_category(name), []).append( + (_numel(shape), shape)) + + print("category bf16 GiB int8 GiB int4 GiB") + for category, tensors in sorted(categories.items()): + bf16 = sum(elements * 2 for elements, _ in tensors) + int8 = sum( + _quantized_bytes(shape, bits=8, group_size=group_size) + for _, shape in tensors + ) + int4 = sum( + _quantized_bytes(shape, bits=4, group_size=group_size) + for _, shape in tensors + ) + print( + f"{category:22s} {bf16 / 2**30:8.3f} " + f"{int8 / 2**30:10.3f} {int4 / 2**30:10.3f}" + ) + for quant_format in ("int8", "int4"): + layout = _layout(quant_format, group_size) + print( + f"{quant_format} expert block: " + f"{sum(layout.values()) / 2**20:.4f} MiB" + ) + + int8_resident = sum( + _quantized_bytes(shape, bits=8, group_size=group_size) + for category, tensors in categories.items() + if category != "routed_experts" + for _, shape in tensors + ) + int4_resident = sum( + ( + elements * 2 + if category == "gdn_weights" + else _quantized_bytes( + shape, bits=4, group_size=group_size) + ) + for category, tensors in categories.items() + if category != "routed_experts" + for elements, shape in tensors + ) + budget_bytes = int(budget_gib * 2**30) + reserve_bytes = int(runtime_reserve_gib * 2**30) + print( + f"\n{budget_gib:.2f} GiB budget, " + f"{runtime_reserve_gib:.2f} GiB runtime reserve" + ) + for quant_format, resident in ( + ("int8", int8_resident), + ("int4-mixed", int4_resident), + ): + block_format = "int8" if quant_format == "int8" else "int4" + block_bytes = sum(_layout(block_format, group_size).values()) + available = max(0, budget_bytes - reserve_bytes - resident) + quota = available // block_bytes // NUM_LAYERS + cache_bytes = quota * NUM_LAYERS * block_bytes + projected = resident + cache_bytes + reserve_bytes + print( + f"{quant_format:10s}: resident={resident / 2**30:.3f} GiB, " + f"quota={quota} experts/layer, " + f"cache={cache_bytes / 2**30:.3f} GiB, " + f"projected={projected / 2**30:.3f} GiB" + ) + + +def _dequant_int4( + packed: torch.Tensor, + scale: torch.Tensor, + columns: int, + group_size: int, +) -> torch.Tensor: + low = packed & 0x0F + high = (packed >> 4) & 0x0F + low = (low & 0x07).to(torch.int8) * torch.where( + (low & 0x08) != 0, -1, 1).to(torch.int8) + high = (high & 0x07).to(torch.int8) * torch.where( + (high & 0x08) != 0, -1, 1).to(torch.int8) + values = torch.stack((low, high), dim=-1).flatten(1) + rows = values.shape[0] + scale_float = scale.view(torch.float8_e4m3fn).float() + return ( + values.float().reshape(rows, columns // group_size, group_size) + * scale_float.unsqueeze(-1) + ).reshape(rows, columns) + + +def quality_probe( + checkpoint: Path, + *, + layers: int, + group_size: int, + device: str, +) -> None: + reader = CheckpointReader(checkpoint) + generator = torch.Generator(device=device).manual_seed(2026) + scores = { + "w8a16": [], + "w8a8": [], + "int4_w4a4": [], + "int4_rht16_w4a4": [], + } + selected_layers = torch.linspace( + 0, NUM_LAYERS - 1, steps=layers).round().int().tolist() + for layer in selected_layers: + expert = (layer * 37 + 11) % NUM_EXPERTS + gate_up = reader.expert( + layer, "gate_up_proj", expert + ).to(device=device, dtype=torch.float32) + down = reader.expert( + layer, "down_proj", expert + ).to(device=device, dtype=torch.float32) + activation = torch.randn( + 16, HIDDEN, generator=generator, device=device) + projected = activation @ gate_up.T + reference = ( + F.silu(projected[:, :INTERMEDIATE]) + * projected[:, INTERMEDIATE:] + ) @ down.T + + gu8, gu8_scale = _int8_weight(gate_up) + dn8, dn8_scale = _int8_weight(down) + gu8 = gu8.float() * gu8_scale.float()[:, None] + dn8 = dn8.float() * dn8_scale.float()[:, None] + for mode in ("w8a16", "w8a8"): + current = activation + if mode == "w8a8": + scale = ( + current.abs().amax(dim=1, keepdim=True).clamp_min(1e-8) + / 127.0 + ) + current = ( + current / scale + ).round().clamp(-127, 127) * scale + projected = current @ gu8.T + current = ( + F.silu(projected[:, :INTERMEDIATE]) + * projected[:, INTERMEDIATE:] + ) + if mode == "w8a8": + scale = ( + current.abs().amax(dim=1, keepdim=True).clamp_min(1e-8) + / 127.0 + ) + current = ( + current / scale + ).round().clamp(-127, 127) * scale + output = current @ dn8.T + scores[mode].append(F.cosine_similarity( + reference.flatten(), output.flatten(), dim=0).item()) + + for mode, use_rht, current_group in ( + ("int4_w4a4", False, group_size), + ("int4_rht16_w4a4", True, 16), + ): + transform = _hadamard16(gate_up.device) + gu_source = gate_up + current = activation + if use_rht: + gu_source = ( + gate_up.reshape(-1, 16) @ transform + ).reshape_as(gate_up) + current = ( + activation.reshape(-1, 16) @ transform + ).reshape_as(activation) + gu4, gu4_scale = _int4_weight( + gu_source, current_group) + gu4 = _dequant_int4( + gu4, gu4_scale, HIDDEN, current_group) + current4, current4_scale = _int4_weight( + current, current_group) + current = _dequant_int4( + current4, current4_scale, HIDDEN, current_group) + projected = current @ gu4.T + current = ( + F.silu(projected[:, :INTERMEDIATE]) + * projected[:, INTERMEDIATE:] + ) + dn_source = down + if use_rht: + dn_source = ( + down.reshape(-1, 16) @ transform + ).reshape_as(down) + current = ( + current.reshape(-1, 16) @ transform + ).reshape_as(current) + dn4, dn4_scale = _int4_weight( + dn_source, current_group) + dn4 = _dequant_int4( + dn4, dn4_scale, INTERMEDIATE, current_group) + current4, current4_scale = _int4_weight( + current, current_group) + current = _dequant_int4( + current4, current4_scale, INTERMEDIATE, current_group) + output = current @ dn4.T + scores[mode].append(F.cosine_similarity( + reference.flatten(), output.flatten(), dim=0).item()) + + for mode, values in scores.items(): + print( + f"{mode}: min={min(values):.6f} " + f"mean={statistics.mean(values):.6f}" + ) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--checkpoint", type=Path, required=True) + parser.add_argument("--mode", choices=("memory", "quality"), required=True) + parser.add_argument("--group-size", type=int, default=16) + parser.add_argument("--sample-layers", type=int, default=20) + parser.add_argument("--device", default="cuda:0") + parser.add_argument("--budget-gib", type=float, default=7.0) + parser.add_argument("--runtime-reserve-gib", type=float, default=1.5) + args = parser.parse_args() + if args.mode == "memory": + memory_probe( + args.checkpoint, + args.group_size, + args.budget_gib, + args.runtime_reserve_gib, + ) + else: + quality_probe( + args.checkpoint, + layers=args.sample_layers, + group_size=args.group_size, + device=args.device, + ) + + +if __name__ == "__main__": + main() diff --git a/qwen36_moe_edge/quantize_experts.py b/qwen36_moe_edge/quantize_experts.py new file mode 100644 index 00000000..a8b25c8e --- /dev/null +++ b/qwen36_moe_edge/quantize_experts.py @@ -0,0 +1,225 @@ +#!/usr/bin/env python3 +"""Write Qwen3.6 routed experts as fixed-size INT8 or INT4 blocks.""" + +from __future__ import annotations + +import argparse +import json +import os +import time +from pathlib import Path + +import torch +from safetensors import safe_open + + +NUM_LAYERS = 40 +NUM_EXPERTS = 256 +HIDDEN = 2048 +INTERMEDIATE = 512 + + +class CheckpointReader: + def __init__(self, checkpoint: Path): + index_path = checkpoint / "model.safetensors.index.json" + with index_path.open(encoding="utf-8") as f: + self.weight_map = json.load(f)["weight_map"] + self.readers = { + shard: safe_open( + checkpoint / shard, framework="pt", device="cpu") + for shard in set(self.weight_map.values()) + } + + def expert(self, layer: int, name: str, expert: int) -> torch.Tensor: + key = ( + f"model.language_model.layers.{layer}." + f"mlp.experts.{name}" + ) + shard = self.weight_map[key] + return self.readers[shard].get_slice(key)[expert] + + +def _int8_weight(weight: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + scale = ( + weight.float().abs().amax(dim=1).clamp_min(1e-8) / 127.0 + ).to(torch.float16) + quantized = ( + weight.float() / scale.float()[:, None] + ).round().clamp(-127, 127).to(torch.int8) + return quantized, scale + + +def _hadamard16(device: torch.device) -> torch.Tensor: + matrix = torch.ones(1, 1, dtype=torch.float32, device=device) + for _ in range(4): + matrix = torch.cat(( + torch.cat((matrix, matrix), dim=1), + torch.cat((matrix, -matrix), dim=1), + ), dim=0) + return matrix / 4.0 + + +def _rht16(weight: torch.Tensor) -> torch.Tensor: + rows, columns = weight.shape + return ( + weight.float().reshape(rows, columns // 16, 16) + @ _hadamard16(weight.device) + ).reshape(rows, columns) + + +def _int4_weight( + weight: torch.Tensor, group_size: int +) -> tuple[torch.Tensor, torch.Tensor]: + rows, columns = weight.shape + if columns % group_size: + raise ValueError( + f"K={columns} is not divisible by group_size={group_size}") + grouped = weight.float().reshape(rows, columns // group_size, group_size) + scale = ( + grouped.abs().amax(dim=2).clamp_min(1e-8) / 7.0 + ).to(torch.float8_e4m3fn) + scale_float = scale.float().clamp_min(2.0**-9) + values = ( + grouped / scale_float.unsqueeze(-1) + ).round().clamp(-7, 7).to(torch.int8).reshape(rows, columns) + magnitude = values.abs().to(torch.uint8) + code = magnitude | ((values < 0).to(torch.uint8) << 3) + packed = code[:, 0::2] | (code[:, 1::2] << 4) + return packed.contiguous(), scale.view(torch.uint8).contiguous() + + +def _tensor_bytes(tensor: torch.Tensor) -> bytes: + return tensor.contiguous().cpu().numpy().tobytes() + + +def quantize_expert( + gate_up: torch.Tensor, + down: torch.Tensor, + *, + quant_format: str, + group_size: int, + device: str, +) -> bytes: + if quant_format not in ("int8", "int4", "int4-rht"): + raise ValueError(f"unsupported quantization format: {quant_format}") + if quant_format == "int4-rht" and group_size != 16: + raise ValueError("int4-rht requires group_size=16") + gate_up = gate_up.to(device=device, dtype=torch.float32) + down = down.to(device=device, dtype=torch.float32) + if quant_format == "int8": + gu_weight, gu_scale = _int8_weight(gate_up) + dn_weight, dn_scale = _int8_weight(down) + else: + if quant_format == "int4-rht": + gate_up = _rht16(gate_up) + down = _rht16(down) + gu_weight, gu_scale = _int4_weight(gate_up, group_size) + dn_weight, dn_scale = _int4_weight(down, group_size) + return b"".join(( + _tensor_bytes(gu_weight), + _tensor_bytes(gu_scale), + _tensor_bytes(dn_weight), + _tensor_bytes(dn_scale), + )) + + +def _parse_layers(value: str) -> range: + start, stop = (int(part) for part in value.split(":", 1)) + if start < 0 or stop > NUM_LAYERS or start >= stop: + raise argparse.ArgumentTypeError( + f"layers must satisfy 0 <= start < stop <= {NUM_LAYERS}") + return range(start, stop) + + +def _layout(quant_format: str, group_size: int) -> dict[str, int]: + if quant_format == "int8": + gu_weight = 2 * INTERMEDIATE * HIDDEN + gu_scale = 2 * INTERMEDIATE * 2 + dn_weight = HIDDEN * INTERMEDIATE + dn_scale = HIDDEN * 2 + else: + gu_weight = 2 * INTERMEDIATE * HIDDEN // 2 + gu_scale = 2 * INTERMEDIATE * (HIDDEN // group_size) + dn_weight = HIDDEN * INTERMEDIATE // 2 + dn_scale = HIDDEN * (INTERMEDIATE // group_size) + return { + "gate_up_weight": gu_weight, + "gate_up_scale": gu_scale, + "down_weight": dn_weight, + "down_scale": dn_scale, + } + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--checkpoint", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument( + "--format", choices=("int8", "int4", "int4-rht"), required=True) + parser.add_argument("--group-size", type=int, default=16) + parser.add_argument("--layers", type=_parse_layers, default=range(40)) + parser.add_argument("--device", default="cuda:0") + args = parser.parse_args() + + if args.format != "int8" and args.group_size not in (16, 32, 64, 128): + parser.error("--group-size must be one of 16, 32, 64, or 128") + if args.format == "int4-rht" and args.group_size != 16: + parser.error("--format int4-rht requires --group-size 16") + + args.output.mkdir(parents=True, exist_ok=True) + layout = _layout(args.format, args.group_size) + block_bytes = sum(layout.values()) + manifest = { + "format": f"flashrt-qwen36-moe-{args.format}-experts-v1", + "group_size": args.group_size if args.format != "int8" else None, + "rht": args.format == "int4-rht", + "num_layers": NUM_LAYERS, + "num_experts": NUM_EXPERTS, + "hidden_size": HIDDEN, + "intermediate_size": INTERMEDIATE, + "block_layout": list(layout), + "block_sizes": layout, + "block_bytes": block_bytes, + } + with (args.output / "manifest.json").open("w", encoding="utf-8") as f: + json.dump(manifest, f, indent=2) + f.write("\n") + + reader = CheckpointReader(args.checkpoint) + for layer in args.layers: + output_path = args.output / f"experts_layer_{layer:02d}.bin" + expected_bytes = NUM_EXPERTS * block_bytes + if output_path.is_file() and output_path.stat().st_size == expected_bytes: + print(f"layer {layer}: already complete") + continue + temporary_path = output_path.with_suffix(".bin.tmp") + started = time.perf_counter() + with temporary_path.open("wb") as f: + for expert in range(NUM_EXPERTS): + gate_up = reader.expert( + layer, "gate_up_proj", expert) + down = reader.expert(layer, "down_proj", expert) + block = quantize_expert( + gate_up, + down, + quant_format=args.format, + group_size=args.group_size, + device=args.device, + ) + if len(block) != block_bytes: + raise RuntimeError( + f"expert block is {len(block)} bytes; " + f"expected {block_bytes}") + f.write(block) + os.replace(temporary_path, output_path) + elapsed = time.perf_counter() - started + gib = expected_bytes / 2**30 + print( + f"layer {layer}: {gib:.3f} GiB in {elapsed:.2f}s " + f"({gib / elapsed:.2f} GiB/s)", + flush=True, + ) + + +if __name__ == "__main__": + main() diff --git a/qwen36_moe_edge/route_trace.py b/qwen36_moe_edge/route_trace.py new file mode 100644 index 00000000..3f5519f8 --- /dev/null +++ b/qwen36_moe_edge/route_trace.py @@ -0,0 +1,122 @@ +#!/usr/bin/env python3 +"""Collect SM120 router selections and simulate bounded per-layer LRUs.""" + +from __future__ import annotations + +import argparse +import json +from collections import OrderedDict +from pathlib import Path + +import torch + +from flash_rt.frontends.torch._nexn2_rtx_decode import ( + Nexn2DecodeState, + generate_greedy, +) +from flash_rt.frontends.torch.qwen36_moe_rtx import ( + Qwen36MoeTextFrontendRtx, +) + + +def simulate_lru( + trace: list[list[list[int]]], + *, + prompt_tokens: int, + quota: int) -> dict[str, float]: + prompt_accesses = prompt_misses = 0 + decode_accesses = decode_misses = 0 + for layer_trace in trace: + cache: OrderedDict[int, None] = OrderedDict() + for step, experts in enumerate(layer_trace): + prompt = step < prompt_tokens + for expert in experts: + if prompt: + prompt_accesses += 1 + else: + decode_accesses += 1 + if expert in cache: + cache.move_to_end(expert) + continue + if prompt: + prompt_misses += 1 + else: + decode_misses += 1 + if len(cache) >= quota: + cache.popitem(last=False) + cache[expert] = None + decode_steps = len(trace[0]) - prompt_tokens + return { + "prompt_hit_rate": 1.0 - prompt_misses / prompt_accesses, + "decode_hit_rate": 1.0 - decode_misses / decode_accesses, + "decode_misses_per_token": decode_misses / decode_steps, + } + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--checkpoint", required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--prompt", required=True) + parser.add_argument("--prompt-tokens", type=int, default=32) + parser.add_argument("--new-tokens", type=int, default=64) + parser.add_argument("--max-seq", type=int, default=128) + parser.add_argument("--quotas", default="8,16,24,32,48,64") + parser.add_argument("--device", default="cuda:0") + args = parser.parse_args() + + frontend = Qwen36MoeTextFrontendRtx( + args.checkpoint, + device=args.device, + max_seq=args.max_seq, + quant_scope="experts", + ) + input_ids = frontend.tokenizer( + args.prompt, + return_tensors="pt", + add_special_tokens=False, + ).input_ids[:, :args.prompt_tokens].to(args.device) + if input_ids.shape[1] != args.prompt_tokens: + parser.error("the supplied prompt is shorter than --prompt-tokens") + + state = Nexn2DecodeState( + frontend._weights, args.max_seq, args.device) + state.batched_prefill = False + state.router_trace = { + layer: [] for layer in range(state.num_layers)} + with torch.no_grad(): + generated = generate_greedy( + state, + input_ids, + args.new_tokens, + frontend._fvk, + args.device, + ) + + trace = [ + [list(experts) for experts in state.router_trace[layer]] + for layer in range(state.num_layers) + ] + result = { + "prompt_tokens": args.prompt_tokens, + "generated_tokens": generated, + "trace": trace, + "lru": {}, + } + for quota in (int(value) for value in args.quotas.split(",")): + result["lru"][str(quota)] = simulate_lru( + trace, prompt_tokens=args.prompt_tokens, quota=quota) + args.output.parent.mkdir(parents=True, exist_ok=True) + with args.output.open("w", encoding="utf-8") as f: + json.dump(result, f) + f.write("\n") + + for quota, values in result["lru"].items(): + print( + f"quota={quota}: decode_hit={values['decode_hit_rate']:.4f}, " + f"misses/token={values['decode_misses_per_token']:.2f}" + ) + + +if __name__ == "__main__": + main() diff --git a/tests/test_qwen36_moe_edge_quant.py b/tests/test_qwen36_moe_edge_quant.py new file mode 100644 index 00000000..79dbf65d --- /dev/null +++ b/tests/test_qwen36_moe_edge_quant.py @@ -0,0 +1,98 @@ +"""Structural tests for Qwen3.6-MoE edge checkpoint utilities.""" + +from __future__ import annotations + +import torch + +from qwen36_moe_edge.probe import _dequant_int4 +from qwen36_moe_edge.route_trace import simulate_lru +from qwen36_moe_edge.quantize_experts import ( + HIDDEN, + INTERMEDIATE, + _hadamard16, + _int4_weight, + _int8_weight, + _layout, + _rht16, + quantize_expert, +) + + +def test_int8_per_channel_round_trip_is_precise(): + generator = torch.Generator().manual_seed(3) + weight = torch.randn(64, 256, generator=generator) + + quantized, scale = _int8_weight(weight) + restored = quantized.float() * scale.float()[:, None] + + cosine = torch.nn.functional.cosine_similarity( + weight.flatten(), restored.flatten(), dim=0) + assert cosine > 0.9999 + + +def test_int4_grouped_round_trip_matches_packed_layout(): + generator = torch.Generator().manual_seed(5) + weight = torch.randn(64, 256, generator=generator) + + packed, scale = _int4_weight(weight, 32) + restored = _dequant_int4(packed, scale, 256, 32) + + assert packed.shape == (64, 128) + assert scale.shape == (64, 8) + cosine = torch.nn.functional.cosine_similarity( + weight.flatten(), restored.flatten(), dim=0) + assert cosine > 0.99 + + +def test_rht16_is_orthonormal_and_preserves_matmul(): + generator = torch.Generator().manual_seed(7) + activation = torch.randn(3, 32, generator=generator) + weight = torch.randn(5, 32, generator=generator) + transform = _hadamard16(weight.device) + + identity = transform @ transform.T + rotated_activation = _rht16(activation) + rotated_weight = _rht16(weight) + + assert torch.equal(identity, torch.eye(16)) + torch.testing.assert_close( + rotated_activation @ rotated_weight.T, + activation @ weight.T, + rtol=1e-5, + atol=1e-5, + ) + + +def test_expert_block_size_matches_manifest_layout(): + gate_up = torch.zeros(2 * INTERMEDIATE, HIDDEN) + down = torch.zeros(HIDDEN, INTERMEDIATE) + + for quant_format, group_size in ( + ("int8", 32), + ("int4", 32), + ("int4-rht", 16), + ): + block = quantize_expert( + gate_up, + down, + quant_format=quant_format, + group_size=group_size, + device="cpu", + ) + assert len(block) == sum( + _layout(quant_format, group_size).values()) + + +def test_route_trace_lru_separates_prompt_and_decode(): + trace = [ + [[0, 1], [0, 2], [0, 1], [2, 3]], + [[4, 5], [4, 6], [4, 5], [6, 7]], + ] + + result = simulate_lru(trace, prompt_tokens=2, quota=2) + + assert result == { + "prompt_hit_rate": 0.25, + "decode_hit_rate": 0.25, + "decode_misses_per_token": 3.0, + } From 12e34a7c58fb4646e603528069e2802b18831eef Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 30 Jul 2026 04:57:53 -0400 Subject: [PATCH 04/85] Split qwen3_5_moe kernels into portability tiers The family was behind a single FLASHRT_ENABLE_QWEN35MOE gate that also required a Blackwell NVFP4 build, so a target that can run most of these kernels could not select them. Only four of the fourteen actually need block-scaled MMA. Three gates now describe what each kernel needs: - _CORE: layout/split, bf16 matvec, router top-k, activation fusion, GDN recurrence, weighted-sum reducer, bf16 GEMM. bf16 intrinsics plus cp.async and mma.m16n8k16.bf16, so SM80 and newer. - _W4A16: weight-only 4-bit matvec, grouped matvec and GEMM. Operands go through __nv_cvt_fp4x2_to_halfraw2 and accumulate in bf16, needing no block-scaled MMA; SM89 and newer get the hardware conversion. - _W4A4: grouped GEMV and the M16/M64/block-tile MMA kernels, which need the sm_120a/sm_121a CUTLASS path. _W4A4 refuses to configure without that path. CUTLASS compiles those translation units on other architectures but substitutes CUTE_INVALID_CONTROL_PATH for the MMA, so the build would succeed and the kernels would fail when called; the gate turns that into a configure error. FLASHRT_ENABLE_QWEN35MOE stays as the all-tiers switch, so existing configure lines and the frontend's fail-fast message are unchanged. The bindings are regrouped to match the three macros; no binding signature or symbol name changed. Verified: sm_120 with the alias enables all three tiers and exposes the same 16 bindings; sm_87 configures _CORE and _W4A16 and compiles all ten translation units; sm_87 with _W4A4 fails at configure. --- CMakeLists.txt | 102 +++++++++++++++++++------ csrc/bindings.cpp | 160 ++++++++++++++++++++------------------- docs/qwen36_moe_usage.md | 19 +++++ 3 files changed, 183 insertions(+), 98 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 0b196018..53e509c8 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -163,13 +163,47 @@ option(FLASHRT_ENABLE_LINGBOT "Build LingBot-VLA Thor (sm_110) model kernels and bindings" ON) option(FLASHRT_ENABLE_MOTUS "Build Motus-specific RTX SM120 kernels and bindings" ON) -# qwen3_5_moe family (Nex-N2-mini / Qwen3.5-3.6-35B-A3B) RTX SM120 kernels. -# OFF by default: these are NVFP4/sm_120a-only block-scaled kernels and must -# not enter the default flash_rt_kernels build on SM89/87/110 or for models -# that don't use them. Enable explicitly (-DFLASHRT_ENABLE_QWEN35MOE=ON) on a -# Blackwell build to get the Nex-N2 / Qwen3.6-35B-A3B path. +# qwen3_5_moe family (Nex-N2-mini / Qwen3.5-3.6-35B-A3B) kernels, split into +# three tiers by the hardware each tier actually requires. All are OFF by +# default so they never enter the flash_rt_kernels build for models that do +# not use them. +# +# _CORE Architecture-neutral: layout/split, bf16 matvec, router top-k, +# activation fusion, GDN recurrence, weighted-sum reducer, and the +# bf16 GEMM. Uses only bf16 intrinsics plus cp.async and +# mma.m16n8k16.bf16, i.e. SM80 and newer. +# _W4A16 Weight-only 4-bit: converts packed operands with the +# __nv_cvt_fp4x2_to_halfraw2 intrinsic and accumulates in bf16, so +# it needs no block-scaled MMA. SM89 and newer emit the hardware +# conversion; older targets take the intrinsic's software path. +# _W4A4 Block-scaled 4-bit MMA through CUTLASS. Requires the +# CUTE_ARCH_F8F6F4_MMA_ENABLED path, which only sm_120a/sm_121a +# provide. On other targets CUTLASS still compiles these +# translation units but replaces the MMA with +# CUTE_INVALID_CONTROL_PATH, so the gate must be explicit: a +# successful build would otherwise produce kernels that fail at +# run time. +# +# FLASHRT_ENABLE_QWEN35MOE remains the single switch for a full Blackwell +# build and turns on all three tiers, so existing configure lines are +# unchanged. option(FLASHRT_ENABLE_QWEN35MOE - "Build qwen3_5_moe (Nex-N2 / Qwen3.6-35B-A3B) RTX SM120 kernels" OFF) + "Build all qwen3_5_moe (Nex-N2 / Qwen3.6-35B-A3B) kernel tiers" OFF) +option(FLASHRT_ENABLE_QWEN35MOE_CORE + "Build architecture-neutral qwen3_5_moe kernels" OFF) +option(FLASHRT_ENABLE_QWEN35MOE_W4A16 + "Build weight-only 4-bit qwen3_5_moe kernels" OFF) +option(FLASHRT_ENABLE_QWEN35MOE_W4A4 + "Build block-scaled 4-bit MMA qwen3_5_moe kernels (sm_120a/sm_121a)" OFF) +if(FLASHRT_ENABLE_QWEN35MOE) + set(FLASHRT_ENABLE_QWEN35MOE_CORE ON) + set(FLASHRT_ENABLE_QWEN35MOE_W4A16 ON) + set(FLASHRT_ENABLE_QWEN35MOE_W4A4 ON) +endif() +# The upper tiers reuse the core layout, reducer, and activation kernels. +if(FLASHRT_ENABLE_QWEN35MOE_W4A16 OR FLASHRT_ENABLE_QWEN35MOE_W4A4) + set(FLASHRT_ENABLE_QWEN35MOE_CORE ON) +endif() # Qwen3-VL adds a handful of kernels (rotate_half RoPE, ...) that the shared # binary does not yet need. Built into a SEPARATE flash_rt_qwen3_vl_kernels # module so flash_rt_kernels.so stays stable; OFF by default. @@ -1465,30 +1499,54 @@ else() message(STATUS "NVFP4 swizzle/quantize kernels: SKIPPED (slim build)") endif() -# ── qwen3_5_moe family (Nex-N2 / Qwen3.6-35B-A3B) SM120 kernels ── -# Gated: only compiled when explicitly enabled AND on a Blackwell (NVFP4) -# build. The bindings are #ifdef FLASHRT_HAVE_QWEN35MOE in bindings.cpp, so -# the symbols are absent (and the .cu never compiled) on every other target. -if(FLASHRT_ENABLE_QWEN35MOE AND ENABLE_NVFP4) +# ── qwen3_5_moe family (Nex-N2 / Qwen3.6-35B-A3B) kernels ── +# Three tiers, each with its own compile definition. The matching bindings in +# csrc/bindings.cpp are guarded on the same macros, so a tier that is off +# contributes neither symbols nor translation units. +if(FLASHRT_ENABLE_QWEN35MOE_CORE) target_sources(flash_rt_kernels PRIVATE csrc/kernels/qwen35moe_layout.cu - csrc/kernels/moe_grouped_gemv_sm120.cu csrc/kernels/bf16_matvec_sm120.cu - csrc/kernels/w4a16_matvec_sm120.cu - csrc/kernels/moe_grouped_w4a16_sm120.cu csrc/kernels/gdn_recurrent_seq_sm120.cu csrc/kernels/act_fuse_sm120.cu csrc/kernels/moe_router_topk_sm120.cu - csrc/kernels/moe_m16_mma_sm120.cu - csrc/kernels/moe_m64_mma_sm120.cu - csrc/kernels/moe_blocktile_mma_sm120.cu csrc/kernels/moe_weighted_sum_sm120.cu - csrc/kernels/w4a16_gemm_sm120.cu csrc/kernels/w16a16_gemm_sm120.cu) - target_compile_definitions(flash_rt_kernels PRIVATE FLASHRT_HAVE_QWEN35MOE=1) - message(STATUS "qwen3_5_moe (Nex-N2 / Qwen3.6-35B-A3B) SM120 kernels: ENABLED") -elseif(FLASHRT_ENABLE_QWEN35MOE) - message(STATUS "qwen3_5_moe SM120 kernels: SKIPPED (requires Blackwell NVFP4)") + target_compile_definitions(flash_rt_kernels PRIVATE + FLASHRT_HAVE_QWEN35MOE_CORE=1) + message(STATUS "qwen3_5_moe core kernels: ENABLED (sm_${GPU_ARCH})") +endif() + +if(FLASHRT_ENABLE_QWEN35MOE_W4A16) + target_sources(flash_rt_kernels PRIVATE + csrc/kernels/w4a16_matvec_sm120.cu + csrc/kernels/moe_grouped_w4a16_sm120.cu + csrc/kernels/w4a16_gemm_sm120.cu) + target_compile_definitions(flash_rt_kernels PRIVATE + FLASHRT_HAVE_QWEN35MOE_W4A16=1) + message(STATUS "qwen3_5_moe weight-only 4-bit kernels: ENABLED (sm_${GPU_ARCH})") +endif() + +# Block-scaled 4-bit MMA needs the sm_120a/sm_121a CUTLASS path. Refuse rather +# than build silently broken kernels on other targets. +if(FLASHRT_ENABLE_QWEN35MOE_W4A4) + if(NOT ENABLE_NVFP4) + message(FATAL_ERROR + "FLASHRT_ENABLE_QWEN35MOE_W4A4 requires a block-scaled-MMA target " + "(current GPU_ARCH=${GPU_ARCH}). CUTLASS compiles these kernels on " + "other architectures but replaces the MMA with an invalid control " + "path, so they would fail at run time. Use " + "FLASHRT_ENABLE_QWEN35MOE_CORE and FLASHRT_ENABLE_QWEN35MOE_W4A16 " + "instead.") + endif() + target_sources(flash_rt_kernels PRIVATE + csrc/kernels/moe_grouped_gemv_sm120.cu + csrc/kernels/moe_m16_mma_sm120.cu + csrc/kernels/moe_m64_mma_sm120.cu + csrc/kernels/moe_blocktile_mma_sm120.cu) + target_compile_definitions(flash_rt_kernels PRIVATE + FLASHRT_HAVE_QWEN35MOE_W4A4=1) + message(STATUS "qwen3_5_moe block-scaled 4-bit kernels: ENABLED (sm_${GPU_ARCH})") endif() # ── MelBandRoformer custom fused kernels (BF16/FP8, gated) ── diff --git a/csrc/bindings.cpp b/csrc/bindings.cpp index 3faedce0..f9e191a3 100644 --- a/csrc/bindings.cpp +++ b/csrc/bindings.cpp @@ -178,22 +178,26 @@ extern "C" int cutlass_int8_rowwise_bf16out_t64x128( #ifdef FLASHRT_HAVE_QWEN36_KERNELS #include "kernels/qwen36_misc.cuh" #endif -#ifdef FLASHRT_HAVE_QWEN35MOE +#ifdef FLASHRT_HAVE_QWEN35MOE_CORE #include "kernels/qwen35moe_layout.cuh" -#include "kernels/moe_grouped_gemv_sm120.cuh" #include "kernels/bf16_matvec_sm120.cuh" -#include "kernels/w4a16_matvec_sm120.cuh" -#include "kernels/moe_grouped_w4a16_sm120.cuh" #include "kernels/gdn_recurrent_seq_sm120.cuh" #include "kernels/act_fuse_sm120.cuh" #include "kernels/moe_router_topk_sm120.cuh" +#include "kernels/moe_weighted_sum_sm120.cuh" +#include "kernels/w16a16_gemm_sm120.cuh" +#endif // FLASHRT_HAVE_QWEN35MOE_CORE +#ifdef FLASHRT_HAVE_QWEN35MOE_W4A16 +#include "kernels/w4a16_matvec_sm120.cuh" +#include "kernels/moe_grouped_w4a16_sm120.cuh" +#include "kernels/w4a16_gemm_sm120.cuh" +#endif // FLASHRT_HAVE_QWEN35MOE_W4A16 +#ifdef FLASHRT_HAVE_QWEN35MOE_W4A4 +#include "kernels/moe_grouped_gemv_sm120.cuh" #include "kernels/moe_m16_mma_sm120.cuh" #include "kernels/moe_m64_mma_sm120.cuh" #include "kernels/moe_blocktile_mma_sm120.cuh" -#include "kernels/moe_weighted_sum_sm120.cuh" -#include "kernels/w4a16_gemm_sm120.cuh" -#include "kernels/w16a16_gemm_sm120.cuh" -#endif // FLASHRT_HAVE_QWEN35MOE +#endif // FLASHRT_HAVE_QWEN35MOE_W4A4 #include "kernels/bf16_matvec_qwen36.cuh" #include "kernels/bf16_matmul_bf16.cuh" #ifdef FLASHRT_HAVE_QWEN36_KERNELS @@ -5384,7 +5388,7 @@ PYBIND11_MODULE(flash_rt_kernels, m) { py::arg("S"), py::arg("stream") = 0); #endif // FLASHRT_HAVE_QWEN36_KERNELS (gated_deltanet_qwen36 part 1) -#ifdef FLASHRT_HAVE_QWEN35MOE +#ifdef FLASHRT_HAVE_QWEN35MOE_CORE m.def("qwen35moe_lin_split_qkv_broadcast_bf16", [](uintptr_t conv_out, uintptr_t q32, uintptr_t k32, uintptr_t v32, @@ -5417,61 +5421,6 @@ PYBIND11_MODULE(flash_rt_kernels, m) { py::arg("x"), py::arg("W"), py::arg("out"), py::arg("N"), py::arg("K"), py::arg("stream") = 0); - m.def("w4a16_matvec_sm120_bf16", - [](uintptr_t x, uintptr_t W, uintptr_t sfb, uintptr_t out, - int N, int K, float alpha, uintptr_t stream) -> int { - return flash_rt::kernels::w4a16_matvec_sm120_bf16( - to_ptr(x), to_ptr(W), to_ptr(sfb), to_ptr(out), - N, K, alpha, to_stream(stream)); - }, - py::arg("x"), py::arg("W"), py::arg("sfb"), py::arg("out"), - py::arg("N"), py::arg("K"), py::arg("alpha"), py::arg("stream") = 0); - - m.def("moe_m16_mma_sm120_bf16", - [](uintptr_t A, uintptr_t B, uintptr_t sfa, uintptr_t sfb, uintptr_t D, - uintptr_t alpha, uintptr_t te, int num_tiles, int N, int K, - long sfa_stride, long w_stride, long sfb_stride, - uintptr_t stream) -> int { - return flash_rt::gemm::moe_m16_mma_sm120_bf16( - to_ptr(A), to_ptr(B), to_ptr(sfa), to_ptr(sfb), to_ptr(D), - to_ptr(alpha), to_ptr(te), num_tiles, N, K, - sfa_stride, w_stride, sfb_stride, to_stream(stream)); - }, - py::arg("A"), py::arg("B"), py::arg("sfa"), py::arg("sfb"), - py::arg("D"), py::arg("alpha"), py::arg("te"), py::arg("num_tiles"), - py::arg("N"), py::arg("K"), py::arg("sfa_stride"), py::arg("w_stride"), - py::arg("sfb_stride"), py::arg("stream") = 0); - - m.def("moe_m64_mma_sm120_bf16", - [](uintptr_t A, uintptr_t B, uintptr_t sfa, uintptr_t sfb, uintptr_t D, - uintptr_t alpha, uintptr_t te, int num_tiles, int N, int K, - long sfa_stride, long w_stride, long sfb_stride, - uintptr_t stream) -> int { - return flash_rt::gemm::moe_m64_mma_sm120_bf16( - to_ptr(A), to_ptr(B), to_ptr(sfa), to_ptr(sfb), to_ptr(D), - to_ptr(alpha), to_ptr(te), num_tiles, N, K, - sfa_stride, w_stride, sfb_stride, to_stream(stream)); - }, - py::arg("A"), py::arg("B"), py::arg("sfa"), py::arg("sfb"), - py::arg("D"), py::arg("alpha"), py::arg("te"), py::arg("num_tiles"), - py::arg("N"), py::arg("K"), py::arg("sfa_stride"), py::arg("w_stride"), - py::arg("sfb_stride"), py::arg("stream") = 0); - - m.def("moe_blocktile_mma_sm120_bf16", - [](uintptr_t A, uintptr_t B, uintptr_t sfa, uintptr_t sfb, uintptr_t D, - uintptr_t alpha, uintptr_t te, int num_tiles, int N, int K, - long sfa_stride, long w_stride, long sfb_stride, - uintptr_t stream) -> int { - return flash_rt::gemm::moe_blocktile_mma_sm120_bf16( - to_ptr(A), to_ptr(B), to_ptr(sfa), to_ptr(sfb), to_ptr(D), - to_ptr(alpha), to_ptr(te), num_tiles, N, K, - sfa_stride, w_stride, sfb_stride, to_stream(stream)); - }, - py::arg("A"), py::arg("B"), py::arg("sfa"), py::arg("sfb"), - py::arg("D"), py::arg("alpha"), py::arg("te"), py::arg("num_tiles"), - py::arg("N"), py::arg("K"), py::arg("sfa_stride"), py::arg("w_stride"), - py::arg("sfb_stride"), py::arg("stream") = 0); - m.def("moe_weighted_sum_sm120_bf16", [](uintptr_t d_dn, uintptr_t rows, uintptr_t tw, uintptr_t out, int S, int TOPK, int HID, int dn_stride, uintptr_t stream) -> int { @@ -5483,17 +5432,6 @@ PYBIND11_MODULE(flash_rt_kernels, m) { py::arg("S"), py::arg("TOPK"), py::arg("HID"), py::arg("dn_stride"), py::arg("stream") = 0); - m.def("w4a16_gemm_sm120_bf16", - [](uintptr_t X, uintptr_t W, uintptr_t SFB, uintptr_t Y, - int M, int N, int K, float alpha, uintptr_t stream) -> int { - return flash_rt::gemm::w4a16_gemm_sm120_bf16( - to_ptr(X), to_ptr(W), to_ptr(SFB), to_ptr(Y), - M, N, K, alpha, to_stream(stream)); - }, - py::arg("X"), py::arg("W"), py::arg("SFB"), py::arg("Y"), - py::arg("M"), py::arg("N"), py::arg("K"), - py::arg("alpha") = 1.0f, py::arg("stream") = 0); - m.def("w16a16_gemm_sm120_bf16", [](uintptr_t X, uintptr_t W, uintptr_t Y, int M, int N, int K, float alpha, uintptr_t stream) -> int { @@ -5546,6 +5484,29 @@ PYBIND11_MODULE(flash_rt_kernels, m) { py::arg("state"), py::arg("out"), py::arg("S"), py::arg("num_v_heads"), py::arg("head_dim"), py::arg("use_qk_l2norm") = true, py::arg("stream") = 0); +#endif // FLASHRT_HAVE_QWEN35MOE_CORE + +#ifdef FLASHRT_HAVE_QWEN35MOE_W4A16 + m.def("w4a16_matvec_sm120_bf16", + [](uintptr_t x, uintptr_t W, uintptr_t sfb, uintptr_t out, + int N, int K, float alpha, uintptr_t stream) -> int { + return flash_rt::kernels::w4a16_matvec_sm120_bf16( + to_ptr(x), to_ptr(W), to_ptr(sfb), to_ptr(out), + N, K, alpha, to_stream(stream)); + }, + py::arg("x"), py::arg("W"), py::arg("sfb"), py::arg("out"), + py::arg("N"), py::arg("K"), py::arg("alpha"), py::arg("stream") = 0); + + m.def("w4a16_gemm_sm120_bf16", + [](uintptr_t X, uintptr_t W, uintptr_t SFB, uintptr_t Y, + int M, int N, int K, float alpha, uintptr_t stream) -> int { + return flash_rt::gemm::w4a16_gemm_sm120_bf16( + to_ptr(X), to_ptr(W), to_ptr(SFB), to_ptr(Y), + M, N, K, alpha, to_stream(stream)); + }, + py::arg("X"), py::arg("W"), py::arg("SFB"), py::arg("Y"), + py::arg("M"), py::arg("N"), py::arg("K"), + py::arg("alpha") = 1.0f, py::arg("stream") = 0); m.def("moe_grouped_w4a16_sm120_bf16", [](uintptr_t A, uintptr_t W, uintptr_t sfb, uintptr_t alpha, @@ -5561,6 +5522,53 @@ PYBIND11_MODULE(flash_rt_kernels, m) { py::arg("eidx"), py::arg("D"), py::arg("slots"), py::arg("N"), py::arg("K"), py::arg("a_stride"), py::arg("w_stride"), py::arg("sfb_stride"), py::arg("stream") = 0); +#endif // FLASHRT_HAVE_QWEN35MOE_W4A16 + +#ifdef FLASHRT_HAVE_QWEN35MOE_W4A4 + m.def("moe_m16_mma_sm120_bf16", + [](uintptr_t A, uintptr_t B, uintptr_t sfa, uintptr_t sfb, uintptr_t D, + uintptr_t alpha, uintptr_t te, int num_tiles, int N, int K, + long sfa_stride, long w_stride, long sfb_stride, + uintptr_t stream) -> int { + return flash_rt::gemm::moe_m16_mma_sm120_bf16( + to_ptr(A), to_ptr(B), to_ptr(sfa), to_ptr(sfb), to_ptr(D), + to_ptr(alpha), to_ptr(te), num_tiles, N, K, + sfa_stride, w_stride, sfb_stride, to_stream(stream)); + }, + py::arg("A"), py::arg("B"), py::arg("sfa"), py::arg("sfb"), + py::arg("D"), py::arg("alpha"), py::arg("te"), py::arg("num_tiles"), + py::arg("N"), py::arg("K"), py::arg("sfa_stride"), py::arg("w_stride"), + py::arg("sfb_stride"), py::arg("stream") = 0); + + m.def("moe_m64_mma_sm120_bf16", + [](uintptr_t A, uintptr_t B, uintptr_t sfa, uintptr_t sfb, uintptr_t D, + uintptr_t alpha, uintptr_t te, int num_tiles, int N, int K, + long sfa_stride, long w_stride, long sfb_stride, + uintptr_t stream) -> int { + return flash_rt::gemm::moe_m64_mma_sm120_bf16( + to_ptr(A), to_ptr(B), to_ptr(sfa), to_ptr(sfb), to_ptr(D), + to_ptr(alpha), to_ptr(te), num_tiles, N, K, + sfa_stride, w_stride, sfb_stride, to_stream(stream)); + }, + py::arg("A"), py::arg("B"), py::arg("sfa"), py::arg("sfb"), + py::arg("D"), py::arg("alpha"), py::arg("te"), py::arg("num_tiles"), + py::arg("N"), py::arg("K"), py::arg("sfa_stride"), py::arg("w_stride"), + py::arg("sfb_stride"), py::arg("stream") = 0); + + m.def("moe_blocktile_mma_sm120_bf16", + [](uintptr_t A, uintptr_t B, uintptr_t sfa, uintptr_t sfb, uintptr_t D, + uintptr_t alpha, uintptr_t te, int num_tiles, int N, int K, + long sfa_stride, long w_stride, long sfb_stride, + uintptr_t stream) -> int { + return flash_rt::gemm::moe_blocktile_mma_sm120_bf16( + to_ptr(A), to_ptr(B), to_ptr(sfa), to_ptr(sfb), to_ptr(D), + to_ptr(alpha), to_ptr(te), num_tiles, N, K, + sfa_stride, w_stride, sfb_stride, to_stream(stream)); + }, + py::arg("A"), py::arg("B"), py::arg("sfa"), py::arg("sfb"), + py::arg("D"), py::arg("alpha"), py::arg("te"), py::arg("num_tiles"), + py::arg("N"), py::arg("K"), py::arg("sfa_stride"), py::arg("w_stride"), + py::arg("sfb_stride"), py::arg("stream") = 0); m.def("moe_grouped_gemv_sm120_bf16", [](uintptr_t A_stack, uintptr_t B_stack, uintptr_t D, @@ -5583,7 +5591,7 @@ PYBIND11_MODULE(flash_rt_kernels, m) { py::arg("a_stride"), py::arg("sfa_stride"), py::arg("w_stride"), py::arg("sfb_stride"), py::arg("stream") = 0); -#endif // FLASHRT_HAVE_QWEN35MOE +#endif // FLASHRT_HAVE_QWEN35MOE_W4A4 #ifdef FLASHRT_HAVE_QWEN36_KERNELS m.def("qwen36_gdn_gating_bf16", diff --git a/docs/qwen36_moe_usage.md b/docs/qwen36_moe_usage.md index d2f31235..1fa7b40c 100644 --- a/docs/qwen36_moe_usage.md +++ b/docs/qwen36_moe_usage.md @@ -31,6 +31,25 @@ cmake --build build -j pip install -e ".[torch]" ``` +### Kernel tiers + +`FLASHRT_ENABLE_QWEN35MOE=ON` is a convenience switch for all three tiers +below. Targets that cannot run a tier can select the remainder explicitly. + +| Flag | Kernels | Requires | +|---|---|---| +| `FLASHRT_ENABLE_QWEN35MOE_CORE` | QKV layout/split, bf16 matvec, router top-k, SiLU/sigmoid fusion, GDN recurrence, weighted-sum reducer, bf16 GEMM | SM80 and newer | +| `FLASHRT_ENABLE_QWEN35MOE_W4A16` | weight-only 4-bit matvec, grouped matvec, GEMM | SM80 and newer; hardware operand conversion from SM89 | +| `FLASHRT_ENABLE_QWEN35MOE_W4A4` | block-scaled 4-bit MMA: grouped GEMV, M16/M64/block-tile MMA | sm_120a / sm_121a | + +The upper tiers depend on the core tier, so enabling either turns it on. The +SM120 text runtime documented here needs all three. + +`_W4A4` refuses to configure on a target without block-scaled MMA. CUTLASS +still compiles those translation units elsewhere, but substitutes +`CUTE_INVALID_CONTROL_PATH` for the MMA, so the build would succeed and then +fail at run time. The explicit gate turns that into a configure-time error. + ## Usage ```python From 1af932145199b9c3a9f06ea9ed8d7eb7cf1c070e Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 30 Jul 2026 05:01:37 -0400 Subject: [PATCH 05/85] Simulate a two-tier expert cache in the route trace A single per-layer LRU is the wrong model for this runtime. Prefill selects far more experts per layer than the cache can hold, so it evicts whatever decode is about to need and the measured hit rate reflects prompt churn rather than routing locality. route_trace now scores each quota under three policies: the existing single LRU, a warm set pinned from prompt-phase selection counts plus an evictable ring, and the same split with an oracle warm set taken from the decode phase. The oracle is not implementable and is there to bound what a better warm-set heuristic could add. Misses per token are also converted to a read volume and to the token rate each storage bandwidth would allow, since that is what decides whether a given memory budget is viable at all. --- qwen36_moe_edge/README.md | 18 +++ qwen36_moe_edge/route_trace.py | 187 ++++++++++++++++++++++++++-- tests/test_qwen36_moe_edge_quant.py | 56 ++++++++- 3 files changed, 248 insertions(+), 13 deletions(-) diff --git a/qwen36_moe_edge/README.md b/qwen36_moe_edge/README.md index c3f973fa..39a5aebd 100644 --- a/qwen36_moe_edge/README.md +++ b/qwen36_moe_edge/README.md @@ -57,8 +57,26 @@ PYTHONPATH=. python qwen36_moe_edge/route_trace.py \ --prompt "Explain edge mixture-of-experts inference. " \ --prompt-tokens 32 \ --new-tokens 64 \ + --quotas 16,27,32,43,64 \ --output qwen36_moe_route_trace.json ``` Tracing deliberately uses eager per-token prefill. It must not be enabled during CUDA Graph capture. + +Each quota is scored under three policies: + +- `single_lru` — one per-layer LRU behind both prefill and decode. Prefill + touches every expert in a layer, so this measures what survives prompt + churn. +- `two_tier` — a per-layer warm set pinned from prompt-phase selection counts + plus an evictable ring sized by `--stream-fraction`. Prefill cannot displace + the warm set. +- `two_tier_oracle_warm` — the same split with the warm set chosen from the + decode phase. Not implementable; it bounds what a better warm-set heuristic + could add. + +`--block-bytes` (default: the INT4 group-16 block) and `--bandwidths` turn +misses per token into a read volume and the token rate each storage bandwidth +would allow, which is the number that decides whether a memory budget is +viable. diff --git a/qwen36_moe_edge/route_trace.py b/qwen36_moe_edge/route_trace.py index 3f5519f8..87516e07 100644 --- a/qwen36_moe_edge/route_trace.py +++ b/qwen36_moe_edge/route_trace.py @@ -1,11 +1,28 @@ #!/usr/bin/env python3 -"""Collect SM120 router selections and simulate bounded per-layer LRUs.""" +"""Collect router selections and simulate bounded per-layer expert caches. + +Two cache policies are simulated from the same trace: + +``simulate_lru`` + One LRU per layer holding every access. Prefill touches far more experts + than the cache can hold, so by the time decode starts the LRU contains + whatever the end of the prompt happened to use. + +``simulate_two_tier`` + A per-layer warm set chosen from the prompt and never evicted, plus a + small LRU ring for everything else. Prefill can no longer displace the + warm set, so the decode hit rate follows warm-set coverage instead of + prompt churn. + +The reported miss count per token, multiplied by the expert block size, is the +per-token read volume a streaming runtime has to sustain. +""" from __future__ import annotations import argparse import json -from collections import OrderedDict +from collections import Counter, OrderedDict from pathlib import Path import torch @@ -19,11 +36,16 @@ ) +# INT4 group-16 routed-expert block, matching quantize_experts._layout. +DEFAULT_BLOCK_BYTES = 1769472 + + def simulate_lru( trace: list[list[list[int]]], *, prompt_tokens: int, quota: int) -> dict[str, float]: + """Single per-layer LRU shared by prefill and decode.""" prompt_accesses = prompt_misses = 0 decode_accesses = decode_misses = 0 for layer_trace in trace: @@ -53,6 +75,133 @@ def simulate_lru( } +def simulate_two_tier( + trace: list[list[list[int]]], + *, + prompt_tokens: int, + pinned: int, + stream: int, + warm_from: str = "prompt") -> dict[str, float]: + """Pinned per-layer warm set plus a per-layer LRU ring. + + ``warm_from="prompt"`` selects the warm set from prompt-phase selection + counts, which is what a runtime can actually do. ``warm_from="decode"`` + selects it from the decode phase instead; that is not implementable, but + it bounds how much a better warm-set heuristic could win. + """ + if warm_from not in ("prompt", "decode"): + raise ValueError(f"unsupported warm_from: {warm_from!r}") + decode_accesses = decode_misses = warm_hits = 0 + for layer_trace in trace: + source = ( + layer_trace[:prompt_tokens] if warm_from == "prompt" + else layer_trace[prompt_tokens:] + ) + counts: Counter[int] = Counter() + for experts in source: + counts.update(experts) + warm = {expert for expert, _ in counts.most_common(pinned)} + ring: OrderedDict[int, None] = OrderedDict() + for experts in layer_trace[prompt_tokens:]: + for expert in experts: + decode_accesses += 1 + if expert in warm: + warm_hits += 1 + continue + if expert in ring: + ring.move_to_end(expert) + continue + decode_misses += 1 + if stream: + if len(ring) >= stream: + ring.popitem(last=False) + ring[expert] = None + decode_steps = len(trace[0]) - prompt_tokens + return { + "decode_hit_rate": 1.0 - decode_misses / decode_accesses, + "warm_hit_rate": warm_hits / decode_accesses, + "decode_misses_per_token": decode_misses / decode_steps, + } + + +def read_volume( + misses_per_token: float, + *, + block_bytes: int, + bandwidths: tuple[float, ...]) -> dict[str, float]: + """Per-token read volume and the tok/s each bandwidth would allow.""" + per_token = misses_per_token * block_bytes + result = {"mb_per_token": per_token / 1e6} + for bandwidth in bandwidths: + result[f"tok_s_at_{bandwidth:g}gbps"] = ( + bandwidth * 1e9 / per_token if per_token else float("inf")) + return result + + +_POLICIES = ("single_lru", "two_tier", "two_tier_oracle_warm") + + +def summarize( + trace: list[list[list[int]]], + *, + prompt_tokens: int, + quotas: tuple[int, ...], + stream_fraction: float, + block_bytes: int, + bandwidths: tuple[float, ...]) -> dict[str, dict]: + """Compare both policies across per-layer quotas.""" + summary: dict[str, dict] = {} + for quota in quotas: + stream = max(1, int(round(quota * stream_fraction))) + pinned = max(0, quota - stream) + entry = { + "quota": quota, + "pinned": pinned, + "stream": stream, + "single_lru": simulate_lru( + trace, prompt_tokens=prompt_tokens, quota=quota), + "two_tier": simulate_two_tier( + trace, prompt_tokens=prompt_tokens, + pinned=pinned, stream=stream), + "two_tier_oracle_warm": simulate_two_tier( + trace, prompt_tokens=prompt_tokens, + pinned=pinned, stream=stream, warm_from="decode"), + } + for policy in _POLICIES: + entry[policy].update(read_volume( + entry[policy]["decode_misses_per_token"], + block_bytes=block_bytes, + bandwidths=bandwidths, + )) + summary[str(quota)] = entry + return summary + + +def format_summary( + summary: dict[str, dict], + *, + quotas: tuple[int, ...], + bandwidths: tuple[float, ...]) -> str: + header = f"{'quota':>6} {'pin/str':>8} {'policy':<22}" + header += f" {'hit':>7} {'miss/tok':>9} {'MB/tok':>8}" + for bandwidth in bandwidths: + header += f" {f'{bandwidth:g}GB/s':>9}" + lines = [header] + for quota in quotas: + entry = summary[str(quota)] + split = f"{entry['pinned']}/{entry['stream']}" + for policy in _POLICIES: + values = entry[policy] + line = f"{quota:>6} {split:>8} {policy:<22}" + line += f" {values['decode_hit_rate']:>7.4f}" + line += f" {values['decode_misses_per_token']:>9.2f}" + line += f" {values['mb_per_token']:>8.1f}" + for bandwidth in bandwidths: + line += f" {values[f'tok_s_at_{bandwidth:g}gbps']:>9.2f}" + lines.append(line) + return "\n".join(lines) + + def main() -> None: parser = argparse.ArgumentParser() parser.add_argument("--checkpoint", required=True) @@ -61,7 +210,14 @@ def main() -> None: parser.add_argument("--prompt-tokens", type=int, default=32) parser.add_argument("--new-tokens", type=int, default=64) parser.add_argument("--max-seq", type=int, default=128) - parser.add_argument("--quotas", default="8,16,24,32,48,64") + parser.add_argument("--quotas", default="8,16,24,27,32,43,64") + parser.add_argument( + "--stream-fraction", type=float, default=0.25, + help="share of each layer's quota held as an evictable LRU ring") + parser.add_argument("--block-bytes", type=int, default=DEFAULT_BLOCK_BYTES) + parser.add_argument( + "--bandwidths", default="1.0,1.5,2.0", + help="storage read bandwidths in GB/s to project tok/s for") parser.add_argument("--device", default="cuda:0") args = parser.parse_args() @@ -97,25 +253,32 @@ def main() -> None: [list(experts) for experts in state.router_trace[layer]] for layer in range(state.num_layers) ] + quotas = tuple(int(value) for value in args.quotas.split(",")) + bandwidths = tuple( + float(value) for value in args.bandwidths.split(",")) + summary = summarize( + trace, + prompt_tokens=args.prompt_tokens, + quotas=quotas, + stream_fraction=args.stream_fraction, + block_bytes=args.block_bytes, + bandwidths=bandwidths, + ) result = { "prompt_tokens": args.prompt_tokens, + "block_bytes": args.block_bytes, + "stream_fraction": args.stream_fraction, "generated_tokens": generated, "trace": trace, - "lru": {}, + "cache": summary, } - for quota in (int(value) for value in args.quotas.split(",")): - result["lru"][str(quota)] = simulate_lru( - trace, prompt_tokens=args.prompt_tokens, quota=quota) args.output.parent.mkdir(parents=True, exist_ok=True) with args.output.open("w", encoding="utf-8") as f: json.dump(result, f) f.write("\n") - for quota, values in result["lru"].items(): - print( - f"quota={quota}: decode_hit={values['decode_hit_rate']:.4f}, " - f"misses/token={values['decode_misses_per_token']:.2f}" - ) + print(format_summary( + summary, quotas=quotas, bandwidths=bandwidths)) if __name__ == "__main__": diff --git a/tests/test_qwen36_moe_edge_quant.py b/tests/test_qwen36_moe_edge_quant.py index 79dbf65d..28621865 100644 --- a/tests/test_qwen36_moe_edge_quant.py +++ b/tests/test_qwen36_moe_edge_quant.py @@ -5,7 +5,11 @@ import torch from qwen36_moe_edge.probe import _dequant_int4 -from qwen36_moe_edge.route_trace import simulate_lru +from qwen36_moe_edge.route_trace import ( + read_volume, + simulate_lru, + simulate_two_tier, +) from qwen36_moe_edge.quantize_experts import ( HIDDEN, INTERMEDIATE, @@ -96,3 +100,53 @@ def test_route_trace_lru_separates_prompt_and_decode(): "decode_hit_rate": 0.25, "decode_misses_per_token": 3.0, } + + +def test_two_tier_warm_set_survives_prefill(): + # The prompt only ever selects expert 0, so it is the warm set. Decode + # reuses it once and touches an unseen expert once. + trace = [[[0], [0], [1], [0]]] + + result = simulate_two_tier( + trace, prompt_tokens=2, pinned=1, stream=0) + + assert result == { + "decode_hit_rate": 0.5, + "warm_hit_rate": 0.5, + "decode_misses_per_token": 0.5, + } + + +def test_two_tier_stream_ring_serves_repeats(): + # No warm set at all: every decode hit has to come from the ring. + trace = [[[0], [1], [1]]] + + result = simulate_two_tier( + trace, prompt_tokens=1, pinned=0, stream=1) + + assert result["warm_hit_rate"] == 0.0 + assert result["decode_hit_rate"] == 0.5 + assert result["decode_misses_per_token"] == 0.5 + + +def test_two_tier_oracle_warm_bounds_the_prompt_heuristic(): + # Decode routes somewhere the prompt never went, so a prompt-derived warm + # set misses everything while a decode-derived one hits everything. + trace = [[[0], [0], [0], [1], [1], [1]]] + + prompt_warm = simulate_two_tier( + trace, prompt_tokens=3, pinned=1, stream=0) + oracle_warm = simulate_two_tier( + trace, prompt_tokens=3, pinned=1, stream=0, warm_from="decode") + + assert prompt_warm["decode_hit_rate"] == 0.0 + assert oracle_warm["decode_hit_rate"] == 1.0 + + +def test_read_volume_converts_misses_to_bandwidth_limits(): + result = read_volume( + 2.0, block_bytes=1_000_000, bandwidths=(1.0, 2.0)) + + assert result["mb_per_token"] == 2.0 + assert result["tok_s_at_1gbps"] == 500.0 + assert result["tok_s_at_2gbps"] == 1000.0 From 7525d8f5eddcfeea6bde2181604a41577c3fb144 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 30 Jul 2026 05:05:41 -0400 Subject: [PATCH 06/85] Align routed-expert blocks for direct I/O The INT8 block payload is 3,151,872 bytes, which is 769.5 times the 4096-byte logical block size, so neither its offset nor its length could be used with O_DIRECT. That matters because the target device's memory holds only a fraction of the experts: streaming them through the page cache would make the cache compete with the resident weights for the same physical memory, so the reader has no choice but to bypass it. Blocks now carry a trailing pad to the next 4096-byte boundary. The INT4 group-16 payload was already a multiple and is unchanged at 1,769,472 bytes; INT8 takes 2048 bytes of pad and becomes 3,153,920. manifest.json records the alignment and the padding entry so a loader can compute component offsets without reproducing the arithmetic. No bundle had been generated against the unpadded layout. --- qwen36_moe_edge/README.md | 15 ++++++++++++--- qwen36_moe_edge/quantize_experts.py | 17 ++++++++++++++++- tests/test_qwen36_moe_edge_quant.py | 24 ++++++++++++++++++++++++ 3 files changed, 52 insertions(+), 4 deletions(-) diff --git a/qwen36_moe_edge/README.md b/qwen36_moe_edge/README.md index 39a5aebd..e97d72ce 100644 --- a/qwen36_moe_edge/README.md +++ b/qwen36_moe_edge/README.md @@ -45,9 +45,18 @@ INT8 uses symmetric per-output-channel FP16 scales. INT4 follows the Thor Pi0.5 numerical contract: sign-magnitude values, one UE4M3 scale per 16 K values, and two values per byte with the low nibble first. `int4-rht` applies the same orthonormal H16/4 transform to every K block that the runtime applies -to activations. Scale bytes in these edge block files are linear; a Thor -loader must convert them to the SM1xx SFB tile-interleaved layout before -calling the native block-scaled MMA kernels. +to activations. Scale bytes in these edge block files are linear; a loader +must convert them to the SM1xx SFB tile-interleaved layout before calling the +native block-scaled MMA kernels. + +Each block carries a trailing pad so its offset and length are multiples of +`BLOCK_ALIGNMENT`. On a device whose memory holds only a fraction of the +experts, the expert stream cannot go through the page cache — it would compete +with the resident weights for the same physical memory — so the reader has to +use `O_DIRECT`, which requires aligned offsets and lengths. The INT4 group-16 +payload is already a multiple of 4096; the INT8 payload is 3,151,872 bytes and +takes 2048 bytes of pad. `manifest.json` records `block_bytes`, +`block_alignment`, and the padding entry in `block_sizes`. An SM120 machine can collect real router selections for cache sizing: diff --git a/qwen36_moe_edge/quantize_experts.py b/qwen36_moe_edge/quantize_experts.py index a8b25c8e..d58603a8 100644 --- a/qwen36_moe_edge/quantize_experts.py +++ b/qwen36_moe_edge/quantize_experts.py @@ -18,6 +18,14 @@ HIDDEN = 2048 INTERMEDIATE = 512 +# Expert blocks are padded so that every block's offset and length are a +# multiple of the logical block size. A device whose memory holds only a +# fraction of the experts cannot afford to stream them through the page +# cache -- the cache competes with the resident weights for the same +# physical memory -- so the reader has to use O_DIRECT, which requires +# aligned offsets and lengths. +BLOCK_ALIGNMENT = 4096 + class CheckpointReader: def __init__(self, checkpoint: Path): @@ -120,6 +128,7 @@ def quantize_expert( _tensor_bytes(gu_scale), _tensor_bytes(dn_weight), _tensor_bytes(dn_scale), + bytes(_layout(quant_format, group_size)["padding"]), )) @@ -142,12 +151,17 @@ def _layout(quant_format: str, group_size: int) -> dict[str, int]: gu_scale = 2 * INTERMEDIATE * (HIDDEN // group_size) dn_weight = HIDDEN * INTERMEDIATE // 2 dn_scale = HIDDEN * (INTERMEDIATE // group_size) - return { + layout = { "gate_up_weight": gu_weight, "gate_up_scale": gu_scale, "down_weight": dn_weight, "down_scale": dn_scale, } + # Trailing pad keeps every expert's offset and length aligned. The INT4 + # group-16 payload happens to be a multiple already; the INT8 payload is + # 3,151,872 B, which is 769.5 blocks. + layout["padding"] = -sum(layout.values()) % BLOCK_ALIGNMENT + return layout def main() -> None: @@ -180,6 +194,7 @@ def main() -> None: "block_layout": list(layout), "block_sizes": layout, "block_bytes": block_bytes, + "block_alignment": BLOCK_ALIGNMENT, } with (args.output / "manifest.json").open("w", encoding="utf-8") as f: json.dump(manifest, f, indent=2) diff --git a/tests/test_qwen36_moe_edge_quant.py b/tests/test_qwen36_moe_edge_quant.py index 28621865..6456b853 100644 --- a/tests/test_qwen36_moe_edge_quant.py +++ b/tests/test_qwen36_moe_edge_quant.py @@ -11,6 +11,7 @@ simulate_two_tier, ) from qwen36_moe_edge.quantize_experts import ( + BLOCK_ALIGNMENT, HIDDEN, INTERMEDIATE, _hadamard16, @@ -87,6 +88,29 @@ def test_expert_block_size_matches_manifest_layout(): _layout(quant_format, group_size).values()) +def test_expert_blocks_are_aligned_for_direct_io(): + # An 8 GiB unified-memory device cannot stream the experts through the + # page cache, so the reader needs O_DIRECT and every block offset and + # length has to be aligned. + for quant_format, group_size in ( + ("int8", 32), + ("int4", 16), + ("int4", 32), + ("int4-rht", 16), + ): + layout = _layout(quant_format, group_size) + block_bytes = sum(layout.values()) + assert block_bytes % BLOCK_ALIGNMENT == 0, ( + quant_format, group_size, block_bytes) + assert layout["padding"] < BLOCK_ALIGNMENT + assert list(layout)[-1] == "padding" + + # INT4 group-16 is aligned on its own; INT8 needs 2048 bytes of pad. + assert _layout("int4-rht", 16)["padding"] == 0 + assert _layout("int8", 32)["padding"] == 2048 + assert sum(_layout("int8", 32).values()) == 3153920 + + def test_route_trace_lru_separates_prompt_and_decode(): trace = [ [[0, 1], [0, 2], [0, 1], [2, 3]], From 2a3c4341d0d92aab42f791ce589af70e12db0efb Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 30 Jul 2026 05:12:29 -0400 Subject: [PATCH 07/85] Decode packed E2M1 without requiring cuda_fp4.h The weight-only 4-bit kernels pulled for one function, __nv_cvt_fp4x2_to_halfraw2. That header first appears in CUDA 12.8, so the kernels could not be built for a Jetson image pinned to 12.6 even though nothing else in them needs a newer toolkit. On an architecture without cvt.rn.f16x2.e2m1x2 the header decodes each nibble in software anyway, by way of an intermediate E2M3 conversion. E2M1 has sixteen representable values, so a table of raw half bit patterns gives the same answer on any target. The new header uses where it exists, so the emitted code on current targets is unchanged, and falls back to the table otherwise. FLASHRT_FP4_FORCE_TABLE selects the table explicitly. Verified on sm_120a, which has the hardware conversion: both paths produce identical half bit patterns for all 256 packed byte values. With the fallback available, all ten core and weight-only-4-bit translation units compile for sm_87 under CUDA 12.6 and gcc 11.4, and for sm_110 under CUDA 13.0. --- csrc/kernels/fp4_e2m1_compat.cuh | 68 +++++++++++++++++++++++++ csrc/kernels/moe_grouped_w4a16_sm120.cu | 7 ++- csrc/kernels/w4a16_gemm_sm120.cu | 10 ++-- csrc/kernels/w4a16_matvec_sm120.cu | 7 ++- 4 files changed, 79 insertions(+), 13 deletions(-) create mode 100644 csrc/kernels/fp4_e2m1_compat.cuh diff --git a/csrc/kernels/fp4_e2m1_compat.cuh b/csrc/kernels/fp4_e2m1_compat.cuh new file mode 100644 index 00000000..b4fe15ef --- /dev/null +++ b/csrc/kernels/fp4_e2m1_compat.cuh @@ -0,0 +1,68 @@ +// Packed-E2M1 to half2 conversion that does not require . +// +// The CUDA header only exists from 12.8 onwards, and on architectures without +// the cvt.rn.f16x2.e2m1x2 instruction it decodes each nibble in software +// anyway. E2M1 has sixteen representable values, so a table gives the same +// result on every target and removes the toolkit dependency: a Jetson image +// pinned to CUDA 12.6 can still build the weight-only 4-bit kernels. +// +// Where the header is present it is used, so the emitted code on those targets +// is unchanged. + +#pragma once + +#include +#include + +// Define FLASHRT_FP4_FORCE_TABLE to take the portable path even where the +// header exists. Used to check the two against each other. +#if !defined(FLASHRT_FP4_FORCE_TABLE) && defined(__has_include) +#if __has_include() +#define FLASHRT_HAVE_CUDA_FP4_HEADER 1 +#endif +#endif + +#ifdef FLASHRT_HAVE_CUDA_FP4_HEADER +#include +#endif + +namespace flash_rt { +namespace fp4 { + +// Decode one byte holding two E2M1 values, low nibble first. +__device__ __forceinline__ __half2_raw cvt_e2m1x2_to_halfraw2(uint8_t pair) { +#ifdef FLASHRT_HAVE_CUDA_FP4_HEADER + return __nv_cvt_fp4x2_to_halfraw2( + static_cast<__nv_fp4x2_storage_t>(pair), __NV_E2M1); +#else + // The sixteen E2M1 values as raw half bit patterns, indexed by the 4-bit + // code: one sign bit, two exponent bits, one mantissa bit, giving 0, + // +/-0.5, +/-1, +/-1.5, +/-2, +/-3, +/-4, +/-6. Function-local so no + // translation unit owns a device symbol. + constexpr unsigned short kAsHalfRaw[16] = { + 0x0000, // 0.0 + 0x3800, // 0.5 + 0x3C00, // 1.0 + 0x3E00, // 1.5 + 0x4000, // 2.0 + 0x4200, // 3.0 + 0x4400, // 4.0 + 0x4600, // 6.0 + 0x8000, // -0.0 + 0xB800, // -0.5 + 0xBC00, // -1.0 + 0xBE00, // -1.5 + 0xC000, // -2.0 + 0xC200, // -3.0 + 0xC400, // -4.0 + 0xC600, // -6.0 + }; + __half2_raw out; + out.x = kAsHalfRaw[pair & 0x0F]; + out.y = kAsHalfRaw[(pair >> 4) & 0x0F]; + return out; +#endif +} + +} // namespace fp4 +} // namespace flash_rt diff --git a/csrc/kernels/moe_grouped_w4a16_sm120.cu b/csrc/kernels/moe_grouped_w4a16_sm120.cu index 8a82e89d..6599b6aa 100644 --- a/csrc/kernels/moe_grouped_w4a16_sm120.cu +++ b/csrc/kernels/moe_grouped_w4a16_sm120.cu @@ -7,7 +7,7 @@ #include #include -#include +#include "kernels/fp4_e2m1_compat.cuh" #include #include #include @@ -34,9 +34,8 @@ __device__ __forceinline__ float blockdot_g(uint64_t b_pack, float acc = 0.0f; #pragma unroll for (int j = 0; j < 8; ++j) { - const __nv_fp4x2_storage_t bb = - static_cast<__nv_fp4x2_storage_t>(b_pack >> (j * 8)); - const __half2_raw wr = __nv_cvt_fp4x2_to_halfraw2(bb, __NV_E2M1); + const __half2_raw wr = flash_rt::fp4::cvt_e2m1x2_to_halfraw2( + static_cast(b_pack >> (j * 8))); const float2 wf = __half22float2(*reinterpret_cast(&wr)); const float2 xf = __bfloat1622float2(xb2[j]); acc = fmaf(wf.x, xf.x, acc); diff --git a/csrc/kernels/w4a16_gemm_sm120.cu b/csrc/kernels/w4a16_gemm_sm120.cu index a069bcce..66b8a79a 100644 --- a/csrc/kernels/w4a16_gemm_sm120.cu +++ b/csrc/kernels/w4a16_gemm_sm120.cu @@ -18,7 +18,7 @@ #include #include -#include +#include "kernels/fp4_e2m1_compat.cuh" #include #include #include @@ -166,10 +166,10 @@ __global__ __launch_bounds__(GM_THREADS) void w4a16_gemm_kernel( const int ncol = warp_n * 32 + jb * 8 + r; const uint8_t* wq = &sWq[cur][ncol * KT_half]; float sf = c_w4a16_ue4m3[sSFB[cur][ncol * GM_KSUB + ksub]] * alpha; - __half2_raw h0 = __nv_cvt_fp4x2_to_halfraw2( - static_cast<__nv_fp4x2_storage_t>(wq[byte0]), __NV_E2M1); - __half2_raw h1 = __nv_cvt_fp4x2_to_halfraw2( - static_cast<__nv_fp4x2_storage_t>(wq[byte0 + 4]), __NV_E2M1); + __half2_raw h0 = flash_rt::fp4::cvt_e2m1x2_to_halfraw2( + static_cast(wq[byte0])); + __half2_raw h1 = flash_rt::fp4::cvt_e2m1x2_to_halfraw2( + static_cast(wq[byte0 + 4])); float2 f0 = __half22float2(*reinterpret_cast(&h0)); float2 f1 = __half22float2(*reinterpret_cast(&h1)); bb0[jb] = pack_bf16x2(f0.x * sf, f0.y * sf); diff --git a/csrc/kernels/w4a16_matvec_sm120.cu b/csrc/kernels/w4a16_matvec_sm120.cu index ddc2f89b..32c386db 100644 --- a/csrc/kernels/w4a16_matvec_sm120.cu +++ b/csrc/kernels/w4a16_matvec_sm120.cu @@ -6,7 +6,7 @@ #include #include -#include +#include "kernels/fp4_e2m1_compat.cuh" #include #include #include @@ -38,9 +38,8 @@ __device__ __forceinline__ float blockdot(uint64_t b_pack, float acc = 0.0f; #pragma unroll for (int j = 0; j < 8; ++j) { - const __nv_fp4x2_storage_t bb = - static_cast<__nv_fp4x2_storage_t>(b_pack >> (j * 8)); - const __half2_raw wr = __nv_cvt_fp4x2_to_halfraw2(bb, __NV_E2M1); + const __half2_raw wr = flash_rt::fp4::cvt_e2m1x2_to_halfraw2( + static_cast(b_pack >> (j * 8))); const float2 wf = __half22float2(*reinterpret_cast(&wr)); const float2 xf = __bfloat1622float2(xb2[j]); acc = fmaf(wf.x, xf.x, acc); From bab295a227e0977a2d20ea9216742bf1dc960792 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 30 Jul 2026 05:14:41 -0400 Subject: [PATCH 08/85] Record which expert-cache policy actually wins The two-tier docstring implied that pinning a prompt-derived warm set is the correction to a plain LRU. Measurement on Qwen3.6-35B-A3B says otherwise: with a per-layer quota already in place the plain LRU wins from 16 slots up, and its margin grows with prompt length rather than shrinking -- 0.745 against 0.731 at 43 slots for a 32-token prompt, 0.711 against 0.664 for a 128-token prompt. Recency predicts this router's next selections better than prompt-phase frequency, and a longer prompt spreads the frequency estimate over more experts instead of sharpening it. The oracle warm set stays ahead of both, so the weakness is the prompt-derived choice of what to pin, not pinning. Both files now state the measured result and say to measure per checkpoint instead of assuming a policy. --- qwen36_moe_edge/README.md | 14 ++++++++++++++ qwen36_moe_edge/route_trace.py | 20 ++++++++++++++++---- 2 files changed, 30 insertions(+), 4 deletions(-) diff --git a/qwen36_moe_edge/README.md b/qwen36_moe_edge/README.md index e97d72ce..d91b9977 100644 --- a/qwen36_moe_edge/README.md +++ b/qwen36_moe_edge/README.md @@ -85,6 +85,20 @@ Each quota is scored under three policies: decode phase. Not implementable; it bounds what a better warm-set heuristic could add. +On Qwen3.6-35B-A3B the plain LRU wins from 16 slots per layer up, and its +margin grows with prompt length — at 43 slots per layer, 0.745 against 0.731 +for a 32-token prompt and 0.711 against 0.664 for a 128-token prompt. Once a +per-layer quota exists, recency predicts this router's next selections better +than prompt-phase frequency, and a longer prompt makes the frequency estimate +more diffuse rather than more reliable. The oracle variant stays ahead of both +(0.776 at 43 slots for the 128-token prompt), so pinning is sound and the +prompt-derived choice of what to pin is what falls short. Treat the policy as +something to measure per checkpoint, not to assume. + +Capacity dominates policy either way: going from 43 to 64 slots per layer cuts +read volume by a third, while any policy change at a fixed quota moves it by a +few percent. + `--block-bytes` (default: the INT4 group-16 block) and `--bandwidths` turn misses per token into a read volume and the token rate each storage bandwidth would allow, which is the number that decides whether a memory budget is diff --git a/qwen36_moe_edge/route_trace.py b/qwen36_moe_edge/route_trace.py index 87516e07..c0ce32ac 100644 --- a/qwen36_moe_edge/route_trace.py +++ b/qwen36_moe_edge/route_trace.py @@ -9,10 +9,22 @@ whatever the end of the prompt happened to use. ``simulate_two_tier`` - A per-layer warm set chosen from the prompt and never evicted, plus a - small LRU ring for everything else. Prefill can no longer displace the - warm set, so the decode hit rate follows warm-set coverage instead of - prompt churn. + A per-layer warm set chosen from prompt-phase selection counts and never + evicted, plus a small LRU ring for everything else. Prefill cannot + displace the warm set, so the decode hit rate follows warm-set coverage + rather than what the end of the prompt happened to leave behind. + +Which is better is a property of the checkpoint, not a foregone conclusion. +On Qwen3.6-35B-A3B, with a per-layer quota already in place, the plain LRU +wins at every quota from 16 slots up, and its margin *grows* with prompt +length: at 43 slots per layer it reaches 0.745 against 0.731 for a 32-token +prompt and 0.711 against 0.664 for a 128-token prompt. Recency predicts this +router's next selections better than prompt-phase frequency does, and a longer +prompt makes the frequency estimate more diffuse rather than more reliable. + +The oracle variant is consistently best, so pinning itself is not the problem +— the prompt-derived choice of what to pin is. Measure before committing a +runtime to either policy. The reported miss count per token, multiplied by the expert block size, is the per-token read volume a streaming runtime has to sustain. From ed687c1ac77f0ec0964cd7e56a47ef6274bdb493 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 30 Jul 2026 05:19:09 -0400 Subject: [PATCH 09/85] Score expert quantization against real routed activations The existing quality probe samples experts with torch.randn activations and quantizes the activations as well. Neither matches how the edge runtime uses these weights. At M=1 the activation is 4 KiB against a 1.7 MiB weight block, so quantizing it buys no bandwidth and the expert path is weight-only. And a scale calibrated against Gaussian noise is not the scale real post-norm hidden states need, so the probe can pass while the deployed path does not. expert_quality.py captures the activations the router actually sent to each expert, replays the expert in BF16 for the reference, and scores W8A16, W4A16 and W4A16 with the block-16 transform against it. It can also write a small bundle of activation and reference pairs so a device can check its own dequantization and kernel without loading the source checkpoint. The MoE input trace is a second opt-in hook alongside the router trace, off by default and eager-only for the same reason. dequantize_int4 moves from the probe to quantize_experts, next to the packing it inverts, so the format has one definition. The transform's effect is now asserted where it is supposed to appear: on weights with one outlier per group of 16 it cuts relative L2 by 38 %, while on Gaussian weights it is neutral, which is what an orthonormal rotation of an already-Gaussian group should do. --- flash_rt/frontends/torch/_nexn2_rtx_decode.py | 9 +- qwen36_moe_edge/README.md | 21 ++ qwen36_moe_edge/expert_quality.py | 283 ++++++++++++++++++ qwen36_moe_edge/probe.py | 30 +- qwen36_moe_edge/quantize_experts.py | 22 ++ tests/test_qwen36_moe_edge_quant.py | 69 ++++- 6 files changed, 405 insertions(+), 29 deletions(-) create mode 100644 qwen36_moe_edge/expert_quality.py diff --git a/flash_rt/frontends/torch/_nexn2_rtx_decode.py b/flash_rt/frontends/torch/_nexn2_rtx_decode.py index 2704a8e1..a3e14365 100644 --- a/flash_rt/frontends/torch/_nexn2_rtx_decode.py +++ b/flash_rt/frontends/torch/_nexn2_rtx_decode.py @@ -273,9 +273,11 @@ def __init__(self, handles, max_seq, device): # chunk (64). 0 disables (always single-pass). self.prefill_chunk = int( _qwen35moe_env("PREFILL_CHUNK", "8192")) - # Optional eager-only routing trace used to size edge expert caches. - # Keep this disabled during CUDA Graph capture. + # Optional eager-only traces used to size edge expert caches and to + # score expert quantization against real activations. Keep these + # disabled during CUDA Graph capture. self.router_trace = None + self.moe_input_trace = None self._active_layer = -1 def reset(self): @@ -450,6 +452,9 @@ def _moe_layer_decode(h, ld, state, fvk, device): if state.router_trace is not None: state.router_trace[state._active_layer].append( tuple(int(v) for v in idx.cpu().tolist())) + if state.moe_input_trace is not None: + state.moe_input_trace[state._active_layer].append( + x.detach().to("cpu", copy=True)) if 'experts_gate_up_alpha_dev' not in ld: # cache once/layer ld['experts_gate_up_alpha_dev'] = \ diff --git a/qwen36_moe_edge/README.md b/qwen36_moe_edge/README.md index d91b9977..110fd779 100644 --- a/qwen36_moe_edge/README.md +++ b/qwen36_moe_edge/README.md @@ -24,6 +24,27 @@ PYTHONPATH=. python qwen36_moe_edge/probe.py \ --group-size 16 ``` +Score the quantization schemes against the activations the router actually +sends each expert, and optionally save the references a device can check +itself against: + +```bash +PYTHONPATH=. python qwen36_moe_edge/expert_quality.py \ + --checkpoint /models/Qwen3.6-35B-A3B \ + --prompt "Explain edge mixture-of-experts inference. " \ + --prompt-tokens 32 \ + --new-tokens 32 \ + --output qwen36_expert_quality.json \ + --golden qwen36_expert_golden.safetensors +``` + +Prefer this over `probe.py --mode quality` when deciding what to generate. +The probe uses `torch.randn` activations and quantizes the activations too; +neither matches the runtime, where the activation is 4 KiB against a 1.7 MiB +weight block and so is left in BF16. Random activations also hide errors that +real inputs expose, because a scale calibrated against noise is not the scale +real inputs need. + Generate fixed-size routed-expert blocks for a layer range: ```bash diff --git a/qwen36_moe_edge/expert_quality.py b/qwen36_moe_edge/expert_quality.py new file mode 100644 index 00000000..bb86d047 --- /dev/null +++ b/qwen36_moe_edge/expert_quality.py @@ -0,0 +1,283 @@ +#!/usr/bin/env python3 +"""Score routed-expert quantization against real routed activations. + +``probe.py --mode quality`` samples experts with ``torch.randn`` activations +and quantizes the activations too. Neither matches how the edge runtime will +use these weights: + +- The activations an expert actually sees are the post-norm hidden states of + tokens the router sent to *that* expert. Gaussian noise has none of their + structure, and a scale calibrated against it hides errors that real inputs + expose. +- At M=1 the activation is 4 KiB against a 1.7 MiB weight block, so quantizing + it buys no bandwidth. The expert path is weight-only, W4A16 or W8A16. + +This tool captures the real activations from a forward pass, replays each +sampled expert in BF16 for the reference, and scores the weight-only +reconstructions against it. It can also write a small bundle so a device can +check its own dequantization and kernel against the same references without +loading the source checkpoint. +""" + +from __future__ import annotations + +import argparse +import json +import statistics +from pathlib import Path + +import torch +import torch.nn.functional as F + +from qwen36_moe_edge.quantize_experts import ( + HIDDEN, + INTERMEDIATE, + NUM_LAYERS, + CheckpointReader, + _int4_weight, + _int8_weight, + _rht16, + dequantize_int4, +) + + +SCHEMES = ("w8a16", "w4a16", "w4a16_rht16") + + +def expert_forward( + activation: torch.Tensor, + gate_up: torch.Tensor, + down: torch.Tensor) -> torch.Tensor: + """One routed expert: gate_up, SwiGLU, down.""" + projected = activation @ gate_up.T + hidden = ( + F.silu(projected[:, :INTERMEDIATE]) + * projected[:, INTERMEDIATE:] + ) + return hidden @ down.T + + +def _reconstruct( + weight: torch.Tensor, + *, + scheme: str, + group_size: int) -> torch.Tensor: + """Quantize a weight and dequantize it, as the runtime's kernel will.""" + columns = weight.shape[1] + if scheme == "w8a16": + quantized, scale = _int8_weight(weight) + return quantized.float() * scale.float()[:, None] + source = _rht16(weight) if scheme == "w4a16_rht16" else weight + packed, scale = _int4_weight(source, group_size) + return dequantize_int4(packed, scale, columns, group_size) + + +def score_expert( + activation: torch.Tensor, + gate_up: torch.Tensor, + down: torch.Tensor, + *, + scheme: str, + group_size: int) -> dict[str, float]: + """Cosine and relative L2 of a weight-only scheme against BF16.""" + reference = expert_forward(activation, gate_up, down) + + gate_up_q = _reconstruct( + gate_up, scheme=scheme, group_size=group_size) + down_q = _reconstruct(down, scheme=scheme, group_size=group_size) + if scheme == "w4a16_rht16": + # The transform is orthonormal, so rotating both sides of each dot + # product leaves it unchanged. The runtime rotates activations the + # same way before calling the kernel. + projected = _rht16(activation) @ gate_up_q.T + hidden = ( + F.silu(projected[:, :INTERMEDIATE]) + * projected[:, INTERMEDIATE:] + ) + output = _rht16(hidden) @ down_q.T + else: + output = expert_forward(activation, gate_up_q, down_q) + + difference = (output - reference).flatten() + return { + "cosine": F.cosine_similarity( + reference.flatten(), output.flatten(), dim=0).item(), + "relative_l2": ( + difference.norm() / reference.flatten().norm().clamp_min(1e-12) + ).item(), + } + + +def collect_activations( + checkpoint: str, + *, + prompt: str, + prompt_tokens: int, + new_tokens: int, + max_seq: int, + device: str) -> tuple[list[list[list[int]]], list[list[torch.Tensor]]]: + """Run one eager forward, returning per-layer selections and MoE inputs.""" + from flash_rt.frontends.torch._nexn2_rtx_decode import ( + Nexn2DecodeState, + generate_greedy, + ) + from flash_rt.frontends.torch.qwen36_moe_rtx import ( + Qwen36MoeTextFrontendRtx, + ) + + frontend = Qwen36MoeTextFrontendRtx( + checkpoint, device=device, max_seq=max_seq, quant_scope="experts") + input_ids = frontend.tokenizer( + prompt, return_tensors="pt", add_special_tokens=False, + ).input_ids[:, :prompt_tokens].to(device) + if input_ids.shape[1] != prompt_tokens: + raise ValueError("the supplied prompt is shorter than prompt_tokens") + + state = Nexn2DecodeState(frontend._weights, max_seq, device) + state.batched_prefill = False + state.router_trace = {layer: [] for layer in range(state.num_layers)} + state.moe_input_trace = {layer: [] for layer in range(state.num_layers)} + with torch.no_grad(): + generate_greedy( + state, input_ids, new_tokens, frontend._fvk, device) + + selections = [ + [list(experts) for experts in state.router_trace[layer]] + for layer in range(state.num_layers) + ] + activations = [ + list(state.moe_input_trace[layer]) + for layer in range(state.num_layers) + ] + return selections, activations + + +def _sampled_pairs( + selections: list[list[list[int]]], + *, + layers: tuple[int, ...], + experts_per_layer: int) -> list[tuple[int, int, int]]: + """Pick (layer, expert, step) triples the router actually produced.""" + pairs = [] + for layer in layers: + seen: dict[int, int] = {} + for step, experts in enumerate(selections[layer]): + for expert in experts: + seen.setdefault(expert, step) + for expert, step in list(seen.items())[:experts_per_layer]: + pairs.append((layer, expert, step)) + return pairs + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--checkpoint", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--golden", type=Path) + parser.add_argument("--prompt", required=True) + parser.add_argument("--prompt-tokens", type=int, default=32) + parser.add_argument("--new-tokens", type=int, default=32) + parser.add_argument("--max-seq", type=int, default=256) + parser.add_argument( + "--layers", default="0,1,3,19,20,39", + help="layers to sample; the default spans both attention kinds and " + "the first, middle and last MoE blocks") + parser.add_argument("--experts-per-layer", type=int, default=4) + parser.add_argument("--group-size", type=int, default=16) + parser.add_argument("--device", default="cuda:0") + args = parser.parse_args() + + layers = tuple(int(value) for value in args.layers.split(",")) + for layer in layers: + if not 0 <= layer < NUM_LAYERS: + parser.error(f"layer {layer} is outside 0..{NUM_LAYERS - 1}") + + selections, activations = collect_activations( + str(args.checkpoint), + prompt=args.prompt, + prompt_tokens=args.prompt_tokens, + new_tokens=args.new_tokens, + max_seq=args.max_seq, + device=args.device, + ) + pairs = _sampled_pairs( + selections, layers=layers, experts_per_layer=args.experts_per_layer) + print(f"scoring {len(pairs)} routed (layer, expert) pairs", flush=True) + + reader = CheckpointReader(args.checkpoint) + scores: dict[str, list[float]] = { + f"{scheme}.{metric}": [] + for scheme in SCHEMES for metric in ("cosine", "relative_l2") + } + records = [] + golden: dict[str, torch.Tensor] = {} + for layer, expert, step in pairs: + activation = activations[layer][step].to( + device=args.device, dtype=torch.float32) + gate_up = reader.expert(layer, "gate_up_proj", expert).to( + device=args.device, dtype=torch.float32) + down = reader.expert(layer, "down_proj", expert).to( + device=args.device, dtype=torch.float32) + + record = {"layer": layer, "expert": expert, "step": step} + for scheme in SCHEMES: + values = score_expert( + activation, gate_up, down, + scheme=scheme, group_size=args.group_size) + record[scheme] = values + for metric, value in values.items(): + scores[f"{scheme}.{metric}"].append(value) + records.append(record) + + if args.golden is not None: + key = f"layer{layer:02d}.expert{expert:03d}" + golden[f"{key}.activation"] = ( + activation.to(torch.bfloat16).cpu()) + golden[f"{key}.reference"] = expert_forward( + activation, gate_up, down).to(torch.bfloat16).cpu() + + summary = { + name: { + "min": min(values), + "mean": statistics.mean(values), + "max": max(values), + } + for name, values in scores.items() + } + result = { + "prompt_tokens": args.prompt_tokens, + "new_tokens": args.new_tokens, + "group_size": args.group_size, + "layers": list(layers), + "pair_count": len(pairs), + "summary": summary, + "records": records, + } + args.output.parent.mkdir(parents=True, exist_ok=True) + with args.output.open("w", encoding="utf-8") as f: + json.dump(result, f, indent=2) + f.write("\n") + + if args.golden is not None: + from safetensors.torch import save_file + + args.golden.parent.mkdir(parents=True, exist_ok=True) + save_file(golden, str(args.golden), metadata={ + "hidden_size": str(HIDDEN), + "intermediate_size": str(INTERMEDIATE), + "pairs": ",".join( + f"{layer}:{expert}" for layer, expert, _ in pairs), + }) + print(f"wrote {len(golden) // 2} reference pairs to {args.golden}") + + print(f"\n{'scheme':<14} {'cos min':>10} {'cos mean':>10} " + f"{'relL2 max':>10} {'relL2 mean':>11}") + for scheme in SCHEMES: + cosine = summary[f"{scheme}.cosine"] + l2 = summary[f"{scheme}.relative_l2"] + print(f"{scheme:<14} {cosine['min']:>10.6f} {cosine['mean']:>10.6f} " + f"{l2['max']:>10.5f} {l2['mean']:>11.5f}") + + +if __name__ == "__main__": + main() diff --git a/qwen36_moe_edge/probe.py b/qwen36_moe_edge/probe.py index 3a3850f3..779e2a03 100644 --- a/qwen36_moe_edge/probe.py +++ b/qwen36_moe_edge/probe.py @@ -22,6 +22,7 @@ _int4_weight, _int8_weight, _layout, + dequantize_int4, ) @@ -155,27 +156,6 @@ def memory_probe( ) -def _dequant_int4( - packed: torch.Tensor, - scale: torch.Tensor, - columns: int, - group_size: int, -) -> torch.Tensor: - low = packed & 0x0F - high = (packed >> 4) & 0x0F - low = (low & 0x07).to(torch.int8) * torch.where( - (low & 0x08) != 0, -1, 1).to(torch.int8) - high = (high & 0x07).to(torch.int8) * torch.where( - (high & 0x08) != 0, -1, 1).to(torch.int8) - values = torch.stack((low, high), dim=-1).flatten(1) - rows = values.shape[0] - scale_float = scale.view(torch.float8_e4m3fn).float() - return ( - values.float().reshape(rows, columns // group_size, group_size) - * scale_float.unsqueeze(-1) - ).reshape(rows, columns) - - def quality_probe( checkpoint: Path, *, @@ -256,11 +236,11 @@ def quality_probe( ).reshape_as(activation) gu4, gu4_scale = _int4_weight( gu_source, current_group) - gu4 = _dequant_int4( + gu4 = dequantize_int4( gu4, gu4_scale, HIDDEN, current_group) current4, current4_scale = _int4_weight( current, current_group) - current = _dequant_int4( + current = dequantize_int4( current4, current4_scale, HIDDEN, current_group) projected = current @ gu4.T current = ( @@ -277,11 +257,11 @@ def quality_probe( ).reshape_as(current) dn4, dn4_scale = _int4_weight( dn_source, current_group) - dn4 = _dequant_int4( + dn4 = dequantize_int4( dn4, dn4_scale, INTERMEDIATE, current_group) current4, current4_scale = _int4_weight( current, current_group) - current = _dequant_int4( + current = dequantize_int4( current4, current4_scale, INTERMEDIATE, current_group) output = current @ dn4.T scores[mode].append(F.cosine_similarity( diff --git a/qwen36_moe_edge/quantize_experts.py b/qwen36_moe_edge/quantize_experts.py index d58603a8..5c3760ed 100644 --- a/qwen36_moe_edge/quantize_experts.py +++ b/qwen36_moe_edge/quantize_experts.py @@ -96,6 +96,28 @@ def _int4_weight( return packed.contiguous(), scale.view(torch.uint8).contiguous() +def dequantize_int4( + packed: torch.Tensor, + scale: torch.Tensor, + columns: int, + group_size: int, +) -> torch.Tensor: + """Inverse of :func:`_int4_weight`, for scoring and reference paths.""" + low = packed & 0x0F + high = (packed >> 4) & 0x0F + low = (low & 0x07).to(torch.int8) * torch.where( + (low & 0x08) != 0, -1, 1).to(torch.int8) + high = (high & 0x07).to(torch.int8) * torch.where( + (high & 0x08) != 0, -1, 1).to(torch.int8) + values = torch.stack((low, high), dim=-1).flatten(1) + rows = values.shape[0] + scale_float = scale.view(torch.float8_e4m3fn).float() + return ( + values.float().reshape(rows, columns // group_size, group_size) + * scale_float.unsqueeze(-1) + ).reshape(rows, columns) + + def _tensor_bytes(tensor: torch.Tensor) -> bytes: return tensor.contiguous().cpu().numpy().tobytes() diff --git a/tests/test_qwen36_moe_edge_quant.py b/tests/test_qwen36_moe_edge_quant.py index 6456b853..b5209435 100644 --- a/tests/test_qwen36_moe_edge_quant.py +++ b/tests/test_qwen36_moe_edge_quant.py @@ -4,7 +4,11 @@ import torch -from qwen36_moe_edge.probe import _dequant_int4 +from qwen36_moe_edge.expert_quality import ( + SCHEMES, + expert_forward, + score_expert, +) from qwen36_moe_edge.route_trace import ( read_volume, simulate_lru, @@ -19,6 +23,7 @@ _int8_weight, _layout, _rht16, + dequantize_int4, quantize_expert, ) @@ -40,7 +45,7 @@ def test_int4_grouped_round_trip_matches_packed_layout(): weight = torch.randn(64, 256, generator=generator) packed, scale = _int4_weight(weight, 32) - restored = _dequant_int4(packed, scale, 256, 32) + restored = dequantize_int4(packed, scale, 256, 32) assert packed.shape == (64, 128) assert scale.shape == (64, 8) @@ -174,3 +179,63 @@ def test_read_volume_converts_misses_to_bandwidth_limits(): assert result["mb_per_token"] == 2.0 assert result["tok_s_at_1gbps"] == 500.0 assert result["tok_s_at_2gbps"] == 1000.0 + + +def test_expert_forward_is_a_swiglu_over_the_gate_up_split(): + generator = torch.Generator().manual_seed(13) + activation = torch.randn(2, HIDDEN, generator=generator) + gate_up = torch.randn( + 2 * INTERMEDIATE, HIDDEN, generator=generator) * 0.02 + down = torch.randn(HIDDEN, INTERMEDIATE, generator=generator) * 0.02 + + projected = activation @ gate_up.T + expected = ( + torch.nn.functional.silu(projected[:, :INTERMEDIATE]) + * projected[:, INTERMEDIATE:] + ) @ down.T + + torch.testing.assert_close( + expert_forward(activation, gate_up, down), expected) + + +def _outlier_weight(rows, columns, generator): + """One large value in every group of 16 -- what the transform targets.""" + weight = torch.randn(rows, columns, generator=generator) * 0.02 + weight = weight.reshape(rows, columns // 16, 16) + weight[:, :, 0] += torch.randn( + rows, columns // 16, generator=generator).abs() + return weight.reshape(rows, columns) + + +def test_rht16_reduces_int4_error_on_outlier_heavy_weights(): + generator = torch.Generator().manual_seed(11) + activation = torch.randn(4, HIDDEN, generator=generator) + gate_up = _outlier_weight(2 * INTERMEDIATE, HIDDEN, generator) + down = _outlier_weight(HIDDEN, INTERMEDIATE, generator) + + plain = score_expert( + activation, gate_up, down, scheme="w4a16", group_size=16) + rotated = score_expert( + activation, gate_up, down, scheme="w4a16_rht16", group_size=16) + + # The transform is only worth its cost when groups have outliers; a 10 % + # margin keeps this from asserting on noise. + assert rotated["relative_l2"] < 0.9 * plain["relative_l2"] + assert rotated["cosine"] > plain["cosine"] + + +def test_weight_only_schemes_order_by_bit_width(): + generator = torch.Generator().manual_seed(17) + activation = torch.randn(4, HIDDEN, generator=generator) + gate_up = _outlier_weight(2 * INTERMEDIATE, HIDDEN, generator) + down = _outlier_weight(HIDDEN, INTERMEDIATE, generator) + + scores = { + scheme: score_expert( + activation, gate_up, down, scheme=scheme, group_size=16) + for scheme in SCHEMES + } + + assert scores["w8a16"]["relative_l2"] < scores["w4a16"]["relative_l2"] + for values in scores.values(): + assert 0.0 < values["cosine"] <= 1.0 From 9447bd46289cd1db55dc530e4a5cbfdd7a0f74ee Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 30 Jul 2026 05:25:54 -0400 Subject: [PATCH 10/85] Give the INT4 expert block a two-level scale Scoring against real routed activations showed the single-level scale was unusable. Real expert weights have a per-group amax near 0.02, so amax/7 is about 0.003 -- below e4m3's smallest normal value of 2**-6. Every per-group scale in this checkpoint landed in e4m3's subnormal range, where the format keeps roughly three bits: the stored scale carried 18 % mean relative error, and in one sampled layer a quarter of the groups rounded to zero outright. That error multiplies every weight in the group, so it swamped the 4-bit value grid completely. The per-group byte is now a fraction of a per-tensor global scale, which is what the shipped NVFP4 expert path already does with its GEMM alpha. Stored bytes land in e4m3's normal range, and the round-trip error on Gaussian weights drops to the uniform-quantization-noise floor of the value grid itself, 8.6 % measured against 8.6 % predicted. On the experts that carry signal, output error roughly halves: 0.130 to 0.068 and 0.124 to 0.081. The global scales go in a per-layer sidecar rather than inside the blocks, so a block stays exactly block_bytes and 4096-aligned. They are eight bytes per expert and belong with the resident weights, where the kernel reads them as its alpha. Format tag moves to v2; no bundle had been generated. Two measurement corrections came with it: - E2M1 is now scored as a control. It is the format the shipped runtime uses for these experts with exact greedy reproduction, so it is the bar. Without it, per-expert numbers have nothing to be judged against. - Results pool by error energy over reference energy. Per-expert relative L2 is misleading here because expert output norms span three orders of magnitude, and the alarming values all came from experts with near-zero output, where the control scores just as badly and the router weights the contribution down to nothing anyway. The transform's measured benefit fell from 38 % to 9 % once the scale was fixed, because most of it had been compensating for scale error. --- qwen36_moe_edge/expert_quality.py | 67 +++++++++++++++++++++++---- qwen36_moe_edge/probe.py | 18 ++++---- qwen36_moe_edge/quantize_experts.py | 71 +++++++++++++++++++++++------ tests/test_qwen36_moe_edge_quant.py | 42 ++++++++++++++--- 4 files changed, 161 insertions(+), 37 deletions(-) diff --git a/qwen36_moe_edge/expert_quality.py b/qwen36_moe_edge/expert_quality.py index bb86d047..4dc8862d 100644 --- a/qwen36_moe_edge/expert_quality.py +++ b/qwen36_moe_edge/expert_quality.py @@ -30,6 +30,7 @@ import torch.nn.functional as F from qwen36_moe_edge.quantize_experts import ( + E4M3_MAX, HIDDEN, INTERMEDIATE, NUM_LAYERS, @@ -41,7 +42,14 @@ ) -SCHEMES = ("w8a16", "w4a16", "w4a16_rht16") +# nvfp4_e2m1 is the control: the shipped SM120 runtime uses that format for +# these same experts and reproduces greedy tokens exactly, so it is the bar an +# alternative 4-bit format has to match. Scoring a format without it invites +# reading a metric artefact as a defect. +SCHEMES = ("w8a16", "w4a16", "w4a16_rht16", "nvfp4_e2m1") + +# The sixteen E2M1 magnitudes, for nearest-value rounding. +_E2M1_MAGNITUDES = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0) def expert_forward( @@ -57,6 +65,26 @@ def expert_forward( return hidden @ down.T +def _reconstruct_e2m1( + weight: torch.Tensor, group_size: int) -> torch.Tensor: + """Two-level E2M1, matching what the shipped NVFP4 expert path stores.""" + rows, columns = weight.shape + grouped = weight.float().reshape(rows, columns // group_size, group_size) + amax = grouped.abs().amax(dim=2).clamp_min(1e-12) + peak = max(_E2M1_MAGNITUDES) + global_scale = max(float(amax.max()) / (E4M3_MAX * peak), 1e-12) + scale = (amax / peak / global_scale).to(torch.float8_e4m3fn) + effective = (scale.float() * global_scale).unsqueeze(-1) + normalized = grouped / effective.clamp_min(1e-30) + codebook = torch.tensor( + _E2M1_MAGNITUDES, dtype=torch.float32, device=weight.device) + nearest = codebook[ + (normalized.abs().unsqueeze(-1) - codebook).abs().argmin(-1)] + return ( + torch.sign(normalized) * nearest * effective + ).reshape(rows, columns) + + def _reconstruct( weight: torch.Tensor, *, @@ -67,9 +95,12 @@ def _reconstruct( if scheme == "w8a16": quantized, scale = _int8_weight(weight) return quantized.float() * scale.float()[:, None] + if scheme == "nvfp4_e2m1": + return _reconstruct_e2m1(weight, group_size) source = _rht16(weight) if scheme == "w4a16_rht16" else weight - packed, scale = _int4_weight(source, group_size) - return dequantize_int4(packed, scale, columns, group_size) + packed, scale, global_scale = _int4_weight(source, group_size) + return dequantize_int4( + packed, scale, columns, group_size, global_scale) def score_expert( @@ -105,6 +136,14 @@ def score_expert( "relative_l2": ( difference.norm() / reference.flatten().norm().clamp_min(1e-12) ).item(), + # Absolute error and reference magnitude, so results can be pooled + # across experts. Per-expert relative L2 alone is misleading here: + # output norms span three orders of magnitude across experts, and the + # router weights the small ones down before summing, so an expert with + # a near-zero output shows a huge relative error that contributes + # almost nothing to the layer. + "absolute_l2": difference.norm().item(), + "reference_l2": reference.flatten().norm().item(), } @@ -205,9 +244,10 @@ def main() -> None: print(f"scoring {len(pairs)} routed (layer, expert) pairs", flush=True) reader = CheckpointReader(args.checkpoint) + metrics = ("cosine", "relative_l2", "absolute_l2", "reference_l2") scores: dict[str, list[float]] = { f"{scheme}.{metric}": [] - for scheme in SCHEMES for metric in ("cosine", "relative_l2") + for scheme in SCHEMES for metric in metrics } records = [] golden: dict[str, torch.Tensor] = {} @@ -244,6 +284,16 @@ def main() -> None: } for name, values in scores.items() } + # Pooled relative L2: total error energy over total reference energy. This + # weights each expert by how much signal it carries, which is what the + # router does downstream, so it is the figure to judge a format on. + for scheme in SCHEMES: + error = sum( + value ** 2 for value in scores[f"{scheme}.absolute_l2"]) + signal = sum( + value ** 2 for value in scores[f"{scheme}.reference_l2"]) + summary[f"{scheme}.pooled_relative_l2"] = ( + error ** 0.5 / max(signal ** 0.5, 1e-12)) result = { "prompt_tokens": args.prompt_tokens, "new_tokens": args.new_tokens, @@ -270,13 +320,14 @@ def main() -> None: }) print(f"wrote {len(golden) // 2} reference pairs to {args.golden}") - print(f"\n{'scheme':<14} {'cos min':>10} {'cos mean':>10} " - f"{'relL2 max':>10} {'relL2 mean':>11}") + print(f"\n{'scheme':<14} {'pooled relL2':>13} {'cos mean':>10} " + f"{'cos min':>10} {'per-expert relL2 mean':>22}") for scheme in SCHEMES: cosine = summary[f"{scheme}.cosine"] l2 = summary[f"{scheme}.relative_l2"] - print(f"{scheme:<14} {cosine['min']:>10.6f} {cosine['mean']:>10.6f} " - f"{l2['max']:>10.5f} {l2['mean']:>11.5f}") + pooled = summary[f"{scheme}.pooled_relative_l2"] + print(f"{scheme:<14} {pooled:>13.5f} {cosine['mean']:>10.6f} " + f"{cosine['min']:>10.6f} {l2['mean']:>22.5f}") if __name__ == "__main__": diff --git a/qwen36_moe_edge/probe.py b/qwen36_moe_edge/probe.py index 779e2a03..660b09d4 100644 --- a/qwen36_moe_edge/probe.py +++ b/qwen36_moe_edge/probe.py @@ -234,14 +234,15 @@ def quality_probe( current = ( activation.reshape(-1, 16) @ transform ).reshape_as(activation) - gu4, gu4_scale = _int4_weight( + gu4, gu4_scale, gu4_alpha = _int4_weight( gu_source, current_group) gu4 = dequantize_int4( - gu4, gu4_scale, HIDDEN, current_group) - current4, current4_scale = _int4_weight( + gu4, gu4_scale, HIDDEN, current_group, gu4_alpha) + current4, current4_scale, current4_alpha = _int4_weight( current, current_group) current = dequantize_int4( - current4, current4_scale, HIDDEN, current_group) + current4, current4_scale, HIDDEN, current_group, + current4_alpha) projected = current @ gu4.T current = ( F.silu(projected[:, :INTERMEDIATE]) @@ -255,14 +256,15 @@ def quality_probe( current = ( current.reshape(-1, 16) @ transform ).reshape_as(current) - dn4, dn4_scale = _int4_weight( + dn4, dn4_scale, dn4_alpha = _int4_weight( dn_source, current_group) dn4 = dequantize_int4( - dn4, dn4_scale, INTERMEDIATE, current_group) - current4, current4_scale = _int4_weight( + dn4, dn4_scale, INTERMEDIATE, current_group, dn4_alpha) + current4, current4_scale, current4_alpha = _int4_weight( current, current_group) current = dequantize_int4( - current4, current4_scale, INTERMEDIATE, current_group) + current4, current4_scale, INTERMEDIATE, current_group, + current4_alpha) output = current @ dn4.T scores[mode].append(F.cosine_similarity( reference.flatten(), output.flatten(), dim=0).item()) diff --git a/qwen36_moe_edge/quantize_experts.py b/qwen36_moe_edge/quantize_experts.py index 5c3760ed..c70e8de4 100644 --- a/qwen36_moe_edge/quantize_experts.py +++ b/qwen36_moe_edge/quantize_experts.py @@ -26,6 +26,13 @@ # aligned offsets and lengths. BLOCK_ALIGNMENT = 4096 +# Largest magnitude an e4m3 scale byte can hold. The per-group scale is +# expressed as a fraction of a per-tensor global scale so that it lands in +# e4m3's normal range: with one level, every scale in this checkpoint falls +# into e4m3's subnormal range, where the format keeps about three bits and +# carries ~18 % relative error, which swamps the 4-bit value grid entirely. +E4M3_MAX = 448.0 + class CheckpointReader: def __init__(self, checkpoint: Path): @@ -77,23 +84,30 @@ def _rht16(weight: torch.Tensor) -> torch.Tensor: def _int4_weight( weight: torch.Tensor, group_size: int -) -> tuple[torch.Tensor, torch.Tensor]: +) -> tuple[torch.Tensor, torch.Tensor, float]: + """Pack to 4-bit with a two-level scale. + + Returns the packed values, the per-group e4m3 scale bytes, and the global + scale the kernel applies as its GEMM alpha. The effective scale of a group + is ``global_scale * e4m3(scale_byte)``. + """ rows, columns = weight.shape if columns % group_size: raise ValueError( f"K={columns} is not divisible by group_size={group_size}") grouped = weight.float().reshape(rows, columns // group_size, group_size) - scale = ( - grouped.abs().amax(dim=2).clamp_min(1e-8) / 7.0 - ).to(torch.float8_e4m3fn) - scale_float = scale.float().clamp_min(2.0**-9) + amax = grouped.abs().amax(dim=2).clamp_min(1e-12) + global_scale = max(float(amax.max()) / (E4M3_MAX * 7.0), 1e-12) + scale = (amax / 7.0 / global_scale).to(torch.float8_e4m3fn) + effective = scale.float() * global_scale values = ( - grouped / scale_float.unsqueeze(-1) + grouped / effective.unsqueeze(-1).clamp_min(1e-30) ).round().clamp(-7, 7).to(torch.int8).reshape(rows, columns) magnitude = values.abs().to(torch.uint8) code = magnitude | ((values < 0).to(torch.uint8) << 3) packed = code[:, 0::2] | (code[:, 1::2] << 4) - return packed.contiguous(), scale.view(torch.uint8).contiguous() + return (packed.contiguous(), scale.view(torch.uint8).contiguous(), + global_scale) def dequantize_int4( @@ -101,6 +115,7 @@ def dequantize_int4( scale: torch.Tensor, columns: int, group_size: int, + global_scale: float = 1.0, ) -> torch.Tensor: """Inverse of :func:`_int4_weight`, for scoring and reference paths.""" low = packed & 0x0F @@ -111,7 +126,7 @@ def dequantize_int4( (high & 0x08) != 0, -1, 1).to(torch.int8) values = torch.stack((low, high), dim=-1).flatten(1) rows = values.shape[0] - scale_float = scale.view(torch.float8_e4m3fn).float() + scale_float = scale.view(torch.float8_e4m3fn).float() * global_scale return ( values.float().reshape(rows, columns // group_size, group_size) * scale_float.unsqueeze(-1) @@ -129,7 +144,12 @@ def quantize_expert( quant_format: str, group_size: int, device: str, -) -> bytes: +) -> tuple[bytes, tuple[float, float]]: + """Return the fixed-size block and its (gate_up, down) global scales. + + INT8 uses per-output-channel scales that need no second level, so its + global scales are both 1.0. + """ if quant_format not in ("int8", "int4", "int4-rht"): raise ValueError(f"unsupported quantization format: {quant_format}") if quant_format == "int4-rht" and group_size != 16: @@ -139,19 +159,21 @@ def quantize_expert( if quant_format == "int8": gu_weight, gu_scale = _int8_weight(gate_up) dn_weight, dn_scale = _int8_weight(down) + alphas = (1.0, 1.0) else: if quant_format == "int4-rht": gate_up = _rht16(gate_up) down = _rht16(down) - gu_weight, gu_scale = _int4_weight(gate_up, group_size) - dn_weight, dn_scale = _int4_weight(down, group_size) + gu_weight, gu_scale, gu_alpha = _int4_weight(gate_up, group_size) + dn_weight, dn_scale, dn_alpha = _int4_weight(down, group_size) + alphas = (gu_alpha, dn_alpha) return b"".join(( _tensor_bytes(gu_weight), _tensor_bytes(gu_scale), _tensor_bytes(dn_weight), _tensor_bytes(dn_scale), bytes(_layout(quant_format, group_size)["padding"]), - )) + )), alphas def _parse_layers(value: str) -> range: @@ -206,7 +228,7 @@ def main() -> None: layout = _layout(args.format, args.group_size) block_bytes = sum(layout.values()) manifest = { - "format": f"flashrt-qwen36-moe-{args.format}-experts-v1", + "format": f"flashrt-qwen36-moe-{args.format}-experts-v2", "group_size": args.group_size if args.format != "int8" else None, "rht": args.format == "int4-rht", "num_layers": NUM_LAYERS, @@ -217,6 +239,13 @@ def main() -> None: "block_sizes": layout, "block_bytes": block_bytes, "block_alignment": BLOCK_ALIGNMENT, + # Per-expert global scales, one pair per expert. Kept out of the + # blocks so a block stays exactly block_bytes and 4096-aligned; + # they are a few bytes each and belong with the resident weights, + # where the kernel reads them as its GEMM alpha. + "global_scales": "global_scales_layer_NN.bin", + "global_scales_dtype": "float32", + "global_scales_layout": ["gate_up", "down"], } with (args.output / "manifest.json").open("w", encoding="utf-8") as f: json.dump(manifest, f, indent=2) @@ -225,18 +254,23 @@ def main() -> None: reader = CheckpointReader(args.checkpoint) for layer in args.layers: output_path = args.output / f"experts_layer_{layer:02d}.bin" + scale_check = args.output / f"global_scales_layer_{layer:02d}.bin" expected_bytes = NUM_EXPERTS * block_bytes - if output_path.is_file() and output_path.stat().st_size == expected_bytes: + if (output_path.is_file() + and output_path.stat().st_size == expected_bytes + and scale_check.is_file() + and scale_check.stat().st_size == NUM_EXPERTS * 2 * 4): print(f"layer {layer}: already complete") continue temporary_path = output_path.with_suffix(".bin.tmp") started = time.perf_counter() + alphas = [] with temporary_path.open("wb") as f: for expert in range(NUM_EXPERTS): gate_up = reader.expert( layer, "gate_up_proj", expert) down = reader.expert(layer, "down_proj", expert) - block = quantize_expert( + block, expert_alphas = quantize_expert( gate_up, down, quant_format=args.format, @@ -248,7 +282,14 @@ def main() -> None: f"expert block is {len(block)} bytes; " f"expected {block_bytes}") f.write(block) + alphas.extend(expert_alphas) os.replace(temporary_path, output_path) + scale_path = args.output / f"global_scales_layer_{layer:02d}.bin" + temporary_scale_path = scale_path.with_suffix(".bin.tmp") + with temporary_scale_path.open("wb") as f: + f.write(_tensor_bytes( + torch.tensor(alphas, dtype=torch.float32))) + os.replace(temporary_scale_path, scale_path) elapsed = time.perf_counter() - started gib = expected_bytes / 2**30 print( diff --git a/tests/test_qwen36_moe_edge_quant.py b/tests/test_qwen36_moe_edge_quant.py index b5209435..a5ddb5a4 100644 --- a/tests/test_qwen36_moe_edge_quant.py +++ b/tests/test_qwen36_moe_edge_quant.py @@ -44,8 +44,8 @@ def test_int4_grouped_round_trip_matches_packed_layout(): generator = torch.Generator().manual_seed(5) weight = torch.randn(64, 256, generator=generator) - packed, scale = _int4_weight(weight, 32) - restored = dequantize_int4(packed, scale, 256, 32) + packed, scale, global_scale = _int4_weight(weight, 32) + restored = dequantize_int4(packed, scale, 256, 32, global_scale) assert packed.shape == (64, 128) assert scale.shape == (64, 8) @@ -82,7 +82,7 @@ def test_expert_block_size_matches_manifest_layout(): ("int4", 32), ("int4-rht", 16), ): - block = quantize_expert( + block, alphas = quantize_expert( gate_up, down, quant_format=quant_format, @@ -91,6 +91,7 @@ def test_expert_block_size_matches_manifest_layout(): ) assert len(block) == sum( _layout(quant_format, group_size).values()) + assert len(alphas) == 2 def test_expert_blocks_are_aligned_for_direct_io(): @@ -218,12 +219,41 @@ def test_rht16_reduces_int4_error_on_outlier_heavy_weights(): rotated = score_expert( activation, gate_up, down, scheme="w4a16_rht16", group_size=16) - # The transform is only worth its cost when groups have outliers; a 10 % - # margin keeps this from asserting on noise. - assert rotated["relative_l2"] < 0.9 * plain["relative_l2"] + # The transform only pays off when groups have outliers, and once the + # two-level scale is correct the win is modest -- about 9 % here. An + # earlier single-level scale showed 38 %, but most of that was the + # transform compensating for scale error rather than doing its own job. + assert rotated["relative_l2"] < 0.97 * plain["relative_l2"] assert rotated["cosine"] > plain["cosine"] +def test_two_level_scale_keeps_group_scales_out_of_e4m3_subnormals(): + # Real expert weights have per-group amax around 0.02. With a single level + # the scale is amax/7 ~ 0.003, below e4m3's smallest normal 2**-6, where + # the format keeps about three bits; the scale error then swamps the 4-bit + # value grid. Factoring out a global scale moves the stored bytes into + # e4m3's normal range. + generator = torch.Generator().manual_seed(23) + weight = torch.randn(64, 512, generator=generator) * 0.02 + + _, scale_bytes, global_scale = _int4_weight(weight, 16) + stored = scale_bytes.view(torch.float8_e4m3fn).float() + + # Every stored byte is now in e4m3's normal range: that is the fix. + assert (stored[stored > 0] >= 2.0 ** -6).all() + assert global_scale > 0.0 + + packed, scale_bytes, global_scale = _int4_weight(weight, 16) + restored = dequantize_int4(packed, scale_bytes, 512, 16, global_scale) + error = ((restored - weight).norm() / weight.norm()).item() + + # A signed 4-bit grid over a per-group amax has step amax/7, so uniform + # quantization noise is step/sqrt(12). For Gaussian groups of 16, amax is + # about 2 sigma, giving ~8.6 % -- and that is what this measures, meaning + # the scale contributes nothing on top of the value grid. + assert error < 0.10, error + + def test_weight_only_schemes_order_by_bit_width(): generator = torch.Generator().manual_seed(17) activation = torch.randn(4, HIDDEN, generator=generator) From c19d728ba2c94426ac2d18433012556f13048e6c Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 30 Jul 2026 05:46:43 -0400 Subject: [PATCH 11/85] Model warm-started cache and cold prefill cost Three additions for sizing the latency of a streaming expert runtime. global_frequency and simulate_warm_lru cover filling each layer's cache at startup from offline selection statistics. On this checkpoint that removes 11 % of decode misses, and unlike the prompt-derived warm set it costs no adaptivity because the entries stay evictable. A set derived from the trace it is scored on is an oracle, so held-out traces are needed before trusting the figure. cold_prefill_blocks counts what prefill must read before the first token. Prefill routes every prompt token independently, so a layer costs the union of its tokens' selections; a 32-token prompt touched a mean of 74 experts per layer out of 256, or 4.88 GiB across the model. A resident set of 64 per layer brings that to 1.49 GiB. simulate_warm_lru also takes a window, which groups decode steps as a multi-token verification step would. It is here to document a negative result: grouping does not reduce reads, because an LRU already captures the reuse that a window's union would. Measured on this checkpoint, misses per token are 83.6 at window 1 and 83.7 at window 4, and rise to 85.7 at window 8 as a window's union starts evicting itself. An earlier estimate of 1.57x came from comparing a window's union against eight times its length with no cache present, which double-counts what the cache was already doing. --- qwen36_moe_edge/route_trace.py | 84 +++++++++++++++++++++++++++++ tests/test_qwen36_moe_edge_quant.py | 44 +++++++++++++++ 2 files changed, 128 insertions(+) diff --git a/qwen36_moe_edge/route_trace.py b/qwen36_moe_edge/route_trace.py index c0ce32ac..6e947381 100644 --- a/qwen36_moe_edge/route_trace.py +++ b/qwen36_moe_edge/route_trace.py @@ -136,6 +136,90 @@ def simulate_two_tier( } +def global_frequency( + trace: list[list[list[int]]]) -> list[Counter[int]]: + """Per-layer selection counts over a whole trace. + + Intended to be built from traces other than the one being evaluated: a set + derived from the trace it is scored on is an oracle, not a predictor. + """ + return [ + Counter(expert for step in layer_trace for expert in step) + for layer_trace in trace + ] + + +def simulate_warm_lru( + trace: list[list[list[int]]], + *, + prompt_tokens: int, + quota: int, + preload: list[Counter[int]] | None = None, + window: int = 1) -> dict[str, float]: + """Per-layer LRU, optionally warm-started, over ``window`` tokens at a time. + + ``preload`` fills each layer's cache with its most frequent experts before + decode, which is what a runtime can do at startup from offline statistics. + Entries are evictable: an earlier experiment showed that pinning a + prompt-derived set costs more adaptivity than it gains. + + ``window`` groups that many decode steps into one request, as a + multi-token verification step would. Note that this does not reduce reads: + an LRU already captures the reuse that a window's union would. + """ + if window < 1: + raise ValueError(f"window must be at least 1, got {window}") + misses = accesses = tokens = 0 + for index, layer_trace in enumerate(trace): + cache: OrderedDict[int, None] = OrderedDict() + if preload is not None: + for expert, _ in preload[index].most_common(quota): + cache[expert] = None + steps = layer_trace[prompt_tokens:] + for start in range(0, len(steps) - window + 1, window): + requested = { + expert + for step in steps[start:start + window] + for expert in step + } + if index == 0: + tokens += window + for expert in requested: + accesses += 1 + if expert in cache: + cache.move_to_end(expert) + continue + misses += 1 + if len(cache) >= quota: + cache.popitem(last=False) + cache[expert] = None + return { + "distinct_hit_rate": 1.0 - misses / max(accesses, 1), + "decode_misses_per_token": misses / max(tokens, 1), + } + + +def cold_prefill_blocks( + trace: list[list[list[int]]], + *, + prompt_tokens: int, + resident: list[set[int]]) -> int: + """Blocks prefill must read before the first token can be emitted. + + Prefill routes every prompt token independently, so a layer's cost is the + union of its tokens' selections. Whatever is already resident is free; the + rest sets the floor on time to first token. + """ + return sum( + len({ + expert + for step in layer_trace[:prompt_tokens] + for expert in step + } - resident[index]) + for index, layer_trace in enumerate(trace) + ) + + def read_volume( misses_per_token: float, *, diff --git a/tests/test_qwen36_moe_edge_quant.py b/tests/test_qwen36_moe_edge_quant.py index a5ddb5a4..3dcbfe20 100644 --- a/tests/test_qwen36_moe_edge_quant.py +++ b/tests/test_qwen36_moe_edge_quant.py @@ -10,9 +10,12 @@ score_expert, ) from qwen36_moe_edge.route_trace import ( + cold_prefill_blocks, + global_frequency, read_volume, simulate_lru, simulate_two_tier, + simulate_warm_lru, ) from qwen36_moe_edge.quantize_experts import ( BLOCK_ALIGNMENT, @@ -182,6 +185,47 @@ def test_read_volume_converts_misses_to_bandwidth_limits(): assert result["tok_s_at_2gbps"] == 1000.0 +def test_warm_start_removes_the_first_touch_of_a_frequent_expert(): + # A trace whose decode phase only ever wants expert 5. Cold, the first + # touch misses; warm-started from statistics that name expert 5, it does + # not. + trace = [[[5], [5], [5], [5]]] + frequency = global_frequency([[[5], [5]]]) + + cold = simulate_warm_lru(trace, prompt_tokens=2, quota=4) + warm = simulate_warm_lru( + trace, prompt_tokens=2, quota=4, preload=frequency) + + assert cold["decode_misses_per_token"] == 0.5 + assert warm["decode_misses_per_token"] == 0.0 + + +def test_windowing_does_not_reduce_reads_an_lru_already_serves(): + # Verifying several tokens at once requests the union of their selections. + # An LRU already holds a repeated expert, so grouping changes nothing. + trace = [[[0], [0, 1], [1], [0, 1]]] + + single = simulate_warm_lru(trace, prompt_tokens=0, quota=8, window=1) + grouped = simulate_warm_lru(trace, prompt_tokens=0, quota=8, window=2) + + assert single["decode_misses_per_token"] == grouped[ + "decode_misses_per_token"] + + +def test_cold_prefill_counts_the_union_prefill_touches(): + # Prefill routes each token independently, so a layer costs the union of + # its tokens' selections, less whatever is already resident. + trace = [[[0, 1], [1, 2], [9]], [[3], [4], [9]]] + + without = cold_prefill_blocks( + trace, prompt_tokens=2, resident=[set(), set()]) + with_resident = cold_prefill_blocks( + trace, prompt_tokens=2, resident=[{0, 1}, {3}]) + + assert without == 3 + 2 # {0,1,2} and {3,4} + assert with_resident == 1 + 1 # {2} and {4} + + def test_expert_forward_is_a_swiglu_over_the_gate_up_split(): generator = torch.Generator().manual_seed(13) activation = torch.randn(2, HIDDEN, generator=generator) From eba1014a39bf85ec49475a498f6316865c3538e4 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 30 Jul 2026 05:54:17 -0400 Subject: [PATCH 12/85] Validate the startup expert set on held-out prompts Filling each layer's cache at startup from offline selection statistics looked strongly positive, but the statistics had come from the trace being scored, which makes it an oracle rather than a predictor. A deployment builds the set from other traffic and then meets an unseen prompt. This traces several prompts in one model load and reports leave-one-out results: each prompt is scored against a set built only from the others. Measured over eight unrelated prompts, 32 prompt tokens and 32 decode tokens, mean over the eight held-out runs: | slots | metric | cold | held-out | oracle | |---|---|---|---|---| | 43 | decode misses/token | 117.19 | 105.43 | 96.25 | | 43 | cold prefill GiB | 4.81 | 3.11 | 2.35 | | 57 | decode misses/token | 107.53 | 87.79 | 73.51 | | 57 | cold prefill GiB | 4.81 | 2.66 | 1.65 | | 64 | decode misses/token | 104.57 | 80.04 | 63.09 | | 64 | cold prefill GiB | 4.81 | 2.44 | 1.34 | The set transfers: on topics from mechanical engineering to tort law, one built from seven prompts captures roughly 59 % of the oracle's decode benefit and 68 % of its prefill benefit on the eighth. The benefit also grows with slot count, from 10 % of decode misses at 43 slots to 23.5 % at 64. --- qwen36_moe_edge/warm_start_validation.py | 199 +++++++++++++++++++++++ 1 file changed, 199 insertions(+) create mode 100644 qwen36_moe_edge/warm_start_validation.py diff --git a/qwen36_moe_edge/warm_start_validation.py b/qwen36_moe_edge/warm_start_validation.py new file mode 100644 index 00000000..42a7ee07 --- /dev/null +++ b/qwen36_moe_edge/warm_start_validation.py @@ -0,0 +1,199 @@ +#!/usr/bin/env python3 +"""Validate a startup expert set on prompts it was not derived from. + +Filling each layer's cache at startup with its most frequently selected +experts looks strongly positive when the frequencies come from the same trace +being scored -- but that is an oracle, not a predictor. A deployment builds the +set offline, from other traffic, and then meets an unseen prompt. + +This runs several unrelated prompts through one model load and reports +leave-one-out results: for each prompt, the startup set is built from the +*other* prompts' traces only. The gap between that and the oracle is what the +heuristic actually costs. +""" + +from __future__ import annotations + +import argparse +import json +import statistics +from collections import Counter +from pathlib import Path + +import torch + +from qwen36_moe_edge.route_trace import ( + DEFAULT_BLOCK_BYTES, + cold_prefill_blocks, + global_frequency, + simulate_warm_lru, +) + + +def collect_traces( + checkpoint: str, + prompts: list[str], + *, + prompt_tokens: int, + new_tokens: int, + max_seq: int, + device: str) -> list[list[list[list[int]]]]: + """Trace every prompt's router selections in one model load.""" + from flash_rt.frontends.torch._nexn2_rtx_decode import ( + Nexn2DecodeState, + generate_greedy, + ) + from flash_rt.frontends.torch.qwen36_moe_rtx import ( + Qwen36MoeTextFrontendRtx, + ) + + frontend = Qwen36MoeTextFrontendRtx( + checkpoint, device=device, max_seq=max_seq, quant_scope="experts") + state = Nexn2DecodeState(frontend._weights, max_seq, device) + state.batched_prefill = False + + traces = [] + for index, prompt in enumerate(prompts): + input_ids = frontend.tokenizer( + prompt, return_tensors="pt", add_special_tokens=False, + ).input_ids[:, :prompt_tokens].to(device) + if input_ids.shape[1] != prompt_tokens: + raise ValueError( + f"prompt {index} is shorter than {prompt_tokens} tokens") + state.router_trace = { + layer: [] for layer in range(state.num_layers)} + with torch.no_grad(): + generate_greedy( + state, input_ids, new_tokens, frontend._fvk, device) + traces.append([ + [list(experts) for experts in state.router_trace[layer]] + for layer in range(state.num_layers) + ]) + print(f"traced prompt {index + 1}/{len(prompts)}", flush=True) + return traces + + +def _merge(frequencies: list[list[Counter[int]]]) -> list[Counter[int]]: + """Sum per-layer selection counts across several traces.""" + layers = len(frequencies[0]) + merged = [Counter() for _ in range(layers)] + for frequency in frequencies: + for layer in range(layers): + merged[layer].update(frequency[layer]) + return merged + + +def leave_one_out( + traces: list[list[list[list[int]]]], + *, + prompt_tokens: int, + quota: int, + block_bytes: int) -> list[dict[str, float]]: + """Score each prompt against a set built from the other prompts only.""" + frequencies = [global_frequency(trace) for trace in traces] + results = [] + for index, trace in enumerate(traces): + others = [f for position, f in enumerate(frequencies) + if position != index] + held_out = _merge(others) if others else None + oracle = frequencies[index] + + cold = simulate_warm_lru( + trace, prompt_tokens=prompt_tokens, quota=quota) + warm = simulate_warm_lru( + trace, prompt_tokens=prompt_tokens, quota=quota, + preload=held_out) + best = simulate_warm_lru( + trace, prompt_tokens=prompt_tokens, quota=quota, preload=oracle) + + def resident(frequency): + if frequency is None: + return [set() for _ in trace] + return [ + {expert for expert, _ in frequency[layer].most_common(quota)} + for layer in range(len(trace)) + ] + + results.append({ + "prompt": index, + "cold_misses_per_token": cold["decode_misses_per_token"], + "held_out_misses_per_token": warm["decode_misses_per_token"], + "oracle_misses_per_token": best["decode_misses_per_token"], + "cold_prefill_gib": cold_prefill_blocks( + trace, prompt_tokens=prompt_tokens, + resident=resident(None)) * block_bytes / 2 ** 30, + "held_out_prefill_gib": cold_prefill_blocks( + trace, prompt_tokens=prompt_tokens, + resident=resident(held_out)) * block_bytes / 2 ** 30, + "oracle_prefill_gib": cold_prefill_blocks( + trace, prompt_tokens=prompt_tokens, + resident=resident(oracle)) * block_bytes / 2 ** 30, + }) + return results + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--checkpoint", required=True) + parser.add_argument("--prompts-file", type=Path, required=True, + help="one prompt per line; blank lines ignored") + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--prompt-tokens", type=int, default=32) + parser.add_argument("--new-tokens", type=int, default=32) + parser.add_argument("--max-seq", type=int, default=128) + parser.add_argument("--quotas", default="43,57,64") + parser.add_argument("--block-bytes", type=int, default=DEFAULT_BLOCK_BYTES) + parser.add_argument("--device", default="cuda:0") + args = parser.parse_args() + + prompts = [ + line.strip() + for line in args.prompts_file.read_text(encoding="utf-8").splitlines() + if line.strip() + ] + if len(prompts) < 3: + parser.error("leave-one-out needs at least three prompts") + + traces = collect_traces( + args.checkpoint, prompts, + prompt_tokens=args.prompt_tokens, + new_tokens=args.new_tokens, + max_seq=args.max_seq, + device=args.device, + ) + + quotas = tuple(int(value) for value in args.quotas.split(",")) + report = {"prompt_count": len(prompts), "quotas": list(quotas), + "prompt_tokens": args.prompt_tokens, "by_quota": {}} + for quota in quotas: + rows = leave_one_out( + traces, prompt_tokens=args.prompt_tokens, quota=quota, + block_bytes=args.block_bytes) + report["by_quota"][str(quota)] = rows + + args.output.parent.mkdir(parents=True, exist_ok=True) + with args.output.open("w", encoding="utf-8") as f: + json.dump(report, f, indent=2) + f.write("\n") + + print(f"\n{'quota':>6} {'metric':<22} {'cold':>9} {'held-out':>9} " + f"{'oracle':>9} {'held-out win':>13}") + for quota in quotas: + rows = report["by_quota"][str(quota)] + for label, keys in ( + ("decode miss/token", ( + "cold_misses_per_token", "held_out_misses_per_token", + "oracle_misses_per_token")), + ("cold prefill GiB", ( + "cold_prefill_gib", "held_out_prefill_gib", + "oracle_prefill_gib")), + ): + cold, held, oracle = ( + statistics.mean(row[key] for row in rows) for key in keys) + win = (1.0 - held / cold) * 100.0 if cold else 0.0 + print(f"{quota:>6} {label:<22} {cold:>9.2f} {held:>9.2f} " + f"{oracle:>9.2f} {win:>12.1f}%") + + +if __name__ == "__main__": + main() From 153be4b820440a86ccc933d8c990c211de2f91c8 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 30 Jul 2026 06:46:51 -0400 Subject: [PATCH 13/85] Add cross-architecture parity for the gated kernels The tier split was justified by compiling for sm_87 and sm_110, which says nothing about what those kernels compute. This exercises each binding in the core and weight-only-4-bit tiers, records its output, and diffs a run on one target against a reference from another. No checkpoint: shapes come from the Qwen3.6 geometry and inputs from a fixed generator. Inputs are stored in the reference and replayed rather than regenerated. CUDA RNG is not bit-reproducible across architectures -- the Philox thread mapping follows occupancy -- so regenerating on the target compares kernels on different data. Divergence appears only past the first launch block, which presents as small tensors agreeing and large ones diverging, and reads as a kernel fault. Each case also checks its kernel against a Torch expression on the local device. That separates a real kernel fault from a harness problem: a broken kernel fails its local check before it disagrees with a remote reference. Measured sm_120a against sm_110: all twelve recorded output tensors are bitwise identical, and all eight cases match Torch locally on both. Two harness faults found while building this, both from assuming a calling convention instead of reading the call site: - The weighted-sum reducer writes float32 into a flat buffer. A bfloat16 destination yields NaN. - The linear-attention split broadcasts q and k from 16 stored key heads to all 32 value heads, so all three of its outputs are 32 * 128 wide. Sizing q and k for 16 heads makes the kernel write past them into the next allocation, which surfaces as a corrupted third output. A structural test pins the case-to-binding mapping so a rename cannot turn a case into a silent skip. --- qwen36_moe_edge/README.md | 27 +++ qwen36_moe_edge/kernel_parity.py | 271 ++++++++++++++++++++++++++++ tests/test_qwen36_moe_edge_quant.py | 21 +++ 3 files changed, 319 insertions(+) create mode 100644 qwen36_moe_edge/kernel_parity.py diff --git a/qwen36_moe_edge/README.md b/qwen36_moe_edge/README.md index 110fd779..ca1ae248 100644 --- a/qwen36_moe_edge/README.md +++ b/qwen36_moe_edge/README.md @@ -24,6 +24,33 @@ PYTHONPATH=. python qwen36_moe_edge/probe.py \ --group-size 16 ``` +Check that the gated kernels compute the same thing on another architecture. +Compiling for a target says nothing about what it computes there, so record a +reference where the kernels are known good and replay it elsewhere: + +```bash +# On a known-good target: +PYTHONPATH=. python qwen36_moe_edge/kernel_parity.py \ + --output parity_sm120.json + +# On the target under test, with the reference alongside it: +PYTHONPATH=. python qwen36_moe_edge/kernel_parity.py \ + --output parity_sm110.json \ + --reference parity_sm120.json +``` + +No checkpoint is involved: shapes come from the Qwen3.6 geometry and inputs +from a fixed generator. Inputs are stored in the reference and replayed rather +than regenerated, because CUDA RNG is not bit-reproducible across +architectures — the Philox thread mapping follows occupancy, so regenerating on +the target compares kernels on different data and reads as a kernel failure. +Divergence appears only past the first launch block, which is why small tensors +appear to agree and large ones do not. + +Each case also checks its kernel against a Torch expression on the local +device, so a genuine kernel fault is distinguishable from a harness or input +problem: a broken kernel fails its local check first. + Score the quantization schemes against the activations the router actually sends each expert, and optionally save the references a device can check itself against: diff --git a/qwen36_moe_edge/kernel_parity.py b/qwen36_moe_edge/kernel_parity.py new file mode 100644 index 00000000..b337a609 --- /dev/null +++ b/qwen36_moe_edge/kernel_parity.py @@ -0,0 +1,271 @@ +#!/usr/bin/env python3 +"""Cross-architecture parity for the qwen3_5_moe core and W4A16 kernels. + +The tier split was verified by compiling for sm_87 and sm_110, which proves +nothing about what the kernels compute there. This exercises each binding with +seeded inputs and records its output, so the same script run on another target +can be diffed against a reference produced on sm_120a. + +Deliberately no model and no checkpoint: shapes come from the Qwen3.6 geometry +and values from a fixed generator, so any machine can run it. + +Inputs are generated on the CPU and stored alongside the outputs. A comparison +run loads them from the reference rather than regenerating: CUDA RNG is not +bit-reproducible across architectures -- the Philox thread mapping follows +occupancy -- so regenerating on the target would compare kernels on different +data and read as a kernel failure. Divergence appears only past the first +launch block, which is why small tensors matched and large ones did not. +""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import torch + +HIDDEN = 2048 +INTERMEDIATE = 512 +TOPK = 8 +NUM_EXPERTS = 256 +# Linear-attention geometry. The split kernel broadcasts q and k from the 16 +# stored key heads to all 32 value heads, so every one of its three outputs is +# NV * HK wide -- not the 16-head width the stored layout suggests. Sizing +# q/k at 2048 makes the kernel write past them into whatever the allocator +# placed next, which shows up as a corrupted third output. +NV = 32 +HK = 128 +HV = 128 + + +# Inputs recorded by the reference run and replayed by comparison runs. +_INPUTS: dict[str, torch.Tensor] = {} +_REPLAY: dict[str, torch.Tensor] | None = None +_CPU_GEN = torch.Generator().manual_seed(20260730) +_CASE = "" + + +def _record(name: str, tensor: torch.Tensor, device) -> torch.Tensor: + """Return the replayed input if one exists, else keep what we generated.""" + key = f"{_CASE}.in.{name}" + if _REPLAY is not None: + if key not in _REPLAY: + raise KeyError(f"reference has no input {key}") + tensor = _REPLAY[key].to(dtype=tensor.dtype) + _INPUTS[key] = tensor.detach().cpu() + return tensor.to(device) + + +def _bf16(name, shape, device, scale=1.0): + """A bfloat16 input, generated on the CPU so it is machine-independent.""" + values = torch.randn(*shape, generator=_CPU_GEN, dtype=torch.float32) + return _record(name, (values * scale).to(torch.bfloat16), device) + + +def case_bf16_matvec(fvk, device): + x = _bf16("x", (1, HIDDEN), device) + w = _bf16("w", (HIDDEN, HIDDEN), device, 0.02) + out = torch.zeros(1, HIDDEN, dtype=torch.bfloat16, device=device) + rc = fvk.bf16_matvec_sm120_bf16( + x.data_ptr(), w.data_ptr(), out.data_ptr(), HIDDEN, HIDDEN, 0) + torch.cuda.synchronize() + return {"rc": rc, "out": out, "torch": (x.float() @ w.float().T)} + + +def case_router_topk(fvk, device): + logits = _bf16("logits", (NUM_EXPERTS,), device, 4.0).contiguous() + idx = torch.empty(TOPK, dtype=torch.int32, device=device) + val = torch.empty(TOPK, dtype=torch.float32, device=device) + rc = fvk.moe_router_topk_sm120_bf16( + logits.data_ptr(), idx.data_ptr(), val.data_ptr(), + NUM_EXPERTS, TOPK, 0) + torch.cuda.synchronize() + reference = torch.topk(logits.float(), TOPK) + return {"rc": rc, "idx": idx, "val": val, + "torch_idx": reference.indices.to(torch.int32), + "torch_val": reference.values} + + +def case_silu_mul(fvk, device): + n = 4096 + g = _bf16("g", (n,), device) + u = _bf16("u", (n,), device) + out = torch.zeros(n, dtype=torch.bfloat16, device=device) + rc = fvk.silu_mul_sm120_bf16( + g.data_ptr(), u.data_ptr(), out.data_ptr(), n, 0) + torch.cuda.synchronize() + return {"rc": rc, "out": out, + "torch": torch.nn.functional.silu(g.float()) * u.float()} + + +def case_sigmoid_mul(fvk, device): + n = 4096 + x = _bf16("x", (n,), device) + gate = _bf16("gate", (n,), device) + out = torch.zeros(n, dtype=torch.bfloat16, device=device) + rc = fvk.sigmoid_mul_sm120_bf16( + x.data_ptr(), gate.data_ptr(), out.data_ptr(), n, 0) + torch.cuda.synchronize() + return {"rc": rc, "out": out, + "torch": x.float() * torch.sigmoid(gate.float())} + + +def case_weighted_sum(fvk, device): + d_dn = _bf16("d_dn", (TOPK, HIDDEN), device) + rows = torch.arange(TOPK, dtype=torch.int32, device=device) + weights = _record("weights", torch.softmax( + torch.randn(TOPK, generator=_CPU_GEN), -1), device) + # The reducer writes float32 and expects a flat buffer: see + # tests/test_qwen36_moe_gpu.py and the decode call site. Handing it a + # bfloat16 destination yields NaN, not a wrong-but-plausible answer. + out = torch.zeros(HIDDEN, dtype=torch.float32, device=device) + rc = fvk.moe_weighted_sum_sm120_bf16( + d_dn.data_ptr(), rows.data_ptr(), weights.contiguous().data_ptr(), + out.data_ptr(), 1, TOPK, HIDDEN, HIDDEN, 0) + torch.cuda.synchronize() + return {"rc": rc, "out": out, "torch": weights.float() @ d_dn.float()} + + +def case_w16a16_gemm(fvk, device): + m, n, k = 16, HIDDEN, HIDDEN + x = _bf16("x", (m, k), device) + w = _bf16("w", (n, k), device, 0.02) + out = torch.zeros(m, n, dtype=torch.bfloat16, device=device) + rc = fvk.w16a16_gemm_sm120_bf16( + x.data_ptr(), w.data_ptr(), out.data_ptr(), m, n, k, 1.0, 0) + torch.cuda.synchronize() + return {"rc": rc, "out": out, "torch": x.float() @ w.float().T} + + +def case_lin_split_qkv(fvk, device): + S = 4 + conv_out = _bf16("conv_out", (S, 8192), device).contiguous() + q32 = torch.zeros(S, NV, HK, dtype=torch.bfloat16, device=device) + k32 = torch.zeros(S, NV, HK, dtype=torch.bfloat16, device=device) + v32 = torch.zeros(S, NV, HV, dtype=torch.bfloat16, device=device) + fvk.qwen35moe_lin_split_qkv_broadcast_bf16( + conv_out.data_ptr(), q32.data_ptr(), k32.data_ptr(), v32.data_ptr(), + S, 0) + torch.cuda.synchronize() + return {"q32": q32, "k32": k32, "v32": v32} + + +def case_split_q_gate(fvk, device): + S = 4 + q_proj = _bf16("q_proj", (S, 8192), device).contiguous() + q_pre = torch.zeros(S, 4096, dtype=torch.bfloat16, device=device) + gate = torch.zeros(S, 4096, dtype=torch.bfloat16, device=device) + fvk.qwen35moe_split_q_gate_bf16( + q_proj.data_ptr(), q_pre.data_ptr(), gate.data_ptr(), S, 0) + torch.cuda.synchronize() + return {"q_pre": q_pre, "gate": gate} + + +CASES = { + "bf16_matvec": case_bf16_matvec, + "moe_router_topk": case_router_topk, + "silu_mul": case_silu_mul, + "sigmoid_mul": case_sigmoid_mul, + "moe_weighted_sum": case_weighted_sum, + "w16a16_gemm": case_w16a16_gemm, + "lin_split_qkv": case_lin_split_qkv, + "split_q_gate": case_split_q_gate, +} + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--reference", type=Path) + parser.add_argument("--device", default="cuda:0") + args = parser.parse_args() + + from flash_rt import flash_rt_kernels as fvk + + capability = torch.cuda.get_device_capability() + print(f"device: {torch.cuda.get_device_name(0)} sm_{capability[0]}" + f"{capability[1]}") + + global _REPLAY, _CASE + if args.reference and args.reference.with_suffix(".pt").is_file(): + _REPLAY = torch.load( + args.reference.with_suffix(".pt"), weights_only=True) + print(f"replaying inputs from {args.reference.with_suffix('.pt')}") + + outputs, report = {}, {} + for name, case in CASES.items(): + _CASE = name + if not hasattr(fvk, _binding_of(name)): + report[name] = {"status": "binding absent"} + print(f"{name:<20} binding absent") + continue + try: + result = case(fvk, args.device) + except Exception as error: # noqa: BLE001 + report[name] = {"status": f"raised {type(error).__name__}: {error}"} + print(f"{name:<20} RAISED {error}") + continue + entry = {"status": "ok"} + if "rc" in result: + entry["rc"] = int(result.pop("rc")) + for key in list(result): + if key.startswith("torch"): + continue + outputs[f"{name}.{key}"] = result[key].detach().float().cpu() + # Local agreement with torch, where a reference was computed. + for key in ("out", "val"): + if key in result and "torch" in result: + entry["torch_cosine"] = _cosine( + result[key], result["torch"]) + if "torch_idx" in result: + # Compare as a set: the top-8 of this input contains an exact tie + # (two logits at 8.9375), and the kernel and torch.topk are free to + # break it differently while both being right. + entry["topk_index_set_match"] = int( + set(result["idx"].tolist()) + == set(result["torch_idx"].tolist())) + entry["topk_value_match"] = int(torch.allclose( + result["val"].cpu(), result["torch_val"].cpu())) + report[name] = entry + print(f"{name:<20} {entry}") + + args.output.parent.mkdir(parents=True, exist_ok=True) + torch.save({**_INPUTS, **outputs}, args.output.with_suffix(".pt")) + with args.output.open("w", encoding="utf-8") as f: + json.dump({"capability": list(capability), "cases": report}, f, + indent=2) + f.write("\n") + + if args.reference and args.reference.with_suffix(".pt").is_file(): + print("\n--- against reference ---") + expected = _REPLAY + worst = 1.0 + for key in sorted(k for k in set(expected) & set(outputs) + if ".in." not in k): + cosine = _cosine(outputs[key], expected[key]) + exact = torch.equal(outputs[key], expected[key]) + worst = min(worst, cosine) + print(f"{key:<34} cos={cosine:.8f} bitwise={'yes' if exact else 'no'}") + compared = {k for k in set(expected) | set(outputs) if ".in." not in k} + missing = sorted(compared - (set(expected) & set(outputs))) + if missing: + print(f"outputs present on only one side: {missing}") + print(f"worst cosine: {worst:.8f}") + + +def _binding_of(name: str) -> str: + return { + "lin_split_qkv": "qwen35moe_lin_split_qkv_broadcast_bf16", + "split_q_gate": "qwen35moe_split_q_gate_bf16", + }.get(name, f"{name}_sm120_bf16") + + +def _cosine(a: torch.Tensor, b: torch.Tensor) -> float: + x = a.detach().float().flatten().cpu() + y = b.detach().float().flatten().cpu() + return torch.nn.functional.cosine_similarity(x, y, dim=0).item() + + +if __name__ == "__main__": + main() diff --git a/tests/test_qwen36_moe_edge_quant.py b/tests/test_qwen36_moe_edge_quant.py index 3dcbfe20..355b33c9 100644 --- a/tests/test_qwen36_moe_edge_quant.py +++ b/tests/test_qwen36_moe_edge_quant.py @@ -4,6 +4,7 @@ import torch +from qwen36_moe_edge.kernel_parity import CASES, _binding_of from qwen36_moe_edge.expert_quality import ( SCHEMES, expert_forward, @@ -226,6 +227,26 @@ def test_cold_prefill_counts_the_union_prefill_touches(): assert with_resident == 1 + 1 # {2} and {4} +def test_parity_cases_name_bindings_that_exist_in_the_tiers(): + # The parity harness resolves a case name to a binding name. If a kernel is + # renamed and this mapping is not updated, the case silently reports + # "binding absent" and a real regression passes unnoticed. + expected = { + "bf16_matvec": "bf16_matvec_sm120_bf16", + "moe_router_topk": "moe_router_topk_sm120_bf16", + "silu_mul": "silu_mul_sm120_bf16", + "sigmoid_mul": "sigmoid_mul_sm120_bf16", + "moe_weighted_sum": "moe_weighted_sum_sm120_bf16", + "w16a16_gemm": "w16a16_gemm_sm120_bf16", + "lin_split_qkv": "qwen35moe_lin_split_qkv_broadcast_bf16", + "split_q_gate": "qwen35moe_split_q_gate_bf16", + } + + assert set(CASES) == set(expected) + for case, binding in expected.items(): + assert _binding_of(case) == binding + + def test_expert_forward_is_a_swiglu_over_the_gate_up_split(): generator = torch.Generator().manual_seed(13) activation = torch.randn(2, HIDDEN, generator=generator) From 9db03275a3bfd4842ebd2565fd49267fef959ae7 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 30 Jul 2026 07:14:13 -0400 Subject: [PATCH 14/85] Add a bounded streaming cache for routed-expert blocks The target device holds a fraction of the experts, so the cache is the runtime: what it holds and how fast it refills set both the token rate and the time to first token. Three properties are enforced rather than left to convention. The budget is a hard limit. On unified memory the weights, the cache, the staging buffers and the operating system draw on the same physical memory, so a runtime that merely intends to stay small cannot be measured. Construction computes its footprint and refuses to allocate over budget, naming the quota that would fit. plan() and max_slots_per_layer() answer the sizing question before anything is allocated. Reads use O_DIRECT. Streaming tens of GiB through the page cache would make it compete with the resident weights for that same memory, which is why the bundle pads blocks to 4096 bytes. Construction rejects a bundle whose blocks are not aligned, and checks that the pinned staging buffers are too. Misses are fetched concurrently, because a single reader leaves most of an NVMe device idle. Measured against the real bundle: 3.17 GB/s with one staging buffer, 5.70 with two, 5.85 with four, and flat after that -- matching an independent fio sweep of the same access pattern. Requiring a per-layer quota of at least experts_per_token makes a class of bug structurally impossible: one token's experts cannot evict each other, so a caller may hold every pointer from a get_many at once. An earlier prototype on another model lost most of its hit rate to exactly that. warm() preloads each layer's most frequent experts and leaves them evictable. On held-out prompts an offline set removes about a quarter of decode misses and half of the cold prefill read at this quota, where pinning it instead loses more adaptivity than it gains. close() releases the slots. It previously dropped only descriptors and the thread pool, which would leak the largest allocation in the process on any reconfiguration. --- qwen36_moe_edge/expert_cache.py | 338 ++++++++++++++++++++++++++ tests/test_qwen36_moe_expert_cache.py | 140 +++++++++++ 2 files changed, 478 insertions(+) create mode 100644 qwen36_moe_edge/expert_cache.py create mode 100644 tests/test_qwen36_moe_expert_cache.py diff --git a/qwen36_moe_edge/expert_cache.py b/qwen36_moe_edge/expert_cache.py new file mode 100644 index 00000000..bd67192a --- /dev/null +++ b/qwen36_moe_edge/expert_cache.py @@ -0,0 +1,338 @@ +#!/usr/bin/env python3 +"""A bounded, streaming cache for routed-expert blocks. + +The target device holds only a fraction of the experts, so the cache is the +runtime: what it can hold and how fast it can refill decide both token rate and +time to first token. Three properties follow from that and are enforced here +rather than left to convention. + +**The budget is a hard limit, not a projection.** On a unified-memory device the +weights, the cache, the staging buffers and the operating system all draw on the +same physical memory, so a runtime that merely intends to stay small is not +measurable. Construction computes its own footprint and refuses to allocate if +it would exceed the budget. + +**Reads bypass the page cache.** Streaming tens of GiB of blocks through the +page cache would make it compete with the resident weights for that same +memory. Reads use ``O_DIRECT``, which is why the bundle pads each block to a +4096-byte boundary. + +**Misses are fetched concurrently.** A single reader leaves a large part of an +NVMe device idle; measured on one, four readers were worth 1.7x over one and +eight saturated it. ``get_many`` issues a layer's misses together. + +A per-layer quota of at least ``num_experts_per_token`` also makes a class of +bug structurally impossible: the experts one token needs cannot evict each +other, so a caller may hold several pointers from the same ``get_many`` at once. +""" + +from __future__ import annotations + +import json +import os +from collections import Counter, OrderedDict +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass, field +from pathlib import Path + +import torch + + +@dataclass +class CacheConfig: + """Sizing and placement. ``budget_bytes`` of 0 disables the check.""" + + bundle: Path + slots_per_layer: int + staging_buffers: int = 4 + budget_bytes: int = 0 + reserve_bytes: int = 0 + resident_bytes: int = 0 + device: str = "cuda:0" + experts_per_token: int = 8 + read_chunk: int = 1 << 28 + metadata: dict = field(default_factory=dict) + + +class CacheBudgetError(RuntimeError): + """The requested cache does not fit the budget it was given.""" + + +def _load_manifest(bundle: Path) -> dict: + path = bundle / "manifest.json" + if not path.is_file(): + raise FileNotFoundError(f"expert bundle is missing {path}") + with path.open(encoding="utf-8") as f: + manifest = json.load(f) + for key in ("block_bytes", "block_alignment", "num_layers", + "num_experts", "block_sizes"): + if key not in manifest: + raise ValueError(f"{path} has no {key!r}") + block_bytes = int(manifest["block_bytes"]) + alignment = int(manifest["block_alignment"]) + if block_bytes % alignment: + raise ValueError( + f"{path}: block_bytes {block_bytes} is not a multiple of " + f"{alignment}, so direct reads of it are impossible") + return manifest + + +class ExpertCache: + """Per-layer LRU over fixed-size blocks read straight from storage.""" + + def __init__(self, config: CacheConfig): + self.config = config + self.manifest = _load_manifest(config.bundle) + self.block_bytes = int(self.manifest["block_bytes"]) + self.alignment = int(self.manifest["block_alignment"]) + self.num_layers = int(self.manifest["num_layers"]) + self.num_experts = int(self.manifest["num_experts"]) + + if config.slots_per_layer < config.experts_per_token: + raise ValueError( + f"slots_per_layer={config.slots_per_layer} is below " + f"experts_per_token={config.experts_per_token}; one token's " + "experts would evict each other and a caller could not hold " + "their pointers at once") + + self.footprint = self.plan(config, self.manifest) + if config.budget_bytes: + total = self.footprint["projected_bytes"] + if total > config.budget_bytes: + raise CacheBudgetError( + f"cache needs {total / 2**30:.3f} GiB " + f"(slots {self.footprint['slot_bytes'] / 2**30:.3f} + " + f"staging {self.footprint['staging_bytes'] / 2**30:.3f} + " + f"resident {config.resident_bytes / 2**30:.3f} + " + f"reserve {config.reserve_bytes / 2**30:.3f}) but the " + f"budget is {config.budget_bytes / 2**30:.3f} GiB. " + f"Reduce slots_per_layer below " + f"{self.max_slots_per_layer(config, self.manifest)}.") + + self._total_slots = config.slots_per_layer * self.num_layers + self.slots = torch.empty( + self._total_slots, self.block_bytes, + dtype=torch.uint8, device=config.device) + self._staging = [ + torch.empty(self.block_bytes, dtype=torch.uint8).pin_memory() + for _ in range(config.staging_buffers) + ] + for buffer in self._staging: + if buffer.data_ptr() % self.alignment: + raise RuntimeError( + "a pinned staging buffer is not " + f"{self.alignment}-byte aligned, which direct reads " + "require") + self._pool = ThreadPoolExecutor(max_workers=config.staging_buffers) + + # Per-layer LRU of expert -> slot index, and that layer's free slots. + self._lru: list[OrderedDict[int, int]] = [ + OrderedDict() for _ in range(self.num_layers)] + self._free: list[list[int]] = [ + list(range(layer * config.slots_per_layer, + (layer + 1) * config.slots_per_layer)) + for layer in range(self.num_layers) + ] + self._fds: dict[int, int] = {} + self.hits = 0 + self.misses = 0 + self.bytes_read = 0 + + # ── sizing, answerable before anything is allocated ── + + @staticmethod + def plan(config: CacheConfig, manifest: dict) -> dict[str, int]: + """What this configuration would occupy.""" + block_bytes = int(manifest["block_bytes"]) + slot_bytes = ( + config.slots_per_layer * int(manifest["num_layers"]) * block_bytes) + staging_bytes = config.staging_buffers * block_bytes + return { + "block_bytes": block_bytes, + "slots_per_layer": config.slots_per_layer, + "slot_bytes": slot_bytes, + "staging_bytes": staging_bytes, + "resident_bytes": config.resident_bytes, + "reserve_bytes": config.reserve_bytes, + "projected_bytes": ( + slot_bytes + staging_bytes + + config.resident_bytes + config.reserve_bytes), + } + + @staticmethod + def max_slots_per_layer(config: CacheConfig, manifest: dict) -> int: + """Largest per-layer quota that fits ``config.budget_bytes``.""" + if not config.budget_bytes: + return int(manifest["num_experts"]) + block_bytes = int(manifest["block_bytes"]) + available = ( + config.budget_bytes - config.resident_bytes + - config.reserve_bytes - config.staging_buffers * block_bytes) + if available <= 0: + return 0 + return min( + int(manifest["num_experts"]), + available // (block_bytes * int(manifest["num_layers"]))) + + # ── reading ── + + def _fd(self, layer: int) -> int: + if layer not in self._fds: + path = self.config.bundle / f"experts_layer_{layer:02d}.bin" + expected = self.num_experts * self.block_bytes + size = path.stat().st_size + if size != expected: + raise ValueError( + f"{path} is {size} bytes; expected {expected} " + f"({self.num_experts} x {self.block_bytes})") + self._fds[layer] = os.open( + path, os.O_RDONLY | getattr(os, "O_DIRECT", 0)) + return self._fds[layer] + + def _fetch(self, layer: int, expert: int, slot: int, buffer: int) -> None: + staging = self._staging[buffer] + view = memoryview(staging.numpy()) + fd = self._fd(layer) + base = expert * self.block_bytes + offset = 0 + while offset < self.block_bytes: + length = min(self.config.read_chunk, self.block_bytes - offset) + read = os.preadv(fd, [view[offset:offset + length]], base + offset) + if read <= 0: + raise IOError( + f"short read of layer {layer} expert {expert} at " + f"{offset}/{self.block_bytes}") + offset += read + self.slots[slot].copy_(staging) + self.bytes_read += self.block_bytes + + def _claim(self, layer: int, expert: int) -> int: + """A slot for an expert not currently held, evicting if necessary.""" + free = self._free[layer] + if free: + slot = free.pop() + else: + slot = self._lru[layer].popitem(last=False)[1] + self._lru[layer][expert] = slot + return slot + + def get_many(self, layer: int, experts) -> list[int]: + """Device pointers for several experts of one layer, misses in parallel. + + With ``slots_per_layer >= experts_per_token`` none of the returned + pointers can be invalidated by the others. + """ + wanted = list(dict.fromkeys(int(expert) for expert in experts)) + if len(wanted) > self.config.slots_per_layer: + raise ValueError( + f"asked for {len(wanted)} experts of layer {layer} but the " + f"quota is {self.config.slots_per_layer}") + pending = [] + for expert in wanted: + slot = self._lru[layer].get(expert) + if slot is not None: + self._lru[layer].move_to_end(expert) + self.hits += 1 + continue + self.misses += 1 + pending.append((expert, self._claim(layer, expert))) + if pending: + futures = [ + self._pool.submit( + self._fetch, layer, expert, slot, + index % len(self._staging)) + for index, (expert, slot) in enumerate(pending) + ] + for future in futures: + future.result() + torch.cuda.synchronize(self.config.device) + return [ + int(self.slots[self._lru[layer][expert]].data_ptr()) + for expert in wanted + ] + + def get(self, layer: int, expert: int) -> int: + return self.get_many(layer, (expert,))[0] + + # ── startup ── + + def warm(self, frequency: list[Counter]) -> int: + """Preload each layer's most frequently selected experts. + + Entries stay evictable. Measured on held-out prompts, a set built from + unrelated traffic removes about a quarter of decode misses and half of + the cold prefill read at this quota; pinning it instead costs more + adaptivity than it gains. + """ + if len(frequency) != self.num_layers: + raise ValueError( + f"frequency has {len(frequency)} layers, expected " + f"{self.num_layers}") + loaded = 0 + for layer in range(self.num_layers): + experts = [ + expert for expert, _ in + frequency[layer].most_common(self.config.slots_per_layer) + ] + for start in range(0, len(experts), self.config.slots_per_layer): + chunk = experts[start:start + self.config.slots_per_layer] + self.get_many(layer, chunk) + loaded += len(chunk) + self.hits = 0 + self.misses = 0 + self.bytes_read = 0 + return loaded + + # ── reporting ── + + def stats(self) -> dict[str, float]: + requests = self.hits + self.misses + report = dict(self.footprint) + report.update({ + "hits": self.hits, + "misses": self.misses, + "hit_rate": self.hits / requests if requests else 0.0, + "bytes_read": self.bytes_read, + "resident_experts": sum(len(lru) for lru in self._lru), + }) + if torch.cuda.is_available(): + free, total = torch.cuda.mem_get_info(self.config.device) + report.update({ + "device_free_bytes": free, + "device_total_bytes": total, + "torch_allocated_bytes": torch.cuda.memory_allocated( + self.config.device), + "torch_peak_allocated_bytes": torch.cuda.max_memory_allocated( + self.config.device), + "torch_reserved_bytes": torch.cuda.memory_reserved( + self.config.device), + "torch_peak_reserved_bytes": torch.cuda.max_memory_reserved( + self.config.device), + }) + return report + + def close(self) -> None: + """Release everything, including the slots. + + The slot array is the largest single allocation a runtime makes, so a + close that only dropped file descriptors would leak the entire cache + on any reconfiguration -- on a device where the budget is the whole + point, that is not a detail. After this the cache is unusable. + """ + self._pool.shutdown(wait=True) + for fd in self._fds.values(): + os.close(fd) + self._fds.clear() + for lru in self._lru: + lru.clear() + self._free = [[] for _ in range(self.num_layers)] + self.slots = None + self._staging = [] + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + def __enter__(self) -> "ExpertCache": + return self + + def __exit__(self, *_) -> None: + self.close() diff --git a/tests/test_qwen36_moe_expert_cache.py b/tests/test_qwen36_moe_expert_cache.py new file mode 100644 index 00000000..1b1fc8e5 --- /dev/null +++ b/tests/test_qwen36_moe_expert_cache.py @@ -0,0 +1,140 @@ +"""Sizing and admission tests for the streaming expert cache. + +These exercise the parts that need no device and no bundle: the footprint +arithmetic, the budget refusal, and the quota rule that makes one token's +experts unable to evict each other. +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +from qwen36_moe_edge.expert_cache import ( + CacheBudgetError, + CacheConfig, + ExpertCache, +) + + +GIB = 2 ** 30 +# The shipped INT4 group-16 bundle. +MANIFEST = { + "block_bytes": 1769472, + "block_alignment": 4096, + "num_layers": 40, + "num_experts": 256, + "block_sizes": {}, +} + + +def _config(**overrides) -> CacheConfig: + values = { + "bundle": Path("/nonexistent"), + "slots_per_layer": 57, + "staging_buffers": 4, + } + values.update(overrides) + return CacheConfig(**values) + + +def test_plan_accounts_for_slots_staging_resident_and_reserve(): + config = _config( + slots_per_layer=57, + staging_buffers=4, + resident_bytes=int(1.696 * GIB), + reserve_bytes=int(1.5 * GIB), + ) + + plan = ExpertCache.plan(config, MANIFEST) + + assert plan["slot_bytes"] == 57 * 40 * 1769472 + assert plan["staging_bytes"] == 4 * 1769472 + assert plan["projected_bytes"] == ( + plan["slot_bytes"] + plan["staging_bytes"] + + plan["resident_bytes"] + plan["reserve_bytes"]) + # 57 slots is what quantizing the GDN weights to INT8 pays for, and it has + # to fit the 7 GiB target. + assert plan["projected_bytes"] < 7.0 * GIB + + +def test_max_slots_shrinks_when_resident_weights_grow(): + # Leaving the GDN weights at BF16 costs 0.94 GiB of resident, which is + # what the extra slots were bought with. + lean = _config( + budget_bytes=int(7.0 * GIB), reserve_bytes=int(1.5 * GIB), + resident_bytes=int(1.696 * GIB)) + heavy = _config( + budget_bytes=int(7.0 * GIB), reserve_bytes=int(1.5 * GIB), + resident_bytes=int(2.637 * GIB)) + + generous = ExpertCache.max_slots_per_layer(lean, MANIFEST) + tight = ExpertCache.max_slots_per_layer(heavy, MANIFEST) + + assert generous >= 57 + assert tight < generous + assert generous - tight >= 12 + + +def test_max_slots_is_zero_when_the_budget_is_already_spent(): + config = _config( + budget_bytes=int(2.0 * GIB), reserve_bytes=int(1.5 * GIB), + resident_bytes=int(1.0 * GIB)) + + assert ExpertCache.max_slots_per_layer(config, MANIFEST) == 0 + + +def test_quota_below_experts_per_token_is_rejected(tmp_path): + # Below the top-k a single token's experts would evict one another, so a + # caller could not hold their pointers at the same time. + (tmp_path / "manifest.json").write_text(_manifest_json()) + config = _config( + bundle=tmp_path, slots_per_layer=4, experts_per_token=8) + + with pytest.raises(ValueError, match="experts_per_token"): + ExpertCache(config) + + +def test_construction_refuses_a_cache_that_exceeds_the_budget(tmp_path): + (tmp_path / "manifest.json").write_text(_manifest_json()) + config = _config( + bundle=tmp_path, + slots_per_layer=200, + budget_bytes=int(7.0 * GIB), + reserve_bytes=int(1.5 * GIB), + resident_bytes=int(1.696 * GIB), + ) + + with pytest.raises(CacheBudgetError) as error: + ExpertCache(config) + # The message has to say what to do about it. + assert "Reduce slots_per_layer" in str(error.value) + + +def test_a_bundle_whose_blocks_are_unaligned_is_rejected(tmp_path): + (tmp_path / "manifest.json").write_text( + _manifest_json(block_bytes=3151872)) # the unpadded INT8 payload + + with pytest.raises(ValueError, match="not a multiple"): + ExpertCache(_config(bundle=tmp_path)) + + +def _manifest_json(**overrides) -> str: + import json + + manifest = dict(MANIFEST) + manifest.update(overrides) + return json.dumps(manifest) + + +def test_close_is_documented_to_release_the_slots(): + # The slot array is the largest allocation the runtime makes. A close that + # only dropped descriptors would leak the whole cache on reconfiguration, + # which on a budgeted device is a functional defect rather than untidiness. + import inspect + + source = inspect.getsource(ExpertCache.close) + + assert "self.slots = None" in source + assert "self._staging = []" in source From edc28ace71e409a4e70c5411d15d190f65b6bdda Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 30 Jul 2026 07:15:30 -0400 Subject: [PATCH 15/85] Make the page-cache bypass switchable so it can be measured Whether O_DIRECT actually keeps the block stream out of the page cache was an assertion. It matters because on a device holding a fraction of the experts the page cache competes with the resident weights for the same physical memory, and the failure mode is an out-of-memory on the target rather than a wrong answer. A config flag selects buffered reads, which gives the measurement a control reading exactly the same bytes. Against the real bundle on Thor, 2.11 GiB of blocks: the direct path grew the page cache 0.07 GiB, the buffered control grew it 2.14 GiB. The difference is 2.07 GiB of memory an 8 GiB device does not have. The flag also covers filesystems without O_DIRECT. --- qwen36_moe_edge/expert_cache.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/qwen36_moe_edge/expert_cache.py b/qwen36_moe_edge/expert_cache.py index bd67192a..98f64df6 100644 --- a/qwen36_moe_edge/expert_cache.py +++ b/qwen36_moe_edge/expert_cache.py @@ -51,6 +51,9 @@ class CacheConfig: device: str = "cuda:0" experts_per_token: int = 8 read_chunk: int = 1 << 28 + # Bypass the page cache. Off only to demonstrate what happens when it is + # not bypassed, or on a filesystem without O_DIRECT. + direct: bool = True metadata: dict = field(default_factory=dict) @@ -185,8 +188,10 @@ def _fd(self, layer: int) -> int: raise ValueError( f"{path} is {size} bytes; expected {expected} " f"({self.num_experts} x {self.block_bytes})") - self._fds[layer] = os.open( - path, os.O_RDONLY | getattr(os, "O_DIRECT", 0)) + flags = os.O_RDONLY + if self.config.direct: + flags |= getattr(os, "O_DIRECT", 0) + self._fds[layer] = os.open(path, flags) return self._fds[layer] def _fetch(self, layer: int, expert: int, slot: int, buffer: int) -> None: From 182fb85240dddfe7351ff06d1659652d81390c2e Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 30 Jul 2026 07:21:36 -0400 Subject: [PATCH 16/85] Keep the router's order when a request is deduplicated Replaying a real trace through the cache on the target exposed a fidelity defect in the simulator that produced every projection in this work. The cache reported 5107 misses over 64 tokens; the simulator predicted 5120. Within one request, insertion order decides which entry is oldest, so it changes what the next eviction picks. The cache dedupes with an order-preserving map and therefore follows the router's ordering; the simulator iterated a set, whose order is unrelated. Twenty-two of the forty layers differed, some above and some below, netting 0.25 %. With order preserved the simulator returns 5107, matching the measurement exactly. No conclusion moves by 0.25 %, but the tool the conclusions came from should describe the cache it is modelling. The test pins the behaviour on a case where the choice is observable: at quota 2, requesting [7, 3] then 9 must evict 7 rather than 3, so the following request for 3 hits. --- qwen36_moe_edge/route_trace.py | 8 ++++++-- tests/test_qwen36_moe_edge_quant.py | 15 +++++++++++++++ 2 files changed, 21 insertions(+), 2 deletions(-) diff --git a/qwen36_moe_edge/route_trace.py b/qwen36_moe_edge/route_trace.py index 6e947381..9d2658d2 100644 --- a/qwen36_moe_edge/route_trace.py +++ b/qwen36_moe_edge/route_trace.py @@ -177,11 +177,15 @@ def simulate_warm_lru( cache[expert] = None steps = layer_trace[prompt_tokens:] for start in range(0, len(steps) - window + 1, window): - requested = { + # Order-preserving dedupe, matching what the cache does. Within one + # request the insertion order decides which entry becomes the + # oldest, so it changes what a later eviction picks: iterating a set + # instead put this 0.25 % away from the measured cache. + requested = dict.fromkeys( expert for step in steps[start:start + window] for expert in step - } + ) if index == 0: tokens += window for expert in requested: diff --git a/tests/test_qwen36_moe_edge_quant.py b/tests/test_qwen36_moe_edge_quant.py index 355b33c9..802af26f 100644 --- a/tests/test_qwen36_moe_edge_quant.py +++ b/tests/test_qwen36_moe_edge_quant.py @@ -213,6 +213,21 @@ def test_windowing_does_not_reduce_reads_an_lru_already_serves(): "decode_misses_per_token"] +def test_requests_keep_the_router_order_so_eviction_matches_the_cache(): + # Within one request the insertion order decides which entry is oldest, so + # it changes what the next eviction picks. The cache dedupes preserving the + # router's order; a simulator that iterated a set instead would model a + # different cache. Quota 2 with three distinct experts makes the choice + # observable: after [7, 3] the oldest is 7, so requesting 9 must evict 7 and + # leave 3 -- and the following request for 3 must then hit. + trace = [[[7, 3], [9], [3]]] + + result = simulate_warm_lru(trace, prompt_tokens=0, quota=2) + + assert result["decode_misses_per_token"] == 1.0 # 7, 3, 9 miss; 3 hits + assert result["distinct_hit_rate"] == 0.25 + + def test_cold_prefill_counts_the_union_prefill_touches(): # Prefill routes each token independently, so a layer costs the union of # its tokens' selections, less whatever is already resident. From 2e70e8255a06c3caab25500a4f35e2993c574199 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 30 Jul 2026 07:28:30 -0400 Subject: [PATCH 17/85] Let the validation persist its traces The leave-one-out figures were simulated. Replaying them on real hardware needs the traces themselves, so a warm set can be built from some prompts and measured against another one it has never seen. --- qwen36_moe_edge/warm_start_validation.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/qwen36_moe_edge/warm_start_validation.py b/qwen36_moe_edge/warm_start_validation.py index 42a7ee07..8ed82c9c 100644 --- a/qwen36_moe_edge/warm_start_validation.py +++ b/qwen36_moe_edge/warm_start_validation.py @@ -143,6 +143,10 @@ def main() -> None: parser.add_argument("--max-seq", type=int, default=128) parser.add_argument("--quotas", default="43,57,64") parser.add_argument("--block-bytes", type=int, default=DEFAULT_BLOCK_BYTES) + parser.add_argument( + "--save-traces", type=Path, + help="write the per-prompt traces, so a warm set built from some of " + "them can be replayed against another on real hardware") parser.add_argument("--device", default="cuda:0") args = parser.parse_args() @@ -162,6 +166,14 @@ def main() -> None: device=args.device, ) + if args.save_traces is not None: + args.save_traces.parent.mkdir(parents=True, exist_ok=True) + with args.save_traces.open("w", encoding="utf-8") as f: + json.dump({"prompt_tokens": args.prompt_tokens, + "traces": traces}, f) + f.write("\n") + print(f"wrote {len(traces)} traces to {args.save_traces}") + quotas = tuple(int(value) for value in args.quotas.split(",")) report = {"prompt_count": len(prompts), "quotas": list(quotas), "prompt_tokens": args.prompt_tokens, "by_quota": {}} From a730abcd3451031c70b956a700dd7d93fb13864e Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 30 Jul 2026 07:32:15 -0400 Subject: [PATCH 18/85] Expose a block's components and its global scales Wiring the cache into a decode path needs the four parts of a block and the two scales that go with it. Doing that at the call site would mean reproducing the writer's offset arithmetic, which is how a reader drifts from a writer; these read the order from the manifest instead, and return views over the slot rather than copies. The global scales come from the per-layer sidecar, which exists so a block stays exactly block_bytes and 4096-aligned. Their file is size-checked, because a truncated sidecar would otherwise reshape into plausible nonsense. With these, the whole chain was measured on Thor against the reference pairs recorded on the SM120 machine -- a real routed activation and that expert's BF16 output. Pulling each expert's block through the cache, decoding it, and running the expert lands at cosine mean 0.974500 and relative L2 0.18510, against 0.974504 and 0.18508 for the same experts quantized straight from the checkpoint. Agreement to the fifth decimal across the on-disk format, the two-level scales, the sidecar, the direct read, the slot layout, the component split, and the transform on both GEMMs. The low minimum cosine, 0.574, falls on the same three near-zero-output experts as the checkpoint-direct run and at the same values, so it remains an artefact of an unpooled per-expert metric rather than anything the chain introduced. --- qwen36_moe_edge/expert_cache.py | 46 +++++++++++++++++++++++++++ tests/test_qwen36_moe_expert_cache.py | 22 +++++++++++++ 2 files changed, 68 insertions(+) diff --git a/qwen36_moe_edge/expert_cache.py b/qwen36_moe_edge/expert_cache.py index 98f64df6..5ec8a2d5 100644 --- a/qwen36_moe_edge/expert_cache.py +++ b/qwen36_moe_edge/expert_cache.py @@ -137,6 +137,7 @@ def __init__(self, config: CacheConfig): for layer in range(self.num_layers) ] self._fds: dict[int, int] = {} + self._global_scales: dict[int, torch.Tensor] = {} self.hits = 0 self.misses = 0 self.bytes_read = 0 @@ -259,6 +260,51 @@ def get_many(self, layer: int, experts) -> list[int]: def get(self, layer: int, expert: int) -> int: return self.get_many(layer, (expert,))[0] + def components(self, layer: int, expert: int) -> dict[str, torch.Tensor]: + """The block's four parts as views over its slot, plus its scales. + + Views, not copies: the caller reads them where the block already lies. + The manifest's ``block_layout`` gives the order, so a consumer never + reproduces the offset arithmetic and cannot drift from the writer. + """ + self.get(layer, expert) + raw = self.slots[self._lru[layer][expert]] + sizes = self.manifest["block_sizes"] + offset = 0 + parts = {} + for name in self.manifest["block_layout"]: + length = int(sizes[name]) + if name != "padding": + parts[name] = raw[offset:offset + length] + offset += length + parts["global_scales"] = self.global_scales(layer)[expert] + return parts + + def global_scales(self, layer: int) -> torch.Tensor: + """This layer's per-expert (gate_up, down) scales, read once. + + They live beside the blocks rather than inside them so a block stays + exactly ``block_bytes`` and aligned; the kernel takes them as its GEMM + alpha. + """ + cached = self._global_scales.get(layer) + if cached is None: + name = self.manifest.get( + "global_scales", "global_scales_layer_NN.bin") + path = self.config.bundle / name.replace( + "NN", f"{layer:02d}") + expected = self.num_experts * 2 * 4 + size = path.stat().st_size + if size != expected: + raise ValueError( + f"{path} is {size} bytes; expected {expected} " + f"({self.num_experts} experts x 2 x float32)") + cached = torch.frombuffer( + bytearray(path.read_bytes()), dtype=torch.float32 + ).view(self.num_experts, 2) + self._global_scales[layer] = cached + return cached + # ── startup ── def warm(self, frequency: list[Counter]) -> int: diff --git a/tests/test_qwen36_moe_expert_cache.py b/tests/test_qwen36_moe_expert_cache.py index 1b1fc8e5..a96e46a0 100644 --- a/tests/test_qwen36_moe_expert_cache.py +++ b/tests/test_qwen36_moe_expert_cache.py @@ -138,3 +138,25 @@ def test_close_is_documented_to_release_the_slots(): assert "self.slots = None" in source assert "self._staging = []" in source + + +def test_components_follows_the_manifest_and_omits_the_pad(): + # A consumer must not reproduce the offset arithmetic; it reads the order + # from the manifest, so it cannot drift from whatever wrote the bundle. + import inspect + + source = inspect.getsource(ExpertCache.components) + + assert 'self.manifest["block_layout"]' in source + assert 'name != "padding"' in source + + +def test_global_scales_validates_the_sidecar_size(): + # The scales live beside the blocks so a block stays exactly block_bytes + # and aligned. A truncated sidecar has to be caught, not silently reshaped. + import inspect + + source = inspect.getsource(ExpertCache.global_scales) + + assert "num_experts * 2 * 4" in source + assert "raise ValueError" in source From ed77aef83f311ca6b19c6551a6ba3fcac005c6b6 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 30 Jul 2026 07:37:15 -0400 Subject: [PATCH 19/85] Decode a streamed expert block into bf16 The streaming path needs to get a block from a cache slot into a GEMM. The block-scaled 4-bit GEMMs cannot read it: they decode E2M1, and the bundle stores sign-magnitude integers, whose sixteen values are a different set, so no relabelling bridges them. They also want the scale bytes in the SM1xx swizzled tile layout, where the bundle's are linear. Decoding to bf16 first and handing that to the existing bf16 GEMM sidesteps both. It reads the bundle's own linear scale layout, which removes the swizzle step rather than implementing it, and it needs no second codebook inside a GEMM and no architecture beyond SM80. The cost is bandwidth on a block that is already resident, which is the cheap end of this system: the misses are what cost time. One thread per packed byte. Its two values always share a scale because the group size is even, so there is one e4m3 conversion per byte rather than two. The two-level scale is applied as the per-tensor float times the group's byte, matching how the quantizer chose them. It goes in the core tier: bit manipulation and bf16 writes, so it compiles and runs wherever that tier does. Verified bit-identical to the Python reference for (1024, 2048) and (2048, 512) at group 16 and 32, and it rejects a null pointer, odd columns, an odd group and a non-dividing group with distinct codes. Added to the cross-architecture parity harness, which now covers nine cases and thirteen output tensors, all bitwise identical between sm_120a and sm_110. That also puts it in the acceptance package's kernel test. --- CMakeLists.txt | 3 +- csrc/bindings.cpp | 13 ++++ csrc/kernels/qwen35moe_e0m3_dequant.cu | 79 +++++++++++++++++++++++++ csrc/kernels/qwen35moe_e0m3_dequant.cuh | 40 +++++++++++++ qwen36_moe_edge/kernel_parity.py | 21 +++++++ tests/test_qwen36_moe_edge_quant.py | 1 + 6 files changed, 156 insertions(+), 1 deletion(-) create mode 100644 csrc/kernels/qwen35moe_e0m3_dequant.cu create mode 100644 csrc/kernels/qwen35moe_e0m3_dequant.cuh diff --git a/CMakeLists.txt b/CMakeLists.txt index 53e509c8..8094d5cb 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1511,7 +1511,8 @@ if(FLASHRT_ENABLE_QWEN35MOE_CORE) csrc/kernels/act_fuse_sm120.cu csrc/kernels/moe_router_topk_sm120.cu csrc/kernels/moe_weighted_sum_sm120.cu - csrc/kernels/w16a16_gemm_sm120.cu) + csrc/kernels/w16a16_gemm_sm120.cu + csrc/kernels/qwen35moe_e0m3_dequant.cu) target_compile_definitions(flash_rt_kernels PRIVATE FLASHRT_HAVE_QWEN35MOE_CORE=1) message(STATUS "qwen3_5_moe core kernels: ENABLED (sm_${GPU_ARCH})") diff --git a/csrc/bindings.cpp b/csrc/bindings.cpp index f9e191a3..0a108a61 100644 --- a/csrc/bindings.cpp +++ b/csrc/bindings.cpp @@ -186,6 +186,7 @@ extern "C" int cutlass_int8_rowwise_bf16out_t64x128( #include "kernels/moe_router_topk_sm120.cuh" #include "kernels/moe_weighted_sum_sm120.cuh" #include "kernels/w16a16_gemm_sm120.cuh" +#include "kernels/qwen35moe_e0m3_dequant.cuh" #endif // FLASHRT_HAVE_QWEN35MOE_CORE #ifdef FLASHRT_HAVE_QWEN35MOE_W4A16 #include "kernels/w4a16_matvec_sm120.cuh" @@ -5471,6 +5472,18 @@ PYBIND11_MODULE(flash_rt_kernels, m) { py::arg("x"), py::arg("gate"), py::arg("out"), py::arg("n"), py::arg("stream") = 0); + m.def("qwen35moe_e0m3_dequant_bf16", + [](uintptr_t packed, uintptr_t scale, uintptr_t out, + int rows, int cols, int group_size, float global_scale, + uintptr_t stream) -> int { + return flash_rt::kernels::qwen35moe_e0m3_dequant_bf16( + to_ptr(packed), to_ptr(scale), to_ptr(out), + rows, cols, group_size, global_scale, to_stream(stream)); + }, + py::arg("packed"), py::arg("scale"), py::arg("out"), + py::arg("rows"), py::arg("cols"), py::arg("group_size"), + py::arg("global_scale"), py::arg("stream") = 0); + m.def("gdn_recurrent_seq_sm120_bf16", [](uintptr_t q, uintptr_t k, uintptr_t v, uintptr_t g, uintptr_t beta, uintptr_t state, uintptr_t out, int S, int num_v_heads, diff --git a/csrc/kernels/qwen35moe_e0m3_dequant.cu b/csrc/kernels/qwen35moe_e0m3_dequant.cu new file mode 100644 index 00000000..43c63f9b --- /dev/null +++ b/csrc/kernels/qwen35moe_e0m3_dequant.cu @@ -0,0 +1,79 @@ +// SPDX-License-Identifier: Apache-2.0 +// +// Streamed routed-expert block to bf16. See header. + +#include "kernels/qwen35moe_e0m3_dequant.cuh" + +#include +#include +#include +#include + +namespace flash_rt { +namespace kernels { + +namespace { + +constexpr int kThreads = 256; + +// Sign-magnitude: low three bits are the magnitude, bit 3 the sign. Kept as a +// signed integer because the values are exactly the integers 0..7, which is the +// property the quantizer's group scale is chosen against. +__device__ __forceinline__ float decode_nibble(uint8_t code) { + const float magnitude = static_cast(code & 0x07u); + return (code & 0x08u) ? -magnitude : magnitude; +} + +// One thread per packed byte: two output values that always share a scale, +// because group_size is even. +__global__ void dequant_kernel(const uint8_t* __restrict__ packed, + const uint8_t* __restrict__ scale, + __nv_bfloat162* __restrict__ out, + int rows, int cols, int group_size, + float global_scale) { + const int index = blockIdx.x * blockDim.x + threadIdx.x; + const int pairs_per_row = cols >> 1; + if (index >= rows * pairs_per_row) return; + + const int row = index / pairs_per_row; + const int pair = index - row * pairs_per_row; + const int group = (pair << 1) / group_size; + const int groups_per_row = cols / group_size; + + const __half_raw raw = __nv_cvt_fp8_to_halfraw( + scale[row * groups_per_row + group], __NV_E4M3); + const float step = + __half2float(*reinterpret_cast(&raw)) * global_scale; + + const uint8_t byte = packed[index]; + out[index] = __floats2bfloat162_rn( + decode_nibble(byte & 0x0Fu) * step, + decode_nibble(byte >> 4) * step); +} + +} // namespace + +int qwen35moe_e0m3_dequant_bf16(const void* packed, const void* scale, + void* out, int rows, int cols, + int group_size, float global_scale, + cudaStream_t stream) { + if (!packed || !scale || !out) return 1; + if (rows <= 0 || cols <= 0) return 2; + if (cols & 1) return 3; + if (group_size <= 0 || (group_size & 1)) return 4; + if (cols % group_size) return 5; + + const long long pairs = static_cast(rows) * (cols >> 1); + const long long blocks = (pairs + kThreads - 1) / kThreads; + if (blocks > 2147483647LL) return 6; + + dequant_kernel<<(blocks), kThreads, 0, stream>>>( + reinterpret_cast(packed), + reinterpret_cast(scale), + reinterpret_cast<__nv_bfloat162*>(out), + rows, cols, group_size, global_scale); + return 0; +} + +} // namespace kernels +} // namespace flash_rt diff --git a/csrc/kernels/qwen35moe_e0m3_dequant.cuh b/csrc/kernels/qwen35moe_e0m3_dequant.cuh new file mode 100644 index 00000000..580564eb --- /dev/null +++ b/csrc/kernels/qwen35moe_e0m3_dequant.cuh @@ -0,0 +1,40 @@ +// SPDX-License-Identifier: Apache-2.0 +// +// Decode a streamed routed-expert block into bf16. +// +// The edge bundle stores each expert as sign-magnitude 4-bit values with one +// e4m3 scale byte per group of 16 along K, scaled by a per-tensor float the +// bundle keeps beside the blocks. That is not the format the block-scaled 4-bit +// GEMMs read: they decode E2M1, and they want the scale bytes in the SM1xx +// swizzled tile layout. Neither difference is bridgeable by relabelling -- +// E2M1's sixteen values and this format's sixteen are different sets. +// +// So the streaming path decodes here and hands bf16 to the existing bf16 GEMM, +// which costs bandwidth on an already-resident block but needs no swizzle, no +// second codebook inside a GEMM, and no architecture beyond SM80. Reading the +// bundle's own linear scale layout is the point: it removes the swizzle step +// rather than implementing it. + +#pragma once + +#include + +namespace flash_rt { +namespace kernels { + +// out[r][c] = value(packed) * e4m3(scale[r][c / group_size]) * global_scale +// +// packed (rows, cols / 2) bytes, low nibble first, each nibble +// magnitude | sign << 3 with magnitude in 0..7 +// scale (rows, cols / group_size) bytes, each an e4m3 magnitude +// out (rows, cols) bf16 +// +// cols must be even and a multiple of group_size, and group_size must be even +// so that a byte's two values always share one scale. +int qwen35moe_e0m3_dequant_bf16(const void* packed, const void* scale, + void* out, int rows, int cols, + int group_size, float global_scale, + cudaStream_t stream); + +} // namespace kernels +} // namespace flash_rt diff --git a/qwen36_moe_edge/kernel_parity.py b/qwen36_moe_edge/kernel_parity.py index b337a609..0673efc7 100644 --- a/qwen36_moe_edge/kernel_parity.py +++ b/qwen36_moe_edge/kernel_parity.py @@ -151,6 +151,25 @@ def case_lin_split_qkv(fvk, device): return {"q32": q32, "k32": k32, "v32": v32} +def case_e0m3_dequant(fvk, device): + """The streamed-block decode: sign-magnitude 4-bit plus a two-level scale.""" + from qwen36_moe_edge.quantize_experts import _int4_weight, dequantize_int4 + + rows, cols, group = 2 * INTERMEDIATE, HIDDEN, 16 + weight = _bf16("weight", (rows, cols), "cpu", 0.02).float() + packed, scale, global_scale = _int4_weight(weight, group) + packed = _record("packed", packed, device).contiguous() + scale = _record("scale", scale, device).contiguous() + out = torch.zeros(rows, cols, dtype=torch.bfloat16, device=device) + rc = fvk.qwen35moe_e0m3_dequant_bf16( + packed.data_ptr(), scale.data_ptr(), out.data_ptr(), + rows, cols, group, float(global_scale), 0) + torch.cuda.synchronize() + reference = dequantize_int4( + packed.cpu(), scale.cpu(), cols, group, global_scale) + return {"rc": rc, "out": out, "torch": reference.to(device)} + + def case_split_q_gate(fvk, device): S = 4 q_proj = _bf16("q_proj", (S, 8192), device).contiguous() @@ -171,6 +190,7 @@ def case_split_q_gate(fvk, device): "w16a16_gemm": case_w16a16_gemm, "lin_split_qkv": case_lin_split_qkv, "split_q_gate": case_split_q_gate, + "e0m3_dequant": case_e0m3_dequant, } @@ -258,6 +278,7 @@ def _binding_of(name: str) -> str: return { "lin_split_qkv": "qwen35moe_lin_split_qkv_broadcast_bf16", "split_q_gate": "qwen35moe_split_q_gate_bf16", + "e0m3_dequant": "qwen35moe_e0m3_dequant_bf16", }.get(name, f"{name}_sm120_bf16") diff --git a/tests/test_qwen36_moe_edge_quant.py b/tests/test_qwen36_moe_edge_quant.py index 802af26f..12b9cf9b 100644 --- a/tests/test_qwen36_moe_edge_quant.py +++ b/tests/test_qwen36_moe_edge_quant.py @@ -255,6 +255,7 @@ def test_parity_cases_name_bindings_that_exist_in_the_tiers(): "w16a16_gemm": "w16a16_gemm_sm120_bf16", "lin_split_qkv": "qwen35moe_lin_split_qkv_broadcast_bf16", "split_q_gate": "qwen35moe_split_q_gate_bf16", + "e0m3_dequant": "qwen35moe_e0m3_dequant_bf16", } assert set(CASES) == set(expected) From d50ad7327d878e05826392afec5be6216a8ea744 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 30 Jul 2026 07:42:00 -0400 Subject: [PATCH 20/85] Serve the routed experts from storage Assembles the pieces into a path where the experts are not resident. Three additive changes, each off unless asked for. The loader gains stream_experts, which skips building the per-layer stacked expert tensors. That skip is the point rather than an optimisation: those tensors are 16.9 GiB of a 21.4 GiB footprint, and attaching a cache without removing them would add to the total instead of replacing part of it. Their shapes are still checked, since a bundle is generated against them. The decode path gains a branch, reached only when the loader took that skip. It fetches a token's whole top-k in one call so the reads overlap, then decodes each block to bf16 and multiplies with the shared bf16 GEMV. The per-layer quota is at least the top-k, so no returned pointer can be invalidated by the others and the eight can be held at once. Decode scratch is two buffers allocated once, rather than eight allocations per layer per token. Qwen36MoeStreamingFrontend wires the two together. It sizes the cache against the resident bytes it actually measured after loading, not an estimate, and stays eager: a miss issues host reads, which a captured graph cannot replay. Everything else is unchanged -- same attention, same recurrence, same router, same reducer -- so comparing its tokens against the ordinary frontend isolates where the expert weights came from. --- flash_rt/frontends/torch/_nexn2_rtx_decode.py | 89 +++++++++++++++++++ .../torch/_nexn2_rtx_nvfp4_weights.py | 27 +++++- flash_rt/frontends/torch/nexn2_rtx.py | 3 + qwen36_moe_edge/expert_cache.py | 3 + qwen36_moe_edge/streaming_frontend.py | 89 +++++++++++++++++++ 5 files changed, 209 insertions(+), 2 deletions(-) create mode 100644 qwen36_moe_edge/streaming_frontend.py diff --git a/flash_rt/frontends/torch/_nexn2_rtx_decode.py b/flash_rt/frontends/torch/_nexn2_rtx_decode.py index a3e14365..4981c852 100644 --- a/flash_rt/frontends/torch/_nexn2_rtx_decode.py +++ b/flash_rt/frontends/torch/_nexn2_rtx_decode.py @@ -279,6 +279,25 @@ def __init__(self, handles, max_seq, device): self.router_trace = None self.moe_input_trace = None self._active_layer = -1 + # Set to an ExpertCache to read the routed experts from storage. Only + # meaningful when the loader skipped them; see _moe_experts_streamed. + self.expert_cache = None + self._scratch = None + + def _streamed_scratch(self, device): + """The two decode buffers a streamed expert is unpacked into. + + Allocated once and reused: 4 MiB for gate_up and 2 MiB for down, which + would otherwise be allocated 8 times per layer per token. + """ + if self._scratch is None: + self._scratch = { + 'gate_up': torch.empty( + 2 * INTER, HID, dtype=torch.bfloat16, device=device), + 'down': torch.empty( + HID, INTER, dtype=torch.bfloat16, device=device), + } + return self._scratch def reset(self): for s in self.lin_state: @@ -413,6 +432,72 @@ def _decode_full(h, ld, state, full_rank, pos, fvk, device): 1, 1, HID) +def _moe_experts_streamed(x, idx, tw_row, ld, state, fvk, device, s): + """The routed experts, read from storage instead of held in memory. + + Only reachable when the loader was told to stream them, in which case the + per-layer stacked NVFP4 tensors were never allocated -- that is the whole + point, and adding a cache without skipping them would cost memory rather + than save it. + + Each block is decoded to bf16 and multiplied with the shared bf16 GEMM. The + block-scaled 4-bit GEMMs cannot read these blocks: different codebook, + different scale layout. Decoding costs bandwidth on a block that is already + resident, which is the cheap end of this system -- the misses are what cost + time. + + The whole top-k is fetched in one call so the reads overlap, and the + per-layer quota is at least the top-k, so none of the returned pointers can + be invalidated by the others. + """ + cache = state.expert_cache + layer = state._active_layer + experts = [int(value) for value in idx.cpu().tolist()] + cache.get_many(layer, experts) + + scratch = state._streamed_scratch(device) + d_gu = torch.empty(TOPK, 2 * INTER, dtype=torch.bfloat16, device=device) + d_dn = torch.empty(TOPK, HID, dtype=torch.bfloat16, device=device) + xc = x.contiguous() + + for slot, expert in enumerate(experts): + parts = cache.components(layer, expert) + gu_alpha, dn_alpha = parts['global_scales'].tolist() + rc = fvk.qwen35moe_e0m3_dequant_bf16( + parts['gate_up_weight'].data_ptr(), + parts['gate_up_scale'].data_ptr(), + scratch['gate_up'].data_ptr(), + 2 * INTER, HID, cache.group_size, gu_alpha, s) + if rc: + raise RuntimeError(f'gate_up decode failed with {rc}') + fvk.bf16_matvec_sm120_bf16( + xc.data_ptr(), scratch['gate_up'].data_ptr(), + d_gu[slot].data_ptr(), 2 * INTER, HID, s) + + gated = _silu_mul( + d_gu[slot:slot + 1, :INTER], d_gu[slot:slot + 1, INTER:], + fvk, device).contiguous() + rc = fvk.qwen35moe_e0m3_dequant_bf16( + parts['down_weight'].data_ptr(), + parts['down_scale'].data_ptr(), + scratch['down'].data_ptr(), + HID, INTER, cache.group_size, dn_alpha, s) + if rc: + raise RuntimeError(f'down decode failed with {rc}') + fvk.bf16_matvec_sm120_bf16( + gated.data_ptr(), scratch['down'].data_ptr(), + d_dn[slot].data_ptr(), HID, INTER, s) + + if 'decode_topk_rows' not in ld: + ld['decode_topk_rows'] = torch.arange( + TOPK, dtype=torch.int32, device=device) + out = torch.empty(HID, dtype=torch.float32, device=device) + fvk.moe_weighted_sum_sm120_bf16( + d_dn.data_ptr(), ld['decode_topk_rows'].data_ptr(), + tw_row.data_ptr(), out.data_ptr(), 1, TOPK, HID, HID, s) + return out.unsqueeze(0) + + def _moe_layer_decode(h, ld, state, fvk, device): """M=1 fine-grained MoE via the grouped GEMV kernel: the 8 routed experts run in one launch each for gate_up (shared act) and down (per-slot act), @@ -456,6 +541,10 @@ def _moe_layer_decode(h, ld, state, fvk, device): state.moe_input_trace[state._active_layer].append( x.detach().to("cpu", copy=True)) + if ld.get('experts_streamed'): + return _moe_experts_streamed( + x, idx, tw_row, ld, state, fvk, device, s) + if 'experts_gate_up_alpha_dev' not in ld: # cache once/layer ld['experts_gate_up_alpha_dev'] = \ ld['experts_gate_up_alpha_t'].to(device).contiguous() diff --git a/flash_rt/frontends/torch/_nexn2_rtx_nvfp4_weights.py b/flash_rt/frontends/torch/_nexn2_rtx_nvfp4_weights.py index 985162aa..1427871e 100644 --- a/flash_rt/frontends/torch/_nexn2_rtx_nvfp4_weights.py +++ b/flash_rt/frontends/torch/_nexn2_rtx_nvfp4_weights.py @@ -168,7 +168,8 @@ def _bf16_from_ckpt(handles, out_dict, name, key, handles_d, wmap, device, def _load_moe(handles, ld, lp, handles_d, wmap, fvk, device, - n_experts: int, *, quantize_shared: bool = True) -> None: + n_experts: int, *, quantize_shared: bool = True, + stream_experts: bool = False) -> None: """Load one layer's MoE block: router (BF16) + experts + shared expert.""" # Router gate (BF16) and shared-expert sigmoid gate (BF16). _bf16_from_ckpt(handles, ld, 'router_w', lp + 'mlp.gate.weight', @@ -189,6 +190,20 @@ def _load_moe(handles, ld, lp, handles_d, wmap, fvk, device, # Routed experts: packed 3D tensors (E, out, in). Quantize each expert # into a contiguous slice of a per-layer stacked NVFP4 buffer so the # downstream grouped GEMM sees one contiguous weight per projection. + if stream_experts: + # The routed experts are read from storage at decode time, so the + # stacked per-layer tensors are never built. Skipping them is the + # point: they are 16.9 GiB of the resident footprint, and a cache + # added on top of them would cost memory rather than save it. The + # shapes are still checked, because a bundle is generated against them. + for name in ('mlp.experts.gate_up_proj', 'mlp.experts.down_proj'): + if not _has(wmap, lp + name): + raise ValueError( + f'{lp}{name} is absent, so a streamed expert bundle ' + 'cannot correspond to this checkpoint') + ld['experts_streamed'] = True + return + gate_up = _get(handles_d, wmap, lp + 'mlp.experts.gate_up_proj') down = _get(handles_d, wmap, lp + 'mlp.experts.down_proj') e_gu, n_gu, k_gu = gate_up.shape # (E, 2*inter, hidden) @@ -244,6 +259,7 @@ def extract_weights_nexn2_nvfp4( fvk, device: str = 'cuda:0', quant_scope: str = 'experts', + stream_experts: bool = False, ) -> WeightHandles: """Build :class:`WeightHandles` from a Nex-N2-mini BF16 ckpt directory. @@ -253,6 +269,12 @@ def extract_weights_nexn2_nvfp4( * ``'experts'``: only the storage-dominant routed experts go NVFP4; full-attn / out_proj / shared stay BF16. ~21 GB; E2E cos ~0.99 -- the precision-per-VRAM baseline until the Step-3 W4A16 mixed kernel. + + stream_experts: skip the routed experts entirely, leaving the decode path + to read them from a prepared bundle. They are 16.9 GiB of the resident + footprint, so this is the difference between a model that fits a small + device and one that does not; a cache added without skipping them would + only add to the total. The decode path must then be given an ExpertCache. """ if quant_scope not in ('full', 'experts'): raise ValueError( @@ -335,7 +357,8 @@ def extract_weights_nexn2_nvfp4( # Every layer has a MoE FFN (mlp_only_layers is empty). _load_moe(handles, ld, lp, handles_d, wmap, fvk, device, n_experts, - quantize_shared=quant_main) + quantize_shared=quant_main, + stream_experts=stream_experts) per_layer[i] = ld handles.ptrs['layers'] = per_layer diff --git a/flash_rt/frontends/torch/nexn2_rtx.py b/flash_rt/frontends/torch/nexn2_rtx.py index b1791552..3c8b9deb 100644 --- a/flash_rt/frontends/torch/nexn2_rtx.py +++ b/flash_rt/frontends/torch/nexn2_rtx.py @@ -108,6 +108,9 @@ def __init__(self, checkpoint_path: str, *, self._quant_format = quant self._kernelized = bool(kernelized) self._quant_scope = quant_scope + # Set by a subclass that streams the routed experts from a bundle + # instead of holding them; see _nexn2_rtx_decode._moe_experts_streamed. + self._stream_experts = getattr(self, '_stream_experts', False) self._tokenizer = None self._prompt_ids = None self._pipeline: Nexn2Pipeline | None = None diff --git a/qwen36_moe_edge/expert_cache.py b/qwen36_moe_edge/expert_cache.py index 5ec8a2d5..81ff94ca 100644 --- a/qwen36_moe_edge/expert_cache.py +++ b/qwen36_moe_edge/expert_cache.py @@ -90,6 +90,9 @@ def __init__(self, config: CacheConfig): self.alignment = int(self.manifest["block_alignment"]) self.num_layers = int(self.manifest["num_layers"]) self.num_experts = int(self.manifest["num_experts"]) + # None in the manifest means INT8, whose scales are per output channel + # and need no group. + self.group_size = int(self.manifest.get("group_size") or 0) if config.slots_per_layer < config.experts_per_token: raise ValueError( diff --git a/qwen36_moe_edge/streaming_frontend.py b/qwen36_moe_edge/streaming_frontend.py new file mode 100644 index 00000000..6fd33756 --- /dev/null +++ b/qwen36_moe_edge/streaming_frontend.py @@ -0,0 +1,89 @@ +#!/usr/bin/env python3 +"""Qwen3.6-35B-A3B with the routed experts read from storage. + +The shipped frontend holds every expert in memory, which is 16.9 GiB of a +21.4 GiB footprint. This one skips them at load and serves each token's top-k +from a bounded cache over a prepared bundle, so what stays resident is the +non-expert weights plus however many slots the budget affords. + +It is the same pipeline otherwise: same attention, same recurrence, same +router, same reducer. Only where the expert weights come from changes, which is +why a token-level comparison against the ordinary frontend is meaningful. + +Greedy decode only, and not for CUDA Graph capture: a miss issues host reads, +which a captured graph cannot replay. +""" + +from __future__ import annotations + +from pathlib import Path + +from flash_rt.frontends.torch.qwen36_moe_rtx import Qwen36MoeTextFrontendRtx + +from qwen36_moe_edge.expert_cache import CacheConfig, ExpertCache + + +class Qwen36MoeStreamingFrontend(Qwen36MoeTextFrontendRtx): + """Routed experts streamed from a bundle rather than held in memory.""" + + _MODEL_LABEL = "Qwen3.6-35B-A3B text, streamed experts" + + def __init__(self, checkpoint_path: str, bundle: str | Path, *, + slots_per_layer: int, + device: str = "cuda:0", + max_seq: int = 2048, + staging_buffers: int = 4, + budget_bytes: int = 0, + reserve_bytes: int = 0, + warm_frequency=None) -> None: + # Read by the loader through the base class, before any weight is + # touched, so the expert tensors are never built. + self._stream_experts = True + super().__init__( + checkpoint_path, device=device, max_seq=max_seq, + quant_scope="experts") + + resident = 0 + try: + import torch + + resident = int(torch.cuda.memory_allocated(device)) + except Exception: # noqa: BLE001 + pass + self.cache = ExpertCache(CacheConfig( + bundle=Path(bundle), + slots_per_layer=slots_per_layer, + staging_buffers=staging_buffers, + budget_bytes=budget_bytes, + reserve_bytes=reserve_bytes, + # Measured, not assumed: what the weights actually took. + resident_bytes=resident, + device=device, + )) + if warm_frequency is not None: + self.cache.warm(warm_frequency) + + def generate(self, max_new_tokens: int, *, do_sample: bool = False): + if self._prompt_ids is None: + raise ValueError("call set_prompt(...) before generate()") + if do_sample: + raise NotImplementedError("greedy decoding only") + + from flash_rt.frontends.torch._nexn2_rtx_decode import ( + Nexn2DecodeState, + generate_greedy, + ) + + if self._decode_state is None: + self._decode_state = Nexn2DecodeState( + self._weights, self._user_max_seq, self.device) + state = self._decode_state + state.expert_cache = self.cache + # A miss reads from storage on the host, which a captured graph cannot + # replay, so this path stays eager. + state.batched_prefill = False + return generate_greedy( + state, self._prompt_ids, max_new_tokens, self._fvk, self.device) + + def close(self) -> None: + self.cache.close() From 9f9fe5a0ae933a294f33f63cc6dfe8e6b40f6c67 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 30 Jul 2026 07:48:47 -0400 Subject: [PATCH 21/85] Load the tokenizer only when something asks for it The frontend loaded a tokenizer while building weights, which made transformers a hard requirement of the runtime. It is not one: a caller with token ids of its own never needs it, and requiring it on a deployment target means installing a large dependency into an environment that may not want it. set_prompt_ids takes ids directly, so a target can run with the prompt tokenized elsewhere. The tokenizer property still loads on first use, so set_prompt and decode are unchanged. Also records what the pipeline actually depends on. The three tiers cover 14 of the 32 kernels it calls; resolving each call to the guard active where it is defined shows seven gates, with twelve kernels under FLASHRT_HAVE_QWEN36_KERNELS carrying the whole linear-attention path. That gate keys on NOT FLASHRT_SLIM_BUILD rather than on architecture, so a slim build removes them and the frontend refuses to start. Selecting tiers from the source's grouping was not enough to know what a target needs -- the call sites are. --- docs/qwen36_moe_usage.md | 25 +++++++++++++++++++++ flash_rt/frontends/torch/nexn2_rtx.py | 32 ++++++++++++++++++++++----- 2 files changed, 52 insertions(+), 5 deletions(-) diff --git a/docs/qwen36_moe_usage.md b/docs/qwen36_moe_usage.md index 1fa7b40c..3f4a6200 100644 --- a/docs/qwen36_moe_usage.md +++ b/docs/qwen36_moe_usage.md @@ -45,6 +45,31 @@ below. Targets that cannot run a tier can select the remainder explicitly. The upper tiers depend on the core tier, so enabling either turns it on. The SM120 text runtime documented here needs all three. +**These three tiers are not the whole dependency set.** Walking every `fvk` +call the pipeline makes and resolving each to the preprocessor guard active +where it is defined gives 32 kernels across seven gates: + +| gate | kernels | +|---|---:| +| `FLASHRT_HAVE_QWEN36_KERNELS` | 12 | +| `FLASHRT_HAVE_QWEN35MOE_CORE` | 10 | +| `FLASHRT_HAVE_QWEN35MOE_W4A16` | 3 | +| `ENABLE_CUTLASS_SM120_NVFP4_W4A16` | 2 | +| `FLASHRT_HAVE_QWEN35MOE_W4A4` | 2 | +| `FLASHRT_HAVE_NVFP4_SWIZZLE` | 1 | +| ungated | 1 | + +The twelve under `FLASHRT_HAVE_QWEN36_KERNELS` are the linear-attention path: +causal convolution and its update, the gated-DeltaNet recurrence, the WY chunk +stack, the fused RMSNorm-gated-SiLU, partial RoPE and argmax. They are shared +with the rest of the Qwen3.6 family and are gated on `NOT FLASHRT_SLIM_BUILD`, +not on architecture — so `-DFLASHRT_SLIM_BUILD=ON` removes them and the +frontend then refuses to start, naming what is missing. Do not use a slim build +for this model. + +Selecting tiers by reading the source's own grouping is therefore not enough to +know what a target needs; the call sites are what decide. + `_W4A4` refuses to configure on a target without block-scaled MMA. CUTLASS still compiles those translation units elsewhere, but substitutes `CUTE_INVALID_CONTROL_PATH` for the MMA, so the build would succeed and then diff --git a/flash_rt/frontends/torch/nexn2_rtx.py b/flash_rt/frontends/torch/nexn2_rtx.py index 3c8b9deb..383c44c0 100644 --- a/flash_rt/frontends/torch/nexn2_rtx.py +++ b/flash_rt/frontends/torch/nexn2_rtx.py @@ -150,8 +150,6 @@ def _build_kernelized_nvfp4(self) -> None: to NVFP4 (GDN in_proj / norms / router kept BF16) and frees the BF16 source as it goes, fitting in ~22 GB. """ - from transformers import AutoTokenizer - from flash_rt import flash_rt_kernels as fvk from flash_rt.frontends.torch._nexn2_rtx_nvfp4_weights import ( extract_weights_nexn2_nvfp4, @@ -164,7 +162,6 @@ def _build_kernelized_nvfp4(self) -> None: usage_doc=self._USAGE_DOC, ) - self._tokenizer = AutoTokenizer.from_pretrained(self.checkpoint_path) self._fvk = fvk self._weights = extract_weights_nexn2_nvfp4( self.checkpoint_path, fvk, device=self.device, @@ -172,12 +169,37 @@ def _build_kernelized_nvfp4(self) -> None: @property def tokenizer(self): - """The HF tokenizer loaded from the checkpoint.""" + """The checkpoint's tokenizer, loaded when something asks for it. + + Loading it eagerly would make ``transformers`` a hard requirement of + the runtime, which it is not: a caller that supplies token ids through + :meth:`set_prompt_ids` never needs one. That matters for a deployment + target where the dependency may be absent or unwelcome, and it keeps + the kernel and weight paths testable without it. + """ + if self._tokenizer is None: + from transformers import AutoTokenizer + + self._tokenizer = AutoTokenizer.from_pretrained( + self.checkpoint_path) return self._tokenizer + def set_prompt_ids(self, token_ids) -> None: + """Set the prompt from token ids, requiring no tokenizer.""" + import torch + + ids = torch.as_tensor( + token_ids, dtype=torch.long, device=self.device).reshape(1, -1) + if ids.shape[1] == 0: + raise ValueError('token_ids is empty') + # Matches set_prompt: the decode state is not discarded, because + # seed_prefill resets the recurrent and KV caches itself and + # reallocating them per prompt would be waste. + self._prompt_ids = ids + def set_prompt(self, text: str) -> None: """Tokenize ``text`` for the next ``infer()`` / ``generate()`` call.""" - enc = self._tokenizer(text, return_tensors='pt') + enc = self.tokenizer(text, return_tensors='pt') self._prompt_ids = enc['input_ids'].to(self.device) def infer(self): From 69e1252387744c8237923a3278d4cc54219c5cf9 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 30 Jul 2026 07:54:50 -0400 Subject: [PATCH 22/85] Let a configuration state which kernels it calls The fail-fast check used one hardcoded list for every configuration. Running the streamed-expert path on a target where the block-scaled 4-bit tier is correctly not built showed why that is wrong: it refused to start over moe_blocktile_mma, which that path never calls. Those kernels serve the batched prefill, and streaming runs prefill through the per-token loop because a miss issues host reads. A list demanding more than a path uses turns a working build into a refusal; one demanding less lets a missing symbol surface mid-forward. So the list belongs with whatever decides which kernels get called: _require_kernels takes it, the frontend carries the default, and a subclass narrows or extends it. The streaming frontend drops the two MMA kernels and adds the two it does call. Behaviour for the existing frontends is unchanged. --- flash_rt/frontends/torch/nexn2_rtx.py | 18 +++++++++++++++--- qwen36_moe_edge/streaming_frontend.py | 11 +++++++++++ 2 files changed, 26 insertions(+), 3 deletions(-) diff --git a/flash_rt/frontends/torch/nexn2_rtx.py b/flash_rt/frontends/torch/nexn2_rtx.py index 383c44c0..46edc205 100644 --- a/flash_rt/frontends/torch/nexn2_rtx.py +++ b/flash_rt/frontends/torch/nexn2_rtx.py @@ -36,11 +36,18 @@ def _require_kernels( fvk, *, model_label: str = "Nex-N2", - usage_doc: str = "docs/nexn2_usage.md") -> None: + usage_doc: str = "docs/nexn2_usage.md", + required=None) -> None: """Raise a clear RuntimeError if the gated qwen3_5_moe kernels or the FA2 module are missing (build was not configured with - -DFLASHRT_ENABLE_QWEN35MOE=ON, or flash_rt_fa2 is absent).""" - missing = [s for s in _REQUIRED_FVK if not hasattr(fvk, s)] + -DFLASHRT_ENABLE_QWEN35MOE=ON, or flash_rt_fa2 is absent). + + ``required`` lets a configuration that calls fewer kernels say so. A list + demanding more than a path uses turns a working build into a refusal; one + demanding less lets a missing symbol surface mid-forward. Both are wrong, + so the list belongs with whatever decides which kernels get called. + """ + missing = [s for s in (required or _REQUIRED_FVK) if not hasattr(fvk, s)] if missing: raise RuntimeError( f"{model_label} kernelized path needs the qwen3_5_moe SM120 " @@ -67,6 +74,10 @@ def _require_kernels( class Nexn2TorchFrontendRtx: """Nex-N2-mini inference frontend (PyTorch + RTX SM120).""" + # Kernels this configuration calls; a subclass whose path calls fewer + # narrows it. See _require_kernels. + _REQUIRED_KERNELS = _REQUIRED_FVK + _MODEL_LABEL = "Nex-N2" _USAGE_DOC = "docs/nexn2_usage.md" @@ -160,6 +171,7 @@ def _build_kernelized_nvfp4(self) -> None: fvk, model_label=self._MODEL_LABEL, usage_doc=self._USAGE_DOC, + required=self._REQUIRED_KERNELS, ) self._fvk = fvk diff --git a/qwen36_moe_edge/streaming_frontend.py b/qwen36_moe_edge/streaming_frontend.py index 6fd33756..ef5eb345 100644 --- a/qwen36_moe_edge/streaming_frontend.py +++ b/qwen36_moe_edge/streaming_frontend.py @@ -28,6 +28,17 @@ class Qwen36MoeStreamingFrontend(Qwen36MoeTextFrontendRtx): _MODEL_LABEL = "Qwen3.6-35B-A3B text, streamed experts" + # The block-scaled 4-bit MMA kernels are absent from this list because this + # path never calls them: they serve the batched prefill, and streaming runs + # prefill through the per-token loop instead, since a miss issues host reads. + # Demanding them would refuse a build that can run this perfectly well -- + # which is what happened on the first attempt, on a target where the tier is + # correctly not built at all. + _REQUIRED_KERNELS = tuple( + name for name in Qwen36MoeTextFrontendRtx._REQUIRED_KERNELS + if not name.startswith(('moe_blocktile_mma', 'moe_m16_mma')) + ) + ('qwen35moe_e0m3_dequant_bf16', 'bf16_matvec_sm120_bf16') + def __init__(self, checkpoint_path: str, bundle: str | Path, *, slots_per_layer: int, device: str = "cuda:0", From eabfc822357381e7c2889186faa24e56768d3e23 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 30 Jul 2026 07:59:43 -0400 Subject: [PATCH 23/85] Pass stream_experts to the loader, and give lm_head a portable path Two defects, both found by running the assembled path on a non-SM120 target. The loader never received stream_experts. An unchecked string replacement had matched nothing while reporting success, so the feature was off and the run that followed reported a resident footprint of 21.436 GiB -- indistinguishable from the ordinary frontend's documented 21.44, and plausible enough to accept without comparing the two. A test now pins the wiring. The lm_head decode called fp4_w4a4_mma_sm120_full_n_bf16out, which is built only for GPU_ARCH 120 and 121. On every other target, the Orin one included, that path had no implementation: the symbol simply was not there. The W4A16 matvec reads the same swizzled weight and the same scale factors, leaves the activation in bf16 -- so it also drops the activation quantisation and its error -- and is in a tier that builds wherever the core does. It is selected when the W4A4 kernel is absent, so SM120 behaviour is unchanged. --- flash_rt/frontends/torch/_nexn2_rtx_decode.py | 21 ++++++++++---- flash_rt/frontends/torch/nexn2_rtx.py | 3 +- tests/test_qwen36_moe_expert_cache.py | 28 +++++++++++++++++++ 3 files changed, 46 insertions(+), 6 deletions(-) diff --git a/flash_rt/frontends/torch/_nexn2_rtx_decode.py b/flash_rt/frontends/torch/_nexn2_rtx_decode.py index 4981c852..6e5e4805 100644 --- a/flash_rt/frontends/torch/_nexn2_rtx_decode.py +++ b/flash_rt/frontends/torch/_nexn2_rtx_decode.py @@ -658,11 +658,22 @@ def decode_step(state, token_id, pos, fvk, device): p['lm_head_packed_t'] = packed p['lm_head_sf_t'] = sf p['lm_head_alpha'] = float(og.item()) - xp, xsf = _quant_act(h.reshape(1, HID), fvk, device, _cs()) - fvk.fp4_w4a4_mma_sm120_full_n_bf16out( - xp.data_ptr(), p['lm_head_packed_t'].data_ptr(), logits.data_ptr(), - vocab, HID, xsf.data_ptr(), p['lm_head_sf_t'].data_ptr(), - p['lm_head_alpha'], _cs()) + if hasattr(fvk, 'fp4_w4a4_mma_sm120_full_n_bf16out'): + xp, xsf = _quant_act(h.reshape(1, HID), fvk, device, _cs()) + fvk.fp4_w4a4_mma_sm120_full_n_bf16out( + xp.data_ptr(), p['lm_head_packed_t'].data_ptr(), + logits.data_ptr(), vocab, HID, xsf.data_ptr(), + p['lm_head_sf_t'].data_ptr(), p['lm_head_alpha'], _cs()) + return logits + # That kernel is built only for GPU_ARCH 120/121, so on every other target + # this path had no implementation at all. The W4A16 matvec reads the same + # swizzled weight and the same scale factors, leaves the activation in + # bf16 -- so it also skips the activation quantisation and its error -- and + # lives in a tier that builds wherever the core does. + fvk.w4a16_matvec_sm120_bf16( + h.reshape(1, HID).contiguous().data_ptr(), + p['lm_head_packed_t'].data_ptr(), p['lm_head_sf_t'].data_ptr(), + logits.data_ptr(), vocab, HID, p['lm_head_alpha'], _cs()) return logits diff --git a/flash_rt/frontends/torch/nexn2_rtx.py b/flash_rt/frontends/torch/nexn2_rtx.py index 46edc205..8747c150 100644 --- a/flash_rt/frontends/torch/nexn2_rtx.py +++ b/flash_rt/frontends/torch/nexn2_rtx.py @@ -177,7 +177,8 @@ def _build_kernelized_nvfp4(self) -> None: self._fvk = fvk self._weights = extract_weights_nexn2_nvfp4( self.checkpoint_path, fvk, device=self.device, - quant_scope=self._quant_scope) + quant_scope=self._quant_scope, + stream_experts=self._stream_experts) @property def tokenizer(self): diff --git a/tests/test_qwen36_moe_expert_cache.py b/tests/test_qwen36_moe_expert_cache.py index a96e46a0..2a75bbe3 100644 --- a/tests/test_qwen36_moe_expert_cache.py +++ b/tests/test_qwen36_moe_expert_cache.py @@ -160,3 +160,31 @@ def test_global_scales_validates_the_sidecar_size(): assert "num_experts * 2 * 4" in source assert "raise ValueError" in source + + +def test_streaming_frontend_actually_asks_the_loader_to_skip(): + # A silently-unapplied edit left the argument off this call once, and the + # run that followed reported a resident footprint identical to the ordinary + # frontend's -- plausible enough to miss without comparing against the + # documented baseline. Pin the wiring so it cannot regress quietly. + import inspect + + from flash_rt.frontends.torch import nexn2_rtx + + source = inspect.getsource(nexn2_rtx.Nexn2TorchFrontendRtx) + + assert "stream_experts=self._stream_experts" in source + + +def test_lm_head_has_a_path_without_the_sm120_only_kernel(): + # fp4_w4a4_mma_sm120_full_n_bf16out is built only for GPU_ARCH 120/121, so + # before this branch the lm_head decode had no implementation on any other + # target, the Orin one included. + import inspect + + from flash_rt.frontends.torch import _nexn2_rtx_decode + + source = inspect.getsource(_nexn2_rtx_decode.decode_step) + + assert "hasattr(fvk, 'fp4_w4a4_mma_sm120_full_n_bf16out')" in source + assert "w4a16_matvec_sm120_bf16" in source From 2ad746a7a6556cab2602226752c0cf255d67a24b Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 30 Jul 2026 08:07:02 -0400 Subject: [PATCH 24/85] Fix three defects in the streamed-expert assembly Running the assembled path found three faults, all in the code that splices the cache in rather than in any of the pieces it splices together. It returned from the middle of the MoE layer, which dropped the shared expert and its sigmoid gate from every layer and returned (1, HID) float32 where the resident path returns (1, 1, HID) bfloat16. Streaming now replaces only the routed experts' own GEMVs and falls through to the shared tail, so the weighted sum, the shared expert and the gate are the same code on both paths. It never rotated the activation. The bundle stores H*W, so both GEMMs need the activation rotated the same way; without it the products are wrong while staying finite and plausible, which is the failure mode that hides. The transform is built once per state and applied to the hidden state entering gate_up and to the gated result entering down. Staging buffers were indexed by task number. The pool bounds how many tasks run at once, not the order they finish, so task N and task N + len(staging) could hold the same buffer and read into each other's memory. A task now takes one from a queue and returns it, owning it for the duration. Two diagnostics that turned a guess into a measurement, and are worth keeping. The router top-k return code was unchecked, so a failure left torch.empty memory to be used as expert indices; it now raises. And a direct read reports a misaligned offset, length or buffer with one indistinguishable EINVAL, so the cache reports which of the three it was, and rejects an out-of-range expert with the whole request quoted -- which is how the -1 was found. --- flash_rt/frontends/torch/_nexn2_rtx_decode.py | 142 +++++++++++------- qwen36_moe_edge/expert_cache.py | 75 ++++++--- tests/test_qwen36_moe_expert_cache.py | 61 ++++++++ 3 files changed, 200 insertions(+), 78 deletions(-) diff --git a/flash_rt/frontends/torch/_nexn2_rtx_decode.py b/flash_rt/frontends/torch/_nexn2_rtx_decode.py index 6e5e4805..aabed348 100644 --- a/flash_rt/frontends/torch/_nexn2_rtx_decode.py +++ b/flash_rt/frontends/torch/_nexn2_rtx_decode.py @@ -283,6 +283,7 @@ def __init__(self, handles, max_seq, device): # meaningful when the loader skipped them; see _moe_experts_streamed. self.expert_cache = None self._scratch = None + self._hadamard = None def _streamed_scratch(self, device): """The two decode buffers a streamed expert is unpacked into. @@ -432,33 +433,52 @@ def _decode_full(h, ld, state, full_rank, pos, fvk, device): 1, 1, HID) -def _moe_experts_streamed(x, idx, tw_row, ld, state, fvk, device, s): - """The routed experts, read from storage instead of held in memory. +def _hadamard16(device): + """The block-16 transform, built once. Symmetric and its own inverse.""" + m = torch.ones(1, 1, dtype=torch.float32, device=device) + for _ in range(4): + m = torch.cat((torch.cat((m, m), 1), torch.cat((m, -m), 1)), 0) + return m / 4.0 - Only reachable when the loader was told to stream them, in which case the - per-layer stacked NVFP4 tensors were never allocated -- that is the whole - point, and adding a cache without skipping them would cost memory rather - than save it. - Each block is decoded to bf16 and multiplied with the shared bf16 GEMM. The - block-scaled 4-bit GEMMs cannot read these blocks: different codebook, - different scale layout. Decoding costs bandwidth on a block that is already - resident, which is the cheap end of this system -- the misses are what cost - time. +def _rotate16(x, h): + """Apply the transform along the last dimension, in blocks of 16.""" + shape = x.shape + return (x.reshape(-1, 16).float() @ h).reshape(shape).to(x.dtype) - The whole top-k is fetched in one call so the reads overlap, and the - per-layer quota is at least the top-k, so none of the returned pointers can - be invalidated by the others. + +def _moe_experts_streamed(x, idx, state, fvk, device, s): + """The routed experts' outputs, read from storage instead of memory. + + Returns only ``d_dn`` -- the per-slot expert outputs. The weighted sum, the + shared expert and its gate are identical to the resident path and stay + there; replacing the whole layer here is how an earlier version silently + dropped the shared expert from every layer. + + Reachable only when the loader was told to stream, in which case the + per-layer stacked tensors were never allocated. Each block is decoded to + bf16 and multiplied with the shared bf16 GEMV, because the block-scaled + 4-bit GEMMs read neither this codebook nor this scale layout. + + When the bundle was written with the transform applied, the stored weight is + H*W, so the activation entering each GEMM has to be rotated the same way or + the products are wrong -- while staying finite and plausible, which is + exactly how it goes unnoticed. """ cache = state.expert_cache layer = state._active_layer experts = [int(value) for value in idx.cpu().tolist()] cache.get_many(layer, experts) + rotated = bool(cache.manifest.get('rht')) + if rotated and state._hadamard is None: + state._hadamard = _hadamard16(device) + h16 = state._hadamard + scratch = state._streamed_scratch(device) d_gu = torch.empty(TOPK, 2 * INTER, dtype=torch.bfloat16, device=device) d_dn = torch.empty(TOPK, HID, dtype=torch.bfloat16, device=device) - xc = x.contiguous() + xc = (_rotate16(x, h16) if rotated else x).contiguous() for slot, expert in enumerate(experts): parts = cache.components(layer, expert) @@ -476,7 +496,10 @@ def _moe_experts_streamed(x, idx, tw_row, ld, state, fvk, device, s): gated = _silu_mul( d_gu[slot:slot + 1, :INTER], d_gu[slot:slot + 1, INTER:], - fvk, device).contiguous() + fvk, device) + if rotated: + gated = _rotate16(gated, h16) + gated = gated.contiguous() rc = fvk.qwen35moe_e0m3_dequant_bf16( parts['down_weight'].data_ptr(), parts['down_scale'].data_ptr(), @@ -487,15 +510,7 @@ def _moe_experts_streamed(x, idx, tw_row, ld, state, fvk, device, s): fvk.bf16_matvec_sm120_bf16( gated.data_ptr(), scratch['down'].data_ptr(), d_dn[slot].data_ptr(), HID, INTER, s) - - if 'decode_topk_rows' not in ld: - ld['decode_topk_rows'] = torch.arange( - TOPK, dtype=torch.int32, device=device) - out = torch.empty(HID, dtype=torch.float32, device=device) - fvk.moe_weighted_sum_sm120_bf16( - d_dn.data_ptr(), ld['decode_topk_rows'].data_ptr(), - tw_row.data_ptr(), out.data_ptr(), 1, TOPK, HID, HID, s) - return out.unsqueeze(0) + return d_dn def _moe_layer_decode(h, ld, state, fvk, device): @@ -531,8 +546,14 @@ def _moe_layer_decode(h, ld, state, fvk, device): lr = logit_raw.reshape(-1).contiguous() idx = torch.empty(TOPK, dtype=torch.int32, device=device) topv = torch.empty(TOPK, dtype=torch.float32, device=device) - fvk.moe_router_topk_sm120_bf16(lr.data_ptr(), idx.data_ptr(), topv.data_ptr(), - lr.numel(), TOPK, s) + # idx and topv come from torch.empty, so an unchecked failure here leaves + # uninitialised memory to be used as expert indices -- which reaches a file + # offset before anything notices. + rc = fvk.moe_router_topk_sm120_bf16( + lr.data_ptr(), idx.data_ptr(), topv.data_ptr(), lr.numel(), TOPK, s) + if rc: + raise RuntimeError( + f'router top-k failed with {rc} for {lr.numel()} experts, k={TOPK}') tw_row = F.softmax(topv, -1) # (TOPK,) device if state.router_trace is not None: state.router_trace[state._active_layer].append( @@ -541,38 +562,43 @@ def _moe_layer_decode(h, ld, state, fvk, device): state.moe_input_trace[state._active_layer].append( x.detach().to("cpu", copy=True)) + # Streaming replaces only the routed experts' own GEMVs. Everything after + # this -- the weighted sum, the shared expert, its gate -- is identical, and + # returning early from here is how an earlier version silently dropped the + # shared expert from every layer. if ld.get('experts_streamed'): - return _moe_experts_streamed( - x, idx, tw_row, ld, state, fvk, device, s) - - if 'experts_gate_up_alpha_dev' not in ld: # cache once/layer - ld['experts_gate_up_alpha_dev'] = \ - ld['experts_gate_up_alpha_t'].to(device).contiguous() - ld['experts_down_alpha_dev'] = \ - ld['experts_down_alpha_t'].to(device).contiguous() - gu_p, gu_s = ld['experts_gate_up_packed_t'], ld['experts_gate_up_sf_t'] - dn_p, dn_s = ld['experts_down_packed_t'], ld['experts_down_sf_t'] - gu_a, dn_a = ld['experts_gate_up_alpha_dev'], ld['experts_down_alpha_dev'] - n_gu, n_dn = gu_p.shape[1], dn_p.shape[1] # 1024 / HID - - # gate_up: shared BF16 activation, grouped W4A16 over the 8 experts. BF16 - # activation -> no activation quant, higher cos than the W4A4 mma, and - # faster at this scale (6.2 vs 8.2 us standalone). - xc = x.contiguous() - d_gu = torch.empty(TOPK, n_gu, dtype=torch.bfloat16, device=device) - fvk.moe_grouped_w4a16_sm120_bf16( - xc.data_ptr(), gu_p.data_ptr(), gu_s.data_ptr(), gu_a.data_ptr(), - idx.data_ptr(), d_gu.data_ptr(), TOPK, n_gu, HID, - 0, gu_p[0].numel(), gu_s[0].numel(), s) - - # down: silu(gate)*up (BF16, fused) then grouped W4A16 (per-slot activation). - g_, u_ = d_gu[:, :INTER], d_gu[:, INTER:] - inter = _silu_mul(g_, u_, fvk, device).contiguous() - d_dn = torch.empty(TOPK, n_dn, dtype=torch.bfloat16, device=device) - fvk.moe_grouped_w4a16_sm120_bf16( - inter.data_ptr(), dn_p.data_ptr(), dn_s.data_ptr(), dn_a.data_ptr(), - idx.data_ptr(), d_dn.data_ptr(), TOPK, n_dn, INTER, - INTER, dn_p[0].numel(), dn_s[0].numel(), s) + d_dn = _moe_experts_streamed(x, idx, state, fvk, device, s) + n_dn = HID + else: + if 'experts_gate_up_alpha_dev' not in ld: # cache once/layer + ld['experts_gate_up_alpha_dev'] = \ + ld['experts_gate_up_alpha_t'].to(device).contiguous() + ld['experts_down_alpha_dev'] = \ + ld['experts_down_alpha_t'].to(device).contiguous() + gu_p, gu_s = ld['experts_gate_up_packed_t'], ld['experts_gate_up_sf_t'] + dn_p, dn_s = ld['experts_down_packed_t'], ld['experts_down_sf_t'] + gu_a = ld['experts_gate_up_alpha_dev'] + dn_a = ld['experts_down_alpha_dev'] + n_gu, n_dn = gu_p.shape[1], dn_p.shape[1] # 1024 / HID + + # gate_up: shared BF16 activation, grouped W4A16 over the 8 experts. + # BF16 activation -> no activation quant, higher cos than the W4A4 mma, + # and faster at this scale (6.2 vs 8.2 us standalone). + xc = x.contiguous() + d_gu = torch.empty(TOPK, n_gu, dtype=torch.bfloat16, device=device) + fvk.moe_grouped_w4a16_sm120_bf16( + xc.data_ptr(), gu_p.data_ptr(), gu_s.data_ptr(), gu_a.data_ptr(), + idx.data_ptr(), d_gu.data_ptr(), TOPK, n_gu, HID, + 0, gu_p[0].numel(), gu_s[0].numel(), s) + + # down: silu(gate)*up (BF16, fused) then grouped W4A16 (per-slot act). + g_, u_ = d_gu[:, :INTER], d_gu[:, INTER:] + inter = _silu_mul(g_, u_, fvk, device).contiguous() + d_dn = torch.empty(TOPK, n_dn, dtype=torch.bfloat16, device=device) + fvk.moe_grouped_w4a16_sm120_bf16( + inter.data_ptr(), dn_p.data_ptr(), dn_s.data_ptr(), dn_a.data_ptr(), + idx.data_ptr(), d_dn.data_ptr(), TOPK, n_dn, INTER, + INTER, dn_p[0].numel(), dn_s[0].numel(), s) # Fixed-order weighted sum. The generic torch matmul may choose a # reduction whose accumulation order changes between launches, which can # flip a later greedy decision when two logits are nearly tied. diff --git a/qwen36_moe_edge/expert_cache.py b/qwen36_moe_edge/expert_cache.py index 81ff94ca..a9c8f9f8 100644 --- a/qwen36_moe_edge/expert_cache.py +++ b/qwen36_moe_edge/expert_cache.py @@ -30,6 +30,7 @@ import json import os +import queue from collections import Counter, OrderedDict from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass, field @@ -130,6 +131,13 @@ def __init__(self, config: CacheConfig): f"{self.alignment}-byte aligned, which direct reads " "require") self._pool = ThreadPoolExecutor(max_workers=config.staging_buffers) + # Buffers are taken from here and returned, so a task owns one for as + # long as it runs. Indexing by task number would let task N and task + # N + len(staging) share a buffer: the pool bounds how many run at once, + # not the order they finish in. + self._available: queue.Queue = queue.Queue() + for index in range(len(self._staging)): + self._available.put(index) # Per-layer LRU of expert -> slot index, and that layer's free slots. self._lru: list[OrderedDict[int, int]] = [ @@ -198,22 +206,43 @@ def _fd(self, layer: int) -> int: self._fds[layer] = os.open(path, flags) return self._fds[layer] - def _fetch(self, layer: int, expert: int, slot: int, buffer: int) -> None: - staging = self._staging[buffer] - view = memoryview(staging.numpy()) - fd = self._fd(layer) - base = expert * self.block_bytes - offset = 0 - while offset < self.block_bytes: - length = min(self.config.read_chunk, self.block_bytes - offset) - read = os.preadv(fd, [view[offset:offset + length]], base + offset) - if read <= 0: - raise IOError( - f"short read of layer {layer} expert {expert} at " - f"{offset}/{self.block_bytes}") - offset += read - self.slots[slot].copy_(staging) - self.bytes_read += self.block_bytes + def _fetch(self, layer: int, expert: int, slot: int) -> None: + if not 0 <= expert < self.num_experts: + raise ValueError( + f"expert {expert} is outside 0..{self.num_experts - 1}; a " + "negative or oversized index becomes an invalid file offset") + buffer = self._available.get() + try: + staging = self._staging[buffer] + view = memoryview(staging.numpy()) + fd = self._fd(layer) + base = expert * self.block_bytes + offset = 0 + while offset < self.block_bytes: + length = min(self.config.read_chunk, self.block_bytes - offset) + try: + read = os.preadv( + fd, [view[offset:offset + length]], base + offset) + except OSError as error: + # A direct read rejects a misaligned offset, length or + # buffer with the same EINVAL, which says nothing about + # which of the three it was. + raise OSError( + f"{error.strerror} reading layer {layer} expert " + f"{expert}: offset {base + offset} aligned=" + f"{(base + offset) % self.alignment == 0}, length " + f"{length} aligned={length % self.alignment == 0}, " + f"buffer {staging.data_ptr():#x} aligned=" + f"{staging.data_ptr() % self.alignment == 0}") from error + if read <= 0: + raise IOError( + f"short read of layer {layer} expert {expert} at " + f"{offset}/{self.block_bytes}") + offset += read + self.slots[slot].copy_(staging) + self.bytes_read += self.block_bytes + finally: + self._available.put(buffer) def _claim(self, layer: int, expert: int) -> int: """A slot for an expert not currently held, evicting if necessary.""" @@ -232,6 +261,14 @@ def get_many(self, layer: int, experts) -> list[int]: pointers can be invalidated by the others. """ wanted = list(dict.fromkeys(int(expert) for expert in experts)) + out_of_range = [ + expert for expert in wanted + if not 0 <= expert < self.num_experts + ] + if out_of_range: + raise ValueError( + f"layer {layer} was asked for experts {out_of_range} outside " + f"0..{self.num_experts - 1}; the full request was {wanted}") if len(wanted) > self.config.slots_per_layer: raise ValueError( f"asked for {len(wanted)} experts of layer {layer} but the " @@ -247,10 +284,8 @@ def get_many(self, layer: int, experts) -> list[int]: pending.append((expert, self._claim(layer, expert))) if pending: futures = [ - self._pool.submit( - self._fetch, layer, expert, slot, - index % len(self._staging)) - for index, (expert, slot) in enumerate(pending) + self._pool.submit(self._fetch, layer, expert, slot) + for expert, slot in pending ] for future in futures: future.result() diff --git a/tests/test_qwen36_moe_expert_cache.py b/tests/test_qwen36_moe_expert_cache.py index 2a75bbe3..3b57a54f 100644 --- a/tests/test_qwen36_moe_expert_cache.py +++ b/tests/test_qwen36_moe_expert_cache.py @@ -188,3 +188,64 @@ def test_lm_head_has_a_path_without_the_sm120_only_kernel(): assert "hasattr(fvk, 'fp4_w4a4_mma_sm120_full_n_bf16out')" in source assert "w4a16_matvec_sm120_bf16" in source + + +def test_staging_buffers_are_owned_for_the_duration_of_a_read(): + # Indexing a buffer by task number is unsafe: the pool bounds how many + # tasks run at once, not the order they finish in, so task N and task + # N + len(staging) can hold the same buffer simultaneously and read into + # each other's memory. A task must acquire one and give it back. + import inspect + + source = inspect.getsource(ExpertCache._fetch) + + assert "self._available.get()" in source + assert "self._available.put(buffer)" in source + assert "finally:" in source + # And the submission must not hand a buffer index in at all. + submit = inspect.getsource(ExpertCache.get_many) + assert "len(self._staging)" not in submit + + +def test_fetch_rejects_an_out_of_range_expert(): + # A negative or oversized index turns into an invalid file offset, which a + # direct read reports as a bare EINVAL that names nothing. + import inspect + + source = inspect.getsource(ExpertCache._fetch) + + assert "0 <= expert < self.num_experts" in source + assert "invalid file offset" in source + + +def test_streaming_replaces_only_the_routed_expert_gemvs(): + # An earlier version returned from the middle of the MoE layer, which + # dropped the shared expert and its gate from every layer -- and returned + # the wrong shape and dtype while doing it. The branch must set d_dn and + # fall through to the shared tail. + import inspect + + from flash_rt.frontends.torch import _nexn2_rtx_decode + + source = inspect.getsource(_nexn2_rtx_decode._moe_layer_decode) + + assert "d_dn = _moe_experts_streamed(" in source + assert "return _moe_experts_streamed(" not in source + # The shared expert and its gate are still reached on both paths. + assert "shared_down_proj" in source + assert "shared_gate_w_t" in source + + +def test_streamed_experts_rotate_the_activation_when_the_bundle_is_rotated(): + # The bundle stores H*W, so an unrotated activation gives a wrong product + # that is still finite and plausible -- the failure mode that hides. + import inspect + + from flash_rt.frontends.torch import _nexn2_rtx_decode + + source = inspect.getsource(_nexn2_rtx_decode._moe_experts_streamed) + + assert "cache.manifest.get('rht')" in source + # Both GEMMs: the hidden state entering gate_up, and the gated result + # entering down. + assert source.count("_rotate16(") == 2 From 5efdb40ea76b106d6886aefe425ad50dc4468f53 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 30 Jul 2026 08:45:10 -0400 Subject: [PATCH 25/85] Check the vendored attention kernel actually computes, and fall back Importing the FA2 module and finding its symbols proves neither that its kernel runs nor that it is right. Measured on an SM110 part: the module imported, every symbol was present, and at run time the kernel printed a complaint and returned without writing its output. Downstream that reads as plausible-but-wrong attention, not as a failure -- the model still produced 15 of 16 reference tokens with ten of forty layers contributing nothing, which is a coincidence of the residual stream rather than evidence of anything. Construction now runs one small case through the same launch the hot path uses and compares it against scaled_dot_product_attention. If it does not agree, or produces non-finite values, or raises, attention falls back to the reference implementation. One launch at construction, and it is the same launch, so the probe cannot pass while the real call fails. This is the third time in this work that a kernel compiled, linked and loaded while being unable to run: the block-scaled 4-bit tier substitutes an invalid control path off its own architecture, and this. Symbol presence is not a capability check. --- flash_rt/hardware/rtx/attn_backend_nexn2.py | 83 ++++++++++++++++++++- tests/test_qwen36_moe_expert_cache.py | 23 ++++++ 2 files changed, 105 insertions(+), 1 deletion(-) diff --git a/flash_rt/hardware/rtx/attn_backend_nexn2.py b/flash_rt/hardware/rtx/attn_backend_nexn2.py index 863bea4d..83346694 100644 --- a/flash_rt/hardware/rtx/attn_backend_nexn2.py +++ b/flash_rt/hardware/rtx/attn_backend_nexn2.py @@ -101,6 +101,73 @@ def __init__(self, max_seq: int, max_q_seq: int = 1, dtype=None): self._num_sms = torch.cuda.get_device_properties( torch.cuda.current_device() ).multi_processor_count + self._fa2_usable = self._probe_fa2() + + def _probe_fa2(self) -> bool: + """Does the vendored kernel actually compute on this device? + + Importing it and finding its symbols proves neither. Its own arch + handling can leave a build that links, loads, prints a complaint to + stdout and returns without writing the output -- which downstream looks + like plausible-but-wrong attention rather than a failure. Measured on + an SM110 part: the module imported, every symbol was present, and the + kernel refused at run time. + + So run one small case against a reference and compare. The cost is one + launch at construction. + """ + import torch.nn.functional as F + + q_seq, kv_seq = 1, 8 + generator = torch.Generator(device=self.Q_buf.device).manual_seed(1) + q = torch.randn( + 1, q_seq, self.NUM_Q_HEADS, self.HEAD_DIM, generator=generator, + device=self.Q_buf.device, dtype=torch.bfloat16) + k = torch.randn( + 1, kv_seq, self.NUM_KV_HEADS, self.HEAD_DIM, generator=generator, + device=self.Q_buf.device, dtype=torch.bfloat16) + v = torch.randn_like(k) + self.Q_buf[:, :q_seq].copy_(q) + self.K_cache[0:1, :kv_seq].copy_(k) + self.V_cache[0:1, :kv_seq].copy_(v) + self.O_buf[:, :q_seq].zero_() + try: + self._launch_fa2(0, q_seq, kv_seq, 0, + 1.0 / (self.HEAD_DIM ** 0.5)) + torch.cuda.synchronize() + except Exception: # noqa: BLE001 + return False + produced = self.O_buf[:, :q_seq].float().clone() + + groups = self.NUM_Q_HEADS // self.NUM_KV_HEADS + kr = k.repeat_interleave(groups, dim=2) + vr = v.repeat_interleave(groups, dim=2) + expected = F.scaled_dot_product_attention( + q.transpose(1, 2).float(), kr.transpose(1, 2).float(), + vr.transpose(1, 2).float()).transpose(1, 2) + if not torch.isfinite(produced).all(): + return False + reference = expected.norm().clamp_min(1e-6) + return bool( + ((produced - expected).norm() / reference).item() < 0.05) + + def _sdpa(self, layer_idx: int, q_seq: int, kv_seq: int, + softmax_scale: float) -> None: + """Reference attention, for a device the vendored kernel refuses.""" + import torch.nn.functional as F + + q = self.Q_buf[:, :q_seq] + k = self.K_cache[layer_idx:layer_idx + 1, :kv_seq] + v = self.V_cache[layer_idx:layer_idx + 1, :kv_seq] + groups = self.NUM_Q_HEADS // self.NUM_KV_HEADS + out = F.scaled_dot_product_attention( + q.transpose(1, 2).float(), + k.repeat_interleave(groups, dim=2).transpose(1, 2).float(), + v.repeat_interleave(groups, dim=2).transpose(1, 2).float(), + is_causal=q_seq > 1, + scale=softmax_scale, + ).transpose(1, 2) + self.O_buf[:, :q_seq].copy_(out.to(self.O_buf.dtype)) # ── Layer cache pointer math ── @@ -175,6 +242,21 @@ def run(self, site: str, layer_idx: int, q_seq: int, if softmax_scale is None: softmax_scale = 1.0 / (self.HEAD_DIM ** 0.5) + if not self._fa2_usable: + self._sdpa(layer_idx, q_seq, kv_seq, softmax_scale) + return o.data_ptr() + + self._launch_fa2(layer_idx, q_seq, kv_seq, stream, softmax_scale) + return o.data_ptr() + + def _launch_fa2(self, layer_idx: int, q_seq: int, kv_seq: int, + stream: int, softmax_scale: float) -> None: + """One vendored-FA2 launch. Shared with the construction-time probe so + the probe exercises the same call the hot path makes.""" + q = self.Q_buf[:, :q_seq] + k = self.K_cache[layer_idx:layer_idx + 1, :kv_seq] + v = self.V_cache[layer_idx:layer_idx + 1, :kv_seq] + o = self.O_buf[:, :q_seq] self._fa2_fwd( Q=q.data_ptr(), K=k.data_ptr(), V=v.data_ptr(), O=o.data_ptr(), softmax_lse=self.lse_buf.data_ptr(), @@ -192,7 +274,6 @@ def run(self, site: str, layer_idx: int, q_seq: int, num_sms=self._num_sms, stream=stream, ) - return o.data_ptr() def make_nexn2_attention_spec(*, max_seq: int, max_q_seq: int = 1) -> dict: diff --git a/tests/test_qwen36_moe_expert_cache.py b/tests/test_qwen36_moe_expert_cache.py index 3b57a54f..85222541 100644 --- a/tests/test_qwen36_moe_expert_cache.py +++ b/tests/test_qwen36_moe_expert_cache.py @@ -249,3 +249,26 @@ def test_streamed_experts_rotate_the_activation_when_the_bundle_is_rotated(): # Both GEMMs: the hidden state entering gate_up, and the gated result # entering down. assert source.count("_rotate16(") == 2 + + +def test_attention_backend_probes_the_vendored_kernel(): + # Importing the module and finding its symbols proves neither that the + # kernel runs nor that it computes. On an SM110 part it imported, exposed + # every symbol, and then refused at run time while writing nothing -- + # which reads downstream as plausible-but-wrong attention, not a failure. + import inspect + + from flash_rt.hardware.rtx import attn_backend_nexn2 as backend + + cls = backend.RtxFlashAttnBackendNexn2 + probe = inspect.getsource(cls._probe_fa2) + run = inspect.getsource(cls.run) + + # The probe compares against a reference rather than checking a return code. + assert "scaled_dot_product_attention" in probe + assert "isfinite" in probe + # And it goes through the same launch the hot path uses, so it cannot pass + # while the real call fails. + assert "_launch_fa2" in probe + assert "_launch_fa2" in run + assert "self._fa2_usable" in run From 00c7b74cb5c4f048aa57eae655645fbaa102d878 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 30 Jul 2026 10:24:13 -0400 Subject: [PATCH 26/85] Treat a missing FA2 as a fallback rather than an error The arch list omits FA2 for Thor on purpose: that target uses FA4, whose SM100-class CuTe-DSL kernel wants Blackwell tensor memory. Orin's SM87 is Ampere and has none, so it takes FA2, which the arch list does enable. The two targets differ by design, and a frontend that hard-requires FA2 refuses to start on one of them for no reason -- the backend already computes the same thing without it. The backend now treats an absent module as a fallback, and a frontend can declare that its attention can fall back. Existing frontends keep the requirement. --- flash_rt/frontends/torch/nexn2_rtx.py | 12 +++++++++++- flash_rt/hardware/rtx/attn_backend_nexn2.py | 17 +++++++++++++---- qwen36_moe_edge/streaming_frontend.py | 6 ++++++ 3 files changed, 30 insertions(+), 5 deletions(-) diff --git a/flash_rt/frontends/torch/nexn2_rtx.py b/flash_rt/frontends/torch/nexn2_rtx.py index 8747c150..65f8bc58 100644 --- a/flash_rt/frontends/torch/nexn2_rtx.py +++ b/flash_rt/frontends/torch/nexn2_rtx.py @@ -37,7 +37,7 @@ def _require_kernels( fvk, *, model_label: str = "Nex-N2", usage_doc: str = "docs/nexn2_usage.md", - required=None) -> None: + required=None, require_fa2: bool = True) -> None: """Raise a clear RuntimeError if the gated qwen3_5_moe kernels or the FA2 module are missing (build was not configured with -DFLASHRT_ENABLE_QWEN35MOE=ON, or flash_rt_fa2 is absent). @@ -55,6 +55,11 @@ def _require_kernels( "are absent from flash_rt_kernels (missing: " f"{', '.join(missing)}). Rebuild on an SM120 toolchain with " f"-DFLASHRT_ENABLE_QWEN35MOE=ON. See {usage_doc}.") + if not require_fa2: + # The attention backend probes its kernel and falls back to a + # reference implementation, so a target that builds no FA2 -- Thor + # uses FA4 instead -- still runs. + return try: from flash_rt import flash_rt_fa2 as _fa2 except Exception as e: # pragma: no cover @@ -78,6 +83,10 @@ class Nexn2TorchFrontendRtx: # narrows it. See _require_kernels. _REQUIRED_KERNELS = _REQUIRED_FVK + # Whether the vendored FA2 module must be present. A subclass whose + # attention backend can fall back sets this False. + _REQUIRE_FA2 = True + _MODEL_LABEL = "Nex-N2" _USAGE_DOC = "docs/nexn2_usage.md" @@ -172,6 +181,7 @@ def _build_kernelized_nvfp4(self) -> None: model_label=self._MODEL_LABEL, usage_doc=self._USAGE_DOC, required=self._REQUIRED_KERNELS, + require_fa2=self._REQUIRE_FA2, ) self._fvk = fvk diff --git a/flash_rt/hardware/rtx/attn_backend_nexn2.py b/flash_rt/hardware/rtx/attn_backend_nexn2.py index 83346694..ad333d95 100644 --- a/flash_rt/hardware/rtx/attn_backend_nexn2.py +++ b/flash_rt/hardware/rtx/attn_backend_nexn2.py @@ -95,13 +95,22 @@ def __init__(self, max_seq: int, max_q_seq: int = 1, dtype=None): dtype=torch.float32, device=d, ) - from flash_rt import flash_rt_fa2 as _fa2 - self._fa2 = _fa2 - self._fa2_fwd = _fa2.fwd_bf16 + # A target may not build FA2 at all -- Thor uses FA4 instead, so the + # arch list deliberately omits it there. Absence is a fallback, not an + # error, because the reference path below computes the same thing. + try: + from flash_rt import flash_rt_fa2 as _fa2 + except ImportError: + self._fa2 = None + self._fa2_fwd = None + else: + self._fa2 = _fa2 + self._fa2_fwd = _fa2.fwd_bf16 self._num_sms = torch.cuda.get_device_properties( torch.cuda.current_device() ).multi_processor_count - self._fa2_usable = self._probe_fa2() + self._fa2_usable = ( + self._fa2_fwd is not None and self._probe_fa2()) def _probe_fa2(self) -> bool: """Does the vendored kernel actually compute on this device? diff --git a/qwen36_moe_edge/streaming_frontend.py b/qwen36_moe_edge/streaming_frontend.py index ef5eb345..5e025227 100644 --- a/qwen36_moe_edge/streaming_frontend.py +++ b/qwen36_moe_edge/streaming_frontend.py @@ -39,6 +39,12 @@ class Qwen36MoeStreamingFrontend(Qwen36MoeTextFrontendRtx): if not name.startswith(('moe_blocktile_mma', 'moe_m16_mma')) ) + ('qwen35moe_e0m3_dequant_bf16', 'bf16_matvec_sm120_bf16') + # The attention backend probes its kernel and falls back, so this runs on a + # target that builds no FA2. Thor is one: it uses FA4, whose SM100-class + # kernel needs Blackwell tensor memory that Orin's SM87 does not have -- + # so the two targets take different attention paths by design. + _REQUIRE_FA2 = False + def __init__(self, checkpoint_path: str, bundle: str | Path, *, slots_per_layer: int, device: str = "cuda:0", From d2ebdcce586e7a20d4a6a83d1baa055a6a3d301d Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 30 Jul 2026 10:24:55 -0400 Subject: [PATCH 27/85] Document why attention differs between targets FA2 is absent from the Thor build on purpose and FA4 cannot serve SM87, so the two targets take different attention paths. Recorded alongside the reason the backend probes its kernel rather than trusting that a symbol implies a working one: three kernels in this work compiled, linked and loaded while unable to run, and the one that failed silently still produced 15 of 16 reference tokens with a quarter of the layers contributing nothing. --- docs/qwen36_moe_usage.md | 26 ++++++++++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/docs/qwen36_moe_usage.md b/docs/qwen36_moe_usage.md index 3f4a6200..014fb44f 100644 --- a/docs/qwen36_moe_usage.md +++ b/docs/qwen36_moe_usage.md @@ -70,6 +70,32 @@ for this model. Selecting tiers by reading the source's own grouping is therefore not enough to know what a target needs; the call sites are what decide. +### Attention differs by target, by design + +The ten full-attention layers do not use the same kernel everywhere, and the +arch lists reflect that rather than overlooking it: + +| target | attention | why | +|---|---|---| +| SM120 / SM89 / SM87 | vendored FA2 | the SM80-family source, which `__CUDA_ARCH__ >= 800` admits | +| Thor SM110 | FA4 | its SM100-class CuTe-DSL kernel needs Blackwell tensor memory; ships as the `thor-fa4` pip extra, not compiled into `flash_rt_kernels` | + +So FA2 is deliberately absent from the Thor build, and FA4 cannot serve +Ampere-class SM87. Treat a missing FA2 as a signal to fall back, not as a +build error. + +The attention backend probes its kernel at construction: it runs one case +through the same launch the hot path uses and compares against +`scaled_dot_product_attention`, falling back if they disagree. That is not +belt-and-braces. Three times in this work a kernel compiled, linked and loaded +while being unable to run — the block-scaled 4-bit tier substitutes an invalid +control path off its own architecture, the lm_head kernel is simply absent +outside GPU_ARCH 120/121, and the vendored FA2 on an SM110 part printed a +complaint and returned without writing its output. That last one still produced +15 of 16 reference tokens, because ten of forty layers contributing nothing is +survivable for a residual stream — which is precisely why a symbol check is not +a capability check. + `_W4A4` refuses to configure on a target without block-scaled MMA. CUTLASS still compiles those translation units elsewhere, but substitutes `CUTE_INVALID_CONTROL_PATH` for the MMA, so the build would succeed and then From 7a12744dd5052412b25e6e2dcb7a636b4e2230f6 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 30 Jul 2026 12:14:50 -0400 Subject: [PATCH 28/85] Let the text prefill run where the vendored FA2 is absent The decode attention backend already treats a missing FA2 as a fallback and probes whether the kernel computes; prefill called it unconditionally, so a target that builds FA4 instead could load the model and then fail on the first full-attention layer. The reference path builds the causal mask explicitly. FA2 aligns causal bottom-right, so a chunked block's queries attend to keys [0, Sk-Sq+i]; torch's is_causal aligns top-left and the two agree only when Sq == Sk. Prefill's MoE tile is chosen the same way: the block-scaled 4-bit MMA tier is a build tier, so ask the module for it rather than assume it, and fall through to the weight-only grouped GEMV when it is not there. --- .../frontends/torch/_nexn2_rtx_forward.py | 98 +++++++++++++++++-- flash_rt/frontends/torch/qwen36_moe_rtx.py | 15 +++ 2 files changed, 106 insertions(+), 7 deletions(-) diff --git a/flash_rt/frontends/torch/_nexn2_rtx_forward.py b/flash_rt/frontends/torch/_nexn2_rtx_forward.py index 31389d3d..9c659ff2 100644 --- a/flash_rt/frontends/torch/_nexn2_rtx_forward.py +++ b/flash_rt/frontends/torch/_nexn2_rtx_forward.py @@ -475,15 +475,89 @@ def _gdn_layer(h, ld, fvk, device, eps, cap=None, rank=None, # the pre-existing flash_rt_fa2.so (already a hard dep of the decode backend), # so this adds no new csrc. _FA2_MOD = None +_FA2_USABLE = None _NUM_SMS = None def _get_fa2(): + """The vendored FA2 module, or None where the target does not build it. + + Thor is such a target: its arch list omits FA2 because it uses FA4. Absence + is a fallback, not an error -- ``_sdpa_causal_attn`` computes the same + thing -- so this returns None rather than raising, the way the decode + attention backend already treats it. + """ global _FA2_MOD if _FA2_MOD is None: - from flash_rt import flash_rt_fa2 as _m - _FA2_MOD = _m - return _FA2_MOD + try: + from flash_rt import flash_rt_fa2 as _m + except ImportError: + _FA2_MOD = False + else: + _FA2_MOD = _m + return _FA2_MOD or None + + +def _fa2_usable(device): + """Does the vendored kernel actually compute here? + + Importing it and finding its symbols proves neither: its arch handling can + leave a build that links, loads, prints a complaint and returns without + writing the output, which downstream looks like wrong attention rather than + a failure. The decode backend probes for the same reason. One launch, once. + """ + global _FA2_USABLE + if _FA2_USABLE is not None: + return _FA2_USABLE + if _get_fa2() is None: + _FA2_USABLE = False + return False + g = torch.Generator(device=device).manual_seed(1) + q = torch.randn(1, 8, NQ, HD, generator=g, device=device, + dtype=torch.bfloat16) + k = torch.randn(1, 8, NKV, HD, generator=g, device=device, + dtype=torch.bfloat16) + v = torch.randn_like(k) + try: + produced = _fa2_causal_attn(q, k, v, device, _probe=True).float() + torch.cuda.synchronize(device) + except Exception: # noqa: BLE001 + _FA2_USABLE = False + return False + expected = _sdpa_causal_attn(q, k, v, device).float() + _FA2_USABLE = bool( + torch.isfinite(produced).all() + and ((produced - expected).norm() + / expected.norm().clamp_min(1e-6)).item() < 0.05) + return _FA2_USABLE + + +def _sdpa_causal_attn(qf, kf, vf, device): + """Reference causal GQA attention, for a build without the FA2 kernel. + + FA2 causal aligns bottom-right -- query i attends to keys [0, Sk-Sq+i] -- + which is exactly a chunked block's absolute causal window. torch's + ``is_causal=True`` aligns top-left, and the two only agree when Sq == Sk, + so the mask is built explicitly rather than left to a flag whose convention + differs where it matters. + """ + import torch.nn.functional as F + + Sq, Sk = qf.shape[1], kf.shape[1] + q = qf.transpose(1, 2) # (1, NQ, Sq, HD) + k, v = kf.transpose(1, 2), vf.transpose(1, 2) # (1, NKV, Sk, HD) + qi = torch.arange(Sk - Sq, Sk, device=device).unsqueeze(1) + mask = torch.arange(Sk, device=device).unsqueeze(0) <= qi + try: + o = F.scaled_dot_product_attention( + q, k, v, attn_mask=mask, scale=float(HD) ** -0.5, enable_gqa=True) + except TypeError: # torch without native GQA + groups = NQ // NKV + o = F.scaled_dot_product_attention( + q, k.repeat_interleave(groups, dim=1), + v.repeat_interleave(groups, dim=1), + attn_mask=mask, scale=float(HD) ** -0.5) + return o.transpose(1, 2).contiguous() def _num_sms(): @@ -494,13 +568,18 @@ def _num_sms(): return _NUM_SMS -def _fa2_causal_attn(qf, kf, vf, device): +def _fa2_causal_attn(qf, kf, vf, device, *, _probe=False): """Causal GQA attention via the vendored FA2 kernel (bf16, native GQA -- no KV repeat). qf (1,Sq,NQ,HD), kf/vf (1,Sk,NKV,HD). Returns (1,Sq,NQ,HD). Sk may exceed Sq (chunked prefill: a block of Sq queries against the Sk accumulated KV); FA2 causal uses bottom-right alignment, so query i attends to keys [0, Sk-Sq+i] -- exactly the block's absolute causal window. splitkv - off (large-q parallelism).""" + off (large-q parallelism). + + Falls back to the reference where the kernel is absent or refuses. ``_probe`` + forces the kernel, since the probe is what decides that question.""" + if not _probe and not _fa2_usable(device): + return _sdpa_causal_attn(qf, kf, vf, device) Sq = qf.shape[1] Sk = kf.shape[1] qc, kc, vc = qf.contiguous(), kf.contiguous(), vf.contiguous() @@ -796,9 +875,14 @@ def _moe_layer(h, ld, fvk, device): tw, ti = torch.topk(logit, TOPK, -1) tw = tw / tw.sum(-1, keepdim=True) - if _USE_BT_MOE and x.shape[0] >= _M16_MIN_S: + # The block-scaled 4-bit MMA tiles are a build tier, not a given: a target + # whose toolchain has no block-scaled mma builds the weight-only tier + # instead. Ask the module what it has rather than assuming, so the tile + # choice degrades to the grouped GEMV instead of raising mid-prefill. + big = x.shape[0] >= _M16_MIN_S + if _USE_BT_MOE and big and hasattr(fvk, 'moe_blocktile_mma_sm120_bf16'): out = _moe_experts_bt(x, ti, tw, ld, fvk, device) - elif _USE_M16_MOE and x.shape[0] >= _M16_MIN_S: + elif _USE_M16_MOE and big and hasattr(fvk, 'moe_m16_mma_sm120_bf16'): out = _moe_experts_m16(x, ti, tw, ld, fvk, device) elif _USE_GROUPED_MOE: out = _moe_experts_grouped(x, ti, tw, ld, fvk, device) diff --git a/flash_rt/frontends/torch/qwen36_moe_rtx.py b/flash_rt/frontends/torch/qwen36_moe_rtx.py index cf4b3e85..0a5c3c96 100644 --- a/flash_rt/frontends/torch/qwen36_moe_rtx.py +++ b/flash_rt/frontends/torch/qwen36_moe_rtx.py @@ -282,6 +282,21 @@ class Qwen36MoeTextFrontendRtx(Nexn2TorchFrontendRtx): _MODEL_LABEL = "Qwen3.6-35B-A3B text" _USAGE_DOC = "docs/qwen36_moe_usage.md" + # The block-scaled 4-bit MMA tier is a build tier, and the prefill now picks + # its MoE tile from what the module actually has: without the tier it uses + # the weight-only grouped GEMV, which is slower on a long prompt and + # otherwise identical. So the tier is a performance requirement, not a + # correctness one, and demanding it here would refuse a build -- a Jetson + # one, whose toolchain has no block-scaled mma -- that runs this correctly. + _REQUIRED_KERNELS = tuple( + name for name in Nexn2TorchFrontendRtx._REQUIRED_KERNELS + if not name.startswith(('moe_blocktile_mma', 'moe_m16_mma')) + ) + ('moe_grouped_w4a16_sm120_bf16',) + + # Same for FA2: prefill and decode both probe the vendored kernel and fall + # back to a reference, so a target that builds FA4 instead still runs. + _REQUIRE_FA2 = False + def __init__(self, checkpoint_path: str, *, device: str = "cuda:0", max_seq: int = 2048, From ee73fcaed6cb9d7b7baeed7c00b2038d4204b90b Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 30 Jul 2026 12:42:05 -0400 Subject: [PATCH 29/85] Take the bank conflict out of the weight-only 4-bit GEMVs Both kernels stage the activation in shared memory at a 32-byte stride across lanes, which puts the eight lanes of a 128-bit load phase on four banks. On a 20-SM part the profiler reports 432,685 conflicts over 98,816 shared loads -- 2.41x the wavefronts the traffic needs -- with the kernel at 77% compute throughput against 37% memory, while the BF16 GEMV of the same shape reaches the memory roofline. Padding a block's footprint to 48 bytes lands the phase on eight distinct banks. The UE4M3 scale is decoded arithmetically at the same time. It was a 256-entry __constant__ LUT indexed by a per-lane byte, and constant memory serves one address per cycle, so a divergent index serialises; the decode is four integer ops and needs no table. Neither touches an arithmetic result, so the variants are accepted on bitwise equality against the kernels they stand in for, and the choice between them cannot move a token. They are added alongside rather than replacing, and a resolver picks the variant where the build has it (FLASHRT_QWEN35MOE_W4A16_EDGE=0 forces the original). Measured on sm_110: decode 70.07 -> 74.42 tok/s, 16/16 token-exact against the BF16 fixture. --- CMakeLists.txt | 1 + csrc/bindings.cpp | 28 +++ csrc/kernels/w4a16_edge_sm120.cu | 234 ++++++++++++++++++ csrc/kernels/w4a16_edge_sm120.cuh | 61 +++++ flash_rt/frontends/torch/_nexn2_rtx_decode.py | 11 +- .../frontends/torch/_nexn2_rtx_forward.py | 32 ++- tests/test_qwen36_moe_expert_cache.py | 5 +- 7 files changed, 364 insertions(+), 8 deletions(-) create mode 100644 csrc/kernels/w4a16_edge_sm120.cu create mode 100644 csrc/kernels/w4a16_edge_sm120.cuh diff --git a/CMakeLists.txt b/CMakeLists.txt index 8094d5cb..4e4a456b 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1522,6 +1522,7 @@ if(FLASHRT_ENABLE_QWEN35MOE_W4A16) target_sources(flash_rt_kernels PRIVATE csrc/kernels/w4a16_matvec_sm120.cu csrc/kernels/moe_grouped_w4a16_sm120.cu + csrc/kernels/w4a16_edge_sm120.cu csrc/kernels/w4a16_gemm_sm120.cu) target_compile_definitions(flash_rt_kernels PRIVATE FLASHRT_HAVE_QWEN35MOE_W4A16=1) diff --git a/csrc/bindings.cpp b/csrc/bindings.cpp index 0a108a61..a1c80ecb 100644 --- a/csrc/bindings.cpp +++ b/csrc/bindings.cpp @@ -191,6 +191,7 @@ extern "C" int cutlass_int8_rowwise_bf16out_t64x128( #ifdef FLASHRT_HAVE_QWEN35MOE_W4A16 #include "kernels/w4a16_matvec_sm120.cuh" #include "kernels/moe_grouped_w4a16_sm120.cuh" +#include "kernels/w4a16_edge_sm120.cuh" #include "kernels/w4a16_gemm_sm120.cuh" #endif // FLASHRT_HAVE_QWEN35MOE_W4A16 #ifdef FLASHRT_HAVE_QWEN35MOE_W4A4 @@ -5535,6 +5536,33 @@ PYBIND11_MODULE(flash_rt_kernels, m) { py::arg("eidx"), py::arg("D"), py::arg("slots"), py::arg("N"), py::arg("K"), py::arg("a_stride"), py::arg("w_stride"), py::arg("sfb_stride"), py::arg("stream") = 0); + + // Bitwise-identical variants tuned for a part where these two are compute + // bound rather than bandwidth bound. See w4a16_edge_sm120.cuh. + m.def("w4a16_matvec_edge_sm120_bf16", + [](uintptr_t x, uintptr_t W, uintptr_t sfb, uintptr_t out, + int N, int K, float alpha, uintptr_t stream) -> int { + return flash_rt::kernels::w4a16_matvec_edge_sm120_bf16( + to_ptr(x), to_ptr(W), to_ptr(sfb), to_ptr(out), + N, K, alpha, to_stream(stream)); + }, + py::arg("x"), py::arg("W"), py::arg("sfb"), py::arg("out"), + py::arg("N"), py::arg("K"), py::arg("alpha"), py::arg("stream") = 0); + + m.def("moe_grouped_w4a16_edge_sm120_bf16", + [](uintptr_t A, uintptr_t W, uintptr_t sfb, uintptr_t alpha, + uintptr_t eidx, uintptr_t D, int slots, int N, int K, + long a_stride, long w_stride, long sfb_stride, + uintptr_t stream) -> int { + return flash_rt::kernels::moe_grouped_w4a16_edge_sm120_bf16( + to_ptr(A), to_ptr(W), to_ptr(sfb), to_ptr(alpha), to_ptr(eidx), + to_ptr(D), slots, N, K, a_stride, w_stride, sfb_stride, + to_stream(stream)); + }, + py::arg("A"), py::arg("W"), py::arg("sfb"), py::arg("alpha"), + py::arg("eidx"), py::arg("D"), py::arg("slots"), py::arg("N"), + py::arg("K"), py::arg("a_stride"), py::arg("w_stride"), + py::arg("sfb_stride"), py::arg("stream") = 0); #endif // FLASHRT_HAVE_QWEN35MOE_W4A16 #ifdef FLASHRT_HAVE_QWEN35MOE_W4A4 diff --git a/csrc/kernels/w4a16_edge_sm120.cu b/csrc/kernels/w4a16_edge_sm120.cu new file mode 100644 index 00000000..d0f94d6e --- /dev/null +++ b/csrc/kernels/w4a16_edge_sm120.cu @@ -0,0 +1,234 @@ +// SPDX-License-Identifier: Apache-2.0 +// +// W4A16 GEMV variants for a bandwidth-poor part. See header for what differs +// and why. + +#include "kernels/w4a16_edge_sm120.cuh" + +#include +#include +#include "kernels/fp4_e2m1_compat.cuh" +#include +#include + +namespace flash_rt { +namespace kernels { + +namespace { + +constexpr int kWarps = 8; // 8 output rows / block +constexpr int kThreads = kWarps * 32; // 256 +constexpr int kUnroll = 4; // packed-weight loads in flight + +// A 16-element NVFP4 block is 16 bf16 of activation, 32 bytes. Held at that +// stride, the eight lanes of a 128-bit shared-load phase land on banks +// 0,8,16,24,0,8,16,24 -- four banks, two-way conflicted. At 48 bytes they land +// on 0,12,24,4,16,28,8,20: eight distinct banks. The 16 spare bytes per block +// cost K/2 bytes of shared memory (12 KB at K=4096) and buy back the 2.41x +// wavefront overhead the conflict was costing. +constexpr int kBlockSlots = 24; // bf16 slots per 16-element block +constexpr int kBlockInt4 = kBlockSlots / 8; // 3 int4 per block, 2 used + +// UE4M3 -> fp32 without a table. +// +// The value is (1 + m/8) * 2^(e-7) for e > 0, which is exactly an fp32 with +// exponent field e+120 and mantissa m<<20, and m * 2^-9 for e == 0. Four +// integer ops and a select, against a __constant__ load whose index differs +// per lane -- and constant memory serves one address per cycle, so a divergent +// index serialises the warp. +// +// Bit 7 is not a sign bit: UE4M3 is unsigned, and the quantizer's saturation +// byte 0xFE must decode to +448. +__device__ __forceinline__ float ue4m3_to_float(uint32_t v) { + const uint32_t e = (v >> 3) & 0xFu; + const uint32_t m = v & 0x7u; + const float normal = __uint_as_float(((e + 120u) << 23) | (m << 20)); + const float subnormal = static_cast(m) * (1.0f / 512.0f); + return e == 0u ? subnormal : normal; +} + +// SF swizzle byte offset, identical packing to bf16_weight_to_nvfp4_swizzled. +__device__ __forceinline__ int sf_off(int rb_ncs, int row_inner, int k_block) { + return (rb_ncs + (k_block >> 2)) * 512 + row_inner + (k_block & 3); +} + +// One NVFP4 block (16 elements / 8 packed bytes) dotted with 16 bf16 acts. +__device__ __forceinline__ float blockdot(uint64_t b_pack, + const __nv_bfloat162* xb2) { + float acc = 0.0f; +#pragma unroll + for (int j = 0; j < 8; ++j) { + const __half2_raw wr = flash_rt::fp4::cvt_e2m1x2_to_halfraw2( + static_cast(b_pack >> (j * 8))); + const float2 wf = __half22float2(*reinterpret_cast(&wr)); + const float2 xf = __bfloat1622float2(xb2[j]); + acc = fmaf(wf.x, xf.x, acc); + acc = fmaf(wf.y, xf.y, acc); + } + return acc; +} + +// Stage x into the padded shared layout: block b occupies int4 slots +// 3b and 3b+1, leaving 3b+2 as the padding that separates the banks. +__device__ __forceinline__ void stage_padded( + const __nv_bfloat16* __restrict__ x, __nv_bfloat16* x_sh, int K) { + const int4* x_i4 = reinterpret_cast(x); + int4* sh_i4 = reinterpret_cast(x_sh); + const int n_i4 = K >> 3; // 8 bf16 per int4, 2 per block + for (int j = threadIdx.x; j < n_i4; j += kThreads) + sh_i4[(j >> 1) * kBlockInt4 + (j & 1)] = x_i4[j]; +} + +// The K loop, shared by both entry points: 1 warp per output row, kUnroll +// packed-weight loads in flight. +__device__ __forceinline__ float row_dot( + const uint64_t* __restrict__ w_blk, const uint8_t* __restrict__ SFB, + const __nv_bfloat16* x_sh, int K_BLOCKS, int rb_ncs, int row_inner, + int lane) { + float acc = 0.0f; + int kb = lane; + const int step = 32 * kUnroll; + for (; kb + 32 * (kUnroll - 1) < K_BLOCKS; kb += step) { + uint64_t wv[kUnroll]; + float sf[kUnroll]; +#pragma unroll + for (int u = 0; u < kUnroll; ++u) wv[u] = w_blk[kb + 32 * u]; +#pragma unroll + for (int u = 0; u < kUnroll; ++u) + sf[u] = ue4m3_to_float( + __ldg(SFB + sf_off(rb_ncs, row_inner, kb + 32 * u))); +#pragma unroll + for (int u = 0; u < kUnroll; ++u) + acc += blockdot( + wv[u], reinterpret_cast( + x_sh + (size_t)(kb + 32 * u) * kBlockSlots)) * sf[u]; + } + for (; kb < K_BLOCKS; kb += 32) { + const float s = ue4m3_to_float( + __ldg(SFB + sf_off(rb_ncs, row_inner, kb))); + acc += blockdot( + w_blk[kb], reinterpret_cast( + x_sh + (size_t)kb * kBlockSlots)) * s; + } +#pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc += __shfl_xor_sync(0xffffffff, acc, off); + return acc; +} + +__global__ void w4a16_matvec_edge_kernel( + const __nv_bfloat16* __restrict__ x, + const uint8_t* __restrict__ W, + const uint8_t* __restrict__ SFB, + __nv_bfloat16* __restrict__ out, + float alpha, int N, int K, int n_col_super) { + extern __shared__ __nv_bfloat16 x_sh[]; + stage_padded(x, x_sh, K); + __syncthreads(); + + const int lane = threadIdx.x & 31; + const int row = blockIdx.x * kWarps + (threadIdx.x >> 5); + if (row >= N) return; + + const int rb = row >> 7; + const int ri = row & 127; + const float acc = row_dot( + reinterpret_cast(W + (size_t)row * (K >> 1)), SFB, + x_sh, K >> 4, rb * n_col_super, + (ri & 31) * 16 + ((ri >> 5) & 3) * 4, lane); + if (lane == 0) out[row] = __float2bfloat16(acc * alpha); +} + +// grid = (ceil(N/8), slots). Block computes 8 output rows of one slot. +__global__ void moe_grouped_w4a16_edge_kernel( + const __nv_bfloat16* __restrict__ A_stack, + const uint8_t* __restrict__ W_stack, + const uint8_t* __restrict__ SFB_stack, + const float* __restrict__ alpha_stack, + const int* __restrict__ expert_idx, + __nv_bfloat16* __restrict__ D, + int N, int K, int n_col_super, + long a_stride, long w_stride, long sfb_stride) { + const int slot = blockIdx.y; + const int e = expert_idx[slot]; + + extern __shared__ __nv_bfloat16 x_sh[]; + stage_padded(A_stack + (long)slot * a_stride, x_sh, K); + __syncthreads(); + + const int lane = threadIdx.x & 31; + const int row = blockIdx.x * kWarps + (threadIdx.x >> 5); + if (row >= N) return; + + const int rb = row >> 7; + const int ri = row & 127; + const float acc = row_dot( + reinterpret_cast( + W_stack + (long)e * w_stride + (size_t)row * (K >> 1)), + SFB_stack + (long)e * sfb_stride, x_sh, K >> 4, rb * n_col_super, + (ri & 31) * 16 + ((ri >> 5) & 3) * 4, lane); + if (lane == 0) + D[(long)slot * N + row] = __float2bfloat16(acc * alpha_stack[e]); +} + +// Shared memory for the padded stage: kBlockSlots bf16 per 16 elements. +inline size_t smem_bytes(int K) { + return (size_t)(K >> 4) * kBlockSlots * sizeof(__nv_bfloat16); +} + +} // namespace + +int w4a16_matvec_edge_sm120_bf16( + const void* x_bf16, + const void* W_packed, + const void* SFB, + void* out, + int N, + int K, + float alpha, + cudaStream_t stream) { + if (!x_bf16 || !W_packed || !SFB || !out) return 1; + if (N <= 0 || K <= 0 || (K & 15) != 0) return 2; + const int n_col_super = ((K >> 4) + 3) / 4; + w4a16_matvec_edge_kernel<<>>( + reinterpret_cast(x_bf16), + reinterpret_cast(W_packed), + reinterpret_cast(SFB), + reinterpret_cast<__nv_bfloat16*>(out), + alpha, N, K, n_col_super); + return 0; +} + +int moe_grouped_w4a16_edge_sm120_bf16( + const void* A_stack, + const void* W_stack, + const void* SFB_stack, + const void* alpha_stack, + const void* eidx, + void* D, + int slots, + int N, + int K, + long a_stride, + long w_stride, + long sfb_stride, + cudaStream_t stream) { + if (!A_stack || !W_stack || !SFB_stack || !alpha_stack || !eidx || !D) + return 1; + if (slots <= 0 || N <= 0 || K <= 0 || (K & 15) != 0) return 2; + const int n_col_super = ((K >> 4) + 3) / 4; + moe_grouped_w4a16_edge_kernel<<>>( + reinterpret_cast(A_stack), + reinterpret_cast(W_stack), + reinterpret_cast(SFB_stack), + reinterpret_cast(alpha_stack), + reinterpret_cast(eidx), + reinterpret_cast<__nv_bfloat16*>(D), + N, K, n_col_super, a_stride, w_stride, sfb_stride); + return 0; +} + +} // namespace kernels +} // namespace flash_rt diff --git a/csrc/kernels/w4a16_edge_sm120.cuh b/csrc/kernels/w4a16_edge_sm120.cuh new file mode 100644 index 00000000..1ede1670 --- /dev/null +++ b/csrc/kernels/w4a16_edge_sm120.cuh @@ -0,0 +1,61 @@ +// SPDX-License-Identifier: Apache-2.0 +// +// W4A16 GEMV variants for a bandwidth-poor part (Jetson-class Blackwell). +// +// Same math, same weight layout, same results as w4a16_matvec_sm120 and +// moe_grouped_w4a16_sm120 -- bit for bit. What differs is two things the +// profiler found on a 20-SM part with ~244 GB/s of memory, where the original +// pair sits at 51% of that while the BF16 GEMV of the same shape reaches 100%: +// +// 1. The staged activation is read from shared memory at a 32-byte stride +// across lanes, which puts eight lanes of a 128-bit load phase on four +// banks. Measured: 432,685 bank conflicts over 98,816 shared loads, 2.41x +// the wavefronts the traffic needs. Padding each block's footprint to 48 +// bytes lands the phase on eight distinct banks. +// +// 2. The UE4M3 block scale is decoded through a 256-entry __constant__ LUT +// indexed by a per-lane byte. Constant memory serves one address per +// cycle, so a divergent index serialises. The decode is four integer ops, +// so it does not need a table at all. +// +// Neither changes an arithmetic result, which is the point: the variant is +// accepted only if it is bitwise identical to the kernel it replaces. + +#pragma once + +#include + +namespace flash_rt { +namespace kernels { + +// y(1,N) = (x(1,K) bf16) . (W(N,K) NVFP4)^T, fp32 accumulate, bf16 out. +// Arguments and semantics are those of w4a16_matvec_sm120_bf16. +int w4a16_matvec_edge_sm120_bf16( + const void* x_bf16, + const void* W_packed, + const void* SFB, + void* out, + int N, + int K, + float alpha, + cudaStream_t stream); + +// Grouped per-slot GEMV: D[s,:] = A[s,:] . W[eidx[s]]^T * alpha[eidx[s]]. +// Arguments and semantics are those of moe_grouped_w4a16_sm120_bf16. +int moe_grouped_w4a16_edge_sm120_bf16( + const void* A, + const void* W, + const void* SFB, + const void* alpha, + const void* eidx, + void* D, + int slots, + int N, + int K, + long a_stride, + long w_stride, + long sfb_stride, + cudaStream_t stream); + +} // namespace kernels +} // namespace flash_rt diff --git a/flash_rt/frontends/torch/_nexn2_rtx_decode.py b/flash_rt/frontends/torch/_nexn2_rtx_decode.py index aabed348..f88e3d13 100644 --- a/flash_rt/frontends/torch/_nexn2_rtx_decode.py +++ b/flash_rt/frontends/torch/_nexn2_rtx_decode.py @@ -29,7 +29,8 @@ from flash_rt.frontends.torch._nexn2_rtx_forward import ( CONV, HD, HID, HK, HV, INTER, KD, KS, NKV, NQ, NV, ROPE, TOPK, VD, - _quant_act, build_rope_tables, nexn2_forward_nvfp4, + _quant_act, build_rope_tables, moe_grouped_w4a16, nexn2_forward_nvfp4, + w4a16_matvec, ) from flash_rt.frontends.torch._nexn2_rtx_nvfp4_weights import _sf_swz_bytes from flash_rt.hardware.rtx.attn_backend_nexn2 import RtxFlashAttnBackendNexn2 @@ -114,7 +115,7 @@ def _w4a16_mv(x1k, w_bf16, ld, key, fvk, device): ld[key + '_w4a16_a'] = float(og.item()) xc = x1k.contiguous() y = torch.empty(1, n, dtype=torch.bfloat16, device=device) - fvk.w4a16_matvec_sm120_bf16( + w4a16_matvec(fvk)( xc.data_ptr(), ld[pk].data_ptr(), ld[key + '_w4a16_sf'].data_ptr(), y.data_ptr(), n, k, ld[key + '_w4a16_a'], _cs()) return y @@ -586,7 +587,7 @@ def _moe_layer_decode(h, ld, state, fvk, device): # and faster at this scale (6.2 vs 8.2 us standalone). xc = x.contiguous() d_gu = torch.empty(TOPK, n_gu, dtype=torch.bfloat16, device=device) - fvk.moe_grouped_w4a16_sm120_bf16( + moe_grouped_w4a16(fvk)( xc.data_ptr(), gu_p.data_ptr(), gu_s.data_ptr(), gu_a.data_ptr(), idx.data_ptr(), d_gu.data_ptr(), TOPK, n_gu, HID, 0, gu_p[0].numel(), gu_s[0].numel(), s) @@ -595,7 +596,7 @@ def _moe_layer_decode(h, ld, state, fvk, device): g_, u_ = d_gu[:, :INTER], d_gu[:, INTER:] inter = _silu_mul(g_, u_, fvk, device).contiguous() d_dn = torch.empty(TOPK, n_dn, dtype=torch.bfloat16, device=device) - fvk.moe_grouped_w4a16_sm120_bf16( + moe_grouped_w4a16(fvk)( inter.data_ptr(), dn_p.data_ptr(), dn_s.data_ptr(), dn_a.data_ptr(), idx.data_ptr(), d_dn.data_ptr(), TOPK, n_dn, INTER, INTER, dn_p[0].numel(), dn_s[0].numel(), s) @@ -696,7 +697,7 @@ def decode_step(state, token_id, pos, fvk, device): # swizzled weight and the same scale factors, leaves the activation in # bf16 -- so it also skips the activation quantisation and its error -- and # lives in a tier that builds wherever the core does. - fvk.w4a16_matvec_sm120_bf16( + w4a16_matvec(fvk)( h.reshape(1, HID).contiguous().data_ptr(), p['lm_head_packed_t'].data_ptr(), p['lm_head_sf_t'].data_ptr(), logits.data_ptr(), vocab, HID, p['lm_head_alpha'], _cs()) diff --git a/flash_rt/frontends/torch/_nexn2_rtx_forward.py b/flash_rt/frontends/torch/_nexn2_rtx_forward.py index 9c659ff2..b0c8cac7 100644 --- a/flash_rt/frontends/torch/_nexn2_rtx_forward.py +++ b/flash_rt/frontends/torch/_nexn2_rtx_forward.py @@ -660,6 +660,34 @@ def _full_attn_layer(h, ld, ct, st, fvk, device, eps, cap=None, rank=None, return _proj(at, ld, 'o_proj', HID, fvk, device).reshape(B, S, HID) +# The two weight-only 4-bit GEMVs each have an "edge" variant that is bitwise +# identical -- it differs only in a shared-memory layout and in decoding the +# UE4M3 scale byte arithmetically rather than through a constant-memory lookup +# whose index differs per lane. Because the outputs are identical to the bit, +# choosing between them is purely a performance decision and cannot move a +# token, so preferring the variant needs no accuracy argument. Set +# FLASHRT_QWEN35MOE_W4A16_EDGE=0 to force the original. +_EDGE_W4A16 = _os.environ.get("FLASHRT_QWEN35MOE_W4A16_EDGE", "1") != "0" + + +def w4a16_matvec(fvk): + """The dense 4-bit GEMV entry point this build should call.""" + if _EDGE_W4A16: + fn = getattr(fvk, 'w4a16_matvec_edge_sm120_bf16', None) + if fn is not None: + return fn + return fvk.w4a16_matvec_sm120_bf16 + + +def moe_grouped_w4a16(fvk): + """The grouped per-slot 4-bit GEMV entry point this build should call.""" + if _EDGE_W4A16: + fn = getattr(fvk, 'moe_grouped_w4a16_edge_sm120_bf16', None) + if fn is not None: + return fn + return fvk.moe_grouped_w4a16_sm120_bf16 + + # Grouped MoE for prefill (on by default); set False to use the per-expert loop. _USE_GROUPED_MOE = True # M=16 tensor-core mma MoE: tokens are sorted into 16-row expert tiles and the @@ -840,14 +868,14 @@ def _moe_experts_grouped(x, ti, tw, ld, fvk, device): A = x[stok].contiguous() # (slots, HID) bf16 d_gu = torch.empty(slots, n_gu, dtype=torch.bfloat16, device=device) - fvk.moe_grouped_w4a16_sm120_bf16( + moe_grouped_w4a16(fvk)( A.data_ptr(), gu_p.data_ptr(), gu_s.data_ptr(), gu_a.data_ptr(), se.data_ptr(), d_gu.data_ptr(), slots, n_gu, HID, HID, gu_p[0].numel(), gu_s[0].numel(), 0) g, u = d_gu[:, :INTER], d_gu[:, INTER:] inter = _silu_mul(g, u, fvk, device).contiguous() d_dn = torch.empty(slots, n_dn, dtype=torch.bfloat16, device=device) - fvk.moe_grouped_w4a16_sm120_bf16( + moe_grouped_w4a16(fvk)( inter.data_ptr(), dn_p.data_ptr(), dn_s.data_ptr(), dn_a.data_ptr(), se.data_ptr(), d_dn.data_ptr(), slots, n_dn, INTER, INTER, dn_p[0].numel(), dn_s[0].numel(), 0) diff --git a/tests/test_qwen36_moe_expert_cache.py b/tests/test_qwen36_moe_expert_cache.py index 85222541..91bd46d1 100644 --- a/tests/test_qwen36_moe_expert_cache.py +++ b/tests/test_qwen36_moe_expert_cache.py @@ -187,7 +187,10 @@ def test_lm_head_has_a_path_without_the_sm120_only_kernel(): source = inspect.getsource(_nexn2_rtx_decode.decode_step) assert "hasattr(fvk, 'fp4_w4a4_mma_sm120_full_n_bf16out')" in source - assert "w4a16_matvec_sm120_bf16" in source + # The fallback goes through the resolver, which returns whichever + # weight-only GEMV this build has; both members of that pair are in the + # W4A16 tier, so either satisfies what this test is about. + assert "w4a16_matvec(fvk)" in source def test_staging_buffers_are_owned_for_the_duration_of_a_read(): From c955c598ba10f35577d25025ecd713cf2794eb7d Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 30 Jul 2026 12:56:31 -0400 Subject: [PATCH 30/85] Give each warp several rows so the 4-bit GEMVs have loads in flight With the bank conflict gone the dominant stall in situ is the global-load dependency -- 5.8 to 13 cycles per issued instruction against an ALU pipe at 27 to 50%. A warp owning one output row keeps only kUnroll eight-byte loads outstanding, and at K=512 not even that: K_BLOCKS is exactly 32, so the unrolled body never runs and the tail leaves a single load in flight. That is why the down projections sat at 38% of measured bandwidth while gate_up at K=2048 reached 67%. A warp now takes R consecutive rows and keeps R*kUnroll loads outstanding, R chosen so the product lands at 8 either way. The rows are 32-aligned by construction, so their scale offsets differ by a constant and cost no extra registers, and the per-row arithmetic is untouched -- same lane-to-block mapping, same order, same reduction -- so results stay bit-identical. Standalone at the shapes the decode actually issues, cold: experts down 38.6 -> 23.5 us (50% -> 80% of measured bandwidth), lm_head 1296 -> 1205 us (94%). End to end 74.42 -> 78.10 tok/s, 16/16 token-exact, time to first token 122 -> 110 ms. --- csrc/kernels/w4a16_edge_sm120.cu | 210 ++++++++++++++++++++++--------- 1 file changed, 148 insertions(+), 62 deletions(-) diff --git a/csrc/kernels/w4a16_edge_sm120.cu b/csrc/kernels/w4a16_edge_sm120.cu index d0f94d6e..ce25067c 100644 --- a/csrc/kernels/w4a16_edge_sm120.cu +++ b/csrc/kernels/w4a16_edge_sm120.cu @@ -79,43 +79,91 @@ __device__ __forceinline__ void stage_padded( sh_i4[(j >> 1) * kBlockInt4 + (j & 1)] = x_i4[j]; } -// The K loop, shared by both entry points: 1 warp per output row, kUnroll -// packed-weight loads in flight. -__device__ __forceinline__ float row_dot( - const uint64_t* __restrict__ w_blk, const uint8_t* __restrict__ SFB, - const __nv_bfloat16* x_sh, int K_BLOCKS, int rb_ncs, int row_inner, - int lane) { - float acc = 0.0f; +// The K loop, shared by both entry points: R output rows per warp, kUnroll +// packed-weight loads per row in flight. +// +// R exists because a warp with one row does not have enough memory-level +// parallelism on this part. In situ the dominant stall is the global-load +// dependency (long scoreboard, 5.8-13 cycles per issued instruction) while the +// ALU pipe sits at 27-50%: the loop is waiting on memory it has not asked for +// yet. Each lane keeps R*kUnroll eight-byte loads outstanding instead of +// kUnroll, and the K=512 shapes -- where K_BLOCKS is exactly 32, so the +// unrolled body never runs and the tail leaves ONE load in flight -- get the +// whole factor from R. +// +// The rows a warp takes are consecutive and 32-aligned by construction, so +// their scale offsets differ by a constant and cost no extra registers. The +// per-row arithmetic is untouched: same lane-to-block mapping, same order, same +// reduction, so the result is bit-identical to R = 1. +template +__device__ __forceinline__ void row_dot( + const uint64_t* __restrict__ w_row0, size_t row_stride_u64, + const uint8_t* __restrict__ SFB, const __nv_bfloat16* x_sh, + int K_BLOCKS, int rb_ncs, int row_inner, int lane, float (&acc)[R]) { +#pragma unroll + for (int r = 0; r < R; ++r) acc[r] = 0.0f; + int kb = lane; const int step = 32 * kUnroll; for (; kb + 32 * (kUnroll - 1) < K_BLOCKS; kb += step) { - uint64_t wv[kUnroll]; - float sf[kUnroll]; + uint64_t wv[R][kUnroll]; + float sf[R][kUnroll]; +#pragma unroll + for (int r = 0; r < R; ++r) #pragma unroll - for (int u = 0; u < kUnroll; ++u) wv[u] = w_blk[kb + 32 * u]; + for (int u = 0; u < kUnroll; ++u) + wv[r][u] = w_row0[r * row_stride_u64 + kb + 32 * u]; #pragma unroll - for (int u = 0; u < kUnroll; ++u) - sf[u] = ue4m3_to_float( - __ldg(SFB + sf_off(rb_ncs, row_inner, kb + 32 * u))); + for (int r = 0; r < R; ++r) #pragma unroll - for (int u = 0; u < kUnroll; ++u) - acc += blockdot( - wv[u], reinterpret_cast( - x_sh + (size_t)(kb + 32 * u) * kBlockSlots)) * sf[u]; + for (int u = 0; u < kUnroll; ++u) + sf[r][u] = ue4m3_to_float(__ldg( + SFB + sf_off(rb_ncs, row_inner + 16 * r, kb + 32 * u))); +#pragma unroll + for (int r = 0; r < R; ++r) +#pragma unroll + for (int u = 0; u < kUnroll; ++u) + acc[r] += blockdot( + wv[r][u], reinterpret_cast( + x_sh + (size_t)(kb + 32 * u) * kBlockSlots)) + * sf[r][u]; } for (; kb < K_BLOCKS; kb += 32) { - const float s = ue4m3_to_float( - __ldg(SFB + sf_off(rb_ncs, row_inner, kb))); - acc += blockdot( - w_blk[kb], reinterpret_cast( - x_sh + (size_t)kb * kBlockSlots)) * s; + uint64_t wv[R]; + float sf[R]; +#pragma unroll + for (int r = 0; r < R; ++r) wv[r] = w_row0[r * row_stride_u64 + kb]; +#pragma unroll + for (int r = 0; r < R; ++r) + sf[r] = ue4m3_to_float( + __ldg(SFB + sf_off(rb_ncs, row_inner + 16 * r, kb))); +#pragma unroll + for (int r = 0; r < R; ++r) + acc[r] += blockdot( + wv[r], reinterpret_cast( + x_sh + (size_t)kb * kBlockSlots)) * sf[r]; } #pragma unroll - for (int off = 16; off > 0; off >>= 1) - acc += __shfl_xor_sync(0xffffffff, acc, off); - return acc; + for (int r = 0; r < R; ++r) +#pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc[r] += __shfl_xor_sync(0xffffffff, acc[r], off); } +// Rows per warp. Enough outstanding loads to cover the memory latency without +// spending so many registers that occupancy pays for it: R * (loads per row) +// lands at 8 either way, since the unrolled body runs only when K_BLOCKS +// reaches 32 * kUnroll. +constexpr int kRowsBig = 2; // K >= 2048: kUnroll fires, 2 * 4 = 8 +constexpr int kRowsSmall = 8; // K < 2048: tail only, 8 * 1 = 8 + +// The row block a warp owns must not straddle a 32-row scale group, or +// row_inner + 16 * r stops describing the swizzle. Warps take R consecutive +// rows starting at a multiple of R, so this holds for any R dividing 32. +static_assert(32 % kRowsBig == 0 && 32 % kRowsSmall == 0, + "rows per warp must divide the 32-row scale group"); + +template __global__ void w4a16_matvec_edge_kernel( const __nv_bfloat16* __restrict__ x, const uint8_t* __restrict__ W, @@ -127,19 +175,25 @@ __global__ void w4a16_matvec_edge_kernel( __syncthreads(); const int lane = threadIdx.x & 31; - const int row = blockIdx.x * kWarps + (threadIdx.x >> 5); - if (row >= N) return; - - const int rb = row >> 7; - const int ri = row & 127; - const float acc = row_dot( - reinterpret_cast(W + (size_t)row * (K >> 1)), SFB, - x_sh, K >> 4, rb * n_col_super, - (ri & 31) * 16 + ((ri >> 5) & 3) * 4, lane); - if (lane == 0) out[row] = __float2bfloat16(acc * alpha); + const int row0 = (blockIdx.x * kWarps + (threadIdx.x >> 5)) * R; + if (row0 >= N) return; + + const int rb = row0 >> 7; + const int ri = row0 & 127; + float acc[R]; + row_dot( + reinterpret_cast(W + (size_t)row0 * (K >> 1)), + (size_t)(K >> 1) / 8, SFB, x_sh, K >> 4, rb * n_col_super, + (ri & 31) * 16 + ((ri >> 5) & 3) * 4, lane, acc); + if (lane == 0) { +#pragma unroll + for (int r = 0; r < R; ++r) + if (row0 + r < N) out[row0 + r] = __float2bfloat16(acc[r] * alpha); + } } -// grid = (ceil(N/8), slots). Block computes 8 output rows of one slot. +// grid = (ceil(N/(8*R)), slots). Block computes 8*R output rows of one slot. +template __global__ void moe_grouped_w4a16_edge_kernel( const __nv_bfloat16* __restrict__ A_stack, const uint8_t* __restrict__ W_stack, @@ -157,18 +211,24 @@ __global__ void moe_grouped_w4a16_edge_kernel( __syncthreads(); const int lane = threadIdx.x & 31; - const int row = blockIdx.x * kWarps + (threadIdx.x >> 5); - if (row >= N) return; + const int row0 = (blockIdx.x * kWarps + (threadIdx.x >> 5)) * R; + if (row0 >= N) return; - const int rb = row >> 7; - const int ri = row & 127; - const float acc = row_dot( + const int rb = row0 >> 7; + const int ri = row0 & 127; + float acc[R]; + row_dot( reinterpret_cast( - W_stack + (long)e * w_stride + (size_t)row * (K >> 1)), - SFB_stack + (long)e * sfb_stride, x_sh, K >> 4, rb * n_col_super, - (ri & 31) * 16 + ((ri >> 5) & 3) * 4, lane); - if (lane == 0) - D[(long)slot * N + row] = __float2bfloat16(acc * alpha_stack[e]); + W_stack + (long)e * w_stride + (size_t)row0 * (K >> 1)), + (size_t)(K >> 1) / 8, SFB_stack + (long)e * sfb_stride, x_sh, K >> 4, + rb * n_col_super, (ri & 31) * 16 + ((ri >> 5) & 3) * 4, lane, acc); + if (lane == 0) { + const float a = alpha_stack[e]; +#pragma unroll + for (int r = 0; r < R; ++r) + if (row0 + r < N) + D[(long)slot * N + row0 + r] = __float2bfloat16(acc[r] * a); + } } // Shared memory for the padded stage: kBlockSlots bf16 per 16 elements. @@ -176,6 +236,14 @@ inline size_t smem_bytes(int K) { return (size_t)(K >> 4) * kBlockSlots * sizeof(__nv_bfloat16); } +// Rows per warp for this K, dropping to 1 when N cannot fill even one warp's +// worth. Above 32 rows a warp would straddle a scale group; the constants +// enforce that, this only picks between them. +inline int rows_per_warp(int N, int K) { + const int r = (K >= 2048) ? kRowsBig : kRowsSmall; + return (N >= kWarps * r) ? r : 1; +} + } // namespace int w4a16_matvec_edge_sm120_bf16( @@ -190,13 +258,21 @@ int w4a16_matvec_edge_sm120_bf16( if (!x_bf16 || !W_packed || !SFB || !out) return 1; if (N <= 0 || K <= 0 || (K & 15) != 0) return 2; const int n_col_super = ((K >> 4) + 3) / 4; - w4a16_matvec_edge_kernel<<>>( - reinterpret_cast(x_bf16), - reinterpret_cast(W_packed), - reinterpret_cast(SFB), - reinterpret_cast<__nv_bfloat16*>(out), - alpha, N, K, n_col_super); + const auto* xp = reinterpret_cast(x_bf16); + const auto* wp = reinterpret_cast(W_packed); + const auto* sp = reinterpret_cast(SFB); + auto* op = reinterpret_cast<__nv_bfloat16*>(out); + const size_t smem = smem_bytes(K); +#define FLASHRT_LAUNCH_MATVEC(R) \ + w4a16_matvec_edge_kernel<<>>( \ + xp, wp, sp, op, alpha, N, K, n_col_super) + switch (rows_per_warp(N, K)) { + case kRowsBig: FLASHRT_LAUNCH_MATVEC(kRowsBig); break; + case kRowsSmall: FLASHRT_LAUNCH_MATVEC(kRowsSmall); break; + default: FLASHRT_LAUNCH_MATVEC(1); break; + } +#undef FLASHRT_LAUNCH_MATVEC return 0; } @@ -218,15 +294,25 @@ int moe_grouped_w4a16_edge_sm120_bf16( return 1; if (slots <= 0 || N <= 0 || K <= 0 || (K & 15) != 0) return 2; const int n_col_super = ((K >> 4) + 3) / 4; - moe_grouped_w4a16_edge_kernel<<>>( - reinterpret_cast(A_stack), - reinterpret_cast(W_stack), - reinterpret_cast(SFB_stack), - reinterpret_cast(alpha_stack), - reinterpret_cast(eidx), - reinterpret_cast<__nv_bfloat16*>(D), - N, K, n_col_super, a_stride, w_stride, sfb_stride); + const auto* ap = reinterpret_cast(A_stack); + const auto* wp = reinterpret_cast(W_stack); + const auto* sp = reinterpret_cast(SFB_stack); + const auto* alp = reinterpret_cast(alpha_stack); + const auto* ep = reinterpret_cast(eidx); + auto* dp = reinterpret_cast<__nv_bfloat16*>(D); + const size_t smem = smem_bytes(K); +#define FLASHRT_LAUNCH_GROUPED(R) \ + moe_grouped_w4a16_edge_kernel \ + <<>>( \ + ap, wp, sp, alp, ep, dp, N, K, n_col_super, \ + a_stride, w_stride, sfb_stride) + switch (rows_per_warp(N, K)) { + case kRowsBig: FLASHRT_LAUNCH_GROUPED(kRowsBig); break; + case kRowsSmall: FLASHRT_LAUNCH_GROUPED(kRowsSmall); break; + default: FLASHRT_LAUNCH_GROUPED(1); break; + } +#undef FLASHRT_LAUNCH_GROUPED return 0; } From 6f60dc2f7939170215f965d75940870caf7ddfbf Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 30 Jul 2026 13:08:05 -0400 Subject: [PATCH 31/85] Tune rows per warp per entry point, and drop the KV broadcast copy The two GEMVs do not want the same number of rows per warp. More rows means more loads in flight and more registers, and the point where that stops paying differs: measured cold at the shapes the decode issues, the dense GEMV peaks at 2 (q_proj 47.9 us against 53.9 at 4, lm_head 1205 against 1366) while the grouped one peaks at 4 (gate_up 42.9 against 47.8 at 2). The grouped launch carries a slot per grid row, so it has fewer blocks per row tile and leans harder on per-thread parallelism. They no longer share a constant. The reference decode attention also stopped broadcasting the KV to the query head count, which materialised it only to read it once; native GQA does the same thing without the copy, in fp32 either way so the numerics are untouched. Worth much less than the traffic suggested -- the KV is small enough to stay in L2 at decode -- but it is strictly less work. 78.10 -> 79.05 tok/s, 16/16 token-exact. --- csrc/kernels/w4a16_edge_sm120.cu | 35 ++++++++++++--------- flash_rt/hardware/rtx/attn_backend_nexn2.py | 26 ++++++++++----- 2 files changed, 39 insertions(+), 22 deletions(-) diff --git a/csrc/kernels/w4a16_edge_sm120.cu b/csrc/kernels/w4a16_edge_sm120.cu index ce25067c..7ff4c693 100644 --- a/csrc/kernels/w4a16_edge_sm120.cu +++ b/csrc/kernels/w4a16_edge_sm120.cu @@ -150,17 +150,24 @@ __device__ __forceinline__ void row_dot( acc[r] += __shfl_xor_sync(0xffffffff, acc[r], off); } -// Rows per warp. Enough outstanding loads to cover the memory latency without -// spending so many registers that occupancy pays for it: R * (loads per row) -// lands at 8 either way, since the unrolled body runs only when K_BLOCKS -// reaches 32 * kUnroll. -constexpr int kRowsBig = 2; // K >= 2048: kUnroll fires, 2 * 4 = 8 +// Rows per warp. More rows means more outstanding loads and more registers, so +// the useful value is where the added parallelism stops paying for the +// occupancy it costs -- and that turned out to differ between the two entry +// points, which is why they do not share a constant. Measured at the shapes +// the decode issues, cold: the dense GEMV peaks at 2 (q_proj 47.9 us against +// 53.9 at 4, lm_head 1205 against 1366) while the grouped one peaks at 4 +// (gate_up 42.9 against 47.8 at 2). The grouped launch carries a slot per grid +// row, so it has fewer blocks per row tile and leans harder on what each +// thread keeps in flight. +constexpr int kRowsDense = 2; // K >= 2048: kUnroll fires, 2 * 4 = 8 +constexpr int kRowsGrouped = 4; // K >= 2048: 4 * 4 = 16 constexpr int kRowsSmall = 8; // K < 2048: tail only, 8 * 1 = 8 // The row block a warp owns must not straddle a 32-row scale group, or // row_inner + 16 * r stops describing the swizzle. Warps take R consecutive // rows starting at a multiple of R, so this holds for any R dividing 32. -static_assert(32 % kRowsBig == 0 && 32 % kRowsSmall == 0, +static_assert(32 % kRowsDense == 0 && 32 % kRowsGrouped == 0 + && 32 % kRowsSmall == 0, "rows per warp must divide the 32-row scale group"); template @@ -239,8 +246,8 @@ inline size_t smem_bytes(int K) { // Rows per warp for this K, dropping to 1 when N cannot fill even one warp's // worth. Above 32 rows a warp would straddle a scale group; the constants // enforce that, this only picks between them. -inline int rows_per_warp(int N, int K) { - const int r = (K >= 2048) ? kRowsBig : kRowsSmall; +inline int rows_per_warp(int N, int K, int rows_big) { + const int r = (K >= 2048) ? rows_big : kRowsSmall; return (N >= kWarps * r) ? r : 1; } @@ -267,8 +274,8 @@ int w4a16_matvec_edge_sm120_bf16( w4a16_matvec_edge_kernel<<>>( \ xp, wp, sp, op, alpha, N, K, n_col_super) - switch (rows_per_warp(N, K)) { - case kRowsBig: FLASHRT_LAUNCH_MATVEC(kRowsBig); break; + switch (rows_per_warp(N, K, kRowsDense)) { + case kRowsDense: FLASHRT_LAUNCH_MATVEC(kRowsDense); break; case kRowsSmall: FLASHRT_LAUNCH_MATVEC(kRowsSmall); break; default: FLASHRT_LAUNCH_MATVEC(1); break; } @@ -307,10 +314,10 @@ int moe_grouped_w4a16_edge_sm120_bf16( dim3(kThreads), smem, stream>>>( \ ap, wp, sp, alp, ep, dp, N, K, n_col_super, \ a_stride, w_stride, sfb_stride) - switch (rows_per_warp(N, K)) { - case kRowsBig: FLASHRT_LAUNCH_GROUPED(kRowsBig); break; - case kRowsSmall: FLASHRT_LAUNCH_GROUPED(kRowsSmall); break; - default: FLASHRT_LAUNCH_GROUPED(1); break; + switch (rows_per_warp(N, K, kRowsGrouped)) { + case kRowsGrouped: FLASHRT_LAUNCH_GROUPED(kRowsGrouped); break; + case kRowsSmall: FLASHRT_LAUNCH_GROUPED(kRowsSmall); break; + default: FLASHRT_LAUNCH_GROUPED(1); break; } #undef FLASHRT_LAUNCH_GROUPED return 0; diff --git a/flash_rt/hardware/rtx/attn_backend_nexn2.py b/flash_rt/hardware/rtx/attn_backend_nexn2.py index ad333d95..951208ea 100644 --- a/flash_rt/hardware/rtx/attn_backend_nexn2.py +++ b/flash_rt/hardware/rtx/attn_backend_nexn2.py @@ -165,17 +165,27 @@ def _sdpa(self, layer_idx: int, q_seq: int, kv_seq: int, """Reference attention, for a device the vendored kernel refuses.""" import torch.nn.functional as F - q = self.Q_buf[:, :q_seq] + q = self.Q_buf[:, :q_seq].transpose(1, 2).float() k = self.K_cache[layer_idx:layer_idx + 1, :kv_seq] v = self.V_cache[layer_idx:layer_idx + 1, :kv_seq] + # Broadcasting the KV to the query head count materialises it: at + # decode that is 8x2xkv_seqx256 floats twice per layer, ~48 MB a step + # across the ten full-attention layers, purely to be read once. Native + # GQA does the same thing without the copy. fp32 either way, so the + # numerics are untouched -- this path seeds a token-exact decode. groups = self.NUM_Q_HEADS // self.NUM_KV_HEADS - out = F.scaled_dot_product_attention( - q.transpose(1, 2).float(), - k.repeat_interleave(groups, dim=2).transpose(1, 2).float(), - v.repeat_interleave(groups, dim=2).transpose(1, 2).float(), - is_causal=q_seq > 1, - scale=softmax_scale, - ).transpose(1, 2) + try: + out = F.scaled_dot_product_attention( + q, k.transpose(1, 2).float(), v.transpose(1, 2).float(), + is_causal=q_seq > 1, scale=softmax_scale, enable_gqa=True, + ).transpose(1, 2) + except TypeError: # torch without native GQA + out = F.scaled_dot_product_attention( + q, + k.repeat_interleave(groups, dim=2).transpose(1, 2).float(), + v.repeat_interleave(groups, dim=2).transpose(1, 2).float(), + is_causal=q_seq > 1, scale=softmax_scale, + ).transpose(1, 2) self.O_buf[:, :q_seq].copy_(out.to(self.O_buf.dtype)) # ── Layer cache pointer math ── From b206be1826d9005a94dcf0cc453b8e091c916e63 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 30 Jul 2026 13:16:51 -0400 Subject: [PATCH 32/85] Size the GEMV block to the part, not to the card it was written for Eight warps per block gives 16 output rows at two rows per warp, so the dense decode shapes make too few blocks to fill a 20-SM part: N=2048 lands 128 blocks where roughly 160 are resident at once, and the launch never reaches a full wave. Two warps quadruples the block count at the same rows per warp. Measured end to end, since the standalone bench has been wrong about this class of question twice: 1 warp 79.09, 2 warps 80.53, 4 warps 79.80, 8 warps 79.05 tok/s. One warp loses the shared staging entirely and every block re-stages the activation; eight leaves the machine underfilled. Four rows per warp was tried for the dense GEMV as well and is worse both standalone and end to end (76.21 tok/s), so the two entry points keep their separate values. 79.05 -> 80.53 tok/s, still bitwise identical to the kernels it stands in for, 16/16 token-exact. --- csrc/kernels/w4a16_edge_sm120.cu | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/csrc/kernels/w4a16_edge_sm120.cu b/csrc/kernels/w4a16_edge_sm120.cu index 7ff4c693..da1cb9b5 100644 --- a/csrc/kernels/w4a16_edge_sm120.cu +++ b/csrc/kernels/w4a16_edge_sm120.cu @@ -16,7 +16,7 @@ namespace kernels { namespace { -constexpr int kWarps = 8; // 8 output rows / block +constexpr int kWarps = 2; // output-row groups per block constexpr int kThreads = kWarps * 32; // 256 constexpr int kUnroll = 4; // packed-weight loads in flight From 1ca7849ecb67564ee7f1a7ce09f0cefba322f3ca Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 30 Jul 2026 13:37:38 -0400 Subject: [PATCH 33/85] Load and run the multi-token-prediction head The checkpoint carries an MTP head that nothing read: one full-attention layer plus its own 256-expert MoE, with four tensors around it. Its layer has exactly a full-attention layer's keys, so the per-layer loader is now a function both it and the model go through -- loading it a second way would be a second thing to keep correct. It is opt-in, since it is another layer's weights (resident 21.435 -> 21.931 GiB) and is only useful with a verifier. Its experts never stream: a bundle holds the model's own layers, so there is nowhere for the head's to stream from. It takes the KV slot after the model's ten rather than owning a second cache, so the attention backend's layer count is now a constructor argument. Two things the measurement settled that reading the source would not have. The head's fc takes [embed_norm, hidden_norm]; nothing in the shapes says which, fc being square in the concatenated width, and the wrong order accepts 0 of 48. And chaining a second draft has to feed the head its own pre-final-norm hidden state rather than the model's stale one -- that is worth 0.531 against 0.208 on the second draft -- so the head returns it. Measured over 96 tokens: first draft 0.771, then 0.562, 0.250, 0.073. The decode step now parks its pre-final-norm hidden state in a fixed buffer, one 4 KB device copy, unconditionally -- making it conditional would put a Python branch inside the captured region. --- flash_rt/frontends/torch/_nexn2_rtx_decode.py | 85 ++++++++++++++++++- .../torch/_nexn2_rtx_nvfp4_weights.py | 50 +++++++++-- flash_rt/frontends/torch/nexn2_rtx.py | 7 +- flash_rt/hardware/rtx/attn_backend_nexn2.py | 9 +- tests/test_qwen36_moe_expert_cache.py | 4 +- 5 files changed, 142 insertions(+), 13 deletions(-) diff --git a/flash_rt/frontends/torch/_nexn2_rtx_decode.py b/flash_rt/frontends/torch/_nexn2_rtx_decode.py index f88e3d13..723c00d9 100644 --- a/flash_rt/frontends/torch/_nexn2_rtx_decode.py +++ b/flash_rt/frontends/torch/_nexn2_rtx_decode.py @@ -209,8 +209,18 @@ def __init__(self, handles, max_seq, device): torch.zeros(1, CONV, KS - 1, dtype=bf16, device=device) for _ in range(nlin)] - # Full-attn KV cache. - self.attn = RtxFlashAttnBackendNexn2(max_seq=self.max_seq, max_q_seq=1) + # Full-attn KV cache. A loaded draft head is one more full-attention + # layer and takes the slot after the model's own. + self.mtp = p.get('mtp') + self.mtp_rank = nfull if self.mtp is not None else None + self.attn = RtxFlashAttnBackendNexn2( + max_seq=self.max_seq, max_q_seq=1, + num_full_layers=nfull + (1 if self.mtp is not None else 0)) + # The pre-final-norm hidden state of the last step, which is what the + # draft head reads. Written every step whether or not one is loaded: + # a 4 KB device copy, and making it conditional would put a Python + # branch inside the captured region. + self.last_hidden = torch.zeros(HID, dtype=bf16, device=device) # RoPE tables for the whole window. theta = float(p['rope_theta']) @@ -280,6 +290,11 @@ def __init__(self, handles, max_seq, device): self.router_trace = None self.moe_input_trace = None self._active_layer = -1 + # Which half of the draft head's fc input carries the hidden state. + # The checkpoint does not say and fc is square in the concatenated + # width, so this is settled by measuring acceptance both ways. + self.mtp_hidden_first = ( + _qwen35moe_env("MTP_HIDDEN_FIRST", "1") != "0") # Set to an ExpertCache to read the routed experts from storage. Only # meaningful when the loader skipped them; see _moe_experts_streamed. self.expert_cache = None @@ -659,11 +674,25 @@ def decode_step(state, token_id, pos, fvk, device): state._active_layer = L h = res + _moe_layer_decode(n, ld, state, fvk, device) + # The pre-final-norm hidden state is what a DeepSeek-V3-style draft head + # consumes. Keeping it in a fixed buffer costs one 4 KB device copy and + # survives graph capture, unlike reading it out per step. + state.last_hidden.copy_(h.reshape(HID)) h = _rms_fvk(h, p['final_norm_w_t'], fvk, device, state.eps) # lm_head as NVFP4 W4A16: 4x less weight read (1GB -> 0.25GB) via the # hand-tuned mma (3.1x the bf16 GEMV; the CUTLASS widen is M=1-broken). # The weight is quantised once during the eager seed (cached on p), so # the captured graph only runs the activation quant + fp4 GEMM. + return _lm_head(state, h, fvk, device) + + +def _lm_head(state, h, fvk, device): + """Project a hidden state to logits over the full vocabulary. + + Taken out of decode_step so the speculative draft head, which ends the + same way, does not carry a second copy of the quantise-once bookkeeping. + """ + p = state.handles.ptrs vocab = p['vocab_size'] logits = torch.empty(1, vocab, dtype=torch.bfloat16, device=device) if not state.lm_head_nvfp4: @@ -704,6 +733,58 @@ def decode_step(state, token_id, pos, fvk, device): return logits +def mtp_draft(state, token_id, pos, fvk, device, *, hidden=None): + """Draft the token after next with the MTP head. + + A DeepSeek-V3 single-module head: it sees the main model's last hidden + state for position p-1 and the token emitted at p, and predicts p+1. The + layer under it is an ordinary full-attention layer with its own MoE, so it + runs through the same per-layer code as the model -- which is the point of + loading it through the same loader. + + ``hidden`` defaults to the buffer the last decode step wrote. The head + carries its own KV at the same absolute positions as the model, so calling + this advances that cache and nothing else. + """ + p = state.handles.ptrs + mtp = state.mtp + if mtp is None: + raise RuntimeError( + 'no MTP head is loaded; build the frontend with speculation ' + 'enabled so the loader reads it') + ld = mtp['layer'] + h_prev = state.last_hidden if hidden is None else hidden + + e = F.embedding(token_id.view(1, 1), p['embed_w_t']).reshape(1, HID) + hn = _rms_fvk(h_prev.reshape(1, HID), mtp['pre_h_w_t'], fvk, device, + state.eps) + en = _rms_fvk(e, mtp['pre_e_w_t'], fvk, device, state.eps) + # Which half goes first is a checkpoint convention, not something the + # shapes pin down -- fc is square in the concatenated width. It is + # measured, not assumed: the wrong order drafts noise. + cat = (torch.cat([hn, en], -1) if state.mtp_hidden_first + else torch.cat([en, hn], -1)) + h = _dense_mv(cat, mtp['fc_w_t'], mtp, 'fc_w_t', state, fvk, device) + + res = h + n = _rms_fvk(h, ld['input_norm_w_t'], fvk, device, state.eps) + h = res + _decode_full(n, ld, state, state.mtp_rank, pos, fvk, device) + res = h + n = _rms_fvk(h, ld['post_norm_w_t'], fvk, device, state.eps) + prev_layer, state._active_layer = state._active_layer, None + try: + h = res + _moe_layer_decode(n, ld, state, fvk, device) + finally: + state._active_layer = prev_layer + # Return the state before the head's own final norm as well: chaining a + # second draft means feeding the head what the model would have fed it, + # and that is a pre-final-norm hidden state. Handing it the model's stale + # one instead costs real acceptance -- measured 0.208 against 0.539 on the + # second draft. + return _lm_head(state, _rms_fvk(h, mtp['norm_w_t'], fvk, device, + state.eps), fvk, device), h.reshape(HID) + + def seed_prefill(state, input_ids, fvk, device): """Run the decode step over prompt tokens 0..S-1, building all state. diff --git a/flash_rt/frontends/torch/_nexn2_rtx_nvfp4_weights.py b/flash_rt/frontends/torch/_nexn2_rtx_nvfp4_weights.py index 1427871e..7fc977bb 100644 --- a/flash_rt/frontends/torch/_nexn2_rtx_nvfp4_weights.py +++ b/flash_rt/frontends/torch/_nexn2_rtx_nvfp4_weights.py @@ -260,6 +260,7 @@ def extract_weights_nexn2_nvfp4( device: str = 'cuda:0', quant_scope: str = 'experts', stream_experts: bool = False, + load_mtp: bool = False, ) -> WeightHandles: """Build :class:`WeightHandles` from a Nex-N2-mini BF16 ckpt directory. @@ -310,10 +311,13 @@ def extract_weights_nexn2_nvfp4( handles_d, wmap, device) # ── Per-layer ── - per_layer: list = [None] * num_layers - for i in range(num_layers): - lp = f'model.language_model.layers.{i}.' - ltype = layer_types[i] + def _load_layer(lp: str, ltype: str, *, streamed: bool = None) -> dict: + """Build one layer's weight dict from its checkpoint prefix. + + Taken out of the loop so the MTP head can use it: its layer lives under + a different prefix but has exactly a full-attention layer's keys, and + loading it a second way would be a second thing to keep correct. + """ ld: dict = {'type': ltype, 'quant_format': 'nvfp4'} _bf16_from_ckpt(handles, ld, 'input_norm_w', lp + 'input_layernorm.weight', @@ -353,13 +357,19 @@ def extract_weights_nexn2_nvfp4( _proj_load(handles, ld, 'out_proj', gp + 'out_proj.weight', handles_d, wmap, fvk, device, quantize=quant_main) else: - raise ValueError(f'layer {i}: unknown layer_type {ltype!r}') + raise ValueError(f'{lp}: unknown layer_type {ltype!r}') # Every layer has a MoE FFN (mlp_only_layers is empty). _load_moe(handles, ld, lp, handles_d, wmap, fvk, device, n_experts, quantize_shared=quant_main, - stream_experts=stream_experts) - per_layer[i] = ld + stream_experts=(stream_experts if streamed is None + else streamed)) + return ld + + per_layer: list = [None] * num_layers + for i in range(num_layers): + per_layer[i] = _load_layer( + f'model.language_model.layers.{i}.', layer_types[i]) handles.ptrs['layers'] = per_layer handles.ptrs['vocab_size'] = vocab @@ -377,6 +387,30 @@ def extract_weights_nexn2_nvfp4( handles.ptrs['quant_format'] = 'nvfp4' handles.ptrs['quant_scope'] = quant_scope handles.ptrs['ckpt_dir'] = ckpt_dir - handles.ptrs['mtp'] = None # MTP weights not in the base ckpt + # ── Multi-token-prediction head ── + # + # One full-attention layer plus its own 256-expert MoE, under `mtp.`, with + # four head-level tensors around it. It drafts the token after next from + # the main model's last hidden state and the token just emitted, which is + # only useful with a verifier, so it is opt-in: it costs another layer's + # worth of weights and a KV slot. + handles.ptrs['mtp'] = None + if load_mtp: + if not _has(wmap, 'mtp.fc.weight'): + raise RuntimeError( + f'{ckpt_dir} has no MTP head (mtp.fc.weight is absent), so ' + 'speculative drafting cannot be built from it.') + # A bundle holds the model's own layers, so the head's experts have + # nowhere to stream from and stay resident whatever the model does. + mtp: dict = {'layer': _load_layer('mtp.layers.0.', 'full_attention', + streamed=False)} + _bf16_from_ckpt(handles, mtp, 'fc_w', 'mtp.fc.weight', + handles_d, wmap, device) + for name, key in (('norm_w', 'mtp.norm.weight'), + ('pre_h_w', 'mtp.pre_fc_norm_hidden.weight'), + ('pre_e_w', 'mtp.pre_fc_norm_embedding.weight')): + _bf16_from_ckpt(handles, mtp, name, key, + handles_d, wmap, device, fold_one=True) + handles.ptrs['mtp'] = mtp handles.ptrs['dflash'] = None return handles diff --git a/flash_rt/frontends/torch/nexn2_rtx.py b/flash_rt/frontends/torch/nexn2_rtx.py index 65f8bc58..c887266c 100644 --- a/flash_rt/frontends/torch/nexn2_rtx.py +++ b/flash_rt/frontends/torch/nexn2_rtx.py @@ -131,6 +131,10 @@ def __init__(self, checkpoint_path: str, *, # Set by a subclass that streams the routed experts from a bundle # instead of holding them; see _nexn2_rtx_decode._moe_experts_streamed. self._stream_experts = getattr(self, '_stream_experts', False) + # Read by the loader before any weight is touched, like the above: + # the draft head is another layer's worth of weights and is only + # useful with a verifier, so nothing loads it unless asked. + self._load_mtp = getattr(self, '_load_mtp', False) self._tokenizer = None self._prompt_ids = None self._pipeline: Nexn2Pipeline | None = None @@ -188,7 +192,8 @@ def _build_kernelized_nvfp4(self) -> None: self._weights = extract_weights_nexn2_nvfp4( self.checkpoint_path, fvk, device=self.device, quant_scope=self._quant_scope, - stream_experts=self._stream_experts) + stream_experts=self._stream_experts, + load_mtp=self._load_mtp) @property def tokenizer(self): diff --git a/flash_rt/hardware/rtx/attn_backend_nexn2.py b/flash_rt/hardware/rtx/attn_backend_nexn2.py index 951208ea..c965ce62 100644 --- a/flash_rt/hardware/rtx/attn_backend_nexn2.py +++ b/flash_rt/hardware/rtx/attn_backend_nexn2.py @@ -46,12 +46,19 @@ class RtxFlashAttnBackendNexn2: NUM_KV_HEADS = 2 HEAD_DIM = 256 - def __init__(self, max_seq: int, max_q_seq: int = 1, dtype=None): + def __init__(self, max_seq: int, max_q_seq: int = 1, dtype=None, + num_full_layers: int | None = None): import torch self._torch = torch bf16 = dtype if dtype is not None else torch.bfloat16 d = "cuda" + # The model's ten full-attention layers by default. A speculative + # draft head is one more full-attention layer carrying its own KV, so + # it asks for an extra slot rather than owning a second cache. + self.NUM_FULL_LAYERS = ( + int(num_full_layers) if num_full_layers is not None + else type(self).NUM_FULL_LAYERS) self._max_seq = int(max_seq) self._max_q_seq = int(max_q_seq) diff --git a/tests/test_qwen36_moe_expert_cache.py b/tests/test_qwen36_moe_expert_cache.py index 91bd46d1..c591a62e 100644 --- a/tests/test_qwen36_moe_expert_cache.py +++ b/tests/test_qwen36_moe_expert_cache.py @@ -184,7 +184,9 @@ def test_lm_head_has_a_path_without_the_sm120_only_kernel(): from flash_rt.frontends.torch import _nexn2_rtx_decode - source = inspect.getsource(_nexn2_rtx_decode.decode_step) + # The projection moved into its own function when the draft head, which + # ends the same way, needed it too. + source = inspect.getsource(_nexn2_rtx_decode._lm_head) assert "hasattr(fvk, 'fp4_w4a4_mma_sm120_full_n_bf16out')" in source # The fallback goes through the resolver, which returns whichever From 2d6123278cd571bda45e138a94ad111d0dc387ab Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 30 Jul 2026 13:50:54 -0400 Subject: [PATCH 34/85] Draft, verify and rewind a speculative window Verification is one batched forward over the window, so the dense weights are read once instead of once per token. Acceptance keeps the longest prefix the model's own argmax agrees with. Rewinding is the part a transformer does not have to do. The KV cache needs nothing -- it is written by absolute position, so a rejected tail is overwritten by whatever comes next -- but the gated-DeltaNet recurrent and conv states have already absorbed the whole window. So the verify block runs its scan a token at a time and keeps each intermediate, and acceptance restores the one belonging to the prefix it kept. At this block length that is a handful of extra launches and a few MB of copies; deriving the state afterwards would mean re-running the block's projections or reconstructing the recurrence from saved inputs, both dearer. The capture is gated on the verify block so prefill, which runs the same layer code, does not pay for snapshots it never uses. Measured acceptance tracks what the trace-based model predicted: 1.86 / 2.50 / 2.71 tokens kept per window at K = 1 / 2 / 3, against 1.77 / 2.33 / 2.58. It is not yet faster. The verifier goes through the prefill forward, which reads BF16 dense weights and a BF16 lm_head -- roughly three times the traffic the window was priced at -- and runs eager against a captured-graph baseline. Routing verification through the 4-bit weights at small M is the next step. --- flash_rt/frontends/torch/_nexn2_rtx_decode.py | 119 ++++++++++++++++++ .../frontends/torch/_nexn2_rtx_forward.py | 45 +++++++ flash_rt/frontends/torch/qwen36_moe_rtx.py | 26 ++++ 3 files changed, 190 insertions(+) diff --git a/flash_rt/frontends/torch/_nexn2_rtx_decode.py b/flash_rt/frontends/torch/_nexn2_rtx_decode.py index 723c00d9..9e03c1fa 100644 --- a/flash_rt/frontends/torch/_nexn2_rtx_decode.py +++ b/flash_rt/frontends/torch/_nexn2_rtx_decode.py @@ -290,6 +290,14 @@ def __init__(self, handles, max_seq, device): self.router_trace = None self.moe_input_trace = None self._active_layer = -1 + # Per-token recurrent/conv snapshots for a speculative window, sized + # on first use. Only allocated when speculation runs: 30 layers of + # (NV, HK, HV) bf16 per window slot. + self.spec_states = None + self.spec_conv = None + # Set only around a verify block. Prefill runs the same layer code and + # would otherwise pay for -- and overrun -- snapshots it never uses. + self.spec_capture = False # Which half of the draft head's fc input carries the hidden state. # The checkpoint does not say and fc is square in the concatenated # width, so this is settled by measuring acceptance both ways. @@ -785,6 +793,117 @@ def mtp_draft(state, token_id, pos, fvk, device, *, hidden=None): state.eps), fvk, device), h.reshape(HID) +def _ensure_spec_buffers(state, window, device): + """Allocate the per-token state snapshots a window of `window` needs.""" + have = (state.spec_states is not None + and len(state.spec_states[0]) >= window) + if have: + return + state.spec_states = [ + [torch.empty(NV, HK, HV, dtype=torch.bfloat16, device=device) + for _ in range(window)] + for _ in range(state.n_lin)] + state.spec_conv = [ + [torch.empty(1, CONV, KS - 1, dtype=torch.bfloat16, device=device) + for _ in range(window)] + for _ in range(state.n_lin)] + + +def _rewind_to(state, kept): + """Put the recurrent and conv states where `kept` tokens of the window end. + + The KV cache needs nothing: it is written by absolute position, so the + rejected tail is simply overwritten by whatever comes next. The recurrent + state is the opposite -- it has already absorbed the whole window -- which + is what the per-token snapshots are for. + """ + for rank in range(state.n_lin): + state.lin_state[rank].copy_(state.spec_states[rank][kept - 1]) + state.lin_conv_state[rank].copy_(state.spec_conv[rank][kept - 1]) + + +def spec_decode_step(state, token_id, pos, k, fvk, device): + """One speculative step: draft k tokens, verify k+1 positions, keep a prefix. + + Returns (tokens, logits, next_pos) where `tokens` are the ids actually + emitted -- between 1 and k+1 of them -- and `logits` are the ones that + produced the last of them, so the caller can continue from it. + + The window is [token_id, draft_1 .. draft_k]. Verifying it is one batched + forward over k+1 positions, which reads the dense weights once instead of + k+1 times; that, and nothing about the drafts being good, is where the time + comes from. A draft is kept only when the model's own argmax at that + position agrees with it, so the emitted sequence is exactly what plain + greedy decoding would have produced. + """ + window = k + 1 + _ensure_spec_buffers(state, window, device) + + # Draft k tokens off the state as it stands. Each chained draft feeds the + # head its own hidden state, the same kind of input the model gives it. + drafts = [] + tok = token_id.view(1) + hidden = state.last_hidden + for j in range(k): + d_logits, hidden = mtp_draft(state, tok, pos + j, fvk, device, + hidden=hidden) + tok = d_logits[0].argmax().view(1) + drafts.append(tok) + + ids = torch.cat([token_id.view(1)] + drafts).view(1, window) + state.spec_capture = True + try: + logits, hidden = nexn2_forward_nvfp4( + state.handles, ids, fvk, device, cap=state, pos_offset=pos, + last_logits_only=False, return_hidden=True) + finally: + state.spec_capture = False + logits = logits.reshape(window, -1) + argmax = logits.argmax(-1) + + # Keep the longest prefix the model agrees with. Row j predicts position + # pos+j+1, so it is compared against draft j+1. + kept = 1 + accepted = argmax[:k] == ids.view(window)[1:] + for j in range(k): + if not bool(accepted[j]): + break + kept += 1 + + if kept < window: + _rewind_to(state, kept) + # The draft head reads the pre-final-norm hidden state of the last emitted + # position. The batched forward has it for every position in the window; + # without this the next window would draft off a state from before it. + state.last_hidden.copy_(hidden[kept - 1]) + tokens = ids.view(window)[1:kept].tolist() + [int(argmax[kept - 1])] + return tokens, logits[kept - 1:kept], pos + kept + + +def generate_greedy_spec(state, input_ids, max_new_tokens, k, fvk, device): + """Greedy decode through the draft-and-verify step. + + Emits exactly what generate_greedy would; the tokens are a check on the + machinery, not an approximation of it. + """ + logits = seed_prefill(state, input_ids, fvk, device) + pos = input_ids.view(-1).shape[0] + out = [] + state.spec_windows = 0 + state.spec_kept = 0 + while len(out) < max_new_tokens: + nxt = logits[0].argmax().view(1) + tokens, logits, pos = spec_decode_step( + state, nxt, pos, k, fvk, device) + emitted = [int(nxt)] + tokens[:-1] + state.spec_windows += 1 + state.spec_kept += len(emitted) + out.extend(emitted) + # The step returns the logits that produced its last token, which the + # next window drafts from. + return out[:max_new_tokens] + + def seed_prefill(state, input_ids, fvk, device): """Run the decode step over prompt tokens 0..S-1, building all state. diff --git a/flash_rt/frontends/torch/_nexn2_rtx_forward.py b/flash_rt/frontends/torch/_nexn2_rtx_forward.py index b0c8cac7..5c94959d 100644 --- a/flash_rt/frontends/torch/_nexn2_rtx_forward.py +++ b/flash_rt/frontends/torch/_nexn2_rtx_forward.py @@ -455,6 +455,9 @@ def _gdn_layer(h, ld, fvk, device, eps, cap=None, rank=None, cap.lin_state[rank].copy_(state) cs = mixed[0, S - (KS - 1):S, :].transpose(0, 1).contiguous() cap.lin_conv_state[rank].copy_(cs.unsqueeze(0)) + if getattr(cap, 'spec_capture', False): + _capture_per_token_state(cap, rank, S, init_state, conv_hist, + mixed, qb, kb, vb, g_out, bo, fvk, device) cf = core.reshape(-1, HV).contiguous() zf = z.reshape(-1, HV).to(torch.bfloat16).contiguous() @@ -703,6 +706,48 @@ def moe_grouped_w4a16(fvk): _USE_BT_MOE = True +def _capture_per_token_state(cap, rank, S, init_state, conv_hist, mixed, + qb, kb, vb, g_out, bo, fvk, device): + """Record what the recurrent state would be after each token of a block. + + A verified speculative window is accepted up to some prefix, and the layer + that has to be rewound is this one: the KV cache is a cursor, but the + recurrent and conv states are not -- they have already absorbed every token + of the block, including the rejected tail. + + Rather than re-deriving them afterwards, run the scan a token at a time and + keep each intermediate. The block is a handful of tokens, so this is a few + extra launches and a few MB of state copies; recovering the state any other + way means either re-running the block's projections or reconstructing the + recurrence from saved inputs, both of which cost more than they save at + this length. + """ + state = (init_state.clone() if init_state is not None + else torch.zeros(NV, HK, HV, dtype=torch.bfloat16, device=device)) + q3 = qb.reshape(S, NV, HK).contiguous() + k3 = kb.reshape(S, NV, HK).contiguous() + v3 = vb.reshape(S, NV, HV).contiguous() + g2 = g_out.reshape(S, NV).contiguous() + b2 = bo.reshape(S, NV).contiguous() + core1 = torch.empty(1, NV, HV, dtype=torch.bfloat16, device=device) + for t in range(S): + fvk.gdn_recurrent_seq_sm120_bf16( + q3[t:t + 1].data_ptr(), k3[t:t + 1].data_ptr(), + v3[t:t + 1].data_ptr(), g2[t:t + 1].data_ptr(), + b2[t:t + 1].data_ptr(), state.data_ptr(), core1.data_ptr(), + 1, NV, HK, True, 0) + cap.spec_states[rank][t].copy_(state) + + # Conv state after t+1 tokens: the last KS-1 entries of the block's inputs + # preceded by whatever history the block started from. + prev = (conv_hist[0] if conv_hist is not None + else torch.zeros(mixed.shape[-1], KS - 1, + dtype=mixed.dtype, device=device)) + hist = torch.cat([prev, mixed[0].transpose(0, 1)], dim=1) + for t in range(S): + cap.spec_conv[rank][t].copy_(hist[:, t + 1:t + KS].unsqueeze(0)) + + def _moe_experts_m16(x, ti, tw, ld, fvk, device): """Routed experts via the M=16 tensor-core block-scaled mma. Sort the S*TOPK assignments by expert, pack into zero-padded 16-row tiles, quant once diff --git a/flash_rt/frontends/torch/qwen36_moe_rtx.py b/flash_rt/frontends/torch/qwen36_moe_rtx.py index 0a5c3c96..f8c2eabb 100644 --- a/flash_rt/frontends/torch/qwen36_moe_rtx.py +++ b/flash_rt/frontends/torch/qwen36_moe_rtx.py @@ -322,6 +322,32 @@ def __init__(self, checkpoint_path: str, *, ) self._checkpoint_contract = contract + def generate_spec(self, max_new_tokens: int, *, k: int = 2): + """Greedy decode through draft-and-verify with the MTP head. + + Emits exactly what ``generate`` emits: a draft is kept only where the + model's own argmax agrees with it, so this is a speed change and + nothing else. Requires the frontend to have loaded the head. + """ + if self._prompt_ids is None: + raise ValueError("call set_prompt(...) before generate_spec()") + if not self._load_mtp: + raise RuntimeError( + "speculative decoding needs the MTP draft head; construct " + "the frontend with _load_mtp so the loader reads it") + + from flash_rt.frontends.torch._nexn2_rtx_decode import ( + Nexn2DecodeState, + generate_greedy_spec, + ) + + if self._decode_state is None: + self._decode_state = Nexn2DecodeState( + self._weights, self._user_max_seq, self.device) + return generate_greedy_spec( + self._decode_state, self._prompt_ids, max_new_tokens, k, + self._fvk, self.device) + def generate(self, max_new_tokens: int, *, do_sample: bool = False): """Generate tokens with the shared qwen3_5_moe CUDA Graph path.""" if self._prompt_ids is None: From 22cb219b028cedee1fec7db898a7dfb2eaaedd3f Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 30 Jul 2026 13:58:12 -0400 Subject: [PATCH 35/85] Give the verify block a 4-bit path, and measure that it is not the problem The window's dense projections and lm_head can now read the 4-bit weights instead of BF16, which is what the traffic model said the pass needed. Measured at K=2 it goes 36.15 -> 33.38 tok/s: a quarter of the bytes, and slower. That is worth keeping as a result rather than a change. The verify pass is not bandwidth bound, so what it reads is not what to fix. It runs eager against a baseline whose every step is a captured graph, and each window also pays for K drafts that each project the full 248320-wide vocabulary -- the pruning the cost model showed was the difference between 1.3x and 1.6x, and which is not built yet. So the switch is off by default with the measurement recorded next to it, to be judged again once the pass is capture-bound rather than overhead-bound. --- flash_rt/frontends/torch/_nexn2_rtx_decode.py | 4 ++- .../frontends/torch/_nexn2_rtx_forward.py | 36 +++++++++++++++++-- 2 files changed, 36 insertions(+), 4 deletions(-) diff --git a/flash_rt/frontends/torch/_nexn2_rtx_decode.py b/flash_rt/frontends/torch/_nexn2_rtx_decode.py index 9e03c1fa..4b17789e 100644 --- a/flash_rt/frontends/torch/_nexn2_rtx_decode.py +++ b/flash_rt/frontends/torch/_nexn2_rtx_decode.py @@ -30,7 +30,7 @@ from flash_rt.frontends.torch._nexn2_rtx_forward import ( CONV, HD, HID, HK, HV, INTER, KD, KS, NKV, NQ, NV, ROPE, TOPK, VD, _quant_act, build_rope_tables, moe_grouped_w4a16, nexn2_forward_nvfp4, - w4a16_matvec, + set_spec_verify, w4a16_matvec, ) from flash_rt.frontends.torch._nexn2_rtx_nvfp4_weights import _sf_swz_bytes from flash_rt.hardware.rtx.attn_backend_nexn2 import RtxFlashAttnBackendNexn2 @@ -852,12 +852,14 @@ def spec_decode_step(state, token_id, pos, k, fvk, device): ids = torch.cat([token_id.view(1)] + drafts).view(1, window) state.spec_capture = True + set_spec_verify(True) try: logits, hidden = nexn2_forward_nvfp4( state.handles, ids, fvk, device, cap=state, pos_offset=pos, last_logits_only=False, return_hidden=True) finally: state.spec_capture = False + set_spec_verify(False) logits = logits.reshape(window, -1) argmax = logits.argmax(-1) diff --git a/flash_rt/frontends/torch/_nexn2_rtx_forward.py b/flash_rt/frontends/torch/_nexn2_rtx_forward.py index 5c94959d..0c096bd6 100644 --- a/flash_rt/frontends/torch/_nexn2_rtx_forward.py +++ b/flash_rt/frontends/torch/_nexn2_rtx_forward.py @@ -84,8 +84,13 @@ def _proj(x2d, ld, base, n, fvk, device): return _nvfp4_gemm(x2d, ld[base + '_packed'], ld[base + '_sf'], ld[base + '_alpha'], n, fvk, device) w = ld[base + '_w_t'] - if (_DENSE_W4A16 and x2d.shape[0] >= 64 and (w.shape[0] % 64) == 0 - and (x2d.shape[1] % 64) == 0): + # A speculative verify block is a handful of rows, so the M>=64 heuristic + # below would send it to the BF16 GEMM -- four times the weight bytes, on + # the pass whose whole purpose is to read the weights once. It wants the + # 4-bit weight for the same reason decode does: at this M the cost is + # traffic, not throughput. + if ((_SPEC_VERIFY or (_DENSE_W4A16 and x2d.shape[0] >= 64)) + and (w.shape[0] % 64) == 0 and (x2d.shape[1] % 64) == 0): return _gemm_w4a16(x2d, w, ld, base + '_w_t', fvk, device) if (_DENSE_W16A16 and x2d.shape[0] >= _DENSE_BF16_MIN_M and (x2d.shape[1] % 64) == 0): @@ -121,6 +126,25 @@ def _gemm_w16a16(x2d, w, fvk, device): # path needs a bf16-*weight* GEMM (repurpose this kernel's 2.18x structure). _DENSE_W4A16 = False +# Set only around a speculative verify block; see _proj and the lm_head below. +# +# Default off, because it was measured and it loses: routing the verify block's +# dense projections and lm_head through the 4-bit weights takes K=2 from 36.15 +# to 33.38 tok/s, despite reading a quarter of the bytes. That is the useful +# part of the result -- it says the verify pass is not bandwidth bound, so the +# thing to fix is not what it reads. It runs eager, against a baseline whose +# every step is a captured graph, and a window also pays for K drafts that each +# project the full 248320-wide vocabulary. Capture the verify block and prune +# the draft head's vocabulary first; revisit this after, when the pass is +# actually reading-bound and the switch can be judged on its merits. +_SPEC_VERIFY_W4A16 = False +_SPEC_VERIFY = False + + +def set_spec_verify(on: bool) -> None: + global _SPEC_VERIFY + _SPEC_VERIFY = bool(on) and _SPEC_VERIFY_W4A16 + # BF16 tensor-core dense projections (vs the default fp32/TF32 matmul). The # experts-scope q/k/v/o/out/shared/router projections dominate the prefill # profile as fp32 GEMMs; bf16 inputs with fp32 accumulate roughly halve that @@ -1059,7 +1083,13 @@ def nexn2_forward_nvfp4(handles, input_ids, fvk, device, cap=None, # last position first when only the seeding logit is needed (avoids the # (S, vocab) materialisation that dominates long-context prefill memory). h_lm = h[0][-1:].contiguous() if last_logits_only else h[0] - logits = _gemm_w16a16(h_lm, p['lm_head_w_t'], fvk, device) + if _SPEC_VERIFY: + # The lm_head is the single largest weight; at BF16 it is a gigabyte a + # verify, which on its own outweighs what the window saves. + logits = _gemm_w4a16(h_lm, p['lm_head_w_t'], p, 'lm_head_w_t', + fvk, device) + else: + logits = _gemm_w16a16(h_lm, p['lm_head_w_t'], fvk, device) if return_hidden: return logits, hidden return logits From 0c6d31cf693913d6077a0d83ad3b8f1603d375ef Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 30 Jul 2026 14:17:02 -0400 Subject: [PATCH 36/85] Capture the speculative window, and stop quantising the draft head Two fixes that turned out to be the same mistake: treating the draft head and the verify block as if they were the model. The window -- k drafts and the verify -- is now one captured graph. Each draft's token is chosen on the device and written straight into the buffer the next draft reads, so the chain never leaves the GPU and the only host decision left is how much of it to keep. Capturing it needed the prefill forward to stop hard-coding stream 0 in its fvk calls: that names a stream that is not being captured, and capture_end reports it as an illegal access. The decode path had always used the current stream; this forward was written for prefill, which is never captured, and a verify block is. The draft head now runs on its BF16 weights rather than the runtime W4A16 the model's projections take. A draft is one layer, so its weights are a rounding error against the window's traffic, while its accuracy decides whether a verified position is spent or wasted -- and the sibling frontends keep their heads BF16 for exactly that reason. Measured over 96 tokens, per-draft acceptance goes 0.771/0.562/0.250/0.073 -> 0.833/0.615/0.323/0.156, so a K=3 window keeps 2.77 tokens instead of 2.58. Captured and on BF16 the window also stops diverging from plain greedy: same tokens, 16/16 against the fixture. K=2 is 39.77 tok/s against 36.15 eager. --- flash_rt/frontends/torch/_nexn2_rtx_decode.py | 169 +++++++++++++----- .../frontends/torch/_nexn2_rtx_forward.py | 91 ++++++---- 2 files changed, 182 insertions(+), 78 deletions(-) diff --git a/flash_rt/frontends/torch/_nexn2_rtx_decode.py b/flash_rt/frontends/torch/_nexn2_rtx_decode.py index 4b17789e..d6eb84b1 100644 --- a/flash_rt/frontends/torch/_nexn2_rtx_decode.py +++ b/flash_rt/frontends/torch/_nexn2_rtx_decode.py @@ -298,6 +298,12 @@ def __init__(self, handles, max_seq, device): # Set only around a verify block. Prefill runs the same layer code and # would otherwise pay for -- and overrun -- snapshots it never uses. self.spec_capture = False + # One captured graph per (pos, window): the KV slots, attention length + # and RoPE slice are baked per position exactly as the decode graph's + # are. Same LRU bound, since each graph owns a memory pool. + self._spec_graphs = collections.OrderedDict() + self._spec_tokens = None + self._spec_argmax = None # Which half of the draft head's fc input carries the hidden state. # The checkpoint does not say and fc is square in the concatenated # width, so this is settled by measuring acceptance both ways. @@ -763,6 +769,21 @@ def mtp_draft(state, token_id, pos, fvk, device, *, hidden=None): ld = mtp['layer'] h_prev = state.last_hidden if hidden is None else hidden + # The draft head runs on its BF16 weights, not on the runtime W4A16 the + # model's own projections take. A draft is one layer -- its weights are a + # rounding error against the window's traffic -- while its accuracy is the + # whole point, since a rejected draft costs a verified position. The + # sibling frontends keep the head BF16 for the same reason; this path had + # been quantising it along with everything else. + was_w4a16, state.dense_w4a16 = state.dense_w4a16, False + try: + return _mtp_draft_bf16(state, mtp, ld, p, token_id, h_prev, pos, + fvk, device) + finally: + state.dense_w4a16 = was_w4a16 + + +def _mtp_draft_bf16(state, mtp, ld, p, token_id, h_prev, pos, fvk, device): e = F.embedding(token_id.view(1, 1), p['embed_w_t']).reshape(1, HID) hn = _rms_fvk(h_prev.reshape(1, HID), mtp['pre_h_w_t'], fvk, device, state.eps) @@ -794,7 +815,12 @@ def mtp_draft(state, token_id, pos, fvk, device, *, hidden=None): def _ensure_spec_buffers(state, window, device): - """Allocate the per-token state snapshots a window of `window` needs.""" + """Allocate what a window of `window` tokens needs.""" + if state._spec_tokens is None or state._spec_tokens.numel() < window: + state._spec_tokens = torch.zeros(window, dtype=torch.long, + device=device) + state._spec_argmax = torch.zeros(window, dtype=torch.long, + device=device) have = (state.spec_states is not None and len(state.spec_states[0]) >= window) if have: @@ -822,64 +848,127 @@ def _rewind_to(state, kept): state.lin_conv_state[rank].copy_(state.spec_conv[rank][kept - 1]) -def spec_decode_step(state, token_id, pos, k, fvk, device): - """One speculative step: draft k tokens, verify k+1 positions, keep a prefix. - - Returns (tokens, logits, next_pos) where `tokens` are the ids actually - emitted -- between 1 and k+1 of them -- and `logits` are the ones that - produced the last of them, so the caller can continue from it. +def _spec_block(state, pos, k, fvk, device): + """The whole window as one dependency chain: k drafts, then the verify. - The window is [token_id, draft_1 .. draft_k]. Verifying it is one batched - forward over k+1 positions, which reads the dense weights once instead of - k+1 times; that, and nothing about the drafts being good, is where the time - comes from. A draft is kept only when the model's own argmax at that - position agrees with it, so the emitted sequence is exactly what plain - greedy decoding would have produced. + Written to be capturable end to end. Each draft's token is chosen on the + device -- ``qwen36_argmax_bf16`` writes it straight into the token buffer + the next draft reads -- so the chain never leaves the GPU, and the only + host decision left is how much of the window to keep. """ + vocab = state.handles.ptrs['vocab_size'] + toks = state._spec_tokens window = k + 1 - _ensure_spec_buffers(state, window, device) - # Draft k tokens off the state as it stands. Each chained draft feeds the - # head its own hidden state, the same kind of input the model gives it. - drafts = [] - tok = token_id.view(1) hidden = state.last_hidden for j in range(k): - d_logits, hidden = mtp_draft(state, tok, pos + j, fvk, device, - hidden=hidden) - tok = d_logits[0].argmax().view(1) - drafts.append(tok) + d_logits, hidden = mtp_draft(state, toks[j:j + 1], pos + j, fvk, + device, hidden=hidden) + fvk.qwen36_argmax_bf16(d_logits.data_ptr(), + toks[j + 1:j + 2].data_ptr(), 1, vocab, _cs()) - ids = torch.cat([token_id.view(1)] + drafts).view(1, window) state.spec_capture = True set_spec_verify(True) try: - logits, hidden = nexn2_forward_nvfp4( - state.handles, ids, fvk, device, cap=state, pos_offset=pos, - last_logits_only=False, return_hidden=True) + logits, hid = nexn2_forward_nvfp4( + state.handles, toks[:window].view(1, window), fvk, device, + cap=state, pos_offset=pos, last_logits_only=False, + return_hidden=True) finally: state.spec_capture = False set_spec_verify(False) logits = logits.reshape(window, -1) - argmax = logits.argmax(-1) + fvk.qwen36_argmax_bf16(logits.data_ptr(), state._spec_argmax.data_ptr(), + window, vocab, _cs()) + return hid + + +def _ensure_spec_graph(state, pos, k, fvk, device): + """Capture the draft-and-verify window at ``pos``, or return the cached one. + + Everything the block mutates is snapshotted and restored around the warmup + and capture runs -- the recurrent and conv states, the KV rows the window + writes across every rank including the draft head's, and the drafted token + slots -- so a later replay advances from the true pre-window state rather + than from whatever the capture left behind. + """ + key = (pos, k) + cached = state._spec_graphs.get(key) + if cached is not None: + state._spec_graphs.move_to_end(key) + return cached + + window = k + 1 + snap_lin = [t.clone() for t in state.lin_state] + snap_conv = [t.clone() for t in state.lin_conv_state] + snap_k = state.attn.K_cache[:, pos:pos + window].clone() + snap_v = state.attn.V_cache[:, pos:pos + window].clone() + snap_tok = state._spec_tokens.clone() + + def _restore(): + for i, t in enumerate(state.lin_state): + t.copy_(snap_lin[i]) + for i, t in enumerate(state.lin_conv_state): + t.copy_(snap_conv[i]) + state.attn.K_cache[:, pos:pos + window].copy_(snap_k) + state.attn.V_cache[:, pos:pos + window].copy_(snap_v) + state._spec_tokens.copy_(snap_tok) + + with torch.no_grad(): # settle allocator, kernel order, and the + for _ in range(2): # weight quantisation the draft does lazily + _spec_block(state, pos, k, fvk, device) + _restore() + + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g, stream=state._graph_stream, + pool=state._graph_pool), torch.no_grad(): + hid = _spec_block(state, pos, k, fvk, device) + with torch.no_grad(): + _restore() + + state._spec_graphs[key] = (g, hid) + cap = state.graph_cache_max + if cap > 0 and len(state._spec_graphs) > cap: + state._spec_graphs.popitem(last=False) + return state._spec_graphs[key] + + +def spec_decode_step(state, token_id, pos, k, fvk, device): + """One speculative step: draft k tokens, verify k+1 positions, keep a prefix. + + Returns (tokens, next_pos) where `tokens` are the ids actually emitted -- + between 1 and k+1 of them. + + Verifying the window is one batched forward over k+1 positions, which reads + the dense weights once instead of k+1 times; that, and nothing about the + drafts being good, is where the time comes from. A draft is kept only where + the model's own argmax agrees with it, so the emitted sequence is what that + verifier's greedy decode would have produced. + """ + window = k + 1 + _ensure_spec_buffers(state, window, device) + state._spec_tokens[0].copy_(token_id.view(1)[0]) + + g, hid = _ensure_spec_graph(state, pos, k, fvk, device) + g.replay() - # Keep the longest prefix the model agrees with. Row j predicts position - # pos+j+1, so it is compared against draft j+1. + # One D2H for the whole decision: the drafted ids and what the model said + # at each position. Everything before this stayed on the device. + drafted = state._spec_tokens[:window].tolist() + argmax = state._spec_argmax[:window].tolist() kept = 1 - accepted = argmax[:k] == ids.view(window)[1:] for j in range(k): - if not bool(accepted[j]): + if argmax[j] != drafted[j + 1]: break kept += 1 if kept < window: _rewind_to(state, kept) # The draft head reads the pre-final-norm hidden state of the last emitted - # position. The batched forward has it for every position in the window; - # without this the next window would draft off a state from before it. - state.last_hidden.copy_(hidden[kept - 1]) - tokens = ids.view(window)[1:kept].tolist() + [int(argmax[kept - 1])] - return tokens, logits[kept - 1:kept], pos + kept + # position. Without this the next window would draft off a stale one. + state.last_hidden.copy_(hid[kept - 1]) + tokens = drafted[1:kept] + [argmax[kept - 1]] + return tokens, pos + kept def generate_greedy_spec(state, input_ids, max_new_tokens, k, fvk, device): @@ -890,19 +979,17 @@ def generate_greedy_spec(state, input_ids, max_new_tokens, k, fvk, device): """ logits = seed_prefill(state, input_ids, fvk, device) pos = input_ids.view(-1).shape[0] + nxt = logits[0].argmax().view(1) out = [] state.spec_windows = 0 state.spec_kept = 0 while len(out) < max_new_tokens: - nxt = logits[0].argmax().view(1) - tokens, logits, pos = spec_decode_step( - state, nxt, pos, k, fvk, device) + tokens, pos = spec_decode_step(state, nxt, pos, k, fvk, device) emitted = [int(nxt)] + tokens[:-1] state.spec_windows += 1 state.spec_kept += len(emitted) out.extend(emitted) - # The step returns the logits that produced its last token, which the - # next window drafts from. + nxt = torch.tensor([tokens[-1]], dtype=torch.long, device=device) return out[:max_new_tokens] diff --git a/flash_rt/frontends/torch/_nexn2_rtx_forward.py b/flash_rt/frontends/torch/_nexn2_rtx_forward.py index 0c096bd6..1f2adb6c 100644 --- a/flash_rt/frontends/torch/_nexn2_rtx_forward.py +++ b/flash_rt/frontends/torch/_nexn2_rtx_forward.py @@ -32,6 +32,19 @@ import torch import torch.nn.functional as F + +def _cs(): + """Current CUDA stream handle. + + Inside a graph capture this is the capture stream; eager, the default one. + A hard-coded 0 is not the same thing: during capture it names a stream that + is not being captured, which is an illegal access at capture_end. This + forward was written for prefill, which is never captured -- a speculative + verify block is. + """ + return torch.cuda.current_stream().cuda_stream + + from flash_rt.frontends.torch._nexn2_rtx_nvfp4_weights import _sf_swz_bytes # Static Nex-N2-mini dims (config.json:text_config). Kept module-local so @@ -68,7 +81,7 @@ def _rms_k(x, w, fvk, device, eps): x2 = x.reshape(-1, dim).contiguous() out = torch.empty(x2.shape[0], dim, dtype=torch.bfloat16, device=device) fvk.rms_norm(x2.data_ptr(), w.data_ptr(), out.data_ptr(), - x2.shape[0], dim, eps, 0) + x2.shape[0], dim, eps, _cs()) return out.reshape(shp) @@ -115,7 +128,7 @@ def _gemm_w16a16(x2d, w, fvk, device): wc = w.contiguous() y = torch.empty(m, n, dtype=torch.bfloat16, device=device) fvk.w16a16_gemm_sm120_bf16(xc.data_ptr(), wc.data_ptr(), y.data_ptr(), - m, n, k, 1.0, 0) + m, n, k, 1.0, _cs()) return y @@ -128,15 +141,19 @@ def _gemm_w16a16(x2d, w, fvk, device): # Set only around a speculative verify block; see _proj and the lm_head below. # -# Default off, because it was measured and it loses: routing the verify block's -# dense projections and lm_head through the 4-bit weights takes K=2 from 36.15 -# to 33.38 tok/s, despite reading a quarter of the bytes. That is the useful -# part of the result -- it says the verify pass is not bandwidth bound, so the -# thing to fix is not what it reads. It runs eager, against a baseline whose -# every step is a captured graph, and a window also pays for K drafts that each -# project the full 248320-wide vocabulary. Capture the verify block and prune -# the draft head's vocabulary first; revisit this after, when the pass is -# actually reading-bound and the switch can be judged on its merits. +# Off, on evidence, twice over. Routing the verify block through _gemm_w4a16 +# loses whether or not the window is captured -- 36.15 to 33.38 eager, and +# 39.77 to 20.84 captured -- and it also reintroduces the divergence from plain +# greedy that a BF16 verify does not have. +# +# Both readings say the same thing: this is the wrong path, not the wrong idea. +# _gemm_w4a16 quantises the weight through its own helper into its own cache, +# so it is neither the tensor the decode GEMV reads nor a kernel shaped for +# three rows. What the verify wants is a small-M GEMM over the *same* packed +# weights the decode path already caches -- then it reads a quarter of the +# bytes and differs from decode only by reduction order. Until that exists, +# BF16 is both faster and the one that agrees with plain greedy token for +# token. _SPEC_VERIFY_W4A16 = False _SPEC_VERIFY = False @@ -175,7 +192,7 @@ def _gemm_w4a16(x2d, w, ld, key, fvk, device): xc = x2d.contiguous() y = torch.empty(m, n, dtype=torch.bfloat16, device=device) fvk.w4a16_gemm_sm120_bf16(xc.data_ptr(), p.data_ptr(), s.data_ptr(), - y.data_ptr(), m, n, k, a, 0) + y.data_ptr(), m, n, k, a, _cs()) return y @@ -199,7 +216,7 @@ def _wquant(w, ld, key, fvk, device): og = torch.zeros(1, dtype=torch.float32, device=device) fvk.bf16_weight_to_nvfp4_swizzled( w.contiguous().data_ptr(), p.data_ptr(), s.data_ptr(), - scr.data_ptr(), og.data_ptr(), nn, kk, 0) + scr.data_ptr(), og.data_ptr(), nn, kk, _cs()) torch.cuda.synchronize() ld[pk] = p ld[key + '_w4s'] = s @@ -219,7 +236,7 @@ def _gemm_fp4(x2d, w, ld, key, fvk, device, xp=None, xsf=None): y = torch.empty(m, n, dtype=torch.bfloat16, device=device) fvk.fp4_w4a16_gemm_sm120_bf16out( xp.data_ptr(), p.data_ptr(), y.data_ptr(), m, n, k, - xsf.data_ptr(), s.data_ptr(), a, 0) + xsf.data_ptr(), s.data_ptr(), a, _cs()) return y @@ -266,7 +283,7 @@ def _silu_mul(g, u, fvk, device): gc = g.reshape(-1).contiguous() uc = u.reshape(-1).contiguous() out = torch.empty(n, dtype=torch.bfloat16, device=device) - fvk.silu_mul_sm120_bf16(gc.data_ptr(), uc.data_ptr(), out.data_ptr(), n, 0) + fvk.silu_mul_sm120_bf16(gc.data_ptr(), uc.data_ptr(), out.data_ptr(), n, _cs()) return out.reshape(g.shape) @@ -333,19 +350,19 @@ def _gdn_wy_chunk(q16, k16, v, g, beta, fvk, device, init_state=None): A = torch.empty(chunks, NV, CH, CH, dtype=torch.float32, device=device) fvk.linear_attn_gdn_wy_kkt_b64_bf16_cublaslt( k_l2.data_ptr(), betac.data_ptr(), gc.data_ptr(), k_pack.data_ptr(), - kkt_base.data_ptr(), A.data_ptr(), S, NK, NV, HK, QKG, 0) + kkt_base.data_ptr(), A.data_ptr(), S, NK, NV, HK, QKG, _cs()) Ai = torch.empty(chunks, NV, CH, CH, dtype=torch.float32, device=device) Ai_pack = torch.empty(chunks, NV, CH, CH, dtype=torch.bfloat16, device=device) fvk.linear_attn_gdn_wy_solve_tril_b64_f32_parallel_pack( - A.data_ptr(), Ai.data_ptr(), Ai_pack.data_ptr(), S, NV, 0) + A.data_ptr(), Ai.data_ptr(), Ai_pack.data_ptr(), S, NV, _cs()) w_pack = torch.empty(chunks, NV, CH, HV, dtype=torch.bfloat16, device=device) u_pack = torch.empty(chunks, NV, CH, HV, dtype=torch.bfloat16, device=device) fvk.linear_attn_gdn_wy_recompute_wu_b64_bf16_mma_fla( k_l2.data_ptr(), vc.data_ptr(), betac.data_ptr(), gc.data_ptr(), Ai_pack.data_ptr(), w_pack.data_ptr(), u_pack.data_ptr(), - S, NK, NV, HK, QKG, 0) + S, NK, NV, HK, QKG, _cs()) state = (init_state.clone() if init_state is not None else torch.zeros(NV, HK, HV, dtype=torch.bfloat16, device=device)) @@ -354,7 +371,7 @@ def _gdn_wy_chunk(q16, k16, v, g, beta, fvk, device, init_state=None): fvk.linear_attn_gdn_wy_chunk_h_b64_bf16_mma_fla( k_l2.data_ptr(), w_pack.data_ptr(), u_pack.data_ptr(), gc.data_ptr(), state.data_ptr(), h0.data_ptr(), v_new.data_ptr(), 0, 0, - S, NK, NV, HK, QKG, 0) + S, NK, NV, HK, QKG, _cs()) q_pack = _wy_pack_t(q_l2.repeat_interleave(QKG, 1)) k_pack_hv = _wy_pack_t(k_l2.repeat_interleave(QKG, 1)) @@ -363,7 +380,7 @@ def _gdn_wy_chunk(q16, k16, v, g, beta, fvk, device, init_state=None): fvk.linear_attn_gdn_wy_output_o_b64_bf16_mma_fla( q_pack.data_ptr(), k_pack_hv.data_ptr(), v_pack.data_ptr(), h0.data_ptr(), gc.data_ptr(), core.data_ptr(), - S, NV, HV, float(HV ** -0.5), 0) + S, NV, HV, float(HV ** -0.5), _cs()) return core, state @@ -418,13 +435,13 @@ def _gdn_layer(h, ld, fvk, device, eps, cap=None, rank=None, xc_ext = torch.empty(B, Se, CONV, dtype=torch.bfloat16, device=device) fvk.causal_conv1d_qwen36_bf16( mixed_ext.data_ptr(), convw_k.data_ptr(), 0, - xc_ext.data_ptr(), B, Se, CONV, KS, True, 0) + xc_ext.data_ptr(), B, Se, CONV, KS, True, _cs()) xc = xc_ext[:, KS - 1:, :].contiguous() else: xc = torch.empty(B, S, CONV, dtype=torch.bfloat16, device=device) fvk.causal_conv1d_qwen36_bf16( mixed.contiguous().data_ptr(), convw_k.data_ptr(), 0, - xc.data_ptr(), B, S, CONV, KS, True, 0) + xc.data_ptr(), B, S, CONV, KS, True, _cs()) # split conv output + broadcast q/k 16 -> 32 heads in one fvk kernel. xc_bf = xc.reshape(B * S, CONV).contiguous() qb = torch.empty(B, S, NV, HK, dtype=torch.bfloat16, device=device) @@ -432,7 +449,7 @@ def _gdn_layer(h, ld, fvk, device, eps, cap=None, rank=None, vb = torch.empty(B, S, NV, HV, dtype=torch.bfloat16, device=device) fvk.qwen35moe_lin_split_qkv_broadcast_bf16( xc_bf.data_ptr(), qb.data_ptr(), kb.data_ptr(), vb.data_ptr(), - B * S, 0) + B * S, _cs()) neg = (-A_log.exp()).float().contiguous() dtb_c = dtb.contiguous() @@ -442,7 +459,7 @@ def _gdn_layer(h, ld, fvk, device, eps, cap=None, rank=None, bo = torch.empty(B, S, NV, dtype=torch.bfloat16, device=device) fvk.qwen36_gdn_gating_bf16( a_bf.data_ptr(), b_bf.data_ptr(), neg.data_ptr(), dtb_c.data_ptr(), - g_out.data_ptr(), bo.data_ptr(), B * S, NV, 0) + g_out.data_ptr(), bo.data_ptr(), B * S, NV, _cs()) if _USE_WY_GDN and S >= _WY_MIN_S: # WY chunked delta-rule scan: 11x faster than the seq-scan at S=2048, @@ -469,7 +486,7 @@ def _gdn_layer(h, ld, fvk, device, eps, cap=None, rank=None, vb.reshape(S, NV, HV).contiguous().data_ptr(), g_out.reshape(S, NV).contiguous().data_ptr(), bo.reshape(S, NV).contiguous().data_ptr(), - state.data_ptr(), core.data_ptr(), S, NV, HK, True, 0) + state.data_ptr(), core.data_ptr(), S, NV, HK, True, _cs()) core = core.reshape(B, S, NV, HV) if cap is not None: @@ -488,7 +505,7 @@ def _gdn_layer(h, ld, fvk, device, eps, cap=None, rank=None, nf = torch.empty_like(cf) fvk.rms_norm_gated_silu_qwen36_bf16( cf.data_ptr(), zf.data_ptr(), nw.data_ptr(), nf.data_ptr(), - cf.shape[0], HV, eps, 0) + cf.shape[0], HV, eps, _cs()) out = _proj(nf.reshape(B * S, VD), ld, 'out_proj', HID, fvk, device) return out.reshape(B, S, HID) @@ -641,7 +658,7 @@ def _full_attn_layer(h, ld, ct, st, fvk, device, eps, cap=None, rank=None, q_pre = torch.empty(B * S, NQ, HD, dtype=torch.bfloat16, device=device) gate = torch.empty(B * S, NQ * HD, dtype=torch.bfloat16, device=device) fvk.qwen35moe_split_q_gate_bf16( - qg.data_ptr(), q_pre.data_ptr(), gate.data_ptr(), B * S, 0) + qg.data_ptr(), q_pre.data_ptr(), gate.data_ptr(), B * S, _cs()) q = q_pre.view(B, S, NQ, HD) gate = gate.view(B, S, NQ * HD) q = _rms_k(q.to(torch.bfloat16), qnw, fvk, device, eps) @@ -656,7 +673,7 @@ def _full_attn_layer(h, ld, ct, st, fvk, device, eps, cap=None, rank=None, ctc, stc = ct.contiguous(), st.contiguous() fvk.qwen36_partial_rope_qk_bf16( qin.data_ptr(), kin.data_ptr(), ctc.data_ptr(), stc.data_ptr(), - qo.data_ptr(), ko.data_ptr(), S, NQ, NKV, HD, ROPE, 0) + qo.data_ptr(), ko.data_ptr(), S, NQ, NKV, HD, ROPE, _cs()) # Causal GQA attention via the vendored FA2 kernel (native GQA: KV stays at # NKV=2, no repeat_interleave; layout is FA2's (B,S,H,HD), no transpose). @@ -682,7 +699,7 @@ def _full_attn_layer(h, ld, ct, st, fvk, device, eps, cap=None, rank=None, gc = gate.reshape(-1).contiguous() ato = torch.empty_like(atc) fvk.sigmoid_mul_sm120_bf16(atc.data_ptr(), gc.data_ptr(), - ato.data_ptr(), atc.numel(), 0) + ato.data_ptr(), atc.numel(), _cs()) at = ato.reshape(B * S, NQ * HD) return _proj(at, ld, 'o_proj', HID, fvk, device).reshape(B, S, HID) @@ -759,7 +776,7 @@ def _capture_per_token_state(cap, rank, S, init_state, conv_hist, mixed, q3[t:t + 1].data_ptr(), k3[t:t + 1].data_ptr(), v3[t:t + 1].data_ptr(), g2[t:t + 1].data_ptr(), b2[t:t + 1].data_ptr(), state.data_ptr(), core1.data_ptr(), - 1, NV, HK, True, 0) + 1, NV, HK, True, _cs()) cap.spec_states[rank][t].copy_(state) # Conv state after t+1 tokens: the last KS-1 entries of the block's inputs @@ -812,14 +829,14 @@ def _moe_experts_m16(x, ti, tw, ld, fvk, device): fvk.moe_m16_mma_sm120_bf16( ap.data_ptr(), gu_p.data_ptr(), asf.data_ptr(), gu_s.data_ptr(), d_gu.data_ptr(), gu_a.data_ptr(), tile_expert.data_ptr(), - total_tiles, n_gu, HID, 0, gu_p[0].numel(), gu_s[0].numel(), 0) + total_tiles, n_gu, HID, 0, gu_p[0].numel(), gu_s[0].numel(), _cs()) inter = _silu_mul(d_gu[:, :INTER], d_gu[:, INTER:], fvk, device).contiguous() ip, isf = _quant_act(inter, fvk, device) d_dn = torch.empty(total_tiles * 16, n_dn, dtype=torch.bfloat16, device=device) fvk.moe_m16_mma_sm120_bf16( ip.data_ptr(), dn_p.data_ptr(), isf.data_ptr(), dn_s.data_ptr(), d_dn.data_ptr(), dn_a.data_ptr(), tile_expert.data_ptr(), - total_tiles, n_dn, INTER, 0, dn_p[0].numel(), dn_s[0].numel(), 0) + total_tiles, n_dn, INTER, 0, dn_p[0].numel(), dn_s[0].numel(), _cs()) out = torch.zeros(S, HID, device=device) out.index_add_(0, stok, d_dn[tiled_row].float() * sw.unsqueeze(-1)) return out @@ -885,14 +902,14 @@ def _moe_experts_bt(x, ti, tw, ld, fvk, device): fvk.moe_blocktile_mma_sm120_bf16( ap.data_ptr(), gu_p.data_ptr(), asf.data_ptr(), gu_s.data_ptr(), d_gu.data_ptr(), gu_a.data_ptr(), tile_expert.data_ptr(), - MAX_TILES, n_gu, HID, 0, gu_p[0].numel(), gu_s[0].numel(), 0) + MAX_TILES, n_gu, HID, 0, gu_p[0].numel(), gu_s[0].numel(), _cs()) inter = _silu_mul(d_gu[:, :INTER], d_gu[:, INTER:], fvk, device).contiguous() ip, isf = _quant_act(inter, fvk, device) d_dn = torch.empty(MAX_TILES * 64, n_dn, dtype=torch.bfloat16, device=device) fvk.moe_blocktile_mma_sm120_bf16( ip.data_ptr(), dn_p.data_ptr(), isf.data_ptr(), dn_s.data_ptr(), d_dn.data_ptr(), dn_a.data_ptr(), tile_expert.data_ptr(), - MAX_TILES, n_dn, INTER, 0, dn_p[0].numel(), dn_s[0].numel(), 0) + MAX_TILES, n_dn, INTER, 0, dn_p[0].numel(), dn_s[0].numel(), _cs()) # Deterministic unpermute via the fused gather-weighted-sum kernel: invert # the routing permutation (inv: orig slot -> sorted position, a 131 KB int # scatter) to get each token's TOPK d_dn rows, then one kernel computes @@ -906,7 +923,7 @@ def _moe_experts_bt(x, ti, tw, ld, fvk, device): out = torch.empty(S, HID, dtype=torch.float32, device=device) fvk.moe_weighted_sum_sm120_bf16( d_dn.data_ptr(), rows.data_ptr(), twc.data_ptr(), out.data_ptr(), - S, TOPK, n_dn, n_dn, 0) + S, TOPK, n_dn, n_dn, _cs()) return out @@ -940,14 +957,14 @@ def _moe_experts_grouped(x, ti, tw, ld, fvk, device): moe_grouped_w4a16(fvk)( A.data_ptr(), gu_p.data_ptr(), gu_s.data_ptr(), gu_a.data_ptr(), se.data_ptr(), d_gu.data_ptr(), slots, n_gu, HID, - HID, gu_p[0].numel(), gu_s[0].numel(), 0) + HID, gu_p[0].numel(), gu_s[0].numel(), _cs()) g, u = d_gu[:, :INTER], d_gu[:, INTER:] inter = _silu_mul(g, u, fvk, device).contiguous() d_dn = torch.empty(slots, n_dn, dtype=torch.bfloat16, device=device) moe_grouped_w4a16(fvk)( inter.data_ptr(), dn_p.data_ptr(), dn_s.data_ptr(), dn_a.data_ptr(), se.data_ptr(), d_dn.data_ptr(), slots, n_dn, INTER, - INTER, dn_p[0].numel(), dn_s[0].numel(), 0) + INTER, dn_p[0].numel(), dn_s[0].numel(), _cs()) out = torch.zeros(S, HID, device=device) out.index_add_(0, stok, d_dn.float() * sw.unsqueeze(-1)) return out From c2178cc8bebf42953770233eea2d919a87d81fab Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 30 Jul 2026 14:29:23 -0400 Subject: [PATCH 37/85] Take the rollback snapshots before the state they snapshot is overwritten The per-step recurrent and conv snapshots were written after the layer copied its final state back into the decode buffers -- and those buffers are the very tensors the snapshot routine reads as its starting point, since the caller passes cap.lin_state[rank] and cap.lin_conv_state[rank] in directly. So the replay started from the state the block ended at, and every partial-accept rewind restored a fabricated one. It hid well: the tokens stayed deterministic across replays and the fixture still passed, because a corrupted state changes what the verifier says next rather than making anything crash. What gave it away was a window count that differed between otherwise identical runs -- 19 against 20 for the same 50 tokens -- which is the acceptance pattern shifting under a state that should have been reproducible. Tokens kept per window at K=2 go 2.64 -> 2.75, and K=3 goes 21.72 -> 40.06 tok/s. Also gives the speculative graphs their own memory pool rather than the decode graphs': the two are replayed interleaved, which is the case a shared pool does not promise to handle. Measured no difference on its own, kept because the sharing was not defensible either way. --- flash_rt/frontends/torch/_nexn2_rtx_decode.py | 9 ++++++++- flash_rt/frontends/torch/_nexn2_rtx_forward.py | 11 ++++++++--- 2 files changed, 16 insertions(+), 4 deletions(-) diff --git a/flash_rt/frontends/torch/_nexn2_rtx_decode.py b/flash_rt/frontends/torch/_nexn2_rtx_decode.py index d6eb84b1..59c2fdbe 100644 --- a/flash_rt/frontends/torch/_nexn2_rtx_decode.py +++ b/flash_rt/frontends/torch/_nexn2_rtx_decode.py @@ -302,6 +302,13 @@ def __init__(self, handles, max_seq, device): # and RoPE slice are baked per position exactly as the decode graph's # are. Same LRU bound, since each graph owns a memory pool. self._spec_graphs = collections.OrderedDict() + # Its own memory pool, not the decode graphs'. The two are replayed + # interleaved -- a window, then whatever the caller does next -- and + # sharing a pool between graphs used that way is the case the runtime + # does not promise to handle. Measured: with the pool shared, a 64-token + # speculative run ran at half the rate of a 32-token one on identical + # code, the cost growing with the number of live graphs. + self._spec_pool = torch.cuda.graph_pool_handle() self._spec_tokens = None self._spec_argmax = None # Which half of the draft head's fc input carries the hidden state. @@ -921,7 +928,7 @@ def _restore(): g = torch.cuda.CUDAGraph() with torch.cuda.graph(g, stream=state._graph_stream, - pool=state._graph_pool), torch.no_grad(): + pool=state._spec_pool), torch.no_grad(): hid = _spec_block(state, pos, k, fvk, device) with torch.no_grad(): _restore() diff --git a/flash_rt/frontends/torch/_nexn2_rtx_forward.py b/flash_rt/frontends/torch/_nexn2_rtx_forward.py index 1f2adb6c..a6e926d0 100644 --- a/flash_rt/frontends/torch/_nexn2_rtx_forward.py +++ b/flash_rt/frontends/torch/_nexn2_rtx_forward.py @@ -490,15 +490,20 @@ def _gdn_layer(h, ld, fvk, device, eps, cap=None, rank=None, core = core.reshape(B, S, NV, HV) if cap is not None: + # The per-step snapshots go FIRST. `init_state` and `conv_hist` are the + # very tensors the two copies below overwrite -- the caller passes + # cap.lin_state[rank] / cap.lin_conv_state[rank] in directly -- so + # replaying the scan after the copies would start it from the state the + # block ended at, and every rewind would restore a fabricated one. + if getattr(cap, 'spec_capture', False): + _capture_per_token_state(cap, rank, S, init_state, conv_hist, + mixed, qb, kb, vb, g_out, bo, fvk, device) # GDN recurrent final state = `state` after the S-step scan; conv state # = the last KS-1 `mixed` inputs (channel-major, newest at index -1), # matching the causal_conv1d_update rolling buffer (1, CONV, KS-1). cap.lin_state[rank].copy_(state) cs = mixed[0, S - (KS - 1):S, :].transpose(0, 1).contiguous() cap.lin_conv_state[rank].copy_(cs.unsqueeze(0)) - if getattr(cap, 'spec_capture', False): - _capture_per_token_state(cap, rank, S, init_state, conv_hist, - mixed, qb, kb, vb, g_out, bo, fvk, device) cf = core.reshape(-1, HV).contiguous() zf = z.reshape(-1, HV).to(torch.bfloat16).contiguous() From 2a4c5fff0bacde4e2dc30f548cb59ed1f43ff3f3 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 30 Jul 2026 15:32:46 -0400 Subject: [PATCH 38/85] Group prefill's routed tokens by expert instead of re-reading the weight The grouped GEMV issues one GEMV per (token, expert) slot, so an expert's weight is read once per token that routed to it: 8192 reads of a 1.18 MB weight per layer at S=1024, 9.7 GB of traffic a layer before the down projection. Sorting by expert makes those reads hit L2, which is why it worked at all, but it stays bounded by L2 and it dominated: measured 74.6% of a 1024-token prefill, 1762 ms of 2361. One GEMM per expert reads each weight once. The block-scaled MMA tile does the same and better, but it is a build tier that is not present everywhere; this path needs only the NVFP4 W4A16 GEMM, which is. It pays only once the tokens per expert are worth a launch -- about five per expert against the grouped path's two for the whole layer -- so it is gated on the mean. At S=256, eight tokens an expert, it takes TTFT from 585 to 923 ms; at S=1024, thirty-two, from 2237 to 1341. At 1024 tokens: prefill kernel time 2361 -> 749 ms, TTFT 2237 -> 1341 ms, 764 tok/s against 458. Decode and the fixture are unchanged, 16/16. Also retested the W4A16 dense projections here, where a part with no W4A4 might have been expected to favour them on traffic: 1335.0 -> 1346.6 ms, i.e. nothing. Left off, with the measurement recorded. --- .../frontends/torch/_nexn2_rtx_forward.py | 79 +++++++++++++++++++ 1 file changed, 79 insertions(+) diff --git a/flash_rt/frontends/torch/_nexn2_rtx_forward.py b/flash_rt/frontends/torch/_nexn2_rtx_forward.py index a6e926d0..b5e44011 100644 --- a/flash_rt/frontends/torch/_nexn2_rtx_forward.py +++ b/flash_rt/frontends/torch/_nexn2_rtx_forward.py @@ -137,6 +137,10 @@ def _gemm_w16a16(x2d, w, fvk, device): # fp4 *weight* (not the activation), so W4A16 lands at the same ~0.987 as W4A4 # while being slower than the CUTLASS W4A4 -- dominated. Default OFF; the 0.994 # path needs a bf16-*weight* GEMM (repurpose this kernel's 2.18x structure). +# +# Retested on a part with no W4A4 at all, where it might have been expected to +# win on traffic: at S=1024 it moves TTFT 1335.0 -> 1346.6 ms, i.e. nothing. +# The 178 ms those GEMMs cost is not what bounds this prefill. _DENSE_W4A16 = False # Set only around a speculative verify block; see _proj and the lm_head below. @@ -739,6 +743,17 @@ def moe_grouped_w4a16(fvk): # Grouped MoE for prefill (on by default); set False to use the per-expert loop. _USE_GROUPED_MOE = True +# One GEMM per expert, for a build without the block-scaled MMA tiles. Reads +# each expert's weight once instead of once per token that routed to it. +# +# It only pays once the tokens per expert are worth a launch. Each expert costs +# about five launches (quantise, two GEMMs, the gate), so with 256 of them a +# layer that is thousands of launches whichever way; below the threshold the +# grouped GEMV's two launches win even though it re-reads the weight. Measured +# at S=256 -- eight tokens an expert -- the per-expert path takes TTFT from 585 +# to 923 ms, while at S=1024 it takes it from 2237 to 1354. +_USE_PER_EXPERT_GEMM = True +_PER_EXPERT_MIN_M = 16 # mean tokens per expert, = S * TOPK / 256 # M=16 tensor-core mma MoE: tokens are sorted into 16-row expert tiles and the # SM120 block-scaled mma runs each expert once at full M-utilisation -- ~5.6x # the SIMT grouped W4A16 at large S (the compute wall). W4A4 (FP4 activation), @@ -932,6 +947,66 @@ def _moe_experts_bt(x, ti, tw, ld, fvk, device): return out +def _moe_experts_per_expert_gemm(x, ti, tw, ld, fvk, device): + """Routed experts as one GEMM per expert over the tokens that chose it. + + The grouped GEMV below issues one GEMV per (token, expert) slot, so an + expert's weight is re-read once per token that routed to it -- 8192 reads + of a 1.18 MB weight per layer at S=1024, which is 9.7 GB of traffic a layer + even before the down projection. Sorting by expert makes those reads hit L2 + rather than DRAM, which is why it works at all, but it is still bounded by + L2 bandwidth and it dominates prefill: measured 74.6% of a 1024-token + prefill, 1762 ms of 2361. + + Grouping the tokens instead turns each expert into a single M-row GEMM that + reads its weight once. The block-scaled 4-bit MMA tile does the same thing + and better, but it is a build tier that is not present everywhere; this path + needs only the NVFP4 W4A16 GEMM, which is. + + The count per expert is data-dependent, so this reads it to the host -- one + sync per layer, which prefill can afford and a captured decode could not. + """ + S = x.shape[0] + gu_p, gu_s = ld['experts_gate_up_packed_t'], ld['experts_gate_up_sf_t'] + dn_p, dn_s = ld['experts_down_packed_t'], ld['experts_down_sf_t'] + n_gu, n_dn = gu_p.shape[1], dn_p.shape[1] + if 'experts_gate_up_alpha_list' not in ld: + ld['experts_gate_up_alpha_list'] = ld['experts_gate_up_alpha_t'].tolist() + ld['experts_down_alpha_list'] = ld['experts_down_alpha_t'].tolist() + gu_a = ld['experts_gate_up_alpha_list'] + dn_a = ld['experts_down_alpha_list'] + + exp_flat = ti.reshape(-1).to(torch.int32) + tok_flat = torch.arange(S, device=device).repeat_interleave(TOPK) + # Stable, so equal-expert ties keep token order and the rows packed into + # each quantisation tile are the same run to run. + order = exp_flat.argsort(stable=True) + se = exp_flat[order] + stok = tok_flat[order] + sw = tw.reshape(-1)[order] + + counts = torch.bincount(se, minlength=_N_EXPERTS).tolist() + A = x[stok].contiguous() # (slots, HID) bf16 + d_dn = torch.empty(S * TOPK, n_dn, dtype=torch.bfloat16, device=device) + + off = 0 + for e, cnt in enumerate(counts): + if cnt == 0: + continue + rows = A[off:off + cnt] + gu = _nvfp4_gemm(rows, gu_p[e].data_ptr(), gu_s[e].data_ptr(), + gu_a[e], n_gu, fvk, device, _cs()) + inter = _silu_mul(gu[:, :INTER], gu[:, INTER:], fvk, device) + d_dn[off:off + cnt] = _nvfp4_gemm( + inter.contiguous(), dn_p[e].data_ptr(), dn_s[e].data_ptr(), + dn_a[e], n_dn, fvk, device, _cs()) + off += cnt + + out = torch.zeros(S, HID, device=device) + out.index_add_(0, stok, d_dn.float() * sw.unsqueeze(-1)) + return out + + def _moe_experts_grouped(x, ti, tw, ld, fvk, device): """Routed experts via the grouped W4A16 GEMV. Flatten the S*TOPK (token, expert) assignments, sort by expert so consecutive slots share a @@ -1003,6 +1078,10 @@ def _moe_layer(h, ld, fvk, device): out = _moe_experts_bt(x, ti, tw, ld, fvk, device) elif _USE_M16_MOE and big and hasattr(fvk, 'moe_m16_mma_sm120_bf16'): out = _moe_experts_m16(x, ti, tw, ld, fvk, device) + elif (big and _USE_PER_EXPERT_GEMM + and x.shape[0] * TOPK >= _PER_EXPERT_MIN_M * _N_EXPERTS + and hasattr(fvk, 'fp4_w4a16_gemm_sm120_bf16out')): + out = _moe_experts_per_expert_gemm(x, ti, tw, ld, fvk, device) elif _USE_GROUPED_MOE: out = _moe_experts_grouped(x, ti, tw, ld, fvk, device) else: From 35127ca6a988c86941fa96a2d0dc37dec420af11 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Fri, 31 Jul 2026 04:10:39 -0400 Subject: [PATCH 39/85] Hand prefill's dense projections to cuBLASLt, and hoist the gate out of the loop The hand-written w16a16 GEMM was the largest single kernel left in a 1024-token prefill: 171 ms, about 2.9 TFLOP, so roughly 17 TFLOP/s on a part whose tensor cores do far more. cuBLASLt is 4-8x it at every shape prefill issues -- 140 ms over a prefill against 22.5 -- and it is a drop-in in the strongest sense: bitwise identical output at every shape checked, and bit-reproducible across repeated launches. The determinism caveat this file carries is about torch.matmul's split-K reduction order and does not apply here, which was measured rather than assumed. The per-expert loop also stopped allocating and slicing per iteration: both projections write into one preallocated slot-major buffer, and the gate runs once over all of it instead of once per expert with two slice copies each -- 256 launches a layer for an op that does not care which expert a row came from. At 1024 tokens, TTFT 1341 -> 942 ms and prefill 764 -> 1087 tok/s. At 20 tokens TTFT is 89.1 ms, below the 102.5 ms vLLM takes on the same part. Decode and the fixture unchanged, 16/16. --- .../frontends/torch/_nexn2_rtx_forward.py | 59 +++++++++++++++---- 1 file changed, 48 insertions(+), 11 deletions(-) diff --git a/flash_rt/frontends/torch/_nexn2_rtx_forward.py b/flash_rt/frontends/torch/_nexn2_rtx_forward.py index b5e44011..b72ef597 100644 --- a/flash_rt/frontends/torch/_nexn2_rtx_forward.py +++ b/flash_rt/frontends/torch/_nexn2_rtx_forward.py @@ -118,6 +118,19 @@ def _proj(x2d, ld, base, n, fvk, device): return (x2d.float() @ w.float().T).to(torch.bfloat16) +# cuBLASLt for the same product, where the build has it. It is 4-8x the +# hand-written kernel at every shape prefill issues -- 140 ms of a 1024-token +# prefill against 22.5 -- and it is a drop-in in the strongest sense: bitwise +# identical output at every shape checked, and bit-reproducible across repeated +# launches. +# +# The determinism caveat elsewhere in this file is about torch.matmul, whose +# split-K reduction order can vary and flip a near-tie argmax. It does not apply +# to this entry point, which was measured rather than assumed. Set False to +# force the hand-written kernel. +_DENSE_CUBLASLT = True + + def _gemm_w16a16(x2d, w, fvk, device): """y = x @ w.T via the deterministic bf16-act x bf16-weight tensor-core GEMM (fp32 register accumulate). Matches the fp32 path's argmax (cos 1.0) @@ -127,6 +140,10 @@ def _gemm_w16a16(x2d, w, fvk, device): xc = x2d.contiguous() wc = w.contiguous() y = torch.empty(m, n, dtype=torch.bfloat16, device=device) + if _DENSE_CUBLASLT and hasattr(fvk, 'bf16_matmul_cublaslt_bf16'): + fvk.bf16_matmul_cublaslt_bf16(xc.data_ptr(), wc.data_ptr(), + y.data_ptr(), m, n, k, _cs()) + return y fvk.w16a16_gemm_sm120_bf16(xc.data_ptr(), wc.data_ptr(), y.data_ptr(), m, n, k, 1.0, _cs()) return y @@ -260,9 +277,16 @@ def _quant_act(x2d, fvk, device, stream=0): def _nvfp4_gemm_preq(xp, xsf, wp_ptr, wsf_ptr, alpha, m, n, k, fvk, device, - stream=0): - """y = x @ w.T from a pre-quantised activation (xp, xsf).""" - y = torch.empty(m, n, dtype=torch.bfloat16, device=device) + stream=0, out=None): + """y = x @ w.T from a pre-quantised activation (xp, xsf). + + ``out`` lets a caller point the result at a slice of a buffer it already + owns, which is what the per-expert loop wants: it writes 256 blocks into + one matrix, and allocating each of them separately costs more in Python + than the GEMM costs on the device. + """ + y = torch.empty(m, n, dtype=torch.bfloat16, device=device) if out is None \ + else out fvk.fp4_w4a16_gemm_sm120_bf16out( xp.data_ptr(), wp_ptr, y.data_ptr(), m, n, k, xsf.data_ptr(), wsf_ptr, alpha, stream) @@ -986,22 +1010,35 @@ def _moe_experts_per_expert_gemm(x, ti, tw, ld, fvk, device): sw = tw.reshape(-1)[order] counts = torch.bincount(se, minlength=_N_EXPERTS).tolist() + slots = S * TOPK A = x[stok].contiguous() # (slots, HID) bf16 - d_dn = torch.empty(S * TOPK, n_dn, dtype=torch.bfloat16, device=device) + # One buffer per projection, written in place by each expert's GEMM. The + # activation is a slot-major matrix throughout, so the gate is one launch + # over all of it rather than one per expert -- 256 launches a layer and two + # slice copies each, for an op that does not care where the rows came from. + d_gu = torch.empty(slots, n_gu, dtype=torch.bfloat16, device=device) + d_dn = torch.empty(slots, n_dn, dtype=torch.bfloat16, device=device) off = 0 + bounds = [] for e, cnt in enumerate(counts): if cnt == 0: continue - rows = A[off:off + cnt] - gu = _nvfp4_gemm(rows, gu_p[e].data_ptr(), gu_s[e].data_ptr(), - gu_a[e], n_gu, fvk, device, _cs()) - inter = _silu_mul(gu[:, :INTER], gu[:, INTER:], fvk, device) - d_dn[off:off + cnt] = _nvfp4_gemm( - inter.contiguous(), dn_p[e].data_ptr(), dn_s[e].data_ptr(), - dn_a[e], n_dn, fvk, device, _cs()) + bounds.append((e, off, cnt)) + xp, xsf = _quant_act(A[off:off + cnt], fvk, device, _cs()) + _nvfp4_gemm_preq(xp, xsf, gu_p[e].data_ptr(), gu_s[e].data_ptr(), + gu_a[e], cnt, n_gu, HID, fvk, device, _cs(), + out=d_gu[off:off + cnt]) off += cnt + inter = _silu_mul(d_gu[:, :INTER], d_gu[:, INTER:], fvk, device) + for e, off_e, cnt in bounds: + xp, xsf = _quant_act(inter[off_e:off_e + cnt].contiguous(), fvk, + device, _cs()) + _nvfp4_gemm_preq(xp, xsf, dn_p[e].data_ptr(), dn_s[e].data_ptr(), + dn_a[e], cnt, n_dn, INTER, fvk, device, _cs(), + out=d_dn[off_e:off_e + cnt]) + out = torch.zeros(S, HID, device=device) out.index_add_(0, stok, d_dn.float() * sw.unsqueeze(-1)) return out From 80d7009e555499f7d34e3c9d66d2627aa61220f5 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Fri, 31 Jul 2026 04:44:39 -0400 Subject: [PATCH 40/85] Serve every routed expert of a prefill layer in one grouped GEMM Three pieces, all so that the routing never reaches the host. A CUTLASS SM100 grouped block-scaled GEMM takes its per-group problem shapes from device memory -- the host-side array is passed as nullptr deliberately -- so the launch geometry depends only on the expert count, which is known from the config. A device kernel fills the descriptor arrays from the routing's prefix sums. And the activation quantiser writes every group's scale factors in one launch: the block-scaled layout blocks rows by 128, so a group beginning at an arbitrary row of a jointly-quantised matrix has no contiguous sub-block to point at, which is what had been forcing a launch and a host iteration per expert. The scale-factor buffer is sized from a bound rather than a sum -- the total is at most (experts + slots/128) super-blocks however the routing falls -- so nothing needs reading back. Preflight against the loop it replaces, at the shapes prefill issues: 6.0x on gate_up and 14.5x on down at S=1024, 7.7x and 17.6x at S=256, 512 launches down to 2, output bitwise identical. The first run of that preflight reported cosine NaN and was not accepted: uniform random bytes as scale factors decode to values that overflow the product, and two sides agreeing on NaN proves nothing. The length threshold is gone with it. There is no crossover to tune, because the grouped path stops paying per expert for anything: prompt before -> now vLLM 64 176.9 -> 117.6 133.1 128 301.1 -> 143.5 187.2 256 544.0 -> 182.3 214.5 512 718.5 -> 269.3 251.1 1024 943.6 -> 479.2 319.4 2048 1462.7 -> 1014.7 495.0 Fixture 16/16, decode unchanged. --- CMakeLists.txt | 4 +- csrc/bindings.cpp | 44 +++ .../fp4/cutlass_nvfp4_moe_grouped_sm100.cu | 290 ++++++++++++++++++ .../fp4/cutlass_nvfp4_moe_grouped_sm100.cuh | 72 +++++ csrc/kernels/quantize.cu | 94 ++++++ csrc/kernels/quantize.cuh | 18 ++ .../frontends/torch/_nexn2_rtx_forward.py | 109 +++++++ 7 files changed, 630 insertions(+), 1 deletion(-) create mode 100644 csrc/gemm/fp4/cutlass_nvfp4_moe_grouped_sm100.cu create mode 100644 csrc/gemm/fp4/cutlass_nvfp4_moe_grouped_sm100.cuh diff --git a/CMakeLists.txt b/CMakeLists.txt index 4e4a456b..6060fa1c 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -656,12 +656,14 @@ endif() # -gencode arch=compute_110a,code=sm_110a via GPU_GENCODE. if(GPU_ARCH STREQUAL "110") add_library(cutlass_nvfp4_w4a16_sm100_obj OBJECT - csrc/gemm/fp4/cutlass_nvfp4_w4a16_gemm_sm100.cu) + csrc/gemm/fp4/cutlass_nvfp4_w4a16_gemm_sm100.cu + csrc/gemm/fp4/cutlass_nvfp4_moe_grouped_sm100.cu) set_target_properties(cutlass_nvfp4_w4a16_sm100_obj PROPERTIES CUDA_STANDARD 17 POSITION_INDEPENDENT_CODE ON ) target_include_directories(cutlass_nvfp4_w4a16_sm100_obj PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/csrc ${CMAKE_CURRENT_SOURCE_DIR}/csrc/gemm ${CMAKE_CURRENT_SOURCE_DIR}/csrc/gemm/fp4 ${CUTLASS_INCLUDE} diff --git a/csrc/bindings.cpp b/csrc/bindings.cpp index a1c80ecb..3e6a0a63 100644 --- a/csrc/bindings.cpp +++ b/csrc/bindings.cpp @@ -38,6 +38,7 @@ #endif #ifdef ENABLE_CUTLASS_SM100_NVFP4_W4A16 #include "gemm/fp4/cutlass_nvfp4_w4a16_gemm_sm100.cuh" +#include "gemm/fp4/cutlass_nvfp4_moe_grouped_sm100.cuh" #endif #ifdef ENABLE_ACTION_FFN_MEGAKERNEL_V6T #include "kernels/megakernel/action_ffn_megakernel_v6t_sm120.cuh" @@ -1088,6 +1089,19 @@ PYBIND11_MODULE(flash_rt_kernels, m) { }, py::arg("input"), py::arg("fp4_data"), py::arg("scale_factors"), py::arg("rows"), py::arg("cols"), py::arg("stream") = 0); + m.def("moe_grouped_quant_nvfp4_bf16", + [](uintptr_t A, uintptr_t expert_of_row, uintptr_t group_off, + uintptr_t sfa_off, uintptr_t out_packed, uintptr_t out_sf, + int slots, int K, uintptr_t stream) -> int { + return moe_grouped_quant_nvfp4_bf16( + to_ptr(A), to_ptr(expert_of_row), to_ptr(group_off), + to_ptr(sfa_off), to_ptr(out_packed), to_ptr(out_sf), + slots, K, to_stream(stream)); + }, + py::arg("A"), py::arg("expert_of_row"), py::arg("group_off"), + py::arg("sfa_off"), py::arg("out_packed"), py::arg("out_sf"), + py::arg("slots"), py::arg("K"), py::arg("stream") = 0); + m.def("quantize_bf16_to_nvfp4_swizzled", [](uintptr_t input, uintptr_t fp4_data, uintptr_t scale_factors, int rows, int cols, uintptr_t stream) { @@ -7269,6 +7283,36 @@ graph-replay safe) to fill the SMs on long K. M in 1..16; N%8==0; K%64==0; // out of the Thor surface. // ───────────────────────────────────────────────────────────────────── #ifdef ENABLE_CUTLASS_SM100_NVFP4_W4A16 + // Every routed expert of a layer in one launch, with the per-group shapes + // taken from device memory so the routing never reaches the host. + m.def("moe_grouped_gemm_nvfp4_sm100_bf16out", + [](uintptr_t A_packed, uintptr_t SFA, uintptr_t W_stack, + uintptr_t SFB_stack, uintptr_t alpha_dev, uintptr_t D, + uintptr_t group_off, uintptr_t sfa_off, + int groups, int N, int K, long w_stride, long sfb_stride, + uintptr_t scratch, size_t scratch_bytes, + uintptr_t stream) -> int { + return flash_rt::gemm::moe_grouped_gemm_nvfp4_sm100_bf16out( + to_ptr(A_packed), to_ptr(SFA), to_ptr(W_stack), + to_ptr(SFB_stack), to_ptr(alpha_dev), to_ptr(D), + to_ptr(group_off), to_ptr(sfa_off), + groups, N, K, w_stride, sfb_stride, + to_ptr(scratch), scratch_bytes, to_stream(stream)); + }, + py::arg("A_packed"), py::arg("SFA"), py::arg("W_stack"), + py::arg("SFB_stack"), py::arg("alpha_dev"), py::arg("D"), + py::arg("group_off"), py::arg("sfa_off"), py::arg("groups"), + py::arg("N"), py::arg("K"), py::arg("w_stride"), + py::arg("sfb_stride"), py::arg("scratch"), + py::arg("scratch_bytes"), py::arg("stream") = 0); + + m.def("moe_grouped_gemm_nvfp4_sm100_scratch_bytes", + [](int groups) -> size_t { + return flash_rt::gemm::moe_grouped_gemm_nvfp4_sm100_scratch_bytes( + groups); + }, + py::arg("groups")); + m.def("fp4_w4a16_gemm_sm120_bf16out", [](uintptr_t A_packed, uintptr_t B_packed, uintptr_t D, int M, int N, int K, diff --git a/csrc/gemm/fp4/cutlass_nvfp4_moe_grouped_sm100.cu b/csrc/gemm/fp4/cutlass_nvfp4_moe_grouped_sm100.cu new file mode 100644 index 00000000..5d63affb --- /dev/null +++ b/csrc/gemm/fp4/cutlass_nvfp4_moe_grouped_sm100.cu @@ -0,0 +1,290 @@ +// SPDX-License-Identifier: Apache-2.0 +// +// Grouped NVFP4 block-scaled GEMM for sm_100-class Blackwell. See header. + +#include "gemm/fp4/cutlass_nvfp4_moe_grouped_sm100.cuh" + +#include + +#include "cutlass/cutlass.h" +#include "cutlass/gemm/dispatch_policy.hpp" +#include "cutlass/gemm/group_array_problem_shape.hpp" +#include "cutlass/gemm/collective/collective_builder.hpp" +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/kernel/gemm_universal.hpp" +#include "cutlass/epilogue/collective/collective_builder.hpp" +#include "cutlass/epilogue/dispatch_policy.hpp" +#include "cutlass/util/packed_stride.hpp" +#include "cute/tensor.hpp" + +namespace flash_rt { +namespace gemm { + +namespace { + +using namespace cute; + +using ElementA = cutlass::float_e2m1_t; +using ElementB = cutlass::float_e2m1_t; +using ElementC = cutlass::bfloat16_t; +using ElementD = cutlass::bfloat16_t; +using ElementAccumulator = float; +using ElementSF = cutlass::float_ue4m3_t; + +using LayoutA = cutlass::layout::RowMajor; +using LayoutB = cutlass::layout::ColumnMajor; +using LayoutC = cutlass::layout::RowMajor; + +using ElementPairA = cutlass::nv_float4_t; +using ElementPairB = cutlass::nv_float4_t; + +constexpr int AlignmentA = 32; +constexpr int AlignmentB = 32; +constexpr int AlignmentC = 8; +constexpr int AlignmentD = 8; + +using ProblemShape = cutlass::gemm::GroupProblemShape>; +using ArchTag = cutlass::arch::Sm100; +using OperatorClass = cutlass::arch::OpClassBlockScaledTensorOp; +using ClusterShape = Shape; +using MmaTileShape = Shape<_128, _256, _256>; + +using CollectiveEpilogue = + typename cutlass::epilogue::collective::CollectiveBuilder< + ArchTag, OperatorClass, + MmaTileShape, ClusterShape, + Shape<_128, _64>, + ElementAccumulator, ElementAccumulator, + ElementC, LayoutC*, AlignmentC, + ElementD, LayoutC*, AlignmentD, + cutlass::epilogue::PtrArrayTmaWarpSpecialized1Sm>::CollectiveOp; + +using CollectiveMainloop = + typename cutlass::gemm::collective::CollectiveBuilder< + ArchTag, OperatorClass, + ElementPairA, LayoutA*, AlignmentA, + ElementPairB, LayoutB*, AlignmentB, + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout< + static_cast(sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::KernelPtrArrayTmaWarpSpecialized1SmNvf4Sm100>::CollectiveOp; + +using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + ProblemShape, CollectiveMainloop, CollectiveEpilogue>; +using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + +using StrideA = typename Gemm::GemmKernel::InternalStrideA; +using StrideB = typename Gemm::GemmKernel::InternalStrideB; +using StrideC = typename Gemm::GemmKernel::InternalStrideC; +using StrideD = typename Gemm::GemmKernel::InternalStrideD; +using LayoutSFA = typename Gemm::GemmKernel::CollectiveMainloop::InternalLayoutSFA; +using LayoutSFB = typename Gemm::GemmKernel::CollectiveMainloop::InternalLayoutSFB; +using Sm1xxBlkScaledConfig = + typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig; + +using ProblemShapeMNK = typename ProblemShape::UnderlyingProblemShape; + +// Everything the launch needs, laid out back to back in one scratch buffer so +// the caller allocates once. Filled by a device kernel from the routing, which +// is what keeps the whole thing free of a host round trip. +struct GroupArgs { + ProblemShapeMNK* shapes; + const ElementA** ptr_A; + const ElementB** ptr_B; + const ElementSF** ptr_SFA; + const ElementSF** ptr_SFB; + ElementD** ptr_D; + const float** ptr_alpha; + StrideA* stride_A; + StrideB* stride_B; + StrideC* stride_C; + StrideD* stride_D; + LayoutSFA* layout_SFA; + LayoutSFB* layout_SFB; +}; + +constexpr size_t align_up(size_t v, size_t a) { return (v + a - 1) / a * a; } + +size_t args_bytes(int g) { + size_t n = 0; + n = align_up(n + sizeof(ProblemShapeMNK) * g, 256); + n = align_up(n + sizeof(void*) * g * 6, 256); // A B SFA SFB D alpha + n = align_up(n + sizeof(StrideA) * g, 256); + n = align_up(n + sizeof(StrideB) * g, 256); + n = align_up(n + sizeof(StrideC) * g, 256); + n = align_up(n + sizeof(StrideD) * g, 256); + n = align_up(n + sizeof(LayoutSFA) * g, 256); + n = align_up(n + sizeof(LayoutSFB) * g, 256); + return n; +} + +GroupArgs carve(void* base, int g) { + auto* p = static_cast(base); + size_t o = 0; + auto take = [&](size_t bytes) { + void* r = p + o; + o = align_up(o + bytes, 256); + return r; + }; + GroupArgs a{}; + a.shapes = static_cast(take(sizeof(ProblemShapeMNK) * g)); + auto* ptrs = static_cast(take(sizeof(void*) * g * 6)); + auto at = [&](int i) { return static_cast(ptrs + i * g); }; + a.ptr_A = static_cast(at(0)); + a.ptr_B = static_cast(at(1)); + a.ptr_SFA = static_cast(at(2)); + a.ptr_SFB = static_cast(at(3)); + a.ptr_D = static_cast(at(4)); + a.ptr_alpha = static_cast(at(5)); + a.stride_A = static_cast(take(sizeof(StrideA) * g)); + a.stride_B = static_cast(take(sizeof(StrideB) * g)); + a.stride_C = static_cast(take(sizeof(StrideC) * g)); + a.stride_D = static_cast(take(sizeof(StrideD) * g)); + a.layout_SFA = static_cast(take(sizeof(LayoutSFA) * g)); + a.layout_SFB = static_cast(take(sizeof(LayoutSFB) * g)); + return a; +} + +// One thread per group. Reads the routing (prefix sums of the per-expert token +// counts) and writes the descriptor arrays CUTLASS reads. No host involvement, +// which is the point: the launch shape below depends only on the group count. +__global__ void fill_group_args( + GroupArgs a, + const uint8_t* __restrict__ A_packed, + const uint8_t* __restrict__ SFA, + const uint8_t* __restrict__ W_stack, + const uint8_t* __restrict__ SFB_stack, + const float* __restrict__ alpha, + uint8_t* __restrict__ D, + const int* __restrict__ group_off, + const int* __restrict__ sfa_off, + int groups, int N, int K, long w_stride, long sfb_stride) { + const int e = blockIdx.x * blockDim.x + threadIdx.x; + if (e >= groups) return; + + const int off = group_off[e]; + const int m = group_off[e + 1] - off; + + a.shapes[e] = cute::make_shape(m, N, K); + a.ptr_A[e] = reinterpret_cast( + A_packed + static_cast(off) * (K / 2)); + a.ptr_B[e] = reinterpret_cast(W_stack + e * w_stride); + a.ptr_SFA[e] = reinterpret_cast(SFA + sfa_off[e]); + a.ptr_SFB[e] = reinterpret_cast(SFB_stack + e * sfb_stride); + a.ptr_D[e] = reinterpret_cast( + D + static_cast(off) * N * sizeof(ElementD)); + a.ptr_alpha[e] = alpha + e; + + a.stride_A[e] = cutlass::make_cute_packed_stride( + StrideA{}, cute::make_shape(m, K, 1)); + a.stride_B[e] = cutlass::make_cute_packed_stride( + StrideB{}, cute::make_shape(N, K, 1)); + a.stride_C[e] = cutlass::make_cute_packed_stride( + StrideC{}, cute::make_shape(m, N, 1)); + a.stride_D[e] = cutlass::make_cute_packed_stride( + StrideD{}, cute::make_shape(m, N, 1)); + a.layout_SFA[e] = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFA( + cute::make_shape(m, N, K, 1)); + a.layout_SFB[e] = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB( + cute::make_shape(m, N, K, 1)); +} + +} // namespace + +size_t moe_grouped_gemm_nvfp4_sm100_scratch_bytes(int groups) { + if (groups <= 0) return 0; + // The CUTLASS workspace for a grouped launch scales with the group count and + // the scheduler, not with the token counts, so a bound taken at construction + // stays valid however the routing falls. + return args_bytes(groups) + static_cast(groups) * 1024 + (1u << 20); +} + +int moe_grouped_gemm_nvfp4_sm100_bf16out( + const void* A_packed, + const void* SFA, + const void* W_stack, + const void* SFB_stack, + const void* alpha_dev, + void* D, + const void* group_off, + const void* sfa_off, + int groups, + int N, + int K, + long w_stride, + long sfb_stride, + void* scratch, + size_t scratch_bytes, + cudaStream_t stream) { + if (!A_packed || !SFA || !W_stack || !SFB_stack || !alpha_dev || !D + || !group_off || !sfa_off || !scratch) return 1; + if (groups <= 0 || N <= 0 || K <= 0 || (K & 15) != 0) return 2; + const size_t need = args_bytes(groups); + if (scratch_bytes < need) return 3; + + GroupArgs ga = carve(scratch, groups); + const int threads = 128; + fill_group_args<<<(groups + threads - 1) / threads, threads, 0, stream>>>( + ga, + static_cast(A_packed), + static_cast(SFA), + static_cast(W_stack), + static_cast(SFB_stack), + static_cast(alpha_dev), + static_cast(D), + static_cast(group_off), + static_cast(sfa_off), + groups, N, K, w_stride, sfb_stride); + + cutlass::KernelHardwareInfo hw_info; + hw_info.device_id = 0; + hw_info.sm_count = + cutlass::KernelHardwareInfo::query_device_multiprocessor_count(0); + hw_info.cluster_shape = dim3(1, 1, 1); + hw_info.cluster_shape_fallback = dim3(1, 1, 1); + + typename Gemm::Arguments args_proto{}; + // The fusion argument type is reachable only through an Arguments instance, + // which is how the CUTLASS example spells it too. + decltype(args_proto.epilogue.thread) fusion_args; + fusion_args.alpha = 0.0f; + fusion_args.alpha_ptr_array = ga.ptr_alpha; + fusion_args.dAlpha = {_0{}, _0{}, 1}; + fusion_args.beta = 0.0f; + fusion_args.beta_ptr_array = nullptr; + fusion_args.dBeta = {_0{}, _0{}, 0}; + + typename Gemm::GemmKernel::TileSchedulerArguments scheduler{}; + + // Host-side problem shapes are passed as nullptr deliberately: the shapes + // live only on device, so nothing here depends on the routing and the call + // is safe to capture. + typename Gemm::Arguments args{ + cutlass::gemm::GemmUniversalMode::kGrouped, + {groups, ga.shapes, nullptr}, + {ga.ptr_A, ga.stride_A, ga.ptr_B, ga.stride_B, + ga.ptr_SFA, ga.layout_SFA, ga.ptr_SFB, ga.layout_SFB}, + {fusion_args, nullptr, ga.stride_C, ga.ptr_D, ga.stride_D}, + hw_info, scheduler}; + + Gemm gemm; + const size_t ws = Gemm::get_workspace_size(args); + if (need + ws > scratch_bytes) return 4; + void* ws_ptr = static_cast(scratch) + need; + + auto status = gemm.can_implement(args); + if (status != cutlass::Status::kSuccess) { + std::fprintf(stderr, + "[moe_grouped_gemm_nvfp4_sm100] can_implement FAIL groups=%d N=%d " + "K=%d status=%d\n", groups, N, K, static_cast(status)); + return 5; + } + status = gemm.initialize(args, ws_ptr, stream); + if (status != cutlass::Status::kSuccess) return 6; + status = gemm.run(stream); + return status == cutlass::Status::kSuccess ? 0 : 7; +} + +} // namespace gemm +} // namespace flash_rt diff --git a/csrc/gemm/fp4/cutlass_nvfp4_moe_grouped_sm100.cuh b/csrc/gemm/fp4/cutlass_nvfp4_moe_grouped_sm100.cuh new file mode 100644 index 00000000..84616804 --- /dev/null +++ b/csrc/gemm/fp4/cutlass_nvfp4_moe_grouped_sm100.cuh @@ -0,0 +1,72 @@ +// SPDX-License-Identifier: Apache-2.0 +// +// Grouped NVFP4 block-scaled GEMM for sm_100-class Blackwell (datacenter +// SM100 / Jetson AGX Thor SM110): every routed expert of a MoE layer in one +// launch. +// +// Why this exists. Prefill routes S*8 tokens across 256 experts. Serving that +// as one GEMV per (token, expert) slot re-reads an expert's weight once per +// token that chose it -- 9.7 GB a layer at S=1024 -- and serving it as one +// GEMM per expert costs 256 launches and 256 Python iterations a layer, whose +// host time exceeded the device time. Neither scales: the first is bounded by +// L2 bandwidth, the second by the host. +// +// A grouped GEMM is bounded by neither. One launch covers every expert, each +// weight is read once, and -- because CUTLASS accepts the per-group problem +// shapes from device memory (the host-side array is optional) -- the launch +// geometry is host-known and the routing is not, which is what a CUDA-graph +// capture requires. That is the property this is really for: capture a prefill +// chunk once and replay it for any context length, rather than tuning a +// threshold per prompt length. +// +// Wire format matches cutlass_nvfp4_w4a16_gemm_sm100: e2m1 nibbles, UE4M3 +// block scales of 16, Sm1xx block-scaled atom layout, BF16 out. The per-expert +// global scale enters as the epilogue's per-group alpha. + +#pragma once + +#include +#include + +namespace flash_rt { +namespace gemm { + +// Scratch the entry point needs, in bytes, for `groups` groups. Holds the +// per-group pointer/stride/layout arrays it fills on device, plus the CUTLASS +// workspace. Allocate once and reuse; it does not depend on the token counts. +size_t moe_grouped_gemm_nvfp4_sm100_scratch_bytes(int groups); + +// D[off_e : off_e + cnt_e, :] = A[off_e : off_e + cnt_e, :] @ W[e].T * alpha[e] +// +// A_packed (slots, K/2) u8 rows sorted by expert +// SFA per-group block-scaled atom layouts, group e at byte offset +// sfa_offsets[e] +// W_stack (E, N, K/2) u8 +// SFB_stack (E, sfb_bytes) u8 +// alpha_dev (E,) f32 device +// D (slots, N) bf16 +// group_off (E + 1,) i32 device, prefix sums of the per-expert counts +// sfa_off (E,) i32 device, byte offset of group e's SFA block +// +// Nothing is read to the host: the group shapes are derived on device from +// group_off. Returns 0 on success, nonzero on argument or CUTLASS error. +int moe_grouped_gemm_nvfp4_sm100_bf16out( + const void* A_packed, + const void* SFA, + const void* W_stack, + const void* SFB_stack, + const void* alpha_dev, + void* D, + const void* group_off, + const void* sfa_off, + int groups, + int N, + int K, + long w_stride, + long sfb_stride, + void* scratch, + size_t scratch_bytes, + cudaStream_t stream); + +} // namespace gemm +} // namespace flash_rt diff --git a/csrc/kernels/quantize.cu b/csrc/kernels/quantize.cu index a25a7be7..a7d55ccf 100644 --- a/csrc/kernels/quantize.cu +++ b/csrc/kernels/quantize.cu @@ -2804,3 +2804,97 @@ void dequant_int32_to_bf16(const int32_t* input, __nv_bfloat16* output, dequant_int32_to_bf16_kernel<<>>( input, output, d_act_scale, d_weight_scale, n); } + +// ── Grouped activation quantiser for the MoE grouped GEMM ── +// +// Same math as quantize_bf16_to_nvfp4_swizzled_kernel, block for block; what +// differs is where the scale factors land. The block-scaled GEMM wants each +// group's scales in the Sm1xx atom layout for that group's own row count, and +// that layout blocks rows by 128, so a group beginning at an arbitrary row of a +// jointly-quantised matrix has no contiguous sub-block to point at. Quantising +// per group is correct but costs a launch and a host iteration per expert. +// +// Here a row reads the expert it was sorted by, subtracts its group's first +// row, and indexes its group's own block. Nothing reaches the host, which is +// what lets the surrounding prefill chunk be captured. +__global__ void moe_grouped_quant_nvfp4_kernel( + const __nv_bfloat16* __restrict__ input, + const int* __restrict__ expert_of_row, + const int* __restrict__ group_off, + const int* __restrict__ sfa_off, + uint8_t* __restrict__ fp4_data, + uint8_t* __restrict__ scale_factors, + int cols, int num_blocks, int n_col_blocks) +{ + const int row = blockIdx.x; + const int e = expert_of_row[row]; + const int local = row - group_off[e]; // row index inside its group + const __nv_bfloat16* row_in = input + (size_t)row * cols; + uint8_t* row_fp4 = fp4_data + (size_t)row * cols / 2; + uint8_t* sf_base = scale_factors + sfa_off[e]; + + extern __shared__ float smem[]; + + for (int b = threadIdx.x; b < num_blocks; b += blockDim.x) smem[b] = 0.0f; + __syncthreads(); + + for (int i = threadIdx.x; i < cols; i += blockDim.x) { + float val = fabsf(__bfloat162float(row_in[i])); + atomicMax((int*)&smem[i >> 4], __float_as_int(val)); + } + __syncthreads(); + + const int rb = local / 128; + const int ri = local % 128; + for (int b = threadIdx.x; b < num_blocks; b += blockDim.x) { + float amax = __int_as_float(*(int*)&smem[b]); + uint8_t ue_scale = float_to_ue4m3_ceil(amax / 6.0f); + const int cb = b / 4; + const int ci = b % 4; + sf_base[(rb * n_col_blocks + cb) * 512 + (ri % 32) * 16 + + (ri / 32) * 4 + ci] = ue_scale; + smem[b] = ue4m3_to_float(ue_scale); + } + __syncthreads(); + + const int half_cols = cols >> 1; + for (int p = threadIdx.x; p < half_cols; p += blockDim.x) { + const int i = p * 2; + const int blk = i >> 4; + float scale = smem[blk]; + float inv_scale = (scale > 0.0f) ? (1.0f / scale) : 0.0f; + float v0 = __bfloat162float(row_in[i]) * inv_scale; + float v1 = __bfloat162float(row_in[i + 1]) * inv_scale; + const int blk1 = (i + 1) >> 4; + if (blk1 != blk) { + float s1 = smem[blk1]; + float inv1 = (s1 > 0.0f) ? (1.0f / s1) : 0.0f; + v1 = __bfloat162float(row_in[i + 1]) * inv1; + } + row_fp4[p] = (uint8_t)((float_to_fp4_e2m1(v1) << 4) + | (float_to_fp4_e2m1(v0) & 0x0F)); + } +} + +int moe_grouped_quant_nvfp4_bf16( + const void* A, const void* expert_of_row, const void* group_off, + const void* sfa_off, void* out_packed, void* out_sf, + int slots, int K, cudaStream_t stream) +{ + if (!A || !expert_of_row || !group_off || !sfa_off || !out_packed + || !out_sf) return 1; + if (slots <= 0 || K <= 0 || (K & 15) != 0) return 2; + const int num_blocks = K / 16; + const int n_col_blocks = (num_blocks + 3) / 4; + const int threads = 256; + const size_t smem = (size_t)num_blocks * sizeof(float); + moe_grouped_quant_nvfp4_kernel<<>>( + reinterpret_cast(A), + reinterpret_cast(expert_of_row), + reinterpret_cast(group_off), + reinterpret_cast(sfa_off), + reinterpret_cast(out_packed), + reinterpret_cast(out_sf), + K, num_blocks, n_col_blocks); + return 0; +} diff --git a/csrc/kernels/quantize.cuh b/csrc/kernels/quantize.cuh index cdc7bded..4ba544d0 100644 --- a/csrc/kernels/quantize.cuh +++ b/csrc/kernels/quantize.cuh @@ -325,3 +325,21 @@ void quantize_int8_rowwise_static(const __nv_bfloat16* input, int8_t* output, void dequant_int32_to_bf16(const int32_t* input, __nv_bfloat16* output, const float* d_act_scale, const float* d_weight_scale, int n, cudaStream_t stream = 0); + +// Grouped activation quantiser for the MoE grouped GEMM: every expert's block +// in one launch. Same math as quantize_bf16_to_nvfp4_swizzled; what differs is +// that each group's scale factors go into the Sm1xx atom layout for that +// group's own row count, which is what the block-scaled grouped GEMM reads. +// Quantising per group instead is correct but costs a launch and a host +// iteration per expert -- and a host iteration is what a graph capture cannot +// have. +// +// A (slots, K) bf16, rows already sorted by expert +// expert_of_row (slots,) i32 +// group_off (E + 1,) i32 prefix sums of the per-expert row counts +// sfa_off (E,) i32 byte offset of each group's SF block +// K must be a multiple of 16. Returns 0 on success, nonzero on arg error. +int moe_grouped_quant_nvfp4_bf16( + const void* A, const void* expert_of_row, const void* group_off, + const void* sfa_off, void* out_packed, void* out_sf, + int slots, int K, cudaStream_t stream); diff --git a/flash_rt/frontends/torch/_nexn2_rtx_forward.py b/flash_rt/frontends/torch/_nexn2_rtx_forward.py index b72ef597..308c3fac 100644 --- a/flash_rt/frontends/torch/_nexn2_rtx_forward.py +++ b/flash_rt/frontends/torch/_nexn2_rtx_forward.py @@ -971,6 +971,111 @@ def _moe_experts_bt(x, ti, tw, ld, fvk, device): return out +_GROUPED_SCRATCH = {} + + +def _grouped_scratch(fvk, device): + """One scratch buffer per device for the grouped GEMM's descriptor arrays. + + Its size depends only on the expert count, not on how the routing falls, so + it is allocated once and never resized -- which is also what lets the call + sit inside a captured region. + """ + key = str(device) + got = _GROUPED_SCRATCH.get(key) + if got is None: + nbytes = int(fvk.moe_grouped_gemm_nvfp4_sm100_scratch_bytes(_N_EXPERTS)) + got = (torch.empty(nbytes, dtype=torch.uint8, device=device), nbytes) + _GROUPED_SCRATCH[key] = got + return got + + +def _sf_layout(counts, k, device): + """Per-group scale-factor byte offsets, and a host-known bound on the total. + + The block-scaled layout blocks rows by 128, so a group of c rows needs + ceil(c/128) super-blocks. Summing that needs the counts, which live on the + device -- but the sum is bounded by (experts + slots/128) super-blocks + whatever the routing does, and that bound follows from the shapes alone. + Sizing the buffer from the bound rather than the sum is what keeps this free + of a host read. + """ + n_col = ((k // 16) + 3) // 4 + per_group = ((counts + 127) // 128) * (n_col * 512) + off = torch.zeros(_N_EXPERTS, dtype=torch.int32, device=device) + off[1:] = per_group.cumsum(0)[:-1].to(torch.int32) + return off, n_col + + +def _moe_experts_grouped_gemm(x, ti, tw, ld, fvk, device): + """Every routed expert of the layer in two GEMM launches. + + The per-expert loop below reads each weight once, which is the right amount, + but pays a launch and a host iteration per expert -- and the host iteration + is fatal twice over: it dominated the time at S=1024, and it makes the layer + impossible to capture. A grouped GEMM takes the per-group shapes from device + memory, so the launch geometry depends only on the expert count and the + routing never reaches the host. + + Measured against the loop it replaces, at the shapes prefill issues: 6.0x on + gate_up and 14.5x on down at S=1024, 512 launches down to 2, output bitwise + identical. + """ + S = x.shape[0] + slots = S * TOPK + gu_p, gu_s = ld['experts_gate_up_packed_t'], ld['experts_gate_up_sf_t'] + dn_p, dn_s = ld['experts_down_packed_t'], ld['experts_down_sf_t'] + n_gu, n_dn = gu_p.shape[1], dn_p.shape[1] + if 'experts_gate_up_alpha_dev' not in ld: + ld['experts_gate_up_alpha_dev'] = \ + ld['experts_gate_up_alpha_t'].to(device).contiguous() + ld['experts_down_alpha_dev'] = \ + ld['experts_down_alpha_t'].to(device).contiguous() + gu_a, dn_a = ld['experts_gate_up_alpha_dev'], ld['experts_down_alpha_dev'] + scratch, scratch_bytes = _grouped_scratch(fvk, device) + + exp_flat = ti.reshape(-1).to(torch.int32) + tok_flat = torch.arange(S, device=device).repeat_interleave(TOPK) + order = exp_flat.argsort(stable=True) + se = exp_flat[order].contiguous() + stok = tok_flat[order] + sw = tw.reshape(-1)[order] + + counts = torch.bincount(se, minlength=_N_EXPERTS) + group_off = torch.zeros(_N_EXPERTS + 1, dtype=torch.int32, device=device) + group_off[1:] = counts.cumsum(0).to(torch.int32) + + def project(A, k, n, w_p, w_s, alpha, out): + sfa_off, n_col = _sf_layout(counts, k, device) + bound = (_N_EXPERTS + slots // 128 + 1) * n_col * 512 + packed = torch.empty(slots, k // 2, dtype=torch.uint8, device=device) + sfa = torch.empty(bound, dtype=torch.uint8, device=device) + rc = fvk.moe_grouped_quant_nvfp4_bf16( + A.data_ptr(), se.data_ptr(), group_off.data_ptr(), + sfa_off.data_ptr(), packed.data_ptr(), sfa.data_ptr(), + slots, k, _cs()) + if rc: + raise RuntimeError(f'grouped activation quant failed with {rc}') + rc = fvk.moe_grouped_gemm_nvfp4_sm100_bf16out( + packed.data_ptr(), sfa.data_ptr(), w_p.data_ptr(), + w_s.data_ptr(), alpha.data_ptr(), out.data_ptr(), + group_off.data_ptr(), sfa_off.data_ptr(), + _N_EXPERTS, n, k, w_p[0].numel(), w_s[0].numel(), + scratch.data_ptr(), scratch_bytes, _cs()) + if rc: + raise RuntimeError(f'grouped MoE GEMM failed with {rc}') + + d_gu = torch.empty(slots, n_gu, dtype=torch.bfloat16, device=device) + project(x[stok].contiguous(), HID, n_gu, gu_p, gu_s, gu_a, d_gu) + inter = _silu_mul(d_gu[:, :INTER], d_gu[:, INTER:], fvk, device) + d_dn = torch.empty(slots, n_dn, dtype=torch.bfloat16, device=device) + project(inter, INTER, n_dn, dn_p, dn_s, dn_a, d_dn) + + out = torch.zeros(S, HID, device=device) + out.index_add_(0, stok, d_dn.float() * sw.unsqueeze(-1)) + return out + + def _moe_experts_per_expert_gemm(x, ti, tw, ld, fvk, device): """Routed experts as one GEMM per expert over the tokens that chose it. @@ -1115,6 +1220,10 @@ def _moe_layer(h, ld, fvk, device): out = _moe_experts_bt(x, ti, tw, ld, fvk, device) elif _USE_M16_MOE and big and hasattr(fvk, 'moe_m16_mma_sm120_bf16'): out = _moe_experts_m16(x, ti, tw, ld, fvk, device) + elif big and hasattr(fvk, 'moe_grouped_gemm_nvfp4_sm100_bf16out'): + # No threshold: the grouped path wins at every prefill length measured, + # because it does not pay per expert for anything. + out = _moe_experts_grouped_gemm(x, ti, tw, ld, fvk, device) elif (big and _USE_PER_EXPERT_GEMM and x.shape[0] * TOPK >= _PER_EXPERT_MIN_M * _N_EXPERTS and hasattr(fvk, 'fp4_w4a16_gemm_sm120_bf16out')): From 4c6cdda2540f4a13633ab4c631a62bb716f2e33c Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Fri, 31 Jul 2026 04:50:01 -0400 Subject: [PATCH 41/85] Unpermute the grouped MoE without atomics index_add_ was 37.8 ms of a 1024-token prefill and reduces through atomics, so its accumulation order varies between runs. Prefill seeds the decode state, which has to be reproducible, so that is not a cost worth paying twice. The block-tile path already solved this: invert the routing permutation and let one kernel sum each token's TOPK rows in fixed order. Same reducer, reused. TTFT: 256 182.3 -> 174.8 ms, 1024 479.2 -> 399.4, 2048 1014.7 -> 847.7. Fixture 16/16. --- flash_rt/frontends/torch/_nexn2_rtx_forward.py | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/flash_rt/frontends/torch/_nexn2_rtx_forward.py b/flash_rt/frontends/torch/_nexn2_rtx_forward.py index 308c3fac..f99f02e7 100644 --- a/flash_rt/frontends/torch/_nexn2_rtx_forward.py +++ b/flash_rt/frontends/torch/_nexn2_rtx_forward.py @@ -1071,8 +1071,20 @@ def project(A, k, n, w_p, w_s, alpha, out): d_dn = torch.empty(slots, n_dn, dtype=torch.bfloat16, device=device) project(inter, INTER, n_dn, dn_p, dn_s, dn_a, d_dn) - out = torch.zeros(S, HID, device=device) - out.index_add_(0, stok, d_dn.float() * sw.unsqueeze(-1)) + # Deterministic unpermute, the same one the block-tile path uses: invert the + # routing permutation and let one kernel sum each token's TOPK rows in fixed + # order. index_add_ was 37.8 ms of a 1024-token prefill and reduces through + # atomics, so its order varies -- which prefill cannot afford, since it + # seeds a decode that has to be reproducible. + inv = torch.empty(slots, dtype=torch.long, device=device) + inv[order] = torch.arange(slots, device=device) + rows = (torch.arange(slots, dtype=torch.int32, device=device)[inv] + ).contiguous() + out = torch.empty(S, HID, dtype=torch.float32, device=device) + fvk.moe_weighted_sum_sm120_bf16( + d_dn.data_ptr(), rows.data_ptr(), + tw.reshape(S, TOPK).contiguous().data_ptr(), out.data_ptr(), + S, TOPK, n_dn, n_dn, _cs()) return out From 4a058b377d0922a21efd677752373de872cdb039 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Fri, 31 Jul 2026 04:56:36 -0400 Subject: [PATCH 42/85] Stop forcing prefill attention onto the math backend The explicit causal mask this file builds was mine, added so a chunked block's queries attend to the right window. It also disqualifies the fused attention backends, so a square block fell through to the math one: scores materialised and run through SIMT fp32 GEMMs, 131.7 ms of a 2048-token prefill in twenty launches, growing as S^2. When the block is square the bottom-right and top-left causal conventions coincide, so is_causal says the same thing and the fused backend takes it. The mask stays for Sq < Sk, which is what chunked prefill actually needs it for. TTFT: 1024 399.4 -> 331.3 ms, 2048 847.7 -> 587.6. Fixture 16/16. Found by fixing the profiler rather than the code: the attention bucket read 74.7 ms because its pattern matched `flash` and so claimed every unclaimed flash_rt kernel. The real attention kernel was 6.1 ms, and the quadratic term was sitting in the dense-GEMM bucket under a cutlass SIMT name. A CUTLASS Blackwell FMHA was about to be built on that misreading. --- flash_rt/frontends/torch/_nexn2_rtx_forward.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/flash_rt/frontends/torch/_nexn2_rtx_forward.py b/flash_rt/frontends/torch/_nexn2_rtx_forward.py index f99f02e7..5d3ad7c9 100644 --- a/flash_rt/frontends/torch/_nexn2_rtx_forward.py +++ b/flash_rt/frontends/torch/_nexn2_rtx_forward.py @@ -623,6 +623,15 @@ def _sdpa_causal_attn(qf, kf, vf, device): Sq, Sk = qf.shape[1], kf.shape[1] q = qf.transpose(1, 2) # (1, NQ, Sq, HD) k, v = kf.transpose(1, 2), vf.transpose(1, 2) # (1, NKV, Sk, HD) + # An explicit mask forces the math backend, which materialises the scores + # and runs them through SIMT fp32 GEMMs -- 131.7 ms of a 2048-token prefill + # in twenty launches, growing as S^2. When the block is square the two + # causal conventions coincide, so say is_causal and let the fused backend + # take it; the mask is only needed when Sq < Sk, which is chunked prefill. + if Sq == Sk: + return F.scaled_dot_product_attention( + q, k, v, is_causal=True, scale=float(HD) ** -0.5, enable_gqa=True + ).transpose(1, 2).contiguous() qi = torch.arange(Sk - Sq, Sk, device=device).unsqueeze(1) mask = torch.arange(Sk, device=device).unsqueeze(0) <= qi try: From 8dc4ac4386028a6d83f040f9061340123bd2863b Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Fri, 31 Jul 2026 05:01:47 -0400 Subject: [PATCH 43/85] Take the atomics out of the grouped activation quantiser The grouped quantiser reduced each 16-block's maximum with one shared-memory atomicMax per element. At 2048 tokens that came to 58.3 ms for traffic worth 0.8 -- seventy times its roofline, and the largest single kernel after the GEMM it feeds. Now a thread reduces eight bf16 in registers and pairs with its neighbour through a shuffle, and the pack writes four bytes at a time from the same vector load. Still bitwise identical to the per-expert loop it is checked against, which runs the original quantiser -- so the rewrite kept the numerics exactly. TTFT: 256 169.1 -> 160.5 ms, 1024 331.3 -> 321.3, 2048 587.6 -> 566.3. Fixture 16/16. --- csrc/kernels/quantize.cu | 59 +++++++++++++++++++++++----------------- 1 file changed, 34 insertions(+), 25 deletions(-) diff --git a/csrc/kernels/quantize.cu b/csrc/kernels/quantize.cu index a7d55ccf..117a5cfe 100644 --- a/csrc/kernels/quantize.cu +++ b/csrc/kernels/quantize.cu @@ -2834,21 +2834,29 @@ __global__ void moe_grouped_quant_nvfp4_kernel( uint8_t* sf_base = scale_factors + sfa_off[e]; extern __shared__ float smem[]; - - for (int b = threadIdx.x; b < num_blocks; b += blockDim.x) smem[b] = 0.0f; - __syncthreads(); - - for (int i = threadIdx.x; i < cols; i += blockDim.x) { - float val = fabsf(__bfloat162float(row_in[i])); - atomicMax((int*)&smem[i >> 4], __float_as_int(val)); + const int tid = threadIdx.x; + + // Per-16-block amax without atomics. One thread takes eight bf16 (a half + // block), reduces them in registers, and pairs with its neighbour through a + // shuffle -- j and j^1 land on lanes t and t^1 because the block size is + // even. The first version of this kernel used one atomicMax per element and + // ran at 58.3 ms for traffic worth 0.8; the atomics were all of it. + const int vec8 = cols >> 3; + for (int j = tid; j < vec8; j += blockDim.x) { + uint4 v = *reinterpret_cast(&row_in[j << 3]); + const __nv_bfloat16* bf = reinterpret_cast(&v); + float a = 0.0f; + #pragma unroll + for (int i = 0; i < 8; ++i) a = fmaxf(a, fabsf(__bfloat162float(bf[i]))); + a = fmaxf(a, __shfl_xor_sync(0xffffffffu, a, 1)); + if ((j & 1) == 0) smem[j >> 1] = a; } __syncthreads(); const int rb = local / 128; const int ri = local % 128; - for (int b = threadIdx.x; b < num_blocks; b += blockDim.x) { - float amax = __int_as_float(*(int*)&smem[b]); - uint8_t ue_scale = float_to_ue4m3_ceil(amax / 6.0f); + for (int b = tid; b < num_blocks; b += blockDim.x) { + uint8_t ue_scale = float_to_ue4m3_ceil(smem[b] * (1.0f / 6.0f)); const int cb = b / 4; const int ci = b % 4; sf_base[(rb * n_col_blocks + cb) * 512 + (ri % 32) * 16 @@ -2857,22 +2865,23 @@ __global__ void moe_grouped_quant_nvfp4_kernel( } __syncthreads(); - const int half_cols = cols >> 1; - for (int p = threadIdx.x; p < half_cols; p += blockDim.x) { - const int i = p * 2; - const int blk = i >> 4; - float scale = smem[blk]; - float inv_scale = (scale > 0.0f) ? (1.0f / scale) : 0.0f; - float v0 = __bfloat162float(row_in[i]) * inv_scale; - float v1 = __bfloat162float(row_in[i + 1]) * inv_scale; - const int blk1 = (i + 1) >> 4; - if (blk1 != blk) { - float s1 = smem[blk1]; - float inv1 = (s1 > 0.0f) ? (1.0f / s1) : 0.0f; - v1 = __bfloat162float(row_in[i + 1]) * inv1; + // Pack four bytes at a time: eight bf16 in, one uint32 out, and the eight + // share a 16-block so the scale is read once. + const int quads = cols >> 3; + for (int j = tid; j < quads; j += blockDim.x) { + uint4 v = *reinterpret_cast(&row_in[j << 3]); + const __nv_bfloat16* bf = reinterpret_cast(&v); + const float scale = smem[j >> 1]; + const float inv = (scale > 0.0f) ? (1.0f / scale) : 0.0f; + uint32_t packed = 0; + #pragma unroll + for (int k = 0; k < 4; ++k) { + uint32_t lo = float_to_fp4_e2m1(__bfloat162float(bf[2 * k]) * inv); + uint32_t hi = float_to_fp4_e2m1( + __bfloat162float(bf[2 * k + 1]) * inv); + packed |= ((hi << 4) | (lo & 0xF)) << (k * 8); } - row_fp4[p] = (uint8_t)((float_to_fp4_e2m1(v1) << 4) - | (float_to_fp4_e2m1(v0) & 0x0F)); + *reinterpret_cast(row_fp4 + (j << 2)) = packed; } } From d45a3d8cf08ff1bdb847900aeb05267a047c29df Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Fri, 31 Jul 2026 05:07:38 -0400 Subject: [PATCH 44/85] Gate and quantise the grouped MoE's intermediate in one pass The grouped GEMM leaves gate and up interleaved in one buffer, and slicing the two halves out of it is not free: they are strided, so both get copied -- 67 MB a layer at 2048 tokens, to feed a gate that writes another 17 and has it read straight back by the quantiser. Reading the merged buffer directly costs none of that. The silu is computed and rounded to bf16 exactly as the separate gate kernel does, so the value reaching the quantiser is unchanged, and the fixture still matches 16/16. TTFT: 256 160.5 -> 158.7 ms, 1024 321.3 -> 313.5, 2048 566.3 -> 550.1. --- csrc/bindings.cpp | 13 +++ csrc/kernels/quantize.cu | 87 +++++++++++++++++++ csrc/kernels/quantize.cuh | 9 ++ .../frontends/torch/_nexn2_rtx_forward.py | 22 +++-- 4 files changed, 124 insertions(+), 7 deletions(-) diff --git a/csrc/bindings.cpp b/csrc/bindings.cpp index 3e6a0a63..533ad40a 100644 --- a/csrc/bindings.cpp +++ b/csrc/bindings.cpp @@ -1089,6 +1089,19 @@ PYBIND11_MODULE(flash_rt_kernels, m) { }, py::arg("input"), py::arg("fp4_data"), py::arg("scale_factors"), py::arg("rows"), py::arg("cols"), py::arg("stream") = 0); + m.def("moe_grouped_silu_quant_nvfp4_bf16", + [](uintptr_t merged, uintptr_t expert_of_row, uintptr_t group_off, + uintptr_t sfa_off, uintptr_t out_packed, uintptr_t out_sf, + int slots, int inter, uintptr_t stream) -> int { + return moe_grouped_silu_quant_nvfp4_bf16( + to_ptr(merged), to_ptr(expert_of_row), to_ptr(group_off), + to_ptr(sfa_off), to_ptr(out_packed), to_ptr(out_sf), + slots, inter, to_stream(stream)); + }, + py::arg("merged"), py::arg("expert_of_row"), py::arg("group_off"), + py::arg("sfa_off"), py::arg("out_packed"), py::arg("out_sf"), + py::arg("slots"), py::arg("inter"), py::arg("stream") = 0); + m.def("moe_grouped_quant_nvfp4_bf16", [](uintptr_t A, uintptr_t expert_of_row, uintptr_t group_off, uintptr_t sfa_off, uintptr_t out_packed, uintptr_t out_sf, diff --git a/csrc/kernels/quantize.cu b/csrc/kernels/quantize.cu index 117a5cfe..7aaea5d7 100644 --- a/csrc/kernels/quantize.cu +++ b/csrc/kernels/quantize.cu @@ -2907,3 +2907,90 @@ int moe_grouped_quant_nvfp4_bf16( K, num_blocks, n_col_blocks); return 0; } + +// ── Gate and quantise in one pass, for the grouped MoE's down projection ── +// +// The grouped GEMM produces gate and up interleaved in one (slots, 2*inter) +// buffer, and the gate op wants them as two matrices. Slicing columns out of it +// is not free: the halves are strided, so `.contiguous()` copies both -- 67 MB +// a layer at 2048 tokens, to feed an op that then writes another 17 and has it +// read straight back by the quantiser. +// +// Reading the merged buffer directly costs none of that. The silu is computed +// and rounded to bf16 exactly as silu_mul_sm120_bf16 does, so the value that +// reaches the quantiser is the same one it saw before. +__global__ void moe_grouped_silu_quant_nvfp4_kernel( + const __nv_bfloat16* __restrict__ merged, // (slots, 2 * inter) + const int* __restrict__ expert_of_row, + const int* __restrict__ group_off, + const int* __restrict__ sfa_off, + uint8_t* __restrict__ fp4_data, + uint8_t* __restrict__ scale_factors, + int inter, int num_blocks, int n_col_blocks) +{ + const int row = blockIdx.x; + const int e = expert_of_row[row]; + const int local = row - group_off[e]; + const __nv_bfloat16* g_in = merged + (size_t)row * 2 * inter; + const __nv_bfloat16* u_in = g_in + inter; + uint8_t* row_fp4 = fp4_data + (size_t)row * inter / 2; + uint8_t* sf_base = scale_factors + sfa_off[e]; + + extern __shared__ float smem[]; // inter gated values, then scales + float* gated = smem; + float* scales = smem + inter; + + const int tid = threadIdx.x; + for (int i = tid; i < inter; i += blockDim.x) { + const float gv = __bfloat162float(g_in[i]); + const float uv = __bfloat162float(u_in[i]); + // Rounded to bf16 here, as the separate gate kernel does, so the + // quantiser downstream sees the identical value. + gated[i] = __bfloat162float( + __float2bfloat16_rn(gv / (1.0f + __expf(-gv)) * uv)); + } + __syncthreads(); + + for (int b = tid; b < num_blocks; b += blockDim.x) { + float a = 0.0f; + #pragma unroll 4 + for (int j = 0; j < 16; ++j) a = fmaxf(a, fabsf(gated[b * 16 + j])); + const uint8_t ue = float_to_ue4m3_ceil(a * (1.0f / 6.0f)); + const int rb = local / 128, ri = local % 128; + sf_base[(rb * n_col_blocks + (b >> 2)) * 512 + (ri % 32) * 16 + + (ri / 32) * 4 + (b & 3)] = ue; + scales[b] = ue4m3_to_float(ue); + } + __syncthreads(); + + const int half = inter >> 1; + for (int p = tid; p < half; p += blockDim.x) { + const int i = p * 2; + const float s = scales[i >> 4]; + const float inv = (s > 0.0f) ? (1.0f / s) : 0.0f; + row_fp4[p] = (uint8_t)((float_to_fp4_e2m1(gated[i + 1] * inv) << 4) + | (float_to_fp4_e2m1(gated[i] * inv) & 0x0F)); + } +} + +int moe_grouped_silu_quant_nvfp4_bf16( + const void* merged, const void* expert_of_row, const void* group_off, + const void* sfa_off, void* out_packed, void* out_sf, + int slots, int inter, cudaStream_t stream) +{ + if (!merged || !expert_of_row || !group_off || !sfa_off || !out_packed + || !out_sf) return 1; + if (slots <= 0 || inter <= 0 || (inter & 15) != 0) return 2; + const int num_blocks = inter / 16; + const int n_col_blocks = (num_blocks + 3) / 4; + const size_t smem = ((size_t)inter + num_blocks) * sizeof(float); + moe_grouped_silu_quant_nvfp4_kernel<<>>( + reinterpret_cast(merged), + reinterpret_cast(expert_of_row), + reinterpret_cast(group_off), + reinterpret_cast(sfa_off), + reinterpret_cast(out_packed), + reinterpret_cast(out_sf), + inter, num_blocks, n_col_blocks); + return 0; +} diff --git a/csrc/kernels/quantize.cuh b/csrc/kernels/quantize.cuh index 4ba544d0..ccfc3feb 100644 --- a/csrc/kernels/quantize.cuh +++ b/csrc/kernels/quantize.cuh @@ -343,3 +343,12 @@ int moe_grouped_quant_nvfp4_bf16( const void* A, const void* expert_of_row, const void* group_off, const void* sfa_off, void* out_packed, void* out_sf, int slots, int K, cudaStream_t stream); + +// Gate and quantise in one pass: reads the grouped GEMM's merged (slots, +// 2*inter) gate/up output directly, so the strided column halves are never +// copied out. The silu is rounded to bf16 exactly as silu_mul_sm120_bf16 does, +// so the quantiser sees the same value it did when the two were separate. +int moe_grouped_silu_quant_nvfp4_bf16( + const void* merged, const void* expert_of_row, const void* group_off, + const void* sfa_off, void* out_packed, void* out_sf, + int slots, int inter, cudaStream_t stream); diff --git a/flash_rt/frontends/torch/_nexn2_rtx_forward.py b/flash_rt/frontends/torch/_nexn2_rtx_forward.py index 5d3ad7c9..00e28824 100644 --- a/flash_rt/frontends/torch/_nexn2_rtx_forward.py +++ b/flash_rt/frontends/torch/_nexn2_rtx_forward.py @@ -1054,15 +1054,24 @@ def _moe_experts_grouped_gemm(x, ti, tw, ld, fvk, device): group_off = torch.zeros(_N_EXPERTS + 1, dtype=torch.int32, device=device) group_off[1:] = counts.cumsum(0).to(torch.int32) - def project(A, k, n, w_p, w_s, alpha, out): + def project(A, k, n, w_p, w_s, alpha, out, gate=False): sfa_off, n_col = _sf_layout(counts, k, device) bound = (_N_EXPERTS + slots // 128 + 1) * n_col * 512 packed = torch.empty(slots, k // 2, dtype=torch.uint8, device=device) sfa = torch.empty(bound, dtype=torch.uint8, device=device) - rc = fvk.moe_grouped_quant_nvfp4_bf16( - A.data_ptr(), se.data_ptr(), group_off.data_ptr(), - sfa_off.data_ptr(), packed.data_ptr(), sfa.data_ptr(), - slots, k, _cs()) + if gate: + # A is the merged (slots, 2k) gate/up output: gate it and quantise + # in one pass rather than slicing two strided halves out of it, + # copying both, gating into a third buffer and reading that back. + rc = fvk.moe_grouped_silu_quant_nvfp4_bf16( + A.data_ptr(), se.data_ptr(), group_off.data_ptr(), + sfa_off.data_ptr(), packed.data_ptr(), sfa.data_ptr(), + slots, k, _cs()) + else: + rc = fvk.moe_grouped_quant_nvfp4_bf16( + A.data_ptr(), se.data_ptr(), group_off.data_ptr(), + sfa_off.data_ptr(), packed.data_ptr(), sfa.data_ptr(), + slots, k, _cs()) if rc: raise RuntimeError(f'grouped activation quant failed with {rc}') rc = fvk.moe_grouped_gemm_nvfp4_sm100_bf16out( @@ -1076,9 +1085,8 @@ def project(A, k, n, w_p, w_s, alpha, out): d_gu = torch.empty(slots, n_gu, dtype=torch.bfloat16, device=device) project(x[stok].contiguous(), HID, n_gu, gu_p, gu_s, gu_a, d_gu) - inter = _silu_mul(d_gu[:, :INTER], d_gu[:, INTER:], fvk, device) d_dn = torch.empty(slots, n_dn, dtype=torch.bfloat16, device=device) - project(inter, INTER, n_dn, dn_p, dn_s, dn_a, d_dn) + project(d_gu, INTER, n_dn, dn_p, dn_s, dn_a, d_dn, gate=True) # Deterministic unpermute, the same one the block-tile path uses: invert the # routing permutation and let one kernel sum each token's TOPK rows in fixed From 1d30ed83f21eedc537d32e87f77ede8b20909d6e Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Fri, 31 Jul 2026 05:27:43 -0400 Subject: [PATCH 45/85] Stop upcasting the KV cache to fp32 on every decode step The reference decode attention read the cache and cast it to fp32. The cache is bf16 and holds the whole history, so that materialises the history again, in twice the width, once per layer per token: hundreds of MB of temporaries a token at a long context. It is why decode fell from 26 to 11 tok/s between 4k and 10k while the engine we compare against stayed flat. SDPA accumulates in fp32 whatever the inputs are, so the cast bought nothing. Decode at 10k context: 11.2 -> 81.6 tok/s. Fixture 16/16. --- flash_rt/hardware/rtx/attn_backend_nexn2.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/flash_rt/hardware/rtx/attn_backend_nexn2.py b/flash_rt/hardware/rtx/attn_backend_nexn2.py index c965ce62..07e5de31 100644 --- a/flash_rt/hardware/rtx/attn_backend_nexn2.py +++ b/flash_rt/hardware/rtx/attn_backend_nexn2.py @@ -172,7 +172,11 @@ def _sdpa(self, layer_idx: int, q_seq: int, kv_seq: int, """Reference attention, for a device the vendored kernel refuses.""" import torch.nn.functional as F - q = self.Q_buf[:, :q_seq].transpose(1, 2).float() + # BF16 throughout. The cache is bf16, so upcasting it materialises the + # whole history in fp32 every step -- hundreds of MB of temporaries per + # token at a long context, which is what made decode fall from 26 to 11 + # tok/s between 4k and 10k. SDPA accumulates in fp32 regardless. + q = self.Q_buf[:, :q_seq].transpose(1, 2) k = self.K_cache[layer_idx:layer_idx + 1, :kv_seq] v = self.V_cache[layer_idx:layer_idx + 1, :kv_seq] # Broadcasting the KV to the query head count materialises it: at @@ -183,14 +187,14 @@ def _sdpa(self, layer_idx: int, q_seq: int, kv_seq: int, groups = self.NUM_Q_HEADS // self.NUM_KV_HEADS try: out = F.scaled_dot_product_attention( - q, k.transpose(1, 2).float(), v.transpose(1, 2).float(), + q, k.transpose(1, 2), v.transpose(1, 2), is_causal=q_seq > 1, scale=softmax_scale, enable_gqa=True, ).transpose(1, 2) except TypeError: # torch without native GQA out = F.scaled_dot_product_attention( q, - k.repeat_interleave(groups, dim=2).transpose(1, 2).float(), - v.repeat_interleave(groups, dim=2).transpose(1, 2).float(), + k.repeat_interleave(groups, dim=2).transpose(1, 2), + v.repeat_interleave(groups, dim=2).transpose(1, 2), is_causal=q_seq > 1, scale=softmax_scale, ).transpose(1, 2) self.O_buf[:, :q_seq].copy_(out.to(self.O_buf.dtype)) From b19fc75c31716abb3ea2e3971b1a81f139173d9b Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Fri, 31 Jul 2026 05:41:20 -0400 Subject: [PATCH 46/85] Record what chunked prefill's attention can and cannot use Three candidates for the bottom-right causal window a chunked block needs, measured rather than reasoned about: is_causal wrong. It means top-left, so it silently drops the history a chunk is supposed to attend to: cos 0.24 to 0.30 against the mask that states the real window. explicit mask right, and what the code uses, but it forces the math backend, which materialises the scores. That is what it costs today and what stops a long context outright. flex_attention right (cos 0.999997) and linear in memory, but it compiles per shape and a chunked prefill hands it a new (Sq, Sk) per chunk: 10240 tokens 4071 -> 4330 ms, 16384 -> 13858. So the mask stays, with the alternatives written down next to it and the flex builder left in place for when the chunk shapes are fixed and warmed. Also found the ceiling on single-chunk prefill: causal_conv1d launches dim3(_, S, B), and gridDim.y stops at 65535, so S = 65536 fails with an invalid argument rather than anything numerical. --- .../frontends/torch/_nexn2_rtx_forward.py | 48 +++++++++++++++++++ 1 file changed, 48 insertions(+) diff --git a/flash_rt/frontends/torch/_nexn2_rtx_forward.py b/flash_rt/frontends/torch/_nexn2_rtx_forward.py index 00e28824..2edefbc0 100644 --- a/flash_rt/frontends/torch/_nexn2_rtx_forward.py +++ b/flash_rt/frontends/torch/_nexn2_rtx_forward.py @@ -609,6 +609,45 @@ def _fa2_usable(device): return _FA2_USABLE +_FLEX_CACHE = {} + + +def _flex_causal(sq, sk, device): + """A bottom-right-causal flex_attention closure for this block shape. + + Block masks are built per (Sq, Sk) and cached, because a chunked prefill + revisits the same shapes as it walks the prompt. Returns None where flex is + unavailable, leaving the explicit-mask path to handle it. + """ + key = (sq, sk, str(device)) + got = _FLEX_CACHE.get(key) + if got is not None: + return got + if key in _FLEX_CACHE: + return None + try: + from torch.nn.attention.flex_attention import ( + create_block_mask, flex_attention, + ) + + off = sk - sq + + def mask_mod(b, h, q_idx, kv_idx): + return kv_idx <= q_idx + off + + block_mask = create_block_mask(mask_mod, 1, NQ, sq, sk, device=device) + + def run(q, k, v): + return flex_attention(q, k, v, block_mask=block_mask, + scale=float(HD) ** -0.5, enable_gqa=True) + + _FLEX_CACHE[key] = run + return run + except Exception: # noqa: BLE001 + _FLEX_CACHE[key] = None + return None + + def _sdpa_causal_attn(qf, kf, vf, device): """Reference causal GQA attention, for a build without the FA2 kernel. @@ -632,6 +671,15 @@ def _sdpa_causal_attn(qf, kf, vf, device): return F.scaled_dot_product_attention( q, k, v, is_causal=True, scale=float(HD) ** -0.5, enable_gqa=True ).transpose(1, 2).contiguous() + # A chunked block's window is bottom-right causal, which is_causal does not + # mean (measured: cos 0.24 against this mask, i.e. it silently truncates the + # history) and which a boolean mask only expresses by materialising the + # scores. flex_attention states it as a predicate and skips fully-masked + # blocks -- numerically right, cos 0.999997 -- but it compiles per shape, + # and a chunked prefill hands it a new (Sq, Sk) for every chunk: measured + # 10240 tokens 4071 -> 4330 ms, 16384 tokens 13858. Left out on that + # evidence; it becomes the right answer once the chunk shapes are fixed and + # warmed, which is where the long-context work goes next. qi = torch.arange(Sk - Sq, Sk, device=device).unsqueeze(1) mask = torch.arange(Sk, device=device).unsqueeze(0) <= qi try: From cc662b624ef851212e0dbba08d102406bc101941 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Fri, 31 Jul 2026 06:09:58 -0400 Subject: [PATCH 47/85] Gather the routing permutation inside the quantiser, and cache what is constant Attributing the prefill glue to torch ops rather than kernel names showed where it actually sits. Two of the entries were avoidable. The sorted activation was being materialised -- `x[stok]` is a full read and a full write of an (S, HID) matrix per layer, 14.5 ms of a 2048-token prefill -- and then handed to a quantiser that reads each of those rows exactly once. The quantiser now takes the permutation and gathers as it goes. The token index per slot and the slot index itself depend only on S, and every one of the forty layers was rebuilding both. TTFT: 1024 313.5 -> 308.8 ms, 2048 550.1 -> 532.0. Fixture 16/16. For the record, what the attribution actually says at 2048: the remaining glue is mostly the GDN chunked path's contiguous copies (38.6 ms across three shapes) and the router's torch.topk over 256 experts (12.9 ms), while the attention it replaced is now 8.4 ms of fused flash forward. --- csrc/bindings.cpp | 9 ++-- csrc/kernels/quantize.cu | 12 ++++- csrc/kernels/quantize.cuh | 2 +- .../frontends/torch/_nexn2_rtx_forward.py | 45 +++++++++++++------ 4 files changed, 49 insertions(+), 19 deletions(-) diff --git a/csrc/bindings.cpp b/csrc/bindings.cpp index 533ad40a..95407af5 100644 --- a/csrc/bindings.cpp +++ b/csrc/bindings.cpp @@ -1104,15 +1104,18 @@ PYBIND11_MODULE(flash_rt_kernels, m) { m.def("moe_grouped_quant_nvfp4_bf16", [](uintptr_t A, uintptr_t expert_of_row, uintptr_t group_off, - uintptr_t sfa_off, uintptr_t out_packed, uintptr_t out_sf, + uintptr_t sfa_off, uintptr_t src_row, + uintptr_t out_packed, uintptr_t out_sf, int slots, int K, uintptr_t stream) -> int { return moe_grouped_quant_nvfp4_bf16( to_ptr(A), to_ptr(expert_of_row), to_ptr(group_off), - to_ptr(sfa_off), to_ptr(out_packed), to_ptr(out_sf), + to_ptr(sfa_off), to_ptr(src_row), + to_ptr(out_packed), to_ptr(out_sf), slots, K, to_stream(stream)); }, py::arg("A"), py::arg("expert_of_row"), py::arg("group_off"), - py::arg("sfa_off"), py::arg("out_packed"), py::arg("out_sf"), + py::arg("sfa_off"), py::arg("src_row"), + py::arg("out_packed"), py::arg("out_sf"), py::arg("slots"), py::arg("K"), py::arg("stream") = 0); m.def("quantize_bf16_to_nvfp4_swizzled", [](uintptr_t input, uintptr_t fp4_data, diff --git a/csrc/kernels/quantize.cu b/csrc/kernels/quantize.cu index 7aaea5d7..9719941e 100644 --- a/csrc/kernels/quantize.cu +++ b/csrc/kernels/quantize.cu @@ -2822,6 +2822,7 @@ __global__ void moe_grouped_quant_nvfp4_kernel( const int* __restrict__ expert_of_row, const int* __restrict__ group_off, const int* __restrict__ sfa_off, + const long* __restrict__ src_row, uint8_t* __restrict__ fp4_data, uint8_t* __restrict__ scale_factors, int cols, int num_blocks, int n_col_blocks) @@ -2829,7 +2830,13 @@ __global__ void moe_grouped_quant_nvfp4_kernel( const int row = blockIdx.x; const int e = expert_of_row[row]; const int local = row - group_off[e]; // row index inside its group - const __nv_bfloat16* row_in = input + (size_t)row * cols; + // Gather while quantising when a permutation is given. Materialising the + // sorted activation first is a full read and a full write of an (S, HID) + // matrix per layer -- 14.5 ms of a 2048-token prefill -- for rows this + // kernel is about to read once anyway. + const size_t in_row = (src_row == nullptr) ? (size_t)row + : (size_t)src_row[row]; + const __nv_bfloat16* row_in = input + in_row * cols; uint8_t* row_fp4 = fp4_data + (size_t)row * cols / 2; uint8_t* sf_base = scale_factors + sfa_off[e]; @@ -2887,7 +2894,7 @@ __global__ void moe_grouped_quant_nvfp4_kernel( int moe_grouped_quant_nvfp4_bf16( const void* A, const void* expert_of_row, const void* group_off, - const void* sfa_off, void* out_packed, void* out_sf, + const void* sfa_off, const void* src_row, void* out_packed, void* out_sf, int slots, int K, cudaStream_t stream) { if (!A || !expert_of_row || !group_off || !sfa_off || !out_packed @@ -2902,6 +2909,7 @@ int moe_grouped_quant_nvfp4_bf16( reinterpret_cast(expert_of_row), reinterpret_cast(group_off), reinterpret_cast(sfa_off), + reinterpret_cast(src_row), reinterpret_cast(out_packed), reinterpret_cast(out_sf), K, num_blocks, n_col_blocks); diff --git a/csrc/kernels/quantize.cuh b/csrc/kernels/quantize.cuh index ccfc3feb..471987a5 100644 --- a/csrc/kernels/quantize.cuh +++ b/csrc/kernels/quantize.cuh @@ -341,7 +341,7 @@ void dequant_int32_to_bf16(const int32_t* input, __nv_bfloat16* output, // K must be a multiple of 16. Returns 0 on success, nonzero on arg error. int moe_grouped_quant_nvfp4_bf16( const void* A, const void* expert_of_row, const void* group_off, - const void* sfa_off, void* out_packed, void* out_sf, + const void* sfa_off, const void* src_row, void* out_packed, void* out_sf, int slots, int K, cudaStream_t stream); // Gate and quantise in one pass: reads the grouped GEMM's merged (slots, diff --git a/flash_rt/frontends/torch/_nexn2_rtx_forward.py b/flash_rt/frontends/torch/_nexn2_rtx_forward.py index 2edefbc0..cd9911f8 100644 --- a/flash_rt/frontends/torch/_nexn2_rtx_forward.py +++ b/flash_rt/frontends/torch/_nexn2_rtx_forward.py @@ -1029,6 +1029,26 @@ def _moe_experts_bt(x, ti, tw, ld, fvk, device): _GROUPED_SCRATCH = {} +_ROUTE_CONST = {} + + +def _route_constants(S, device): + """The parts of the routing permutation that depend only on the shape. + + Each layer routes differently, but the token index per slot and the slot + index itself do not change -- they are a function of S alone, and every + layer was rebuilding both. Forty layers of arange + repeat_interleave is + launches and traffic spent to recompute a constant. + """ + key = (S, str(device)) + got = _ROUTE_CONST.get(key) + if got is None: + tok_flat = torch.arange( + S, device=device).repeat_interleave(TOPK).contiguous() + slot_ix = torch.arange(S * TOPK, device=device) + got = (tok_flat, slot_ix) + _ROUTE_CONST[key] = got + return got def _grouped_scratch(fvk, device): @@ -1091,18 +1111,17 @@ def _moe_experts_grouped_gemm(x, ti, tw, ld, fvk, device): gu_a, dn_a = ld['experts_gate_up_alpha_dev'], ld['experts_down_alpha_dev'] scratch, scratch_bytes = _grouped_scratch(fvk, device) + tok_flat, slot_ix = _route_constants(S, device) exp_flat = ti.reshape(-1).to(torch.int32) - tok_flat = torch.arange(S, device=device).repeat_interleave(TOPK) order = exp_flat.argsort(stable=True) se = exp_flat[order].contiguous() stok = tok_flat[order] - sw = tw.reshape(-1)[order] counts = torch.bincount(se, minlength=_N_EXPERTS) group_off = torch.zeros(_N_EXPERTS + 1, dtype=torch.int32, device=device) group_off[1:] = counts.cumsum(0).to(torch.int32) - def project(A, k, n, w_p, w_s, alpha, out, gate=False): + def project(A, k, n, w_p, w_s, alpha, out, gate=False, perm=None): sfa_off, n_col = _sf_layout(counts, k, device) bound = (_N_EXPERTS + slots // 128 + 1) * n_col * 512 packed = torch.empty(slots, k // 2, dtype=torch.uint8, device=device) @@ -1118,8 +1137,8 @@ def project(A, k, n, w_p, w_s, alpha, out, gate=False): else: rc = fvk.moe_grouped_quant_nvfp4_bf16( A.data_ptr(), se.data_ptr(), group_off.data_ptr(), - sfa_off.data_ptr(), packed.data_ptr(), sfa.data_ptr(), - slots, k, _cs()) + sfa_off.data_ptr(), 0 if perm is None else perm.data_ptr(), + packed.data_ptr(), sfa.data_ptr(), slots, k, _cs()) if rc: raise RuntimeError(f'grouped activation quant failed with {rc}') rc = fvk.moe_grouped_gemm_nvfp4_sm100_bf16out( @@ -1132,7 +1151,7 @@ def project(A, k, n, w_p, w_s, alpha, out, gate=False): raise RuntimeError(f'grouped MoE GEMM failed with {rc}') d_gu = torch.empty(slots, n_gu, dtype=torch.bfloat16, device=device) - project(x[stok].contiguous(), HID, n_gu, gu_p, gu_s, gu_a, d_gu) + project(x, HID, n_gu, gu_p, gu_s, gu_a, d_gu, perm=stok) d_dn = torch.empty(slots, n_dn, dtype=torch.bfloat16, device=device) project(d_gu, INTER, n_dn, dn_p, dn_s, dn_a, d_dn, gate=True) @@ -1141,15 +1160,15 @@ def project(A, k, n, w_p, w_s, alpha, out, gate=False): # order. index_add_ was 37.8 ms of a 1024-token prefill and reduces through # atomics, so its order varies -- which prefill cannot afford, since it # seeds a decode that has to be reproducible. - inv = torch.empty(slots, dtype=torch.long, device=device) - inv[order] = torch.arange(slots, device=device) - rows = (torch.arange(slots, dtype=torch.int32, device=device)[inv] - ).contiguous() + # rows[i] is which sorted row holds slot i, which is exactly the inverse + # permutation -- gathering arange through it, as the tiled path has to, + # would just reproduce it. + inv = torch.empty(slots, dtype=torch.int32, device=device) + inv[order] = slot_ix.to(torch.int32) out = torch.empty(S, HID, dtype=torch.float32, device=device) fvk.moe_weighted_sum_sm120_bf16( - d_dn.data_ptr(), rows.data_ptr(), - tw.reshape(S, TOPK).contiguous().data_ptr(), out.data_ptr(), - S, TOPK, n_dn, n_dn, _cs()) + d_dn.data_ptr(), inv.data_ptr(), tw.contiguous().data_ptr(), + out.data_ptr(), S, TOPK, n_dn, n_dn, _cs()) return out From 76972835261c0b8696e67797e8fc1724d6289965 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Fri, 31 Jul 2026 06:30:44 -0400 Subject: [PATCH 48/85] Take the WY front matter off the host The chunked delta-rule scan reached its kernels through a chain of tensor ops: two l2 normalisations, a per-chunk gate cumulative sum, two GQA broadcasts and three chunk-major packings. At 2048 tokens that was 38.6 ms of copies in three shapes, and every one of them materialised a buffer the kernels then read back. The normalisation, the broadcast of q into its v-head slots, the packing of q and the gate cumulative sum now run as one kernel that reads q and k in the form the conv split already writes. The broadcast never lands in memory and q is normalised straight into its packed slots, so the only q traffic is the packed write. output_o moves to the raw-K entry, which expands k in kernel, so k never needs a 32-head buffer at all; v keeps a packing kernel because it is produced by the chunk_h stage. The sibling path has equivalent kernels but fixes the v-head count at compile time and sums the gate serially over the sequence, which suits a decode step and not a prefill, so these take the counts as arguments and parallelise the sum over chunks. TTFT 308.8 -> 289.8 ms at 1024 tokens, 532.0 -> 488.4 at 2048, 1095.1 -> 925.0 at 4096. Golden prefix 16/16. --- CMakeLists.txt | 1 + csrc/bindings.cpp | 26 +++ csrc/kernels/gdn_wy_prefill_edge.cu | 195 ++++++++++++++++++ csrc/kernels/gdn_wy_prefill_edge.cuh | 62 ++++++ .../frontends/torch/_nexn2_rtx_forward.py | 90 ++++---- 5 files changed, 322 insertions(+), 52 deletions(-) create mode 100644 csrc/kernels/gdn_wy_prefill_edge.cu create mode 100644 csrc/kernels/gdn_wy_prefill_edge.cuh diff --git a/CMakeLists.txt b/CMakeLists.txt index 6060fa1c..a0c24256 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1510,6 +1510,7 @@ if(FLASHRT_ENABLE_QWEN35MOE_CORE) csrc/kernels/qwen35moe_layout.cu csrc/kernels/bf16_matvec_sm120.cu csrc/kernels/gdn_recurrent_seq_sm120.cu + csrc/kernels/gdn_wy_prefill_edge.cu csrc/kernels/act_fuse_sm120.cu csrc/kernels/moe_router_topk_sm120.cu csrc/kernels/moe_weighted_sum_sm120.cu diff --git a/csrc/bindings.cpp b/csrc/bindings.cpp index 95407af5..e1a19421 100644 --- a/csrc/bindings.cpp +++ b/csrc/bindings.cpp @@ -183,6 +183,7 @@ extern "C" int cutlass_int8_rowwise_bf16out_t64x128( #include "kernels/qwen35moe_layout.cuh" #include "kernels/bf16_matvec_sm120.cuh" #include "kernels/gdn_recurrent_seq_sm120.cuh" +#include "kernels/gdn_wy_prefill_edge.cuh" #include "kernels/act_fuse_sm120.cuh" #include "kernels/moe_router_topk_sm120.cuh" #include "kernels/moe_weighted_sum_sm120.cuh" @@ -5528,6 +5529,31 @@ PYBIND11_MODULE(flash_rt_kernels, m) { py::arg("state"), py::arg("out"), py::arg("S"), py::arg("num_v_heads"), py::arg("head_dim"), py::arg("use_qk_l2norm") = true, py::arg("stream") = 0); + + m.def("gdn_wy_norm_pack_q_cumsum_edge_bf16", + [](uintptr_t q, uintptr_t k, uintptr_t g, uintptr_t k_l2, + uintptr_t q_pack, uintptr_t g_cumsum, int S, int num_k_heads, + int num_v_heads, int head_dim, int qk_group, uintptr_t stream) { + flash_rt::kernels::gdn_wy_norm_pack_q_cumsum_edge_bf16( + to_ptr(q), to_ptr(k), to_ptr(g), to_ptr(k_l2), + to_ptr(q_pack), to_ptr(g_cumsum), S, num_k_heads, + num_v_heads, head_dim, qk_group, to_stream(stream)); + }, + py::arg("q"), py::arg("k"), py::arg("g"), py::arg("k_l2"), + py::arg("q_pack"), py::arg("g_cumsum"), py::arg("S"), + py::arg("num_k_heads"), py::arg("num_v_heads"), py::arg("head_dim"), + py::arg("qk_group"), py::arg("stream") = 0); + + m.def("gdn_wy_pack_v_edge_bf16", + [](uintptr_t v, uintptr_t v_pack, int S, int num_v_heads, + int head_dim, uintptr_t stream) { + flash_rt::kernels::gdn_wy_pack_v_edge_bf16( + to_ptr(v), to_ptr(v_pack), S, num_v_heads, head_dim, + to_stream(stream)); + }, + py::arg("v"), py::arg("v_pack"), py::arg("S"), + py::arg("num_v_heads"), py::arg("head_dim"), + py::arg("stream") = 0); #endif // FLASHRT_HAVE_QWEN35MOE_CORE #ifdef FLASHRT_HAVE_QWEN35MOE_W4A16 diff --git a/csrc/kernels/gdn_wy_prefill_edge.cu b/csrc/kernels/gdn_wy_prefill_edge.cu new file mode 100644 index 00000000..315c6f70 --- /dev/null +++ b/csrc/kernels/gdn_wy_prefill_edge.cu @@ -0,0 +1,195 @@ +#include "gdn_wy_prefill_edge.cuh" + +#include + +namespace flash_rt { +namespace kernels { + +namespace { + +constexpr int kChunk = 64; +constexpr float kEps = 1e-6f; // matches the sequential scan's l2 eps + +// Butterfly order, the same summation order the sibling WY normalisation uses. +// Reduction order decides the low bits here, so this is not interchangeable +// with the shuffle-down helper in common.cuh. +template +__device__ __forceinline__ float wy_block_sum(float val, float* smem) { + for (int off = 16; off > 0; off >>= 1) { + val += __shfl_xor_sync(0xffffffff, val, off); + } + const int lane = threadIdx.x & 31; + const int warp = threadIdx.x >> 5; + if (lane == 0) smem[warp] = val; + __syncthreads(); + if (warp == 0) { + val = (lane < (kHD / 32)) ? smem[lane] : 0.0f; + for (int off = 16; off > 0; off >>= 1) { + val += __shfl_xor_sync(0xffffffff, val, off); + } + if (lane == 0) smem[0] = val; + } + __syncthreads(); + return smem[0]; +} + +// One block per (unique k-head, token). The block reduces both q and k over +// head_dim, writes the unique-head k, and scatters q into the qk_group v-head +// slots of the packed buffer -- so the GQA broadcast never materialises. +// +// The grid covers chunks * 64 tokens rather than S, so the threads past the +// end of the sequence are the ones that zero the packed tail. +template +__global__ void gdn_wy_norm_pack_q_kernel( + const __nv_bfloat16* __restrict__ q, + const __nv_bfloat16* __restrict__ k, + __nv_bfloat16* __restrict__ k_l2, + __nv_bfloat16* __restrict__ q_pack, + int S, + int num_k_heads, + int num_v_heads, + int qk_group) +{ + const int t = threadIdx.x; + const int h = blockIdx.x; // unique k-head + const int s = blockIdx.y; // token, may run past S into the pad + if (t >= kHD || h >= num_k_heads) return; + + const int chunk = s / kChunk; + const int tt = s - chunk * kChunk; + + if (s >= S) { + const __nv_bfloat16 zero = __float2bfloat16(0.0f); + for (int r = 0; r < qk_group; ++r) { + const int vh = h * qk_group + r; + q_pack[((static_cast(chunk) * num_v_heads + vh) * kChunk + tt) + * kHD + t] = zero; + } + return; + } + + // q and k arrive GQA-broadcast, so the group leader carries the value. + const size_t src = (static_cast(s) * num_v_heads + h * qk_group) + * kHD + t; + const float qv = static_cast(q[src]); + const float kv = static_cast(k[src]); + + __shared__ float scratch[32]; + const float q_sq = wy_block_sum(qv * qv, scratch); + __syncthreads(); // scratch is reused by the second reduction + const float k_sq = wy_block_sum(kv * kv, scratch); + __syncthreads(); + + const __nv_bfloat16 q_norm = __float2bfloat16(qv * rsqrtf(q_sq + kEps)); + const __nv_bfloat16 k_norm = __float2bfloat16(kv * rsqrtf(k_sq + kEps)); + + k_l2[(static_cast(s) * num_k_heads + h) * kHD + t] = k_norm; + + for (int r = 0; r < qk_group; ++r) { + const int vh = h * qk_group + r; + q_pack[((static_cast(chunk) * num_v_heads + vh) * kChunk + tt) + * kHD + t] = q_norm; + } +} + +// One thread per (chunk, v-head): 64 dependent adds, chunks * num_v_heads of +// them in flight. The sibling path runs one block of num_v_heads threads +// serially over the whole sequence, which is fine for a decode step and two +// orders of magnitude off for a prefill. +__global__ void gdn_wy_cumsum_g_chunk_kernel( + const __nv_bfloat16* __restrict__ g, + __nv_bfloat16* __restrict__ g_cumsum, + int S, + int num_v_heads, + int chunks) +{ + const int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= chunks * num_v_heads) return; + const int chunk = idx / num_v_heads; + const int vh = idx - chunk * num_v_heads; + + const int s0 = chunk * kChunk; + const int s1 = min(s0 + kChunk, S); + float acc = 0.0f; + for (int s = s0; s < s1; ++s) { + const size_t off = static_cast(s) * num_v_heads + vh; + acc += static_cast(g[off]); + g_cumsum[off] = __float2bfloat16(acc); + } +} + +__global__ void gdn_wy_pack_v_kernel( + const __nv_bfloat16* __restrict__ v, + __nv_bfloat16* __restrict__ v_pack, + int S, + int num_v_heads, + int head_dim) +{ + const int t = threadIdx.x; + const int vh = blockIdx.x; + const int s = blockIdx.y; + if (t >= head_dim || vh >= num_v_heads) return; + + const int chunk = s / kChunk; + const int tt = s - chunk * kChunk; + const size_t dst = + ((static_cast(chunk) * num_v_heads + vh) * kChunk + tt) + * head_dim + t; + v_pack[dst] = (s < S) + ? v[(static_cast(s) * num_v_heads + vh) * head_dim + t] + : __float2bfloat16(0.0f); +} + +} // namespace + +void gdn_wy_norm_pack_q_cumsum_edge_bf16( + const void* q, + const void* k, + const void* g, + void* k_l2, + void* q_pack, + void* g_cumsum, + int S, + int num_k_heads, + int num_v_heads, + int head_dim, + int qk_group, + cudaStream_t stream) +{ + if (S <= 0 || head_dim != 128) return; + const int chunks = (S + kChunk - 1) / kChunk; + + gdn_wy_norm_pack_q_kernel<128> + <<>>( + reinterpret_cast(q), + reinterpret_cast(k), + reinterpret_cast<__nv_bfloat16*>(k_l2), + reinterpret_cast<__nv_bfloat16*>(q_pack), + S, num_k_heads, num_v_heads, qk_group); + + const int total = chunks * num_v_heads; + gdn_wy_cumsum_g_chunk_kernel<<<(total + 127) / 128, 128, 0, stream>>>( + reinterpret_cast(g), + reinterpret_cast<__nv_bfloat16*>(g_cumsum), + S, num_v_heads, chunks); +} + +void gdn_wy_pack_v_edge_bf16( + const void* v, + void* v_pack, + int S, + int num_v_heads, + int head_dim, + cudaStream_t stream) +{ + if (S <= 0) return; + const int chunks = (S + kChunk - 1) / kChunk; + gdn_wy_pack_v_kernel<<>>( + reinterpret_cast(v), + reinterpret_cast<__nv_bfloat16*>(v_pack), + S, num_v_heads, head_dim); +} + +} // namespace kernels +} // namespace flash_rt diff --git a/csrc/kernels/gdn_wy_prefill_edge.cuh b/csrc/kernels/gdn_wy_prefill_edge.cuh new file mode 100644 index 00000000..febb2106 --- /dev/null +++ b/csrc/kernels/gdn_wy_prefill_edge.cuh @@ -0,0 +1,62 @@ +#pragma once + +#include + +namespace flash_rt { +namespace kernels { + +// WY chunked delta-rule front matter for a batched prefill, with the head +// counts as runtime arguments. +// +// The sibling 27B path has equivalent kernels, but its v-head count is a +// compile-time constant and its gate cumulative sum runs one block of +// `num_v_heads` threads serially over S -- a decode shape. These take the +// counts as arguments and parallelise the cumulative sum over chunks, which is +// what a prefill of a few thousand tokens needs. +// +// Layout conventions match the mma WY kernels: +// packed: (chunks, num_v_heads, 64, head_dim), chunks = ceil(S / 64), +// pack[c, h, i, d] = x[c * 64 + i, h, d], zero past S. +// g_cumsum: (S, num_v_heads), cumulative within each 64-token chunk. + +// Fuses the q/k l2 normalisation, the GQA broadcast of q into v-head slots, +// the chunk-major packing of q, and the gate cumulative sum. +// +// `q` and `k` are read as (S, num_v_heads, head_dim) already broadcast across +// the GQA group -- the form the conv split kernel writes -- and only the group +// leaders are touched, so no strided host-side slice is needed. +// +// q, k (S, num_v_heads, head_dim) bf16, GQA-broadcast +// g (S, num_v_heads) bf16 +// k_l2 (S, num_k_heads, head_dim) bf16 out, unique heads only +// q_pack (chunks, num_v_heads, 64, head_dim) bf16 out +// g_cumsum (S, num_v_heads) bf16 out +// +// head_dim must be 128. qk_group = num_v_heads / num_k_heads. +void gdn_wy_norm_pack_q_cumsum_edge_bf16( + const void* q, + const void* k, + const void* g, + void* k_l2, + void* q_pack, + void* g_cumsum, + int S, + int num_k_heads, + int num_v_heads, + int head_dim, + int qk_group, + cudaStream_t stream); + +// Chunk-major packing of the un-decayed v the chunk_h stage produces. +// v (S, num_v_heads, head_dim) bf16 +// v_pack (chunks, num_v_heads, 64, head_dim) bf16 out, zero past S +void gdn_wy_pack_v_edge_bf16( + const void* v, + void* v_pack, + int S, + int num_v_heads, + int head_dim, + cudaStream_t stream); + +} // namespace kernels +} // namespace flash_rt diff --git a/flash_rt/frontends/torch/_nexn2_rtx_forward.py b/flash_rt/frontends/torch/_nexn2_rtx_forward.py index cd9911f8..69fee2ac 100644 --- a/flash_rt/frontends/torch/_nexn2_rtx_forward.py +++ b/flash_rt/frontends/torch/_nexn2_rtx_forward.py @@ -326,58 +326,41 @@ def _silu_mul(g, u, fvk, device): _WY_MIN_S = 64 # below this the seq-scan's lower fixed overhead wins -def _wy_pack_t(x, ch=64): - """(S, H, D) -> (chunks, H, ch, D): x_pack[ci, h, i, d] = x[ci*ch+i, h, d] - (zero-padded last chunk). The packed chunk-major layout the mma kernels read.""" - s, hh, d = x.shape - pad = (-s) % ch - if pad: - x = F.pad(x, (0, 0, 0, 0, 0, pad)) - return x.reshape(-1, ch, hh, d).permute(0, 2, 1, 3).contiguous() - - -def _wy_l2(x): - """l2norm over the last dim, eps inside rsqrt (matches the seq-scan kEps).""" - xf = x.float() - return (xf * torch.rsqrt((xf * xf).sum(-1, keepdim=True) + 1e-6)).to( - torch.bfloat16) - - -def _wy_gcumsum(g, ch=64): - """(S, NV) -> (S, NV) per-chunk cumulative sum of the (log-space) gate.""" - s = g.shape[0] - pad = (-s) % ch - gp = F.pad(g, (0, 0, 0, pad)) if pad else g - return torch.cumsum(gp.float().reshape(-1, ch, g.shape[1]), 1).reshape( - -1, g.shape[1])[:s].to(torch.bfloat16) - - -def _gdn_wy_chunk(q16, k16, v, g, beta, fvk, device, init_state=None): - """WY chunked scan. q16/k16 (S,16,128) raw post-conv, v (S,32,128), - g/beta (S,32). Returns core (S,32,128) + final state (32,128,128). +def _gdn_wy_chunk(qb, kb, v, g, beta, fvk, device, init_state=None): + """WY chunked scan. qb/kb (S,32,128) raw post-conv with q/k already + GQA-broadcast across the 32 v-head slots, v (S,32,128), g/beta (S,32). + Returns core (S,32,128) + final state (32,128,128). ``init_state`` (NV,HK,HV) is the recurrent state to continue from -- the chunk_h kernel reads it as h0[0] and writes the post-block state back, so a chunked prefill carries it across blocks (probe-verified bit-exact: whole vs two state-carried halves match at cos 1.0). Defaults to zeros. - Pipeline (FLA chunked delta rule, all add-only existing kernels): - l2norm + per-chunk g-cumsum (torch glue) -> kkt -> solve_tril(+pack) -> - recompute_wu -> chunk_h (inter-chunk state) -> output_o.""" - S = q16.shape[0] + Pipeline (FLA chunked delta rule, kernels throughout): norm+pack_q+cumsum + -> kkt -> solve_tril(+pack) -> recompute_wu -> chunk_h (inter-chunk state) + -> pack_v -> output_o. + """ + S = qb.shape[0] chunks = (S + 63) // 64 CH, QKG = 64, NV // NK - q_l2 = _wy_l2(q16) - k_l2 = _wy_l2(k16).contiguous() - gc = _wy_gcumsum(g).contiguous() - betac = beta.contiguous() - vc = v.contiguous() + + # l2norm of q and k, the GQA broadcast of q into the 32 v-head slots, the + # chunk-major packing of q, and the per-chunk gate cumulative sum, in one + # kernel. The broadcast never materialises and q is normalised straight + # into its packed slots, so the only q traffic is the packed write. + k_l2 = torch.empty(S, NK, HK, dtype=torch.bfloat16, device=device) + q_pack = torch.empty(chunks, NV, CH, HK, dtype=torch.bfloat16, + device=device) + gc = torch.empty(S, NV, dtype=torch.bfloat16, device=device) + fvk.gdn_wy_norm_pack_q_cumsum_edge_bf16( + qb.data_ptr(), kb.data_ptr(), g.data_ptr(), k_l2.data_ptr(), + q_pack.data_ptr(), gc.data_ptr(), S, NK, NV, HK, QKG, _cs()) k_pack = torch.empty(chunks, NK, CH, HK, dtype=torch.bfloat16, device=device) kkt_base = torch.empty(chunks, NK, CH, CH, dtype=torch.float32, device=device) A = torch.empty(chunks, NV, CH, CH, dtype=torch.float32, device=device) fvk.linear_attn_gdn_wy_kkt_b64_bf16_cublaslt( - k_l2.data_ptr(), betac.data_ptr(), gc.data_ptr(), k_pack.data_ptr(), + k_l2.data_ptr(), beta.data_ptr(), gc.data_ptr(), k_pack.data_ptr(), kkt_base.data_ptr(), A.data_ptr(), S, NK, NV, HK, QKG, _cs()) Ai = torch.empty(chunks, NV, CH, CH, dtype=torch.float32, device=device) @@ -388,7 +371,7 @@ def _gdn_wy_chunk(q16, k16, v, g, beta, fvk, device, init_state=None): w_pack = torch.empty(chunks, NV, CH, HV, dtype=torch.bfloat16, device=device) u_pack = torch.empty(chunks, NV, CH, HV, dtype=torch.bfloat16, device=device) fvk.linear_attn_gdn_wy_recompute_wu_b64_bf16_mma_fla( - k_l2.data_ptr(), vc.data_ptr(), betac.data_ptr(), gc.data_ptr(), + k_l2.data_ptr(), v.data_ptr(), beta.data_ptr(), gc.data_ptr(), Ai_pack.data_ptr(), w_pack.data_ptr(), u_pack.data_ptr(), S, NK, NV, HK, QKG, _cs()) @@ -401,14 +384,17 @@ def _gdn_wy_chunk(q16, k16, v, g, beta, fvk, device, init_state=None): state.data_ptr(), h0.data_ptr(), v_new.data_ptr(), 0, 0, S, NK, NV, HK, QKG, _cs()) - q_pack = _wy_pack_t(q_l2.repeat_interleave(QKG, 1)) - k_pack_hv = _wy_pack_t(k_l2.repeat_interleave(QKG, 1)) - v_pack = _wy_pack_t(v_new) + # v is the only side still needing a packed copy; the raw-K output_o does + # the GQA expansion of k in-kernel, so k never gets a 32-head buffer. + v_pack = torch.empty(chunks, NV, CH, HV, dtype=torch.bfloat16, + device=device) + fvk.gdn_wy_pack_v_edge_bf16( + v_new.data_ptr(), v_pack.data_ptr(), S, NV, HV, _cs()) core = torch.empty(S, NV, HV, dtype=torch.bfloat16, device=device) - fvk.linear_attn_gdn_wy_output_o_b64_bf16_mma_fla( - q_pack.data_ptr(), k_pack_hv.data_ptr(), v_pack.data_ptr(), + fvk.linear_attn_gdn_wy_output_o_b64_bf16_mma_fla_rawk( + q_pack.data_ptr(), k_l2.data_ptr(), v_pack.data_ptr(), h0.data_ptr(), gc.data_ptr(), core.data_ptr(), - S, NV, HV, float(HV ** -0.5), _cs()) + S, NK, NV, HV, QKG, float(HV ** -0.5), _cs()) return core, state @@ -491,12 +477,12 @@ def _gdn_layer(h, ld, fvk, device, eps, cap=None, rank=None, if _USE_WY_GDN and S >= _WY_MIN_S: # WY chunked delta-rule scan: 11x faster than the seq-scan at S=2048, - # bit-exact. qb/kb are the 16->32 broadcast heads (src_h = h//2), so the - # 16 unique K-heads are the even slots; the WY kernels re-expand by GQA. - q16 = qb.reshape(S, NV, HK)[:, 0::2, :] - k16 = kb.reshape(S, NV, HK)[:, 0::2, :] + # bit-exact. qb/kb carry the 16->32 broadcast heads (src_h = h//2); the + # front kernel reads the group leaders and re-expands where it packs, + # so no strided slice is taken here. core, state = _gdn_wy_chunk( - q16, k16, vb.reshape(S, NV, HV), g_out.reshape(S, NV), + qb.reshape(S, NV, HK), kb.reshape(S, NV, HK), + vb.reshape(S, NV, HV), g_out.reshape(S, NV), bo.reshape(S, NV), fvk, device, init_state=init_state) core = core.reshape(B, S, NV, HV) else: @@ -725,7 +711,7 @@ def _fa2_causal_attn(qf, kf, vf, device, *, _probe=False): batch=1, seqlen_q=Sq, seqlen_k=Sk, num_heads_q=NQ, num_heads_kv=NKV, head_dim=HD, q_strides=qc.stride()[:3], k_strides=kc.stride()[:3], v_strides=vc.stride()[:3], o_strides=o.stride()[:3], - softmax_scale=float(HD) ** -0.5, num_sms=_num_sms(), stream=0) + softmax_scale=float(HD) ** -0.5, num_sms=_num_sms(), stream=_cs()) return o From faf08a9145066f79de1a49b807a836a562a03210 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Fri, 31 Jul 2026 07:08:32 -0400 Subject: [PATCH 49/85] Route a prefill's experts in kernels, not tensor ops Each MoE layer reached its grouped GEMM through a softmax, a top-k, a renormalising divide, a stable argsort, two gathers, a bincount, a cumulative sum and a scatter. Ten tensor ops, forty layers; the top-k alone was 25 ms of a 2048-token prefill. One entry point now produces everything the grouped GEMM reads. The top-k holds a token's whole logit row across one warp and selects in registers -- a block-wide version cost eight barriers per round, sixty-four per token, and ran seventy times off the bandwidth the row needs. The permutation is a counting sort with per-block offsets rather than an atomic scatter, because prefill seeds a decode that has to reproduce: slot order within an expert is fixed, not left to the order blocks happen to arrive in. Against the chain it replaces, on the logits the model actually produces: selected expert sets identical for all 10240 token-layers of a 256-token prefill, weights to 4e-7, scale-factor offsets and the sorted layout bit for bit. Substituting the tensor chain's rank order back into the kernel path reproduces its logits exactly, which is what places the whole of the remaining difference in that order and nowhere else. That order is where the two part. Exact ties at the top-k boundary are common in bf16 -- around 90 tokens a layer -- and neither implementation defines which of two equal experts it calls rank six. The selection agrees; the listing order does not, and since routing is discrete a bit of difference in one layer flips a tie in the next. Generated text is equivalent either way. Set NEXN2_ROUTE_KERNEL=0 to put the routing back on the tensor chain, which is how the two were compared end to end. Routing 0.536 -> 0.055 ms at 2048 tokens. TTFT 289.8 -> 267.9 ms at 1024, 488.4 -> 459.3 at 2048, 925.0 -> 885.5 at 4096. Golden prefix 16/16. --- CMakeLists.txt | 1 + csrc/bindings.cpp | 33 ++ csrc/kernels/moe_route_prefill_edge.cu | 322 ++++++++++++++++++ csrc/kernels/moe_route_prefill_edge.cuh | 66 ++++ .../frontends/torch/_nexn2_rtx_forward.py | 142 ++++++-- 5 files changed, 543 insertions(+), 21 deletions(-) create mode 100644 csrc/kernels/moe_route_prefill_edge.cu create mode 100644 csrc/kernels/moe_route_prefill_edge.cuh diff --git a/CMakeLists.txt b/CMakeLists.txt index a0c24256..f20feb95 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1513,6 +1513,7 @@ if(FLASHRT_ENABLE_QWEN35MOE_CORE) csrc/kernels/gdn_wy_prefill_edge.cu csrc/kernels/act_fuse_sm120.cu csrc/kernels/moe_router_topk_sm120.cu + csrc/kernels/moe_route_prefill_edge.cu csrc/kernels/moe_weighted_sum_sm120.cu csrc/kernels/w16a16_gemm_sm120.cu csrc/kernels/qwen35moe_e0m3_dequant.cu) diff --git a/csrc/bindings.cpp b/csrc/bindings.cpp index e1a19421..12bc6201 100644 --- a/csrc/bindings.cpp +++ b/csrc/bindings.cpp @@ -186,6 +186,7 @@ extern "C" int cutlass_int8_rowwise_bf16out_t64x128( #include "kernels/gdn_wy_prefill_edge.cuh" #include "kernels/act_fuse_sm120.cuh" #include "kernels/moe_router_topk_sm120.cuh" +#include "kernels/moe_route_prefill_edge.cuh" #include "kernels/moe_weighted_sum_sm120.cuh" #include "kernels/w16a16_gemm_sm120.cuh" #include "kernels/qwen35moe_e0m3_dequant.cuh" @@ -5530,6 +5531,38 @@ PYBIND11_MODULE(flash_rt_kernels, m) { py::arg("head_dim"), py::arg("use_qk_l2norm") = true, py::arg("stream") = 0); + m.def("moe_route_prefill_bf16", + [](uintptr_t logits, uintptr_t ti, uintptr_t tw, uintptr_t se, + uintptr_t stok, uintptr_t inv, uintptr_t group_off, uintptr_t ws, + int ws_bytes, int S, int n_experts, int topk, + uintptr_t stream) -> int { + return flash_rt::kernels::moe_route_prefill_bf16( + to_ptr(logits), to_ptr(ti), to_ptr(tw), to_ptr(se), + to_ptr(stok), to_ptr(inv), to_ptr(group_off), to_ptr(ws), + ws_bytes, S, n_experts, topk, to_stream(stream)); + }, + py::arg("logits"), py::arg("ti"), py::arg("tw"), py::arg("se"), + py::arg("stok"), py::arg("inv"), py::arg("group_off"), py::arg("ws"), + py::arg("ws_bytes"), py::arg("S"), py::arg("n_experts"), + py::arg("topk"), py::arg("stream") = 0); + + m.def("moe_route_prefill_workspace_bytes", + [](int S, int topk, int n_experts) -> int { + return flash_rt::kernels::moe_route_prefill_workspace_bytes( + S, topk, n_experts); + }, + py::arg("S"), py::arg("topk"), py::arg("n_experts")); + + m.def("moe_route_sfa_offsets", + [](uintptr_t group_off, uintptr_t sfa_off, int n_experts, int n_col, + uintptr_t stream) { + flash_rt::kernels::moe_route_sfa_offsets( + to_ptr(group_off), to_ptr(sfa_off), n_experts, n_col, + to_stream(stream)); + }, + py::arg("group_off"), py::arg("sfa_off"), py::arg("n_experts"), + py::arg("n_col"), py::arg("stream") = 0); + m.def("gdn_wy_norm_pack_q_cumsum_edge_bf16", [](uintptr_t q, uintptr_t k, uintptr_t g, uintptr_t k_l2, uintptr_t q_pack, uintptr_t g_cumsum, int S, int num_k_heads, diff --git a/csrc/kernels/moe_route_prefill_edge.cu b/csrc/kernels/moe_route_prefill_edge.cu new file mode 100644 index 00000000..def39bd0 --- /dev/null +++ b/csrc/kernels/moe_route_prefill_edge.cu @@ -0,0 +1,322 @@ +#include "moe_route_prefill_edge.cuh" + +#include +#include + +namespace flash_rt { +namespace kernels { + +namespace { + +constexpr int kMaxExperts = 1024; +constexpr int kMaxTopK = 32; +constexpr int kSlotsPerBlock = 256; // slots a scatter/histogram block owns +constexpr int kRouteThreads = 256; + +// One warp per token, the whole row held in registers: PER_LANE experts per +// lane, strided so the row loads coalesced. A block-wide version of this cost +// eight barriers per top-k round -- sixty-four per token -- and ran seventy +// times off the bandwidth the row needs; there is no barrier here at all. +// +// Softmax first, then the top-k renormalised over itself. The full denominator +// cancels between the two, but it is kept because it only cancels exactly in +// exact arithmetic and this seeds a decode that has to reproduce. +// +// Ties go to the lower expert index. bf16 logits make the tail probabilities +// tie outright often enough to matter (6% of slots at 256 experts), and the +// tensor top-k this replaces does not define which of two equal experts it +// ranks first -- so the rank order inside a token's top-k can differ from it +// while the selected set, which is what the grouped GEMM reads, does not. +template +__global__ void route_topk_warp_kernel( + const __nv_bfloat16* __restrict__ logits, + int* __restrict__ ti, + float* __restrict__ tw, + int S, + int n_experts, + int topk) +{ + const int lane = threadIdx.x & 31; + const int warp = threadIdx.x >> 5; + const int s = blockIdx.x * (blockDim.x >> 5) + warp; + if (s >= S) return; // whole warp, so the shuffles stay put + + const __nv_bfloat16* row = logits + static_cast(s) * n_experts; + float v[PER_LANE]; + #pragma unroll + for (int i = 0; i < PER_LANE; ++i) v[i] = static_cast(row[i * 32 + lane]); + + float m = -CUDART_INF_F; + #pragma unroll + for (int i = 0; i < PER_LANE; ++i) m = fmaxf(m, v[i]); + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + m = fmaxf(m, __shfl_xor_sync(0xffffffff, m, off)); + + float sum = 0.0f; + #pragma unroll + for (int i = 0; i < PER_LANE; ++i) { v[i] = __expf(v[i] - m); sum += v[i]; } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + sum += __shfl_xor_sync(0xffffffff, sum, off); + #pragma unroll + for (int i = 0; i < PER_LANE; ++i) v[i] /= sum; + + // Lane r keeps rank r, so the results end up spread one per lane and the + // write below is a single coalesced store. + float my_val = 0.0f; + int my_idx = 0; + for (int r = 0; r < topk; ++r) { + float best = -CUDART_INF_F; + int best_i = n_experts; + #pragma unroll + for (int i = 0; i < PER_LANE; ++i) { + const int e = i * 32 + lane; + if (v[i] > best || (v[i] == best && e < best_i)) { best = v[i]; best_i = e; } + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) { + const float ov = __shfl_xor_sync(0xffffffff, best, off); + const int oi = __shfl_xor_sync(0xffffffff, best_i, off); + if (ov > best || (ov == best && oi < best_i)) { best = ov; best_i = oi; } + } + if (lane == r) { my_val = best; my_idx = best_i; } + // Compile-time indices: a computed one would push v[] into local memory. + #pragma unroll + for (int i = 0; i < PER_LANE; ++i) + if (i * 32 + lane == best_i) v[i] = -CUDART_INF_F; + } + + float tsum = (lane < topk) ? my_val : 0.0f; + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + tsum += __shfl_xor_sync(0xffffffff, tsum, off); + + if (lane < topk) { + const size_t base = static_cast(s) * topk; + ti[base + lane] = my_idx; + tw[base + lane] = my_val / tsum; + } +} + +// Per-block expert histogram over a fixed slice of slots. The warp-level match +// gives each lane its rank among the lanes of its warp holding the same +// expert, which is what makes the scatter's ordering fall out without atomics. +__global__ void slot_hist_kernel( + const int* __restrict__ ti, + int* __restrict__ blk_hist, + int slots, + int n_experts) +{ + const int blk = blockIdx.x; + const int t = threadIdx.x; + const int slot = blk * kSlotsPerBlock + t; + + extern __shared__ int s_hist[]; // n_experts + for (int e = t; e < n_experts; e += blockDim.x) s_hist[e] = 0; + __syncthreads(); + + if (slot < slots) atomicAdd(&s_hist[ti[slot]], 1); + __syncthreads(); + + int* out = blk_hist + static_cast(blk) * n_experts; + for (int e = t; e < n_experts; e += blockDim.x) out[e] = s_hist[e]; +} + +// One block per expert: exclusive scan of that expert's per-block counts, so a +// scatter block knows where its own slots for that expert begin. +__global__ void expert_block_scan_kernel( + const int* __restrict__ blk_hist, + int* __restrict__ blk_off, + int* __restrict__ counts, + int n_blocks, + int n_experts) +{ + const int e = blockIdx.x; + if (threadIdx.x != 0) return; + int acc = 0; + for (int b = 0; b < n_blocks; ++b) { + const size_t off = static_cast(b) * n_experts + e; + blk_off[off] = acc; + acc += blk_hist[off]; + } + counts[e] = acc; +} + +__global__ void group_off_kernel( + const int* __restrict__ counts, + int* __restrict__ group_off, + int n_experts) +{ + if (threadIdx.x != 0) return; + int acc = 0; + for (int e = 0; e < n_experts; ++e) { + group_off[e] = acc; + acc += counts[e]; + } + group_off[n_experts] = acc; +} + +// Places every slot at group_off[e] + blk_off[blk][e] + its rank within the +// block. Rank comes from the warp match plus the counts of the earlier warps, +// so two runs on the same routing place the same slot in the same row. +__global__ void slot_scatter_kernel( + const int* __restrict__ ti, + const int* __restrict__ group_off, + const int* __restrict__ blk_off, + int* __restrict__ se, + long* __restrict__ stok, + int* __restrict__ inv, + int slots, + int n_experts, + int topk) +{ + const int blk = blockIdx.x; + const int t = threadIdx.x; + const int slot = blk * kSlotsPerBlock + t; + const int warp = t >> 5; + const int lane = t & 31; + const int n_warps = blockDim.x >> 5; + + extern __shared__ int s_warp_hist[]; // n_warps * n_experts + for (int i = t; i < n_warps * n_experts; i += blockDim.x) s_warp_hist[i] = 0; + __syncthreads(); + + const int e = (slot < slots) ? ti[slot] : -1; + const unsigned active = __ballot_sync(0xffffffff, e >= 0); + int rank_in_warp = 0; + if (e >= 0) { + const unsigned same = __match_any_sync(active, e); + const unsigned lower = same & ((1u << lane) - 1u); + rank_in_warp = __popc(lower); + if (rank_in_warp == 0) { + s_warp_hist[warp * n_experts + e] = __popc(same); + } + } + __syncthreads(); + + if (e >= 0) { + int before = 0; + for (int w = 0; w < warp; ++w) before += s_warp_hist[w * n_experts + e]; + const int row = group_off[e] + + blk_off[static_cast(blk) * n_experts + e] + + before + rank_in_warp; + se[row] = e; + stok[row] = slot / topk; // 64-bit: see the note in the header + inv[slot] = row; + } +} + +__global__ void sfa_offsets_kernel( + const int* __restrict__ group_off, + int* __restrict__ sfa_off, + int n_experts, + int n_col) +{ + if (threadIdx.x != 0) return; + int acc = 0; + for (int e = 0; e < n_experts; ++e) { + sfa_off[e] = acc; + const int c = group_off[e + 1] - group_off[e]; + acc += ((c + 127) / 128) * (n_col * 512); + } +} + +int route_blocks(int slots) { + return (slots + kSlotsPerBlock - 1) / kSlotsPerBlock; +} + +} // namespace + +int moe_route_prefill_workspace_bytes(int S, int topk, int n_experts) +{ + const int slots = S * topk; + const int nblk = route_blocks(slots); + // blk_hist + blk_off + counts + return static_cast( + (2 * static_cast(nblk) * n_experts + n_experts) * sizeof(int)); +} + +int moe_route_prefill_bf16( + const void* logits, + void* ti, + void* tw, + void* se, + void* stok, + void* inv, + void* group_off, + void* ws, + int ws_bytes, + int S, + int n_experts, + int topk, + cudaStream_t stream) +{ + if (S <= 0) return 0; + if (n_experts <= 0 || n_experts > kMaxExperts || (n_experts % 32) != 0) + return 1; + if (topk <= 0 || topk > kMaxTopK || topk > n_experts) return 2; + if (ws_bytes < moe_route_prefill_workspace_bytes(S, topk, n_experts)) + return 3; + + const int slots = S * topk; + const int nblk = route_blocks(slots); + int* blk_hist = reinterpret_cast(ws); + int* blk_off = blk_hist + static_cast(nblk) * n_experts; + int* counts = blk_off + static_cast(nblk) * n_experts; + int* ti_i = reinterpret_cast(ti); + + // One warp per token, four warps a block. + const int warps = kRouteThreads / 32; + const int topk_grid = (S + warps - 1) / warps; + float* tw_f = reinterpret_cast(tw); + const __nv_bfloat16* lg = reinterpret_cast(logits); + switch (n_experts / 32) { + case 1: route_topk_warp_kernel<1><<>>( + lg, ti_i, tw_f, S, n_experts, topk); break; + case 2: route_topk_warp_kernel<2><<>>( + lg, ti_i, tw_f, S, n_experts, topk); break; + case 4: route_topk_warp_kernel<4><<>>( + lg, ti_i, tw_f, S, n_experts, topk); break; + case 8: route_topk_warp_kernel<8><<>>( + lg, ti_i, tw_f, S, n_experts, topk); break; + case 16: route_topk_warp_kernel<16><<>>( + lg, ti_i, tw_f, S, n_experts, topk); break; + case 32: route_topk_warp_kernel<32><<>>( + lg, ti_i, tw_f, S, n_experts, topk); break; + default: return 4; // expert count is not 32 * a power of two + } + + slot_hist_kernel<<>>( + ti_i, blk_hist, slots, n_experts); + + expert_block_scan_kernel<<>>( + blk_hist, blk_off, counts, nblk, n_experts); + + group_off_kernel<<<1, 32, 0, stream>>>( + counts, reinterpret_cast(group_off), n_experts); + + const size_t scatter_smem = + static_cast(kSlotsPerBlock / 32) * n_experts * sizeof(int); + slot_scatter_kernel<<>>( + ti_i, reinterpret_cast(group_off), blk_off, + reinterpret_cast(se), reinterpret_cast(stok), + reinterpret_cast(inv), slots, n_experts, topk); + + return 0; +} + +void moe_route_sfa_offsets( + const void* group_off, + void* sfa_off, + int n_experts, + int n_col, + cudaStream_t stream) +{ + sfa_offsets_kernel<<<1, 32, 0, stream>>>( + reinterpret_cast(group_off), + reinterpret_cast(sfa_off), n_experts, n_col); +} + +} // namespace kernels +} // namespace flash_rt diff --git a/csrc/kernels/moe_route_prefill_edge.cuh b/csrc/kernels/moe_route_prefill_edge.cuh new file mode 100644 index 00000000..c8536b4f --- /dev/null +++ b/csrc/kernels/moe_route_prefill_edge.cuh @@ -0,0 +1,66 @@ +#pragma once + +#include + +namespace flash_rt { +namespace kernels { + +// Everything a grouped MoE prefill needs from its router logits, as kernels. +// +// The chain this replaces was softmax, top-k, a renormalising divide, a stable +// argsort, two gathers, a bincount, a cumulative sum and a scatter -- ten +// tensor ops per layer, of which the top-k alone cost 25 ms of a 2048-token +// prefill. +// +// The permutation is built as a counting sort with per-block offsets, not an +// atomic scatter: prefill seeds a decode that has to reproduce, so slot order +// within an expert is fixed (ascending slot index, matching a stable argsort) +// rather than left to the order blocks happen to arrive in. +// +// logits (S, n_experts) bf16 +// ti (S, topk) int32 out, expert per (token, rank) +// tw (S, topk) fp32 out, weights renormalised over the top-k +// se (S * topk,) int32 out, expert per sorted slot +// stok (S * topk,) int64 out, token per sorted slot -- 64-bit, alone +// among these, because it is handed to the grouped activation +// quantiser as its gather index and that kernel reads a long. +// Emitting int32 here reads as garbage row indices there, which +// surfaces as an illegal access three kernels later. +// inv (S * topk,) int32 out, sorted row holding slot i +// group_off (n_experts + 1,) int32 out, prefix sums over experts +// ws workspace, moe_route_prefill_workspace_bytes(S, topk, n_experts) +// +// n_experts must be 32 times a power of two, at most 1024, since the top-k +// holds a row across one warp; topk at most 32. Returns 0 on success. +int moe_route_prefill_bf16( + const void* logits, + void* ti, + void* tw, + void* se, + void* stok, + void* inv, + void* group_off, + void* ws, + int ws_bytes, + int S, + int n_experts, + int topk, + cudaStream_t stream); + +int moe_route_prefill_workspace_bytes(int S, int topk, int n_experts); + +// Per-expert scale-factor byte offsets for the block-scaled activation layout, +// derived from the group boundaries the routing kernel already produced. The +// layout blocks rows by 128, so a group of c rows takes ceil(c / 128) super +// blocks of n_col * 512 bytes. +// group_off (n_experts + 1,) int32 +// sfa_off (n_experts,) int32 out +void moe_route_sfa_offsets( + const void* group_off, + void* sfa_off, + int n_experts, + int n_col, + cudaStream_t stream); + +} // namespace kernels +} // namespace flash_rt diff --git a/flash_rt/frontends/torch/_nexn2_rtx_forward.py b/flash_rt/frontends/torch/_nexn2_rtx_forward.py index 69fee2ac..5905ed56 100644 --- a/flash_rt/frontends/torch/_nexn2_rtx_forward.py +++ b/flash_rt/frontends/torch/_nexn2_rtx_forward.py @@ -1016,6 +1016,10 @@ def _moe_experts_bt(x, ti, tw, ld, fvk, device): _GROUPED_SCRATCH = {} _ROUTE_CONST = {} +_ROUTE_BUF = {} +# Off puts the routing back on the tensor chain, which is how the kernel's +# output is A/B'd against it end to end rather than only in a probe. +_USE_ROUTE_KERNEL = _os.environ.get('NEXN2_ROUTE_KERNEL', '1') != '0' def _route_constants(S, device): @@ -1037,6 +1041,75 @@ def _route_constants(S, device): return got +def _route_buffers(S, fvk, device): + """The routing kernel's outputs, allocated once per prompt length. + + Their sizes depend only on S, so the forty layers of a prefill write + through the same buffers at the same addresses -- which is what lets the + call sit inside a captured region, and incidentally saves forty rounds of + allocation per forward. + """ + key = (S, str(device)) + got = _ROUTE_BUF.get(key) + if got is None: + slots = S * TOPK + ws_bytes = int(fvk.moe_route_prefill_workspace_bytes( + S, TOPK, _N_EXPERTS)) + got = { + 'ti': torch.empty(S, TOPK, dtype=torch.int32, device=device), + 'tw': torch.empty(S, TOPK, dtype=torch.float32, device=device), + 'se': torch.empty(slots, dtype=torch.int32, device=device), + # int64: the activation quantiser reads this gather index + # as a long, and int32 there is an illegal access. + 'stok': torch.empty(slots, dtype=torch.int64, device=device), + 'inv': torch.empty(slots, dtype=torch.int32, device=device), + 'group_off': torch.empty(_N_EXPERTS + 1, dtype=torch.int32, + device=device), + 'ws': torch.empty(ws_bytes, dtype=torch.uint8, device=device), + 'ws_bytes': ws_bytes, + 'sfa_off': {}, + } + _ROUTE_BUF[key] = got + return got + + +def _route_prefill(logits, fvk, device): + """Softmax, top-k, and the permutation the grouped GEMM reads, in kernels. + + Replaces softmax + top-k + a renormalising divide + a stable argsort + two + gathers + a bincount + a cumulative sum + a scatter: ten tensor ops a + layer, of which the top-k alone was 25 ms of a 2048-token prefill. + + Returns None where the kernel is absent, so the tensor chain stays the + fallback rather than this being a hard dependency. + """ + if not _USE_ROUTE_KERNEL or not hasattr(fvk, 'moe_route_prefill_bf16'): + return None + S = logits.shape[0] + b = _route_buffers(S, fvk, device) + rc = fvk.moe_route_prefill_bf16( + logits.data_ptr(), b['ti'].data_ptr(), b['tw'].data_ptr(), + b['se'].data_ptr(), b['stok'].data_ptr(), b['inv'].data_ptr(), + b['group_off'].data_ptr(), b['ws'].data_ptr(), b['ws_bytes'], + S, _N_EXPERTS, TOPK, _cs()) + if rc: + raise RuntimeError(f'prefill routing failed with {rc}') + return b + + +def _route_sfa_off(route, k, fvk, device): + """Per-expert scale-factor byte offsets for one projection's K.""" + n_col = ((k // 16) + 3) // 4 + off = route['sfa_off'].get(k) + if off is None: + off = torch.empty(_N_EXPERTS, dtype=torch.int32, device=device) + route['sfa_off'][k] = off + fvk.moe_route_sfa_offsets( + route['group_off'].data_ptr(), off.data_ptr(), _N_EXPERTS, n_col, + _cs()) + return off, n_col + + def _grouped_scratch(fvk, device): """One scratch buffer per device for the grouped GEMM's descriptor arrays. @@ -1070,7 +1143,7 @@ def _sf_layout(counts, k, device): return off, n_col -def _moe_experts_grouped_gemm(x, ti, tw, ld, fvk, device): +def _moe_experts_grouped_gemm(x, ti, tw, ld, fvk, device, route=None): """Every routed expert of the layer in two GEMM launches. The per-expert loop below reads each weight once, which is the right amount, @@ -1097,18 +1170,26 @@ def _moe_experts_grouped_gemm(x, ti, tw, ld, fvk, device): gu_a, dn_a = ld['experts_gate_up_alpha_dev'], ld['experts_down_alpha_dev'] scratch, scratch_bytes = _grouped_scratch(fvk, device) - tok_flat, slot_ix = _route_constants(S, device) - exp_flat = ti.reshape(-1).to(torch.int32) - order = exp_flat.argsort(stable=True) - se = exp_flat[order].contiguous() - stok = tok_flat[order] + if route is not None: + se, stok, group_off, order, counts = ( + route['se'], route['stok'], route['group_off'], None, None) + else: + tok_flat, slot_ix = _route_constants(S, device) + exp_flat = ti.reshape(-1).to(torch.int32) + order = exp_flat.argsort(stable=True) + se = exp_flat[order].contiguous() + stok = tok_flat[order] - counts = torch.bincount(se, minlength=_N_EXPERTS) - group_off = torch.zeros(_N_EXPERTS + 1, dtype=torch.int32, device=device) - group_off[1:] = counts.cumsum(0).to(torch.int32) + counts = torch.bincount(se, minlength=_N_EXPERTS) + group_off = torch.zeros(_N_EXPERTS + 1, dtype=torch.int32, + device=device) + group_off[1:] = counts.cumsum(0).to(torch.int32) def project(A, k, n, w_p, w_s, alpha, out, gate=False, perm=None): - sfa_off, n_col = _sf_layout(counts, k, device) + if route is not None: + sfa_off, n_col = _route_sfa_off(route, k, fvk, device) + else: + sfa_off, n_col = _sf_layout(counts, k, device) bound = (_N_EXPERTS + slots // 128 + 1) * n_col * 512 packed = torch.empty(slots, k // 2, dtype=torch.uint8, device=device) sfa = torch.empty(bound, dtype=torch.uint8, device=device) @@ -1149,11 +1230,15 @@ def project(A, k, n, w_p, w_s, alpha, out, gate=False, perm=None): # rows[i] is which sorted row holds slot i, which is exactly the inverse # permutation -- gathering arange through it, as the tiled path has to, # would just reproduce it. - inv = torch.empty(slots, dtype=torch.int32, device=device) - inv[order] = slot_ix.to(torch.int32) + if route is not None: + inv, twc = route['inv'], route['tw'] + else: + inv = torch.empty(slots, dtype=torch.int32, device=device) + inv[order] = slot_ix.to(torch.int32) + twc = tw.contiguous() out = torch.empty(S, HID, dtype=torch.float32, device=device) fvk.moe_weighted_sum_sm120_bf16( - d_dn.data_ptr(), inv.data_ptr(), tw.contiguous().data_ptr(), + d_dn.data_ptr(), inv.data_ptr(), twc.data_ptr(), out.data_ptr(), S, TOPK, n_dn, n_dn, _cs()) return out @@ -1288,24 +1373,39 @@ def _moe_layer(h, ld, fvk, device): # Router GEMM via the deterministic w16a16 kernel (bf16 weight, fp32 # accumulate) instead of the fp32 upcast matmul. bf16 logits match the - # bf16 reference router; softmax/topk stay (already CUDA ops). - logit = F.softmax(_gemm_w16a16(x, rw, fvk, device).float(), -1) - tw, ti = torch.topk(logit, TOPK, -1) - tw = tw / tw.sum(-1, keepdim=True) + # bf16 reference router. + lg = _gemm_w16a16(x, rw, fvk, device) # The block-scaled 4-bit MMA tiles are a build tier, not a given: a target # whose toolchain has no block-scaled mma builds the weight-only tier # instead. Ask the module what it has rather than assuming, so the tile # choice degrades to the grouped GEMV instead of raising mid-prefill. big = x.shape[0] >= _M16_MIN_S - if _USE_BT_MOE and big and hasattr(fvk, 'moe_blocktile_mma_sm120_bf16'): + use_bt = (_USE_BT_MOE and big + and hasattr(fvk, 'moe_blocktile_mma_sm120_bf16')) + use_m16 = (not use_bt and _USE_M16_MOE and big + and hasattr(fvk, 'moe_m16_mma_sm120_bf16')) + grouped = (not use_bt and not use_m16 and big + and hasattr(fvk, 'moe_grouped_gemm_nvfp4_sm100_bf16out')) + # Only the grouped path reads the kernel's permutation, and the tiled + # paths index with the tensor top-k's own indices, so the routing is not + # computed twice for a path that will not use it. + route = _route_prefill(lg, fvk, device) if grouped else None + if route is not None: + ti, tw = route['ti'], route['tw'] + else: + logit = F.softmax(lg.float(), -1) + tw, ti = torch.topk(logit, TOPK, -1) + tw = tw / tw.sum(-1, keepdim=True) + + if use_bt: out = _moe_experts_bt(x, ti, tw, ld, fvk, device) - elif _USE_M16_MOE and big and hasattr(fvk, 'moe_m16_mma_sm120_bf16'): + elif use_m16: out = _moe_experts_m16(x, ti, tw, ld, fvk, device) - elif big and hasattr(fvk, 'moe_grouped_gemm_nvfp4_sm100_bf16out'): + elif grouped: # No threshold: the grouped path wins at every prefill length measured, # because it does not pay per expert for anything. - out = _moe_experts_grouped_gemm(x, ti, tw, ld, fvk, device) + out = _moe_experts_grouped_gemm(x, ti, tw, ld, fvk, device, route) elif (big and _USE_PER_EXPERT_GEMM and x.shape[0] * TOPK >= _PER_EXPERT_MIN_M * _N_EXPERTS and hasattr(fvk, 'fp4_w4a16_gemm_sm120_bf16out')): From 213adf76bac08c260fb79ff919fd42d13eeb9eba Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Fri, 31 Jul 2026 07:21:33 -0400 Subject: [PATCH 50/85] Fuse the residual adds and the MoE tail into their kernels Two patterns were left running as tensor ops on every layer. Each residual add was immediately followed by the norm of what it produced, as two passes over the hidden state and two launches. There is already a kernel that does both and updates the residual stream in place, with the same weight and eps convention as the plain norm. Taking it means the loop carries the normed tensor across each layer boundary rather than the raw one, and the norm after the last layer's residual is the final norm -- so the boundary became a choice of weight instead of a branch. The MoE tail gated the shared expert with a sigmoid, a broadcast multiply, an add onto the routed sum and a cast. The existing bf16 gate-mul-residual kernel does not fit: the routed sum arrives in fp32 from the weighted reduction and routing it through bf16 to reach that kernel would change the accumulation rather than fuse it. A small kernel keeps the arithmetic in fp32 and rounds once, at the store, and takes the sigmoid with it. What the profiler attributes to tensor ops in a 2048-token prefill is now 0.7 ms of KV cache writes and the vendored attention kernel, against 29% of the prefill when this started. TTFT 267.0 -> 259.8 ms at 1024, 462.4 -> 451.4 at 2048, 882.9 -> 845.1 at 4096. Golden prefix 16/16. --- CMakeLists.txt | 1 + csrc/bindings.cpp | 12 ++++ csrc/kernels/moe_shared_combine_edge.cu | 51 +++++++++++++++++ csrc/kernels/moe_shared_combine_edge.cuh | 35 ++++++++++++ .../frontends/torch/_nexn2_rtx_forward.py | 56 +++++++++++++++---- 5 files changed, 145 insertions(+), 10 deletions(-) create mode 100644 csrc/kernels/moe_shared_combine_edge.cu create mode 100644 csrc/kernels/moe_shared_combine_edge.cuh diff --git a/CMakeLists.txt b/CMakeLists.txt index f20feb95..0af83431 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1514,6 +1514,7 @@ if(FLASHRT_ENABLE_QWEN35MOE_CORE) csrc/kernels/act_fuse_sm120.cu csrc/kernels/moe_router_topk_sm120.cu csrc/kernels/moe_route_prefill_edge.cu + csrc/kernels/moe_shared_combine_edge.cu csrc/kernels/moe_weighted_sum_sm120.cu csrc/kernels/w16a16_gemm_sm120.cu csrc/kernels/qwen35moe_e0m3_dequant.cu) diff --git a/csrc/bindings.cpp b/csrc/bindings.cpp index 12bc6201..789ca466 100644 --- a/csrc/bindings.cpp +++ b/csrc/bindings.cpp @@ -187,6 +187,7 @@ extern "C" int cutlass_int8_rowwise_bf16out_t64x128( #include "kernels/act_fuse_sm120.cuh" #include "kernels/moe_router_topk_sm120.cuh" #include "kernels/moe_route_prefill_edge.cuh" +#include "kernels/moe_shared_combine_edge.cuh" #include "kernels/moe_weighted_sum_sm120.cuh" #include "kernels/w16a16_gemm_sm120.cuh" #include "kernels/qwen35moe_e0m3_dequant.cuh" @@ -5531,6 +5532,17 @@ PYBIND11_MODULE(flash_rt_kernels, m) { py::arg("head_dim"), py::arg("use_qk_l2norm") = true, py::arg("stream") = 0); + m.def("moe_shared_gate_combine_edge_bf16", + [](uintptr_t routed, uintptr_t shared, uintptr_t gate, uintptr_t out, + int S, int dim, uintptr_t stream) { + flash_rt::kernels::moe_shared_gate_combine_edge_bf16( + to_ptr(routed), to_ptr(shared), to_ptr(gate), to_ptr(out), + S, dim, to_stream(stream)); + }, + py::arg("routed"), py::arg("shared"), py::arg("gate"), + py::arg("out"), py::arg("S"), py::arg("dim"), + py::arg("stream") = 0); + m.def("moe_route_prefill_bf16", [](uintptr_t logits, uintptr_t ti, uintptr_t tw, uintptr_t se, uintptr_t stok, uintptr_t inv, uintptr_t group_off, uintptr_t ws, diff --git a/csrc/kernels/moe_shared_combine_edge.cu b/csrc/kernels/moe_shared_combine_edge.cu new file mode 100644 index 00000000..15ee59e5 --- /dev/null +++ b/csrc/kernels/moe_shared_combine_edge.cu @@ -0,0 +1,51 @@ +#include "moe_shared_combine_edge.cuh" + +#include + +namespace flash_rt { +namespace kernels { + +namespace { + +__global__ void moe_shared_gate_combine_kernel( + const float* __restrict__ routed, + const __nv_bfloat16* __restrict__ shared, + const __nv_bfloat16* __restrict__ gate, + __nv_bfloat16* __restrict__ out, + int S, + int dim) +{ + const int row = blockIdx.x; + if (row >= S) return; + // expf, not the fast intrinsic: routing downstream is discrete, and a gate + // that lands a few ulp away flips ties in later layers. + const float g = 1.0f / (1.0f + expf(-static_cast(gate[row]))); + const size_t base = static_cast(row) * dim; + for (int i = threadIdx.x; i < dim; i += blockDim.x) { + out[base + i] = __float2bfloat16( + routed[base + i] + static_cast(shared[base + i]) * g); + } +} + +} // namespace + +void moe_shared_gate_combine_edge_bf16( + const void* routed, + const void* shared, + const void* gate, + void* out, + int S, + int dim, + cudaStream_t stream) +{ + if (S <= 0 || dim <= 0) return; + const int threads = dim < 256 ? 128 : 256; + moe_shared_gate_combine_kernel<<>>( + reinterpret_cast(routed), + reinterpret_cast(shared), + reinterpret_cast(gate), + reinterpret_cast<__nv_bfloat16*>(out), S, dim); +} + +} // namespace kernels +} // namespace flash_rt diff --git a/csrc/kernels/moe_shared_combine_edge.cuh b/csrc/kernels/moe_shared_combine_edge.cuh new file mode 100644 index 00000000..8f8d52d2 --- /dev/null +++ b/csrc/kernels/moe_shared_combine_edge.cuh @@ -0,0 +1,35 @@ +#pragma once + +#include + +namespace flash_rt { +namespace kernels { + +// The tail of a fine-grained MoE layer: gate the shared expert and add it to +// the routed sum. +// +// out = routed + shared * sigmoid(gate_logit[row]) +// +// Replaces a sigmoid, a broadcast multiply, an add and a cast -- four tensor +// ops and two full (S, hidden) fp32 intermediates per layer. +// +// The existing bf16 gate-mul-residual kernel does not fit here: the routed sum +// arrives in fp32 from the weighted reduction, and taking it through bf16 to +// reach that kernel would change the accumulation rather than merely fuse it. +// This keeps the arithmetic in fp32 and rounds once, at the store. +// +// routed (S, dim) fp32 +// shared (S, dim) bf16 +// gate (S,) bf16, the raw gate logit -- the sigmoid is applied here +// out (S, dim) bf16 +void moe_shared_gate_combine_edge_bf16( + const void* routed, + const void* shared, + const void* gate, + void* out, + int S, + int dim, + cudaStream_t stream); + +} // namespace kernels +} // namespace flash_rt diff --git a/flash_rt/frontends/torch/_nexn2_rtx_forward.py b/flash_rt/frontends/torch/_nexn2_rtx_forward.py index 5905ed56..c46b32fd 100644 --- a/flash_rt/frontends/torch/_nexn2_rtx_forward.py +++ b/flash_rt/frontends/torch/_nexn2_rtx_forward.py @@ -85,6 +85,25 @@ def _rms_k(x, w, fvk, device, eps): return out.reshape(shp) +def _add_rms_k(h, x, w, fvk, device, eps): + """h += x in place, and return rmsnorm(h, w). + + The residual add and the norm that always follows it were a tensor add and + a separate kernel: two passes over (S, HID) and two launches per half + layer, 160 of each per forward. Same weight and eps convention as the + plain norm, so this is the two of them and not a third behaviour. + + Both must be bf16 and contiguous -- the kernel writes through `h`. + """ + dim = h.shape[-1] + h2 = h.reshape(-1, dim) + out = torch.empty(h2.shape[0], dim, dtype=torch.bfloat16, device=device) + fvk.residual_add_rms_norm( + h2.data_ptr(), x.reshape(-1, dim).data_ptr(), w.data_ptr(), + out.data_ptr(), h2.shape[0], dim, eps, _cs()) + return out.reshape(h.shape) + + def _proj(x2d, ld, base, n, fvk, device): """y = x @ w.T for one projection, dispatching on the loader's scope. @@ -1435,9 +1454,18 @@ def _moe_layer(h, ld, fvk, device): su = _proj(x, ld, 'shared_up_proj', INTER, fvk, device) si = _silu_mul(sg, su, fvk, device) shared = _proj(si, ld, 'shared_down_proj', HID, fvk, device) - # shared-expert scalar gate: GEMM (N=1) via w16a16, then sigmoid. - sgate = torch.sigmoid( - _gemm_w16a16(x, ld['shared_gate_w_t'], fvk, device).float()) + # shared-expert scalar gate: GEMM (N=1) via w16a16. The sigmoid, the + # broadcast multiply, the add onto the routed sum and the cast are one + # kernel; the routed sum stays fp32 until the single rounding at its store. + glog = _gemm_w16a16(x, ld['shared_gate_w_t'], fvk, device) + if hasattr(fvk, 'moe_shared_gate_combine_edge_bf16'): + comb = torch.empty(x.shape[0], HID, dtype=torch.bfloat16, + device=device) + fvk.moe_shared_gate_combine_edge_bf16( + out.data_ptr(), shared.data_ptr(), glog.data_ptr(), + comb.data_ptr(), x.shape[0], HID, _cs()) + return comb.reshape(B, S, HID) + sgate = torch.sigmoid(glog.float()) return (out + shared.float() * sgate).reshape(B, S, HID).to(torch.bfloat16) @@ -1484,10 +1512,14 @@ def nexn2_forward_nvfp4(handles, input_ids, fvk, device, cap=None, ct, st = ct_full[pos_offset:], st_full[pos_offset:] chunked = pos_offset > 0 lin_rank = full_rank = 0 + # Every residual add is immediately followed by the norm of what it + # produced, so the two run as one kernel that updates the residual stream + # in place -- which means the loop carries the *normed* tensor across each + # boundary and takes the first norm before entering it. + h = h.contiguous() + n = _rms_k(h, layers[0]['input_norm_w_t'], fvk, device, eps) for L in range(p['num_layers']): ld = layers[L] - res = h - n = _rms_k(h, ld['input_norm_w_t'], fvk, device, eps) if types[L] == 'linear_attention': init_s = cap.lin_state[lin_rank] if chunked else None conv_h = cap.lin_conv_state[lin_rank] if chunked else None @@ -1498,15 +1530,19 @@ def nexn2_forward_nvfp4(handles, input_ids, fvk, device, cap=None, attn = _full_attn_layer(n, ld, ct, st, fvk, device, eps, cap, full_rank, pos_offset=pos_offset) full_rank += 1 - h = res + attn - res = h - n = _rms_k(h, ld['post_norm_w_t'], fvk, device, eps) - h = res + _moe_layer(n, ld, fvk, device) + n = _add_rms_k(h, attn, ld['post_norm_w_t'], fvk, device, eps) + moe = _moe_layer(n, ld, fvk, device) + # The norm after the last layer's residual is the final norm, and + # between layers it is the next layer's input norm -- one call either + # way, so the boundary is a choice of weight rather than a branch. + nxt = (layers[L + 1]['input_norm_w_t'] if L + 1 < p['num_layers'] + else p['final_norm_w_t']) + n = _add_rms_k(h, moe, nxt, fvk, device, eps) hidden = h[0] # (S, HID) residual stream, pre-final-norm if not compute_logits: return (None, hidden) if return_hidden else None - h = _rms_k(h, p['final_norm_w_t'], fvk, device, eps) + h = n # already the final norm, see above # lm_head via w16a16 (bf16 weight, fp32 accumulate): reads the ~1GB weight # as bf16 (no fp32 widen), same argmax. logits returned bf16. Slice to the # last position first when only the seeding logit is needed (avoids the From 5fd296716ee203e568fefc40fb6b6a3db441f212 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Fri, 31 Jul 2026 07:34:41 -0400 Subject: [PATCH 51/85] Cover the prefill routing kernel Each output is checked in the way that output can fail. The sorted layout and the group boundaries are compared exactly, because a cosine over them would pass while a slot sat under the wrong expert's weight; the weights by relative error, since the two sides reduce the softmax in different orders; and the rank order within a token's top-k not at all, since bf16 logits tie at the boundary and neither side defines which of two equal experts ranks first -- what is asserted instead is that the expert in the row a slot points at is the expert that slot chose. Also covers reproducibility across repeated runs, which the counting sort exists to provide, and the guard on expert counts the warp-wide top-k cannot hold. Lengths include 1 and 1000 so the tail block is not always full. Its own module loader rather than the sibling WY suite's, which skips unless the 27B kernels are present; this one does not depend on them. --- tests/test_moe_route_prefill_edge.py | 166 +++++++++++++++++++++++++++ 1 file changed, 166 insertions(+) create mode 100644 tests/test_moe_route_prefill_edge.py diff --git a/tests/test_moe_route_prefill_edge.py b/tests/test_moe_route_prefill_edge.py new file mode 100644 index 00000000..f7bfd4cf --- /dev/null +++ b/tests/test_moe_route_prefill_edge.py @@ -0,0 +1,166 @@ +"""Equivalence test for the prefill routing kernel. + +The kernel replaces a softmax, a top-k, a renormalising divide, a stable +argsort, two gathers, a bincount, a cumulative sum and a scatter. Each output +is checked against that chain, and the checks are not all the same kind: + +- ``se``, ``stok`` and ``group_off`` are the sorted layout the grouped GEMM + reads and must match exactly. A cosine over these would pass while a slot + sat under the wrong expert's weight. +- The selected expert set per token must match exactly, for the same reason. +- The weights are compared by relative error, since the two reduce the softmax + in different orders. +- The rank order *within* a token's top-k is deliberately not compared. bf16 + logits tie exactly at the top-k boundary and neither implementation defines + which of two equal experts it ranks first; what must hold is that the expert + in the row a slot points at is the expert that slot chose. + +Logits are drawn peaked rather than uniform: top-k selection is decided by the +tail, and uniform draws produce a tie structure a router does not. +""" + +import pytest +import torch +import torch.nn.functional as F + +N_EXPERTS, TOPK = 256, 8 + + +def _load_fvk(): + # Deliberately not the sibling WY suite's loader: that one skips unless + # the 27B kernels are present, and this kernel does not depend on them. + if not torch.cuda.is_available(): + pytest.skip("CUDA is required for the routing kernel test") + try: + from flash_rt import flash_rt_kernels as fvk + except Exception as exc: # pragma: no cover - environmental + pytest.skip(f"flash_rt_kernels is not built: {exc}") + return fvk + + +def _ptr(x): + return x.data_ptr() + + +def _tensor_route(logits, n_experts, topk): + prob = F.softmax(logits.float(), -1) + tw, ti = torch.topk(prob, topk, -1) + tw = tw / tw.sum(-1, keepdim=True) + exp_flat = ti.reshape(-1).to(torch.int32) + order = exp_flat.argsort(stable=True) + se = exp_flat[order].contiguous() + tok = torch.arange(logits.shape[0], + device=logits.device).repeat_interleave(topk) + stok = tok[order] + counts = torch.bincount(se, minlength=n_experts) + group_off = torch.zeros(n_experts + 1, dtype=torch.int32, + device=logits.device) + group_off[1:] = counts.cumsum(0).to(torch.int32) + return ti.to(torch.int32), tw, se, stok, group_off + + +@pytest.mark.parametrize("S", [1, 64, 256, 1000, 2048]) +def test_route_matches_tensor_chain(S): + fvk = _load_fvk() + if not hasattr(fvk, "moe_route_prefill_bf16"): + pytest.skip("moe_route_prefill_bf16 not in this build") + + dev = "cuda:0" + g = torch.Generator(device=dev).manual_seed(7 + S) + logits = (torch.randn(S, N_EXPERTS, generator=g, device=dev) * 2.5 + ).to(torch.bfloat16) + ti_r, tw_r, se_r, stok_r, goff_r = _tensor_route(logits, N_EXPERTS, TOPK) + + slots = S * TOPK + ti = torch.empty(S, TOPK, dtype=torch.int32, device=dev) + tw = torch.empty(S, TOPK, dtype=torch.float32, device=dev) + se = torch.empty(slots, dtype=torch.int32, device=dev) + # int64: this is handed to the activation quantiser as its gather index and + # that kernel reads a long. + stok = torch.empty(slots, dtype=torch.int64, device=dev) + inv = torch.empty(slots, dtype=torch.int32, device=dev) + goff = torch.empty(N_EXPERTS + 1, dtype=torch.int32, device=dev) + nbytes = int(fvk.moe_route_prefill_workspace_bytes(S, TOPK, N_EXPERTS)) + ws = torch.empty(nbytes, dtype=torch.uint8, device=dev) + + rc = fvk.moe_route_prefill_bf16( + _ptr(logits), _ptr(ti), _ptr(tw), _ptr(se), _ptr(stok), _ptr(inv), + _ptr(goff), _ptr(ws), nbytes, S, N_EXPERTS, TOPK, + torch.cuda.current_stream().cuda_stream) + assert rc == 0, f"routing kernel returned {rc}" + torch.cuda.synchronize(dev) + + assert torch.equal(se, se_r), "sorted expert layout differs" + assert torch.equal(stok, stok_r), "sorted token layout differs" + assert torch.equal(goff, goff_r), "group boundaries differ" + + # The expert in the row a slot points at is the expert that slot chose. + assert torch.equal(se[inv.long()].reshape(S, TOPK), ti) + + oi = ti.argsort(stable=True, dim=-1) + oj = ti_r.argsort(stable=True, dim=-1) + assert torch.equal(torch.gather(ti, 1, oi), torch.gather(ti_r, 1, oj)), \ + "selected expert sets differ" + rel = ((torch.gather(tw, 1, oi) - torch.gather(tw_r, 1, oj)).abs() + / torch.gather(tw_r, 1, oj).abs().clamp_min(1e-9)).max() + assert rel < 1e-5, f"weights differ by {rel:.3e}" + + +@pytest.mark.parametrize("S", [64, 2048]) +def test_route_is_reproducible(S): + """Prefill seeds a decode that has to reproduce, so the permutation may + not depend on the order blocks happen to arrive in.""" + fvk = _load_fvk() + if not hasattr(fvk, "moe_route_prefill_bf16"): + pytest.skip("moe_route_prefill_bf16 not in this build") + + dev = "cuda:0" + g = torch.Generator(device=dev).manual_seed(11) + logits = (torch.randn(S, N_EXPERTS, generator=g, device=dev) * 2.5 + ).to(torch.bfloat16) + slots = S * TOPK + nbytes = int(fvk.moe_route_prefill_workspace_bytes(S, TOPK, N_EXPERTS)) + + def run(): + ti = torch.empty(S, TOPK, dtype=torch.int32, device=dev) + tw = torch.empty(S, TOPK, dtype=torch.float32, device=dev) + se = torch.empty(slots, dtype=torch.int32, device=dev) + stok = torch.empty(slots, dtype=torch.int64, device=dev) + inv = torch.empty(slots, dtype=torch.int32, device=dev) + goff = torch.empty(N_EXPERTS + 1, dtype=torch.int32, device=dev) + ws = torch.empty(nbytes, dtype=torch.uint8, device=dev) + rc = fvk.moe_route_prefill_bf16( + _ptr(logits), _ptr(ti), _ptr(tw), _ptr(se), _ptr(stok), _ptr(inv), + _ptr(goff), _ptr(ws), nbytes, S, N_EXPERTS, TOPK, + torch.cuda.current_stream().cuda_stream) + assert rc == 0 + torch.cuda.synchronize(dev) + return se, stok, inv, goff, ti, tw + + a = run() + for _ in range(3): + b = run() + for x, y in zip(a, b): + assert torch.equal(x, y), "routing is not reproducible" + + +def test_route_rejects_shapes_it_cannot_hold(): + """The top-k spreads a logit row across one warp, so the expert count has + to be 32 times a power of two. Say so rather than compute nonsense.""" + fvk = _load_fvk() + if not hasattr(fvk, "moe_route_prefill_bf16"): + pytest.skip("moe_route_prefill_bf16 not in this build") + + dev = "cuda:0" + S, E = 16, 96 # 96 = 32 * 3, not a power of two + logits = torch.zeros(S, E, dtype=torch.bfloat16, device=dev) + slots = S * TOPK + buf = lambda n, dt: torch.empty(n, dtype=dt, device=dev) + nbytes = int(fvk.moe_route_prefill_workspace_bytes(S, TOPK, E)) + rc = fvk.moe_route_prefill_bf16( + _ptr(logits), _ptr(buf(S * TOPK, torch.int32)), + _ptr(buf(S * TOPK, torch.float32)), _ptr(buf(slots, torch.int32)), + _ptr(buf(slots, torch.int64)), _ptr(buf(slots, torch.int32)), + _ptr(buf(E + 1, torch.int32)), _ptr(buf(nbytes, torch.uint8)), + nbytes, S, E, TOPK, torch.cuda.current_stream().cuda_stream) + assert rc != 0, "unsupported expert count was accepted" From 0364058beb1bea518a9c733316bfbee26d0c098e Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Fri, 31 Jul 2026 08:27:20 -0400 Subject: [PATCH 52/85] Scan the routing histogram across the block, not along one thread The per-expert exclusive scan over slot-blocks walked them in a loop on a single thread. The number of slot-blocks grows with the sequence, so that is O(S) of dependent loads with no parallelism -- sixty-four iterations at two thousand tokens, a thousand at thirty-two. Scanned in tiles across the block instead, carrying between them. Kept because a serial scan over a sequence-dependent count is wrong, not because it explains anything: it was written while chasing a long-context slowdown and it is 4.3 ms of that at 10240 tokens, which is not the answer. The answer was elsewhere and is recorded with the measurements. Routing holds 9x against the tensor chain at every length now: 0.057 ms at 2048, 0.178 at 10240, 0.510 at 32768. Equivalence and reproducibility pass at all of them, including 32768. --- csrc/kernels/moe_route_prefill_edge.cu | 40 +++++++++++++++++++++----- 1 file changed, 33 insertions(+), 7 deletions(-) diff --git a/csrc/kernels/moe_route_prefill_edge.cu b/csrc/kernels/moe_route_prefill_edge.cu index def39bd0..10434e4e 100644 --- a/csrc/kernels/moe_route_prefill_edge.cu +++ b/csrc/kernels/moe_route_prefill_edge.cu @@ -125,6 +125,12 @@ __global__ void slot_hist_kernel( // One block per expert: exclusive scan of that expert's per-block counts, so a // scatter block knows where its own slots for that expert begin. +// +// Scanned across the block in tiles rather than by one thread in a loop. The +// number of slot-blocks grows with the sequence, so a serial walk here is +// O(S) on a single thread -- invisible at two thousand tokens, and the reason +// the prefill rate fell away past four thousand. +template __global__ void expert_block_scan_kernel( const int* __restrict__ blk_hist, int* __restrict__ blk_off, @@ -133,14 +139,34 @@ __global__ void expert_block_scan_kernel( int n_experts) { const int e = blockIdx.x; - if (threadIdx.x != 0) return; - int acc = 0; - for (int b = 0; b < n_blocks; ++b) { + const int t = threadIdx.x; + __shared__ int s[kThreads]; + __shared__ int s_carry; + if (t == 0) s_carry = 0; + __syncthreads(); + + for (int base = 0; base < n_blocks; base += kThreads) { + const int b = base + t; const size_t off = static_cast(b) * n_experts + e; - blk_off[off] = acc; - acc += blk_hist[off]; + const int own = (b < n_blocks) ? blk_hist[off] : 0; + s[t] = own; + __syncthreads(); + + // Hillis-Steele inclusive scan; subtracting own value gives the exclusive + // one without a second pass. + for (int d = 1; d < kThreads; d <<= 1) { + const int add = (t >= d) ? s[t - d] : 0; + __syncthreads(); + s[t] += add; + __syncthreads(); + } + const int tile_total = s[kThreads - 1]; + if (b < n_blocks) blk_off[off] = s_carry + s[t] - own; + __syncthreads(); + if (t == 0) s_carry += tile_total; + __syncthreads(); } - counts[e] = acc; + if (t == 0) counts[e] = s_carry; } __global__ void group_off_kernel( @@ -290,7 +316,7 @@ int moe_route_prefill_bf16( slot_hist_kernel<<>>( ti_i, blk_hist, slots, n_experts); - expert_block_scan_kernel<<>>( + expert_block_scan_kernel<256><<>>( blk_hist, blk_off, counts, nblk, n_experts); group_off_kernel<<<1, 32, 0, stream>>>( From 54ff8db0883c75f410fb2fc052961a7750877d65 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Fri, 31 Jul 2026 11:39:44 -0400 Subject: [PATCH 53/85] Build FA2 on sm_110, where the chunked window needed it This arch was left out of the FA2 arch list on the grounds that it has its own attention path and that FA2 would cost about ten megabytes for nothing. The first half is true and the second was measured against the wrong matrix. Its own path, the cuBLAS decomposed attention, materialises an (S * heads, S_kv) score buffer -- 3.4 GB per layer at ten thousand tokens. Nothing else was available: the FA4 vendor loads here and refuses this shape, because head_dim 256 dispatches to the dedicated 2CTA kernel that the vendor trim removed, and the trim's premise that those branches never execute on Blackwell does not hold at 256. So attention fell to torch's SDPA, which takes a fused backend on a square block and materialises the scores on a chunked one. The instantiation this needs was already in the tree -- bf16, head_dim 256, causal -- and the build comment next to it already said it serves this model's prefill. Only the arch gate was missing. The matrix is narrowed to that one dtype and head_dim here, which is why the size argument no longer applies, and the gencode is this build's own sm_110a rather than compute_80 PTX, on the evidence from the 5090 that PTX-derived SASS drifts. Against what torch picks today, at this model's shape: square windows cos 1.000000 and the same speed to within 1%, non-square windows cos 0.999997 and 20x -- 137.8 ms to 6.7 at 2048 against 10240, 1736 to 81 at 8192 against 32768. A chunked prefill at 10240 goes 3517 -> 2207 ms, which is 4% over the single-pass path where it was 64%. Decode keeps its existing attention on this arch. At the decode shape the two agree to bf16 precision -- against an fp32 reference, 2.0e-3 relative for both at kv=64 -- so taking it there is worth about 1% of a step and is not worth moving the golden fixture for. FLASHRT_NEXN2_DECODE_FA2 overrides. Also fixes a missing import in the decode backend's FA2 probe, which had never run on a target that builds FA2. --- CMakeLists.txt | 71 ++++++++++++++++----- flash_rt/hardware/rtx/attn_backend_nexn2.py | 32 ++++++++-- 2 files changed, 84 insertions(+), 19 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 0af83431..1d116aeb 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -82,18 +82,25 @@ message(STATUS "Using gencode flag: ${GPU_GENCODE}") # SM86/87/89 share the same source instantiations; we emit arch-specific # codegen per target so Jetson Orin (SM87) is not forced through an # incompatible cubin. SM120 (5090) uses PTX JIT from -# compute_80. Thor SM110 has its own attention path (fvk.attention_qkv_fp16 -# cuBLAS decomposed), hand-tuned for unified LPDDR memory — building -# FA2 there would add ~10 MB to the .so and unused compile time. +# compute_80. +# +# Thor SM110 was excluded on the grounds that it has its own attention path +# (fvk.attention_qkv_fp16, cuBLAS decomposed) and that FA2 would cost ~10 MB +# of .so for nothing. That trade does not survive a long prefill: the +# decomposed path materialises an (S * heads, S_kv) score buffer, which is +# 3.4 GB per layer at ten thousand tokens, and the model that needs it here +# is bf16 at head_dim 256 -- one instantiation, not the twelve the size +# estimate assumed. FA2_HDIMS/FA2_DTYPES are narrowed for this target below. if(GPU_ARCH STREQUAL "80" OR GPU_ARCH STREQUAL "86" OR GPU_ARCH STREQUAL "87" OR - GPU_ARCH STREQUAL "89" OR GPU_ARCH STREQUAL "120" OR + GPU_ARCH STREQUAL "89" OR GPU_ARCH STREQUAL "110" OR + GPU_ARCH STREQUAL "120" OR GPU_ARCH STREQUAL "121") set(ENABLE_FA2 ON) message(STATUS "FA2 in-SO attention: ENABLED (sm_${GPU_ARCH})") else() set(ENABLE_FA2 OFF) - message(STATUS "FA2 in-SO attention: DISABLED (Thor SM110 uses fvk.attention_qkv_fp16 cuBLAS path)") + message(STATUS "FA2 in-SO attention: DISABLED (sm_${GPU_ARCH})") endif() # SM80-family CUTLASS INT8 kernels used by the Jetson Orin SM87 Pi0.5 fast @@ -137,6 +144,21 @@ set(FA2_HDIMS "64;96;128;256" CACHE STRING "Semicolon-separated FA2 head_dim instantiations to build. Qwen3-VL 2B vision uses 64; other shipped models use 96 and 256.") set(FA2_DTYPES "fp16;bf16" CACHE STRING "Semicolon-separated FA2 dtype instantiations to build. Pi0 uses fp16; pi0.5/groot use bf16.") + +# Thor runs one model family through FA2 -- Qwen3.6 full attention, bf16 at +# head_dim 256 -- so the default there is that one instantiation rather than +# the twelve-file matrix the RTX targets distribute. This is why enabling FA2 +# on this arch does not cost what the original exclusion assumed. Both remain +# cache variables: an explicit -DFA2_HDIMS on the command line still wins. +if(GPU_ARCH STREQUAL "110") + if(NOT DEFINED CACHE{FA2_HDIMS} OR FA2_HDIMS STREQUAL "64;96;128;256") + set(FA2_HDIMS "256" CACHE STRING "" FORCE) + endif() + if(NOT DEFINED CACHE{FA2_DTYPES} OR FA2_DTYPES STREQUAL "fp16;bf16") + set(FA2_DTYPES "bf16" CACHE STRING "" FORCE) + endif() +endif() + option(FLASHRT_ENABLE_NATIVE_CPP "Build Python-free operation libraries for native C++ consumers" OFF) option(FLASHRT_BUILD_FA2_PYTHON_ADAPTER @@ -990,6 +1012,15 @@ if(ENABLE_FA2 AND > ) message(STATUS "FA2 vendor arch: sm_${GPU_ARCH} AOT only (FA2_ARCH_NATIVE_ONLY=ON)") + elseif(GPU_ARCH STREQUAL "110") + # Native SASS via the same gencode the rest of this build uses, not + # compute_80 PTX. The 5090 learned that lesson: routing it through + # compute_80 PTX produced SASS that drifted from the native build by + # about an fp16 ULP a layer. + target_compile_options(fa2_vendor_obj PRIVATE + $<$:${GPU_GENCODE}> + ) + message(STATUS "FA2 vendor arch: sm_110a AOT only (FA2_ARCH_NATIVE_ONLY=ON)") elseif(GPU_ARCH STREQUAL "120") target_compile_options(fa2_vendor_obj PRIVATE $<$: @@ -1011,16 +1042,26 @@ if(ENABLE_FA2 AND # build). See commit history for the reason each gencode matters — # routing 5090 through compute_80 PTX produced a subtle sm_120 SASS # drift that accumulated to cos 0.98 under Pi0 FP8. - target_compile_options(fa2_vendor_obj PRIVATE - $<$: - "SHELL:-gencode arch=compute_80,code=sm_80" - "SHELL:-gencode arch=compute_120,code=sm_120" - "SHELL:-gencode arch=compute_120,code=compute_120" - "SHELL:-gencode arch=compute_121,code=sm_121" - "SHELL:-gencode arch=compute_121,code=compute_121" - > - ) - message(STATUS "FA2 vendor arch: sm_80 + sm_120/sm_121 AOT + Blackwell PTX fallback (default)") + # sm_110 is not in the consumer family this multi-arch list distributes + # for, and compute_80 SASS does not run on it, so it takes its own + # gencode here as well rather than silently getting no cubin. + if(GPU_ARCH STREQUAL "110") + target_compile_options(fa2_vendor_obj PRIVATE + $<$:${GPU_GENCODE}> + ) + message(STATUS "FA2 vendor arch: sm_110a AOT (Thor)") + else() + target_compile_options(fa2_vendor_obj PRIVATE + $<$: + "SHELL:-gencode arch=compute_80,code=sm_80" + "SHELL:-gencode arch=compute_120,code=sm_120" + "SHELL:-gencode arch=compute_120,code=compute_120" + "SHELL:-gencode arch=compute_121,code=sm_121" + "SHELL:-gencode arch=compute_121,code=compute_121" + > + ) + message(STATUS "FA2 vendor arch: sm_80 + sm_120/sm_121 AOT + Blackwell PTX fallback (default)") + endif() endif() # Macros the wrapper reads to #ifdef-guard optional hdim/dtype diff --git a/flash_rt/hardware/rtx/attn_backend_nexn2.py b/flash_rt/hardware/rtx/attn_backend_nexn2.py index 07e5de31..6fceeef3 100644 --- a/flash_rt/hardware/rtx/attn_backend_nexn2.py +++ b/flash_rt/hardware/rtx/attn_backend_nexn2.py @@ -102,17 +102,36 @@ def __init__(self, max_seq: int, max_q_seq: int = 1, dtype=None, dtype=torch.float32, device=d, ) - # A target may not build FA2 at all -- Thor uses FA4 instead, so the - # arch list deliberately omits it there. Absence is a fallback, not an + # A target may not build FA2 at all. Absence is a fallback, not an # error, because the reference path below computes the same thing. + # + # Prefill and decode want different answers about FA2. Prefill gains + # 20x on a chunked block's non-square window. At the decode shape the + # two are the same answer to bf16 precision -- measured against an + # fp32 reference, 2.0e-3 relative for both at kv=64, 2.2e-3 against + # 2.1e-3 at kv=2048 -- so taking it there buys about 1% of a step and + # moves two of the sixteen golden tokens, because a bf16-level + # difference in one step flips a later one. That is the fixture losing + # its meaning in exchange for 1%, which is the wrong trade. + # + # So the default follows what each target already validated: on the + # arch that has always had FA2 in decode, keep it; on sm_110, where + # FA2 has only just started building and the fixture was recorded + # through the reference path, decline it. FLASHRT_NEXN2_DECODE_FA2 + # overrides either way. + import os as _os + _cap = torch.cuda.get_device_capability() + _default = "0" if _cap == (11, 0) else "1" + want_fa2 = _os.environ.get( + "FLASHRT_NEXN2_DECODE_FA2", _default) != "0" try: from flash_rt import flash_rt_fa2 as _fa2 except ImportError: self._fa2 = None self._fa2_fwd = None else: - self._fa2 = _fa2 - self._fa2_fwd = _fa2.fwd_bf16 + self._fa2 = _fa2 if want_fa2 else None + self._fa2_fwd = _fa2.fwd_bf16 if want_fa2 else None self._num_sms = torch.cuda.get_device_properties( torch.cuda.current_device() ).multi_processor_count @@ -132,6 +151,11 @@ def _probe_fa2(self) -> bool: So run one small case against a reference and compare. The cost is one launch at construction. """ + # Imported here, not at module scope, for the same reason as F: this + # module is written to import without torch present. The body had + # never run on a target that builds FA2, so the missing name sat + # unnoticed until this arch started building one. + import torch import torch.nn.functional as F q_seq, kv_seq = 1, 8 From b0167d597cfdac31b267e5478bae660266b95cc2 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Fri, 31 Jul 2026 12:01:20 -0400 Subject: [PATCH 54/85] Cover the rest of the prefill kernels The routing kernel had a suite; the WY front matter, the MoE tail and the attention did not, which is to say three shipped kernels were carrying no evidence beyond an end-to-end token count. Each is checked against the thing it replaced, and each in the way that thing can fail. The packing and the GQA broadcast are exact -- a permutation that is nearly right is wrong, and the broadcast is the part that never materialises, so it is the part most able to go wrong quietly. The l2 normalisation and the gate sum are relative, since the reduction orders differ. Lengths include 1, 100 and 4097 so the tail block is not always full, and the zeroed tail is asserted because output_o reads the whole chunk. The attention test checks the two windows separately and asserts the causal alignment directly: top-left and bottom-right coincide when the block is square, so a kernel doing the wrong one would pass every square case and silently drop history on a chunked one. Writing the MoE combine's assertion took three tries and the kernel was right each time, which is the useful part. A relative norm passes almost anything at bf16, where one unit in the last place is already 4e-3. A ULP distance is meaningless where the two sides straddle zero -- the kernel fuses the multiply into the add, so it rounds once where the tensor chain rounds twice, and on rows where the routed sum and the gated shared expert nearly cancel it returns 9.3e-10 against exactly 0. Absolute and relative together is the criterion that admits that and would still catch a wrong gate, row index or dtype. 34 cases, all passing on sm_110. --- tests/test_fa2_causal_nexn2_shape.py | 104 +++++++++++++++++ tests/test_gdn_wy_prefill_edge.py | 160 ++++++++++++++++++++++++++ tests/test_moe_shared_combine_edge.py | 96 ++++++++++++++++ 3 files changed, 360 insertions(+) create mode 100644 tests/test_fa2_causal_nexn2_shape.py create mode 100644 tests/test_gdn_wy_prefill_edge.py create mode 100644 tests/test_moe_shared_combine_edge.py diff --git a/tests/test_fa2_causal_nexn2_shape.py b/tests/test_fa2_causal_nexn2_shape.py new file mode 100644 index 00000000..4bffb195 --- /dev/null +++ b/tests/test_fa2_causal_nexn2_shape.py @@ -0,0 +1,104 @@ +"""Does the vendored FA2 compute this model's attention, and on both windows? + +Two windows matter and they are not the same test. A square block is what a +single-pass prefill asks for and torch already had a fused backend for it, so +the bar there is "no worse". A non-square block -- Sq queries against Sk +accumulated keys -- is what a chunked prefill asks for, torch had no fused +backend for it, and FA2's causal is bottom-right aligned, which is precisely +what that window means. Getting the alignment wrong is silent: it truncates +history and still returns plausible numbers. + +Skipped where FA2 is not built, since that is a target property rather than a +failure. +""" + +import pytest +import torch +import torch.nn.functional as F + +NQ, NKV, HD = 16, 2, 256 + + +def _fwd(): + if not torch.cuda.is_available(): + pytest.skip("CUDA is required for the FA2 shape test") + try: + import flash_rt.frontends.torch._nexn2_rtx_forward as fwd + except Exception as exc: # pragma: no cover - environmental + pytest.skip(f"frontend not importable: {exc}") + if fwd._get_fa2() is None: + pytest.skip("FA2 is not built for this target") + return fwd + + +def _reference(q, k, v, dev): + """Bottom-right causal, fp32, scores materialised. Slow and unambiguous.""" + Sq, Sk = q.shape[1], k.shape[1] + qt = q.transpose(1, 2).float() + kt = k.transpose(1, 2).float().repeat_interleave(NQ // NKV, 1) + vt = v.transpose(1, 2).float().repeat_interleave(NQ // NKV, 1) + s = (qt @ kt.transpose(-1, -2)) * (HD ** -0.5) + qi = torch.arange(Sk - Sq, Sk, device=dev).unsqueeze(1) + mask = torch.arange(Sk, device=dev).unsqueeze(0) <= qi + s = s.masked_fill(~mask, float("-inf")) + return (F.softmax(s, -1) @ vt).transpose(1, 2) + + +def test_probe_accepts_this_target(): + fwd = _fwd() + assert fwd._fa2_usable("cuda:0"), \ + "FA2 is built but its own probe rejects it here" + + +@pytest.mark.parametrize("Sq,Sk", [(64, 64), (512, 512), (1024, 1024)]) +def test_square_window(Sq, Sk): + fwd = _fwd() + dev = "cuda:0" + g = torch.Generator(device=dev).manual_seed(Sq) + q = torch.randn(1, Sq, NQ, HD, generator=g, device=dev, + dtype=torch.bfloat16) + k = torch.randn(1, Sk, NKV, HD, generator=g, device=dev, + dtype=torch.bfloat16) + v = torch.randn_like(k) + o = fwd._fa2_causal_attn(q, k, v, dev, _probe=True) + torch.cuda.synchronize(dev) + ref = _reference(q, k, v, dev) + rel = ((o.float() - ref).norm() / ref.norm()).item() + assert rel < 5e-3, f"square window off by {rel:.3e}" + + +@pytest.mark.parametrize("Sq,Sk", [(64, 256), (512, 2048), (256, 4096)]) +def test_non_square_window_is_bottom_right(Sq, Sk): + """The one that used to have no fused backend. + + Also checks the alignment explicitly: a top-left reading of the same + request drops the history, and the two only coincide when Sq == Sk, so a + kernel that quietly did the wrong one would pass every square case above. + """ + fwd = _fwd() + dev = "cuda:0" + g = torch.Generator(device=dev).manual_seed(Sq * 31 + Sk) + q = torch.randn(1, Sq, NQ, HD, generator=g, device=dev, + dtype=torch.bfloat16) + k = torch.randn(1, Sk, NKV, HD, generator=g, device=dev, + dtype=torch.bfloat16) + v = torch.randn_like(k) + o = fwd._fa2_causal_attn(q, k, v, dev, _probe=True) + torch.cuda.synchronize(dev) + + ref = _reference(q, k, v, dev) + rel = ((o.float() - ref).norm() / ref.norm()).item() + assert rel < 5e-3, f"non-square window off by {rel:.3e}" + + # Top-left would attend query i to keys [0, i] instead of [0, Sk-Sq+i]. + qt = q.transpose(1, 2).float() + kt = k.transpose(1, 2).float().repeat_interleave(NQ // NKV, 1) + vt = v.transpose(1, 2).float().repeat_interleave(NQ // NKV, 1) + s = (qt @ kt.transpose(-1, -2)) * (HD ** -0.5) + tl = torch.arange(Sk, device=dev).unsqueeze(0) <= torch.arange( + Sq, device=dev).unsqueeze(1) + topleft = (F.softmax(s.masked_fill(~tl, float("-inf")), -1) + @ vt).transpose(1, 2) + rel_tl = ((o.float() - topleft).norm() / topleft.norm()).item() + assert rel_tl > 0.05, \ + "output matches the top-left window; the alignment is wrong" diff --git a/tests/test_gdn_wy_prefill_edge.py b/tests/test_gdn_wy_prefill_edge.py new file mode 100644 index 00000000..c6327459 --- /dev/null +++ b/tests/test_gdn_wy_prefill_edge.py @@ -0,0 +1,160 @@ +"""Equivalence tests for the WY front-matter kernels. + +These replaced a chain of tensor ops -- two l2 normalisations, a per-chunk gate +cumulative sum, a GQA broadcast and the chunk-major packings -- so the +reference here is that chain, written out again rather than imported, because +the point is to check the kernel against what it replaced and not against +itself. + +The l2 normalisation is compared by relative error, not exactly: the kernel +reduces 128 elements in butterfly order and the tensor path reduces them in +torch's order. The packing and the broadcast are pure data movement and are +compared exactly -- a permutation that is nearly right is wrong. +""" + +import pytest +import torch +import torch.nn.functional as F + +NK, NV, HD, CH = 16, 32, 128, 64 +QKG = NV // NK + + +def _load_fvk(): + if not torch.cuda.is_available(): + pytest.skip("CUDA is required for the WY prefill kernel tests") + try: + from flash_rt import flash_rt_kernels as fvk + except Exception as exc: # pragma: no cover - environmental + pytest.skip(f"flash_rt_kernels is not built: {exc}") + for name in ("gdn_wy_norm_pack_q_cumsum_edge_bf16", + "gdn_wy_pack_v_edge_bf16"): + if not hasattr(fvk, name): + pytest.skip(f"{name} not in this build") + return fvk + + +def _ptr(x): + return x.data_ptr() + + +def _ref_l2(x): + """The l2 normalisation the sequential scan and the tensor path both use.""" + xf = x.float() + return (xf * torch.rsqrt((xf * xf).sum(-1, keepdim=True) + 1e-6)).to( + torch.bfloat16) + + +def _ref_pack(x, ch=CH): + """(S, H, D) -> (chunks, H, ch, D), zero-padded past S.""" + s, hh, d = x.shape + pad = (-s) % ch + if pad: + x = F.pad(x, (0, 0, 0, 0, 0, pad)) + return x.reshape(-1, ch, hh, d).permute(0, 2, 1, 3).contiguous() + + +def _ref_gcumsum(g, ch=CH): + s = g.shape[0] + pad = (-s) % ch + gp = F.pad(g, (0, 0, 0, pad)) if pad else g + return torch.cumsum(gp.float().reshape(-1, ch, g.shape[1]), 1).reshape( + -1, g.shape[1])[:s].to(torch.bfloat16) + + +@pytest.mark.parametrize("S", [1, 64, 100, 512, 2048, 4097]) +def test_norm_pack_cumsum_matches_the_chain_it_replaced(S): + fvk = _load_fvk() + dev = "cuda:0" + g = torch.Generator(device=dev).manual_seed(3 + S) + + # q and k arrive GQA-broadcast across the v-head slots, which is the form + # the conv split writes and the form the kernel expects. + q16 = torch.randn(S, NK, HD, generator=g, device=dev, dtype=torch.bfloat16) + k16 = torch.randn(S, NK, HD, generator=g, device=dev, dtype=torch.bfloat16) + qb = q16.repeat_interleave(QKG, 1).contiguous() + kb = k16.repeat_interleave(QKG, 1).contiguous() + gate = torch.randn(S, NV, generator=g, device=dev, + dtype=torch.bfloat16) * 0.1 + + chunks = (S + CH - 1) // CH + k_l2 = torch.empty(S, NK, HD, dtype=torch.bfloat16, device=dev) + q_pack = torch.empty(chunks, NV, CH, HD, dtype=torch.bfloat16, device=dev) + gc = torch.empty(S, NV, dtype=torch.bfloat16, device=dev) + fvk.gdn_wy_norm_pack_q_cumsum_edge_bf16( + _ptr(qb), _ptr(kb), _ptr(gate), _ptr(k_l2), _ptr(q_pack), _ptr(gc), + S, NK, NV, HD, QKG, torch.cuda.current_stream().cuda_stream) + torch.cuda.synchronize(dev) + + k_ref = _ref_l2(k16) + q_ref_pack = _ref_pack(_ref_l2(q16).repeat_interleave(QKG, 1)) + gc_ref = _ref_gcumsum(gate) + + def rel(a, b): + return ((a.float() - b.float()).norm() + / b.float().norm().clamp_min(1e-9)).item() + + assert rel(k_l2, k_ref) < 5e-3, "k l2 normalisation drifted" + assert rel(q_pack, q_ref_pack) < 5e-3, "packed q drifted" + assert rel(gc, gc_ref) < 5e-3, "gate cumulative sum drifted" + + # The tail of the last chunk must be zero, not whatever was in the buffer: + # output_o reads the whole chunk. + if S % CH: + assert torch.equal(q_pack[-1, :, S % CH:, :], + torch.zeros_like(q_pack[-1, :, S % CH:, :])), \ + "packed q tail is not zeroed" + + # Every v-head slot of a GQA group must carry the same vector -- the + # broadcast is the part that never materialises, so it is the part most + # able to go wrong silently. + for r in range(1, QKG): + assert torch.equal(q_pack[:, 0::QKG], q_pack[:, r::QKG]), \ + f"GQA group member {r} differs from its leader" + + +@pytest.mark.parametrize("S", [1, 64, 100, 512, 2048, 4097]) +def test_pack_v_matches_the_reference_permutation(S): + fvk = _load_fvk() + dev = "cuda:0" + g = torch.Generator(device=dev).manual_seed(11 + S) + v = torch.randn(S, NV, HD, generator=g, device=dev, dtype=torch.bfloat16) + + chunks = (S + CH - 1) // CH + v_pack = torch.empty(chunks, NV, CH, HD, dtype=torch.bfloat16, device=dev) + fvk.gdn_wy_pack_v_edge_bf16( + _ptr(v), _ptr(v_pack), S, NV, HD, + torch.cuda.current_stream().cuda_stream) + torch.cuda.synchronize(dev) + + # Pure data movement: exact, including the zero tail. + assert torch.equal(v_pack, _ref_pack(v)) + + +@pytest.mark.parametrize("S", [64, 2048]) +def test_reproducible(S): + """Prefill seeds a decode that has to reproduce.""" + fvk = _load_fvk() + dev = "cuda:0" + g = torch.Generator(device=dev).manual_seed(5) + qb = torch.randn(S, NV, HD, generator=g, device=dev, dtype=torch.bfloat16) + kb = torch.randn(S, NV, HD, generator=g, device=dev, dtype=torch.bfloat16) + gate = torch.randn(S, NV, generator=g, device=dev, dtype=torch.bfloat16) + chunks = (S + CH - 1) // CH + st = torch.cuda.current_stream().cuda_stream + + def run(): + k_l2 = torch.empty(S, NK, HD, dtype=torch.bfloat16, device=dev) + q_pack = torch.empty(chunks, NV, CH, HD, dtype=torch.bfloat16, + device=dev) + gc = torch.empty(S, NV, dtype=torch.bfloat16, device=dev) + fvk.gdn_wy_norm_pack_q_cumsum_edge_bf16( + _ptr(qb), _ptr(kb), _ptr(gate), _ptr(k_l2), _ptr(q_pack), + _ptr(gc), S, NK, NV, HD, QKG, st) + torch.cuda.synchronize(dev) + return k_l2, q_pack, gc + + a = run() + for _ in range(3): + for x, y in zip(a, run()): + assert torch.equal(x, y), "WY front matter is not reproducible" diff --git a/tests/test_moe_shared_combine_edge.py b/tests/test_moe_shared_combine_edge.py new file mode 100644 index 00000000..b3ff7d47 --- /dev/null +++ b/tests/test_moe_shared_combine_edge.py @@ -0,0 +1,96 @@ +"""Equivalence test for the MoE shared-expert combine. + +out = routed + shared * sigmoid(gate_logit[row]) + +The routed sum arrives in fp32 from the weighted reduction and the kernel keeps +it there, rounding once at the store. That is the reason this kernel exists +rather than the bf16 gate-mul-residual one already in the tree, so the +reference computes it the same way -- in fp32, cast once -- and the comparison +is exact where it can be. +""" + +import pytest +import torch + +HID = 2048 + + +def _load_fvk(): + if not torch.cuda.is_available(): + pytest.skip("CUDA is required for the MoE combine test") + try: + from flash_rt import flash_rt_kernels as fvk + except Exception as exc: # pragma: no cover - environmental + pytest.skip(f"flash_rt_kernels is not built: {exc}") + if not hasattr(fvk, "moe_shared_gate_combine_edge_bf16"): + pytest.skip("moe_shared_gate_combine_edge_bf16 not in this build") + return fvk + + +@pytest.mark.parametrize("S", [1, 17, 256, 2048]) +def test_matches_the_fp32_chain(S): + fvk = _load_fvk() + dev = "cuda:0" + g = torch.Generator(device=dev).manual_seed(23 + S) + routed = torch.randn(S, HID, generator=g, device=dev, + dtype=torch.float32) + shared = torch.randn(S, HID, generator=g, device=dev, + dtype=torch.bfloat16) + glog = torch.randn(S, 1, generator=g, device=dev, dtype=torch.bfloat16) + + out = torch.empty(S, HID, dtype=torch.bfloat16, device=dev) + fvk.moe_shared_gate_combine_edge_bf16( + routed.data_ptr(), shared.data_ptr(), glog.data_ptr(), + out.data_ptr(), S, HID, torch.cuda.current_stream().cuda_stream) + torch.cuda.synchronize(dev) + + sgate = torch.sigmoid(glog.float()) + ref = (routed + shared.float() * sgate).to(torch.bfloat16) + + # Judged on absolute and relative distance together, which is what this + # output needs and what two simpler criteria each got wrong here: + # + # * a relative norm over the whole tensor passes almost anything, since + # one bf16 ULP is already 4e-3 relative; + # * a ULP distance is meaningless where the two sides straddle zero. The + # kernel fuses the multiply into the add, so it rounds once where the + # tensor chain rounds twice, and on rows where `routed` and the gated + # shared expert nearly cancel the results are 9.3e-10 against exactly + # 0 -- an enormous ULP distance for an absolute difference of nothing. + # + # So: an atol that treats a cancellation to zero as agreement, and an rtol + # of two bf16 ULP for everything else. Both are far tighter than a wrong + # gate, a wrong row index or a wrong dtype would produce. + assert torch.allclose(out.float(), ref.float(), rtol=8e-3, atol=1e-6), ( + "combine drifted: max abs " + f"{(out.float() - ref.float()).abs().max().item():.3e}") + + # And the bulk must be bit-identical, which is what catches a systematic + # drift that stays inside the tolerance above. + n_diff = int((out != ref).sum()) + assert n_diff <= out.numel() // 10000, \ + f"{n_diff} of {out.numel()} elements differ; expected a handful" + + +def test_reproducible(): + fvk = _load_fvk() + dev = "cuda:0" + S = 512 + g = torch.Generator(device=dev).manual_seed(29) + routed = torch.randn(S, HID, generator=g, device=dev, dtype=torch.float32) + shared = torch.randn(S, HID, generator=g, device=dev, + dtype=torch.bfloat16) + glog = torch.randn(S, 1, generator=g, device=dev, dtype=torch.bfloat16) + st = torch.cuda.current_stream().cuda_stream + + def run(): + out = torch.empty(S, HID, dtype=torch.bfloat16, device=dev) + fvk.moe_shared_gate_combine_edge_bf16( + routed.data_ptr(), shared.data_ptr(), glog.data_ptr(), + out.data_ptr(), S, HID, st) + torch.cuda.synchronize(dev) + return out + + a = run() + for _ in range(3): + assert torch.equal(a, run()) From 0c38caab1a2101fa52d61571f576cf0286c8ce40 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Fri, 31 Jul 2026 12:14:25 -0400 Subject: [PATCH 55/85] Read the conv1d input once instead of once per output The prefill conv gave one thread one (channel, token). With a kernel width of four that reads every input element four times, once for each output that needs it, and re-reads the channel's weight row for every token of the sequence: 134 MB of traffic to move 67 MB of data, and 4.2x off what that data implies. A thread now walks eight consecutive tokens of one channel with the window held in registers and shifted along, so each input is read once and the weight row once per thread. The silu is written the same way as the existing entry's rather than merely equivalently, because the two are meant to agree to the bit -- and do: exact at every length probed, max difference 0. 3.3x standalone, which lands within 1.3x of the traffic bound. The grid also blocks the sequence, so a single launch is no longer capped at 65535 tokens. TTFT 262.4 -> 221.6 ms at 1024, 457.1 -> 394.3 at 2048, 855.7 -> 753.0 at 4096. Prefill 4621 / 5194 / 5440 tok/s against vLLM's 3206 / 4138 / 4722. The existing entry stays; decode uses its own update variants and is untouched. --- CMakeLists.txt | 1 + csrc/bindings.cpp | 13 +++ csrc/kernels/causal_conv1d_rows_edge.cu | 109 ++++++++++++++++++ csrc/kernels/causal_conv1d_rows_edge.cuh | 42 +++++++ .../frontends/torch/_nexn2_rtx_forward.py | 18 ++- 5 files changed, 177 insertions(+), 6 deletions(-) create mode 100644 csrc/kernels/causal_conv1d_rows_edge.cu create mode 100644 csrc/kernels/causal_conv1d_rows_edge.cuh diff --git a/CMakeLists.txt b/CMakeLists.txt index 1d116aeb..3312f28b 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1552,6 +1552,7 @@ if(FLASHRT_ENABLE_QWEN35MOE_CORE) csrc/kernels/bf16_matvec_sm120.cu csrc/kernels/gdn_recurrent_seq_sm120.cu csrc/kernels/gdn_wy_prefill_edge.cu + csrc/kernels/causal_conv1d_rows_edge.cu csrc/kernels/act_fuse_sm120.cu csrc/kernels/moe_router_topk_sm120.cu csrc/kernels/moe_route_prefill_edge.cu diff --git a/csrc/bindings.cpp b/csrc/bindings.cpp index 789ca466..22d6d396 100644 --- a/csrc/bindings.cpp +++ b/csrc/bindings.cpp @@ -184,6 +184,7 @@ extern "C" int cutlass_int8_rowwise_bf16out_t64x128( #include "kernels/bf16_matvec_sm120.cuh" #include "kernels/gdn_recurrent_seq_sm120.cuh" #include "kernels/gdn_wy_prefill_edge.cuh" +#include "kernels/causal_conv1d_rows_edge.cuh" #include "kernels/act_fuse_sm120.cuh" #include "kernels/moe_router_topk_sm120.cuh" #include "kernels/moe_route_prefill_edge.cuh" @@ -5575,6 +5576,18 @@ PYBIND11_MODULE(flash_rt_kernels, m) { py::arg("group_off"), py::arg("sfa_off"), py::arg("n_experts"), py::arg("n_col"), py::arg("stream") = 0); + m.def("causal_conv1d_qwen36_rows_bf16", + [](uintptr_t x, uintptr_t w, uintptr_t bias, uintptr_t out, + int B, int S, int conv_dim, int k, bool apply_silu, + uintptr_t stream) { + flash_rt::kernels::causal_conv1d_qwen36_rows_bf16( + to_ptr(x), to_ptr(w), to_ptr(bias), to_ptr(out), + B, S, conv_dim, k, apply_silu, to_stream(stream)); + }, + py::arg("x"), py::arg("w"), py::arg("bias"), py::arg("out"), + py::arg("B"), py::arg("S"), py::arg("conv_dim"), py::arg("k"), + py::arg("apply_silu") = true, py::arg("stream") = 0); + m.def("gdn_wy_norm_pack_q_cumsum_edge_bf16", [](uintptr_t q, uintptr_t k, uintptr_t g, uintptr_t k_l2, uintptr_t q_pack, uintptr_t g_cumsum, int S, int num_k_heads, diff --git a/csrc/kernels/causal_conv1d_rows_edge.cu b/csrc/kernels/causal_conv1d_rows_edge.cu new file mode 100644 index 00000000..3603ecdc --- /dev/null +++ b/csrc/kernels/causal_conv1d_rows_edge.cu @@ -0,0 +1,109 @@ +#include "causal_conv1d_rows_edge.cuh" + +#include + +namespace flash_rt { +namespace kernels { + +namespace { + +constexpr int kMaxK = 4; +constexpr int kThreadsX = 256; // matches the existing prefill entry + +// Written the same way as the existing entry's, not merely equivalent to it: +// the two are meant to agree to the bit. +__device__ __forceinline__ float rows_silu(float v) { + return v / (1.0f + __expf(-v)); +} + +// One thread, one channel, `kRows` consecutive tokens. The k-1 inputs a token +// shares with the next are kept in registers and shifted along, so the reads +// are one element per output rather than k. +template +__global__ void causal_conv1d_rows_kernel( + const __nv_bfloat16* __restrict__ x, + const __nv_bfloat16* __restrict__ w, + const __nv_bfloat16* __restrict__ bias, + __nv_bfloat16* __restrict__ out, + int B, int S, int conv_dim, int k, + bool apply_silu) +{ + const int c = blockIdx.x * kThreadsX + threadIdx.x; + if (c >= conv_dim) return; + const int s0 = blockIdx.y * kRows; + if (s0 >= S) return; + const int b = blockIdx.z; + + float wv[kMaxK]; + #pragma unroll + for (int i = 0; i < kMaxK; ++i) { + wv[i] = (i < k) ? static_cast(w[c * k + i]) : 0.0f; + } + const float b0 = (bias != nullptr) ? static_cast(bias[c]) : 0.0f; + + const size_t base = static_cast(b) * S * conv_dim + c; + + // win[j] holds x[s0 - (k-1) + j], the window the first output needs. Reads + // before the start of the sequence are zero, which is what the causal + // convolution means there. + float win[kMaxK]; + #pragma unroll + for (int j = 0; j < kMaxK; ++j) { + const int t = s0 - (k - 1) + j; + win[j] = (j < k && t >= 0 && t < S) + ? static_cast(x[base + static_cast(t) * conv_dim]) + : 0.0f; + } + + #pragma unroll + for (int r = 0; r < kRows; ++r) { + const int s = s0 + r; + if (s >= S) break; + if (r > 0) { + // Shift by one and pull in the token that just became current. + #pragma unroll + for (int j = 0; j < kMaxK - 1; ++j) win[j] = win[j + 1]; + win[k - 1] = static_cast( + x[base + static_cast(s) * conv_dim]); + } + float acc = b0; + #pragma unroll + for (int i = 0; i < kMaxK; ++i) { + if (i < k) acc = fmaf(win[i], wv[i], acc); + } + if (apply_silu) acc = rows_silu(acc); + out[base + static_cast(s) * conv_dim] = __float2bfloat16(acc); + } +} + +} // namespace + +void causal_conv1d_qwen36_rows_bf16( + const void* x, + const void* w, + const void* bias, + void* out, + int B, + int S, + int conv_dim, + int k, + bool apply_silu, + cudaStream_t stream) +{ + if (B <= 0 || S <= 0 || conv_dim <= 0 || k <= 0 || k > kMaxK) return; + + constexpr int kRows = 8; + const dim3 block(kThreadsX); + const dim3 grid((conv_dim + kThreadsX - 1) / kThreadsX, + (S + kRows - 1) / kRows, + B); + causal_conv1d_rows_kernel<<>>( + reinterpret_cast(x), + reinterpret_cast(w), + reinterpret_cast(bias), + reinterpret_cast<__nv_bfloat16*>(out), + B, S, conv_dim, k, apply_silu); +} + +} // namespace kernels +} // namespace flash_rt diff --git a/csrc/kernels/causal_conv1d_rows_edge.cuh b/csrc/kernels/causal_conv1d_rows_edge.cuh new file mode 100644 index 00000000..659f1570 --- /dev/null +++ b/csrc/kernels/causal_conv1d_rows_edge.cuh @@ -0,0 +1,42 @@ +#pragma once + +#include + +namespace flash_rt { +namespace kernels { + +// Causal depthwise conv1d over a whole prompt, several output tokens per +// thread. +// +// The existing prefill entry gives one thread one (channel, token), so each +// input element is fetched once for every output that needs it -- k times -- +// and the channel's weight row is fetched once per token. At the Qwen3.6 +// prefill shape that is 134 MB of reads for 67 MB of data, and the kernel +// measures about four times off what its traffic implies. +// +// Here a thread walks `rows` consecutive tokens of one channel, holding the +// last k inputs in registers, so each input is read once and the weight row +// once per thread rather than once per token. +// +// x (B, S, conv_dim) bf16 +// w (conv_dim, k) bf16 +// bias (conv_dim,) bf16 or null +// out (B, S, conv_dim) bf16 +// +// k must be at most 4. Layout, causality and the optional silu match the +// existing entry exactly; this is the same function computed with less +// traffic. +void causal_conv1d_qwen36_rows_bf16( + const void* x, + const void* w, + const void* bias, + void* out, + int B, + int S, + int conv_dim, + int k, + bool apply_silu, + cudaStream_t stream); + +} // namespace kernels +} // namespace flash_rt diff --git a/flash_rt/frontends/torch/_nexn2_rtx_forward.py b/flash_rt/frontends/torch/_nexn2_rtx_forward.py index c46b32fd..5c8c9e3f 100644 --- a/flash_rt/frontends/torch/_nexn2_rtx_forward.py +++ b/flash_rt/frontends/torch/_nexn2_rtx_forward.py @@ -459,6 +459,14 @@ def _gdn_layer(h, ld, fvk, device, eps, cap=None, rank=None, # Same (B, S, conv_dim) layout the decode update kernel uses; no bias. For a # chunked block, prepend the previous block's last KS-1 inputs (conv_hist) # so the block's first outputs see the right history, then drop them. + # The row-blocked entry walks several tokens per thread with the window in + # registers, so each input is read once instead of once per output that + # needs it. Bit-identical to the per-token entry (probe: exact at every + # length measured) and 3.3x, which puts it within 1.3x of what its traffic + # implies rather than 4.2x off. It also lifts the gridDim.y ceiling that + # capped a single launch at 65535 tokens. + _conv = getattr(fvk, 'causal_conv1d_qwen36_rows_bf16', + fvk.causal_conv1d_qwen36_bf16) convw_k = convw.reshape(CONV, KS).contiguous() if conv_hist is not None: hist = conv_hist[0].transpose(0, 1).reshape(1, KS - 1, CONV) @@ -466,15 +474,13 @@ def _gdn_layer(h, ld, fvk, device, eps, cap=None, rank=None, [hist.to(mixed.dtype), mixed], dim=1).contiguous() Se = mixed_ext.shape[1] xc_ext = torch.empty(B, Se, CONV, dtype=torch.bfloat16, device=device) - fvk.causal_conv1d_qwen36_bf16( - mixed_ext.data_ptr(), convw_k.data_ptr(), 0, - xc_ext.data_ptr(), B, Se, CONV, KS, True, _cs()) + _conv(mixed_ext.data_ptr(), convw_k.data_ptr(), 0, + xc_ext.data_ptr(), B, Se, CONV, KS, True, _cs()) xc = xc_ext[:, KS - 1:, :].contiguous() else: xc = torch.empty(B, S, CONV, dtype=torch.bfloat16, device=device) - fvk.causal_conv1d_qwen36_bf16( - mixed.contiguous().data_ptr(), convw_k.data_ptr(), 0, - xc.data_ptr(), B, S, CONV, KS, True, _cs()) + _conv(mixed.contiguous().data_ptr(), convw_k.data_ptr(), 0, + xc.data_ptr(), B, S, CONV, KS, True, _cs()) # split conv output + broadcast q/k 16 -> 32 heads in one fvk kernel. xc_bf = xc.reshape(B * S, CONV).contiguous() qb = torch.empty(B, S, NV, HK, dtype=torch.bfloat16, device=device) From 0dde39f07f8f512410bf4731212f400ee6342b02 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Fri, 31 Jul 2026 12:18:37 -0400 Subject: [PATCH 56/85] Cover the row-blocked conv1d Compared exactly against the per-token entry, since the two compute the same convolution and write the silu the same way -- a difference there is a bug and not a reduction order. Both are also compared against a torch reference that shares none of their code, so a common misreading of the causal edge cannot pass by the two agreeing with each other. Lengths that are not multiples of the eight rows a thread walks are included, because that tail is where a row-blocked kernel goes wrong, and one case runs past 65535 tokens -- the ceiling the per-token grid had -- to check that shape is reachable at all. --- tests/test_causal_conv1d_rows_edge.py | 111 ++++++++++++++++++++++++++ 1 file changed, 111 insertions(+) create mode 100644 tests/test_causal_conv1d_rows_edge.py diff --git a/tests/test_causal_conv1d_rows_edge.py b/tests/test_causal_conv1d_rows_edge.py new file mode 100644 index 00000000..81531b7c --- /dev/null +++ b/tests/test_causal_conv1d_rows_edge.py @@ -0,0 +1,111 @@ +"""Equivalence test for the row-blocked causal conv1d. + +The row-blocked entry computes the same convolution as the per-token one, and +writes the silu the same way rather than merely equivalently, so this compares +them exactly: any difference is a bug, not a reduction order. Both are checked +against a direct torch reference as well, so a shared misreading of the layout +or the causal edge does not pass by agreeing with itself. + +Lengths that are not multiples of the eight rows a thread walks are included, +since that tail is where a row-blocked kernel goes wrong. +""" + +import pytest +import torch +import torch.nn.functional as F + +CONV, K = 8192, 4 + + +def _load_fvk(): + if not torch.cuda.is_available(): + pytest.skip("CUDA is required for the conv1d test") + try: + from flash_rt import flash_rt_kernels as fvk + except Exception as exc: # pragma: no cover - environmental + pytest.skip(f"flash_rt_kernels is not built: {exc}") + if not hasattr(fvk, "causal_conv1d_qwen36_rows_bf16"): + pytest.skip("causal_conv1d_qwen36_rows_bf16 not in this build") + return fvk + + +def _torch_ref(x, w, k, silu=True): + """y[s, c] = sum_i x[s + i - (k-1), c] * w[c, i], zero before the start.""" + S, C = x.shape[1], x.shape[2] + xt = x[0].t().unsqueeze(0) # (1, C, S) + xp = F.pad(xt.float(), (k - 1, 0)) + y = F.conv1d(xp, w.float().unsqueeze(1), groups=C) + y = y[0].t().unsqueeze(0) # (1, S, C) + if silu: + y = y / (1.0 + torch.exp(-y)) + return y.to(torch.bfloat16) + + +@pytest.mark.parametrize("S", [1, 7, 64, 129, 512, 2048, 4097]) +def test_rows_matches_per_token_entry_exactly(S): + fvk = _load_fvk() + dev = "cuda:0" + g = torch.Generator(device=dev).manual_seed(S) + x = torch.randn(1, S, CONV, generator=g, device=dev, dtype=torch.bfloat16) + w = torch.randn(CONV, K, generator=g, device=dev, dtype=torch.bfloat16) + st = torch.cuda.current_stream().cuda_stream + + a = torch.empty(1, S, CONV, dtype=torch.bfloat16, device=dev) + b = torch.empty_like(a) + fvk.causal_conv1d_qwen36_bf16( + x.data_ptr(), w.data_ptr(), 0, a.data_ptr(), 1, S, CONV, K, True, st) + fvk.causal_conv1d_qwen36_rows_bf16( + x.data_ptr(), w.data_ptr(), 0, b.data_ptr(), 1, S, CONV, K, True, st) + torch.cuda.synchronize(dev) + + assert torch.equal(a, b), ( + f"{int((a != b).sum())} of {a.numel()} elements differ from the " + "per-token entry") + + # And both against a reference that shares none of their code, so a common + # misreading of the causal edge cannot pass. + ref = _torch_ref(x, w, K) + rel = ((b.float() - ref.float()).norm() + / ref.float().norm().clamp_min(1e-9)).item() + assert rel < 1e-2, f"both kernels disagree with the reference by {rel:.3e}" + + +@pytest.mark.parametrize("S", [63, 2048]) +def test_no_silu_variant_matches(S): + fvk = _load_fvk() + dev = "cuda:0" + g = torch.Generator(device=dev).manual_seed(S + 1) + x = torch.randn(1, S, CONV, generator=g, device=dev, dtype=torch.bfloat16) + w = torch.randn(CONV, K, generator=g, device=dev, dtype=torch.bfloat16) + st = torch.cuda.current_stream().cuda_stream + + a = torch.empty(1, S, CONV, dtype=torch.bfloat16, device=dev) + b = torch.empty_like(a) + fvk.causal_conv1d_qwen36_bf16( + x.data_ptr(), w.data_ptr(), 0, a.data_ptr(), 1, S, CONV, K, False, st) + fvk.causal_conv1d_qwen36_rows_bf16( + x.data_ptr(), w.data_ptr(), 0, b.data_ptr(), 1, S, CONV, K, False, st) + torch.cuda.synchronize(dev) + assert torch.equal(a, b) + + +def test_beyond_the_old_grid_ceiling(): + """The per-token entry launches one grid row per token, so a single call + stopped at 65535. The row-blocked grid divides the sequence, so this shape + is reachable at all -- which is the point, not the speed.""" + fvk = _load_fvk() + dev = "cuda:0" + S, C = 70000, 512 # narrow, so the buffer stays small + g = torch.Generator(device=dev).manual_seed(2) + x = torch.randn(1, S, C, generator=g, device=dev, dtype=torch.bfloat16) + w = torch.randn(C, K, generator=g, device=dev, dtype=torch.bfloat16) + out = torch.empty(1, S, C, dtype=torch.bfloat16, device=dev) + fvk.causal_conv1d_qwen36_rows_bf16( + x.data_ptr(), w.data_ptr(), 0, out.data_ptr(), 1, S, C, K, True, + torch.cuda.current_stream().cuda_stream) + torch.cuda.synchronize(dev) + assert torch.isfinite(out.float()).all() + ref = _torch_ref(x, w, K) + rel = ((out.float() - ref.float()).norm() + / ref.float().norm().clamp_min(1e-9)).item() + assert rel < 1e-2, f"past the old ceiling the result drifted {rel:.3e}" From a61e8744ce05488ec1464ecfc19d321bcde7a40f Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Fri, 31 Jul 2026 12:37:15 -0400 Subject: [PATCH 57/85] Let the chunked conv read its history instead of concatenating it A chunked block needs the previous block's last three inputs so its first outputs see the right context. It was getting them by prepending them to the whole block of activations and slicing the result back off -- two copies of the block per layer, to supply three tokens. At 32768 that was 691 ms in and its batched copy, 8% of the prefill. The conv now takes an optional history pointer and reads it where the window reaches before the start. The layout is the one the decode conv state already carries, channel-major with the newest last, so the transpose goes with the concatenation. A null pointer means the sequence starts here, which is the single-pass case and is unchanged. 32768 goes 7688.4 -> 7370.8 ms, 10240 1963.8 -> 1954.3 -- the difference between them is how many block boundaries there are to pay for. Chunked against single-pass: cos 0.9983 and 0.9979 at blocks of 512 and 1024, same seeded token. The conv itself stays bit-identical to the per-token entry. --- csrc/bindings.cpp | 12 ++++++ csrc/kernels/causal_conv1d_rows_edge.cu | 41 ++++++++++++++++--- csrc/kernels/causal_conv1d_rows_edge.cuh | 23 +++++++++++ .../frontends/torch/_nexn2_rtx_forward.py | 15 ++++++- 4 files changed, 83 insertions(+), 8 deletions(-) diff --git a/csrc/bindings.cpp b/csrc/bindings.cpp index 22d6d396..b9d10e05 100644 --- a/csrc/bindings.cpp +++ b/csrc/bindings.cpp @@ -5576,6 +5576,18 @@ PYBIND11_MODULE(flash_rt_kernels, m) { py::arg("group_off"), py::arg("sfa_off"), py::arg("n_experts"), py::arg("n_col"), py::arg("stream") = 0); + m.def("causal_conv1d_qwen36_rows_hist_bf16", + [](uintptr_t x, uintptr_t w, uintptr_t bias, uintptr_t hist, + uintptr_t out, int B, int S, int conv_dim, int k, bool apply_silu, + uintptr_t stream) { + flash_rt::kernels::causal_conv1d_qwen36_rows_hist_bf16( + to_ptr(x), to_ptr(w), to_ptr(bias), to_ptr(hist), to_ptr(out), + B, S, conv_dim, k, apply_silu, to_stream(stream)); + }, + py::arg("x"), py::arg("w"), py::arg("bias"), py::arg("hist"), + py::arg("out"), py::arg("B"), py::arg("S"), py::arg("conv_dim"), + py::arg("k"), py::arg("apply_silu") = true, py::arg("stream") = 0); + m.def("causal_conv1d_qwen36_rows_bf16", [](uintptr_t x, uintptr_t w, uintptr_t bias, uintptr_t out, int B, int S, int conv_dim, int k, bool apply_silu, diff --git a/csrc/kernels/causal_conv1d_rows_edge.cu b/csrc/kernels/causal_conv1d_rows_edge.cu index 3603ecdc..3e92e13d 100644 --- a/csrc/kernels/causal_conv1d_rows_edge.cu +++ b/csrc/kernels/causal_conv1d_rows_edge.cu @@ -24,6 +24,7 @@ __global__ void causal_conv1d_rows_kernel( const __nv_bfloat16* __restrict__ x, const __nv_bfloat16* __restrict__ w, const __nv_bfloat16* __restrict__ bias, + const __nv_bfloat16* __restrict__ hist, __nv_bfloat16* __restrict__ out, int B, int S, int conv_dim, int k, bool apply_silu) @@ -43,16 +44,24 @@ __global__ void causal_conv1d_rows_kernel( const size_t base = static_cast(b) * S * conv_dim + c; - // win[j] holds x[s0 - (k-1) + j], the window the first output needs. Reads - // before the start of the sequence are zero, which is what the causal - // convolution means there. + // win[j] holds x[s0 - (k-1) + j], the window the first output needs. Before + // the start of this block that is the previous block's trailing inputs when + // there are any, and zero when the sequence itself starts here. float win[kMaxK]; #pragma unroll for (int j = 0; j < kMaxK; ++j) { const int t = s0 - (k - 1) + j; - win[j] = (j < k && t >= 0 && t < S) - ? static_cast(x[base + static_cast(t) * conv_dim]) - : 0.0f; + if (j >= k) { win[j] = 0.0f; continue; } + if (t >= 0 && t < S) { + win[j] = static_cast(x[base + static_cast(t) * conv_dim]); + } else if (t < 0 && hist != nullptr) { + // hist is (B, conv_dim, k-1), newest last: t == -1 is the final column. + const int hj = t + (k - 1); + win[j] = static_cast( + hist[(static_cast(b) * conv_dim + c) * (k - 1) + hj]); + } else { + win[j] = 0.0f; + } } #pragma unroll @@ -92,6 +101,25 @@ void causal_conv1d_qwen36_rows_bf16( { if (B <= 0 || S <= 0 || conv_dim <= 0 || k <= 0 || k > kMaxK) return; + causal_conv1d_qwen36_rows_hist_bf16(x, w, bias, nullptr, out, B, S, + conv_dim, k, apply_silu, stream); +} + +void causal_conv1d_qwen36_rows_hist_bf16( + const void* x, + const void* w, + const void* bias, + const void* hist, + void* out, + int B, + int S, + int conv_dim, + int k, + bool apply_silu, + cudaStream_t stream) +{ + if (B <= 0 || S <= 0 || conv_dim <= 0 || k <= 0 || k > kMaxK) return; + constexpr int kRows = 8; const dim3 block(kThreadsX); const dim3 grid((conv_dim + kThreadsX - 1) / kThreadsX, @@ -101,6 +129,7 @@ void causal_conv1d_qwen36_rows_bf16( reinterpret_cast(x), reinterpret_cast(w), reinterpret_cast(bias), + reinterpret_cast(hist), reinterpret_cast<__nv_bfloat16*>(out), B, S, conv_dim, k, apply_silu); } diff --git a/csrc/kernels/causal_conv1d_rows_edge.cuh b/csrc/kernels/causal_conv1d_rows_edge.cuh index 659f1570..82788f32 100644 --- a/csrc/kernels/causal_conv1d_rows_edge.cuh +++ b/csrc/kernels/causal_conv1d_rows_edge.cuh @@ -22,6 +22,15 @@ namespace kernels { // w (conv_dim, k) bf16 // bias (conv_dim,) bf16 or null // out (B, S, conv_dim) bf16 +// hist (B, conv_dim, k-1) bf16 or null -- the previous block's last k-1 +// inputs, channel-major with the newest last, which is the layout the +// decode conv state already carries. Null means the sequence starts +// here and the reads before it are zero. +// +// `hist` is what lets a chunked prefill stop concatenating. Prepending the +// history to the activations and slicing the result back off copies the whole +// block twice per layer -- 691 ms of a 32768-token prefill, in `cat` and its +// batched copy -- to supply three tokens of context. // // k must be at most 4. Layout, causality and the optional silu match the // existing entry exactly; this is the same function computed with less @@ -38,5 +47,19 @@ void causal_conv1d_qwen36_rows_bf16( bool apply_silu, cudaStream_t stream); +// Same, continuing from a previous block's trailing inputs. +void causal_conv1d_qwen36_rows_hist_bf16( + const void* x, + const void* w, + const void* bias, + const void* hist, + void* out, + int B, + int S, + int conv_dim, + int k, + bool apply_silu, + cudaStream_t stream); + } // namespace kernels } // namespace flash_rt diff --git a/flash_rt/frontends/torch/_nexn2_rtx_forward.py b/flash_rt/frontends/torch/_nexn2_rtx_forward.py index 5c8c9e3f..4667a7c7 100644 --- a/flash_rt/frontends/torch/_nexn2_rtx_forward.py +++ b/flash_rt/frontends/torch/_nexn2_rtx_forward.py @@ -468,7 +468,19 @@ def _gdn_layer(h, ld, fvk, device, eps, cap=None, rank=None, _conv = getattr(fvk, 'causal_conv1d_qwen36_rows_bf16', fvk.causal_conv1d_qwen36_bf16) convw_k = convw.reshape(CONV, KS).contiguous() - if conv_hist is not None: + xc = torch.empty(B, S, CONV, dtype=torch.bfloat16, device=device) + _hist_conv = getattr(fvk, 'causal_conv1d_qwen36_rows_hist_bf16', None) + if conv_hist is not None and _hist_conv is not None: + # The conv reads the previous block's trailing inputs where it needs + # them. Prepending them to the activations instead meant concatenating + # and then slicing the whole block back off -- two copies of it per + # layer, 691 ms of a 32768-token prefill, to supply three tokens. + # conv_hist is already (1, CONV, KS-1), newest last, which is the + # layout the kernel reads, so the transpose goes with the copy. + _hist_conv(mixed.contiguous().data_ptr(), convw_k.data_ptr(), 0, + conv_hist.contiguous().data_ptr(), xc.data_ptr(), + B, S, CONV, KS, True, _cs()) + elif conv_hist is not None: hist = conv_hist[0].transpose(0, 1).reshape(1, KS - 1, CONV) mixed_ext = torch.cat( [hist.to(mixed.dtype), mixed], dim=1).contiguous() @@ -478,7 +490,6 @@ def _gdn_layer(h, ld, fvk, device, eps, cap=None, rank=None, xc_ext.data_ptr(), B, Se, CONV, KS, True, _cs()) xc = xc_ext[:, KS - 1:, :].contiguous() else: - xc = torch.empty(B, S, CONV, dtype=torch.bfloat16, device=device) _conv(mixed.contiguous().data_ptr(), convw_k.data_ptr(), 0, xc.data_ptr(), B, S, CONV, KS, True, _cs()) # split conv output + broadcast q/k 16 -> 32 heads in one fvk kernel. From bdf5ddbd066a1b8f70022bc7db44367ddd375536 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Fri, 31 Jul 2026 12:46:32 -0400 Subject: [PATCH 58/85] Give the gate-and-quantise a lane per scale-factor group The block-per-row form handed 256 threads a row of 512 values -- two elements each -- behind three barriers and three passes over shared memory. A block read two kilobytes and then waited, which measured 2.9x off what that traffic implies. A lane now owns one 16-element scale-factor group: it reads its own gate and up values, gates them, takes its own maximum and packs its own eight bytes. Nothing is shared, so there is no shared memory and no barrier, and a lane has sixteen values in flight where it had two. Same arithmetic in the same order, and the output is identical byte for byte -- packed data and scale factors both, over a routing layout with real group boundaries, since the scale-factor offsets depend on where each expert starts. 2.72x at the shape prefill issues: 1.778 -> 0.653 ms at 65536 slots, against a traffic bound of 0.627. That is 1.04x of the bound, so there is nothing left here. TTFT 394.3 -> 388.0 ms at 2048, 753.0 -> 734.7 at 4096, 217.6 at 1024. The block-per-row entry stays; this is a second one. --- csrc/bindings.cpp | 13 +++ csrc/kernels/quantize.cu | 89 +++++++++++++++++++ csrc/kernels/quantize.cuh | 8 ++ .../frontends/torch/_nexn2_rtx_forward.py | 8 +- 4 files changed, 117 insertions(+), 1 deletion(-) diff --git a/csrc/bindings.cpp b/csrc/bindings.cpp index b9d10e05..893d3f29 100644 --- a/csrc/bindings.cpp +++ b/csrc/bindings.cpp @@ -1106,6 +1106,19 @@ PYBIND11_MODULE(flash_rt_kernels, m) { py::arg("sfa_off"), py::arg("out_packed"), py::arg("out_sf"), py::arg("slots"), py::arg("inter"), py::arg("stream") = 0); + m.def("moe_grouped_silu_quant_nvfp4_warp_bf16", + [](uintptr_t merged, uintptr_t expert_of_row, uintptr_t group_off, + uintptr_t sfa_off, uintptr_t out_packed, uintptr_t out_sf, + int slots, int inter, uintptr_t stream) -> int { + return moe_grouped_silu_quant_nvfp4_warp_bf16( + to_ptr(merged), to_ptr(expert_of_row), to_ptr(group_off), + to_ptr(sfa_off), to_ptr(out_packed), to_ptr(out_sf), + slots, inter, to_stream(stream)); + }, + py::arg("merged"), py::arg("expert_of_row"), py::arg("group_off"), + py::arg("sfa_off"), py::arg("out_packed"), py::arg("out_sf"), + py::arg("slots"), py::arg("inter"), py::arg("stream") = 0); + m.def("moe_grouped_quant_nvfp4_bf16", [](uintptr_t A, uintptr_t expert_of_row, uintptr_t group_off, uintptr_t sfa_off, uintptr_t src_row, diff --git a/csrc/kernels/quantize.cu b/csrc/kernels/quantize.cu index 9719941e..98acbb12 100644 --- a/csrc/kernels/quantize.cu +++ b/csrc/kernels/quantize.cu @@ -2981,6 +2981,95 @@ __global__ void moe_grouped_silu_quant_nvfp4_kernel( } } +// Warp-per-row form of the same thing. +// +// The block-per-row kernel above gives 256 threads a row of 512 values -- two +// elements each -- behind three barriers and three passes over shared memory, +// so a block reads two kilobytes and then waits. Measured 2.9x off what that +// traffic implies. +// +// Here a warp owns a row and a lane owns one 16-element scale-factor group: +// it reads its own sixteen gate and up values as vectors, gates them, takes +// its own maximum and packs its own eight bytes. Nothing is shared, so there +// are no barriers and no shared memory at all, and each lane has sixteen +// values in flight instead of two. +// +// The arithmetic is the same in the same order, so the output is identical. +__global__ void moe_grouped_silu_quant_nvfp4_warp_kernel( + const __nv_bfloat16* __restrict__ merged, + const int* __restrict__ expert_of_row, + const int* __restrict__ group_off, + const int* __restrict__ sfa_off, + uint8_t* __restrict__ fp4_data, + uint8_t* __restrict__ scale_factors, + int slots, int inter, int num_blocks, int n_col_blocks) +{ + const int warp_in_blk = threadIdx.x >> 5; + const int lane = threadIdx.x & 31; + const int row = blockIdx.x * (blockDim.x >> 5) + warp_in_blk; + if (row >= slots) return; + + const int e = expert_of_row[row]; + const int local = row - group_off[e]; + const __nv_bfloat16* g_in = merged + (size_t)row * 2 * inter; + const __nv_bfloat16* u_in = g_in + inter; + uint8_t* row_fp4 = fp4_data + (size_t)row * inter / 2; + uint8_t* sf_base = scale_factors + sfa_off[e]; + const int rb = local / 128, ri = local % 128; + + for (int b = lane; b < num_blocks; b += 32) { + float gated[16]; + const int base = b * 16; + #pragma unroll + for (int j = 0; j < 16; ++j) { + const float gv = __bfloat162float(g_in[base + j]); + const float uv = __bfloat162float(u_in[base + j]); + gated[j] = __bfloat162float( + __float2bfloat16_rn(gv / (1.0f + __expf(-gv)) * uv)); + } + float a = 0.0f; + #pragma unroll + for (int j = 0; j < 16; ++j) a = fmaxf(a, fabsf(gated[j])); + + const uint8_t ue = float_to_ue4m3_ceil(a * (1.0f / 6.0f)); + sf_base[(rb * n_col_blocks + (b >> 2)) * 512 + (ri % 32) * 16 + + (ri / 32) * 4 + (b & 3)] = ue; + + const float sc = ue4m3_to_float(ue); + const float inv = (sc > 0.0f) ? (1.0f / sc) : 0.0f; + uint8_t* out8 = row_fp4 + (size_t)b * 8; + #pragma unroll + for (int p = 0; p < 8; ++p) { + out8[p] = (uint8_t)((float_to_fp4_e2m1(gated[2 * p + 1] * inv) << 4) + | (float_to_fp4_e2m1(gated[2 * p] * inv) & 0x0F)); + } + } +} + +int moe_grouped_silu_quant_nvfp4_warp_bf16( + const void* merged, const void* expert_of_row, const void* group_off, + const void* sfa_off, void* out_packed, void* out_sf, + int slots, int inter, cudaStream_t stream) +{ + if (!merged || !expert_of_row || !group_off || !sfa_off || !out_packed + || !out_sf) return 1; + if (slots <= 0 || inter <= 0 || (inter & 15) != 0) return 2; + const int num_blocks = inter / 16; + const int n_col_blocks = (num_blocks + 3) / 4; + constexpr int kThreads = 256; + const int rows_per_block = kThreads / 32; + const int grid = (slots + rows_per_block - 1) / rows_per_block; + moe_grouped_silu_quant_nvfp4_warp_kernel<<>>( + reinterpret_cast(merged), + reinterpret_cast(expert_of_row), + reinterpret_cast(group_off), + reinterpret_cast(sfa_off), + reinterpret_cast(out_packed), + reinterpret_cast(out_sf), + slots, inter, num_blocks, n_col_blocks); + return 0; +} + int moe_grouped_silu_quant_nvfp4_bf16( const void* merged, const void* expert_of_row, const void* group_off, const void* sfa_off, void* out_packed, void* out_sf, diff --git a/csrc/kernels/quantize.cuh b/csrc/kernels/quantize.cuh index 471987a5..88c69118 100644 --- a/csrc/kernels/quantize.cuh +++ b/csrc/kernels/quantize.cuh @@ -352,3 +352,11 @@ int moe_grouped_silu_quant_nvfp4_bf16( const void* merged, const void* expert_of_row, const void* group_off, const void* sfa_off, void* out_packed, void* out_sf, int slots, int inter, cudaStream_t stream); + +// Warp-per-row form of the above: a lane owns one 16-element scale-factor +// group and keeps it in registers, so there is no shared memory and no +// barrier. Same arithmetic in the same order, so the output is identical. +int moe_grouped_silu_quant_nvfp4_warp_bf16( + const void* merged, const void* expert_of_row, const void* group_off, + const void* sfa_off, void* out_packed, void* out_sf, + int slots, int inter, cudaStream_t stream); diff --git a/flash_rt/frontends/torch/_nexn2_rtx_forward.py b/flash_rt/frontends/torch/_nexn2_rtx_forward.py index 4667a7c7..672999aa 100644 --- a/flash_rt/frontends/torch/_nexn2_rtx_forward.py +++ b/flash_rt/frontends/torch/_nexn2_rtx_forward.py @@ -1233,7 +1233,13 @@ def project(A, k, n, w_p, w_s, alpha, out, gate=False, perm=None): # A is the merged (slots, 2k) gate/up output: gate it and quantise # in one pass rather than slicing two strided halves out of it, # copying both, gating into a third buffer and reading that back. - rc = fvk.moe_grouped_silu_quant_nvfp4_bf16( + # Warp-per-row: a lane owns one scale-factor group and keeps + # it in registers, so there is no shared memory and no barrier. + # Byte-identical to the block-per-row form and 2.7x at the shape + # prefill issues, which puts it at 1.04x of its traffic bound. + _sq = getattr(fvk, 'moe_grouped_silu_quant_nvfp4_warp_bf16', + fvk.moe_grouped_silu_quant_nvfp4_bf16) + rc = _sq( A.data_ptr(), se.data_ptr(), group_off.data_ptr(), sfa_off.data_ptr(), packed.data_ptr(), sfa.data_ptr(), slots, k, _cs()) From 772df15d6041142b941b8175206d41babed53d09 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Sat, 1 Aug 2026 04:56:34 -0400 Subject: [PATCH 59/85] Narrow the grouped GEMM's N tile to the shape a prefill routes A prefill sends about sixty-four rows to the average expert, spread over 256 groups of unequal size. With an N tile of 256 the scheduler has 1024 blocks to balance those groups across twenty SMs; with 128 it has twice that. Paired against the wider tile on the same machine state, three runs each: 377.9 / 378.2 / 377.9 ms at 2048 tokens against 429.6 / 386.3 / 387.6. The worst run of the narrow tile beats the best run of the wide one, and the spread falls from 43 ms to 0.3 -- which is the load balance showing up directly. The cluster shape was swept first and does not move this. It is a runtime argument, so (1,1) (2,1) (1,2) (2,2) (4,1) (1,4) were measured without recompiling and the best two alternated three times each; the within-pair differences (+1.4%, -0.4%, +0.8%) came out smaller than the drift between runs of one setting. Left as a knob with that result written next to it rather than as a knob to try again. --- .../fp4/cutlass_nvfp4_moe_grouped_sm100.cu | 26 +++++++++++++++++-- 1 file changed, 24 insertions(+), 2 deletions(-) diff --git a/csrc/gemm/fp4/cutlass_nvfp4_moe_grouped_sm100.cu b/csrc/gemm/fp4/cutlass_nvfp4_moe_grouped_sm100.cu index 5d63affb..49d78997 100644 --- a/csrc/gemm/fp4/cutlass_nvfp4_moe_grouped_sm100.cu +++ b/csrc/gemm/fp4/cutlass_nvfp4_moe_grouped_sm100.cu @@ -3,6 +3,8 @@ // Grouped NVFP4 block-scaled GEMM for sm_100-class Blackwell. See header. #include "gemm/fp4/cutlass_nvfp4_moe_grouped_sm100.cuh" +#include +#include #include @@ -47,7 +49,13 @@ using ProblemShape = cutlass::gemm::GroupProblemShape>; using ArchTag = cutlass::arch::Sm100; using OperatorClass = cutlass::arch::OpClassBlockScaledTensorOp; using ClusterShape = Shape; -using MmaTileShape = Shape<_128, _256, _256>; +// N of 128 rather than 256, measured. A prefill routes about sixty-four rows +// to the average expert across 256 groups of unequal size, and the narrower +// tile gives the scheduler twice as many blocks to balance them across twenty +// SMs. Paired against N=256 on the same machine state: 377.9/378.2/377.9 ms +// against 429.6/386.3/387.6 -- the worst run of this tile beats the best run +// of the other, and the spread goes from 43 ms to 0.3. +using MmaTileShape = Shape<_128, _128, _256>; using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< @@ -241,7 +249,21 @@ int moe_grouped_gemm_nvfp4_sm100_bf16out( hw_info.device_id = 0; hw_info.sm_count = cutlass::KernelHardwareInfo::query_device_multiprocessor_count(0); - hw_info.cluster_shape = dim3(1, 1, 1); + // The cluster is a runtime shape, so it was swept without recompiling: + // (1,1) (2,1) (1,2) (2,2) (4,1) (1,4) at the prefill shape, then the best + // two alternated three times each. The within-pair differences (+1.4%, + // -0.4%, +0.8%) came out smaller than the drift between runs (427 to 385 ms + // for the same setting), so the cluster shape does not move this. Left as a + // knob with the result written down rather than as a knob to try again. + static const dim3 kCluster = [] { + const char* v = std::getenv("FLASHRT_MOE_GROUPED_CLUSTER"); + int x = 1, y = 1; + if (v && std::sscanf(v, "%d,%d", &x, &y) == 2 && x >= 1 && y >= 1) { + return dim3(x, y, 1); + } + return dim3(1, 1, 1); + }(); + hw_info.cluster_shape = kCluster; hw_info.cluster_shape_fallback = dim3(1, 1, 1); typename Gemm::Arguments args_proto{}; From 816ddcae8097c000b61cdceecb05249c5e14157c Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Sat, 1 Aug 2026 04:56:49 -0400 Subject: [PATCH 60/85] Stop letting a timing loop choose the dense GEMM's algorithm The cuBLASLt wrapper picks its algorithm by timing eight candidates at first use. Timing is noisy, so different processes pick different algorithms, and different algorithms reduce in different orders -- which makes the model non-deterministic across processes. One binary gave the golden prefix 16/16 three times and 14/16 three times, flipping between exactly two token streams. Asking the wrapper for one candidate takes the heuristic's own choice instead. Five of five processes then agree, and six of six pass the golden gate. This is not a speed trade in the direction it looks. At 1024 tokens the timed pick is worth 1.6% warm -- 213.6 against 217.1 ms over three runs each -- and costs 25% of the cold time, about 1020 against 770 ms, because the timing loop runs inside the first call. A faster first token and a deterministic model, for 1.6% of the warm path. Set from this frontend rather than in the kernel, whose default is shared. The flip needed two things and neither is a defect alone: FA2 moved token 14 close to a decision boundary (it is bit-reproducible within a process, and with the timed pick left on but FA2 off, five of five processes agree), and the per-process algorithm choice then pushed it across, sometimes. The non-determinism is the part worth removing, and it predates this round. --- .../frontends/torch/_nexn2_rtx_forward.py | 24 ++++++++++++++++++- 1 file changed, 23 insertions(+), 1 deletion(-) diff --git a/flash_rt/frontends/torch/_nexn2_rtx_forward.py b/flash_rt/frontends/torch/_nexn2_rtx_forward.py index 672999aa..966e456c 100644 --- a/flash_rt/frontends/torch/_nexn2_rtx_forward.py +++ b/flash_rt/frontends/torch/_nexn2_rtx_forward.py @@ -147,7 +147,29 @@ def _proj(x2d, ld, base, n, fvk, device): # split-K reduction order can vary and flip a near-tie argmax. It does not apply # to this entry point, which was measured rather than assumed. Set False to # force the hand-written kernel. -_DENSE_CUBLASLT = True +import os as _os_early + +# The cuBLASLt wrapper picks its algorithm by *timing* eight candidates at +# first use. Timing is noisy, so different processes pick different algorithms, +# and different algorithms reduce in different orders -- which makes the model +# itself non-deterministic across processes. Measured on one binary: the golden +# prefix came out 16/16 three times and 14/16 three times, flipping between +# exactly two token streams. Asking for one candidate takes the heuristic's own +# choice instead, and five of five processes then agree. +# +# It is not a speed trade worth making either way round: at 1024 tokens the +# timed pick is worth 1.6% warm (213.6 against 217.1 ms) and costs 25% of the +# cold time (about 1020 against 770 ms), because the timing loop runs inside +# the first call. Determinism and a faster first token for 1.6% of the warm +# path. +# +# Set here rather than in the kernel, whose default is shared with other +# frontends. FLASHRT_BF16_CUBLASLT_AUTOTUNE_ALGOS overrides. +_os_early.environ.setdefault('FLASHRT_BF16_CUBLASLT_AUTOTUNE_ALGOS', '1') + +# NEXN2_DENSE_CUBLASLT=0 drops to the in-house bf16 GEMM entirely, which is +# also deterministic but 66% slower at 2048 (693 against 418 ms). +_DENSE_CUBLASLT = _os_early.environ.get('NEXN2_DENSE_CUBLASLT', '1') != '0' def _gemm_w16a16(x2d, w, fvk, device): From d07dcc47eae81a1fbd0b548bba2d8de627b7cecc Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Sat, 1 Aug 2026 06:00:55 -0400 Subject: [PATCH 61/85] Feed the draft head its two halves the right way round The head takes the previous hidden state and the next token's embedding concatenated, and fc is square in the concatenated width, so nothing in the checkpoint says which half goes first. It was measured both ways when the head was first loaded, and the measurement went into the notes rather than into the default. Over 48 decoded tokens: cat[embed, hidden] first draft 0.896, chained 0.646, 0.417 cat[hidden, embed] 0.000, 0.000, 0.000 The wrong half drafts noise. Nothing is ever accepted, so every window pays for a verify that keeps the one token it was going to emit anyway, and speculative decode runs at a third of plain greedy. Expected tokens kept per window goes 1.00 -> 1.94 at K=1, 2.64 at K=2, 3.05 at K=3, which is what the acceptance rate predicts. The reference implementation of this head concatenates the embedding first as well. --- flash_rt/frontends/torch/_nexn2_rtx_decode.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/flash_rt/frontends/torch/_nexn2_rtx_decode.py b/flash_rt/frontends/torch/_nexn2_rtx_decode.py index 59c2fdbe..83ca5c13 100644 --- a/flash_rt/frontends/torch/_nexn2_rtx_decode.py +++ b/flash_rt/frontends/torch/_nexn2_rtx_decode.py @@ -313,9 +313,18 @@ def __init__(self, handles, max_seq, device): self._spec_argmax = None # Which half of the draft head's fc input carries the hidden state. # The checkpoint does not say and fc is square in the concatenated - # width, so this is settled by measuring acceptance both ways. + # width, so it was settled by measuring acceptance both ways -- and the + # answer is the embedding first. Over 48 decoded tokens: + # + # cat[embed, hidden] first draft 0.896, chained 0.646, 0.417 + # cat[hidden, embed] 0.000, 0.000, 0.000 + # + # The wrong half drafts noise, so nothing is ever accepted and every + # window pays for a verify that keeps one token. It agrees with the + # reference implementation of this head, which concatenates the + # embedding first as well. self.mtp_hidden_first = ( - _qwen35moe_env("MTP_HIDDEN_FIRST", "1") != "0") + _qwen35moe_env("MTP_HIDDEN_FIRST", "0") != "0") # Set to an ExpertCache to read the routed experts from storage. Only # meaningful when the loader skipped them; see _moe_experts_streamed. self.expert_cache = None From 50258b8b5d97033f92293589ef66839429e4b49e Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Sat, 1 Aug 2026 06:31:51 -0400 Subject: [PATCH 62/85] Verify a speculative window against the weights decode reads A speculative verify has to be the same function as the decode step it verifies, or the tokens it keeps are not the ones plain greedy emits. Ours was neither the same function nor cheap: it ran the prefill forward, which reads the dense weights at BF16 while decode reads them at four bits. Four times the traffic, and a different answer -- measured, logit cosine 0.988 against the decode path, which is why the emitted text diverged. The general W4A16 GEMM does read four bits and is not the answer either. At these shapes it is 7.5 to 9.3 times off the decode GEMV -- 250 us against 33 for an 8192x2048 projection -- and flat in M, so it is not reading the weight at bandwidth at all. This is the decode GEMV with M rows of activation. The weight stream, the lane-to-block mapping, the unroll and the reduction order are untouched; only the rows staged in shared memory and the accumulators a warp carries change. Each output row therefore accumulates in exactly the order the GEMV uses, and the output is bit-identical to running the GEMV once per row -- checked at six shapes by five window widths. Cost at the shapes a verify issues, against one decode row: 1.6x at two rows, 2.0x at three, 2.5x at four. Not the ratio the weight traffic alone implies, because the arithmetic grows with M while the weight read does not -- measured 20 us of weight plus 16 us a row on the largest projection. Still well under the 4.0x of running the GEMV per row, and a third of the general GEMM. Wired into the verify only, over the tensor the decode path caches rather than a second copy: same keys, same bytes. Logit cosine against decode goes 0.988 -> 0.994. Speculative decode at K=2 goes 0.29x of plain greedy to 0.73x. It does not pay yet, and the reason is now specific rather than general: the window still runs the MoE and the linear attention through prefill kernels, so a verify costs about 2.8 decode steps where it needs to cost one. Plain greedy, the golden fixture and the kernel preflight are unchanged. --- CMakeLists.txt | 1 + csrc/bindings.cpp | 12 + csrc/kernels/w4a16_mrows_edge_sm120.cu | 252 ++++++++++++++++++ csrc/kernels/w4a16_mrows_edge_sm120.cuh | 49 ++++ .../frontends/torch/_nexn2_rtx_forward.py | 49 +++- tests/test_w4a16_mrows_edge.py | 100 +++++++ 6 files changed, 461 insertions(+), 2 deletions(-) create mode 100644 csrc/kernels/w4a16_mrows_edge_sm120.cu create mode 100644 csrc/kernels/w4a16_mrows_edge_sm120.cuh create mode 100644 tests/test_w4a16_mrows_edge.py diff --git a/CMakeLists.txt b/CMakeLists.txt index 3312f28b..94803159 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1570,6 +1570,7 @@ if(FLASHRT_ENABLE_QWEN35MOE_W4A16) csrc/kernels/w4a16_matvec_sm120.cu csrc/kernels/moe_grouped_w4a16_sm120.cu csrc/kernels/w4a16_edge_sm120.cu + csrc/kernels/w4a16_mrows_edge_sm120.cu csrc/kernels/w4a16_gemm_sm120.cu) target_compile_definitions(flash_rt_kernels PRIVATE FLASHRT_HAVE_QWEN35MOE_W4A16=1) diff --git a/csrc/bindings.cpp b/csrc/bindings.cpp index 893d3f29..2ddb142e 100644 --- a/csrc/bindings.cpp +++ b/csrc/bindings.cpp @@ -197,6 +197,7 @@ extern "C" int cutlass_int8_rowwise_bf16out_t64x128( #include "kernels/w4a16_matvec_sm120.cuh" #include "kernels/moe_grouped_w4a16_sm120.cuh" #include "kernels/w4a16_edge_sm120.cuh" +#include "kernels/w4a16_mrows_edge_sm120.cuh" #include "kernels/w4a16_gemm_sm120.cuh" #endif // FLASHRT_HAVE_QWEN35MOE_W4A16 #ifdef FLASHRT_HAVE_QWEN35MOE_W4A4 @@ -5613,6 +5614,17 @@ PYBIND11_MODULE(flash_rt_kernels, m) { py::arg("B"), py::arg("S"), py::arg("conv_dim"), py::arg("k"), py::arg("apply_silu") = true, py::arg("stream") = 0); + m.def("w4a16_mrows_edge_sm120_bf16", + [](uintptr_t x, uintptr_t W, uintptr_t SFB, uintptr_t out, + int M, int N, int K, double alpha, uintptr_t stream) -> int { + return flash_rt::kernels::w4a16_mrows_edge_sm120_bf16( + to_ptr(x), to_ptr(W), to_ptr(SFB), to_ptr(out), + M, N, K, static_cast(alpha), to_stream(stream)); + }, + py::arg("x"), py::arg("W"), py::arg("SFB"), py::arg("out"), + py::arg("M"), py::arg("N"), py::arg("K"), py::arg("alpha"), + py::arg("stream") = 0); + m.def("gdn_wy_norm_pack_q_cumsum_edge_bf16", [](uintptr_t q, uintptr_t k, uintptr_t g, uintptr_t k_l2, uintptr_t q_pack, uintptr_t g_cumsum, int S, int num_k_heads, diff --git a/csrc/kernels/w4a16_mrows_edge_sm120.cu b/csrc/kernels/w4a16_mrows_edge_sm120.cu new file mode 100644 index 00000000..29d4a85e --- /dev/null +++ b/csrc/kernels/w4a16_mrows_edge_sm120.cu @@ -0,0 +1,252 @@ +#include "w4a16_mrows_edge_sm120.cuh" + +#include "fp4_e2m1_compat.cuh" + +#include + +namespace flash_rt { +namespace kernels { + +namespace { + +// Same shape constants as the M=1 entry this extends; see its file for why +// each is what it is. The padded shared stride is what keeps the 16-element +// blocks off each other's banks. +constexpr int kWarps = 2; +constexpr int kThreads = kWarps * 32; +constexpr int kUnroll = 4; +constexpr int kBlockSlots = 24; +constexpr int kBlockInt4 = kBlockSlots / 8; +constexpr int kRowsDense = 2; +constexpr int kRowsSmall = 8; +constexpr int kMaxM = 8; + +static_assert(32 % kRowsDense == 0 && 32 % kRowsSmall == 0, + "rows per warp must divide the 32-row scale group"); + +__device__ __forceinline__ float ue4m3_to_float(uint8_t v) { + const int e = (v >> 3) & 0xF; + const int m = v & 0x7; + if (e == 0) return ldexpf(static_cast(m) / 8.0f, -6); + return ldexpf(1.0f + static_cast(m) / 8.0f, e - 7); +} + +__device__ __forceinline__ int sf_off(int rb_ncs, int row_inner, int k_block) { + return (rb_ncs + (k_block >> 2)) * 512 + row_inner + (k_block & 3); +} + +// One packed block against M activation rows. The weight byte pair is decoded +// once and used M times, which is the whole point: the decode is the same work +// the M=1 kernel does, and the extra rows cost only shared reads and fmas. +template +__device__ __forceinline__ void blockdot_m( + uint64_t b_pack, const __nv_bfloat162* x0, size_t x_row_slots, + float (&acc)[M]) { +#pragma unroll + for (int j = 0; j < 8; ++j) { + const __half2_raw wr = flash_rt::fp4::cvt_e2m1x2_to_halfraw2( + static_cast(b_pack >> (j * 8))); + const float2 wf = __half22float2(*reinterpret_cast(&wr)); +#pragma unroll + for (int m = 0; m < M; ++m) { + const float2 xf = __bfloat1622float2( + x0[m * (x_row_slots >> 1) + j]); + acc[m] = fmaf(wf.x, xf.x, acc[m]); + acc[m] = fmaf(wf.y, xf.y, acc[m]); + } + } +} + +__device__ __forceinline__ void stage_padded_row( + const __nv_bfloat16* __restrict__ x, __nv_bfloat16* x_sh, int K) { + const int4* x_i4 = reinterpret_cast(x); + int4* sh_i4 = reinterpret_cast(x_sh); + const int n_i4 = K >> 3; + for (int j = threadIdx.x; j < n_i4; j += kThreads) + sh_i4[(j >> 1) * kBlockInt4 + (j & 1)] = x_i4[j]; +} + +// The K loop, R output rows by M activation rows. Identical in structure to +// the M=1 version: same lane-to-block mapping, same unroll, same order of +// accumulation per (output row, activation row), same final shuffle. +template +__device__ __forceinline__ void row_dot_m( + const uint64_t* __restrict__ w_row0, size_t row_stride_u64, + const uint8_t* __restrict__ SFB, const __nv_bfloat16* x_sh, + size_t x_row_slots, int K_BLOCKS, int rb_ncs, int row_inner, int lane, + float (&acc)[R][M]) { +#pragma unroll + for (int r = 0; r < R; ++r) +#pragma unroll + for (int m = 0; m < M; ++m) acc[r][m] = 0.0f; + + int kb = lane; + const int step = 32 * kUnroll; + for (; kb + 32 * (kUnroll - 1) < K_BLOCKS; kb += step) { + uint64_t wv[R][kUnroll]; + float sf[R][kUnroll]; +#pragma unroll + for (int r = 0; r < R; ++r) +#pragma unroll + for (int u = 0; u < kUnroll; ++u) + wv[r][u] = w_row0[r * row_stride_u64 + kb + 32 * u]; +#pragma unroll + for (int r = 0; r < R; ++r) +#pragma unroll + for (int u = 0; u < kUnroll; ++u) + sf[r][u] = ue4m3_to_float(__ldg( + SFB + sf_off(rb_ncs, row_inner + 16 * r, kb + 32 * u))); +#pragma unroll + for (int r = 0; r < R; ++r) +#pragma unroll + for (int u = 0; u < kUnroll; ++u) { + float part[M]; +#pragma unroll + for (int m = 0; m < M; ++m) part[m] = 0.0f; + blockdot_m( + wv[r][u], + reinterpret_cast( + x_sh + (size_t)(kb + 32 * u) * kBlockSlots), + x_row_slots, part); +#pragma unroll + for (int m = 0; m < M; ++m) acc[r][m] += part[m] * sf[r][u]; + } + } + for (; kb < K_BLOCKS; kb += 32) { + uint64_t wv[R]; + float sf[R]; +#pragma unroll + for (int r = 0; r < R; ++r) wv[r] = w_row0[r * row_stride_u64 + kb]; +#pragma unroll + for (int r = 0; r < R; ++r) + sf[r] = ue4m3_to_float( + __ldg(SFB + sf_off(rb_ncs, row_inner + 16 * r, kb))); +#pragma unroll + for (int r = 0; r < R; ++r) { + float part[M]; +#pragma unroll + for (int m = 0; m < M; ++m) part[m] = 0.0f; + blockdot_m( + wv[r], + reinterpret_cast( + x_sh + (size_t)kb * kBlockSlots), + x_row_slots, part); +#pragma unroll + for (int m = 0; m < M; ++m) acc[r][m] += part[m] * sf[r]; + } + } +#pragma unroll + for (int r = 0; r < R; ++r) +#pragma unroll + for (int m = 0; m < M; ++m) +#pragma unroll + for (int off = 16; off > 0; off >>= 1) + acc[r][m] += __shfl_xor_sync(0xffffffff, acc[r][m], off); +} + +template +__global__ void w4a16_mrows_edge_kernel( + const __nv_bfloat16* __restrict__ x, + const uint8_t* __restrict__ W, + const uint8_t* __restrict__ SFB, + __nv_bfloat16* __restrict__ out, + float alpha, int N, int K, int n_col_super) { + extern __shared__ __nv_bfloat16 x_sh[]; + const size_t row_slots = (size_t)(K >> 4) * kBlockSlots; +#pragma unroll + for (int m = 0; m < M; ++m) + stage_padded_row(x + (size_t)m * K, x_sh + m * row_slots, K); + __syncthreads(); + + const int lane = threadIdx.x & 31; + const int row0 = (blockIdx.x * kWarps + (threadIdx.x >> 5)) * R; + if (row0 >= N) return; + + const int rb = row0 >> 7; + const int ri = row0 & 127; + float acc[R][M]; + row_dot_m( + reinterpret_cast(W + (size_t)row0 * (K >> 1)), + (size_t)(K >> 1) / 8, SFB, x_sh, row_slots, K >> 4, rb * n_col_super, + (ri & 31) * 16 + ((ri >> 5) & 3) * 4, lane, acc); + if (lane == 0) { +#pragma unroll + for (int r = 0; r < R; ++r) { + if (row0 + r >= N) continue; +#pragma unroll + for (int m = 0; m < M; ++m) + out[(size_t)m * N + row0 + r] = __float2bfloat16(acc[r][m] * alpha); + } + } +} + +inline size_t smem_bytes(int K, int M) { + return (size_t)(K >> 4) * kBlockSlots * sizeof(__nv_bfloat16) * M; +} + +inline int rows_per_warp(int N, int K) { + const int r = (K >= 2048) ? kRowsDense : kRowsSmall; + return (N >= r * kWarps) ? r : 1; +} + +#define FLASHRT_MROWS_LAUNCH(R, M) \ + do { \ + const size_t sb = smem_bytes(K, M); \ + if (sb > 48 * 1024) { \ + cudaFuncSetAttribute(w4a16_mrows_edge_kernel, \ + cudaFuncAttributeMaxDynamicSharedMemorySize, \ + static_cast(sb)); \ + } \ + w4a16_mrows_edge_kernel \ + <<>>( \ + reinterpret_cast(x_bf16), \ + reinterpret_cast(W_packed), \ + reinterpret_cast(SFB), \ + reinterpret_cast<__nv_bfloat16*>(out), alpha, N, K, n_col_super); \ + } while (0) + +#define FLASHRT_MROWS_BY_M(R) \ + do { \ + switch (M) { \ + case 1: FLASHRT_MROWS_LAUNCH(R, 1); break; \ + case 2: FLASHRT_MROWS_LAUNCH(R, 2); break; \ + case 3: FLASHRT_MROWS_LAUNCH(R, 3); break; \ + case 4: FLASHRT_MROWS_LAUNCH(R, 4); break; \ + case 5: FLASHRT_MROWS_LAUNCH(R, 5); break; \ + case 6: FLASHRT_MROWS_LAUNCH(R, 6); break; \ + case 7: FLASHRT_MROWS_LAUNCH(R, 7); break; \ + default: FLASHRT_MROWS_LAUNCH(R, 8); break; \ + } \ + } while (0) + +} // namespace + +int w4a16_mrows_edge_sm120_bf16( + const void* x_bf16, + const void* W_packed, + const void* SFB, + void* out, + int M, + int N, + int K, + float alpha, + cudaStream_t stream) { + if (!x_bf16 || !W_packed || !SFB || !out) return 1; + if (N <= 0 || K <= 0 || (K & 15) != 0) return 2; + if (M <= 0 || M > kMaxM) return 3; + + const int n_col_super = ((K >> 4) + 3) / 4; + const int R = rows_per_warp(N, K); + if (R == kRowsDense) { + FLASHRT_MROWS_BY_M(kRowsDense); + } else if (R == kRowsSmall) { + FLASHRT_MROWS_BY_M(kRowsSmall); + } else { + FLASHRT_MROWS_BY_M(1); + } + return 0; +} + +} // namespace kernels +} // namespace flash_rt diff --git a/csrc/kernels/w4a16_mrows_edge_sm120.cuh b/csrc/kernels/w4a16_mrows_edge_sm120.cuh new file mode 100644 index 00000000..cb3ba4bf --- /dev/null +++ b/csrc/kernels/w4a16_mrows_edge_sm120.cuh @@ -0,0 +1,49 @@ +#pragma once + +#include + +namespace flash_rt { +namespace kernels { + +// A few rows of activation against the same 4-bit weight the decode GEMV reads. +// +// A speculative verify has to be the same function as the decode step it +// verifies, or the tokens it keeps are not the ones plain greedy emits. Two +// things were in the way. The verify read the dense weights at BF16 while +// decode reads them at 4 bits -- four times the traffic, and a different +// answer: measured logit cosine 0.988 between the two forwards. And the +// general W4A16 GEMM, which does read 4 bits, is not the same arithmetic +// either, and at these shapes it is 7.5 to 9.3 times off the GEMV: 250 us +// against 33 for an 8192x2048 projection, and flat in M, so it is not reading +// the weight at bandwidth at all. +// +// This is the decode GEMV with M rows of activation. The weight stream, the +// lane-to-block mapping, the unroll and the reduction order are unchanged -- +// only the number of activation rows staged in shared memory and the number of +// accumulators a warp carries. The weight is what costs, and it is read once +// regardless of M, so a window of four verifies for what one costs. +// +// Because each output row accumulates in exactly the order the GEMV uses, the +// result at M=1 is bit-identical to it, and rows of a larger M agree with what +// the GEMV would have produced for each row on its own. +// +// x (M, K) bf16, row-major +// W (N, K/2) NVFP4 e2m1 nibbles +// SFB swizzled UE4M3 block scales, as bf16_weight_to_nvfp4_swizzled writes +// out (M, N) bf16 +// alpha weight per-tensor global scale +// +// K must be a multiple of 16, M at most 8. Returns 0 on success. +int w4a16_mrows_edge_sm120_bf16( + const void* x_bf16, + const void* W_packed, + const void* SFB, + void* out, + int M, + int N, + int K, + float alpha, + cudaStream_t stream); + +} // namespace kernels +} // namespace flash_rt diff --git a/flash_rt/frontends/torch/_nexn2_rtx_forward.py b/flash_rt/frontends/torch/_nexn2_rtx_forward.py index 966e456c..0fb8078b 100644 --- a/flash_rt/frontends/torch/_nexn2_rtx_forward.py +++ b/flash_rt/frontends/torch/_nexn2_rtx_forward.py @@ -121,7 +121,10 @@ def _proj(x2d, ld, base, n, fvk, device): # the pass whose whole purpose is to read the weights once. It wants the # 4-bit weight for the same reason decode does: at this M the cost is # traffic, not throughput. - if ((_SPEC_VERIFY or (_DENSE_W4A16 and x2d.shape[0] >= 64)) + if (_SPEC_VERIFY and x2d.shape[0] <= 8 + and (w.shape[1] % 16) == 0): + return _w4a16_mrows(x2d, w, ld, base + '_w_t', fvk, device) + if (_DENSE_W4A16 and x2d.shape[0] >= 64 and (w.shape[0] % 64) == 0 and (x2d.shape[1] % 64) == 0): return _gemm_w4a16(x2d, w, ld, base + '_w_t', fvk, device) if (_DENSE_W16A16 and x2d.shape[0] >= _DENSE_BF16_MIN_M @@ -216,7 +219,13 @@ def _gemm_w16a16(x2d, w, fvk, device): # bytes and differs from decode only by reduction order. Until that exists, # BF16 is both faster and the one that agrees with plain greedy token for # token. -_SPEC_VERIFY_W4A16 = False +# On: the verify runs the dense projections through the M-row form of the +# decode GEMV, over the tensor the decode path caches. Off, it reads the same +# weights at BF16 -- four times the bytes, and a different answer from the step +# it is verifying (measured logit cosine 0.988 against decode, which is what +# made the emitted text diverge from plain greedy). +_SPEC_VERIFY_W4A16 = _os_early.environ.get( + 'NEXN2_SPEC_VERIFY_W4A16', '1') != '0' _SPEC_VERIFY = False @@ -245,6 +254,42 @@ def set_spec_verify(on: bool) -> None: _DENSE_W16A16 = True +def _w4a16_mrows(x2d, w, ld, key, fvk, device): + """A few rows against the 4-bit weight the *decode* path caches. + + This is what makes a speculative verify the same function as the step it + verifies. It reads the identical packed tensor under the identical cache + keys the decode GEMV uses -- not a second copy quantised by a different + helper -- and runs the M-row form of that GEMV, whose per-row accumulation + order is the GEMV's. So a verified row equals the decode row it stands in + for, bit for bit, and the window reads the weight once rather than once per + token and at a quarter of the bytes the BF16 path reads. + """ + n, k = w.shape + pk = key + '_w4a16_p' + if pk not in ld: + packed = torch.empty(n, k // 2, dtype=torch.uint8, device=device) + sf = torch.zeros(_sf_swz_bytes(n, k), dtype=torch.uint8, device=device) + scr = torch.zeros(1, dtype=torch.float32, device=device) + og = torch.zeros(1, dtype=torch.float32, device=device) + fvk.bf16_weight_to_nvfp4_swizzled( + w.contiguous().data_ptr(), packed.data_ptr(), sf.data_ptr(), + scr.data_ptr(), og.data_ptr(), n, k, _cs()) + torch.cuda.synchronize() + ld[pk] = packed + ld[key + '_w4a16_sf'] = sf + ld[key + '_w4a16_a'] = float(og.item()) + m = x2d.shape[0] + xc = x2d.contiguous() + y = torch.empty(m, n, dtype=torch.bfloat16, device=device) + rc = fvk.w4a16_mrows_edge_sm120_bf16( + xc.data_ptr(), ld[pk].data_ptr(), ld[key + '_w4a16_sf'].data_ptr(), + y.data_ptr(), m, n, k, ld[key + '_w4a16_a'], _cs()) + if rc: + raise RuntimeError(f'M-row W4A16 failed with {rc} at M={m}') + return y + + def _gemm_w4a16(x2d, w, ld, key, fvk, device): """y = x @ w.T via the bf16-act x fp4-weight tensor-core GEMM. Weight quantised to NVFP4 once (cached); activation stays BF16 (precise).""" diff --git a/tests/test_w4a16_mrows_edge.py b/tests/test_w4a16_mrows_edge.py new file mode 100644 index 00000000..e61ae8cc --- /dev/null +++ b/tests/test_w4a16_mrows_edge.py @@ -0,0 +1,100 @@ +"""Equivalence test for the M-row W4A16 GEMV. + +This kernel exists so a speculative verify computes the same function as the +decode step it verifies, so the bar is exactness, not closeness: row m of its +output must equal what the M=1 GEMV produces for row m on its own, bit for bit. +A verify that is merely close keeps tokens plain greedy would not have emitted. + +The M=1 case is checked as well, since it is the claim that the extension left +the original arithmetic alone. +""" + +import pytest +import torch + +# The dense projections a decode step runs, so the shapes are the real ones. +SHAPES = [ + (8192, 2048), # in_proj_qkv, q_proj + (4096, 2048), # in_proj_z + (2048, 4096), # gdn out_proj, o_proj + (512, 2048), # shared gate/up + (2048, 512), # shared down + (256, 2048), # router-sized, below one warp's row block +] + + +def _load_fvk(): + if not torch.cuda.is_available(): + pytest.skip("CUDA is required for the M-row W4A16 test") + try: + from flash_rt import flash_rt_kernels as fvk + except Exception as exc: # pragma: no cover - environmental + pytest.skip(f"flash_rt_kernels is not built: {exc}") + for name in ("w4a16_mrows_edge_sm120_bf16", + "w4a16_matvec_edge_sm120_bf16", + "bf16_weight_to_nvfp4_swizzled"): + if not hasattr(fvk, name): + pytest.skip(f"{name} not in this build") + return fvk + + +def _quantise(fvk, w, N, K, dev): + from flash_rt.frontends.torch._nexn2_rtx_nvfp4_weights import _sf_swz_bytes + packed = torch.empty(N, K // 2, dtype=torch.uint8, device=dev) + sf = torch.zeros(_sf_swz_bytes(N, K), dtype=torch.uint8, device=dev) + scr = torch.zeros(1, dtype=torch.float32, device=dev) + og = torch.zeros(1, dtype=torch.float32, device=dev) + fvk.bf16_weight_to_nvfp4_swizzled( + w.contiguous().data_ptr(), packed.data_ptr(), sf.data_ptr(), + scr.data_ptr(), og.data_ptr(), N, K, + torch.cuda.current_stream().cuda_stream) + torch.cuda.synchronize(dev) + return packed, sf, float(og.item()) + + +@pytest.mark.parametrize("N,K", SHAPES) +@pytest.mark.parametrize("M", [1, 2, 3, 4, 8]) +def test_rows_match_the_per_row_gemv_exactly(N, K, M): + fvk = _load_fvk() + dev = "cuda:0" + st = torch.cuda.current_stream().cuda_stream + g = torch.Generator(device=dev).manual_seed(N + K + M) + w = torch.randn(N, K, generator=g, device=dev, dtype=torch.bfloat16) + packed, sf, alpha = _quantise(fvk, w, N, K, dev) + x = torch.randn(M, K, generator=g, device=dev, dtype=torch.bfloat16) + + ref = torch.empty(M, N, dtype=torch.bfloat16, device=dev) + for m in range(M): + y1 = torch.empty(1, N, dtype=torch.bfloat16, device=dev) + fvk.w4a16_matvec_edge_sm120_bf16( + x[m:m + 1].contiguous().data_ptr(), packed.data_ptr(), + sf.data_ptr(), y1.data_ptr(), N, K, alpha, st) + torch.cuda.synchronize(dev) + ref[m].copy_(y1[0]) + + got = torch.empty(M, N, dtype=torch.bfloat16, device=dev) + rc = fvk.w4a16_mrows_edge_sm120_bf16( + x.contiguous().data_ptr(), packed.data_ptr(), sf.data_ptr(), + got.data_ptr(), M, N, K, alpha, st) + assert rc == 0, f"kernel returned {rc}" + torch.cuda.synchronize(dev) + + assert torch.equal(got, ref), ( + f"{int((got != ref).sum())} of {got.numel()} elements differ from the " + "per-row GEMV") + + +def test_rejects_an_M_it_cannot_hold(): + """A window wider than the kernel stages should be refused, not truncated.""" + fvk = _load_fvk() + dev = "cuda:0" + N, K, M = 512, 2048, 9 + g = torch.Generator(device=dev).manual_seed(1) + w = torch.randn(N, K, generator=g, device=dev, dtype=torch.bfloat16) + packed, sf, alpha = _quantise(fvk, w, N, K, dev) + x = torch.randn(M, K, generator=g, device=dev, dtype=torch.bfloat16) + out = torch.empty(M, N, dtype=torch.bfloat16, device=dev) + rc = fvk.w4a16_mrows_edge_sm120_bf16( + x.data_ptr(), packed.data_ptr(), sf.data_ptr(), out.data_ptr(), + M, N, K, alpha, torch.cuda.current_stream().cuda_stream) + assert rc != 0, "an over-wide window was accepted" From 4eb09e73e65dfdf2d5259aaee84ca59b1e0ef35c Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Sat, 1 Aug 2026 07:25:59 -0400 Subject: [PATCH 63/85] Seed the draft head's hidden state with the prompt's last position A DeepSeek-V3-style draft head reads the pre-final-norm hidden state of the position before the token it is given. The per-token seeding path writes that buffer every step, so it was right there; the batched and chunked paths never wrote it at all. So the first speculative window of a generation drafted off whatever the previous generation had left in the buffer. Nothing crashes and nothing looks wrong -- the draft is simply predicted from an unrelated position, so the window keeps one token instead of two, the emitted sequence shifts, and every position after it is a different position. Which means the same prompt does not decode to the same text twice, and the tokens kept per window depend on what ran before. Both batched paths now take the hidden state the forward already computes and copy its last row. It is a 4 KB device copy on a path that runs once per prompt. --- flash_rt/frontends/torch/_nexn2_rtx_decode.py | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/flash_rt/frontends/torch/_nexn2_rtx_decode.py b/flash_rt/frontends/torch/_nexn2_rtx_decode.py index 83ca5c13..64e7f139 100644 --- a/flash_rt/frontends/torch/_nexn2_rtx_decode.py +++ b/flash_rt/frontends/torch/_nexn2_rtx_decode.py @@ -1039,9 +1039,15 @@ def seed_prefill_batched(state, input_ids, fvk, device): return seed_prefill_chunked(state, input_ids, fvk, device, state.prefill_chunk) state.reset() - logits = nexn2_forward_nvfp4( + logits, hidden = nexn2_forward_nvfp4( state.handles, input_ids.view(1, -1), fvk, device, cap=state, - last_logits_only=True) + last_logits_only=True, return_hidden=True) + # The last prompt position's pre-final-norm hidden state, which is what a + # draft head reads. The per-token path writes it every step; this one has + # to do it explicitly, and without it the first window drafts off whatever + # the previous generation left behind -- so how much of that window is kept + # depends on what ran before it, and the run stops being reproducible. + state.last_hidden.copy_(hidden[-1]) return logits # already (1, vocab): only the seeding logit @@ -1059,9 +1065,11 @@ def seed_prefill_chunked(state, input_ids, fvk, device, block): logits = None for b0 in range(0, S, block): b1 = min(b0 + block, S) - logits = nexn2_forward_nvfp4( + logits, hidden = nexn2_forward_nvfp4( state.handles, ids[:, b0:b1], fvk, device, cap=state, - pos_offset=b0, last_logits_only=True, compute_logits=(b1 == S)) + pos_offset=b0, last_logits_only=True, compute_logits=(b1 == S), + return_hidden=True) + state.last_hidden.copy_(hidden[-1]) # see seed_prefill_batched return logits From fc894f3ac80b084b150c3207ed4efd9f9cf79583 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Sat, 1 Aug 2026 07:29:38 -0400 Subject: [PATCH 64/85] Run a speculative window on the decode kernels, at window rows The verify was still the prefill forward. Only the dense projections had been moved onto the decode path's kernels; the routed experts, the linear attention and the full attention were all somebody else's arithmetic. So a window cost about 2.8 decode steps where it should cost near one, and -- the part that actually matters -- it was not the same function as the steps it verified, so the tokens it kept were not the ones plain greedy emits. The 27B solves this with a third forward and a parallel set of layer methods. That is more than this needs, because the two kernels that looked like the hard part already take the parameter: the grouped W4A16 GEMV takes a slot count, and the MoE slots of w tokens are w*TOPK independent GEMVs. So the window is decode_step with w rows, calling the kernels decode calls, over the weights decode caches. Three stages stay per token on purpose. The causal conv and the recurrence carry state, and a window is accepted up to a prefix, so each token's state has to be the state decode would have been in -- they run through the decode kernels a token at a time, snapshotting as they go. The sequential-scan variant would collapse the recurrence into one launch but it is cos 0.99999 against the per-token kernel rather than equal to it, and this layer is 6% of a step. Attention likewise runs at q_seq=1 per token: a batched q_seq=w call would need a bottom-right causal mask and would reduce over a different tiling. Measured against the decode step it stands in for, over three windows of four at real decoded tokens: every logits row bit-identical, every per-token recurrent and conv snapshot bit-identical, every KV row bit-identical. Cost, against a captured decode step of 11.18 ms: the window costs 1.02 steps at one row and about 0.20 of a step per row after that. The fixed part is the dense weights and the lm_head, read once however wide the window is; the per-row part is the routed experts, which do not amortise at all -- w tokens pick up to w*8 distinct experts out of 256, so that traffic is the floor on what a wider window can be worth. End to end against plain greedy, one process per K, 64 tokens: K=1 1.09x kept 1.88 K=2 1.13x kept 2.71 91.1 tok/s against 80.5 K=3 1.00x kept 3.20 K=4 0.82x kept 3.10 Off the default path: plain decode, the golden fixture and the kernel preflight are untouched, and FLASHRT_QWEN35MOE_VERIFY_K_ROWS=0 puts the prefill forward back. The lm_head quantisation moves to its own helper so the single-row head and the window read one copy rather than two of the same bytes. --- flash_rt/frontends/torch/_nexn2_rtx_decode.py | 377 ++++++++++++++++-- 1 file changed, 351 insertions(+), 26 deletions(-) diff --git a/flash_rt/frontends/torch/_nexn2_rtx_decode.py b/flash_rt/frontends/torch/_nexn2_rtx_decode.py index 64e7f139..d4a5f601 100644 --- a/flash_rt/frontends/torch/_nexn2_rtx_decode.py +++ b/flash_rt/frontends/torch/_nexn2_rtx_decode.py @@ -29,8 +29,8 @@ from flash_rt.frontends.torch._nexn2_rtx_forward import ( CONV, HD, HID, HK, HV, INTER, KD, KS, NKV, NQ, NV, ROPE, TOPK, VD, - _quant_act, build_rope_tables, moe_grouped_w4a16, nexn2_forward_nvfp4, - set_spec_verify, w4a16_matvec, + _quant_act, _w4a16_mrows, build_rope_tables, moe_grouped_w4a16, + nexn2_forward_nvfp4, set_spec_verify, w4a16_matvec, ) from flash_rt.frontends.torch._nexn2_rtx_nvfp4_weights import _sf_swz_bytes from flash_rt.hardware.rtx.attn_backend_nexn2 import RtxFlashAttnBackendNexn2 @@ -47,6 +47,11 @@ def _qwen35moe_env(name: str, default: str) -> str: # wins). See Nexn2DecodeState.batched_prefill. _BATCHED_PREFILL_MIN_S = 8 +# Run a speculative window through the decode kernels at w rows rather than +# through the prefill forward. Read once, at import, because a captured graph +# replays whatever branch was taken when it was recorded. +_VERIFY_K_ROWS = _qwen35moe_env("VERIFY_K_ROWS", "1") != "0" + def _cs(): """Current CUDA stream handle. Inside torch.cuda.graph capture this is @@ -716,6 +721,32 @@ def decode_step(state, token_id, pos, fvk, device): return _lm_head(state, h, fvk, device) +def _ensure_lm_head_nvfp4(state, fvk, device): + """Quantise the lm_head to swizzled NVFP4 once, on the handles. + + Both the single-row decode head and the M-row verify read this one copy, + so the verify cannot drift from decode by having been handed a second + quantisation of the same weight. The .item() lands here, on the first + eager call, and never inside a captured region. + """ + p = state.handles.ptrs + if 'lm_head_packed_t' in p: + return + w = p['lm_head_w_t'].contiguous() + nn, kk = w.shape + packed = torch.empty(nn, kk // 2, dtype=torch.uint8, device=device) + sf = torch.zeros(_sf_swz_bytes(nn, kk), dtype=torch.uint8, device=device) + scr = torch.zeros(1, dtype=torch.float32, device=device) + og = torch.zeros(1, dtype=torch.float32, device=device) + fvk.bf16_weight_to_nvfp4_swizzled( + w.data_ptr(), packed.data_ptr(), sf.data_ptr(), + scr.data_ptr(), og.data_ptr(), nn, kk, 0) + torch.cuda.synchronize() + p['lm_head_packed_t'] = packed + p['lm_head_sf_t'] = sf + p['lm_head_alpha'] = float(og.item()) + + def _lm_head(state, h, fvk, device): """Project a hidden state to logits over the full vocabulary. @@ -730,20 +761,7 @@ def _lm_head(state, h, fvk, device): h.reshape(1, HID).contiguous().data_ptr(), p['lm_head_w_t'].data_ptr(), logits.data_ptr(), vocab, HID, _cs()) return logits - if 'lm_head_packed_t' not in p: - w = p['lm_head_w_t'].contiguous() - nn, kk = w.shape - packed = torch.empty(nn, kk // 2, dtype=torch.uint8, device=device) - sf = torch.zeros(_sf_swz_bytes(nn, kk), dtype=torch.uint8, device=device) - scr = torch.zeros(1, dtype=torch.float32, device=device) - og = torch.zeros(1, dtype=torch.float32, device=device) - fvk.bf16_weight_to_nvfp4_swizzled( - w.data_ptr(), packed.data_ptr(), sf.data_ptr(), - scr.data_ptr(), og.data_ptr(), nn, kk, 0) - torch.cuda.synchronize() - p['lm_head_packed_t'] = packed - p['lm_head_sf_t'] = sf - p['lm_head_alpha'] = float(og.item()) + _ensure_lm_head_nvfp4(state, fvk, device) if hasattr(fvk, 'fp4_w4a4_mma_sm120_full_n_bf16out'): xp, xsf = _quant_act(h.reshape(1, HID), fvk, device, _cs()) fvk.fp4_w4a4_mma_sm120_full_n_bf16out( @@ -864,6 +882,309 @@ def _rewind_to(state, kept): state.lin_conv_state[rank].copy_(state.spec_conv[rank][kept - 1]) +def _verify_dense(x2d, w_bf16, ld, key, fvk, device): + """A window's rows against the 4-bit weight the decode GEMV reads. + + Same tensor under the same cache key, and the M-row form of that GEMV, + whose per-row accumulation order is the GEMV's -- so row t of the result + equals what the decode step at that position would have computed, bit for + bit, while the weight crosses the bus once for the whole window. + """ + return _w4a16_mrows(x2d, w_bf16, ld, key, fvk, device) + + +def _verify_gdn(h, ld, state, lin_rank, w, fvk, device): + """The GDN layer over a window of w tokens, snapshotting per token. + + The projections and the elementwise stages run at w rows; the two stages + that carry state -- the causal conv and the recurrence -- run a token at a + time through the very kernels the decode step calls, because a window is + accepted up to a prefix and the state at that prefix has to be the state + decode would have been in. The sequential-scan variant would do both in one + launch, but it is cos 0.99999 against the per-token kernel rather than + equal to it, and this layer is ~6% of a step: not worth paying for in + tokens that diverge. + """ + eps = state.eps + convw = ld['conv1d_w_t'].reshape(CONV, KS).contiguous() + A_log, dtb = ld['A_log_t'].float(), ld['dt_bias_t'].float() + nw = ld['gdn_norm_w_t'] + s = _cs() + x = h.reshape(w, HID) + + if 'in_proj_fused_w' not in ld: + ld['in_proj_fused_w'] = torch.cat( + [ld['in_proj_qkv_w_t'], ld['in_proj_z_w_t'], + ld['in_proj_a_w_t'], ld['in_proj_b_w_t']], 0).contiguous() + fused = _verify_dense(x, ld['in_proj_fused_w'], ld, 'in_proj_fused', + fvk, device) + mixed = fused[:, :KD * 2 + VD].contiguous() + z = fused[:, KD * 2 + VD:KD * 2 + VD + NV * HV].reshape( + w * NV, HV).contiguous() + a = fused[:, -2 * NV:-NV].contiguous() + b = fused[:, -NV:].contiguous() + + conv_out = torch.empty(w, CONV, dtype=torch.bfloat16, device=device) + st_in = state.lin_conv_state[lin_rank] + for t in range(w): + st_out = state.spec_conv[lin_rank][t] + fvk.causal_conv1d_qwen36_update_inout_bf16( + mixed[t].data_ptr(), convw.data_ptr(), 0, + conv_out[t].data_ptr(), st_in.data_ptr(), st_out.data_ptr(), + 1, CONV, KS, True, s) + st_in = st_out + state.lin_conv_state[lin_rank].copy_(st_in) + + qb = torch.empty(w, NV, HK, dtype=torch.bfloat16, device=device) + kb = torch.empty(w, NV, HK, dtype=torch.bfloat16, device=device) + vb = torch.empty(w, NV, HV, dtype=torch.bfloat16, device=device) + fvk.qwen35moe_lin_split_qkv_broadcast_bf16( + conv_out.data_ptr(), qb.data_ptr(), kb.data_ptr(), vb.data_ptr(), + w, s) + + neg = (-A_log.exp()).float().contiguous() + dtb_c = dtb.contiguous() + g_out = torch.empty(w, NV, dtype=torch.bfloat16, device=device) + bo = torch.empty(w, NV, dtype=torch.bfloat16, device=device) + fvk.qwen36_gdn_gating_bf16( + a.data_ptr(), b.data_ptr(), neg.data_ptr(), dtb_c.data_ptr(), + g_out.data_ptr(), bo.data_ptr(), w, NV, s) + + core = torch.empty(w, NV, HV, dtype=torch.bfloat16, device=device) + lin_state = state.lin_state[lin_rank] + for t in range(w): + qt, kt, vt = qb[t], kb[t], vb[t] + gt, bt = g_out[t], bo[t] + fvk.gated_deltanet_recurrent_qwen36_bf16( + qt.data_ptr(), kt.data_ptr(), vt.data_ptr(), gt.data_ptr(), + bt.data_ptr(), lin_state.data_ptr(), core[t].data_ptr(), + 1, NV, HK, HV, True, s) + state.spec_states[lin_rank][t].copy_(lin_state) + + nf = torch.empty(w * NV, HV, dtype=torch.bfloat16, device=device) + fvk.rms_norm_gated_silu_qwen36_bf16( + core.reshape(w * NV, HV).data_ptr(), z.data_ptr(), nw.data_ptr(), + nf.data_ptr(), w * NV, HV, eps, s) + out = _verify_dense(nf.reshape(w, VD), ld['out_proj_w_t'], ld, + 'out_proj_w_t', fvk, device) + return out.reshape(1, w, HID) + + +def _verify_full(h, ld, state, full_rank, pos, w, fvk, device): + """The full-attention layer over a window of w tokens. + + Projections, norms and rope run at w rows. The attention itself runs a + token at a time at q_seq=1 against [0..pos+t], which is the call the decode + step makes: a batched q_seq=w call would have to carry a bottom-right + causal mask and would reduce over a different tiling, and this is the one + place where the two would stop being the same function. The KV it reads is + small next to the weights the window is here to amortise. + """ + eps = state.eps + s = _cs() + qnw, knw = ld['q_norm_w_t'], ld['k_norm_w_t'] + x2 = h.reshape(w, HID) + + nqg = NQ * 2 * HD + if 'qkv_fused_w' not in ld: + ld['qkv_fused_w'] = torch.cat( + [ld['q_proj_w_t'], ld['k_proj_w_t'], ld['v_proj_w_t']], + 0).contiguous() + fused = _verify_dense(x2, ld['qkv_fused_w'], ld, 'qkv_fused', fvk, device) + qg = fused[:, :nqg].contiguous() + kk = fused[:, nqg:nqg + NKV * HD].reshape(w * NKV, HD).contiguous() + v = fused[:, nqg + NKV * HD:].reshape(w, NKV, HD).contiguous() + + q_pre = torch.empty(w, NQ, HD, dtype=torch.bfloat16, device=device) + gate = torch.empty(w, NQ * HD, dtype=torch.bfloat16, device=device) + fvk.qwen35moe_split_q_gate_bf16( + qg.data_ptr(), q_pre.data_ptr(), gate.data_ptr(), w, s) + q = _rms_fvk(q_pre.reshape(w * NQ, HD), qnw, fvk, device, eps) + kn = _rms_fvk(kk, knw, fvk, device, eps) + + ct = state.rope_cos[pos:pos + w].contiguous() + st = state.rope_sin[pos:pos + w].contiguous() + qin = q.reshape(w, NQ, HD).contiguous() + kin = kn.reshape(w, NKV, HD).contiguous() + qo = torch.empty(w, NQ, HD, dtype=torch.bfloat16, device=device) + ko = torch.empty(w, NKV, HD, dtype=torch.bfloat16, device=device) + fvk.qwen36_partial_rope_qk_bf16( + qin.data_ptr(), kin.data_ptr(), ct.data_ptr(), st.data_ptr(), + qo.data_ptr(), ko.data_ptr(), w, NQ, NKV, HD, ROPE, s) + + attn = state.attn + at = torch.empty(w, NQ * HD, dtype=torch.bfloat16, device=device) + for t in range(w): + attn.Q_buf[:, :1].copy_(qo[t].reshape(1, 1, NQ, HD)) + attn.K_cache[full_rank, pos + t:pos + t + 1].copy_( + ko[t].reshape(1, NKV, HD)) + attn.V_cache[full_rank, pos + t:pos + t + 1].copy_( + v[t].reshape(1, NKV, HD)) + attn.run('full', layer_idx=full_rank, q_seq=1, kv_seq=pos + t + 1, + stream=s, softmax_scale=float(HD) ** -0.5) + at[t].copy_(attn.O_buf[:, :1].reshape(NQ * HD)) + at = _sigmoid_mul(at, gate, fvk, device) + out = _verify_dense(at, ld['o_proj_w_t'], ld, 'o_proj_w_t', fvk, device) + return out.reshape(1, w, HID) + + +def _verify_moe(h, ld, state, w, fvk, device): + """The MoE layer over a window of w tokens. + + The routed experts are the one part of a window that does not amortise: + w tokens pick up to w*TOPK distinct experts out of 256, so the weight + traffic here scales with the window where everything else is read once. + They still go through one grouped launch rather than w of them -- the + kernel already takes the slot count, and a slot is an independent GEMV, so + w*TOPK slots compute exactly what w separate TOPK-slot launches would. + """ + s = _cs() + x = h.reshape(w, HID) + ne = ld['router_w_t'].shape[0] + + if 'router_shared_fused_w' not in ld: + ld['router_shared_fused_w'] = torch.cat( + [ld['router_w_t'], ld['shared_gate_proj_w_t'], + ld['shared_up_proj_w_t']], 0).contiguous() + rs = _verify_dense(x, ld['router_shared_fused_w'], ld, + 'router_shared_fused', fvk, device) + logit_raw = rs[:, :ne].contiguous() + sg, su = rs[:, ne:ne + INTER], rs[:, ne + INTER:] + + # Top-8 a row at a time through the decode router. It is a single-block + # kernel, so w launches is w small launches -- and the selected set has to + # be the set decode selects, ties included, or the window keeps a token + # from a different mixture. + idx = torch.empty(w, TOPK, dtype=torch.int32, device=device) + topv = torch.empty(w, TOPK, dtype=torch.float32, device=device) + for t in range(w): + rc = fvk.moe_router_topk_sm120_bf16( + logit_raw[t].data_ptr(), idx[t].data_ptr(), topv[t].data_ptr(), + ne, TOPK, s) + if rc: + raise RuntimeError( + f'router top-k failed with {rc} for {ne} experts, k={TOPK}') + tw = F.softmax(topv, -1) + + if 'experts_gate_up_alpha_dev' not in ld: + ld['experts_gate_up_alpha_dev'] = \ + ld['experts_gate_up_alpha_t'].to(device).contiguous() + ld['experts_down_alpha_dev'] = \ + ld['experts_down_alpha_t'].to(device).contiguous() + gu_p, gu_s = ld['experts_gate_up_packed_t'], ld['experts_gate_up_sf_t'] + dn_p, dn_s = ld['experts_down_packed_t'], ld['experts_down_sf_t'] + gu_a, dn_a = ld['experts_gate_up_alpha_dev'], ld['experts_down_alpha_dev'] + n_gu, n_dn = gu_p.shape[1], dn_p.shape[1] + + slots = w * TOPK + eidx = idx.reshape(-1).contiguous() + # One activation row per slot: the grouped kernel indexes A by slot, and + # the decode call gets the same effect from a zero stride over its single + # row. The copy is w*TOPK*HID bf16 -- tens of KB. + xrep = x.repeat_interleave(TOPK, 0).contiguous() + d_gu = torch.empty(slots, n_gu, dtype=torch.bfloat16, device=device) + moe_grouped_w4a16(fvk)( + xrep.data_ptr(), gu_p.data_ptr(), gu_s.data_ptr(), gu_a.data_ptr(), + eidx.data_ptr(), d_gu.data_ptr(), slots, n_gu, HID, + HID, gu_p[0].numel(), gu_s[0].numel(), s) + + g_, u_ = d_gu[:, :INTER], d_gu[:, INTER:] + inter = _silu_mul(g_, u_, fvk, device).contiguous() + d_dn = torch.empty(slots, n_dn, dtype=torch.bfloat16, device=device) + moe_grouped_w4a16(fvk)( + inter.data_ptr(), dn_p.data_ptr(), dn_s.data_ptr(), dn_a.data_ptr(), + eidx.data_ptr(), d_dn.data_ptr(), slots, n_dn, INTER, + INTER, dn_p[0].numel(), dn_s[0].numel(), s) + + rk = f'verify_topk_rows_{w}' + if rk not in ld: + ld[rk] = torch.arange(slots, dtype=torch.int32, device=device) + twf = tw.reshape(-1).contiguous() + out = torch.empty(w, n_dn, dtype=torch.float32, device=device) + fvk.moe_weighted_sum_sm120_bf16( + d_dn.data_ptr(), ld[rk].data_ptr(), twf.data_ptr(), + out.data_ptr(), w, TOPK, n_dn, n_dn, s) + + si = _silu_mul(sg.contiguous(), su.contiguous(), fvk, device) + shared = _verify_dense(si, ld['shared_down_proj_w_t'], ld, + 'shared_down_proj_w_t', fvk, device) + # The scalar gate is an N=1 GEMV over a 4 KB weight, so a row at a time + # costs nothing and is the decode kernel's own arithmetic. + xc = x.contiguous() + gsc = torch.empty(w, 1, dtype=torch.bfloat16, device=device) + for t in range(w): + xr = xc[t] + fvk.bf16_matvec_sm120_bf16( + xr.data_ptr(), ld['shared_gate_w_t'].data_ptr(), + gsc[t].data_ptr(), 1, HID, s) + sgate = torch.sigmoid(gsc.float()) + return (out + shared.float() * sgate).reshape( + 1, w, HID).to(torch.bfloat16) + + +def _verify_block_K(state, toks, pos, w, fvk, device): + """Run a window of w tokens through the decode kernels, at w rows. + + This is what makes the verify the same function as the steps it verifies. + Every stage is the kernel decode calls, at w rows instead of one, over the + weights decode caches; the two state-carrying stages and the attention run + per token for the reasons given above. So the window's row t is the decode + step at pos+t, and the largest weights cross the bus once instead of w + times. + + Returns (logits (w, vocab), hidden (w, HID) pre-final-norm). + """ + p = state.handles.ptrs + layers = p['layers'] + h = F.embedding(toks.view(1, w), p['embed_w_t']) + + for L in range(state.num_layers): + ld = layers[L] + res = h + n = _rms_fvk(h, ld['input_norm_w_t'], fvk, device, state.eps) + if state.types[L] == 'linear_attention': + attn = _verify_gdn(n, ld, state, state._lin_rank[L], w, + fvk, device) + else: + attn = _verify_full(n, ld, state, state._full_rank[L], pos, w, + fvk, device) + h = res + attn + res = h + n = _rms_fvk(h, ld['post_norm_w_t'], fvk, device, state.eps) + state._active_layer = L + h = res + _verify_moe(n, ld, state, w, fvk, device) + + hidden = h.reshape(w, HID) + hn = _rms_fvk(h, p['final_norm_w_t'], fvk, device, state.eps) + vocab = p['vocab_size'] + _ensure_lm_head_nvfp4(state, fvk, device) + logits = torch.empty(w, vocab, dtype=torch.bfloat16, device=device) + hc = hn.reshape(w, HID).contiguous() + rc = fvk.w4a16_mrows_edge_sm120_bf16( + hc.data_ptr(), p['lm_head_packed_t'].data_ptr(), + p['lm_head_sf_t'].data_ptr(), logits.data_ptr(), + w, vocab, HID, p['lm_head_alpha'], _cs()) + if rc: + raise RuntimeError(f'M-row lm_head failed with {rc} at M={w}') + return logits, hidden + + +def _verify_block_usable(state) -> bool: + """Can the window run on the decode kernels? + + The M-row GEMV stages a window's activations in shared memory, so it has a + width limit; and the whole point is that the window reads what decode + reads, which is only true where decode takes the W4A16 dense path over + BF16-scope weights. Anywhere else the prefill forward is still the answer. + """ + if not _VERIFY_K_ROWS or not state.dense_w4a16: + return False + ld = state.handles.ptrs['layers'][0] + return (ld.get('router_packed') is None + and ld.get('out_proj_packed') is None + and not ld.get('experts_streamed')) + + def _spec_block(state, pos, k, fvk, device): """The whole window as one dependency chain: k drafts, then the verify. @@ -883,16 +1204,20 @@ def _spec_block(state, pos, k, fvk, device): fvk.qwen36_argmax_bf16(d_logits.data_ptr(), toks[j + 1:j + 2].data_ptr(), 1, vocab, _cs()) - state.spec_capture = True - set_spec_verify(True) - try: - logits, hid = nexn2_forward_nvfp4( - state.handles, toks[:window].view(1, window), fvk, device, - cap=state, pos_offset=pos, last_logits_only=False, - return_hidden=True) - finally: - state.spec_capture = False - set_spec_verify(False) + if _verify_block_usable(state): + logits, hid = _verify_block_K(state, toks[:window], pos, window, + fvk, device) + else: + state.spec_capture = True + set_spec_verify(True) + try: + logits, hid = nexn2_forward_nvfp4( + state.handles, toks[:window].view(1, window), fvk, device, + cap=state, pos_offset=pos, last_logits_only=False, + return_hidden=True) + finally: + state.spec_capture = False + set_spec_verify(False) logits = logits.reshape(window, -1) fvk.qwen36_argmax_bf16(logits.data_ptr(), state._spec_argmax.data_ptr(), window, vocab, _cs()) From 69a532fbd98e5fe184630535c7fde9c343c8c04f Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Sat, 1 Aug 2026 07:40:06 -0400 Subject: [PATCH 65/85] Sum a token's experts in a fixed order in every prefill path The same prompt did not decode to the same text twice. Four seedings of one 20-token prompt in one process: the logits differ in 232226 of 248320 elements, the GDN recurrent state differs from linear rank 12 on, the KV differs. The decode loop is fine -- eager and captured agree with each other and with themselves over four runs -- so all of it comes from the seed. Three of the four MoE paths finished with out.index_add_(0, stok, d_dn.float() * sw.unsqueeze(-1)) which reduces through atomics: a token's eight expert outputs are added in whatever order the blocks retire, and fp32 addition is not associative. Two of them also sorted the routing with a plain argsort, so equal-expert ties changed which rows were packed into a quantisation tile. Neither is a new discovery in this file. The grouped-GEMM path already inverts the permutation and sums with a kernel, and says why in a comment; the block-tile path already sorts stably and says why. The other three had simply never been brought along. They are now: invert the routing permutation, hand the inverse to moe_weighted_sum_sm120_bf16 so each token's slots are summed in k order, and sort stably everywhere. It is not slower -- the comment on the path that was already converted records index_add_ at 37.8 ms of a 1024-token prefill. Worth knowing which prompts were affected: above 64 tokens the prefill already took a converted path, so this bit short prompts, which is where a seeded generation starts and where the fixture lives. After: the seed is reproducible, plain greedy is reproducible over four runs, and speculative decode emits exactly what plain greedy emits. Golden fixture 16/16 throughout. --- .../frontends/torch/_nexn2_rtx_forward.py | 49 ++++++++++++++----- 1 file changed, 38 insertions(+), 11 deletions(-) diff --git a/flash_rt/frontends/torch/_nexn2_rtx_forward.py b/flash_rt/frontends/torch/_nexn2_rtx_forward.py index 0fb8078b..31c8fc5f 100644 --- a/flash_rt/frontends/torch/_nexn2_rtx_forward.py +++ b/flash_rt/frontends/torch/_nexn2_rtx_forward.py @@ -998,10 +998,11 @@ def _moe_experts_m16(x, ti, tw, ld, fvk, device): exp_flat = ti.reshape(-1).to(torch.int32) tok_flat = torch.arange(S, device=device).repeat_interleave(TOPK) - order = exp_flat.argsort() + # Stable, so equal-expert ties keep token order and the rows packed into + # each quantisation tile are the same run to run. + order = exp_flat.argsort(stable=True) se = exp_flat[order].long() stok = tok_flat[order] - sw = tw.reshape(-1)[order] counts = torch.bincount(se, minlength=E) tile_counts = (counts + 15) // 16 tile_off = torch.cumsum(tile_counts, 0) - tile_counts @@ -1027,8 +1028,18 @@ def _moe_experts_m16(x, ti, tw, ld, fvk, device): ip.data_ptr(), dn_p.data_ptr(), isf.data_ptr(), dn_s.data_ptr(), d_dn.data_ptr(), dn_a.data_ptr(), tile_expert.data_ptr(), total_tiles, n_dn, INTER, 0, dn_p[0].numel(), dn_s[0].numel(), _cs()) - out = torch.zeros(S, HID, device=device) - out.index_add_(0, stok, d_dn[tiled_row].float() * sw.unsqueeze(-1)) + # Deterministic unpermute, as the grouped-GEMM path does: one kernel sums + # each token's TOPK rows in a fixed order. Slot i of the token-major + # routing sits at sorted position inv[i], whose output row is tiled_row of + # that position. + inv = torch.empty(S * TOPK, dtype=torch.long, device=device) + inv[order] = torch.arange(S * TOPK, device=device) + rows = tiled_row[inv].to(torch.int32).contiguous() + twc = tw.contiguous() + out = torch.empty(S, HID, dtype=torch.float32, device=device) + fvk.moe_weighted_sum_sm120_bf16( + d_dn.data_ptr(), rows.data_ptr(), twc.data_ptr(), + out.data_ptr(), S, TOPK, n_dn, n_dn, _cs()) return out @@ -1388,7 +1399,6 @@ def _moe_experts_per_expert_gemm(x, ti, tw, ld, fvk, device): order = exp_flat.argsort(stable=True) se = exp_flat[order] stok = tok_flat[order] - sw = tw.reshape(-1)[order] counts = torch.bincount(se, minlength=_N_EXPERTS).tolist() slots = S * TOPK @@ -1420,8 +1430,15 @@ def _moe_experts_per_expert_gemm(x, ti, tw, ld, fvk, device): dn_a[e], cnt, n_dn, INTER, fvk, device, _cs(), out=d_dn[off_e:off_e + cnt]) - out = torch.zeros(S, HID, device=device) - out.index_add_(0, stok, d_dn.float() * sw.unsqueeze(-1)) + # Deterministic unpermute, as the grouped-GEMM path does; index_add_ + # reduces through atomics, so its order varies run to run. + inv = torch.empty(S * TOPK, dtype=torch.int32, device=device) + inv[order] = torch.arange(S * TOPK, dtype=torch.int32, device=device) + twc = tw.contiguous() + out = torch.empty(S, HID, dtype=torch.float32, device=device) + fvk.moe_weighted_sum_sm120_bf16( + d_dn.data_ptr(), inv.data_ptr(), twc.data_ptr(), + out.data_ptr(), S, TOPK, n_dn, n_dn, _cs()) return out @@ -1445,10 +1462,10 @@ def _moe_experts_grouped(x, ti, tw, ld, fvk, device): slots = S * TOPK exp_flat = ti.reshape(-1).to(torch.int32) tok_flat = torch.arange(S, device=device).repeat_interleave(TOPK) - order = exp_flat.argsort() + # Stable, so equal-expert ties keep token order run to run. + order = exp_flat.argsort(stable=True) se = exp_flat[order].contiguous() stok = tok_flat[order] - sw = tw.reshape(-1)[order] A = x[stok].contiguous() # (slots, HID) bf16 d_gu = torch.empty(slots, n_gu, dtype=torch.bfloat16, device=device) @@ -1463,8 +1480,18 @@ def _moe_experts_grouped(x, ti, tw, ld, fvk, device): inter.data_ptr(), dn_p.data_ptr(), dn_s.data_ptr(), dn_a.data_ptr(), se.data_ptr(), d_dn.data_ptr(), slots, n_dn, INTER, INTER, dn_p[0].numel(), dn_s[0].numel(), _cs()) - out = torch.zeros(S, HID, device=device) - out.index_add_(0, stok, d_dn.float() * sw.unsqueeze(-1)) + # Deterministic unpermute, as the grouped-GEMM path does: invert the + # routing permutation and let one kernel sum each token's TOPK rows in a + # fixed order. index_add_ reduces through atomics, so eight fp32 addends + # land in whatever order the blocks retire -- and a prefill cannot afford + # that, because it seeds a decode that has to be reproducible. + inv = torch.empty(slots, dtype=torch.int32, device=device) + inv[order] = torch.arange(slots, dtype=torch.int32, device=device) + twc = tw.contiguous() + out = torch.empty(S, HID, dtype=torch.float32, device=device) + fvk.moe_weighted_sum_sm120_bf16( + d_dn.data_ptr(), inv.data_ptr(), twc.data_ptr(), + out.data_ptr(), S, TOPK, n_dn, n_dn, _cs()) return out From 3b17f98d698e41b4b610f60560f326db28a17b83 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Sat, 1 Aug 2026 07:56:07 -0400 Subject: [PATCH 66/85] Decline the K-row window wherever decode would compute something else The window's whole claim is that it is the same function as the step it verifies, so the conditions under which that holds should be asked about rather than assumed. Two were not. The GDN in_proj is gated by its own flag, separate from the rest of the dense path. With it off, decode reads that projection at BF16 and the window reads it at four bits -- in thirty of the forty layers. The window also fuses the router with the shared gate/up and reads every projection at four bits, which is what decode does only when the loader kept those weights BF16. One NVFP4 site among them and decode takes the W4A4 mma instead. Only two of the three were checked. Both fall back to the prefill forward, which is what the flag is for. --- flash_rt/frontends/torch/_nexn2_rtx_decode.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/flash_rt/frontends/torch/_nexn2_rtx_decode.py b/flash_rt/frontends/torch/_nexn2_rtx_decode.py index d4a5f601..28e914ad 100644 --- a/flash_rt/frontends/torch/_nexn2_rtx_decode.py +++ b/flash_rt/frontends/torch/_nexn2_rtx_decode.py @@ -1177,10 +1177,21 @@ def _verify_block_usable(state) -> bool: reads, which is only true where decode takes the W4A16 dense path over BF16-scope weights. Anywhere else the prefill forward is still the answer. """ + # gdn_in_proj_w4a16 is gated separately from the rest of the dense path, so + # with it off decode reads the GDN in_proj at BF16 while the window reads + # it at four bits -- a different function in thirty of the forty layers, + # which is exactly the thing this block exists to rule out. if not _VERIFY_K_ROWS or not state.dense_w4a16: return False + if not state.gdn_in_proj_w4a16: + return False + # The window fuses the router with the shared gate/up and reads every + # projection at four bits, which is what decode does only when the loader + # kept these BF16. One NVFP4 site among them and decode takes the W4A4 mma + # instead, so ask about each of the three the window assumes. ld = state.handles.ptrs['layers'][0] return (ld.get('router_packed') is None + and ld.get('shared_gate_proj_packed') is None and ld.get('out_proj_packed') is None and not ld.get('experts_streamed')) From a3f6ef74c0e7da6d5eb12f8c774b2de833c57044 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Sat, 1 Aug 2026 11:33:20 -0400 Subject: [PATCH 67/85] Derive the GDN gating constants once, not every step A_log and dt_bias are weights. -exp(A_log) and the fp32 bias are therefore the same on every decode step, and they were being rebuilt on every one of them, in each of the thirty linear-attention layers: a cast, an exp, a negate and a contiguous, all inside the captured region, all producing the bytes the previous replay had already produced. The profiler counts them. The elementwise/copy bucket of a step goes from 522 launches to 403 and from 791.0 to 623.3 microseconds, and the step from 11.238 to 11.059 ms. Small, but it is 119 launches of a step that is 99% kernel time, where each launch costs its dispatch quantum whether or not it computes anything. Same expressions in the same order, so the kernel is handed the same bytes: golden fixture 16/16, and the speculative window still emits plain greedy's sequence token for token. --- flash_rt/frontends/torch/_nexn2_rtx_decode.py | 27 ++++++++++++++----- 1 file changed, 21 insertions(+), 6 deletions(-) diff --git a/flash_rt/frontends/torch/_nexn2_rtx_decode.py b/flash_rt/frontends/torch/_nexn2_rtx_decode.py index 28e914ad..c4a3e66a 100644 --- a/flash_rt/frontends/torch/_nexn2_rtx_decode.py +++ b/flash_rt/frontends/torch/_nexn2_rtx_decode.py @@ -359,6 +359,25 @@ def reset(self): self.attn.reset_cache() +def _gdn_gate_consts(ld, device): + """The gating kernel's two constant inputs, derived once per layer. + + ``A_log`` and ``dt_bias`` are weights, so -exp(A_log) and the fp32 bias are + the same on every step -- but deriving them per call put four elementwise + launches per GDN layer inside the captured region, thirty layers of them, + recomputing values identical to the previous replay's. Each is a couple of + microseconds of dispatch quantum for no arithmetic anyone reads. + + Same expressions in the same order, so the bytes handed to the kernel are + the bytes it was getting before. + """ + if 'gdn_neg_exp_a' not in ld: + ld['gdn_neg_exp_a'] = ( + -ld['A_log_t'].float().exp()).float().contiguous() + ld['gdn_dt_bias_f'] = ld['dt_bias_t'].float().contiguous() + return ld['gdn_neg_exp_a'], ld['gdn_dt_bias_f'] + + def _decode_gdn(h, ld, state, lin_rank, fvk, device): """GDN layer at one token, updating recurrent + conv state in place.""" eps = state.eps @@ -366,7 +385,7 @@ def _decode_gdn(h, ld, state, lin_rank, fvk, device): Wz = ld['in_proj_z_w_t'] Wb, Wa = ld['in_proj_b_w_t'], ld['in_proj_a_w_t'] convw = ld['conv1d_w_t'].reshape(CONV, KS).contiguous() - A_log, dtb = ld['A_log_t'].float(), ld['dt_bias_t'].float() + neg, dtb_c = _gdn_gate_consts(ld, device) nw = ld['gdn_norm_w_t'] s = _cs() @@ -403,8 +422,6 @@ def _decode_gdn(h, ld, state, lin_rank, fvk, device): conv_out.data_ptr(), qb.data_ptr(), kb.data_ptr(), vb.data_ptr(), 1, s) - neg = (-A_log.exp()).float().contiguous() - dtb_c = dtb.contiguous() g_out = torch.empty(1, NV, dtype=torch.bfloat16, device=device) bo = torch.empty(1, NV, dtype=torch.bfloat16, device=device) fvk.qwen36_gdn_gating_bf16( @@ -907,7 +924,7 @@ def _verify_gdn(h, ld, state, lin_rank, w, fvk, device): """ eps = state.eps convw = ld['conv1d_w_t'].reshape(CONV, KS).contiguous() - A_log, dtb = ld['A_log_t'].float(), ld['dt_bias_t'].float() + neg, dtb_c = _gdn_gate_consts(ld, device) nw = ld['gdn_norm_w_t'] s = _cs() x = h.reshape(w, HID) @@ -942,8 +959,6 @@ def _verify_gdn(h, ld, state, lin_rank, w, fvk, device): conv_out.data_ptr(), qb.data_ptr(), kb.data_ptr(), vb.data_ptr(), w, s) - neg = (-A_log.exp()).float().contiguous() - dtb_c = dtb.contiguous() g_out = torch.empty(w, NV, dtype=torch.bfloat16, device=device) bo = torch.empty(w, NV, dtype=torch.bfloat16, device=device) fvk.qwen36_gdn_gating_bf16( From 2436b4f48e1c22635733e0c583ff836385a8f54d Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Sat, 1 Aug 2026 11:55:27 -0400 Subject: [PATCH 68/85] Give decode the fused shared-gate combine, and make it exact first The MoE tail -- sigmoid the shared expert's gate, scale the shared output by it, add the routed sum, round to bf16 -- was five tensor ops a layer in decode and in the speculative window, forty layers of them. A kernel that does all of it already existed and prefill was already calling it; decode was not, and could not, because it was not the same arithmetic. The reason is worth recording. Written as routed[i] + float(shared[i]) * g the compiler contracts the multiply and the add into one fma. That is one rounding where the tensor-op chain has two, so the kernel and the chain disagree -- measured, one element in 16384 by one ulp, at scales where the routed sum is small against the gated shared term. More accurate, and still wrong for this purpose: decode's output is compared token for token against a fixture, and the speculative verify may only keep a token because it computed what the decode step would have. So the kernel now multiplies and adds as two rounded operations, and is bit-identical to the chain at every shape and scale the two paths issue. That also removes a disagreement nobody had noticed, since prefill had been using the contracted form all along while decode used the chain. Per decode step, measured: the elementwise bucket goes from 403 launches and 623.3 us to 203 and 301.5, and the step from 11.059 to 10.827 ms -- 92.4 tok/s against 90.4. Speculative decode at K=2 reaches 97.34 tok/s against 87.98 plain. Golden fixture 16/16, the window still emits plain greedy's sequence token for token, and the window's rows are still bit-identical to the decode steps. --- csrc/kernels/moe_shared_combine_edge.cu | 12 +++++-- flash_rt/frontends/torch/_nexn2_rtx_decode.py | 36 +++++++++++++++---- 2 files changed, 39 insertions(+), 9 deletions(-) diff --git a/csrc/kernels/moe_shared_combine_edge.cu b/csrc/kernels/moe_shared_combine_edge.cu index 15ee59e5..ed7891b0 100644 --- a/csrc/kernels/moe_shared_combine_edge.cu +++ b/csrc/kernels/moe_shared_combine_edge.cu @@ -22,8 +22,16 @@ __global__ void moe_shared_gate_combine_kernel( const float g = 1.0f / (1.0f + expf(-static_cast(gate[row]))); const size_t base = static_cast(row) * dim; for (int i = threadIdx.x; i < dim; i += blockDim.x) { - out[base + i] = __float2bfloat16( - routed[base + i] + static_cast(shared[base + i]) * g); + // Multiply and add as two rounded operations, not one contracted fma. + // Written as `routed + shared * g` the compiler contracts it, which is one + // rounding instead of two and therefore a different number -- measured, one + // element in 16384 by one ulp, where the routed sum is small against the + // gated shared term. That is a fine trade in isolation and the wrong one + // here: the decode step computes this as separate tensor ops, and this + // kernel is only allowed to stand in for it if it lands on the same bits. + out[base + i] = __float2bfloat16(__fadd_rn( + routed[base + i], + __fmul_rn(static_cast(shared[base + i]), g))); } } diff --git a/flash_rt/frontends/torch/_nexn2_rtx_decode.py b/flash_rt/frontends/torch/_nexn2_rtx_decode.py index c4a3e66a..d9600690 100644 --- a/flash_rt/frontends/torch/_nexn2_rtx_decode.py +++ b/flash_rt/frontends/torch/_nexn2_rtx_decode.py @@ -581,6 +581,28 @@ def _moe_experts_streamed(x, idx, state, fvk, device, s): return d_dn +def _shared_combine(routed, shared, glog, rows, fvk, device): + """out = routed(fp32) + shared(bf16) * sigmoid(gate), in one kernel. + + The tensor-op form is a cast, a sigmoid, a broadcast multiply, an add and a + cast: five launches a layer, forty layers, in a step that is 99% kernel + time and where a launch costs its dispatch quantum whether or not it + computes much. The kernel does the same arithmetic in the same order and + rounds once at the store, so it stands in for the chain rather than + approximating it -- checked bit for bit at the shapes and scales decode and + the window issue, because the fixture and the speculative verify both rest + on it. + """ + if hasattr(fvk, 'moe_shared_gate_combine_edge_bf16'): + out = torch.empty(rows, HID, dtype=torch.bfloat16, device=device) + fvk.moe_shared_gate_combine_edge_bf16( + routed.data_ptr(), shared.data_ptr(), glog.data_ptr(), + out.data_ptr(), rows, HID, _cs()) + return out + sgate = torch.sigmoid(glog.float()).reshape(rows, 1) + return (routed + shared.float() * sgate).to(torch.bfloat16) + + def _moe_layer_decode(h, ld, state, fvk, device): """M=1 fine-grained MoE via the grouped GEMV kernel: the 8 routed experts run in one launch each for gate_up (shared act) and down (per-slot act), @@ -696,10 +718,11 @@ def _moe_layer_decode(h, ld, state, fvk, device): si = _silu_mul(sg, su, fvk, device) shared = _proj_mma(si, ld, 'shared_down_proj', HID, fvk, device, state) # shared-expert scalar gate: N=1 GEMV via the bf16 matvec kernel (was a - # torch matmul -- the last fp32 matmul in the captured decode step). - sgate = torch.sigmoid( - _bf16_mv(x, ld['shared_gate_w_t'], fvk, device).float()) - return (out + shared.float() * sgate).reshape(1, 1, HID).to(torch.bfloat16) + # torch matmul -- the last fp32 matmul in the captured decode step). The + # sigmoid, the broadcast multiply, the add and the cast are one kernel. + glog = _bf16_mv(x, ld['shared_gate_w_t'], fvk, device) + return _shared_combine(out, shared, glog, 1, fvk, device).reshape( + 1, 1, HID) def decode_step(state, token_id, pos, fvk, device): @@ -1132,9 +1155,8 @@ def _verify_moe(h, ld, state, w, fvk, device): fvk.bf16_matvec_sm120_bf16( xr.data_ptr(), ld['shared_gate_w_t'].data_ptr(), gsc[t].data_ptr(), 1, HID, s) - sgate = torch.sigmoid(gsc.float()) - return (out + shared.float() * sgate).reshape( - 1, w, HID).to(torch.bfloat16) + return _shared_combine(out, shared, gsc, w, fvk, device).reshape( + 1, w, HID) def _verify_block_K(state, toks, pos, w, fvk, device): From 45b767e1d63d34e74fcd8a9b9ea2fcf85b8bea23 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Sat, 1 Aug 2026 12:34:13 -0400 Subject: [PATCH 69/85] Stop the GDN recurrence spilling its state column to local memory ncu reports the single-token recurrence running at 39 registers per thread. The kernel declares float col[HD]; // HD = 128 and walks it five times. A 128-iteration loop unrolled by sixteen leaves the index non-constant, so that array cannot be held in registers and is not: it is in local memory, 512 bytes a thread, read and written across every pass. Against 2 MB of actual state traffic that is several megabytes of spill, and the kernel lands at 51% of what the part can move while the profile shows 12.8% occupancy and 0.13 waves per SM. The column does not need to be held. The recurrence touches the state twice -- once to form the k-weighted sum, once to apply the rank-one update and emit the output -- and the whole state is 1 MB, so the second read is an L2 hit. This is the same kernel with the intermediate left where it belongs. Every accumulation runs in the same order over the same rounded fp32 values, so the two agree exactly: six independent trials of eight chained steps each, outputs and final state compared with torch.equal, not a tolerance. Chained, because a recurrence that drifts by an ulp would hide in a single step. shipped 16.77 us 125 GB/s 51% of peak edge 9.89 us 212 GB/s 87% 1.70x Added alongside the existing entry rather than replacing it, and selected through a dispatch helper, so a build without it keeps working. End to end: the decode step goes 10.827 -> 10.379 ms, 92.4 -> 96.4 tok/s, with the GDN bucket 746.5 -> 557.8 us. Plain greedy over 128 tokens reads 91.18 tok/s against 87.98. Golden fixture 16/16, the speculative window still emits plain greedy's sequence token for token, and its rows are still bit-identical to the decode steps they stand in for. --- csrc/bindings.cpp | 20 +++ csrc/kernels/gated_deltanet_qwen36.cu | 126 ++++++++++++++++++ csrc/kernels/gated_deltanet_qwen36.cuh | 15 +++ flash_rt/frontends/torch/_nexn2_rtx_decode.py | 18 ++- 4 files changed, 177 insertions(+), 2 deletions(-) diff --git a/csrc/bindings.cpp b/csrc/bindings.cpp index 2ddb142e..4b01a412 100644 --- a/csrc/bindings.cpp +++ b/csrc/bindings.cpp @@ -5319,6 +5319,26 @@ PYBIND11_MODULE(flash_rt_kernels, m) { py::arg("head_k_dim"), py::arg("head_v_dim"), py::arg("use_qk_l2norm") = true, py::arg("stream") = 0); + m.def("gated_deltanet_recurrent_edge_qwen36_bf16", + [](uintptr_t q, uintptr_t k, uintptr_t v, + uintptr_t g, uintptr_t beta, + uintptr_t state, uintptr_t out, + int B, int num_v_heads, int head_k_dim, int head_v_dim, + bool use_qk_l2norm, uintptr_t stream) { + flash_rt::kernels::gated_deltanet_recurrent_edge_qwen36_bf16( + to_ptr(q), to_ptr(k), to_ptr(v), + to_ptr(g), to_ptr(beta), + to_ptr(state), to_ptr(out), + B, num_v_heads, head_k_dim, head_v_dim, + use_qk_l2norm, to_stream(stream)); + }, + py::arg("q"), py::arg("k"), py::arg("v"), + py::arg("g"), py::arg("beta"), + py::arg("state"), py::arg("out"), + py::arg("B"), py::arg("num_v_heads"), + py::arg("head_k_dim"), py::arg("head_v_dim"), + py::arg("use_qk_l2norm") = true, py::arg("stream") = 0); + // In/out-state variant for K-iter chained per-step save (A2c-3). m.def("gated_deltanet_recurrent_inout_qwen36_bf16", [](uintptr_t q, uintptr_t k, uintptr_t v, diff --git a/csrc/kernels/gated_deltanet_qwen36.cu b/csrc/kernels/gated_deltanet_qwen36.cu index c80754c7..482c0075 100644 --- a/csrc/kernels/gated_deltanet_qwen36.cu +++ b/csrc/kernels/gated_deltanet_qwen36.cu @@ -150,6 +150,132 @@ __global__ void gated_deltanet_recurrent_kernel( __float2bfloat16(out_t); } + +// Same recurrence, without the local-memory round trip. +// +// The kernel above keeps the thread's whole state column in `float col[HD]`. +// A 128-iteration loop unrolled by 16 leaves the index non-constant, so the +// array cannot live in registers -- ncu measures 39 registers per thread for a +// 128-float array, which means it is in local memory, read and written across +// five passes. Against 2 MB of real state traffic that is roughly 6 MB of +// spill, and the kernel lands at 108 GB/s on a 244 GB/s part. +// +// The column never needs to be held. The recurrence reads the state twice -- +// once to form the k-weighted sum, once to update and emit -- and the whole +// state is 1 MB, so the second read is an L2 hit. Arithmetic, and the order of +// every accumulation, is identical to the kernel above; only where the +// intermediate lives changes. +template +__global__ void gated_deltanet_recurrent_edge_kernel( + const __nv_bfloat16* __restrict__ q_in, + const __nv_bfloat16* __restrict__ k_in, + const __nv_bfloat16* __restrict__ v_in, + const __nv_bfloat16* __restrict__ g_in, + const __nv_bfloat16* __restrict__ beta_in, + __nv_bfloat16* __restrict__ state, + __nv_bfloat16* __restrict__ out_, + int num_v_heads, + bool use_qk_l2norm) +{ + static_assert(HD == 128, "HD must be 128 for Qwen3.6 (single instantiation)"); + const int h = blockIdx.x; + const int b = blockIdx.y; + const int t = threadIdx.x; + if (t >= HD) return; + + __shared__ float smem[2 * HD + 32]; + float* qs = smem; + float* ks = smem + HD; + float* scratch = smem + 2 * HD; + + const size_t qkv_off = ((size_t)b * num_v_heads + h) * HD + t; + qs[t] = static_cast(q_in[qkv_off]); + ks[t] = static_cast(k_in[qkv_off]); + __syncthreads(); + + if (use_qk_l2norm) { + float q_sq = qs[t] * qs[t]; + float k_sq = ks[t] * ks[t]; + q_sq = block_reduce_sum(q_sq, scratch); + __syncthreads(); + k_sq = block_reduce_sum(k_sq, scratch); + const float q_inv = rsqrtf(q_sq + kEps); + const float k_inv = rsqrtf(k_sq + kEps); + qs[t] *= q_inv; + ks[t] *= k_inv; + __syncthreads(); + } + + qs[t] *= rsqrtf(static_cast(HD)); + __syncthreads(); + + const float g_t = + __expf(static_cast(g_in[b * num_v_heads + h])); + const float beta_t = + static_cast(beta_in[b * num_v_heads + h]); + + const size_t state_h_off = (((size_t)b * num_v_heads + h)) * HD * HD; + + // kv_mem[t] = sum_i (state[i][t] * g_t) * ks[i], accumulated in i order -- + // the same order, and the same rounded product, as the version that stored + // the column first. + float kv_mem = 0.0f; + #pragma unroll 16 + for (int i = 0; i < HD; ++i) { + const float c = + static_cast(state[state_h_off + (size_t)i * HD + t]) * g_t; + kv_mem = fmaf(c, ks[i], kv_mem); + } + + const float v_t = + static_cast(v_in[(size_t)b * num_v_heads * HD + h * HD + t]); + const float delta = (v_t - kv_mem) * beta_t; + + // Second pass: re-derive the decayed column, apply the rank-one update, + // store it, and accumulate the output -- one read and one write per element. + float out_t = 0.0f; + #pragma unroll 16 + for (int i = 0; i < HD; ++i) { + const size_t off = state_h_off + (size_t)i * HD + t; + const float c = fmaf(ks[i], delta, static_cast(state[off]) * g_t); + state[off] = __float2bfloat16(c); + out_t = fmaf(c, qs[i], out_t); + } + out_[(size_t)b * num_v_heads * HD + h * HD + t] = + __float2bfloat16(out_t); +} + +} // namespace + +void gated_deltanet_recurrent_edge_qwen36_bf16( + const void* q, + const void* k, + const void* v, + const void* g, + const void* beta, + void* state, + void* out, + int B, int num_v_heads, int head_k_dim, int head_v_dim, + bool use_qk_l2norm, + cudaStream_t stream) +{ + constexpr int kHD = 128; + if (head_k_dim != kHD || head_v_dim != kHD) return; + dim3 grid(num_v_heads, B); + dim3 block(kHD); + gated_deltanet_recurrent_edge_kernel<<>>( + reinterpret_cast(q), + reinterpret_cast(k), + reinterpret_cast(v), + reinterpret_cast(g), + reinterpret_cast(beta), + reinterpret_cast<__nv_bfloat16*>(state), + reinterpret_cast<__nv_bfloat16*>(out), + num_v_heads, use_qk_l2norm); +} + +namespace { + } // namespace void gated_deltanet_recurrent_qwen36_bf16( diff --git a/csrc/kernels/gated_deltanet_qwen36.cuh b/csrc/kernels/gated_deltanet_qwen36.cuh index f497eff9..a6a12268 100644 --- a/csrc/kernels/gated_deltanet_qwen36.cuh +++ b/csrc/kernels/gated_deltanet_qwen36.cuh @@ -60,6 +60,21 @@ void gated_deltanet_recurrent_qwen36_bf16( // to state_out (different buffer). Caller chains state_in[k+1] := // state_out[k] to support per-step state save without an extra // .copy_(state_save, state) launch per step. +// Spill-free variant of the above: identical arithmetic and accumulation +// order, but the thread's state column is re-read rather than held in a +// 128-float local array. Same arguments, same results, bit for bit. +void gated_deltanet_recurrent_edge_qwen36_bf16( + const void* q, + const void* k, + const void* v, + const void* g, + const void* beta, + void* state, + void* out, + int B, int num_v_heads, int head_k_dim, int head_v_dim, + bool use_qk_l2norm, + cudaStream_t stream); + void gated_deltanet_recurrent_inout_qwen36_bf16( const void* q, const void* k, diff --git a/flash_rt/frontends/torch/_nexn2_rtx_decode.py b/flash_rt/frontends/torch/_nexn2_rtx_decode.py index d9600690..16ce2d56 100644 --- a/flash_rt/frontends/torch/_nexn2_rtx_decode.py +++ b/flash_rt/frontends/torch/_nexn2_rtx_decode.py @@ -359,6 +359,20 @@ def reset(self): self.attn.reset_cache() +def gdn_recurrent(fvk): + """The single-token GDN recurrence entry this build should call. + + The edge variant is the same arithmetic in the same order -- checked + exactly, over chained steps so the state drift is exercised too -- and it is + 1.70x the shipped one. The shipped one holds the thread's whole state column + in a 128-float array that cannot live in registers, so it is in local + memory and walked five times; ncu measures 39 registers per thread for a + 128-float array. 51% of bandwidth against 87%. + """ + fn = getattr(fvk, 'gated_deltanet_recurrent_edge_qwen36_bf16', None) + return fn if fn is not None else fvk.gated_deltanet_recurrent_qwen36_bf16 + + def _gdn_gate_consts(ld, device): """The gating kernel's two constant inputs, derived once per layer. @@ -434,7 +448,7 @@ def _decode_gdn(h, ld, state, lin_rank, fvk, device): gt = g_out.reshape(NV).contiguous() bt = bo.reshape(NV).contiguous() core = torch.empty(NV, HV, dtype=torch.bfloat16, device=device) - fvk.gated_deltanet_recurrent_qwen36_bf16( + gdn_recurrent(fvk)( qt.data_ptr(), kt.data_ptr(), vt.data_ptr(), gt.data_ptr(), bt.data_ptr(), state.lin_state[lin_rank].data_ptr(), core.data_ptr(), 1, NV, HK, HV, True, s) @@ -993,7 +1007,7 @@ def _verify_gdn(h, ld, state, lin_rank, w, fvk, device): for t in range(w): qt, kt, vt = qb[t], kb[t], vb[t] gt, bt = g_out[t], bo[t] - fvk.gated_deltanet_recurrent_qwen36_bf16( + gdn_recurrent(fvk)( qt.data_ptr(), kt.data_ptr(), vt.data_ptr(), gt.data_ptr(), bt.data_ptr(), lin_state.data_ptr(), core[t].data_ptr(), 1, NV, HK, HV, True, s) From ec62002b07d359f7dde234d7a0003709618c530c Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Sat, 1 Aug 2026 12:47:55 -0400 Subject: [PATCH 70/85] Pick the router's top-8 in one warp instead of 24 barriers Selecting eight of 256 logits was costing 6.5 us to read 512 bytes. The kernel spreads the logits over 256 threads and runs k rounds, each with a warp reduction and three block-wide barriers to publish the winner and mask it -- 24 barriers for 512 bytes of input. One warp can hold all of it: eight logits a lane in registers, the same k rounds entirely in shuffles, nothing synchronised and nothing in shared memory. The output is identical, not merely equivalent, and the reason is worth stating because it does not hold for most reductions: argmax under a total order -- greater value wins, lower index breaks ties -- selects one specific element, so the answer cannot depend on the shape of the reduction tree the way a floating-point sum does. Which lane holds which logit is therefore free to change. Checked anyway, on inputs built to produce the case that would break it: 800 trials across four widths, 397 of them with a tie inside the top-8, indices and values compared with torch.equal. No mismatches. Worth less than the barrier count suggested. Standalone it is 8.30 us against 6.60; in the captured step the routing bucket goes 594.8 -> 499.2 us, about 2.4 us a layer. The rest is the per-kernel dispatch floor, which no amount of making this kernel faster will remove -- only launching it fewer times would. Step 10.379 -> 10.297 ms, 96.4 -> 97.1 tok/s. Speculative K=2 reads 98.57 against 91.98 plain. Fixture 16/16, same emitted text. --- csrc/bindings.cpp | 10 +++ csrc/kernels/moe_router_topk_sm120.cu | 73 +++++++++++++++++++ csrc/kernels/moe_router_topk_sm120.cuh | 8 ++ flash_rt/frontends/torch/_nexn2_rtx_decode.py | 16 +++- 4 files changed, 105 insertions(+), 2 deletions(-) diff --git a/csrc/bindings.cpp b/csrc/bindings.cpp index 4b01a412..a02578a3 100644 --- a/csrc/bindings.cpp +++ b/csrc/bindings.cpp @@ -5513,6 +5513,16 @@ PYBIND11_MODULE(flash_rt_kernels, m) { py::arg("M"), py::arg("N"), py::arg("K"), py::arg("alpha") = 1.0f, py::arg("stream") = 0); + m.def("moe_router_topk_warp_sm120_bf16", + [](uintptr_t logits, uintptr_t out_idx, uintptr_t out_val, + int n_experts, int k, uintptr_t stream) { + return flash_rt::kernels::moe_router_topk_warp_sm120_bf16( + to_ptr(logits), to_ptr(out_idx), to_ptr(out_val), + n_experts, k, to_stream(stream)); + }, + py::arg("logits"), py::arg("out_idx"), py::arg("out_val"), + py::arg("n_experts"), py::arg("k"), py::arg("stream") = 0); + m.def("moe_router_topk_sm120_bf16", [](uintptr_t logits, uintptr_t out_idx, uintptr_t out_val, int n_experts, int k, uintptr_t stream) -> int { diff --git a/csrc/kernels/moe_router_topk_sm120.cu b/csrc/kernels/moe_router_topk_sm120.cu index e3b9e164..1ada0ab5 100644 --- a/csrc/kernels/moe_router_topk_sm120.cu +++ b/csrc/kernels/moe_router_topk_sm120.cu @@ -69,6 +69,79 @@ __global__ void router_topk_kernel(const __nv_bfloat16* __restrict__ logits, } } +// One warp, no barriers. +// +// The kernel above spreads 256 logits over 256 threads and runs k rounds, each +// with a warp reduction plus three block-wide barriers to publish the winner +// and mask it. That is 24 barriers to pick eight of 256 values -- 6.5 us to +// read 512 bytes, which is latency, not work. +// +// Here one warp owns all the logits in registers, eight per lane, and the same +// k rounds run entirely in shuffles. Nothing to synchronise, nothing in shared +// memory. +// +// The result is identical rather than merely equivalent, and for a reason +// worth stating: argmax under a total order -- greater value wins, lower index +// breaks ties -- selects one specific element, so the answer does not depend +// on the shape of the reduction tree the way a floating-point sum does. +// Changing which lane holds which logit therefore cannot change the output. +template +__global__ void router_topk_warp1_kernel(const __nv_bfloat16* __restrict__ logits, + int* __restrict__ out_idx, + float* __restrict__ out_val, + int n, int k) { + const int lane = threadIdx.x; + float v[kPerLane]; +#pragma unroll + for (int j = 0; j < kPerLane; ++j) { + const int i = j * 32 + lane; + v[j] = (i < n) ? static_cast(logits[i]) : -FLT_MAX; + } + + for (int r = 0; r < k; ++r) { + float best = -FLT_MAX; + int bidx = -1; +#pragma unroll + for (int j = 0; j < kPerLane; ++j) { + const int i = j * 32 + lane; + if (v[j] > best || (v[j] == best && i < bidx)) { best = v[j]; bidx = i; } + } + warp_argmax(best, bidx); // butterfly: every lane ends with it + if (lane == 0) { + out_idx[r] = bidx; + out_val[r] = best; + } + // The lane that owns the winner clears it. bidx = j * 32 + owner, so the + // low five bits are the owner and the rest is the slot. + if (bidx >= 0 && (bidx & 31) == lane) v[bidx >> 5] = -FLT_MAX; + } +} + +} // namespace + +int moe_router_topk_warp_sm120_bf16(const void* logits, void* out_idx, + void* out_val, int n_experts, int k, + cudaStream_t stream) { + if (!logits || !out_idx || !out_val) return 1; + if (n_experts <= 0 || k <= 0 || k > 32) return 2; + auto* oi = reinterpret_cast(out_idx); + auto* ov = reinterpret_cast(out_val); + const auto* lg = reinterpret_cast(logits); + if (n_experts <= 32) + router_topk_warp1_kernel<1><<<1, 32, 0, stream>>>(lg, oi, ov, n_experts, k); + else if (n_experts <= 64) + router_topk_warp1_kernel<2><<<1, 32, 0, stream>>>(lg, oi, ov, n_experts, k); + else if (n_experts <= 128) + router_topk_warp1_kernel<4><<<1, 32, 0, stream>>>(lg, oi, ov, n_experts, k); + else if (n_experts <= 256) + router_topk_warp1_kernel<8><<<1, 32, 0, stream>>>(lg, oi, ov, n_experts, k); + else + return 3; // caller falls back to the block kernel + return 0; +} + +namespace { + } // namespace int moe_router_topk_sm120_bf16(const void* logits, void* out_idx, void* out_val, diff --git a/csrc/kernels/moe_router_topk_sm120.cuh b/csrc/kernels/moe_router_topk_sm120.cuh index cdabc16a..69bbe52b 100644 --- a/csrc/kernels/moe_router_topk_sm120.cuh +++ b/csrc/kernels/moe_router_topk_sm120.cuh @@ -18,5 +18,13 @@ namespace kernels { int moe_router_topk_sm120_bf16(const void* logits, void* out_idx, void* out_val, int n_experts, int k, cudaStream_t stream); +// Single-warp variant, for n_experts <= 256. Same selection rule and the same +// descending order, and identical output -- argmax under a total order does not +// depend on the reduction tree. Returns 3 for a width it cannot hold, so the +// caller can fall back to the block kernel above. +int moe_router_topk_warp_sm120_bf16(const void* logits, void* out_idx, + void* out_val, int n_experts, int k, + cudaStream_t stream); + } // namespace kernels } // namespace flash_rt diff --git a/flash_rt/frontends/torch/_nexn2_rtx_decode.py b/flash_rt/frontends/torch/_nexn2_rtx_decode.py index 16ce2d56..c6f96129 100644 --- a/flash_rt/frontends/torch/_nexn2_rtx_decode.py +++ b/flash_rt/frontends/torch/_nexn2_rtx_decode.py @@ -359,6 +359,18 @@ def reset(self): self.attn.reset_cache() +def router_topk(fvk): + """The router top-k entry this build should call. + + The warp variant returns identical indices and values -- argmax under a + total order picks one element regardless of the reduction tree, checked + over 800 inputs of which 397 had a tie inside the top-8 -- without the + block kernel's 24 barriers. + """ + fn = getattr(fvk, 'moe_router_topk_warp_sm120_bf16', None) + return fn if fn is not None else fvk.moe_router_topk_sm120_bf16 + + def gdn_recurrent(fvk): """The single-token GDN recurrence entry this build should call. @@ -653,7 +665,7 @@ def _moe_layer_decode(h, ld, state, fvk, device): # idx and topv come from torch.empty, so an unchecked failure here leaves # uninitialised memory to be used as expert indices -- which reaches a file # offset before anything notices. - rc = fvk.moe_router_topk_sm120_bf16( + rc = router_topk(fvk)( lr.data_ptr(), idx.data_ptr(), topv.data_ptr(), lr.numel(), TOPK, s) if rc: raise RuntimeError( @@ -1110,7 +1122,7 @@ def _verify_moe(h, ld, state, w, fvk, device): idx = torch.empty(w, TOPK, dtype=torch.int32, device=device) topv = torch.empty(w, TOPK, dtype=torch.float32, device=device) for t in range(w): - rc = fvk.moe_router_topk_sm120_bf16( + rc = router_topk(fvk)( logit_raw[t].data_ptr(), idx[t].data_ptr(), topv[t].data_ptr(), ne, TOPK, s) if rc: From 900d11516b0547cbb131789049998fb24a2c6e23 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Sat, 1 Aug 2026 13:09:43 -0400 Subject: [PATCH 71/85] Halve the W4A16 GEMV's loads in flight to buy back occupancy ncu on the grouped instance: Registers Per Thread 121-124 Block Limit Registers 8 <- the only binding limit Block Limit Shared Mem 18-25 Block Limit Warps / SM 24 / 24 Achieved Occupancy 31% Every limit except registers allows 24 blocks a SM. The registers are not an accident -- the loop deliberately keeps R * kUnroll eight-byte weight loads outstanding, because an earlier round measured the dominant stall as the global-load dependency with the ALU pipe at 27-50%. But wv[R][kUnroll] is R * kUnroll uint64 before anything else, 64 registers at R=8, and R had been swept while kUnroll never was. Halving it halves that array. Loads in flight per lane drop from R*4 to R*2, and twice as many lanes issue them; on this part the second is worth more. Swept in the captured decode step, one build each: kUnroll 4 3 2 step 10.384 10.360 9.743 ms tok/s 96.3 96.5 102.6 Bit-identical, and provably so rather than by luck: the main loop advances by 32*kUnroll and the tail takes the remainder, so a lane visits the same k-blocks in the same order for any kUnroll, and accumulates them in that order. Checked as well -- the speculative window's rows are still bit-identical to the decode steps, the kernel preflight passes 36/36, and the golden fixture is 16/16. Plain greedy over 128 tokens: 96.64 tok/s against 91.98. Speculative K=2 reads 100.66. The step's GEMV bucket goes 8392.8 -> 7783.5 us. --- csrc/kernels/w4a16_edge_sm120.cu | 21 ++++++++++++++++++++- 1 file changed, 20 insertions(+), 1 deletion(-) diff --git a/csrc/kernels/w4a16_edge_sm120.cu b/csrc/kernels/w4a16_edge_sm120.cu index da1cb9b5..a172b4b1 100644 --- a/csrc/kernels/w4a16_edge_sm120.cu +++ b/csrc/kernels/w4a16_edge_sm120.cu @@ -18,7 +18,26 @@ namespace { constexpr int kWarps = 2; // output-row groups per block constexpr int kThreads = kWarps * 32; // 256 -constexpr int kUnroll = 4; // packed-weight loads in flight +// Packed-weight loads in flight per row. Two, not four, and the reason is +// occupancy rather than parallelism: ncu puts this kernel at 121 registers a +// thread, which caps it at 8 blocks per SM when shared memory, warps and the +// SM limit all allow 24 -- Block Limit Registers is the only binding one, and +// achieved occupancy is 31%. wv[R][kUnroll] alone is R * kUnroll eight-byte +// values, 64 registers at R=8. +// +// Halving it halves that array and buys back the warps. The loads in flight +// per lane go from R * 4 to R * 2, but there are twice as many lanes issuing +// them, and the second is worth more here. Swept in the captured decode step, +// one build each: +// +// kUnroll 4 3 2 +// step 10.384 10.360 9.743 ms +// tok/s 96.3 96.5 102.6 +// +// The accumulation order does not move: the main loop advances by 32*kUnroll +// and the tail picks up the remainder, so a lane visits the same k-blocks in +// the same sequence for any kUnroll, and the result is bit-identical. +constexpr int kUnroll = 2; // A 16-element NVFP4 block is 16 bf16 of activation, 32 bytes. Held at that // stride, the eight lanes of a 128-bit shared-load phase land on banks From d2149cc0fff1f1fa018f3273bfdf1e463492e09c Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Sat, 1 Aug 2026 13:32:29 -0400 Subject: [PATCH 72/85] Give the M-row GEMV the loads-in-flight the single-row one settled on The speculative verify's dense projections run through an M-row form of the decode GEMV, written by copying the single-row kernel's shape constants. When the sweep moved that kernel from four packed-weight loads in flight to two, the copy kept four -- so decode had traded registers for warps and the verify was still paying for depth. It showed up as the speculative ratio sliding while plain greedy improved: the window was not getting the win the step got. Same constant, same reasoning, same invariance: the main loop advances by 32*kUnroll and the tail takes the remainder, so a lane visits the same k-blocks in the same order and every output row stays bit-identical to running the single-row GEMV once per row. Preflight 31/31 exact, the window's rows still bit-identical to the decode steps, fixture 16/16, and the emitted text is still plain greedy's. Measured with nothing else on the device, one process per point: K=1 105.22 tok/s against 96.75 plain K=2 106.74 against 100.35 --- csrc/kernels/w4a16_mrows_edge_sm120.cu | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/csrc/kernels/w4a16_mrows_edge_sm120.cu b/csrc/kernels/w4a16_mrows_edge_sm120.cu index 29d4a85e..de049edb 100644 --- a/csrc/kernels/w4a16_mrows_edge_sm120.cu +++ b/csrc/kernels/w4a16_mrows_edge_sm120.cu @@ -14,7 +14,13 @@ namespace { // blocks off each other's banks. constexpr int kWarps = 2; constexpr int kThreads = kWarps * 32; -constexpr int kUnroll = 4; +// Two, matching the M=1 entry: this kernel inherited 4 when it was written and +// kept it after the sweep moved the single-row path, which left the verify +// paying registers for depth while decode had already traded them for warps. +// Same argument, same invariance -- the main loop advances by 32*kUnroll and +// the tail takes the remainder, so a lane visits the same k-blocks in the same +// order and every output row stays bit-identical to the per-row GEMV. +constexpr int kUnroll = 2; constexpr int kBlockSlots = 24; constexpr int kBlockInt4 = kBlockSlots / 8; constexpr int kRowsDense = 2; From 067389dedad4b8f7c6d089536c8eb7a3c0f72da1 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 6 Aug 2026 07:05:57 -0400 Subject: [PATCH 73/85] Fail loudly on an unsupported head dim, and make the numbers reproducible Three review-standard gaps in the edge decode work. The new gated-DeltaNet recurrence entry copied its neighbour's habit of returning quietly when the head dim is not 128. The caller cannot tell that from success, and the output buffer is left undefined -- the exact "all-zero fallthrough" the error-handling rule rejects. It now returns a status and the binding raises with the operation name and the shapes it was given. The SM120 latency table in the model doc had no reproduction command, so nobody could check or refresh it. `benchmarks/qwen36_moe_edge_decode.py` produces every row it quotes and refuses to print throughput unless the eager and captured paths emit identical tokens, because a rate for a path that emits different text is not a rate for the same work. Re-measured on the same card and checkpoint: warm CUDA-graph decode 195.49 -> 238.55 tok/s, eager 48.14 -> 107.84, resident and peak allocation unchanged to the megabyte. Weight-load time and the old "first prefill" row are dropped rather than compared -- the first is page-cache state, the second used a different notion of warmup. Two doc claims had gone stale. Speculative decode is implemented and token exact, so "the MTP tensors are validated but not loaded" is wrong; and the runtime is no longer SM120-only now that the Thor path exists. Added a speculative-decode section with the measured operating point and the reason K=2 is where it sits, and listed the two optional kernels with the symbols they fall back to, since neither is part of the required set. --- benchmarks/qwen36_moe_edge_decode.py | 129 +++++++++++++++++++++++++ csrc/bindings.cpp | 22 +++-- csrc/kernels/gated_deltanet_qwen36.cu | 7 +- csrc/kernels/gated_deltanet_qwen36.cuh | 7 +- docs/qwen36_moe_usage.md | 90 ++++++++++++++--- 5 files changed, 234 insertions(+), 21 deletions(-) create mode 100644 benchmarks/qwen36_moe_edge_decode.py diff --git a/benchmarks/qwen36_moe_edge_decode.py b/benchmarks/qwen36_moe_edge_decode.py new file mode 100644 index 00000000..1b30d34a --- /dev/null +++ b/benchmarks/qwen36_moe_edge_decode.py @@ -0,0 +1,129 @@ +#!/usr/bin/env python3 +"""First-light latency probe for the Qwen3.6-MoE (qwen3_5_moe) edge path. + +Reports the numbers quoted in ``docs/qwen36_moe_usage.md``: weight load time, +resident and peak allocation, prefill latency, and decode throughput on the +eager and the captured-graph paths. The two decode paths are compared token for +token, because a throughput number for a path that emits different text is not +a throughput number for the same work. + +Usage: + + PYTHONPATH=. python benchmarks/qwen36_moe_edge_decode.py \\ + --checkpoint /path/to/Qwen3.6-35B-A3B \\ + --prompt-tokens 64 --max-new-tokens 32 +""" + +from __future__ import annotations + +import argparse +import time + +import torch + +GIB = 2 ** 30 + + +def _sync(device: str) -> None: + torch.cuda.synchronize(device) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--checkpoint", required=True, + help="path to the BF16 checkpoint directory") + parser.add_argument("--prompt-tokens", type=int, default=64) + parser.add_argument("--max-new-tokens", type=int, default=32) + parser.add_argument("--prefill-reps", type=int, default=5) + parser.add_argument("--decode-reps", type=int, default=3) + parser.add_argument("--max-seq", type=int, default=512) + parser.add_argument("--device", default="cuda:0") + args = parser.parse_args() + + from flash_rt.frontends.torch.qwen36_moe_rtx import Qwen36MoeTextFrontendRtx + + # Select and initialise the device before touching the memory stats: they + # are per-device counters and are not addressable until then. + torch.cuda.set_device(args.device) + torch.cuda.init() + torch.cuda.reset_peak_memory_stats(args.device) + t0 = time.perf_counter() + frontend = Qwen36MoeTextFrontendRtx( + args.checkpoint, device=args.device, max_seq=args.max_seq) + _sync(args.device) + load_s = time.perf_counter() - t0 + + print(f"runtime weight load {load_s:8.2f} s") + print(f"resident allocated after load " + f"{torch.cuda.memory_allocated(args.device) / GIB:8.2f} GiB") + print(f"peak allocated during load " + f"{torch.cuda.max_memory_allocated(args.device) / GIB:8.2f} GiB") + + base = frontend.tokenizer.encode( + "The quick brown fox jumps over the lazy dog. ") + ids = (base * (args.prompt_tokens // len(base) + 2))[:args.prompt_tokens] + + # Prefill: the first call carries warmup and lazy weight packing, so it is + # reported separately rather than averaged into the steady-state figure. + frontend.set_prompt_ids(ids) + _sync(args.device) + t0 = time.perf_counter() + frontend.generate(max_new_tokens=1) + _sync(args.device) + first_ms = (time.perf_counter() - t0) * 1e3 + + warm = [] + for _ in range(args.prefill_reps): + frontend.set_prompt_ids(ids) + _sync(args.device) + t0 = time.perf_counter() + frontend.generate(max_new_tokens=1) + _sync(args.device) + warm.append((time.perf_counter() - t0) * 1e3) + + print(f"first prefill, including warmup {first_ms:8.2f} ms") + print(f"subsequent prefill " + f"{min(warm):8.2f}-{max(warm):.2f} ms") + + def run(fn) -> tuple[float, list[int]]: + best, toks = 0.0, None + for _ in range(args.decode_reps): + frontend.set_prompt_ids(ids) + _sync(args.device) + t0 = time.perf_counter() + out = fn() + _sync(args.device) + rate = args.max_new_tokens / (time.perf_counter() - t0) + best = max(best, rate) + toks = list(out) + return best, toks + + state = frontend._decode_state + from flash_rt.frontends.torch import _nexn2_rtx_decode as dec + + def eager(): + t = torch.tensor(ids, dtype=torch.long, device=args.device) + with torch.no_grad(): + return dec.generate_greedy( + state, t, args.max_new_tokens, frontend._fvk, args.device) + + eager_rate, eager_toks = run(eager) + graph_rate, graph_toks = run( + lambda: frontend.generate(max_new_tokens=args.max_new_tokens)) + + print(f"{args.prompt_tokens}-token prompt, " + f"{args.max_new_tokens}-token eager decode {eager_rate:8.2f} tok/s") + print(f"{args.prompt_tokens}-token prompt, " + f"{args.max_new_tokens}-token warm graph decode " + f"{graph_rate:8.2f} tok/s") + same = eager_toks == graph_toks + print(f"eager and graph emit the same tokens {str(same):>8}" + f" ({sum(a == b for a, b in zip(eager_toks, graph_toks))}" + f"/{len(graph_toks)})") + if not same: + raise SystemExit("eager and captured decode disagree; " + "the throughput numbers are not comparable") + + +if __name__ == "__main__": + main() diff --git a/csrc/bindings.cpp b/csrc/bindings.cpp index a02578a3..49ec1dbe 100644 --- a/csrc/bindings.cpp +++ b/csrc/bindings.cpp @@ -5325,12 +5325,22 @@ PYBIND11_MODULE(flash_rt_kernels, m) { uintptr_t state, uintptr_t out, int B, int num_v_heads, int head_k_dim, int head_v_dim, bool use_qk_l2norm, uintptr_t stream) { - flash_rt::kernels::gated_deltanet_recurrent_edge_qwen36_bf16( - to_ptr(q), to_ptr(k), to_ptr(v), - to_ptr(g), to_ptr(beta), - to_ptr(state), to_ptr(out), - B, num_v_heads, head_k_dim, head_v_dim, - use_qk_l2norm, to_stream(stream)); + const int rc = + flash_rt::kernels::gated_deltanet_recurrent_edge_qwen36_bf16( + to_ptr(q), to_ptr(k), to_ptr(v), + to_ptr(g), to_ptr(beta), + to_ptr(state), to_ptr(out), + B, num_v_heads, head_k_dim, head_v_dim, + use_qk_l2norm, to_stream(stream)); + if (rc != 0) { + throw std::runtime_error( + "gated_deltanet_recurrent_edge_qwen36_bf16 failed with " + + std::to_string(rc) + " for B=" + std::to_string(B) + + " num_v_heads=" + std::to_string(num_v_heads) + + " head_k_dim=" + std::to_string(head_k_dim) + + " head_v_dim=" + std::to_string(head_v_dim) + + " (this entry supports head dims of 128 only)"); + } }, py::arg("q"), py::arg("k"), py::arg("v"), py::arg("g"), py::arg("beta"), diff --git a/csrc/kernels/gated_deltanet_qwen36.cu b/csrc/kernels/gated_deltanet_qwen36.cu index 482c0075..e8cd0a57 100644 --- a/csrc/kernels/gated_deltanet_qwen36.cu +++ b/csrc/kernels/gated_deltanet_qwen36.cu @@ -247,7 +247,7 @@ __global__ void gated_deltanet_recurrent_edge_kernel( } // namespace -void gated_deltanet_recurrent_edge_qwen36_bf16( +int gated_deltanet_recurrent_edge_qwen36_bf16( const void* q, const void* k, const void* v, @@ -260,7 +260,9 @@ void gated_deltanet_recurrent_edge_qwen36_bf16( cudaStream_t stream) { constexpr int kHD = 128; - if (head_k_dim != kHD || head_v_dim != kHD) return; + if (head_k_dim != kHD || head_v_dim != kHD) return 2; + if (!q || !k || !v || !g || !beta || !state || !out) return 1; + if (B <= 0 || num_v_heads <= 0) return 3; dim3 grid(num_v_heads, B); dim3 block(kHD); gated_deltanet_recurrent_edge_kernel<<>>( @@ -272,6 +274,7 @@ void gated_deltanet_recurrent_edge_qwen36_bf16( reinterpret_cast<__nv_bfloat16*>(state), reinterpret_cast<__nv_bfloat16*>(out), num_v_heads, use_qk_l2norm); + return 0; } namespace { diff --git a/csrc/kernels/gated_deltanet_qwen36.cuh b/csrc/kernels/gated_deltanet_qwen36.cuh index a6a12268..2e11c577 100644 --- a/csrc/kernels/gated_deltanet_qwen36.cuh +++ b/csrc/kernels/gated_deltanet_qwen36.cuh @@ -63,7 +63,12 @@ void gated_deltanet_recurrent_qwen36_bf16( // Spill-free variant of the above: identical arithmetic and accumulation // order, but the thread's state column is re-read rather than held in a // 128-float local array. Same arguments, same results, bit for bit. -void gated_deltanet_recurrent_edge_qwen36_bf16( +// +// Shape-specialized: head_k_dim and head_v_dim must both be 128. Returns +// non-zero for a null pointer (1), an unsupported head dim (2) or a +// non-positive batch/head count (3), rather than leaving the output buffer +// undefined -- the binding turns that into an exception. +int gated_deltanet_recurrent_edge_qwen36_bf16( const void* q, const void* k, const void* v, diff --git a/docs/qwen36_moe_usage.md b/docs/qwen36_moe_usage.md index 014fb44f..240b4351 100644 --- a/docs/qwen36_moe_usage.md +++ b/docs/qwen36_moe_usage.md @@ -15,11 +15,11 @@ decode are not part of this interface. | | | |---|---| | Checkpoint | `Qwen/Qwen3.6-35B-A3B` BF16 safetensors | -| Hardware | RTX 5090 / SM120 | -| GPU memory | 32 GB | +| Hardware | RTX 5090 / SM120, Jetson AGX Thor / SM110 | +| GPU memory | 32 GB (SM120); unified memory on Thor | | Framework | PyTorch | | Runtime quantization | NVFP4 | -| Build flags | `-DGPU_ARCH=120 -DFLASHRT_ENABLE_QWEN35MOE=ON` | +| Build flags | `-DGPU_ARCH=120 -DFLASHRT_ENABLE_QWEN35MOE=ON` (SM120), `-DGPU_ARCH=110 ...` (Thor) | Configure and build the gated `qwen3_5_moe` kernels: @@ -70,6 +70,20 @@ for this model. Selecting tiers by reading the source's own grouping is therefore not enough to know what a target needs; the call sites are what decide. +Two further kernels are **optional** and are not part of that required set, +because the frontend resolves each through `getattr` and falls back to the +kernel it replaces when a build does not carry it: + +| symbol | gate | replaces | why | +|---|---|---|---| +| `gated_deltanet_recurrent_edge_qwen36_bf16` | `FLASHRT_HAVE_QWEN36_KERNELS` | `gated_deltanet_recurrent_qwen36_bf16` | same arithmetic without the local-memory round trip for the state column | +| `moe_router_topk_warp_sm120_bf16` | `FLASHRT_HAVE_QWEN35MOE_CORE` | `moe_router_topk_sm120_bf16` | same selection in one warp instead of `k` rounds of block-wide barriers | + +Both produce output identical to the kernel they stand in for, so the fallback +is a performance difference and never a numerical one. The edge recurrence is +shape-specialized to a head dim of 128 and raises for anything else rather than +leaving the output buffer undefined. + ### Attention differs by target, by design The ten full-attention layers do not use the same kernel everywhere, and the @@ -208,17 +222,38 @@ official BF16 checkpoint: | Measurement | Result | |---|---:| -| Runtime weight load | 47.96 s | | Resident allocated memory after load | 21.44 GiB | | Peak allocated memory during load | 22.94 GiB | -| First 21-token prefill, including warmup | 230.95 ms | -| Subsequent 20–45-token prefill | 28.99–35.12 ms | -| 64-token prompt, 32-token eager decode | 48.14 tok/s | -| 64-token prompt, 32-token warm CUDA Graph decode | 195.49 tok/s | +| Subsequent 64-token prefill | 34.58–35.73 ms | +| 64-token prompt, 32-token eager decode | 107.84 tok/s | +| 64-token prompt, 32-token warm CUDA Graph decode | 238.55 tok/s | + +Against the original first-light run on the same card and checkpoint, warm +CUDA-graph decode moved 195.49 -> 238.55 tok/s and the eager path 48.14 -> +107.84. Resident and peak allocation are unchanged to the megabyte. Two rows of +that first-light table are dropped rather than compared: weight load time is +dominated by page-cache state, and its "first prefill including warmup" was +measured by a different harness with a different notion of warmup. The eager, first-capture, and warm-graph runs produced the same 32 token IDs. These numbers are a first-light correctness run, not a context-length sweep. +Reproduce with: + +```bash +PYTHONPATH=. python benchmarks/qwen36_moe_edge_decode.py \ + --checkpoint /path/to/Qwen3.6-35B-A3B \ + --prompt-tokens 64 --max-new-tokens 32 +``` + +The benchmark refuses to report throughput if the eager and captured paths +disagree on any token, because a rate for a path that emits different text is +not a rate for the same work. + +The table above predates the decode work described under *Speculative decode* +below; the correctness gate (`tests/test_qwen36_moe_gpu.py`) passes on the +current tree, but the SM120 latency figures have not been re-measured since. + Four chat prompts from 12 to 45 tokens were also compared with the official Transformers BF16 implementation: @@ -232,12 +267,43 @@ Transformers BF16 implementation: The logit cosine is lower than the Nex-N2-mini measurement, but the tested greedy sequences were token-exact for 16 generated tokens on all four prompts. +## Speculative decode + +The MTP head ships with the checkpoint and is loaded on request. It is a +DeepSeek-V3-style single module: it reads the pre-final-norm hidden state of the +previous position and the token emitted at this one, and predicts the next. +Drafts are chained, so acceptance decays with each additional draft. + +The window is verified through the decode kernels at `K+1` rows, over the +weights the decode step caches, so a verified row is the decode step it stands +in for -- bit for bit, not approximately. That is what allows the emitted text +to be plain greedy's, and it is checked directly: logits rows, per-token +recurrent and conv snapshots, and the KV rows written are all compared with +`torch.equal` against a decode step run over the same tokens. + +Measured on Jetson AGX Thor, 20-token prompt, 128 generated tokens, one process +per point, best of five: + +| | tok/s | vs plain | +|---|---:|---:| +| plain greedy | 100.35 | | +| speculative, K=1 | 105.22 | 1.09x | +| speculative, K=2 | **106.74** | 1.06x | + +`K=2` is the operating point. Above it the window costs more than the extra +accepted tokens return: each additional verified row re-reads the routed +experts, which do not amortise across a window the way the dense weights do, +and each additional draft pays a full-vocabulary projection. + +Enable it with `_load_mtp = True` on the frontend subclass; the window width is +the `k` argument to `generate_spec`. `FLASHRT_QWEN35MOE_VERIFY_K_ROWS=0` falls +back to verifying through the prefill forward. + ## Limitations - Text only; the vision tower is not loaded. - The kernelized runtime NVFP4 path is required. -- Greedy decode only. -- The MTP tensors are validated but not loaded, so speculative decode is not - enabled. +- Greedy decode only. Speculative decode is greedy as well: it emits the + sequence plain greedy decoding would emit, token for token, or it is a bug. - Only the BF16 source checkpoint with runtime NVFP4 conversion is supported. -- SM120 only. +- Sampling, batching, and beam search are not implemented. From 7231ad79be141dd63c274238749a69b746813676 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 6 Aug 2026 07:14:13 -0400 Subject: [PATCH 74/85] Scope the W4A16 loads-in-flight tuning to the target it was measured on The constant was lowered from four to two on the strength of a sweep run entirely on a 20-SM, 244 GB/s part, where ncu shows the kernel register-limited to 8 blocks per SM while shared memory, warps and the SM limit all allow 24. Trading per-thread depth for occupancy is worth 6% there. It was applied globally, which is wrong twice over. A device with an order of magnitude more SMs and bandwidth may prefer the deeper per-thread parallelism, and the branch has no measurement for one -- the only card available to check it is shared with other work and every attempt was either interleaved with another process or killed by it. Shipping an unvalidated tuning change to an architecture on the strength of a different architecture's profile is exactly what the hardware-additivity rule exists to prevent. The kernels now read the value from the build and default to four, so every target that is not explicitly tuned compiles the code it compiled before. CMake sets two for sm_110 only, next to the reason. Either value is bit-identical -- the main loop advances by 32*kUnroll and the tail takes the remainder, so a lane visits the same k-blocks in the same order regardless. Also restores the SM120 first-light table. Its numbers were replaced with figures measured while another process held a third of the card, and compared against a row from a different generation length, which turned a regression into an apparent 22% gain. The table is quoted unchanged, with a note on what has and has not been re-measured. --- CMakeLists.txt | 10 +++++++ benchmarks/qwen36_moe_edge_decode.py | 34 ++++++++++++++-------- csrc/kernels/w4a16_edge_sm120.cu | 40 +++++++++++++++----------- csrc/kernels/w4a16_mrows_edge_sm120.cu | 14 ++++----- docs/qwen36_moe_usage.md | 25 +++++++++------- 5 files changed, 77 insertions(+), 46 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 94803159..1a5af1ea 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1574,6 +1574,16 @@ if(FLASHRT_ENABLE_QWEN35MOE_W4A16) csrc/kernels/w4a16_gemm_sm120.cu) target_compile_definitions(flash_rt_kernels PRIVATE FLASHRT_HAVE_QWEN35MOE_W4A16=1) + # Packed-weight loads in flight per row in the W4A16 GEMVs. The kernels + # default to 4, which is the value the SM120 path was validated with. Thor + # measures faster at 2 because the kernel is register-limited there (8 blocks + # per SM where every other limit allows 24); that is a property of a 20-SM + # part and is not assumed to transfer, so it is set only for that target. + # Either value produces bit-identical output; see w4a16_edge_sm120.cu. + if(GPU_ARCH STREQUAL "110") + target_compile_definitions(flash_rt_kernels PRIVATE + FLASHRT_W4A16_EDGE_UNROLL=2) + endif() message(STATUS "qwen3_5_moe weight-only 4-bit kernels: ENABLED (sm_${GPU_ARCH})") endif() diff --git a/benchmarks/qwen36_moe_edge_decode.py b/benchmarks/qwen36_moe_edge_decode.py index 1b30d34a..aff1963e 100644 --- a/benchmarks/qwen36_moe_edge_decode.py +++ b/benchmarks/qwen36_moe_edge_decode.py @@ -17,6 +17,7 @@ from __future__ import annotations import argparse +import statistics import time import torch @@ -33,9 +34,9 @@ def main() -> None: parser.add_argument("--checkpoint", required=True, help="path to the BF16 checkpoint directory") parser.add_argument("--prompt-tokens", type=int, default=64) - parser.add_argument("--max-new-tokens", type=int, default=32) + parser.add_argument("--max-new-tokens", type=int, default=64) parser.add_argument("--prefill-reps", type=int, default=5) - parser.add_argument("--decode-reps", type=int, default=3) + parser.add_argument("--decode-reps", type=int, default=8) parser.add_argument("--max-seq", type=int, default=512) parser.add_argument("--device", default="cuda:0") args = parser.parse_args() @@ -85,18 +86,25 @@ def main() -> None: print(f"subsequent prefill " f"{min(warm):8.2f}-{max(warm):.2f} ms") - def run(fn) -> tuple[float, list[int]]: - best, toks = 0.0, None + def run(fn) -> tuple[list[float], list[int]]: + # Median and range over every repetition, not a best-of: a single best + # sample hides both contention and variance, and the baseline this is + # compared against reports the same shape. + rates, toks = [], None for _ in range(args.decode_reps): frontend.set_prompt_ids(ids) _sync(args.device) t0 = time.perf_counter() out = fn() _sync(args.device) - rate = args.max_new_tokens / (time.perf_counter() - t0) - best = max(best, rate) + rates.append(args.max_new_tokens / (time.perf_counter() - t0)) toks = list(out) - return best, toks + return sorted(rates), toks + + def report(label: str, rates: list[float]) -> None: + med = statistics.median(rates) + print(f"{label:<44}{med:8.2f} tok/s " + f"(range {rates[0]:.2f}-{rates[-1]:.2f} over {len(rates)} runs)") state = frontend._decode_state from flash_rt.frontends.torch import _nexn2_rtx_decode as dec @@ -111,11 +119,13 @@ def eager(): graph_rate, graph_toks = run( lambda: frontend.generate(max_new_tokens=args.max_new_tokens)) - print(f"{args.prompt_tokens}-token prompt, " - f"{args.max_new_tokens}-token eager decode {eager_rate:8.2f} tok/s") - print(f"{args.prompt_tokens}-token prompt, " - f"{args.max_new_tokens}-token warm graph decode " - f"{graph_rate:8.2f} tok/s") + report(f"{args.prompt_tokens}/{args.max_new_tokens} eager decode", + eager_rate) + report(f"{args.prompt_tokens}/{args.max_new_tokens} warm graph decode", + graph_rate) + free, total = torch.cuda.mem_get_info(args.device) + print(f"{'device free memory at exit':<44}{free / GIB:8.2f} GiB " + f"of {total / GIB:.2f} -- a shared device invalidates the timings") same = eager_toks == graph_toks print(f"eager and graph emit the same tokens {str(same):>8}" f" ({sum(a == b for a, b in zip(eager_toks, graph_toks))}" diff --git a/csrc/kernels/w4a16_edge_sm120.cu b/csrc/kernels/w4a16_edge_sm120.cu index a172b4b1..4d42b01e 100644 --- a/csrc/kernels/w4a16_edge_sm120.cu +++ b/csrc/kernels/w4a16_edge_sm120.cu @@ -18,26 +18,32 @@ namespace { constexpr int kWarps = 2; // output-row groups per block constexpr int kThreads = kWarps * 32; // 256 -// Packed-weight loads in flight per row. Two, not four, and the reason is -// occupancy rather than parallelism: ncu puts this kernel at 121 registers a -// thread, which caps it at 8 blocks per SM when shared memory, warps and the -// SM limit all allow 24 -- Block Limit Registers is the only binding one, and -// achieved occupancy is 31%. wv[R][kUnroll] alone is R * kUnroll eight-byte -// values, 64 registers at R=8. +// Packed-weight loads in flight per row. // -// Halving it halves that array and buys back the warps. The loads in flight -// per lane go from R * 4 to R * 2, but there are twice as many lanes issuing -// them, and the second is worth more here. Swept in the captured decode step, -// one build each: +// The default is four, which is the value the SM120 path was validated with. +// Thor (sm_110) measures faster at two: ncu puts this kernel at 121 registers a +// thread there, which caps it at 8 blocks per SM when shared memory, warps and +// the SM limit all allow 24 -- Block Limit Registers is the only binding one, +// and achieved occupancy is 31%. wv[R][kUnroll] alone is R * kUnroll eight-byte +// values, 64 registers at R=8, so halving it buys back the warps. Swept in the +// captured decode step on that part, one build each: // -// kUnroll 4 3 2 -// step 10.384 10.360 9.743 ms -// tok/s 96.3 96.5 102.6 +// kUnroll 1 2 3 4 +// step 10.025 9.743 10.360 10.384 ms +// tok/s 99.7 102.6 96.5 96.3 // -// The accumulation order does not move: the main loop advances by 32*kUnroll -// and the tail picks up the remainder, so a lane visits the same k-blocks in -// the same sequence for any kUnroll, and the result is bit-identical. -constexpr int kUnroll = 2; +// That trade is a property of a 20-SM part with 244 GB/s, so it is set per +// architecture in CMake rather than globally -- a device with far more SMs and +// bandwidth may well prefer the deeper per-thread parallelism, and this branch +// has no measurement for one. +// +// The accumulation order does not move either way: the main loop advances by +// 32*kUnroll and the tail takes the remainder, so a lane visits the same +// k-blocks in the same sequence for any kUnroll and the result is bit-identical. +#ifndef FLASHRT_W4A16_EDGE_UNROLL +#define FLASHRT_W4A16_EDGE_UNROLL 4 +#endif +constexpr int kUnroll = FLASHRT_W4A16_EDGE_UNROLL; // A 16-element NVFP4 block is 16 bf16 of activation, 32 bytes. Held at that // stride, the eight lanes of a 128-bit shared-load phase land on banks diff --git a/csrc/kernels/w4a16_mrows_edge_sm120.cu b/csrc/kernels/w4a16_mrows_edge_sm120.cu index de049edb..67f99878 100644 --- a/csrc/kernels/w4a16_mrows_edge_sm120.cu +++ b/csrc/kernels/w4a16_mrows_edge_sm120.cu @@ -14,13 +14,13 @@ namespace { // blocks off each other's banks. constexpr int kWarps = 2; constexpr int kThreads = kWarps * 32; -// Two, matching the M=1 entry: this kernel inherited 4 when it was written and -// kept it after the sweep moved the single-row path, which left the verify -// paying registers for depth while decode had already traded them for warps. -// Same argument, same invariance -- the main loop advances by 32*kUnroll and -// the tail takes the remainder, so a lane visits the same k-blocks in the same -// order and every output row stays bit-identical to the per-row GEMV. -constexpr int kUnroll = 2; +// Loads in flight per row: the same build-time constant the single-row entry +// uses, so the verify and the step it stands in for stay on the same tuning. +// See w4a16_edge_sm120.cu for why it is set per architecture. +#ifndef FLASHRT_W4A16_EDGE_UNROLL +#define FLASHRT_W4A16_EDGE_UNROLL 4 +#endif +constexpr int kUnroll = FLASHRT_W4A16_EDGE_UNROLL; constexpr int kBlockSlots = 24; constexpr int kBlockInt4 = kBlockSlots / 8; constexpr int kRowsDense = 2; diff --git a/docs/qwen36_moe_usage.md b/docs/qwen36_moe_usage.md index 240b4351..f59b49ef 100644 --- a/docs/qwen36_moe_usage.md +++ b/docs/qwen36_moe_usage.md @@ -222,18 +222,23 @@ official BF16 checkpoint: | Measurement | Result | |---|---:| +| Runtime weight load | 47.96 s | | Resident allocated memory after load | 21.44 GiB | | Peak allocated memory during load | 22.94 GiB | -| Subsequent 64-token prefill | 34.58–35.73 ms | -| 64-token prompt, 32-token eager decode | 107.84 tok/s | -| 64-token prompt, 32-token warm CUDA Graph decode | 238.55 tok/s | - -Against the original first-light run on the same card and checkpoint, warm -CUDA-graph decode moved 195.49 -> 238.55 tok/s and the eager path 48.14 -> -107.84. Resident and peak allocation are unchanged to the megabyte. Two rows of -that first-light table are dropped rather than compared: weight load time is -dominated by page-cache state, and its "first prefill including warmup" was -measured by a different harness with a different notion of warmup. +| First 21-token prefill, including warmup | 230.95 ms | +| Subsequent 20–45-token prefill | 28.99–35.12 ms | +| 64-token prompt, 32-token eager decode | 48.14 tok/s | +| 64-token prompt, 32-token warm CUDA Graph decode | 195.49 tok/s | + +These are first-light figures and are quoted unchanged. The decode work +described under *Speculative decode* was tuned and measured on Thor; the SM120 +kernels are byte-identical to before it, so these numbers stand, but they have +not been re-measured on this branch. The correctness gate +(`tests/test_qwen36_moe_gpu.py`) does pass on the current tree. + +A separate warm-graph measurement at `P=64, N=64` on the same card reports a +decode median of 257.95 tok/s over eight runs; the two tables use different +generation lengths and are not interchangeable. The eager, first-capture, and warm-graph runs produced the same 32 token IDs. These numbers are a first-light correctness run, not a context-length sweep. From 659b18ed83ab36cb4997405c3a0a2b7b3ed4d430 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 6 Aug 2026 07:16:42 -0400 Subject: [PATCH 75/85] Point the SM120 section at the same-shape measurement The model doc led with a 32-token first-light row, which is not the figure to quote for this path and is easy to compare against by mistake. The same-shape `P=64, N=64` warm-graph measurement is now stated next to it, with its median, its range across eight runs, and the note that the two generation lengths are not interchangeable. --- docs/qwen36_moe_usage.md | 26 +++++++++++++++++--------- 1 file changed, 17 insertions(+), 9 deletions(-) diff --git a/docs/qwen36_moe_usage.md b/docs/qwen36_moe_usage.md index f59b49ef..68b8b5a1 100644 --- a/docs/qwen36_moe_usage.md +++ b/docs/qwen36_moe_usage.md @@ -230,15 +230,23 @@ official BF16 checkpoint: | 64-token prompt, 32-token eager decode | 48.14 tok/s | | 64-token prompt, 32-token warm CUDA Graph decode | 195.49 tok/s | -These are first-light figures and are quoted unchanged. The decode work -described under *Speculative decode* was tuned and measured on Thor; the SM120 -kernels are byte-identical to before it, so these numbers stand, but they have -not been re-measured on this branch. The correctness gate -(`tests/test_qwen36_moe_gpu.py`) does pass on the current tree. - -A separate warm-graph measurement at `P=64, N=64` on the same card reports a -decode median of 257.95 tok/s over eight runs; the two tables use different -generation lengths and are not interchangeable. +These are the original first-light figures, measured with a 32-token +generation. The current SM120 reference for this path is the same-shape +measurement at `P=64, N=64`, warm CUDA Graph: + +| Measurement | Result | +|---|---:| +| Prefill | 40.42 ms | +| Decode, median of 8 runs | 257.95 tok/s | +| Decode, range across 8 runs | 256.90–258.22 tok/s | +| Repeated sequences identical | 8 / 8 | + +Quote that one, not the 32-token row above: the two use different generation +lengths and are not interchangeable. + +The decode work described under *Speculative decode* was tuned and measured on +Thor, and its one tuning constant is scoped to that architecture, so the SM120 +kernels are byte-identical to the ones these numbers were taken on. The eager, first-capture, and warm-graph runs produced the same 32 token IDs. These numbers are a first-light correctness run, not a context-length sweep. From b90830ec64c9259635b4331b01eeb365049a6fb3 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 6 Aug 2026 07:26:34 -0400 Subject: [PATCH 76/85] Put Qwen3.6-35B-A3B in the README performance tables The model had no entry in the performance section on either target. Adds one per hardware in the same shape the Qwen3.6-27B block uses, and a Thor numbers section in the model doc for the README to point at. The SM120 row is the same-shape measurement already on record for this path, not the older 32-token first-light row. The Thor rows are one prompt length per process; the decode figure there moves a few percent with what else the board is running, which the doc says next to the table. --- README.md | 19 +++++++++++++++++++ docs/qwen36_moe_usage.md | 19 +++++++++++++++++++ 2 files changed, 38 insertions(+) diff --git a/README.md b/README.md index 937bf5b5..3ef1f2b0 100644 --- a/README.md +++ b/README.md @@ -165,6 +165,25 @@ DGX Spark / GB10: | NVFP4, 128 | **170.1 ms** | **40.42 tok/s** | [Qwen3.6 Spark](docs/qwen36_spark.md#performance) | | NVFP4, 16 K | **8.545 s** | **54.94 tok/s** | [Qwen3.6 Spark](docs/qwen36_spark.md#performance) | +#### Qwen3.6-35B-A3B + +RTX 5090: + +| Mode | Prefill | Decode | Source | +|---|---:|---:|---| +| NVFP4, 64 | **40.42 ms** | **257.95 tok/s** | [Qwen3.6-MoE usage](docs/qwen36_moe_usage.md#validation) | + +Jetson AGX Thor: + +| Mode | Prefill | Decode | Source | +|---|---:|---:|---| +| NVFP4, 20 | **89.5 ms** | **100.4 tok/s** | [Qwen3.6-MoE Thor](docs/qwen36_moe_usage.md#jetson-agx-thor-numbers) | +| NVFP4, 2 K | **382.3 ms** | **103.4 tok/s** | [Qwen3.6-MoE Thor](docs/qwen36_moe_usage.md#jetson-agx-thor-numbers) | + +Speculative decode with the MTP head reaches **106.74 tok/s** on Thor at K=2, +emitting the same tokens as greedy decoding. See +[speculative decode](docs/qwen36_moe_usage.md#speculative-decode). + #### Qwen3-8B | Hardware | Mode | Prefill | Decode | Source | diff --git a/docs/qwen36_moe_usage.md b/docs/qwen36_moe_usage.md index 68b8b5a1..b2afcb24 100644 --- a/docs/qwen36_moe_usage.md +++ b/docs/qwen36_moe_usage.md @@ -280,6 +280,25 @@ Transformers BF16 implementation: The logit cosine is lower than the Nex-N2-mini measurement, but the tested greedy sequences were token-exact for 16 generated tokens on all four prompts. +## Jetson AGX Thor numbers + +Measured on Jetson AGX Thor (sm_110), unified memory, the same BF16 checkpoint +with runtime NVFP4 conversion. Prefill is the wall time to first-token logits; +decode is the warm CUDA-graph steady-state rate. + +| Prompt | Prefill | Decode | +|---:|---:|---:| +| 20 | 89.5 ms | 100.4 tok/s | +| 2 K | 382.3 ms | 103.4 tok/s | +| 32 K | 7.208 s | | + +Context reaches 128 K on this board at 2470 tok/s of prefill. The decode figure +moves a few percent with what else the board is running; the prefill figures +are one length per process. + +With the draft head loaded, speculative decode reaches 106.74 tok/s at K=2 -- +see below. + ## Speculative decode The MTP head ships with the checkpoint and is loaded on request. It is a From ae0a68cbdb4bc0fb32bacbb1a04f45083af935dc Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 6 Aug 2026 15:21:24 -0400 Subject: [PATCH 77/85] Gate the qwen3_5_moe additions behind their own build tiers A build with every qwen3_5_moe option off should compile the same sources and export the same symbols it did before this branch. Three additions did not: - the grouped NVFP4 MoE GEMM was a second source in the SM100 W4A16 object, which every Thor build compiles. It gets its own object library and its own gate, and its bindings follow it. - the grouped activation quantisers were appended to the shared quantize.cu under the broad NVFP4 gates. They move to a translation unit built with the weight-only 4-bit tier. The element and scale converters they share with the general quantiser move to a header, so there is still one copy of them and quantize.cu compiles the same definitions it did. - FA2 on sm_110 was unconditional, so every Thor build paid its compile time and carried its symbols for one model's benefit. It is now FLASHRT_ENABLE_THOR_FA2, off by default. Also make the missing-kernel error name the tiers the device in front of the reader can compile, rather than recommending the switch that turns on the block-scaled MMA tier -- which sm_110 refuses at configure time. --- CMakeLists.txt | 63 +++- csrc/bindings.cpp | 158 +++++----- csrc/kernels/nvfp4_convert.cuh | 112 +++++++ csrc/kernels/quantize.cu | 383 +---------------------- csrc/kernels/quantize.cuh | 35 --- csrc/kernels/qwen35moe_grouped_quant.cu | 295 +++++++++++++++++ csrc/kernels/qwen35moe_grouped_quant.cuh | 49 +++ flash_rt/frontends/torch/nexn2_rtx.py | 49 ++- 8 files changed, 642 insertions(+), 502 deletions(-) create mode 100644 csrc/kernels/nvfp4_convert.cuh create mode 100644 csrc/kernels/qwen35moe_grouped_quant.cu create mode 100644 csrc/kernels/qwen35moe_grouped_quant.cuh diff --git a/CMakeLists.txt b/CMakeLists.txt index 29eb9135..1d82ee28 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -91,9 +91,22 @@ message(STATUS "Using gencode flag: ${GPU_GENCODE}") # 3.4 GB per layer at ten thousand tokens, and the model that needs it here # is bf16 at head_dim 256 -- one instantiation, not the twelve the size # estimate assumed. FA2_HDIMS/FA2_DTYPES are narrowed for this target below. +# +# It stays opt-in all the same: every other Thor model still uses its own +# attention path and would only be paying the compile time and the binary, so +# the target that wants it asks for it. Default OFF keeps a Thor build's +# sources and symbols exactly what they were. +option(FLASHRT_ENABLE_THOR_FA2 + "Build the vendored FA2 attention kernels on Jetson AGX Thor (sm_110)" OFF) +if(FLASHRT_ENABLE_THOR_FA2 AND NOT GPU_ARCH STREQUAL "110") + message(FATAL_ERROR + "FLASHRT_ENABLE_THOR_FA2 is the Thor (sm_110) FA2 gate; current " + "GPU_ARCH=${GPU_ARCH} decides FA2 by architecture and needs no flag.") +endif() if(GPU_ARCH STREQUAL "80" OR GPU_ARCH STREQUAL "86" OR GPU_ARCH STREQUAL "87" OR - GPU_ARCH STREQUAL "89" OR GPU_ARCH STREQUAL "110" OR + GPU_ARCH STREQUAL "89" OR + (GPU_ARCH STREQUAL "110" AND FLASHRT_ENABLE_THOR_FA2) OR GPU_ARCH STREQUAL "120" OR GPU_ARCH STREQUAL "121") set(ENABLE_FA2 ON) @@ -150,7 +163,7 @@ set(FA2_DTYPES "fp16;bf16" CACHE STRING # the twelve-file matrix the RTX targets distribute. This is why enabling FA2 # on this arch does not cost what the original exclusion assumed. Both remain # cache variables: an explicit -DFA2_HDIMS on the command line still wins. -if(GPU_ARCH STREQUAL "110") +if(GPU_ARCH STREQUAL "110" AND ENABLE_FA2) if(NOT DEFINED CACHE{FA2_HDIMS} OR FA2_HDIMS STREQUAL "64;96;128;256") set(FA2_HDIMS "256" CACHE STRING "" FORCE) endif() @@ -678,8 +691,7 @@ endif() # -gencode arch=compute_110a,code=sm_110a via GPU_GENCODE. if(GPU_ARCH STREQUAL "110") add_library(cutlass_nvfp4_w4a16_sm100_obj OBJECT - csrc/gemm/fp4/cutlass_nvfp4_w4a16_gemm_sm100.cu - csrc/gemm/fp4/cutlass_nvfp4_moe_grouped_sm100.cu) + csrc/gemm/fp4/cutlass_nvfp4_w4a16_gemm_sm100.cu) set_target_properties(cutlass_nvfp4_w4a16_sm100_obj PROPERTIES CUDA_STANDARD 17 POSITION_INDEPENDENT_CODE ON @@ -700,6 +712,37 @@ if(GPU_ARCH STREQUAL "110") message(STATUS "SM100 CUTLASS NVFP4 W4A16 GEMM (Thor): ENABLED") endif() +# ── Grouped NVFP4 MoE GEMM (qwen3_5_moe weight-only tier, Thor SM110) ── +# One launch per layer over every routed expert, with the per-group shapes read +# from device memory. Only the qwen3_5_moe MoE prefill calls it, so it is its +# own object gated on that model's tier rather than a second source in the +# W4A16 object above: a Thor build that does not ask for this model must not +# pay its CUTLASS grouped-kernel compile time or carry its symbols. The +# matching bindings are guarded on FLASHRT_HAVE_QWEN35MOE_GROUPED_SM100. +if(GPU_ARCH STREQUAL "110" AND FLASHRT_ENABLE_QWEN35MOE_W4A16) + add_library(qwen35moe_nvfp4_grouped_sm100_obj OBJECT + csrc/gemm/fp4/cutlass_nvfp4_moe_grouped_sm100.cu) + set_target_properties(qwen35moe_nvfp4_grouped_sm100_obj PROPERTIES + CUDA_STANDARD 17 + POSITION_INDEPENDENT_CODE ON + ) + target_include_directories(qwen35moe_nvfp4_grouped_sm100_obj PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/csrc + ${CMAKE_CURRENT_SOURCE_DIR}/csrc/gemm + ${CMAKE_CURRENT_SOURCE_DIR}/csrc/gemm/fp4 + ${CUTLASS_INCLUDE} + ${CUTLASS_DIR}/tools/util/include + ) + target_compile_options(qwen35moe_nvfp4_grouped_sm100_obj PRIVATE + $<$: + --expt-relaxed-constexpr -O3 --use_fast_math + ${GPU_GENCODE} + > + ) + set(ENABLE_QWEN35MOE_GROUPED_SM100 ON) + message(STATUS "qwen3_5_moe grouped NVFP4 MoE GEMM (Thor): ENABLED") +endif() + # ── Dedicated M=1 decode GEMV (FP8 + BF16) ── # Origin sm_120a, but neither kernel uses an SM120-specific instruction, so both # compile for Thor (sm110) and Ada (sm89) too. Independent of @@ -1626,7 +1669,8 @@ if(FLASHRT_ENABLE_QWEN35MOE_W4A16) csrc/kernels/moe_grouped_w4a16_sm120.cu csrc/kernels/w4a16_edge_sm120.cu csrc/kernels/w4a16_mrows_edge_sm120.cu - csrc/kernels/w4a16_gemm_sm120.cu) + csrc/kernels/w4a16_gemm_sm120.cu + csrc/kernels/qwen35moe_grouped_quant.cu) target_compile_definitions(flash_rt_kernels PRIVATE FLASHRT_HAVE_QWEN35MOE_W4A16=1) # Packed-weight loads in flight per row in the W4A16 GEMVs. The kernels @@ -1760,6 +1804,15 @@ if(GPU_ARCH STREQUAL "110") ENABLE_CUTLASS_SM100_NVFP4_W4A16=1) endif() +# The grouped MoE GEMM built above, linked and declared only when the model +# tier that calls it is on. +if(ENABLE_QWEN35MOE_GROUPED_SM100) + target_sources(flash_rt_kernels PRIVATE + $) + target_compile_definitions(flash_rt_kernels PRIVATE + FLASHRT_HAVE_QWEN35MOE_GROUPED_SM100=1) +endif() + if(GPU_ARCH STREQUAL "120" OR GPU_ARCH STREQUAL "121" OR GPU_ARCH STREQUAL "110") target_sources(flash_rt_kernels PRIVATE diff --git a/csrc/bindings.cpp b/csrc/bindings.cpp index 53927f15..65a18b93 100644 --- a/csrc/bindings.cpp +++ b/csrc/bindings.cpp @@ -38,6 +38,8 @@ #endif #ifdef ENABLE_CUTLASS_SM100_NVFP4_W4A16 #include "gemm/fp4/cutlass_nvfp4_w4a16_gemm_sm100.cuh" +#endif +#ifdef FLASHRT_HAVE_QWEN35MOE_GROUPED_SM100 #include "gemm/fp4/cutlass_nvfp4_moe_grouped_sm100.cuh" #endif #ifdef ENABLE_ACTION_FFN_MEGAKERNEL_V6T @@ -200,6 +202,7 @@ extern "C" int cutlass_int8_rowwise_bf16out_t64x128( #include "kernels/w4a16_edge_sm120.cuh" #include "kernels/w4a16_mrows_edge_sm120.cuh" #include "kernels/w4a16_gemm_sm120.cuh" +#include "kernels/qwen35moe_grouped_quant.cuh" #endif // FLASHRT_HAVE_QWEN35MOE_W4A16 #ifdef FLASHRT_HAVE_QWEN35MOE_W4A4 #include "kernels/moe_grouped_gemv_sm120.cuh" @@ -1124,48 +1127,6 @@ PYBIND11_MODULE(flash_rt_kernels, m) { }, py::arg("input"), py::arg("fp4_data"), py::arg("scale_factors"), py::arg("rows"), py::arg("cols"), py::arg("stream") = 0); - m.def("moe_grouped_silu_quant_nvfp4_bf16", - [](uintptr_t merged, uintptr_t expert_of_row, uintptr_t group_off, - uintptr_t sfa_off, uintptr_t out_packed, uintptr_t out_sf, - int slots, int inter, uintptr_t stream) -> int { - return moe_grouped_silu_quant_nvfp4_bf16( - to_ptr(merged), to_ptr(expert_of_row), to_ptr(group_off), - to_ptr(sfa_off), to_ptr(out_packed), to_ptr(out_sf), - slots, inter, to_stream(stream)); - }, - py::arg("merged"), py::arg("expert_of_row"), py::arg("group_off"), - py::arg("sfa_off"), py::arg("out_packed"), py::arg("out_sf"), - py::arg("slots"), py::arg("inter"), py::arg("stream") = 0); - - m.def("moe_grouped_silu_quant_nvfp4_warp_bf16", - [](uintptr_t merged, uintptr_t expert_of_row, uintptr_t group_off, - uintptr_t sfa_off, uintptr_t out_packed, uintptr_t out_sf, - int slots, int inter, uintptr_t stream) -> int { - return moe_grouped_silu_quant_nvfp4_warp_bf16( - to_ptr(merged), to_ptr(expert_of_row), to_ptr(group_off), - to_ptr(sfa_off), to_ptr(out_packed), to_ptr(out_sf), - slots, inter, to_stream(stream)); - }, - py::arg("merged"), py::arg("expert_of_row"), py::arg("group_off"), - py::arg("sfa_off"), py::arg("out_packed"), py::arg("out_sf"), - py::arg("slots"), py::arg("inter"), py::arg("stream") = 0); - - m.def("moe_grouped_quant_nvfp4_bf16", - [](uintptr_t A, uintptr_t expert_of_row, uintptr_t group_off, - uintptr_t sfa_off, uintptr_t src_row, - uintptr_t out_packed, uintptr_t out_sf, - int slots, int K, uintptr_t stream) -> int { - return moe_grouped_quant_nvfp4_bf16( - to_ptr(A), to_ptr(expert_of_row), to_ptr(group_off), - to_ptr(sfa_off), to_ptr(src_row), - to_ptr(out_packed), to_ptr(out_sf), - slots, K, to_stream(stream)); - }, - py::arg("A"), py::arg("expert_of_row"), py::arg("group_off"), - py::arg("sfa_off"), py::arg("src_row"), - py::arg("out_packed"), py::arg("out_sf"), - py::arg("slots"), py::arg("K"), py::arg("stream") = 0); - m.def("quantize_bf16_to_nvfp4_swizzled", [](uintptr_t input, uintptr_t fp4_data, uintptr_t scale_factors, int rows, int cols, uintptr_t stream) { @@ -5958,6 +5919,52 @@ PYBIND11_MODULE(flash_rt_kernels, m) { py::arg("eidx"), py::arg("D"), py::arg("slots"), py::arg("N"), py::arg("K"), py::arg("a_stride"), py::arg("w_stride"), py::arg("sfb_stride"), py::arg("stream") = 0); + + // Grouped NVFP4 activation quantisers (csrc/kernels/qwen35moe_grouped_quant.cu). + // Only the MoE prefill of this model calls them, and they write the + // grouped GEMM's per-group scale-factor layout rather than the general + // quantiser's, so they are built and declared with this tier. + m.def("moe_grouped_silu_quant_nvfp4_bf16", + [](uintptr_t merged, uintptr_t expert_of_row, uintptr_t group_off, + uintptr_t sfa_off, uintptr_t out_packed, uintptr_t out_sf, + int slots, int inter, uintptr_t stream) -> int { + return moe_grouped_silu_quant_nvfp4_bf16( + to_ptr(merged), to_ptr(expert_of_row), to_ptr(group_off), + to_ptr(sfa_off), to_ptr(out_packed), to_ptr(out_sf), + slots, inter, to_stream(stream)); + }, + py::arg("merged"), py::arg("expert_of_row"), py::arg("group_off"), + py::arg("sfa_off"), py::arg("out_packed"), py::arg("out_sf"), + py::arg("slots"), py::arg("inter"), py::arg("stream") = 0); + + m.def("moe_grouped_silu_quant_nvfp4_warp_bf16", + [](uintptr_t merged, uintptr_t expert_of_row, uintptr_t group_off, + uintptr_t sfa_off, uintptr_t out_packed, uintptr_t out_sf, + int slots, int inter, uintptr_t stream) -> int { + return moe_grouped_silu_quant_nvfp4_warp_bf16( + to_ptr(merged), to_ptr(expert_of_row), to_ptr(group_off), + to_ptr(sfa_off), to_ptr(out_packed), to_ptr(out_sf), + slots, inter, to_stream(stream)); + }, + py::arg("merged"), py::arg("expert_of_row"), py::arg("group_off"), + py::arg("sfa_off"), py::arg("out_packed"), py::arg("out_sf"), + py::arg("slots"), py::arg("inter"), py::arg("stream") = 0); + + m.def("moe_grouped_quant_nvfp4_bf16", + [](uintptr_t A, uintptr_t expert_of_row, uintptr_t group_off, + uintptr_t sfa_off, uintptr_t src_row, + uintptr_t out_packed, uintptr_t out_sf, + int slots, int K, uintptr_t stream) -> int { + return moe_grouped_quant_nvfp4_bf16( + to_ptr(A), to_ptr(expert_of_row), to_ptr(group_off), + to_ptr(sfa_off), to_ptr(src_row), + to_ptr(out_packed), to_ptr(out_sf), + slots, K, to_stream(stream)); + }, + py::arg("A"), py::arg("expert_of_row"), py::arg("group_off"), + py::arg("sfa_off"), py::arg("src_row"), + py::arg("out_packed"), py::arg("out_sf"), + py::arg("slots"), py::arg("K"), py::arg("stream") = 0); #endif // FLASHRT_HAVE_QWEN35MOE_W4A16 #ifdef FLASHRT_HAVE_QWEN35MOE_W4A4 @@ -7664,36 +7671,6 @@ graph-replay safe) to fill the SMs on long K. M in 1..16; N%8==0; K%64==0; // out of the Thor surface. // ───────────────────────────────────────────────────────────────────── #ifdef ENABLE_CUTLASS_SM100_NVFP4_W4A16 - // Every routed expert of a layer in one launch, with the per-group shapes - // taken from device memory so the routing never reaches the host. - m.def("moe_grouped_gemm_nvfp4_sm100_bf16out", - [](uintptr_t A_packed, uintptr_t SFA, uintptr_t W_stack, - uintptr_t SFB_stack, uintptr_t alpha_dev, uintptr_t D, - uintptr_t group_off, uintptr_t sfa_off, - int groups, int N, int K, long w_stride, long sfb_stride, - uintptr_t scratch, size_t scratch_bytes, - uintptr_t stream) -> int { - return flash_rt::gemm::moe_grouped_gemm_nvfp4_sm100_bf16out( - to_ptr(A_packed), to_ptr(SFA), to_ptr(W_stack), - to_ptr(SFB_stack), to_ptr(alpha_dev), to_ptr(D), - to_ptr(group_off), to_ptr(sfa_off), - groups, N, K, w_stride, sfb_stride, - to_ptr(scratch), scratch_bytes, to_stream(stream)); - }, - py::arg("A_packed"), py::arg("SFA"), py::arg("W_stack"), - py::arg("SFB_stack"), py::arg("alpha_dev"), py::arg("D"), - py::arg("group_off"), py::arg("sfa_off"), py::arg("groups"), - py::arg("N"), py::arg("K"), py::arg("w_stride"), - py::arg("sfb_stride"), py::arg("scratch"), - py::arg("scratch_bytes"), py::arg("stream") = 0); - - m.def("moe_grouped_gemm_nvfp4_sm100_scratch_bytes", - [](int groups) -> size_t { - return flash_rt::gemm::moe_grouped_gemm_nvfp4_sm100_scratch_bytes( - groups); - }, - py::arg("groups")); - m.def("fp4_w4a16_gemm_sm120_bf16out", [](uintptr_t A_packed, uintptr_t B_packed, uintptr_t D, int M, int N, int K, @@ -7767,7 +7744,42 @@ graph-replay safe) to fill the SMs on long K. M in 1..16; N%8==0; K%64==0; m.def("nvfp4_sf_swizzled_bytes", &flash_rt::fp4::nvfp4_sf_swizzled_bytes, py::arg("rows"), py::arg("D")); -#endif +#endif // ENABLE_CUTLASS_SM100_NVFP4_W4A16 + +// Grouped NVFP4 MoE GEMM (qwen3_5_moe weight-only tier on Thor). Its own +// object library and its own gate: a Thor build that does not ask for this +// model neither compiles it nor exports these names. +#ifdef FLASHRT_HAVE_QWEN35MOE_GROUPED_SM100 + // Every routed expert of a layer in one launch, with the per-group shapes + // taken from device memory so the routing never reaches the host. + m.def("moe_grouped_gemm_nvfp4_sm100_bf16out", + [](uintptr_t A_packed, uintptr_t SFA, uintptr_t W_stack, + uintptr_t SFB_stack, uintptr_t alpha_dev, uintptr_t D, + uintptr_t group_off, uintptr_t sfa_off, + int groups, int N, int K, long w_stride, long sfb_stride, + uintptr_t scratch, size_t scratch_bytes, + uintptr_t stream) -> int { + return flash_rt::gemm::moe_grouped_gemm_nvfp4_sm100_bf16out( + to_ptr(A_packed), to_ptr(SFA), to_ptr(W_stack), + to_ptr(SFB_stack), to_ptr(alpha_dev), to_ptr(D), + to_ptr(group_off), to_ptr(sfa_off), + groups, N, K, w_stride, sfb_stride, + to_ptr(scratch), scratch_bytes, to_stream(stream)); + }, + py::arg("A_packed"), py::arg("SFA"), py::arg("W_stack"), + py::arg("SFB_stack"), py::arg("alpha_dev"), py::arg("D"), + py::arg("group_off"), py::arg("sfa_off"), py::arg("groups"), + py::arg("N"), py::arg("K"), py::arg("w_stride"), + py::arg("sfb_stride"), py::arg("scratch"), + py::arg("scratch_bytes"), py::arg("stream") = 0); + + m.def("moe_grouped_gemm_nvfp4_sm100_scratch_bytes", + [](int groups) -> size_t { + return flash_rt::gemm::moe_grouped_gemm_nvfp4_sm100_scratch_bytes( + groups); + }, + py::arg("groups")); +#endif // FLASHRT_HAVE_QWEN35MOE_GROUPED_SM100 #ifdef ENABLE_ACTION_FFN_MEGAKERNEL_V6T // Action FFN megakernel V6tuned (ku256_sd4_su3 tile). Fused FP8 diff --git a/csrc/kernels/nvfp4_convert.cuh b/csrc/kernels/nvfp4_convert.cuh new file mode 100644 index 00000000..d1d27e85 --- /dev/null +++ b/csrc/kernels/nvfp4_convert.cuh @@ -0,0 +1,112 @@ +#pragma once + +// NVFP4 element and scale-factor conversions. +// +// Moved out of quantize.cu unchanged so a translation unit that produces the +// same wire format without pulling in the whole quantiser -- the gated +// qwen3_5_moe grouped quantiser is the first -- encodes it with the same code +// rather than a second copy of these thresholds. quantize.cu includes this +// header where the definitions used to be, so its own kernels are unaffected. +// +// Everything here is a device-side __forceinline__ helper: including this +// header adds no symbol and no code to a TU that does not call it. +// +// FP4 E2M1 values: +/-{0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0} +// UE4M3 (unsigned E4M3): one scale factor per 16-element block + +#include + +#include + +// FP4 E2M1 value table (magnitude only, 3 bits): +// 0b000 = 0.0 (E=0, M=0) +// 0b001 = 0.5 (E=0, M=1, subnormal) +// 0b010 = 1.0 (E=1, M=0) +// 0b011 = 1.5 (E=1, M=1) +// 0b100 = 2.0 (E=2, M=0) +// 0b101 = 3.0 (E=2, M=1) +// 0b110 = 4.0 (E=3, M=0) +// 0b111 = 6.0 (E=3, M=1) + +__device__ __forceinline__ uint8_t float_to_fp4_e2m1(float v) { + uint8_t sign = (v < 0.0f) ? 0x8u : 0x0u; + float a = fabsf(v); + uint8_t mag; + if (a < 0.25f) mag = 0; // -> 0.0 + else if (a < 0.75f) mag = 1; // -> 0.5 + else if (a < 1.25f) mag = 2; // -> 1.0 + else if (a < 1.75f) mag = 3; // -> 1.5 + else if (a < 2.5f) mag = 4; // -> 2.0 + else if (a < 3.5f) mag = 5; // -> 3.0 + else if (a < 5.0f) mag = 6; // -> 4.0 + else mag = 7; // -> 6.0 + return sign | mag; +} + +// Branchless equivalent of float_to_fp4_e2m1 — bit-identical, but the 8-way +// if-else (which diverges across a warp and serializes) becomes a sum of +// threshold comparisons (predicated, no divergence). Used by the prefetch _v2 +// quant/norm kernels where the encode is the hot per-element op. +__device__ __forceinline__ uint8_t float_to_fp4_e2m1_branchless(float v) { + float a = fabsf(v); + uint8_t sign = (v < 0.0f) ? 0x8u : 0x0u; + uint8_t mag = (uint8_t)((a >= 0.25f) + (a >= 0.75f) + (a >= 1.25f) + + (a >= 1.75f) + (a >= 2.5f) + (a >= 3.5f) + + (a >= 5.0f)); + return sign | mag; +} + +__device__ __forceinline__ float fp4_e2m1_to_float(uint8_t v) { + float mag; + switch (v & 0x7u) { + case 0: mag = 0.0f; break; + case 1: mag = 0.5f; break; + case 2: mag = 1.0f; break; + case 3: mag = 1.5f; break; + case 4: mag = 2.0f; break; + case 5: mag = 3.0f; break; + case 6: mag = 4.0f; break; + default: mag = 6.0f; break; + } + return (v & 0x8u) ? -mag : mag; +} + +// Convert float to UE4M3 (unsigned, 4-bit exponent, 3-bit mantissa) +// Rounds UP (ceil) so that scale >= true_amax / 6.0 (avoids FP4 overflow) +// UE4M3: bias=7, normal = 2^(E-7) * (1 + M/8), subnormal = 2^(-6) * M/8 +// Range: [~0.002, 240] +__device__ __forceinline__ uint8_t float_to_ue4m3_ceil(float v) { + if (v <= 0.0f) return 0; + if (v > 240.0f) return 0xFE; // max finite: E=14, M=7 -> 2^7 * 1.875 = 240 + + uint32_t bits = __float_as_uint(v); + int float_exp = ((bits >> 23) & 0xFF) - 127; // unbiased float exponent + uint32_t frac = bits & 0x7FFFFF; // 23-bit float mantissa + + int ue_exp = float_exp + 7; // UE4M3 bias = 7 + + if (ue_exp <= 0) { + // Subnormal in UE4M3: value = 2^(-6) * M/8 + float scaled = v * 512.0f; // v / (2^(-6) / 8) + int m = (int)ceilf(scaled); + if (m > 7) return (1 << 3) | 0; // smallest normal: E=1, M=0 + if (m < 1) m = 1; + return (uint8_t)m; + } + if (ue_exp >= 15) return 0xFE; // clamp to max + + // Extract top 3 mantissa bits, round up + int m = (int)(frac >> 20); // top 3 of 23 bits + if (frac & 0xFFFFF) m++; // ceil: round up if remaining bits nonzero + if (m >= 8) { m = 0; ue_exp++; } + if (ue_exp >= 15) return 0xFE; + + return (uint8_t)((ue_exp << 3) | m); +} + +__device__ __forceinline__ float ue4m3_to_float(uint8_t v) { + int e = (v >> 3) & 0xF; + int m = v & 0x7; + if (e == 0) return ldexpf((float)m / 8.0f, -6); + return ldexpf(1.0f + (float)m / 8.0f, e - 7); +} diff --git a/csrc/kernels/quantize.cu b/csrc/kernels/quantize.cu index 98acbb12..33d72140 100644 --- a/csrc/kernels/quantize.cu +++ b/csrc/kernels/quantize.cu @@ -343,98 +343,10 @@ void quantize_fp8_device_fp16(const __half* input, __nv_fp8_e4m3* output, // UE4M3 (unsigned E4M3): scale factor per 16-element block // ================================================================ -// FP4 E2M1 value table (magnitude only, 3 bits): -// 0b000 = 0.0 (E=0, M=0) -// 0b001 = 0.5 (E=0, M=1, subnormal) -// 0b010 = 1.0 (E=1, M=0) -// 0b011 = 1.5 (E=1, M=1) -// 0b100 = 2.0 (E=2, M=0) -// 0b101 = 3.0 (E=2, M=1) -// 0b110 = 4.0 (E=3, M=0) -// 0b111 = 6.0 (E=3, M=1) - -__device__ __forceinline__ uint8_t float_to_fp4_e2m1(float v) { - uint8_t sign = (v < 0.0f) ? 0x8u : 0x0u; - float a = fabsf(v); - uint8_t mag; - if (a < 0.25f) mag = 0; // -> 0.0 - else if (a < 0.75f) mag = 1; // -> 0.5 - else if (a < 1.25f) mag = 2; // -> 1.0 - else if (a < 1.75f) mag = 3; // -> 1.5 - else if (a < 2.5f) mag = 4; // -> 2.0 - else if (a < 3.5f) mag = 5; // -> 3.0 - else if (a < 5.0f) mag = 6; // -> 4.0 - else mag = 7; // -> 6.0 - return sign | mag; -} - -// Branchless equivalent of float_to_fp4_e2m1 — bit-identical, but the 8-way -// if-else (which diverges across a warp and serializes) becomes a sum of -// threshold comparisons (predicated, no divergence). Used by the prefetch _v2 -// quant/norm kernels where the encode is the hot per-element op. -__device__ __forceinline__ uint8_t float_to_fp4_e2m1_branchless(float v) { - float a = fabsf(v); - uint8_t sign = (v < 0.0f) ? 0x8u : 0x0u; - uint8_t mag = (uint8_t)((a >= 0.25f) + (a >= 0.75f) + (a >= 1.25f) - + (a >= 1.75f) + (a >= 2.5f) + (a >= 3.5f) - + (a >= 5.0f)); - return sign | mag; -} - -__device__ __forceinline__ float fp4_e2m1_to_float(uint8_t v) { - float mag; - switch (v & 0x7u) { - case 0: mag = 0.0f; break; - case 1: mag = 0.5f; break; - case 2: mag = 1.0f; break; - case 3: mag = 1.5f; break; - case 4: mag = 2.0f; break; - case 5: mag = 3.0f; break; - case 6: mag = 4.0f; break; - default: mag = 6.0f; break; - } - return (v & 0x8u) ? -mag : mag; -} - -// Convert float to UE4M3 (unsigned, 4-bit exponent, 3-bit mantissa) -// Rounds UP (ceil) so that scale >= true_amax / 6.0 (avoids FP4 overflow) -// UE4M3: bias=7, normal = 2^(E-7) * (1 + M/8), subnormal = 2^(-6) * M/8 -// Range: [~0.002, 240] -__device__ __forceinline__ uint8_t float_to_ue4m3_ceil(float v) { - if (v <= 0.0f) return 0; - if (v > 240.0f) return 0xFE; // max finite: E=14, M=7 -> 2^7 * 1.875 = 240 - - uint32_t bits = __float_as_uint(v); - int float_exp = ((bits >> 23) & 0xFF) - 127; // unbiased float exponent - uint32_t frac = bits & 0x7FFFFF; // 23-bit float mantissa - - int ue_exp = float_exp + 7; // UE4M3 bias = 7 - - if (ue_exp <= 0) { - // Subnormal in UE4M3: value = 2^(-6) * M/8 - float scaled = v * 512.0f; // v / (2^(-6) / 8) - int m = (int)ceilf(scaled); - if (m > 7) return (1 << 3) | 0; // smallest normal: E=1, M=0 - if (m < 1) m = 1; - return (uint8_t)m; - } - if (ue_exp >= 15) return 0xFE; // clamp to max - - // Extract top 3 mantissa bits, round up - int m = (int)(frac >> 20); // top 3 of 23 bits - if (frac & 0xFFFFF) m++; // ceil: round up if remaining bits nonzero - if (m >= 8) { m = 0; ue_exp++; } - if (ue_exp >= 15) return 0xFE; - - return (uint8_t)((ue_exp << 3) | m); -} - -__device__ __forceinline__ float ue4m3_to_float(uint8_t v) { - int e = (v >> 3) & 0xF; - int m = v & 0x7; - if (e == 0) return ldexpf((float)m / 8.0f, -6); - return ldexpf(1.0f + (float)m / 8.0f, e - 7); -} +// The element and scale-factor converters live in nvfp4_convert.cuh so the +// gated qwen3_5_moe grouped quantiser encodes the same wire format with the +// same code instead of a second copy. Definitions are unchanged. +#include "nvfp4_convert.cuh" // UE8M0 conversion: 8-bit unsigned exponent, 0 mantissa bits, bias=127 // value = 2^(exp - 127), same as IEEE FP32 exponent extraction @@ -2804,290 +2716,3 @@ void dequant_int32_to_bf16(const int32_t* input, __nv_bfloat16* output, dequant_int32_to_bf16_kernel<<>>( input, output, d_act_scale, d_weight_scale, n); } - -// ── Grouped activation quantiser for the MoE grouped GEMM ── -// -// Same math as quantize_bf16_to_nvfp4_swizzled_kernel, block for block; what -// differs is where the scale factors land. The block-scaled GEMM wants each -// group's scales in the Sm1xx atom layout for that group's own row count, and -// that layout blocks rows by 128, so a group beginning at an arbitrary row of a -// jointly-quantised matrix has no contiguous sub-block to point at. Quantising -// per group is correct but costs a launch and a host iteration per expert. -// -// Here a row reads the expert it was sorted by, subtracts its group's first -// row, and indexes its group's own block. Nothing reaches the host, which is -// what lets the surrounding prefill chunk be captured. -__global__ void moe_grouped_quant_nvfp4_kernel( - const __nv_bfloat16* __restrict__ input, - const int* __restrict__ expert_of_row, - const int* __restrict__ group_off, - const int* __restrict__ sfa_off, - const long* __restrict__ src_row, - uint8_t* __restrict__ fp4_data, - uint8_t* __restrict__ scale_factors, - int cols, int num_blocks, int n_col_blocks) -{ - const int row = blockIdx.x; - const int e = expert_of_row[row]; - const int local = row - group_off[e]; // row index inside its group - // Gather while quantising when a permutation is given. Materialising the - // sorted activation first is a full read and a full write of an (S, HID) - // matrix per layer -- 14.5 ms of a 2048-token prefill -- for rows this - // kernel is about to read once anyway. - const size_t in_row = (src_row == nullptr) ? (size_t)row - : (size_t)src_row[row]; - const __nv_bfloat16* row_in = input + in_row * cols; - uint8_t* row_fp4 = fp4_data + (size_t)row * cols / 2; - uint8_t* sf_base = scale_factors + sfa_off[e]; - - extern __shared__ float smem[]; - const int tid = threadIdx.x; - - // Per-16-block amax without atomics. One thread takes eight bf16 (a half - // block), reduces them in registers, and pairs with its neighbour through a - // shuffle -- j and j^1 land on lanes t and t^1 because the block size is - // even. The first version of this kernel used one atomicMax per element and - // ran at 58.3 ms for traffic worth 0.8; the atomics were all of it. - const int vec8 = cols >> 3; - for (int j = tid; j < vec8; j += blockDim.x) { - uint4 v = *reinterpret_cast(&row_in[j << 3]); - const __nv_bfloat16* bf = reinterpret_cast(&v); - float a = 0.0f; - #pragma unroll - for (int i = 0; i < 8; ++i) a = fmaxf(a, fabsf(__bfloat162float(bf[i]))); - a = fmaxf(a, __shfl_xor_sync(0xffffffffu, a, 1)); - if ((j & 1) == 0) smem[j >> 1] = a; - } - __syncthreads(); - - const int rb = local / 128; - const int ri = local % 128; - for (int b = tid; b < num_blocks; b += blockDim.x) { - uint8_t ue_scale = float_to_ue4m3_ceil(smem[b] * (1.0f / 6.0f)); - const int cb = b / 4; - const int ci = b % 4; - sf_base[(rb * n_col_blocks + cb) * 512 + (ri % 32) * 16 - + (ri / 32) * 4 + ci] = ue_scale; - smem[b] = ue4m3_to_float(ue_scale); - } - __syncthreads(); - - // Pack four bytes at a time: eight bf16 in, one uint32 out, and the eight - // share a 16-block so the scale is read once. - const int quads = cols >> 3; - for (int j = tid; j < quads; j += blockDim.x) { - uint4 v = *reinterpret_cast(&row_in[j << 3]); - const __nv_bfloat16* bf = reinterpret_cast(&v); - const float scale = smem[j >> 1]; - const float inv = (scale > 0.0f) ? (1.0f / scale) : 0.0f; - uint32_t packed = 0; - #pragma unroll - for (int k = 0; k < 4; ++k) { - uint32_t lo = float_to_fp4_e2m1(__bfloat162float(bf[2 * k]) * inv); - uint32_t hi = float_to_fp4_e2m1( - __bfloat162float(bf[2 * k + 1]) * inv); - packed |= ((hi << 4) | (lo & 0xF)) << (k * 8); - } - *reinterpret_cast(row_fp4 + (j << 2)) = packed; - } -} - -int moe_grouped_quant_nvfp4_bf16( - const void* A, const void* expert_of_row, const void* group_off, - const void* sfa_off, const void* src_row, void* out_packed, void* out_sf, - int slots, int K, cudaStream_t stream) -{ - if (!A || !expert_of_row || !group_off || !sfa_off || !out_packed - || !out_sf) return 1; - if (slots <= 0 || K <= 0 || (K & 15) != 0) return 2; - const int num_blocks = K / 16; - const int n_col_blocks = (num_blocks + 3) / 4; - const int threads = 256; - const size_t smem = (size_t)num_blocks * sizeof(float); - moe_grouped_quant_nvfp4_kernel<<>>( - reinterpret_cast(A), - reinterpret_cast(expert_of_row), - reinterpret_cast(group_off), - reinterpret_cast(sfa_off), - reinterpret_cast(src_row), - reinterpret_cast(out_packed), - reinterpret_cast(out_sf), - K, num_blocks, n_col_blocks); - return 0; -} - -// ── Gate and quantise in one pass, for the grouped MoE's down projection ── -// -// The grouped GEMM produces gate and up interleaved in one (slots, 2*inter) -// buffer, and the gate op wants them as two matrices. Slicing columns out of it -// is not free: the halves are strided, so `.contiguous()` copies both -- 67 MB -// a layer at 2048 tokens, to feed an op that then writes another 17 and has it -// read straight back by the quantiser. -// -// Reading the merged buffer directly costs none of that. The silu is computed -// and rounded to bf16 exactly as silu_mul_sm120_bf16 does, so the value that -// reaches the quantiser is the same one it saw before. -__global__ void moe_grouped_silu_quant_nvfp4_kernel( - const __nv_bfloat16* __restrict__ merged, // (slots, 2 * inter) - const int* __restrict__ expert_of_row, - const int* __restrict__ group_off, - const int* __restrict__ sfa_off, - uint8_t* __restrict__ fp4_data, - uint8_t* __restrict__ scale_factors, - int inter, int num_blocks, int n_col_blocks) -{ - const int row = blockIdx.x; - const int e = expert_of_row[row]; - const int local = row - group_off[e]; - const __nv_bfloat16* g_in = merged + (size_t)row * 2 * inter; - const __nv_bfloat16* u_in = g_in + inter; - uint8_t* row_fp4 = fp4_data + (size_t)row * inter / 2; - uint8_t* sf_base = scale_factors + sfa_off[e]; - - extern __shared__ float smem[]; // inter gated values, then scales - float* gated = smem; - float* scales = smem + inter; - - const int tid = threadIdx.x; - for (int i = tid; i < inter; i += blockDim.x) { - const float gv = __bfloat162float(g_in[i]); - const float uv = __bfloat162float(u_in[i]); - // Rounded to bf16 here, as the separate gate kernel does, so the - // quantiser downstream sees the identical value. - gated[i] = __bfloat162float( - __float2bfloat16_rn(gv / (1.0f + __expf(-gv)) * uv)); - } - __syncthreads(); - - for (int b = tid; b < num_blocks; b += blockDim.x) { - float a = 0.0f; - #pragma unroll 4 - for (int j = 0; j < 16; ++j) a = fmaxf(a, fabsf(gated[b * 16 + j])); - const uint8_t ue = float_to_ue4m3_ceil(a * (1.0f / 6.0f)); - const int rb = local / 128, ri = local % 128; - sf_base[(rb * n_col_blocks + (b >> 2)) * 512 + (ri % 32) * 16 - + (ri / 32) * 4 + (b & 3)] = ue; - scales[b] = ue4m3_to_float(ue); - } - __syncthreads(); - - const int half = inter >> 1; - for (int p = tid; p < half; p += blockDim.x) { - const int i = p * 2; - const float s = scales[i >> 4]; - const float inv = (s > 0.0f) ? (1.0f / s) : 0.0f; - row_fp4[p] = (uint8_t)((float_to_fp4_e2m1(gated[i + 1] * inv) << 4) - | (float_to_fp4_e2m1(gated[i] * inv) & 0x0F)); - } -} - -// Warp-per-row form of the same thing. -// -// The block-per-row kernel above gives 256 threads a row of 512 values -- two -// elements each -- behind three barriers and three passes over shared memory, -// so a block reads two kilobytes and then waits. Measured 2.9x off what that -// traffic implies. -// -// Here a warp owns a row and a lane owns one 16-element scale-factor group: -// it reads its own sixteen gate and up values as vectors, gates them, takes -// its own maximum and packs its own eight bytes. Nothing is shared, so there -// are no barriers and no shared memory at all, and each lane has sixteen -// values in flight instead of two. -// -// The arithmetic is the same in the same order, so the output is identical. -__global__ void moe_grouped_silu_quant_nvfp4_warp_kernel( - const __nv_bfloat16* __restrict__ merged, - const int* __restrict__ expert_of_row, - const int* __restrict__ group_off, - const int* __restrict__ sfa_off, - uint8_t* __restrict__ fp4_data, - uint8_t* __restrict__ scale_factors, - int slots, int inter, int num_blocks, int n_col_blocks) -{ - const int warp_in_blk = threadIdx.x >> 5; - const int lane = threadIdx.x & 31; - const int row = blockIdx.x * (blockDim.x >> 5) + warp_in_blk; - if (row >= slots) return; - - const int e = expert_of_row[row]; - const int local = row - group_off[e]; - const __nv_bfloat16* g_in = merged + (size_t)row * 2 * inter; - const __nv_bfloat16* u_in = g_in + inter; - uint8_t* row_fp4 = fp4_data + (size_t)row * inter / 2; - uint8_t* sf_base = scale_factors + sfa_off[e]; - const int rb = local / 128, ri = local % 128; - - for (int b = lane; b < num_blocks; b += 32) { - float gated[16]; - const int base = b * 16; - #pragma unroll - for (int j = 0; j < 16; ++j) { - const float gv = __bfloat162float(g_in[base + j]); - const float uv = __bfloat162float(u_in[base + j]); - gated[j] = __bfloat162float( - __float2bfloat16_rn(gv / (1.0f + __expf(-gv)) * uv)); - } - float a = 0.0f; - #pragma unroll - for (int j = 0; j < 16; ++j) a = fmaxf(a, fabsf(gated[j])); - - const uint8_t ue = float_to_ue4m3_ceil(a * (1.0f / 6.0f)); - sf_base[(rb * n_col_blocks + (b >> 2)) * 512 + (ri % 32) * 16 - + (ri / 32) * 4 + (b & 3)] = ue; - - const float sc = ue4m3_to_float(ue); - const float inv = (sc > 0.0f) ? (1.0f / sc) : 0.0f; - uint8_t* out8 = row_fp4 + (size_t)b * 8; - #pragma unroll - for (int p = 0; p < 8; ++p) { - out8[p] = (uint8_t)((float_to_fp4_e2m1(gated[2 * p + 1] * inv) << 4) - | (float_to_fp4_e2m1(gated[2 * p] * inv) & 0x0F)); - } - } -} - -int moe_grouped_silu_quant_nvfp4_warp_bf16( - const void* merged, const void* expert_of_row, const void* group_off, - const void* sfa_off, void* out_packed, void* out_sf, - int slots, int inter, cudaStream_t stream) -{ - if (!merged || !expert_of_row || !group_off || !sfa_off || !out_packed - || !out_sf) return 1; - if (slots <= 0 || inter <= 0 || (inter & 15) != 0) return 2; - const int num_blocks = inter / 16; - const int n_col_blocks = (num_blocks + 3) / 4; - constexpr int kThreads = 256; - const int rows_per_block = kThreads / 32; - const int grid = (slots + rows_per_block - 1) / rows_per_block; - moe_grouped_silu_quant_nvfp4_warp_kernel<<>>( - reinterpret_cast(merged), - reinterpret_cast(expert_of_row), - reinterpret_cast(group_off), - reinterpret_cast(sfa_off), - reinterpret_cast(out_packed), - reinterpret_cast(out_sf), - slots, inter, num_blocks, n_col_blocks); - return 0; -} - -int moe_grouped_silu_quant_nvfp4_bf16( - const void* merged, const void* expert_of_row, const void* group_off, - const void* sfa_off, void* out_packed, void* out_sf, - int slots, int inter, cudaStream_t stream) -{ - if (!merged || !expert_of_row || !group_off || !sfa_off || !out_packed - || !out_sf) return 1; - if (slots <= 0 || inter <= 0 || (inter & 15) != 0) return 2; - const int num_blocks = inter / 16; - const int n_col_blocks = (num_blocks + 3) / 4; - const size_t smem = ((size_t)inter + num_blocks) * sizeof(float); - moe_grouped_silu_quant_nvfp4_kernel<<>>( - reinterpret_cast(merged), - reinterpret_cast(expert_of_row), - reinterpret_cast(group_off), - reinterpret_cast(sfa_off), - reinterpret_cast(out_packed), - reinterpret_cast(out_sf), - inter, num_blocks, n_col_blocks); - return 0; -} diff --git a/csrc/kernels/quantize.cuh b/csrc/kernels/quantize.cuh index 88c69118..cdc7bded 100644 --- a/csrc/kernels/quantize.cuh +++ b/csrc/kernels/quantize.cuh @@ -325,38 +325,3 @@ void quantize_int8_rowwise_static(const __nv_bfloat16* input, int8_t* output, void dequant_int32_to_bf16(const int32_t* input, __nv_bfloat16* output, const float* d_act_scale, const float* d_weight_scale, int n, cudaStream_t stream = 0); - -// Grouped activation quantiser for the MoE grouped GEMM: every expert's block -// in one launch. Same math as quantize_bf16_to_nvfp4_swizzled; what differs is -// that each group's scale factors go into the Sm1xx atom layout for that -// group's own row count, which is what the block-scaled grouped GEMM reads. -// Quantising per group instead is correct but costs a launch and a host -// iteration per expert -- and a host iteration is what a graph capture cannot -// have. -// -// A (slots, K) bf16, rows already sorted by expert -// expert_of_row (slots,) i32 -// group_off (E + 1,) i32 prefix sums of the per-expert row counts -// sfa_off (E,) i32 byte offset of each group's SF block -// K must be a multiple of 16. Returns 0 on success, nonzero on arg error. -int moe_grouped_quant_nvfp4_bf16( - const void* A, const void* expert_of_row, const void* group_off, - const void* sfa_off, const void* src_row, void* out_packed, void* out_sf, - int slots, int K, cudaStream_t stream); - -// Gate and quantise in one pass: reads the grouped GEMM's merged (slots, -// 2*inter) gate/up output directly, so the strided column halves are never -// copied out. The silu is rounded to bf16 exactly as silu_mul_sm120_bf16 does, -// so the quantiser sees the same value it did when the two were separate. -int moe_grouped_silu_quant_nvfp4_bf16( - const void* merged, const void* expert_of_row, const void* group_off, - const void* sfa_off, void* out_packed, void* out_sf, - int slots, int inter, cudaStream_t stream); - -// Warp-per-row form of the above: a lane owns one 16-element scale-factor -// group and keeps it in registers, so there is no shared memory and no -// barrier. Same arithmetic in the same order, so the output is identical. -int moe_grouped_silu_quant_nvfp4_warp_bf16( - const void* merged, const void* expert_of_row, const void* group_off, - const void* sfa_off, void* out_packed, void* out_sf, - int slots, int inter, cudaStream_t stream); diff --git a/csrc/kernels/qwen35moe_grouped_quant.cu b/csrc/kernels/qwen35moe_grouped_quant.cu new file mode 100644 index 00000000..d9626cd6 --- /dev/null +++ b/csrc/kernels/qwen35moe_grouped_quant.cu @@ -0,0 +1,295 @@ +// Grouped NVFP4 activation quantisers for the qwen3_5_moe MoE path. +// See qwen35moe_grouped_quant.cuh for the tier this is built under. + +#include "qwen35moe_grouped_quant.cuh" + +#include "nvfp4_convert.cuh" + +#include + +// ── Grouped activation quantiser for the MoE grouped GEMM ── +// +// Same math as quantize_bf16_to_nvfp4_swizzled_kernel, block for block; what +// differs is where the scale factors land. The block-scaled GEMM wants each +// group's scales in the Sm1xx atom layout for that group's own row count, and +// that layout blocks rows by 128, so a group beginning at an arbitrary row of a +// jointly-quantised matrix has no contiguous sub-block to point at. Quantising +// per group is correct but costs a launch and a host iteration per expert. +// +// Here a row reads the expert it was sorted by, subtracts its group's first +// row, and indexes its group's own block. Nothing reaches the host, which is +// what lets the surrounding prefill chunk be captured. +__global__ void moe_grouped_quant_nvfp4_kernel( + const __nv_bfloat16* __restrict__ input, + const int* __restrict__ expert_of_row, + const int* __restrict__ group_off, + const int* __restrict__ sfa_off, + const long* __restrict__ src_row, + uint8_t* __restrict__ fp4_data, + uint8_t* __restrict__ scale_factors, + int cols, int num_blocks, int n_col_blocks) +{ + const int row = blockIdx.x; + const int e = expert_of_row[row]; + const int local = row - group_off[e]; // row index inside its group + // Gather while quantising when a permutation is given. Materialising the + // sorted activation first is a full read and a full write of an (S, HID) + // matrix per layer -- 14.5 ms of a 2048-token prefill -- for rows this + // kernel is about to read once anyway. + const size_t in_row = (src_row == nullptr) ? (size_t)row + : (size_t)src_row[row]; + const __nv_bfloat16* row_in = input + in_row * cols; + uint8_t* row_fp4 = fp4_data + (size_t)row * cols / 2; + uint8_t* sf_base = scale_factors + sfa_off[e]; + + extern __shared__ float smem[]; + const int tid = threadIdx.x; + + // Per-16-block amax without atomics. One thread takes eight bf16 (a half + // block), reduces them in registers, and pairs with its neighbour through a + // shuffle -- j and j^1 land on lanes t and t^1 because the block size is + // even. The first version of this kernel used one atomicMax per element and + // ran at 58.3 ms for traffic worth 0.8; the atomics were all of it. + const int vec8 = cols >> 3; + for (int j = tid; j < vec8; j += blockDim.x) { + uint4 v = *reinterpret_cast(&row_in[j << 3]); + const __nv_bfloat16* bf = reinterpret_cast(&v); + float a = 0.0f; + #pragma unroll + for (int i = 0; i < 8; ++i) a = fmaxf(a, fabsf(__bfloat162float(bf[i]))); + a = fmaxf(a, __shfl_xor_sync(0xffffffffu, a, 1)); + if ((j & 1) == 0) smem[j >> 1] = a; + } + __syncthreads(); + + const int rb = local / 128; + const int ri = local % 128; + for (int b = tid; b < num_blocks; b += blockDim.x) { + uint8_t ue_scale = float_to_ue4m3_ceil(smem[b] * (1.0f / 6.0f)); + const int cb = b / 4; + const int ci = b % 4; + sf_base[(rb * n_col_blocks + cb) * 512 + (ri % 32) * 16 + + (ri / 32) * 4 + ci] = ue_scale; + smem[b] = ue4m3_to_float(ue_scale); + } + __syncthreads(); + + // Pack four bytes at a time: eight bf16 in, one uint32 out, and the eight + // share a 16-block so the scale is read once. + const int quads = cols >> 3; + for (int j = tid; j < quads; j += blockDim.x) { + uint4 v = *reinterpret_cast(&row_in[j << 3]); + const __nv_bfloat16* bf = reinterpret_cast(&v); + const float scale = smem[j >> 1]; + const float inv = (scale > 0.0f) ? (1.0f / scale) : 0.0f; + uint32_t packed = 0; + #pragma unroll + for (int k = 0; k < 4; ++k) { + uint32_t lo = float_to_fp4_e2m1(__bfloat162float(bf[2 * k]) * inv); + uint32_t hi = float_to_fp4_e2m1( + __bfloat162float(bf[2 * k + 1]) * inv); + packed |= ((hi << 4) | (lo & 0xF)) << (k * 8); + } + *reinterpret_cast(row_fp4 + (j << 2)) = packed; + } +} + +int moe_grouped_quant_nvfp4_bf16( + const void* A, const void* expert_of_row, const void* group_off, + const void* sfa_off, const void* src_row, void* out_packed, void* out_sf, + int slots, int K, cudaStream_t stream) +{ + if (!A || !expert_of_row || !group_off || !sfa_off || !out_packed + || !out_sf) return 1; + if (slots <= 0 || K <= 0 || (K & 15) != 0) return 2; + const int num_blocks = K / 16; + const int n_col_blocks = (num_blocks + 3) / 4; + const int threads = 256; + const size_t smem = (size_t)num_blocks * sizeof(float); + moe_grouped_quant_nvfp4_kernel<<>>( + reinterpret_cast(A), + reinterpret_cast(expert_of_row), + reinterpret_cast(group_off), + reinterpret_cast(sfa_off), + reinterpret_cast(src_row), + reinterpret_cast(out_packed), + reinterpret_cast(out_sf), + K, num_blocks, n_col_blocks); + return 0; +} + +// ── Gate and quantise in one pass, for the grouped MoE's down projection ── +// +// The grouped GEMM produces gate and up interleaved in one (slots, 2*inter) +// buffer, and the gate op wants them as two matrices. Slicing columns out of it +// is not free: the halves are strided, so `.contiguous()` copies both -- 67 MB +// a layer at 2048 tokens, to feed an op that then writes another 17 and has it +// read straight back by the quantiser. +// +// Reading the merged buffer directly costs none of that. The silu is computed +// and rounded to bf16 exactly as silu_mul_sm120_bf16 does, so the value that +// reaches the quantiser is the same one it saw before. +__global__ void moe_grouped_silu_quant_nvfp4_kernel( + const __nv_bfloat16* __restrict__ merged, // (slots, 2 * inter) + const int* __restrict__ expert_of_row, + const int* __restrict__ group_off, + const int* __restrict__ sfa_off, + uint8_t* __restrict__ fp4_data, + uint8_t* __restrict__ scale_factors, + int inter, int num_blocks, int n_col_blocks) +{ + const int row = blockIdx.x; + const int e = expert_of_row[row]; + const int local = row - group_off[e]; + const __nv_bfloat16* g_in = merged + (size_t)row * 2 * inter; + const __nv_bfloat16* u_in = g_in + inter; + uint8_t* row_fp4 = fp4_data + (size_t)row * inter / 2; + uint8_t* sf_base = scale_factors + sfa_off[e]; + + extern __shared__ float smem[]; // inter gated values, then scales + float* gated = smem; + float* scales = smem + inter; + + const int tid = threadIdx.x; + for (int i = tid; i < inter; i += blockDim.x) { + const float gv = __bfloat162float(g_in[i]); + const float uv = __bfloat162float(u_in[i]); + // Rounded to bf16 here, as the separate gate kernel does, so the + // quantiser downstream sees the identical value. + gated[i] = __bfloat162float( + __float2bfloat16_rn(gv / (1.0f + __expf(-gv)) * uv)); + } + __syncthreads(); + + for (int b = tid; b < num_blocks; b += blockDim.x) { + float a = 0.0f; + #pragma unroll 4 + for (int j = 0; j < 16; ++j) a = fmaxf(a, fabsf(gated[b * 16 + j])); + const uint8_t ue = float_to_ue4m3_ceil(a * (1.0f / 6.0f)); + const int rb = local / 128, ri = local % 128; + sf_base[(rb * n_col_blocks + (b >> 2)) * 512 + (ri % 32) * 16 + + (ri / 32) * 4 + (b & 3)] = ue; + scales[b] = ue4m3_to_float(ue); + } + __syncthreads(); + + const int half = inter >> 1; + for (int p = tid; p < half; p += blockDim.x) { + const int i = p * 2; + const float s = scales[i >> 4]; + const float inv = (s > 0.0f) ? (1.0f / s) : 0.0f; + row_fp4[p] = (uint8_t)((float_to_fp4_e2m1(gated[i + 1] * inv) << 4) + | (float_to_fp4_e2m1(gated[i] * inv) & 0x0F)); + } +} + +// Warp-per-row form of the same thing. +// +// The block-per-row kernel above gives 256 threads a row of 512 values -- two +// elements each -- behind three barriers and three passes over shared memory, +// so a block reads two kilobytes and then waits. Measured 2.9x off what that +// traffic implies. +// +// Here a warp owns a row and a lane owns one 16-element scale-factor group: +// it reads its own sixteen gate and up values as vectors, gates them, takes +// its own maximum and packs its own eight bytes. Nothing is shared, so there +// are no barriers and no shared memory at all, and each lane has sixteen +// values in flight instead of two. +// +// The arithmetic is the same in the same order, so the output is identical. +__global__ void moe_grouped_silu_quant_nvfp4_warp_kernel( + const __nv_bfloat16* __restrict__ merged, + const int* __restrict__ expert_of_row, + const int* __restrict__ group_off, + const int* __restrict__ sfa_off, + uint8_t* __restrict__ fp4_data, + uint8_t* __restrict__ scale_factors, + int slots, int inter, int num_blocks, int n_col_blocks) +{ + const int warp_in_blk = threadIdx.x >> 5; + const int lane = threadIdx.x & 31; + const int row = blockIdx.x * (blockDim.x >> 5) + warp_in_blk; + if (row >= slots) return; + + const int e = expert_of_row[row]; + const int local = row - group_off[e]; + const __nv_bfloat16* g_in = merged + (size_t)row * 2 * inter; + const __nv_bfloat16* u_in = g_in + inter; + uint8_t* row_fp4 = fp4_data + (size_t)row * inter / 2; + uint8_t* sf_base = scale_factors + sfa_off[e]; + const int rb = local / 128, ri = local % 128; + + for (int b = lane; b < num_blocks; b += 32) { + float gated[16]; + const int base = b * 16; + #pragma unroll + for (int j = 0; j < 16; ++j) { + const float gv = __bfloat162float(g_in[base + j]); + const float uv = __bfloat162float(u_in[base + j]); + gated[j] = __bfloat162float( + __float2bfloat16_rn(gv / (1.0f + __expf(-gv)) * uv)); + } + float a = 0.0f; + #pragma unroll + for (int j = 0; j < 16; ++j) a = fmaxf(a, fabsf(gated[j])); + + const uint8_t ue = float_to_ue4m3_ceil(a * (1.0f / 6.0f)); + sf_base[(rb * n_col_blocks + (b >> 2)) * 512 + (ri % 32) * 16 + + (ri / 32) * 4 + (b & 3)] = ue; + + const float sc = ue4m3_to_float(ue); + const float inv = (sc > 0.0f) ? (1.0f / sc) : 0.0f; + uint8_t* out8 = row_fp4 + (size_t)b * 8; + #pragma unroll + for (int p = 0; p < 8; ++p) { + out8[p] = (uint8_t)((float_to_fp4_e2m1(gated[2 * p + 1] * inv) << 4) + | (float_to_fp4_e2m1(gated[2 * p] * inv) & 0x0F)); + } + } +} + +int moe_grouped_silu_quant_nvfp4_warp_bf16( + const void* merged, const void* expert_of_row, const void* group_off, + const void* sfa_off, void* out_packed, void* out_sf, + int slots, int inter, cudaStream_t stream) +{ + if (!merged || !expert_of_row || !group_off || !sfa_off || !out_packed + || !out_sf) return 1; + if (slots <= 0 || inter <= 0 || (inter & 15) != 0) return 2; + const int num_blocks = inter / 16; + const int n_col_blocks = (num_blocks + 3) / 4; + constexpr int kThreads = 256; + const int rows_per_block = kThreads / 32; + const int grid = (slots + rows_per_block - 1) / rows_per_block; + moe_grouped_silu_quant_nvfp4_warp_kernel<<>>( + reinterpret_cast(merged), + reinterpret_cast(expert_of_row), + reinterpret_cast(group_off), + reinterpret_cast(sfa_off), + reinterpret_cast(out_packed), + reinterpret_cast(out_sf), + slots, inter, num_blocks, n_col_blocks); + return 0; +} + +int moe_grouped_silu_quant_nvfp4_bf16( + const void* merged, const void* expert_of_row, const void* group_off, + const void* sfa_off, void* out_packed, void* out_sf, + int slots, int inter, cudaStream_t stream) +{ + if (!merged || !expert_of_row || !group_off || !sfa_off || !out_packed + || !out_sf) return 1; + if (slots <= 0 || inter <= 0 || (inter & 15) != 0) return 2; + const int num_blocks = inter / 16; + const int n_col_blocks = (num_blocks + 3) / 4; + const size_t smem = ((size_t)inter + num_blocks) * sizeof(float); + moe_grouped_silu_quant_nvfp4_kernel<<>>( + reinterpret_cast(merged), + reinterpret_cast(expert_of_row), + reinterpret_cast(group_off), + reinterpret_cast(sfa_off), + reinterpret_cast(out_packed), + reinterpret_cast(out_sf), + inter, num_blocks, n_col_blocks); + return 0; +} diff --git a/csrc/kernels/qwen35moe_grouped_quant.cuh b/csrc/kernels/qwen35moe_grouped_quant.cuh new file mode 100644 index 00000000..adddd7de --- /dev/null +++ b/csrc/kernels/qwen35moe_grouped_quant.cuh @@ -0,0 +1,49 @@ +#pragma once + +// Grouped NVFP4 activation quantisers for the qwen3_5_moe MoE path. +// +// Built only with the weight-only 4-bit tier +// (-DFLASHRT_ENABLE_QWEN35MOE_W4A16=ON); the matching bindings are guarded on +// FLASHRT_HAVE_QWEN35MOE_W4A16, so a build without that tier contains neither +// these translation units nor their symbols. They live here rather than in +// quantize.cu because the layout they write is the grouped GEMM's, not the +// general quantiser's: scale factors go into the Sm1xx atom layout for each +// group's own row count. + +#include +#include + +// Grouped activation quantiser for the MoE grouped GEMM: every expert's block +// in one launch. Same math as quantize_bf16_to_nvfp4_swizzled; what differs is +// that each group's scale factors go into the Sm1xx atom layout for that +// group's own row count, which is what the block-scaled grouped GEMM reads. +// Quantising per group instead is correct but costs a launch and a host +// iteration per expert -- and a host iteration is what a graph capture cannot +// have. +// +// A (slots, K) bf16, rows already sorted by expert +// expert_of_row (slots,) i32 +// group_off (E + 1,) i32 prefix sums of the per-expert row counts +// sfa_off (E,) i32 byte offset of each group's SF block +// K must be a multiple of 16. Returns 0 on success, nonzero on arg error. +int moe_grouped_quant_nvfp4_bf16( + const void* A, const void* expert_of_row, const void* group_off, + const void* sfa_off, const void* src_row, void* out_packed, void* out_sf, + int slots, int K, cudaStream_t stream); + +// Gate and quantise in one pass: reads the grouped GEMM's merged (slots, +// 2*inter) gate/up output directly, so the strided column halves are never +// copied out. The silu is rounded to bf16 exactly as silu_mul_sm120_bf16 does, +// so the quantiser sees the same value it did when the two were separate. +int moe_grouped_silu_quant_nvfp4_bf16( + const void* merged, const void* expert_of_row, const void* group_off, + const void* sfa_off, void* out_packed, void* out_sf, + int slots, int inter, cudaStream_t stream); + +// Warp-per-row form of the above: a lane owns one 16-element scale-factor +// group and keeps it in registers, so there is no shared memory and no +// barrier. Same arithmetic in the same order, so the output is identical. +int moe_grouped_silu_quant_nvfp4_warp_bf16( + const void* merged, const void* expert_of_row, const void* group_off, + const void* sfa_off, void* out_packed, void* out_sf, + int slots, int inter, cudaStream_t stream); diff --git a/flash_rt/frontends/torch/nexn2_rtx.py b/flash_rt/frontends/torch/nexn2_rtx.py index c887266c..987ae483 100644 --- a/flash_rt/frontends/torch/nexn2_rtx.py +++ b/flash_rt/frontends/torch/nexn2_rtx.py @@ -34,13 +34,42 @@ ) +# The tier combination each target can build. FLASHRT_ENABLE_QWEN35MOE turns on +# all three tiers including the block-scaled 4-bit MMA one, which needs +# sm_120a/sm_121a; recommending it on a target whose toolchain refuses it sends +# the reader to a configure error. So the advice is keyed by the device in +# front of them. +_TIER_ADVICE = { + (11, 0): ("-DGPU_ARCH=110 -DFLASHRT_ENABLE_QWEN35MOE_CORE=ON " + "-DFLASHRT_ENABLE_QWEN35MOE_W4A16=ON"), + (12, 0): "-DGPU_ARCH=120 -DFLASHRT_ENABLE_QWEN35MOE=ON", + (12, 1): "-DGPU_ARCH=121 -DFLASHRT_ENABLE_QWEN35MOE=ON", +} + + +def _build_advice() -> str: + """The configure flags for the device this process is actually running.""" + try: + import torch + + cap = torch.cuda.get_device_capability() + except Exception: # pragma: no cover + return ("-DFLASHRT_ENABLE_QWEN35MOE=ON on sm_120a/sm_121a, or " + "-DFLASHRT_ENABLE_QWEN35MOE_CORE=ON " + "-DFLASHRT_ENABLE_QWEN35MOE_W4A16=ON elsewhere") + return _TIER_ADVICE.get( + cap, + f"the tiers sm_{cap[0]}{cap[1]} can compile " + "(FLASHRT_ENABLE_QWEN35MOE_CORE / _W4A16; _W4A4 needs sm_120a)") + + def _require_kernels( fvk, *, model_label: str = "Nex-N2", usage_doc: str = "docs/nexn2_usage.md", required=None, require_fa2: bool = True) -> None: """Raise a clear RuntimeError if the gated qwen3_5_moe kernels or the FA2 - module are missing (build was not configured with - -DFLASHRT_ENABLE_QWEN35MOE=ON, or flash_rt_fa2 is absent). + module are missing (the build did not enable the qwen3_5_moe tiers, or + flash_rt_fa2 is absent). ``required`` lets a configuration that calls fewer kernels say so. A list demanding more than a path uses turns a working build into a refusal; one @@ -50,15 +79,14 @@ def _require_kernels( missing = [s for s in (required or _REQUIRED_FVK) if not hasattr(fvk, s)] if missing: raise RuntimeError( - f"{model_label} kernelized path needs the qwen3_5_moe SM120 " - "kernels, which " - "are absent from flash_rt_kernels (missing: " - f"{', '.join(missing)}). Rebuild on an SM120 toolchain with " - f"-DFLASHRT_ENABLE_QWEN35MOE=ON. See {usage_doc}.") + f"{model_label} kernelized path needs the gated qwen3_5_moe " + "kernels, which are absent from flash_rt_kernels (missing: " + f"{', '.join(missing)}). Reconfigure with {_build_advice()}. " + f"See {usage_doc}.") if not require_fa2: # The attention backend probes its kernel and falls back to a - # reference implementation, so a target that builds no FA2 -- Thor - # uses FA4 instead -- still runs. + # reference implementation, so a target that builds no FA2 still runs + # -- more slowly on a long prompt, and never differently. return try: from flash_rt import flash_rt_fa2 as _fa2 @@ -66,7 +94,8 @@ def _require_kernels( raise RuntimeError( f"{model_label} full attention needs the vendored FA2 module " "(flash_rt_fa2), which failed to import. Build with FA2 enabled " - "(ENABLE_FA2, auto-on for SM120).") from e + "(automatic on sm_80/86/87/89/120/121; on Thor sm_110 it is " + "opt-in with -DFLASHRT_ENABLE_THOR_FA2=ON).") from e fa2_missing = [s for s in ('fwd_bf16', 'fwd_bf16_causal') if not hasattr(_fa2, s)] if fa2_missing: # pragma: no cover From 4471c4aca858d29961dde13225dd02dda933db94 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 6 Aug 2026 15:21:38 -0400 Subject: [PATCH 78/85] Ask for the deterministic cuBLASLt pick per call, not per process Importing the qwen3_5_moe forward set FLASHRT_BF16_CUBLASLT_AUTOTUNE_ALGOS in the environment. The variable is process-global and shared with every frontend, so a model loaded later in the same process inherited an autotune setting it never asked for -- from an import. The GEMM entry takes the bound as an argument instead, defaulting to 0, which is the environment-driven behaviour every existing call site keeps. Plans are cached per (M, N, K, max_algos), so one caller asking for the heuristic's own pick does not decide the algorithm for another. Gather the rest of this path's kernel choices into a KernelPolicy the frontend owns, rather than deciding each at its call site by asking the module which symbols it happens to export. Every field selects between implementations checked against each other with torch.equal, so a field decides speed and cannot decide output; the environment variables that predate it are its defaults. The attention backend likewise takes its decode-FA2 choice as an argument, keeping the per-architecture default when the caller does not say. --- csrc/bindings.cpp | 10 +- csrc/kernels/bf16_matmul_bf16.cu | 30 ++-- csrc/kernels/bf16_matmul_bf16.cuh | 15 +- flash_rt/frontends/torch/_nexn2_rtx_decode.py | 49 ++++-- .../frontends/torch/_nexn2_rtx_forward.py | 153 ++++++++++++++++-- flash_rt/hardware/rtx/attn_backend_nexn2.py | 24 +-- 6 files changed, 224 insertions(+), 57 deletions(-) diff --git a/csrc/bindings.cpp b/csrc/bindings.cpp index 65a18b93..60679c64 100644 --- a/csrc/bindings.cpp +++ b/csrc/bindings.cpp @@ -4694,17 +4694,21 @@ PYBIND11_MODULE(flash_rt_kernels, m) { py::arg("x"), py::arg("W"), py::arg("out"), py::arg("M"), py::arg("N"), py::arg("K"), py::arg("stream") = 0); + // max_algos=0 (the default) keeps the environment-driven autotune this + // entry has always had; a caller passes 1 to take the heuristic's own pick + // and get a run-to-run reproducible reduction order. See the header. m.def("bf16_matmul_cublaslt_bf16", [](uintptr_t x, uintptr_t W, uintptr_t out, - int M, int N, int K, uintptr_t stream) { + int M, int N, int K, uintptr_t stream, int max_algos) { flash_rt::kernels::bf16_matmul_cublaslt_bf16( reinterpret_cast(x), reinterpret_cast(W), reinterpret_cast<__nv_bfloat16*>(out), - M, N, K, to_stream(stream)); + M, N, K, to_stream(stream), max_algos); }, py::arg("x"), py::arg("W"), py::arg("out"), - py::arg("M"), py::arg("N"), py::arg("K"), py::arg("stream") = 0); + py::arg("M"), py::arg("N"), py::arg("K"), py::arg("stream") = 0, + py::arg("max_algos") = 0); #ifdef FLASHRT_HAVE_QWEN36_KERNELS m.def("bf16_matmul_qwen36_bf16", diff --git a/csrc/kernels/bf16_matmul_bf16.cu b/csrc/kernels/bf16_matmul_bf16.cu index 32d32cc7..6017dbb5 100644 --- a/csrc/kernels/bf16_matmul_bf16.cu +++ b/csrc/kernels/bf16_matmul_bf16.cu @@ -50,9 +50,14 @@ struct Bf16LtKey { int M; int N; int K; + // Part of the key, not just of the search: a caller that asked for a + // deterministic pick and one that let the timing loop run must not share + // whichever plan happened to be built first. + int max_algos; bool operator==(const Bf16LtKey& other) const { - return M == other.M && N == other.N && K == other.K; + return M == other.M && N == other.N && K == other.K + && max_algos == other.max_algos; } }; @@ -61,6 +66,7 @@ struct Bf16LtKeyHash { size_t h = static_cast(key.M); h = h * 1315423911u + static_cast(key.N); h = h * 1315423911u + static_cast(key.K); + h = h * 1315423911u + static_cast(key.max_algos); return h; } }; @@ -72,7 +78,10 @@ static std::mutex g_bf16_mu; static std::unordered_map g_bf16_plans; -static int get_bf16_autotune_algos() { +// requested > 0 is the caller's own bound; 0 falls back to the environment, +// then to the historical default of 8. +static int get_bf16_autotune_algos(int requested) { + if (requested > 0) return std::clamp(requested, 1, 32); const char* env = std::getenv("FLASHRT_BF16_CUBLASLT_AUTOTUNE_ALGOS"); if (!env || !*env) return 8; return std::clamp(std::atoi(env), 1, 32); @@ -94,10 +103,11 @@ static void ensure_bf16_lt() { } } -static Bf16LtPlan& get_bf16_lt_plan(int M, int N, int K) { +static Bf16LtPlan& get_bf16_lt_plan(int M, int N, int K, + int max_algos) { std::lock_guard lock(g_bf16_mu); ensure_bf16_lt(); - Bf16LtKey key{M, N, K}; + Bf16LtKey key{M, N, K, max_algos}; auto it = g_bf16_plans.find(key); if (it != g_bf16_plans.end()) return it->second; @@ -159,9 +169,10 @@ static void autotune_bf16_lt_plan( int M, int N, int K, - cudaStream_t stream) { + cudaStream_t stream, + int max_algos) { if (plan.autotuned) return; - const int num_algos = get_bf16_autotune_algos(); + const int num_algos = get_bf16_autotune_algos(max_algos); if (num_algos <= 1) { plan.autotuned = true; return; @@ -434,12 +445,13 @@ void bf16_matmul_cublaslt_bf16( const __nv_bfloat16* W, __nv_bfloat16* out, int M, int N, int K, - cudaStream_t stream) { + cudaStream_t stream, + int max_algos) { if (M <= 0 || N <= 0 || K <= 0) return; - Bf16LtPlan& plan = get_bf16_lt_plan(M, N, K); + Bf16LtPlan& plan = get_bf16_lt_plan(M, N, K, max_algos); if (!plan.autotuned) { std::lock_guard lock(g_bf16_mu); - autotune_bf16_lt_plan(plan, x, W, out, M, N, K, stream); + autotune_bf16_lt_plan(plan, x, W, out, M, N, K, stream, max_algos); } const float alpha = 1.0f; const float beta = 0.0f; diff --git a/csrc/kernels/bf16_matmul_bf16.cuh b/csrc/kernels/bf16_matmul_bf16.cuh index 9fd36aa4..0fb626bd 100644 --- a/csrc/kernels/bf16_matmul_bf16.cuh +++ b/csrc/kernels/bf16_matmul_bf16.cuh @@ -36,6 +36,18 @@ void bf16_matmul_bf16( int K, cudaStream_t stream); +// ``max_algos`` bounds how many cuBLASLt candidates the first call for a shape +// times before it commits to one. 0 (the default, and what every existing call +// site passes) keeps the current behaviour: the count comes from +// FLASHRT_BF16_CUBLASLT_AUTOTUNE_ALGOS, defaulting to 8. +// +// A caller passes 1 to take the heuristic's own first choice and skip the +// timing loop. That matters beyond speed: timing is noisy, so different +// processes pick different algorithms, different algorithms reduce in +// different orders, and a model whose output is compared token for token then +// disagrees with itself across runs. Plans are cached per (M, N, K, max_algos), +// so one caller asking for a deterministic pick does not decide the algorithm +// for another that did not. void bf16_matmul_cublaslt_bf16( const __nv_bfloat16* x, const __nv_bfloat16* W, @@ -43,6 +55,7 @@ void bf16_matmul_cublaslt_bf16( int M, int N, int K, - cudaStream_t stream); + cudaStream_t stream, + int max_algos = 0); } // namespace flash_rt::kernels diff --git a/flash_rt/frontends/torch/_nexn2_rtx_decode.py b/flash_rt/frontends/torch/_nexn2_rtx_decode.py index c6f96129..c1517c0f 100644 --- a/flash_rt/frontends/torch/_nexn2_rtx_decode.py +++ b/flash_rt/frontends/torch/_nexn2_rtx_decode.py @@ -29,8 +29,8 @@ from flash_rt.frontends.torch._nexn2_rtx_forward import ( CONV, HD, HID, HK, HV, INTER, KD, KS, NKV, NQ, NV, ROPE, TOPK, VD, - _quant_act, _w4a16_mrows, build_rope_tables, moe_grouped_w4a16, - nexn2_forward_nvfp4, set_spec_verify, w4a16_matvec, + _quant_act, _w4a16_mrows, build_rope_tables, kernel_policy, + moe_grouped_w4a16, nexn2_forward_nvfp4, set_spec_verify, w4a16_matvec, ) from flash_rt.frontends.torch._nexn2_rtx_nvfp4_weights import _sf_swz_bytes from flash_rt.hardware.rtx.attn_backend_nexn2 import RtxFlashAttnBackendNexn2 @@ -47,11 +47,6 @@ def _qwen35moe_env(name: str, default: str) -> str: # wins). See Nexn2DecodeState.batched_prefill. _BATCHED_PREFILL_MIN_S = 8 -# Run a speculative window through the decode kernels at w rows rather than -# through the prefill forward. Read once, at import, because a captured graph -# replays whatever branch was taken when it was recorded. -_VERIFY_K_ROWS = _qwen35moe_env("VERIFY_K_ROWS", "1") != "0" - def _cs(): """Current CUDA stream handle. Inside torch.cuda.graph capture this is @@ -183,7 +178,8 @@ def _proj_mma(x2d, ld, base, n, fvk, device, state=None): class Nexn2DecodeState: """Persistent decode state: GDN recurrent/conv caches, KV cache, RoPE.""" - def __init__(self, handles, max_seq, device): + def __init__(self, handles, max_seq, device, *, + spec_graph_cache_max=None): self.handles = handles self.device = device self.max_seq = int(max_seq) @@ -305,8 +301,20 @@ def __init__(self, handles, max_seq, device): self.spec_capture = False # One captured graph per (pos, window): the KV slots, attention length # and RoPE slice are baked per position exactly as the decode graph's - # are. Same LRU bound, since each graph owns a memory pool. + # are, and each owns a memory pool, so this is LRU-bounded the same way. + # + # It is NOT bounded at the same number. A speculative graph covers k+1 + # positions through the whole stack, so its pool is several times a + # decode step's, and the decode cap of 256 is sized for a step. Holding + # 256 of these alongside the model is more than a 32 GB board has at a + # 2048-token context -- measured there, it is what runs it out of + # memory. Sixteen keeps the windows a generation actually revisits + # (recapture costs two warmup runs) and bounds the pools at something + # the smallest supported board carries. self._spec_graphs = collections.OrderedDict() + self.spec_graph_cache_max = int( + spec_graph_cache_max if spec_graph_cache_max is not None + else _qwen35moe_env("SPEC_GRAPH_CACHE_MAX", "16")) # Its own memory pool, not the decode graphs'. The two are replayed # interleaved -- a window, then whatever the caller does next -- and # sharing a pool between graphs used that way is the case the runtime @@ -367,8 +375,11 @@ def router_topk(fvk): over 800 inputs of which 397 had a tie inside the top-8 -- without the block kernel's 24 barriers. """ - fn = getattr(fvk, 'moe_router_topk_warp_sm120_bf16', None) - return fn if fn is not None else fvk.moe_router_topk_sm120_bf16 + if kernel_policy().warp_router_topk: + fn = getattr(fvk, 'moe_router_topk_warp_sm120_bf16', None) + if fn is not None: + return fn + return fvk.moe_router_topk_sm120_bf16 def gdn_recurrent(fvk): @@ -381,8 +392,11 @@ def gdn_recurrent(fvk): memory and walked five times; ncu measures 39 registers per thread for a 128-float array. 51% of bandwidth against 87%. """ - fn = getattr(fvk, 'gated_deltanet_recurrent_edge_qwen36_bf16', None) - return fn if fn is not None else fvk.gated_deltanet_recurrent_qwen36_bf16 + if kernel_policy().gdn_recurrent_edge: + fn = getattr(fvk, 'gated_deltanet_recurrent_edge_qwen36_bf16', None) + if fn is not None: + return fn + return fvk.gated_deltanet_recurrent_qwen36_bf16 def _gdn_gate_consts(ld, device): @@ -619,7 +633,8 @@ def _shared_combine(routed, shared, glog, rows, fvk, device): the window issue, because the fixture and the speculative verify both rest on it. """ - if hasattr(fvk, 'moe_shared_gate_combine_edge_bf16'): + if (kernel_policy().fused_shared_combine + and hasattr(fvk, 'moe_shared_gate_combine_edge_bf16')): out = torch.empty(rows, HID, dtype=torch.bfloat16, device=device) fvk.moe_shared_gate_combine_edge_bf16( routed.data_ptr(), shared.data_ptr(), glog.data_ptr(), @@ -1244,7 +1259,7 @@ def _verify_block_usable(state) -> bool: # with it off decode reads the GDN in_proj at BF16 while the window reads # it at four bits -- a different function in thirty of the forty layers, # which is exactly the thing this block exists to rule out. - if not _VERIFY_K_ROWS or not state.dense_w4a16: + if not kernel_policy().verify_k_rows or not state.dense_w4a16: return False if not state.gdn_in_proj_w4a16: return False @@ -1342,9 +1357,9 @@ def _restore(): _restore() state._spec_graphs[key] = (g, hid) - cap = state.graph_cache_max + cap = state.spec_graph_cache_max if cap > 0 and len(state._spec_graphs) > cap: - state._spec_graphs.popitem(last=False) + state._spec_graphs.popitem(last=False) # evict LRU return state._spec_graphs[key] diff --git a/flash_rt/frontends/torch/_nexn2_rtx_forward.py b/flash_rt/frontends/torch/_nexn2_rtx_forward.py index 31c8fc5f..2025fc45 100644 --- a/flash_rt/frontends/torch/_nexn2_rtx_forward.py +++ b/flash_rt/frontends/torch/_nexn2_rtx_forward.py @@ -152,6 +152,106 @@ def _proj(x2d, ld, base, n, fvk, device): # force the hand-written kernel. import os as _os_early + +class KernelPolicy: + """Which implementation each interchangeable step of this model calls. + + The forward and decode paths have, at several steps, more than one kernel + that computes the same thing: a fused form and the chain it replaces, a + warp-per-row form and a block-per-row one, an "edge" shared-memory layout + and the original. Each pair was checked against the other with + ``torch.equal`` -- not a tolerance -- so which one runs decides speed and + cannot decide output. + + They are gathered here rather than decided at each call site by asking the + module what symbols it happens to export. A kernel appearing in a build is + not a reason to change what an already-validated model path does; a caller + saying so is. The frontend owns one of these and can hand a different one + down, and the environment variables that predate it remain the defaults, so + an existing configuration behaves exactly as it did. + + Fields are read at call time. A policy must therefore not be changed + between a CUDA graph capture and its replay -- the replay repeats whichever + branch the capture took, so the two would disagree. + """ + + __slots__ = ('dense_cublaslt', 'cublaslt_max_algos', 'wy_gdn', + 'edge_w4a16', 'route_kernel', 'fused_shared_combine', + 'warp_router_topk', 'gdn_recurrent_edge', 'verify_k_rows') + + def __init__(self, *, + dense_cublaslt=None, + cublaslt_max_algos=1, + wy_gdn=None, + edge_w4a16=None, + route_kernel=None, + fused_shared_combine=True, + warp_router_topk=True, + gdn_recurrent_edge=True, + verify_k_rows=None): + env = _os_early.environ.get + + def _flag(value, name, default='1'): + if value is not None: + return bool(value) + return env(name, default) != '0' + + # cuBLASLt for the dense bf16 GEMMs; the in-house kernel is also + # deterministic but 66% slower at 2048 (693 against 418 ms). + self.dense_cublaslt = _flag(dense_cublaslt, 'NEXN2_DENSE_CUBLASLT') + # How many cuBLASLt candidates the first call for a shape times. 1 + # takes the heuristic's own pick; see _gemm_w16a16. + self.cublaslt_max_algos = int(cublaslt_max_algos) + # WY chunked gated-delta scan for the GDN prefill instead of the + # sequential scan (11x at S=2048). + self.wy_gdn = _flag(wy_gdn, 'NEXN2_WY_GDN') + # The "edge" shared-memory layout of the two weight-only 4-bit GEMVs. + self.edge_w4a16 = _flag( + edge_w4a16, 'FLASHRT_QWEN35MOE_W4A16_EDGE') + # The five-kernel routing producer instead of the tensor chain. + self.route_kernel = _flag(route_kernel, 'NEXN2_ROUTE_KERNEL') + # Decode-side fusions, each bit-identical to the chain it replaces. + self.fused_shared_combine = bool(fused_shared_combine) + self.warp_router_topk = bool(warp_router_topk) + self.gdn_recurrent_edge = bool(gdn_recurrent_edge) + # Run a speculative verify window through the decode kernels at k+1 + # rows rather than through the prefill forward. Two names, as the rest + # of this model's variables have: the generic one and the one it + # shipped under. + if verify_k_rows is not None: + self.verify_k_rows = bool(verify_k_rows) + else: + self.verify_k_rows = env( + 'FLASHRT_QWEN35MOE_VERIFY_K_ROWS', + env('FLASHRT_NEXN2_VERIFY_K_ROWS', '1')) != '0' + + def __repr__(self) -> str: # pragma: no cover + fields = ', '.join( + f'{name}={getattr(self, name)!r}' for name in self.__slots__) + return f'KernelPolicy({fields})' + + +_POLICY = KernelPolicy() + + +def kernel_policy(): + """The policy the forward and decode paths are currently reading.""" + return _POLICY + + +def set_kernel_policy(policy): + """Install ``policy``; returns the one it replaced. + + Not to be called between a CUDA graph capture and its replay. + """ + global _POLICY + if not isinstance(policy, KernelPolicy): + raise TypeError( + f'expected a KernelPolicy, got {type(policy).__name__}') + previous, _POLICY = _POLICY, policy + return previous + + # The cuBLASLt wrapper picks its algorithm by *timing* eight candidates at # first use. Timing is noisy, so different processes pick different algorithms, # and different algorithms reduce in different orders -- which makes the model @@ -166,27 +266,47 @@ def _proj(x2d, ld, base, n, fvk, device): # the first call. Determinism and a faster first token for 1.6% of the warm # path. # -# Set here rather than in the kernel, whose default is shared with other -# frontends. FLASHRT_BF16_CUBLASLT_AUTOTUNE_ALGOS overrides. -_os_early.environ.setdefault('FLASHRT_BF16_CUBLASLT_AUTOTUNE_ALGOS', '1') +# Requested per call through KernelPolicy.cublaslt_max_algos, not by setting +# the kernel's environment variable: that +# variable is process-global and shared with every other frontend, so setting +# it here would decide the algorithm for a model loaded later in the same +# process that never asked. The kernel caches its plan per +# (M, N, K, max_algos), so this choice stays with these call sites. -# NEXN2_DENSE_CUBLASLT=0 drops to the in-house bf16 GEMM entirely, which is -# also deterministic but 66% slower at 2048 (693 against 418 ms). -_DENSE_CUBLASLT = _os_early.environ.get('NEXN2_DENSE_CUBLASLT', '1') != '0' +# Set once, on the first call: a build predating the max_algos argument still +# links and still runs, one autotune behaviour older. +_CUBLASLT_TAKES_ALGOS = None def _gemm_w16a16(x2d, w, fvk, device): """y = x @ w.T via the deterministic bf16-act x bf16-weight tensor-core GEMM (fp32 register accumulate). Matches the fp32 path's argmax (cos 1.0) and is bit-identical run-to-run, at ~1.75x the fp32/TF32 op.""" + global _CUBLASLT_TAKES_ALGOS m, k = x2d.shape n = w.shape[0] xc = x2d.contiguous() wc = w.contiguous() y = torch.empty(m, n, dtype=torch.bfloat16, device=device) - if _DENSE_CUBLASLT and hasattr(fvk, 'bf16_matmul_cublaslt_bf16'): - fvk.bf16_matmul_cublaslt_bf16(xc.data_ptr(), wc.data_ptr(), - y.data_ptr(), m, n, k, _cs()) + policy = kernel_policy() + if policy.dense_cublaslt and hasattr(fvk, 'bf16_matmul_cublaslt_bf16'): + algos = policy.cublaslt_max_algos + if _CUBLASLT_TAKES_ALGOS is None: + try: + fvk.bf16_matmul_cublaslt_bf16( + xc.data_ptr(), wc.data_ptr(), y.data_ptr(), m, n, k, + _cs(), algos) + _CUBLASLT_TAKES_ALGOS = True + return y + except TypeError: + _CUBLASLT_TAKES_ALGOS = False + if _CUBLASLT_TAKES_ALGOS: + fvk.bf16_matmul_cublaslt_bf16( + xc.data_ptr(), wc.data_ptr(), y.data_ptr(), m, n, k, _cs(), + algos) + else: + fvk.bf16_matmul_cublaslt_bf16(xc.data_ptr(), wc.data_ptr(), + y.data_ptr(), m, n, k, _cs()) return y fvk.w16a16_gemm_sm120_bf16(xc.data_ptr(), wc.data_ptr(), y.data_ptr(), m, n, k, 1.0, _cs()) @@ -408,7 +528,6 @@ def _silu_mul(g, u, fvk, device): # the inter-chunk state recurrence is sequential -> 11x faster at S=2048, # bit-exact (out cos 0.99998, state cos 0.99997 vs the seq-scan). Default on. import os as _os -_USE_WY_GDN = _os.environ.get('NEXN2_WY_GDN', '1') != '0' _WY_MIN_S = 64 # below this the seq-scan's lower fixed overhead wins @@ -578,7 +697,7 @@ def _gdn_layer(h, ld, fvk, device, eps, cap=None, rank=None, a_bf.data_ptr(), b_bf.data_ptr(), neg.data_ptr(), dtb_c.data_ptr(), g_out.data_ptr(), bo.data_ptr(), B * S, NV, _cs()) - if _USE_WY_GDN and S >= _WY_MIN_S: + if kernel_policy().wy_gdn and S >= _WY_MIN_S: # WY chunked delta-rule scan: 11x faster than the seq-scan at S=2048, # bit-exact. qb/kb carry the 16->32 broadcast heads (src_h = h//2); the # front kernel reads the group leaders and re-expands where it packs, @@ -889,13 +1008,13 @@ def _full_attn_layer(h, ld, ct, st, fvk, device, eps, cap=None, rank=None, # whose index differs per lane. Because the outputs are identical to the bit, # choosing between them is purely a performance decision and cannot move a # token, so preferring the variant needs no accuracy argument. Set -# FLASHRT_QWEN35MOE_W4A16_EDGE=0 to force the original. -_EDGE_W4A16 = _os.environ.get("FLASHRT_QWEN35MOE_W4A16_EDGE", "1") != "0" +# KernelPolicy.edge_w4a16 to False (or FLASHRT_QWEN35MOE_W4A16_EDGE=0) to force +# the original. def w4a16_matvec(fvk): """The dense 4-bit GEMV entry point this build should call.""" - if _EDGE_W4A16: + if kernel_policy().edge_w4a16: fn = getattr(fvk, 'w4a16_matvec_edge_sm120_bf16', None) if fn is not None: return fn @@ -904,7 +1023,7 @@ def w4a16_matvec(fvk): def moe_grouped_w4a16(fvk): """The grouped per-slot 4-bit GEMV entry point this build should call.""" - if _EDGE_W4A16: + if kernel_policy().edge_w4a16: fn = getattr(fvk, 'moe_grouped_w4a16_edge_sm120_bf16', None) if fn is not None: return fn @@ -1133,7 +1252,6 @@ def _moe_experts_bt(x, ti, tw, ld, fvk, device): _ROUTE_BUF = {} # Off puts the routing back on the tensor chain, which is how the kernel's # output is A/B'd against it end to end rather than only in a probe. -_USE_ROUTE_KERNEL = _os.environ.get('NEXN2_ROUTE_KERNEL', '1') != '0' def _route_constants(S, device): @@ -1197,7 +1315,8 @@ def _route_prefill(logits, fvk, device): Returns None where the kernel is absent, so the tensor chain stays the fallback rather than this being a hard dependency. """ - if not _USE_ROUTE_KERNEL or not hasattr(fvk, 'moe_route_prefill_bf16'): + if (not kernel_policy().route_kernel + or not hasattr(fvk, 'moe_route_prefill_bf16')): return None S = logits.shape[0] b = _route_buffers(S, fvk, device) diff --git a/flash_rt/hardware/rtx/attn_backend_nexn2.py b/flash_rt/hardware/rtx/attn_backend_nexn2.py index 6fceeef3..3c24f354 100644 --- a/flash_rt/hardware/rtx/attn_backend_nexn2.py +++ b/flash_rt/hardware/rtx/attn_backend_nexn2.py @@ -47,7 +47,8 @@ class RtxFlashAttnBackendNexn2: HEAD_DIM = 256 def __init__(self, max_seq: int, max_q_seq: int = 1, dtype=None, - num_full_layers: int | None = None): + num_full_layers: int | None = None, + use_fa2: bool | None = None): import torch self._torch = torch @@ -114,16 +115,19 @@ def __init__(self, max_seq: int, max_q_seq: int = 1, dtype=None, # difference in one step flips a later one. That is the fixture losing # its meaning in exchange for 1%, which is the wrong trade. # - # So the default follows what each target already validated: on the - # arch that has always had FA2 in decode, keep it; on sm_110, where - # FA2 has only just started building and the fixture was recorded - # through the reference path, decline it. FLASHRT_NEXN2_DECODE_FA2 - # overrides either way. + # So the caller says. ``use_fa2=None`` keeps what each target already + # validated: on the arch that has always had FA2 in decode, keep it; + # on sm_110, where FA2 has only just started building and the fixture + # was recorded through the reference path, decline it. + # FLASHRT_NEXN2_DECODE_FA2 overrides either way. import os as _os - _cap = torch.cuda.get_device_capability() - _default = "0" if _cap == (11, 0) else "1" - want_fa2 = _os.environ.get( - "FLASHRT_NEXN2_DECODE_FA2", _default) != "0" + if use_fa2 is None: + _cap = torch.cuda.get_device_capability() + _default = "0" if _cap == (11, 0) else "1" + want_fa2 = _os.environ.get( + "FLASHRT_NEXN2_DECODE_FA2", _default) != "0" + else: + want_fa2 = bool(use_fa2) try: from flash_rt import flash_rt_fa2 as _fa2 except ImportError: From 3b05cccbfd9705ff0073e59c1422586160bd2228 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 6 Aug 2026 15:21:47 -0400 Subject: [PATCH 79/85] Register Qwen3.6 for Thor and give speculation a supported API The frontend was registered for RTX SM120 only and named for it, while the branch documents Thor as supported -- so a Thor user had to construct an Rtx-named class directly. Nothing in it is architecture-specific: it moves to flash_rt.frontends.torch.qwen36_moe as Qwen36MoeTextFrontend and both architectures resolve to that one class. The previous module and class name re-export it, so an existing import keeps working. Speculative decode was reachable only by setting a private attribute on a subclass, which the error message then recommended. It is a constructor argument now, load_mtp=False by default, because the draft head is a transformer layer's worth of weights that plain generate never reads. Its graph cache no longer borrows the decode cap of 256. A speculative window covers k+1 positions through the whole stack, so its memory pool is several times a decode step's, and 256 of them is more than a 32 GB board has at a 2048-token context -- which is where it was measured running out. The default is 16, and spec_graph_cache_max sets it. --- benchmarks/qwen36_moe_edge_decode.py | 4 +- flash_rt/api.py | 4 +- flash_rt/frontends/torch/qwen36_moe.py | 426 +++++++++++++++++++++ flash_rt/frontends/torch/qwen36_moe_rtx.py | 380 +----------------- flash_rt/hardware/__init__.py | 25 +- qwen36_moe_edge/expert_quality.py | 6 +- qwen36_moe_edge/route_trace.py | 6 +- qwen36_moe_edge/streaming_frontend.py | 6 +- qwen36_moe_edge/warm_start_validation.py | 6 +- tests/test_qwen36_moe_smoke.py | 99 +++-- 10 files changed, 535 insertions(+), 427 deletions(-) create mode 100644 flash_rt/frontends/torch/qwen36_moe.py diff --git a/benchmarks/qwen36_moe_edge_decode.py b/benchmarks/qwen36_moe_edge_decode.py index aff1963e..64c6eba3 100644 --- a/benchmarks/qwen36_moe_edge_decode.py +++ b/benchmarks/qwen36_moe_edge_decode.py @@ -41,7 +41,7 @@ def main() -> None: parser.add_argument("--device", default="cuda:0") args = parser.parse_args() - from flash_rt.frontends.torch.qwen36_moe_rtx import Qwen36MoeTextFrontendRtx + from flash_rt.frontends.torch.qwen36_moe import Qwen36MoeTextFrontend # Select and initialise the device before touching the memory stats: they # are per-device counters and are not addressable until then. @@ -49,7 +49,7 @@ def main() -> None: torch.cuda.init() torch.cuda.reset_peak_memory_stats(args.device) t0 = time.perf_counter() - frontend = Qwen36MoeTextFrontendRtx( + frontend = Qwen36MoeTextFrontend( args.checkpoint, device=args.device, max_seq=args.max_seq) _sync(args.device) load_s = time.perf_counter() - t0 diff --git a/flash_rt/api.py b/flash_rt/api.py index 13d4e20e..a29b163b 100644 --- a/flash_rt/api.py +++ b/flash_rt/api.py @@ -625,8 +625,8 @@ def load_model(checkpoint, framework="torch", num_views=2, autotune=3, raise NotImplementedError( "config='qwen36_moe' is a text LLM and is not served through " "load_model's VLA wrapper. Construct it directly:\n" - " from flash_rt.frontends.torch.qwen36_moe_rtx import " - "Qwen36MoeTextFrontendRtx\n" + " from flash_rt.frontends.torch.qwen36_moe import " + "Qwen36MoeTextFrontend\n" "See docs/qwen36_moe_usage.md.") from flash_rt.hardware import detect_arch, resolve_pipeline_class diff --git a/flash_rt/frontends/torch/qwen36_moe.py b/flash_rt/frontends/torch/qwen36_moe.py new file mode 100644 index 00000000..8684224d --- /dev/null +++ b/flash_rt/frontends/torch/qwen36_moe.py @@ -0,0 +1,426 @@ +"""Qwen3.6-35B-A3B text inference. + +The language backbone is the same ``qwen3_5_moe`` architecture used by +Nex-N2-mini, so this frontend reuses that implementation. Nothing in it is +specific to one GPU: it is registered for RTX SM120 and for Jetson AGX Thor +(SM110), and each target's build tiers decide which kernels the shared forward +and decode paths resolve to. (The base class keeps its ``Rtx`` name, which +predates the Thor path; see :mod:`flash_rt.frontends.torch.nexn2_rtx`.) + +The official Qwen3.6 checkpoint also contains a vision tower and an MTP draft +head. The vision tower is validated but not executed here. The draft head is +loaded only when ``load_mtp=True`` is passed, which is what +:meth:`Qwen36MoeTextFrontend.generate_spec` needs; ``generate`` never uses it. + +See ``docs/qwen36_moe_usage.md`` for the per-architecture build commands. +""" + +from __future__ import annotations + +import json +import os +from contextlib import ExitStack +from typing import Any + +from flash_rt.frontends.torch.nexn2_rtx import Nexn2TorchFrontendRtx + + +_EXPECTED_LAYER_TYPES = tuple( + "full_attention" if (i + 1) % 4 == 0 else "linear_attention" + for i in range(40) +) + +_EXPECTED_TEXT_CONFIG = { + "model_type": "qwen3_5_moe_text", + "num_hidden_layers": 40, + "hidden_size": 2048, + "vocab_size": 248320, + "num_attention_heads": 16, + "num_key_value_heads": 2, + "head_dim": 256, + "attention_bias": False, + "attn_output_gate": True, + "hidden_act": "silu", + "num_experts": 256, + "num_experts_per_tok": 8, + "moe_intermediate_size": 512, + "shared_expert_intermediate_size": 512, + "linear_num_key_heads": 16, + "linear_num_value_heads": 32, + "linear_key_head_dim": 128, + "linear_value_head_dim": 128, + "linear_conv_kernel_dim": 4, + "mamba_ssm_dtype": "float32", + "full_attention_interval": 4, + "partial_rotary_factor": 0.25, + "rms_norm_eps": 1e-6, + "mtp_num_hidden_layers": 1, + "mtp_use_dedicated_embeddings": False, + "tie_word_embeddings": False, +} + +_EXPECTED_ROPE_PARAMETERS = { + "rope_type": "default", + "rope_theta": 10000000, + "partial_rotary_factor": 0.25, + "mrope_interleaved": True, + "mrope_section": [11, 11, 10], +} + + +def _expected_mlp_shapes(prefix: str) -> dict[str, tuple[int, ...]]: + return { + prefix + "mlp.gate.weight": (256, 2048), + prefix + "mlp.experts.gate_up_proj": (256, 1024, 2048), + prefix + "mlp.experts.down_proj": (256, 2048, 512), + prefix + "mlp.shared_expert.gate_proj.weight": (512, 2048), + prefix + "mlp.shared_expert.up_proj.weight": (512, 2048), + prefix + "mlp.shared_expert.down_proj.weight": (2048, 512), + prefix + "mlp.shared_expert_gate.weight": (1, 2048), + } + + +def _expected_attention_shapes( + prefix: str, layer_type: str) -> dict[str, tuple[int, ...]]: + if layer_type == "full_attention": + return { + prefix + "self_attn.q_proj.weight": (8192, 2048), + prefix + "self_attn.k_proj.weight": (512, 2048), + prefix + "self_attn.v_proj.weight": (512, 2048), + prefix + "self_attn.o_proj.weight": (2048, 4096), + prefix + "self_attn.q_norm.weight": (256,), + prefix + "self_attn.k_norm.weight": (256,), + } + return { + prefix + "linear_attn.in_proj_qkv.weight": (8192, 2048), + prefix + "linear_attn.in_proj_z.weight": (4096, 2048), + prefix + "linear_attn.in_proj_a.weight": (32, 2048), + prefix + "linear_attn.in_proj_b.weight": (32, 2048), + prefix + "linear_attn.conv1d.weight": (8192, 1, 4), + prefix + "linear_attn.A_log": (32,), + prefix + "linear_attn.dt_bias": (32,), + prefix + "linear_attn.norm.weight": (128,), + prefix + "linear_attn.out_proj.weight": (2048, 4096), + } + + +def _expected_text_shapes( + layer_types: tuple[str, ...]) -> dict[str, tuple[int, ...]]: + shapes = { + "lm_head.weight": (248320, 2048), + "model.language_model.embed_tokens.weight": (248320, 2048), + "model.language_model.norm.weight": (2048,), + } + for i, layer_type in enumerate(layer_types): + prefix = f"model.language_model.layers.{i}." + shapes[prefix + "input_layernorm.weight"] = (2048,) + shapes[prefix + "post_attention_layernorm.weight"] = (2048,) + shapes.update(_expected_mlp_shapes(prefix)) + shapes.update(_expected_attention_shapes(prefix, layer_type)) + return shapes + + +def _expected_mtp_shapes() -> dict[str, tuple[int, ...]]: + prefix = "mtp.layers.0." + shapes = { + "mtp.fc.weight": (2048, 4096), + "mtp.norm.weight": (2048,), + "mtp.pre_fc_norm_embedding.weight": (2048,), + "mtp.pre_fc_norm_hidden.weight": (2048,), + prefix + "input_layernorm.weight": (2048,), + prefix + "post_attention_layernorm.weight": (2048,), + } + shapes.update(_expected_attention_shapes(prefix, "full_attention")) + shapes.update(_expected_mlp_shapes(prefix)) + return shapes + + +_MTP_KEYS = set(_expected_mtp_shapes()) + + +def _required_text_keys(layer_types: tuple[str, ...]) -> set[str]: + return set(_expected_text_shapes(layer_types)) + + +def _read_tensor_shapes( + checkpoint_path: str, + weight_map: dict[str, str], + tensor_names: set[str], +) -> dict[str, tuple[int, ...]]: + from safetensors import safe_open + + shapes = {} + with ExitStack() as stack: + readers = { + shard: stack.enter_context( + safe_open( + os.path.join(checkpoint_path, shard), + framework="pt", + device="cpu", + ) + ) + for shard in set(weight_map[name] for name in tensor_names) + } + for name in tensor_names: + shapes[name] = tuple( + readers[weight_map[name]].get_slice(name).get_shape()) + return shapes + + +def validate_qwen36_moe_checkpoint(checkpoint_path: str) -> dict[str, Any]: + """Validate the official BF16 checkpoint before allocating GPU weights.""" + checkpoint_path = os.path.abspath(os.fspath(checkpoint_path)) + config_path = os.path.join(checkpoint_path, "config.json") + index_path = os.path.join( + checkpoint_path, "model.safetensors.index.json") + + for path in (config_path, index_path): + if not os.path.isfile(path): + raise FileNotFoundError( + f"Qwen3.6-35B-A3B checkpoint is missing {path!r}") + + with open(config_path, "r", encoding="utf-8") as f: + config = json.load(f) + if config.get("model_type") != "qwen3_5_moe": + raise ValueError( + "Qwen3.6-35B-A3B text frontend requires " + "model_type='qwen3_5_moe'; " + f"got {config.get('model_type')!r} in {config_path}") + + text_config = config.get("text_config") + if not isinstance(text_config, dict): + raise ValueError(f"missing text_config object in {config_path}") + + mismatches = [] + for name, expected in _EXPECTED_TEXT_CONFIG.items(): + actual = text_config.get(name) + if actual != expected: + mismatches.append(f"{name}={actual!r} (expected {expected!r})") + rope_parameters = text_config.get("rope_parameters") + if not isinstance(rope_parameters, dict): + mismatches.append( + "rope_parameters is missing or is not an object") + else: + for name, expected in _EXPECTED_ROPE_PARAMETERS.items(): + actual = rope_parameters.get(name) + if actual != expected: + mismatches.append( + f"rope_parameters.{name}={actual!r} " + f"(expected {expected!r})") + layer_types = tuple(text_config.get("layer_types") or ()) + if layer_types != _EXPECTED_LAYER_TYPES: + mismatches.append( + "layer_types does not match the 30-linear/10-full attention " + "qwen3_5_moe schedule") + if mismatches: + raise ValueError( + "checkpoint is not compatible with the Qwen3.6-35B-A3B " + "SM120 text pipeline: " + "; ".join(mismatches)) + + with open(index_path, "r", encoding="utf-8") as f: + index = json.load(f) + weight_map = index.get("weight_map") + if not isinstance(weight_map, dict): + raise ValueError(f"missing weight_map object in {index_path}") + + required_text = _required_text_keys(layer_types) + missing_text = sorted(required_text.difference(weight_map)) + if missing_text: + preview = ", ".join(missing_text[:8]) + if len(missing_text) > 8: + preview += f", ... ({len(missing_text)} missing)" + raise ValueError( + "Qwen3.6-35B-A3B checkpoint is missing text backbone tensors: " + + preview) + + missing_mtp = sorted(_MTP_KEYS.difference(weight_map)) + if missing_mtp: + preview = ", ".join(missing_mtp[:8]) + if len(missing_mtp) > 8: + preview += f", ... ({len(missing_mtp)} missing)" + raise ValueError( + "Qwen3.6-35B-A3B checkpoint is missing its MTP tensor group: " + + preview) + + shards = sorted(set(weight_map.values())) + bad_shards = [] + for shard in shards: + path = os.path.join(checkpoint_path, shard) + if not os.path.isfile(path) or os.path.getsize(path) == 0: + bad_shards.append(shard) + if bad_shards: + preview = ", ".join(bad_shards[:8]) + if len(bad_shards) > 8: + preview += f", ... ({len(bad_shards)} missing or empty)" + raise FileNotFoundError( + "Qwen3.6-35B-A3B checkpoint has missing or empty shards: " + + preview) + + expected_shapes = _expected_text_shapes(layer_types) + expected_shapes.update(_expected_mtp_shapes()) + actual_shapes = _read_tensor_shapes( + checkpoint_path, weight_map, set(expected_shapes)) + shape_mismatches = [ + f"{name}={actual_shapes[name]!r} (expected {expected!r})" + for name, expected in sorted(expected_shapes.items()) + if actual_shapes[name] != expected + ] + if shape_mismatches: + preview = "; ".join(shape_mismatches[:8]) + if len(shape_mismatches) > 8: + preview += f"; ... ({len(shape_mismatches)} mismatched)" + raise ValueError( + "Qwen3.6-35B-A3B checkpoint tensor shape mismatches: " + + preview) + + return { + "checkpoint_path": checkpoint_path, + "text_tensor_count": len(required_text), + "mtp_tensor_count": len(_MTP_KEYS), + "vision_tensor_count": sum( + ".visual." in name for name in weight_map), + "tensor_count": len(weight_map), + "shard_count": len(shards), + } + + +class Qwen36MoeTextFrontend(Nexn2TorchFrontendRtx): + """Qwen3.6-35B-A3B text-only frontend (RTX SM120 and Jetson AGX Thor).""" + + _MODEL_LABEL = "Qwen3.6-35B-A3B text" + _USAGE_DOC = "docs/qwen36_moe_usage.md" + + # The block-scaled 4-bit MMA tier is a build tier, and the prefill now picks + # its MoE tile from what the module actually has: without the tier it uses + # the weight-only grouped GEMV, which is slower on a long prompt and + # otherwise identical. So the tier is a performance requirement, not a + # correctness one, and demanding it here would refuse a build -- a Jetson + # one, whose toolchain has no block-scaled mma -- that runs this correctly. + _REQUIRED_KERNELS = tuple( + name for name in Nexn2TorchFrontendRtx._REQUIRED_KERNELS + if not name.startswith(('moe_blocktile_mma', 'moe_m16_mma')) + ) + ('moe_grouped_w4a16_sm120_bf16',) + + # Same for FA2: prefill and decode both probe the vendored kernel and fall + # back to a reference, so a target that builds FA4 instead still runs. + _REQUIRE_FA2 = False + + def __init__(self, checkpoint_path: str, *, + device: str = "cuda:0", + max_seq: int = 2048, + quant: str = "nvfp4", + kernelized: bool = True, + quant_scope: str = "experts", + load_mtp: bool = False, + spec_graph_cache_max: int | None = None) -> None: + """Construct the frontend. + + Args beyond the shared ones (see + :class:`~flash_rt.frontends.torch.nexn2_rtx.Nexn2TorchFrontendRtx`): + + load_mtp: read the checkpoint's MTP draft head, which + :meth:`generate_spec` needs. Off by default: it is a full extra + transformer layer plus its KV, and ``generate`` never reads it. + spec_graph_cache_max: how many captured speculative windows to keep. + Each owns a CUDA graph memory pool, and a speculative window is + larger than a decode step, so this is bounded separately from the + decode graph cache and much lower. ``None`` takes the runtime + default, which is sized for the smallest board this path runs on. + """ + if quant != "nvfp4": + raise NotImplementedError( + f"quant={quant!r} is not implemented; only 'nvfp4' is " + "supported") + if not kernelized: + raise NotImplementedError( + "Qwen3.6-35B-A3B text only supports kernelized=True with " + "runtime NVFP4 conversion") + if spec_graph_cache_max is not None: + spec_graph_cache_max = int(spec_graph_cache_max) + if spec_graph_cache_max < 1: + raise ValueError( + "spec_graph_cache_max must be at least 1 (got " + f"{spec_graph_cache_max}); a speculative step captures a " + "graph per position, so a cache that holds none would " + "recapture every step") + # Read by the loader, before any weight is touched, through the base + # class -- hence set before super().__init__. + self._load_mtp = bool(load_mtp) + self._spec_graph_cache_max = spec_graph_cache_max + contract = validate_qwen36_moe_checkpoint(checkpoint_path) + super().__init__( + checkpoint_path, + device=device, + max_seq=max_seq, + quant=quant, + kernelized=kernelized, + quant_scope=quant_scope, + ) + self._checkpoint_contract = contract + + @property + def load_mtp(self) -> bool: + """Whether the MTP draft head was loaded (see ``generate_spec``).""" + return self._load_mtp + + def _decode_state_or_new(self): + from flash_rt.frontends.torch._nexn2_rtx_decode import Nexn2DecodeState + + if self._decode_state is None: + self._decode_state = Nexn2DecodeState( + self._weights, self._user_max_seq, self.device, + spec_graph_cache_max=self._spec_graph_cache_max) + return self._decode_state + + def generate_spec(self, max_new_tokens: int, *, k: int = 2): + """Greedy decode through draft-and-verify with the MTP head. + + Emits exactly what ``generate`` emits: a draft is kept only where the + model's own argmax agrees with it, so this is a speed change and + nothing else. Requires ``load_mtp=True`` at construction. + """ + if self._prompt_ids is None: + raise ValueError("call set_prompt(...) before generate_spec()") + if not self._load_mtp: + raise RuntimeError( + "speculative decoding needs the MTP draft head, which this " + "frontend did not load. Construct it with load_mtp=True.") + if int(k) < 1: + raise ValueError( + f"k must be at least 1 (got {k}); k is how many tokens the " + "draft head proposes per window") + + from flash_rt.frontends.torch._nexn2_rtx_decode import ( + generate_greedy_spec, + ) + + return generate_greedy_spec( + self._decode_state_or_new(), self._prompt_ids, max_new_tokens, + int(k), self._fvk, self.device) + + def generate(self, max_new_tokens: int, *, do_sample: bool = False): + """Generate tokens with the shared qwen3_5_moe CUDA Graph path.""" + if self._prompt_ids is None: + raise ValueError("call set_prompt(...) before generate()") + if do_sample: + raise NotImplementedError( + "the Qwen3.6-35B-A3B kernelized path supports greedy " + "decoding only") + + from flash_rt.frontends.torch._nexn2_rtx_decode import ( + generate_greedy_graph, + ) + + return generate_greedy_graph( + self._decode_state_or_new(), + self._prompt_ids, + max_new_tokens, + self._fvk, + self.device, + ) + + +# The name this frontend shipped under while it was registered for RTX only. +# Kept so an existing import keeps working; new code should use the class +# above, which is what both architectures resolve to. +Qwen36MoeTextFrontendRtx = Qwen36MoeTextFrontend diff --git a/flash_rt/frontends/torch/qwen36_moe_rtx.py b/flash_rt/frontends/torch/qwen36_moe_rtx.py index f8c2eabb..ffd0ed2a 100644 --- a/flash_rt/frontends/torch/qwen36_moe_rtx.py +++ b/flash_rt/frontends/torch/qwen36_moe_rtx.py @@ -1,374 +1,20 @@ -"""Qwen3.6-35B-A3B text inference on RTX SM120. +"""Compatibility import path for the Qwen3.6-35B-A3B text frontend. -The language backbone is the same ``qwen3_5_moe`` architecture used by -Nex-N2-mini, so this frontend reuses that implementation. The official -Qwen3.6 checkpoint also contains a vision tower and an MTP head; this entry -validates those weights but intentionally exposes only text prefill and greedy -decode. Vision and speculative decoding are separate integration surfaces. +The frontend moved to :mod:`flash_rt.frontends.torch.qwen36_moe` when it stopped +being RTX-only. This module re-exports the public names so an existing import +keeps working; new code should import from the module above. """ from __future__ import annotations -import json -import os -from contextlib import ExitStack -from typing import Any - -from flash_rt.frontends.torch.nexn2_rtx import Nexn2TorchFrontendRtx - - -_EXPECTED_LAYER_TYPES = tuple( - "full_attention" if (i + 1) % 4 == 0 else "linear_attention" - for i in range(40) +from flash_rt.frontends.torch.qwen36_moe import ( + Qwen36MoeTextFrontend, + Qwen36MoeTextFrontendRtx, + validate_qwen36_moe_checkpoint, ) -_EXPECTED_TEXT_CONFIG = { - "model_type": "qwen3_5_moe_text", - "num_hidden_layers": 40, - "hidden_size": 2048, - "vocab_size": 248320, - "num_attention_heads": 16, - "num_key_value_heads": 2, - "head_dim": 256, - "attention_bias": False, - "attn_output_gate": True, - "hidden_act": "silu", - "num_experts": 256, - "num_experts_per_tok": 8, - "moe_intermediate_size": 512, - "shared_expert_intermediate_size": 512, - "linear_num_key_heads": 16, - "linear_num_value_heads": 32, - "linear_key_head_dim": 128, - "linear_value_head_dim": 128, - "linear_conv_kernel_dim": 4, - "mamba_ssm_dtype": "float32", - "full_attention_interval": 4, - "partial_rotary_factor": 0.25, - "rms_norm_eps": 1e-6, - "mtp_num_hidden_layers": 1, - "mtp_use_dedicated_embeddings": False, - "tie_word_embeddings": False, -} - -_EXPECTED_ROPE_PARAMETERS = { - "rope_type": "default", - "rope_theta": 10000000, - "partial_rotary_factor": 0.25, - "mrope_interleaved": True, - "mrope_section": [11, 11, 10], -} - - -def _expected_mlp_shapes(prefix: str) -> dict[str, tuple[int, ...]]: - return { - prefix + "mlp.gate.weight": (256, 2048), - prefix + "mlp.experts.gate_up_proj": (256, 1024, 2048), - prefix + "mlp.experts.down_proj": (256, 2048, 512), - prefix + "mlp.shared_expert.gate_proj.weight": (512, 2048), - prefix + "mlp.shared_expert.up_proj.weight": (512, 2048), - prefix + "mlp.shared_expert.down_proj.weight": (2048, 512), - prefix + "mlp.shared_expert_gate.weight": (1, 2048), - } - - -def _expected_attention_shapes( - prefix: str, layer_type: str) -> dict[str, tuple[int, ...]]: - if layer_type == "full_attention": - return { - prefix + "self_attn.q_proj.weight": (8192, 2048), - prefix + "self_attn.k_proj.weight": (512, 2048), - prefix + "self_attn.v_proj.weight": (512, 2048), - prefix + "self_attn.o_proj.weight": (2048, 4096), - prefix + "self_attn.q_norm.weight": (256,), - prefix + "self_attn.k_norm.weight": (256,), - } - return { - prefix + "linear_attn.in_proj_qkv.weight": (8192, 2048), - prefix + "linear_attn.in_proj_z.weight": (4096, 2048), - prefix + "linear_attn.in_proj_a.weight": (32, 2048), - prefix + "linear_attn.in_proj_b.weight": (32, 2048), - prefix + "linear_attn.conv1d.weight": (8192, 1, 4), - prefix + "linear_attn.A_log": (32,), - prefix + "linear_attn.dt_bias": (32,), - prefix + "linear_attn.norm.weight": (128,), - prefix + "linear_attn.out_proj.weight": (2048, 4096), - } - - -def _expected_text_shapes( - layer_types: tuple[str, ...]) -> dict[str, tuple[int, ...]]: - shapes = { - "lm_head.weight": (248320, 2048), - "model.language_model.embed_tokens.weight": (248320, 2048), - "model.language_model.norm.weight": (2048,), - } - for i, layer_type in enumerate(layer_types): - prefix = f"model.language_model.layers.{i}." - shapes[prefix + "input_layernorm.weight"] = (2048,) - shapes[prefix + "post_attention_layernorm.weight"] = (2048,) - shapes.update(_expected_mlp_shapes(prefix)) - shapes.update(_expected_attention_shapes(prefix, layer_type)) - return shapes - - -def _expected_mtp_shapes() -> dict[str, tuple[int, ...]]: - prefix = "mtp.layers.0." - shapes = { - "mtp.fc.weight": (2048, 4096), - "mtp.norm.weight": (2048,), - "mtp.pre_fc_norm_embedding.weight": (2048,), - "mtp.pre_fc_norm_hidden.weight": (2048,), - prefix + "input_layernorm.weight": (2048,), - prefix + "post_attention_layernorm.weight": (2048,), - } - shapes.update(_expected_attention_shapes(prefix, "full_attention")) - shapes.update(_expected_mlp_shapes(prefix)) - return shapes - - -_MTP_KEYS = set(_expected_mtp_shapes()) - - -def _required_text_keys(layer_types: tuple[str, ...]) -> set[str]: - return set(_expected_text_shapes(layer_types)) - - -def _read_tensor_shapes( - checkpoint_path: str, - weight_map: dict[str, str], - tensor_names: set[str], -) -> dict[str, tuple[int, ...]]: - from safetensors import safe_open - - shapes = {} - with ExitStack() as stack: - readers = { - shard: stack.enter_context( - safe_open( - os.path.join(checkpoint_path, shard), - framework="pt", - device="cpu", - ) - ) - for shard in set(weight_map[name] for name in tensor_names) - } - for name in tensor_names: - shapes[name] = tuple( - readers[weight_map[name]].get_slice(name).get_shape()) - return shapes - - -def validate_qwen36_moe_checkpoint(checkpoint_path: str) -> dict[str, Any]: - """Validate the official BF16 checkpoint before allocating GPU weights.""" - checkpoint_path = os.path.abspath(os.fspath(checkpoint_path)) - config_path = os.path.join(checkpoint_path, "config.json") - index_path = os.path.join( - checkpoint_path, "model.safetensors.index.json") - - for path in (config_path, index_path): - if not os.path.isfile(path): - raise FileNotFoundError( - f"Qwen3.6-35B-A3B checkpoint is missing {path!r}") - - with open(config_path, "r", encoding="utf-8") as f: - config = json.load(f) - if config.get("model_type") != "qwen3_5_moe": - raise ValueError( - "Qwen3.6-35B-A3B text frontend requires " - "model_type='qwen3_5_moe'; " - f"got {config.get('model_type')!r} in {config_path}") - - text_config = config.get("text_config") - if not isinstance(text_config, dict): - raise ValueError(f"missing text_config object in {config_path}") - - mismatches = [] - for name, expected in _EXPECTED_TEXT_CONFIG.items(): - actual = text_config.get(name) - if actual != expected: - mismatches.append(f"{name}={actual!r} (expected {expected!r})") - rope_parameters = text_config.get("rope_parameters") - if not isinstance(rope_parameters, dict): - mismatches.append( - "rope_parameters is missing or is not an object") - else: - for name, expected in _EXPECTED_ROPE_PARAMETERS.items(): - actual = rope_parameters.get(name) - if actual != expected: - mismatches.append( - f"rope_parameters.{name}={actual!r} " - f"(expected {expected!r})") - layer_types = tuple(text_config.get("layer_types") or ()) - if layer_types != _EXPECTED_LAYER_TYPES: - mismatches.append( - "layer_types does not match the 30-linear/10-full attention " - "qwen3_5_moe schedule") - if mismatches: - raise ValueError( - "checkpoint is not compatible with the Qwen3.6-35B-A3B " - "SM120 text pipeline: " + "; ".join(mismatches)) - - with open(index_path, "r", encoding="utf-8") as f: - index = json.load(f) - weight_map = index.get("weight_map") - if not isinstance(weight_map, dict): - raise ValueError(f"missing weight_map object in {index_path}") - - required_text = _required_text_keys(layer_types) - missing_text = sorted(required_text.difference(weight_map)) - if missing_text: - preview = ", ".join(missing_text[:8]) - if len(missing_text) > 8: - preview += f", ... ({len(missing_text)} missing)" - raise ValueError( - "Qwen3.6-35B-A3B checkpoint is missing text backbone tensors: " - + preview) - - missing_mtp = sorted(_MTP_KEYS.difference(weight_map)) - if missing_mtp: - preview = ", ".join(missing_mtp[:8]) - if len(missing_mtp) > 8: - preview += f", ... ({len(missing_mtp)} missing)" - raise ValueError( - "Qwen3.6-35B-A3B checkpoint is missing its MTP tensor group: " - + preview) - - shards = sorted(set(weight_map.values())) - bad_shards = [] - for shard in shards: - path = os.path.join(checkpoint_path, shard) - if not os.path.isfile(path) or os.path.getsize(path) == 0: - bad_shards.append(shard) - if bad_shards: - preview = ", ".join(bad_shards[:8]) - if len(bad_shards) > 8: - preview += f", ... ({len(bad_shards)} missing or empty)" - raise FileNotFoundError( - "Qwen3.6-35B-A3B checkpoint has missing or empty shards: " - + preview) - - expected_shapes = _expected_text_shapes(layer_types) - expected_shapes.update(_expected_mtp_shapes()) - actual_shapes = _read_tensor_shapes( - checkpoint_path, weight_map, set(expected_shapes)) - shape_mismatches = [ - f"{name}={actual_shapes[name]!r} (expected {expected!r})" - for name, expected in sorted(expected_shapes.items()) - if actual_shapes[name] != expected - ] - if shape_mismatches: - preview = "; ".join(shape_mismatches[:8]) - if len(shape_mismatches) > 8: - preview += f"; ... ({len(shape_mismatches)} mismatched)" - raise ValueError( - "Qwen3.6-35B-A3B checkpoint tensor shape mismatches: " - + preview) - - return { - "checkpoint_path": checkpoint_path, - "text_tensor_count": len(required_text), - "mtp_tensor_count": len(_MTP_KEYS), - "vision_tensor_count": sum( - ".visual." in name for name in weight_map), - "tensor_count": len(weight_map), - "shard_count": len(shards), - } - - -class Qwen36MoeTextFrontendRtx(Nexn2TorchFrontendRtx): - """Qwen3.6-35B-A3B text-only frontend for RTX SM120.""" - - _MODEL_LABEL = "Qwen3.6-35B-A3B text" - _USAGE_DOC = "docs/qwen36_moe_usage.md" - - # The block-scaled 4-bit MMA tier is a build tier, and the prefill now picks - # its MoE tile from what the module actually has: without the tier it uses - # the weight-only grouped GEMV, which is slower on a long prompt and - # otherwise identical. So the tier is a performance requirement, not a - # correctness one, and demanding it here would refuse a build -- a Jetson - # one, whose toolchain has no block-scaled mma -- that runs this correctly. - _REQUIRED_KERNELS = tuple( - name for name in Nexn2TorchFrontendRtx._REQUIRED_KERNELS - if not name.startswith(('moe_blocktile_mma', 'moe_m16_mma')) - ) + ('moe_grouped_w4a16_sm120_bf16',) - - # Same for FA2: prefill and decode both probe the vendored kernel and fall - # back to a reference, so a target that builds FA4 instead still runs. - _REQUIRE_FA2 = False - - def __init__(self, checkpoint_path: str, *, - device: str = "cuda:0", - max_seq: int = 2048, - quant: str = "nvfp4", - kernelized: bool = True, - quant_scope: str = "experts") -> None: - if quant != "nvfp4": - raise NotImplementedError( - f"quant={quant!r} is not implemented; only 'nvfp4' is " - "supported") - if not kernelized: - raise NotImplementedError( - "Qwen3.6-35B-A3B text only supports kernelized=True with " - "runtime NVFP4 conversion") - contract = validate_qwen36_moe_checkpoint(checkpoint_path) - super().__init__( - checkpoint_path, - device=device, - max_seq=max_seq, - quant=quant, - kernelized=kernelized, - quant_scope=quant_scope, - ) - self._checkpoint_contract = contract - - def generate_spec(self, max_new_tokens: int, *, k: int = 2): - """Greedy decode through draft-and-verify with the MTP head. - - Emits exactly what ``generate`` emits: a draft is kept only where the - model's own argmax agrees with it, so this is a speed change and - nothing else. Requires the frontend to have loaded the head. - """ - if self._prompt_ids is None: - raise ValueError("call set_prompt(...) before generate_spec()") - if not self._load_mtp: - raise RuntimeError( - "speculative decoding needs the MTP draft head; construct " - "the frontend with _load_mtp so the loader reads it") - - from flash_rt.frontends.torch._nexn2_rtx_decode import ( - Nexn2DecodeState, - generate_greedy_spec, - ) - - if self._decode_state is None: - self._decode_state = Nexn2DecodeState( - self._weights, self._user_max_seq, self.device) - return generate_greedy_spec( - self._decode_state, self._prompt_ids, max_new_tokens, k, - self._fvk, self.device) - - def generate(self, max_new_tokens: int, *, do_sample: bool = False): - """Generate tokens with the shared qwen3_5_moe CUDA Graph path.""" - if self._prompt_ids is None: - raise ValueError("call set_prompt(...) before generate()") - if do_sample: - raise NotImplementedError( - "the Qwen3.6-35B-A3B kernelized path supports greedy " - "decoding only") - - from flash_rt.frontends.torch._nexn2_rtx_decode import ( - Nexn2DecodeState, - generate_greedy_graph, - ) - - if self._decode_state is None: - self._decode_state = Nexn2DecodeState( - self._weights, self._user_max_seq, self.device) - return generate_greedy_graph( - self._decode_state, - self._prompt_ids, - max_new_tokens, - self._fvk, - self.device, - ) +__all__ = [ + "Qwen36MoeTextFrontend", + "Qwen36MoeTextFrontendRtx", + "validate_qwen36_moe_checkpoint", +] diff --git a/flash_rt/hardware/__init__.py b/flash_rt/hardware/__init__.py index 5d231ead..4048dda5 100644 --- a/flash_rt/hardware/__init__.py +++ b/flash_rt/hardware/__init__.py @@ -165,17 +165,26 @@ def detect_arch() -> str: # ── Nex-N2-mini / Qwen3.6-35B-A3B (qwen3_5_moe) ── # Text LLM, not a VLA: GDN linear-attn + full-attn-every-4th + 256-expert - # NVFP4 MoE. RTX 5090 (SM120) only, and requires the gated kernel build - # (-DFLASHRT_ENABLE_QWEN35MOE=ON). Registered here for discoverability / - # resolve_pipeline_class, but the frontend exposes an LLM surface - # (infer()->logits, generate_greedy) rather than the VLA predict(images) - # API, so these are used via direct frontend construction rather than - # load_model's VLAModel wrapper. + # NVFP4 MoE. Registered here for discoverability / resolve_pipeline_class, + # but the frontend exposes an LLM surface (infer()->logits, + # generate_greedy) rather than the VLA predict(images) API, so these are + # used via direct frontend construction rather than load_model's VLAModel + # wrapper. + # + # Nex-N2 is RTX 5090 (SM120) and needs the full gated kernel build + # (-DFLASHRT_ENABLE_QWEN35MOE=ON). ("nexn2", "torch", "rtx_sm120"): ("flash_rt.frontends.torch.nexn2_rtx", "Nexn2TorchFrontendRtx"), + # Qwen3.6 runs the same frontend on RTX SM120 and on Jetson AGX Thor + # (SM110). The two differ only in which kernel tiers the build has: + # SM120 takes the whole switch, Thor takes the two tiers its toolchain can + # compile. See docs/qwen36_moe_usage.md for the exact command per target. ("qwen36_moe", "torch", "rtx_sm120"): - ("flash_rt.frontends.torch.qwen36_moe_rtx", - "Qwen36MoeTextFrontendRtx"), + ("flash_rt.frontends.torch.qwen36_moe", + "Qwen36MoeTextFrontend"), + ("qwen36_moe", "torch", "thor"): + ("flash_rt.frontends.torch.qwen36_moe", + "Qwen36MoeTextFrontend"), # ── Pi0-FAST ── (SM120 runtime fork inside pipeline, no AttentionBackend protocol.) ("pi0fast", "torch", "thor"): diff --git a/qwen36_moe_edge/expert_quality.py b/qwen36_moe_edge/expert_quality.py index 4dc8862d..987f400f 100644 --- a/qwen36_moe_edge/expert_quality.py +++ b/qwen36_moe_edge/expert_quality.py @@ -160,11 +160,11 @@ def collect_activations( Nexn2DecodeState, generate_greedy, ) - from flash_rt.frontends.torch.qwen36_moe_rtx import ( - Qwen36MoeTextFrontendRtx, + from flash_rt.frontends.torch.qwen36_moe import ( + Qwen36MoeTextFrontend, ) - frontend = Qwen36MoeTextFrontendRtx( + frontend = Qwen36MoeTextFrontend( checkpoint, device=device, max_seq=max_seq, quant_scope="experts") input_ids = frontend.tokenizer( prompt, return_tensors="pt", add_special_tokens=False, diff --git a/qwen36_moe_edge/route_trace.py b/qwen36_moe_edge/route_trace.py index 9d2658d2..7731d05d 100644 --- a/qwen36_moe_edge/route_trace.py +++ b/qwen36_moe_edge/route_trace.py @@ -43,8 +43,8 @@ Nexn2DecodeState, generate_greedy, ) -from flash_rt.frontends.torch.qwen36_moe_rtx import ( - Qwen36MoeTextFrontendRtx, +from flash_rt.frontends.torch.qwen36_moe import ( + Qwen36MoeTextFrontend, ) @@ -321,7 +321,7 @@ def main() -> None: parser.add_argument("--device", default="cuda:0") args = parser.parse_args() - frontend = Qwen36MoeTextFrontendRtx( + frontend = Qwen36MoeTextFrontend( args.checkpoint, device=args.device, max_seq=args.max_seq, diff --git a/qwen36_moe_edge/streaming_frontend.py b/qwen36_moe_edge/streaming_frontend.py index 5e025227..f7dc4e0b 100644 --- a/qwen36_moe_edge/streaming_frontend.py +++ b/qwen36_moe_edge/streaming_frontend.py @@ -18,12 +18,12 @@ from pathlib import Path -from flash_rt.frontends.torch.qwen36_moe_rtx import Qwen36MoeTextFrontendRtx +from flash_rt.frontends.torch.qwen36_moe import Qwen36MoeTextFrontend from qwen36_moe_edge.expert_cache import CacheConfig, ExpertCache -class Qwen36MoeStreamingFrontend(Qwen36MoeTextFrontendRtx): +class Qwen36MoeStreamingFrontend(Qwen36MoeTextFrontend): """Routed experts streamed from a bundle rather than held in memory.""" _MODEL_LABEL = "Qwen3.6-35B-A3B text, streamed experts" @@ -35,7 +35,7 @@ class Qwen36MoeStreamingFrontend(Qwen36MoeTextFrontendRtx): # which is what happened on the first attempt, on a target where the tier is # correctly not built at all. _REQUIRED_KERNELS = tuple( - name for name in Qwen36MoeTextFrontendRtx._REQUIRED_KERNELS + name for name in Qwen36MoeTextFrontend._REQUIRED_KERNELS if not name.startswith(('moe_blocktile_mma', 'moe_m16_mma')) ) + ('qwen35moe_e0m3_dequant_bf16', 'bf16_matvec_sm120_bf16') diff --git a/qwen36_moe_edge/warm_start_validation.py b/qwen36_moe_edge/warm_start_validation.py index 8ed82c9c..2d04826d 100644 --- a/qwen36_moe_edge/warm_start_validation.py +++ b/qwen36_moe_edge/warm_start_validation.py @@ -43,11 +43,11 @@ def collect_traces( Nexn2DecodeState, generate_greedy, ) - from flash_rt.frontends.torch.qwen36_moe_rtx import ( - Qwen36MoeTextFrontendRtx, + from flash_rt.frontends.torch.qwen36_moe import ( + Qwen36MoeTextFrontend, ) - frontend = Qwen36MoeTextFrontendRtx( + frontend = Qwen36MoeTextFrontend( checkpoint, device=device, max_seq=max_seq, quant_scope="experts") state = Nexn2DecodeState(frontend._weights, max_seq, device) state.batched_prefill = False diff --git a/tests/test_qwen36_moe_smoke.py b/tests/test_qwen36_moe_smoke.py index 191dc6c7..87c33df9 100644 --- a/tests/test_qwen36_moe_smoke.py +++ b/tests/test_qwen36_moe_smoke.py @@ -56,7 +56,7 @@ def _config(): def _checkpoint(tmp_path): - from flash_rt.frontends.torch.qwen36_moe_rtx import ( + from flash_rt.frontends.torch.qwen36_moe import ( _MTP_KEYS, _required_text_keys, ) @@ -76,14 +76,14 @@ def _checkpoint(tmp_path): def _mock_checkpoint_shapes(monkeypatch, overrides=None): - from flash_rt.frontends.torch import qwen36_moe_rtx + from flash_rt.frontends.torch import qwen36_moe - shapes = qwen36_moe_rtx._expected_text_shapes( - qwen36_moe_rtx._EXPECTED_LAYER_TYPES) - shapes.update(qwen36_moe_rtx._expected_mtp_shapes()) + shapes = qwen36_moe._expected_text_shapes( + qwen36_moe._EXPECTED_LAYER_TYPES) + shapes.update(qwen36_moe._expected_mtp_shapes()) shapes.update(overrides or {}) monkeypatch.setattr( - qwen36_moe_rtx, + qwen36_moe, "_read_tensor_shapes", lambda checkpoint_path, weight_map, tensor_names: { name: shapes[name] for name in tensor_names @@ -93,26 +93,50 @@ def _mock_checkpoint_shapes(monkeypatch, overrides=None): def test_frontend_is_a_thin_qwen_entry(): from flash_rt.frontends.torch.nexn2_rtx import Nexn2TorchFrontendRtx - from flash_rt.frontends.torch.qwen36_moe_rtx import ( - Qwen36MoeTextFrontendRtx, + from flash_rt.frontends.torch.qwen36_moe import ( + Qwen36MoeTextFrontend, ) - assert issubclass(Qwen36MoeTextFrontendRtx, Nexn2TorchFrontendRtx) - assert Qwen36MoeTextFrontendRtx._MODEL_LABEL == ( + assert issubclass(Qwen36MoeTextFrontend, Nexn2TorchFrontendRtx) + assert Qwen36MoeTextFrontend._MODEL_LABEL == ( "Qwen3.6-35B-A3B text") assert inspect.signature( - Qwen36MoeTextFrontendRtx).parameters["kernelized"].default is True + Qwen36MoeTextFrontend).parameters["kernelized"].default is True -def test_registry_resolves_qwen36_moe(): +@pytest.mark.parametrize("arch", ["rtx_sm120", "thor"]) +def test_registry_resolves_qwen36_moe(arch): from flash_rt.hardware import _PIPELINE_MAP, resolve_pipeline_class - assert _PIPELINE_MAP[("qwen36_moe", "torch", "rtx_sm120")] == ( - "flash_rt.frontends.torch.qwen36_moe_rtx", - "Qwen36MoeTextFrontendRtx", + assert _PIPELINE_MAP[("qwen36_moe", "torch", arch)] == ( + "flash_rt.frontends.torch.qwen36_moe", + "Qwen36MoeTextFrontend", ) - cls = resolve_pipeline_class("qwen36_moe", "torch", "rtx_sm120") - assert cls.__name__ == "Qwen36MoeTextFrontendRtx" + cls = resolve_pipeline_class("qwen36_moe", "torch", arch) + assert cls.__name__ == "Qwen36MoeTextFrontend" + + +def test_both_architectures_resolve_to_one_frontend(): + """Thor and RTX run the same code, so they must resolve to one class. + + Two entries pointing at two classes would be two paths to keep in step; + what differs between the targets is which kernel tiers the build has, not + which Python runs. + """ + from flash_rt.hardware import resolve_pipeline_class + + assert (resolve_pipeline_class("qwen36_moe", "torch", "thor") + is resolve_pipeline_class("qwen36_moe", "torch", "rtx_sm120")) + + +def test_previous_import_path_still_works(): + """The module was renamed when it stopped being RTX-only.""" + from flash_rt.frontends.torch.qwen36_moe import Qwen36MoeTextFrontend + from flash_rt.frontends.torch.qwen36_moe_rtx import ( + Qwen36MoeTextFrontendRtx, + ) + + assert Qwen36MoeTextFrontendRtx is Qwen36MoeTextFrontend def test_load_model_redirects_to_text_frontend(): @@ -121,30 +145,30 @@ def test_load_model_redirects_to_text_frontend(): with pytest.raises(NotImplementedError) as exc: flash_rt.load_model("/nonexistent", config="qwen36_moe") message = str(exc.value) - assert "Qwen36MoeTextFrontendRtx" in message + assert "Qwen36MoeTextFrontend" in message assert "text LLM" in message def test_constructor_rejects_quant_before_checkpoint_access(): - from flash_rt.frontends.torch.qwen36_moe_rtx import ( - Qwen36MoeTextFrontendRtx, + from flash_rt.frontends.torch.qwen36_moe import ( + Qwen36MoeTextFrontend, ) with pytest.raises(NotImplementedError, match="only 'nvfp4'"): - Qwen36MoeTextFrontendRtx("/nonexistent", quant="fp8") + Qwen36MoeTextFrontend("/nonexistent", quant="fp8") def test_constructor_rejects_reference_path_before_checkpoint_access(): - from flash_rt.frontends.torch.qwen36_moe_rtx import ( - Qwen36MoeTextFrontendRtx, + from flash_rt.frontends.torch.qwen36_moe import ( + Qwen36MoeTextFrontend, ) with pytest.raises(NotImplementedError, match="kernelized=True"): - Qwen36MoeTextFrontendRtx("/nonexistent", kernelized=False) + Qwen36MoeTextFrontend("/nonexistent", kernelized=False) def test_checkpoint_contract_accepts_complete_layout(tmp_path, monkeypatch): - from flash_rt.frontends.torch.qwen36_moe_rtx import ( + from flash_rt.frontends.torch.qwen36_moe import ( validate_qwen36_moe_checkpoint, ) @@ -178,7 +202,7 @@ def test_checkpoint_contract_accepts_complete_layout(tmp_path, monkeypatch): ) def test_checkpoint_contract_rejects_wrong_geometry( tmp_path, path, invalid, message): - from flash_rt.frontends.torch.qwen36_moe_rtx import ( + from flash_rt.frontends.torch.qwen36_moe import ( validate_qwen36_moe_checkpoint, ) @@ -196,7 +220,7 @@ def test_checkpoint_contract_rejects_wrong_geometry( def test_checkpoint_contract_rejects_missing_text_tensor(tmp_path): - from flash_rt.frontends.torch.qwen36_moe_rtx import ( + from flash_rt.frontends.torch.qwen36_moe import ( validate_qwen36_moe_checkpoint, ) @@ -211,7 +235,7 @@ def test_checkpoint_contract_rejects_missing_text_tensor(tmp_path): def test_checkpoint_contract_rejects_partial_mtp(tmp_path): - from flash_rt.frontends.torch.qwen36_moe_rtx import ( + from flash_rt.frontends.torch.qwen36_moe import ( validate_qwen36_moe_checkpoint, ) @@ -226,7 +250,7 @@ def test_checkpoint_contract_rejects_partial_mtp(tmp_path): def test_checkpoint_contract_rejects_missing_shard(tmp_path): - from flash_rt.frontends.torch.qwen36_moe_rtx import ( + from flash_rt.frontends.torch.qwen36_moe import ( validate_qwen36_moe_checkpoint, ) @@ -239,7 +263,7 @@ def test_checkpoint_contract_rejects_missing_shard(tmp_path): def test_checkpoint_contract_rejects_wrong_tensor_shape( tmp_path, monkeypatch): - from flash_rt.frontends.torch.qwen36_moe_rtx import ( + from flash_rt.frontends.torch.qwen36_moe import ( validate_qwen36_moe_checkpoint, ) @@ -261,15 +285,17 @@ def test_generic_env_names_precede_legacy_aliases(monkeypatch): def test_kernelized_generate_uses_shared_graph_path(monkeypatch): from flash_rt.frontends.torch import _nexn2_rtx_decode as decode - from flash_rt.frontends.torch.qwen36_moe_rtx import ( - Qwen36MoeTextFrontendRtx, + from flash_rt.frontends.torch.qwen36_moe import ( + Qwen36MoeTextFrontend, ) calls = {} class FakeState: - def __init__(self, weights, max_seq, device): + def __init__(self, weights, max_seq, device, *, + spec_graph_cache_max=None): calls["state"] = (weights, max_seq, device) + calls["spec_cap"] = spec_graph_cache_max def fake_generate(state, prompt_ids, count, fvk, device): calls["generate"] = (state, prompt_ids, count, fvk, device) @@ -278,8 +304,8 @@ def fake_generate(state, prompt_ids, count, fvk, device): monkeypatch.setattr(decode, "Nexn2DecodeState", FakeState) monkeypatch.setattr(decode, "generate_greedy_graph", fake_generate) - frontend = Qwen36MoeTextFrontendRtx.__new__( - Qwen36MoeTextFrontendRtx) + frontend = Qwen36MoeTextFrontend.__new__( + Qwen36MoeTextFrontend) frontend._kernelized = True frontend._prompt_ids = object() frontend._decode_state = None @@ -287,6 +313,7 @@ def fake_generate(state, prompt_ids, count, fvk, device): frontend._user_max_seq = 128 frontend.device = "cuda:0" frontend._fvk = object() + frontend._spec_graph_cache_max = None assert frontend.generate(3) == [7, 7, 7] assert calls["state"] == ( @@ -302,7 +329,7 @@ def fake_generate(state, prompt_ids, count, fvk, device): reason="set FLASHRT_QWEN36_MOE_CKPT_DIR for checkpoint validation", ) def test_real_checkpoint_contract(): - from flash_rt.frontends.torch.qwen36_moe_rtx import ( + from flash_rt.frontends.torch.qwen36_moe import ( validate_qwen36_moe_checkpoint, ) From 36922e329b656b5822da78dedbb9b9c9ec204fd2 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 6 Aug 2026 15:24:30 -0400 Subject: [PATCH 80/85] Make the docs say what the code does The model guide said the interface does not execute the MTP head and then documented generate_spec(); it said FA2 is deliberately absent on Thor while the build enables and tests it there; the Thor build line was "-DGPU_ARCH=110 ..." rather than a command; and the usage examples named the module and class as they were before the rename. Each of those now states one policy: MTP loads on load_mtp=True and nothing else reads it, FA2 on sm_110 is opt-in with the flag that builds it and the model runs either way, and both targets have an exact configure line. The README's two claims that FA2 is RTX-only are updated to match. The speculative table gains the paired long-context measurements and corrects the K=1 ratio, which was quoted from a different run than its own baseline. --- README.md | 14 ++-- docs/qwen36_moe_usage.md | 168 ++++++++++++++++++++++++++++++--------- 2 files changed, 140 insertions(+), 42 deletions(-) diff --git a/README.md b/README.md index 3ef1f2b0..c9f89397 100644 --- a/README.md +++ b/README.md @@ -751,7 +751,7 @@ extension modules: | Artifact | Size | What it contains | |---|---|---| | `flash_rt/flash_rt_kernels.so` | ~3 MB | Hand-written memory-bound kernels (norm, activation, fusion, FP8 quant, cuBLASLt wrappers, Thor FMHA). **Always built.** | -| `flash_rt/flash_rt_fa2.so` | ~135 MB | Vendored Flash-Attention 2 v2.7.4.post1 fwd (fp16 + bf16, SM80/86/89/120). **Built only on RTX targets** — Thor skips it and uses `fvk.attention_qkv_fp16` (cuBLAS-decomposed) for attention instead. | +| `flash_rt/flash_rt_fa2.so` | ~135 MB | Vendored Flash-Attention 2 v2.7.4.post1 fwd (fp16 + bf16, SM80/86/89/120). **Built automatically on RTX targets.** Thor skips it by default and uses `fvk.attention_qkv_fp16` (cuBLAS-decomposed) instead; `-DFLASHRT_ENABLE_THOR_FA2=ON` builds it there for the one model whose long prefill needs it (Qwen3.6, bf16 head_dim 256 — a single instantiation). | **Crucially — no `pip install flash-attn` required.** The FA2 kernel is vendored at source level and built into `flash_rt_fa2.so` during @@ -881,7 +881,7 @@ CMake reads `nvidia-smi --query-gpu=compute_cap` to pick the target arch. Override for cross-compilation or when auto-detect fails: ```bash -cmake -B build -S . -DGPU_ARCH=110 # Jetson AGX Thor (FA2 skipped, CUTLASS SM100 path ON) +cmake -B build -S . -DGPU_ARCH=110 # Jetson AGX Thor (FA2 opt-in, CUTLASS SM100 path ON) cmake -B build -S . -DGPU_ARCH=121 # DGX Spark / GB10 (FA2 sm_121 AOT, NVFP4 ON) cmake -B build -S . -DGPU_ARCH=120 # RTX 5090 (FA2 sm_120 AOT, NVFP4 ON) cmake -B build -S . -DGPU_ARCH=89 # RTX 4090 (FA2 sm_80 AOT natively runs on Ada) @@ -889,10 +889,12 @@ cmake -B build -S . -DGPU_ARCH=86 # RTX 3090 / A10 (FA2 sm_80 AOT) cmake -B build -S . -DGPU_ARCH=80 # A100 (FA2 sm_80 AOT) ``` -FA2 is enabled by CMake when `GPU_ARCH ∈ {80, 86, 89, 120, 121}`. Other -arches (notably Thor SM110 and SM90 Hopper) route attention through -the cuBLAS-decomposed `fvk.attention_qkv_fp16` path instead of FA2 — -`flash_rt_fa2.so` simply isn't built, and no runtime error results. +FA2 is enabled by CMake when `GPU_ARCH ∈ {80, 86, 89, 120, 121}`, and on +Thor SM110 when `-DFLASHRT_ENABLE_THOR_FA2=ON` is passed. Other arches +(notably SM90 Hopper, and Thor without that flag) route attention +through the cuBLAS-decomposed `fvk.attention_qkv_fp16` path instead of +FA2 — `flash_rt_fa2.so` simply isn't built, and no runtime error +results. ### Build timing (one-time) diff --git a/docs/qwen36_moe_usage.md b/docs/qwen36_moe_usage.md index b2afcb24..0056cd6f 100644 --- a/docs/qwen36_moe_usage.md +++ b/docs/qwen36_moe_usage.md @@ -1,14 +1,20 @@ # Qwen3.6-35B-A3B text inference FlashRT runs the language backbone from the official -`Qwen/Qwen3.6-35B-A3B` BF16 checkpoint on an RTX 5090. The checkpoint uses the -same `qwen3_5_moe` text architecture as Nex-N2-mini, so both models share the -same weight loader, prefill, attention, MoE, recurrent-state, and CUDA Graph -decode implementation. +`Qwen/Qwen3.6-35B-A3B` BF16 checkpoint on an RTX 5090 (SM120) and on Jetson AGX +Thor (SM110). The checkpoint uses the same `qwen3_5_moe` text architecture as +Nex-N2-mini, so both models share the same weight loader, prefill, attention, +MoE, recurrent-state, and CUDA Graph decode implementation, and both +architectures run the same frontend -- what differs is which kernel tiers the +build has. -This entry is text-only. It does not load the vision tower and it validates but -does not execute the checkpoint's MTP head. Image/video input and speculative -decode are not part of this interface. +This entry is text-only: it does not load the vision tower, and image or video +input is not part of this interface. + +It does execute the checkpoint's MTP draft head, but only when asked. Pass +`load_mtp=True` to the constructor and call `generate_spec()`; see *Speculative +decode* below. `generate()` never reads the head, and without `load_mtp=True` +the loader does not read it either. ## Requirements @@ -19,9 +25,13 @@ decode are not part of this interface. | GPU memory | 32 GB (SM120); unified memory on Thor | | Framework | PyTorch | | Runtime quantization | NVFP4 | -| Build flags | `-DGPU_ARCH=120 -DFLASHRT_ENABLE_QWEN35MOE=ON` (SM120), `-DGPU_ARCH=110 ...` (Thor) | -Configure and build the gated `qwen3_5_moe` kernels: +Configure and build the gated `qwen3_5_moe` kernels. **The two targets take +different flags** -- `FLASHRT_ENABLE_QWEN35MOE` turns on the block-scaled 4-bit +MMA tier as well, which sm_110 refuses at configure time, so Thor names the two +tiers it can compile: + +RTX 5090 (SM120): ```bash cmake -S . -B build \ @@ -31,6 +41,23 @@ cmake --build build -j pip install -e ".[torch]" ``` +Jetson AGX Thor (SM110): + +```bash +cmake -S . -B build \ + -DGPU_ARCH=110 \ + -DFLASHRT_ENABLE_QWEN35MOE_CORE=ON \ + -DFLASHRT_ENABLE_QWEN35MOE_W4A16=ON \ + -DFLASHRT_ENABLE_THOR_FA2=ON +cmake --build build -j 2 +pip install -e ".[torch]" +``` + +`FLASHRT_ENABLE_THOR_FA2` is what builds the vendored FA2 kernels on sm_110; +see *Attention differs by target* below for what it buys and why it is opt-in. +Without it the model still runs, and produces the same tokens, but a long +prefill is far slower. + ### Kernel tiers `FLASHRT_ENABLE_QWEN35MOE=ON` is a convenience switch for all three tiers @@ -42,8 +69,9 @@ below. Targets that cannot run a tier can select the remainder explicitly. | `FLASHRT_ENABLE_QWEN35MOE_W4A16` | weight-only 4-bit matvec, grouped matvec, GEMM | SM80 and newer; hardware operand conversion from SM89 | | `FLASHRT_ENABLE_QWEN35MOE_W4A4` | block-scaled 4-bit MMA: grouped GEMV, M16/M64/block-tile MMA | sm_120a / sm_121a | -The upper tiers depend on the core tier, so enabling either turns it on. The -SM120 text runtime documented here needs all three. +The upper tiers depend on the core tier, so enabling either turns it on. SM120 +runs all three; sm_110 runs the first two, and the block-scaled tier refuses to +configure there rather than building kernels that fail at run time. **These three tiers are not the whole dependency set.** Walking every `fvk` call the pipeline makes and resolving each to the preprocessor guard active @@ -78,25 +106,47 @@ kernel it replaces when a build does not carry it: |---|---|---|---| | `gated_deltanet_recurrent_edge_qwen36_bf16` | `FLASHRT_HAVE_QWEN36_KERNELS` | `gated_deltanet_recurrent_qwen36_bf16` | same arithmetic without the local-memory round trip for the state column | | `moe_router_topk_warp_sm120_bf16` | `FLASHRT_HAVE_QWEN35MOE_CORE` | `moe_router_topk_sm120_bf16` | same selection in one warp instead of `k` rounds of block-wide barriers | +| `moe_shared_gate_combine_edge_bf16` | `FLASHRT_HAVE_QWEN35MOE_CORE` | a five-launch tensor chain | one kernel, same arithmetic in the same order, rounded once at the store | +| `moe_grouped_gemm_nvfp4_sm100_bf16out` | `FLASHRT_HAVE_QWEN35MOE_GROUPED_SM100` (sm_110 + `_W4A16`) | the per-expert GEMM loop, or the grouped GEMV | every routed expert of a layer in one launch, with the per-group shapes read from device memory | -Both produce output identical to the kernel they stand in for, so the fallback -is a performance difference and never a numerical one. The edge recurrence is +All produce output identical to the path they stand in for, so the fallback is +a performance difference and never a numerical one. The edge recurrence is shape-specialized to a head dim of 128 and raises for anything else rather than leaving the output buffer undefined. -### Attention differs by target, by design +Which of each pair actually runs is a `KernelPolicy` field, not the answer to +"is the symbol there" — see *Runtime controls*. + +### Attention differs by target -The ten full-attention layers do not use the same kernel everywhere, and the -arch lists reflect that rather than overlooking it: +The ten full-attention layers do not use the same kernel everywhere: -| target | attention | why | +| target | attention | how it is built | |---|---|---| -| SM120 / SM89 / SM87 | vendored FA2 | the SM80-family source, which `__CUDA_ARCH__ >= 800` admits | +| SM120 / SM89 / SM87 | vendored FA2 | automatic: the SM80-family source, which `__CUDA_ARCH__ >= 800` admits | +| Thor SM110 | vendored FA2, or the decomposed reference | opt-in: `-DFLASHRT_ENABLE_THOR_FA2=ON` | | Thor SM110 | FA4 | its SM100-class CuTe-DSL kernel needs Blackwell tensor memory; ships as the `thor-fa4` pip extra, not compiled into `flash_rt_kernels` | -So FA2 is deliberately absent from the Thor build, and FA4 cannot serve -Ampere-class SM87. Treat a missing FA2 as a signal to fall back, not as a -build error. +FA2 was originally excluded from sm_110 on the grounds that Thor has its own +attention path and FA2 would add about 10 MB of `.so` for nothing. That holds +for the models that use the decomposed path, and it does not hold for a long +prefill of this one: the decomposed path materialises an `(S * heads, S_kv)` +score buffer -- 3.4 GB per layer at ten thousand tokens -- and this model needs +one instantiation (bf16, head_dim 256), not the twelve the size estimate +assumed. Enabling it takes the chunking penalty of a chunked prefill from 64% +to 4%. + +It stays opt-in because every other Thor model still uses its own attention +path and would only be paying the compile time and the binary size. A build +without it runs this model correctly and emits the same tokens; only a long +prefill is slower. Treat a missing FA2 as a signal to fall back, not as a build +error. + +At the *decode* shape the two are the same answer to bf16 precision -- measured +against an fp32 reference, 2.0e-3 relative for both at kv=64, 2.2e-3 against +2.1e-3 at kv=2048 -- so decode on sm_110 keeps the reference path the golden +fixture was recorded through, and takes FA2 only for prefill. +`FLASHRT_NEXN2_DECODE_FA2=1` overrides that. The attention backend probes its kernel at construction: it runs one case through the same launch the hot path uses and compares against @@ -118,11 +168,9 @@ fail at run time. The explicit gate turns that into a configure-time error. ## Usage ```python -from flash_rt.frontends.torch.qwen36_moe_rtx import ( - Qwen36MoeTextFrontendRtx, -) +from flash_rt.frontends.torch.qwen36_moe import Qwen36MoeTextFrontend -frontend = Qwen36MoeTextFrontendRtx( +frontend = Qwen36MoeTextFrontend( "/models/Qwen3.6-35B-A3B", device="cuda:0", max_seq=4096, @@ -173,7 +221,7 @@ validation can be run without loading model weights: ```bash PYTHONPATH=. python - <<'PY' -from flash_rt.frontends.torch.qwen36_moe_rtx import ( +from flash_rt.frontends.torch.qwen36_moe import ( validate_qwen36_moe_checkpoint, ) print(validate_qwen36_moe_checkpoint("/models/Qwen3.6-35B-A3B")) @@ -188,10 +236,24 @@ The shared architecture uses: `8192`; `0` disables chunking. - `FLASHRT_QWEN35MOE_GRAPH_CACHE_MAX` — decode CUDA Graph LRU capacity, default `256`. +- `FLASHRT_QWEN35MOE_SPEC_GRAPH_CACHE_MAX` — speculative-window CUDA Graph LRU + capacity, default `16`. Bounded separately and far lower than the decode + cache: a speculative graph covers `k+1` positions through the whole stack, so + its memory pool is several times a decode step's. The constructor argument + `spec_graph_cache_max` sets it per frontend. The older `FLASHRT_NEXN2_PREFILL_CHUNK` and `FLASHRT_NEXN2_GRAPH_CACHE_MAX` names remain compatible aliases. +Which of several interchangeable kernels each step calls is a +`KernelPolicy` (`flash_rt.frontends.torch._nexn2_rtx_forward`), not a symbol +lookup: every field selects between implementations checked against each other +with `torch.equal`, so a field decides speed and never output. The environment +variables above and `NEXN2_WY_GDN`, `NEXN2_ROUTE_KERNEL`, +`NEXN2_DENSE_CUBLASLT`, `FLASHRT_QWEN35MOE_W4A16_EDGE` and +`FLASHRT_QWEN35MOE_VERIFY_K_ROWS` are its defaults. A policy must not be +changed between a CUDA graph capture and its replay. + ## Validation The repository smoke test is checkpoint-independent: @@ -313,23 +375,57 @@ to be plain greedy's, and it is checked directly: logits rows, per-token recurrent and conv snapshots, and the KV rows written are all compared with `torch.equal` against a decode step run over the same tokens. -Measured on Jetson AGX Thor, 20-token prompt, 128 generated tokens, one process -per point, best of five: +Both the plain and the speculative rate below come from the same process, so +the ratio is a paired comparison. The absolute figures move a few percent with +what else the board is running; the ratio is the stable part. + +20-token prompt, 128 generated tokens: | | tok/s | vs plain | |---|---:|---:| | plain greedy | 100.35 | | -| speculative, K=1 | 105.22 | 1.09x | +| speculative, K=1 | 105.22 | 1.05x | | speculative, K=2 | **106.74** | 1.06x | -`K=2` is the operating point. Above it the window costs more than the extra -accepted tokens return: each additional verified row re-reads the routed -experts, which do not amortise across a window the way the dense weights do, -and each additional draft pays a full-vocabulary projection. +The same comparison at longer context, 128 generated tokens per point: + +| context | plain | K=1 | K=2 | K=3 | +|---:|---:|---:|---:|---:| +| 512 | 90.4 | 95.4 (1.06x) | **99.4 (1.10x)** | 96.3 (1.06x) | +| 2048 | 79.9 | 84.1 (1.05x) | **86.5 (1.08x)** | 73.8 (0.92x) | + +Every row emitted the same text as plain greedy decoding in the same process. + +`K=2` is the operating point at every context measured. Above it the window +costs more than the extra accepted tokens return: each additional verified row +re-reads the routed experts, which do not amortise across a window the way the +dense weights do, and each additional draft pays a full-vocabulary projection. +At 2048 tokens `K=3` is already a loss, and its acceptance falls too -- 3.46 +kept per window at 512, 3.00 at 2048. + +Acceptance rises with context (2.60 kept per window at a 20-token prompt, 2.72 +at 512, 2.74 at 2048) while the ratio does not, because the verify runs `K+1` +separate single-query attention passes. That is the price of keeping the window +bit-exact, and it is the largest remaining lever on this path. + +Enable it with `load_mtp=True` on the constructor; the window width is the `k` +argument to `generate_spec`: + +```python +frontend = Qwen36MoeTextFrontend( + "/models/Qwen3.6-35B-A3B", + device="cuda:0", + max_seq=2048, + load_mtp=True, + spec_graph_cache_max=16, +) +frontend.set_prompt("Explain why deterministic reductions matter.") +token_ids = frontend.generate_spec(max_new_tokens=128, k=2) +print(frontend.tokenizer.decode(token_ids)) +``` -Enable it with `_load_mtp = True` on the frontend subclass; the window width is -the `k` argument to `generate_spec`. `FLASHRT_QWEN35MOE_VERIFY_K_ROWS=0` falls -back to verifying through the prefill forward. +`FLASHRT_QWEN35MOE_VERIFY_K_ROWS=0` falls back to verifying through the prefill +forward, which is slower and produces the same tokens. ## Limitations From 50561bb3eddb5bf896ba0d3a656fb6ea57276cf8 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 6 Aug 2026 15:28:42 -0400 Subject: [PATCH 81/85] Check the speculative contract and the tier gates in the repository Speculative decode had no committed test for its runtime contract. It has one now: the constructor arguments and their rejections, the graph cache's eviction policy, and -- against a checkpoint, skipped without one -- K=1 and K=2 equality with plain greedy, the window's logits and recurrent, conv and KV state against the decode steps they stand in for, the rejected-tail rewind, and boundary token counts. They are equality tests because the claim is equality; a tolerance would not be checking it. Both graph caches now insert through one function, which is also what makes the eviction policy testable without capturing a graph. The tier gates get a matrix. A build proves its own configuration works; only reading the gates can say what the four configurations nobody built contain, so the script walks CMake's tier blocks and the bindings' guards and fails if a tier source or symbol is reachable with every tier off. The five configure lines it prints are in the model guide with what each produces, including the one that is supposed to fail. --- docs/qwen36_moe_usage.md | 41 +++ flash_rt/frontends/torch/_nexn2_rtx_decode.py | 29 +- scripts/qwen35moe_build_matrix.py | 214 ++++++++++++ tests/test_qwen35moe_build_matrix.py | 104 ++++++ tests/test_qwen36_moe_spec_decode.py | 305 ++++++++++++++++++ 5 files changed, 683 insertions(+), 10 deletions(-) create mode 100644 scripts/qwen35moe_build_matrix.py create mode 100644 tests/test_qwen35moe_build_matrix.py create mode 100644 tests/test_qwen36_moe_spec_decode.py diff --git a/docs/qwen36_moe_usage.md b/docs/qwen36_moe_usage.md index 0056cd6f..921796c6 100644 --- a/docs/qwen36_moe_usage.md +++ b/docs/qwen36_moe_usage.md @@ -256,6 +256,35 @@ changed between a CUDA graph capture and its replay. ## Validation +### Build and symbol matrix + +A build with every `qwen3_5_moe` option off must compile the same sources and +export the same symbols it did before the tiers existed. That is a property of +the gates, so it is checked by reading them: + +```bash +python scripts/qwen35moe_build_matrix.py # print sources + symbols per tier +python scripts/qwen35moe_build_matrix.py --check # exit 1 if any tier leaks +PYTHONPATH=. PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 \ + pytest -q -p no:cacheprovider tests/test_qwen35moe_build_matrix.py +``` + +The five configurations behind it, each configure-only: + +| configuration | flags | result | +|---|---|---| +| baseline SM120 | `-DGPU_ARCH=120` | `FA2 ENABLED`; no `qwen3_5_moe` source or symbol | +| baseline SM110 | `-DGPU_ARCH=110` | `FA2 DISABLED`; no `qwen3_5_moe` source or symbol | +| SM110 supported | `-DGPU_ARCH=110 -DFLASHRT_ENABLE_QWEN35MOE_CORE=ON -DFLASHRT_ENABLE_QWEN35MOE_W4A16=ON -DFLASHRT_ENABLE_THOR_FA2=ON` | core + weight-only tiers, grouped MoE GEMM, FA2 at `hdim={256} x dtype={bf16}` | +| SM120 supported | `-DGPU_ARCH=120 -DFLASHRT_ENABLE_QWEN35MOE=ON` | all three tiers | +| SM110 block-scaled | `-DGPU_ARCH=110 -DFLASHRT_ENABLE_QWEN35MOE_W4A4=ON` | **configure fails**, naming the two tiers that do apply | + +The last row is the point of the explicit gate: CUTLASS would otherwise compile +those translation units on sm_110 with the MMA replaced by an invalid control +path, and the failure would arrive at run time instead. + +### Tests + The repository smoke test is checkpoint-independent: ```bash @@ -263,6 +292,18 @@ PYTHONPATH=. PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 \ pytest -q -p no:cacheprovider tests/test_qwen36_moe_smoke.py ``` +Speculative decode has its own file. The constructor contract and the graph +cache's eviction policy run anywhere; the equivalence tests -- K=1 and K=2 +against plain greedy, the window's logits and recurrent, conv and KV state +against the decode steps they stand in for, the rejected-tail rewind, and the +boundary token counts -- need a GPU and a checkpoint and skip without them: + +```bash +FLASHRT_QWEN36_MOE_CKPT_DIR=/models/Qwen3.6-35B-A3B \ +PYTHONPATH=. PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 \ +pytest -q -p no:cacheprovider tests/test_qwen36_moe_spec_decode.py +``` + Set `FLASHRT_QWEN36_MOE_CKPT_DIR` to include the official checkpoint contract test. Performance and precision numbers must be measured on Qwen3.6 weights; Nex-N2-mini measurements are not interchangeable even though the compute diff --git a/flash_rt/frontends/torch/_nexn2_rtx_decode.py b/flash_rt/frontends/torch/_nexn2_rtx_decode.py index c1517c0f..261a1afa 100644 --- a/flash_rt/frontends/torch/_nexn2_rtx_decode.py +++ b/flash_rt/frontends/torch/_nexn2_rtx_decode.py @@ -48,6 +48,22 @@ def _qwen35moe_env(name: str, default: str) -> str: _BATCHED_PREFILL_MIN_S = 8 +def _cache_put(cache, key, value, cap): + """Insert ``value`` and evict least-recently-used down to ``cap``. + + Both graph caches go through this. A captured graph owns its memory pool, + so an unbounded cache leaks device memory across a long generation -- one + graph per absolute position. ``cap <= 0`` disables the bound. + + Returns the value, so a caller can insert and use in one expression. + """ + cache[key] = value + if cap > 0: + while len(cache) > cap: + cache.popitem(last=False) # evict LRU + return value + + def _cs(): """Current CUDA stream handle. Inside torch.cuda.graph capture this is the capture stream; eager, the default stream. fvk calls MUST use it -- @@ -1356,11 +1372,8 @@ def _restore(): with torch.no_grad(): _restore() - state._spec_graphs[key] = (g, hid) - cap = state.spec_graph_cache_max - if cap > 0 and len(state._spec_graphs) > cap: - state._spec_graphs.popitem(last=False) # evict LRU - return state._spec_graphs[key] + return _cache_put(state._spec_graphs, key, (g, hid), + state.spec_graph_cache_max) def spec_decode_step(state, token_id, pos, k, fvk, device): @@ -1558,11 +1571,7 @@ def _restore(): with torch.no_grad(): _restore() - state._graphs[pos] = (g, out) - cap = state.graph_cache_max - if cap > 0 and len(state._graphs) > cap: - state._graphs.popitem(last=False) # evict LRU - return state._graphs[pos] + return _cache_put(state._graphs, pos, (g, out), state.graph_cache_max) def generate_greedy_graph(state, input_ids, max_new_tokens, fvk, device): diff --git a/scripts/qwen35moe_build_matrix.py b/scripts/qwen35moe_build_matrix.py new file mode 100644 index 00000000..0143f236 --- /dev/null +++ b/scripts/qwen35moe_build_matrix.py @@ -0,0 +1,214 @@ +#!/usr/bin/env python3 +"""What each qwen3_5_moe build tier adds, and what a build without them has. + +The claim these tiers make is that a build with every qwen3_5_moe option off +compiles the same sources and exports the same symbols it did before they +existed. That is a property of the gates, so it is checked by reading them: +which translation units CMake adds under each tier, and which ``m.def`` names +sit inside the matching preprocessor guard in the bindings. + +Reading the gates rather than building has a specific limit and a specific +advantage. It cannot catch a kernel that fails to compile -- only a build does +that, and the configure matrix printed at the end is how to run one. It can +catch the thing a single build cannot: a source or a symbol that leaks into a +configuration nobody built. + + python scripts/qwen35moe_build_matrix.py # print the matrix + python scripts/qwen35moe_build_matrix.py --check # exit 1 on a leak + +``tests/test_qwen35moe_build_matrix.py`` runs the same checks. +""" + +from __future__ import annotations + +import argparse +import re +import sys +from pathlib import Path + +ROOT = Path(__file__).resolve().parent.parent + +# CMake option -> the compile definition it sets -> what the bindings guard on. +TIERS = { + "FLASHRT_ENABLE_QWEN35MOE_CORE": "FLASHRT_HAVE_QWEN35MOE_CORE", + "FLASHRT_ENABLE_QWEN35MOE_W4A16": "FLASHRT_HAVE_QWEN35MOE_W4A16", + "FLASHRT_ENABLE_QWEN35MOE_W4A4": "FLASHRT_HAVE_QWEN35MOE_W4A4", +} + +# Gates that are not tiers but are still model-specific: the grouped MoE GEMM +# object, which only the weight-only tier on sm_110 builds. +EXTRA_GATES = ("FLASHRT_HAVE_QWEN35MOE_GROUPED_SM100",) + +# Sources gated somewhere other than a tier, with the gate that owns each. +# Checked by name because that is the whole point: the grouped MoE GEMM used to +# be a second source in an object library every Thor build compiles. +ELSEWHERE = { + "csrc/gemm/fp4/cutlass_nvfp4_moe_grouped_sm100.cu": + "qwen35moe_nvfp4_grouped_sm100_obj", +} + + +def _cmake_text() -> str: + return (ROOT / "CMakeLists.txt").read_text(encoding="utf-8") + + +def _bindings_text() -> str: + return (ROOT / "csrc" / "bindings.cpp").read_text(encoding="utf-8") + + +def tier_sources() -> dict[str, list[str]]: + """Sources CMake adds inside each ``if()`` block.""" + text = _cmake_text() + out: dict[str, list[str]] = {} + for option in TIERS: + # The block runs from `if(