diff --git a/jenkins/scripts/cbts/cbts/command/main.py b/jenkins/scripts/cbts/cbts/command/main.py index 17ca3e11079c..2c1b045a4c23 100644 --- a/jenkins/scripts/cbts/cbts/command/main.py +++ b/jenkins/scripts/cbts/cbts/command/main.py @@ -101,7 +101,7 @@ def build_rules( TestListRule(yaml_index, stages, repo_root=repo_root), VisualGenRule(yaml_index, stages), SpecDecRule(yaml_index, stages), - ModelingV2Rule(yaml_index, stages), + ModelingV2Rule(yaml_index, stages, repo_root=repo_root), AgentFlowRule(yaml_index, stages), OpenEngineRule(yaml_index, stages), DocsRule(yaml_index, stages), diff --git a/jenkins/scripts/cbts/cbts/rules/modeling_v2.py b/jenkins/scripts/cbts/cbts/rules/modeling_v2.py index 0cb614c423a4..b355d17bb09c 100644 --- a/jenkins/scripts/cbts/cbts/rules/modeling_v2.py +++ b/jenkins/scripts/cbts/cbts/rules/modeling_v2.py @@ -18,17 +18,41 @@ (checkpoint, GPU arch, parallel topology), assembled from a catalog of op wrappers. -Block selection — entry-pattern based only: -It has no `condition.terms.backend` of its own; its entries sit in -`backend: pytorch` blocks beside everything else. A block belongs to -modeling_v2 iff one of its `tests:` entries matches a marker in -`_MV2_ENTRY_PATTERNS`. Those markers are exact by construction rather -than by luck: every unit test file in the subtree is named -`test_modeling_v2_*` precisely so it cannot collide with the upstream +Block selection: a block belongs to modeling_v2 iff one of its `tests:` +entries matches a marker in `_MV2_ENTRY_PATTERNS`. Those markers are exact +by construction rather than by luck: every unit test file in the subtree is +named `test_modeling_v2_*` precisely so it cannot collide with the upstream test of the same op, and the accuracy gates live in their own -`accuracy/modeling_v2/` directory. So there is no substring that could -claim an unrelated entry, and no `mtp_nextn=0`-style carve-out is -needed. +`accuracy/modeling_v2/` directory. + +Entry selection: each changed file maps to the tests it can affect. + +- `models/` is a tree of nodes. A directory is a node when it, or a + directory below it, holds a target's `modeling.py`; a family directory + (`models//`) is always one. Other directories -- `ops/` and the + like -- serve the node above them. A changed file belongs to the nearest + node at or above it, and selects that node's unit tests + (`unittest/_torch/modeling_v2/models/`) plus the accuracy gate of + every target under it. A target's gate is `Test` in + `accuracy/modeling_v2/test_.py`, the class name read from the + import line in the family's `routing.py` that routes to it. +- A target's `cells.py` is input configurations for catalog ops and + nothing else: it changes no forward. It selects only the unit tests of + the catalog ops it imports, at their mirrored paths + (`catalog//.py` -> `unittest/_torch/modeling_v2//test_modeling_v2_.py`). +- Anything else -- the catalog, the resolver, `models/__init__.py` -- can + reach every target and selects every modeling_v2 entry. + +The tests directly under `unittest/_torch/modeling_v2/` check the tree as a +whole (claims, routing, contracts), so any change keeps them. + +A unit test prefix that no entry of a block sits at or under keeps the +longest entry of that block that covers it instead: a family is one +directory entry, and a change to one of its targets runs the family's +unit tests. Whenever the tree on disk cannot answer -- a node directory +deleted with its `routing.py`, a `cells.py` that cannot be parsed or names +an op whose test is not at the mirrored path -- the file falls back to the +wider selection. Outward fallback: not needed, and that is a property of the design rather than an accident. Nothing beyond the resolver module is imported @@ -60,9 +84,12 @@ from __future__ import annotations +import ast +import re +from pathlib import Path from typing import Optional -from cbts.blocks import Stage, YAMLIndex, _entry_target +from cbts.blocks import Stage, YAMLIndex, _entry_target, _target_in_filter_subtree from ._helpers import resolve_affected_stages, stages_by_yaml_stem from .base import PRInputs, Rule, RuleResult @@ -71,20 +98,27 @@ # Source-path prefixes the rule may claim. Tests under tests/** are left # to TestsDefRule; the two scopes combine via _TESTSONLY_FAMILY. _MV2_SRC_PREFIXES: tuple[str, ...] = ("tensorrt_llm/_torch/_experimental/modeling_v2/",) +_MV2_PACKAGE = "tensorrt_llm._torch._experimental.modeling_v2" + +_UNIT_ROOT = "unittest/_torch/modeling_v2" +_ACCURACY_ROOT = "accuracy/modeling_v2" # Substrings that mark a test entry as modeling_v2. All are unambiguous: -# - "unittest/_torch/modeling_v2/" → the op-level catalog matrix, taken -# as whole-directory entries (one on l0_gb300, one on -# l0_gb300_multi_gpus for the 4-rank collectives) +# - "unittest/_torch/modeling_v2/" → the op and target unit tests # - "accuracy/modeling_v2/" → the per-model accuracy gates # - "test_modeling_v2_" → any future unit file named by the subtree's # own convention _MV2_ENTRY_PATTERNS: tuple[str, ...] = ( - "unittest/_torch/modeling_v2/", - "accuracy/modeling_v2/", + f"{_UNIT_ROOT}/", + f"{_ACCURACY_ROOT}/", "test_modeling_v2_", ) +# A family's `routing.py` imports each target it routes to by this line, from +# the target directory directly under the family. The claims test reads the +# same line with the same pattern. +_TARGET_IMPORT = re.compile(r"^\s*from \.(\w+)\.modeling import (\w+)$", re.MULTILINE) + def _is_mv2_claim(path: str) -> bool: """Decide whether ModelingV2Rule claims `path`. @@ -110,31 +144,130 @@ def _mv2_entries(block) -> list[str]: return [t for t in block.tests if _entry_is_mv2(t)] +def _is_whole_tree_test(target: str) -> bool: + """A test directly under the unit root, which checks the tree as a whole.""" + rest = target.removeprefix(f"{_UNIT_ROOT}/").split("::")[0] + return target.startswith(f"{_UNIT_ROOT}/") and "/" not in rest and rest.endswith(".py") + + def _is_perf_sanity_stem(stem: str) -> bool: """True for perf-sanity yaml stems (`l0_*_perf_sanity*`).""" return "perf_sanity" in stem +def _kept_entries(entries: list[str], prefixes: set[str]) -> list[str]: + """The entries of one block that `prefixes` select. + + An entry is kept when it sits at or under a prefix. A prefix with no such + entry in the block keeps the longest entry that covers it instead. + """ + targets = {entry: _entry_target(entry) for entry in entries} + kept = {e for e, t in targets.items() if _is_whole_tree_test(t)} + for prefix in prefixes: + under = {e for e, t in targets.items() if _target_in_filter_subtree(t, prefix)} + if not under: + covering = [e for e, t in targets.items() if _target_in_filter_subtree(prefix, t)] + under = {max(covering, key=lambda e: len(targets[e]))} if covering else set() + kept |= under + return [e for e in entries if e in kept] + + class ModelingV2Rule(Rule): name = "modelingv2" needs_diff_for: tuple[str, ...] = () - def __init__(self, yaml_index: YAMLIndex, stages: dict[str, Stage]) -> None: + def __init__(self, yaml_index: YAMLIndex, stages: dict[str, Stage], repo_root: Path) -> None: self.yaml_index = yaml_index self._stages_by_yaml = stages_by_yaml_stem(stages) + self._src_root = Path(repo_root) / _MV2_SRC_PREFIXES[0] + self._tests_root = Path(repo_root) / "tests" + + def _is_node(self, rel_dir: Path) -> bool: + """Whether `models/` is a node: a family, or holds a target at or below it.""" + path = self._src_root / "models" / rel_dir + return len(rel_dir.parts) == 1 or (path.is_dir() and any(path.rglob("modeling.py"))) + + def _gates_under(self, node: Path) -> Optional[set[str]]: + """Accuracy-gate prefixes of every target at or under `node`; None if unreadable.""" + family = node.parts[0] + routing = self._src_root / "models" / family / "routing.py" + if not routing.is_file(): + return None + gates = set() + for target, target_cls in _TARGET_IMPORT.findall(routing.read_text(encoding="utf-8")): + leaf = Path(family, target) + if leaf == node or node in leaf.parents: + gates.add(f"{_ACCURACY_ROOT}/test_{family}.py::Test{target_cls}") + return gates + + def _cells_tests(self, rel: Path) -> Optional[set[str]]: + """Unit tests of the catalog ops a `cells.py` imports; None if not resolvable.""" + try: + tree = ast.parse((self._src_root / rel).read_text(encoding="utf-8")) + except (OSError, SyntaxError, UnicodeDecodeError): + return None + package = rel.parent.parts + tests = set() + for node in ast.walk(tree): + if not isinstance(node, ast.ImportFrom): + continue + if node.level: + if node.level > len(package) + 1: + return None + base = package[: len(package) - node.level + 1] + module = (*base, *(node.module or "").split(".")) + elif (node.module or "").startswith(_MV2_PACKAGE + "."): + module = tuple(node.module[len(_MV2_PACKAGE) + 1 :].split(".")) + else: + continue + if module[:1] != ("catalog",) or len(module) != 3 or module[1].startswith("_"): + if module == ("catalog", "_op"): + continue + return None + test = f"{_UNIT_ROOT}/{module[1]}/test_modeling_v2_{module[2]}.py" + if not (self._tests_root / test).is_file(): + return None + tests.add(test) + return tests + + def _prefixes_for(self, path: str) -> Optional[set[str]]: + """Filter prefixes a changed file selects; None selects every modeling_v2 entry.""" + rel = Path(path[len(_MV2_SRC_PREFIXES[0]) :]) + if rel.parts[0] != "models" or len(rel.parts) < 3: + return None + if rel.name == "cells.py": + tests = self._cells_tests(rel) + if tests is not None: + return tests + node = Path(*rel.parts[1:-1]) + while not self._is_node(node): + node = node.parent + gates = self._gates_under(node) + if gates is None: + return None + return {f"{_UNIT_ROOT}/models/{node.as_posix()}"} | gates def apply(self, pr: PRInputs) -> Optional[RuleResult]: claimed = {f for f in pr.changed_files if _is_mv2_claim(f)} if not claimed: return None + prefixes: Optional[set[str]] = set() + for path in sorted(claimed): + selected = self._prefixes_for(path) + if selected is None: + prefixes = None + break + prefixes |= selected + block_filters: dict[tuple[str, int], dict[str, set[str]]] = {} for block in self.yaml_index.blocks: entries = _mv2_entries(block) + if prefixes is not None: + entries = _kept_entries(entries, prefixes) if not entries: continue - key = (block.yaml_stem, block.block_index) - prefix_dict = block_filters.setdefault(key, {}) + prefix_dict = block_filters.setdefault((block.yaml_stem, block.block_index), {}) for entry in entries: target = _entry_target(entry) if target: @@ -168,6 +301,7 @@ def apply(self, pr: PRInputs) -> Optional[RuleResult]: perfsanity_relevant=perfsanity_relevant, reason=( f"modelingv2: {len(claimed)} modeling_v2 source file(s) → " - f"{len(block_filters)} modeling_v2 block(s), {len(affected)} stage(s)" + f"{'every modeling_v2 entry' if prefixes is None else f'{len(prefixes)} test prefix(es)'}, " + f"{len(block_filters)} block(s), {len(affected)} stage(s)" ), ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/__init__.py index c5140e7c912b..96e5209c99dc 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/__init__.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/__init__.py @@ -27,7 +27,6 @@ """ from ._router_index import ( - MODELING_V2_ROUTERS, MODELING_V2_VALIDATE_ENV, ModelingV2Mode, NoTarget, @@ -36,7 +35,6 @@ ) __all__ = [ - "MODELING_V2_ROUTERS", "MODELING_V2_VALIDATE_ENV", "ModelingV2Mode", "NoTarget", diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py b/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py index fc1f9a56de1b..790a9ea677e7 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py @@ -1,9 +1,9 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Architecture name -> routing module, and the resolver that reads the table. +"""Architecture name -> ``route``, and the resolver that consults it. The checkpoint's own ``architectures[0]`` (``GptOssForCausalLM``, -``DeepseekV3ForCausalLM``) selects a routing module, never a target: one +``DeepseekV3ForCausalLM``) selects a family's ``route``, never a target: one architecture has many deployments and a target serves exactly one of them. ``modeling_v2_resolve`` returns the target *class*, the model loader records it on ``ModelConfig.modeling_v2_target``, and ``AutoModelForCausalLM._resolve_class`` @@ -11,11 +11,13 @@ Two levels, on purpose: -* This table maps ``architectures[0]`` -> routing module. Its keys are things - that already exist in every checkpoint's config.json, so adding a new - checkpoint, GPU arch or parallel topology never touches it -- only a new - architecture family does. -* The routing module's ``route`` owns the rest of the decision as one +* Each family's ``routing.py`` declares the ``architectures[0]`` values its + ``route`` serves with ``@routes(...)``, and ``models/__init__.py`` imports + every family's ``routing.py``, which is what registers them. The names are + things that already exist in every checkpoint's config.json, so adding a + new checkpoint, GPU arch or parallel topology never touches either -- only + a new architecture family does. +* That ``route`` owns the rest of the decision as one forward-reading tree over the facts that fix which code runs: the checkpoint's config, the parallel mapping, the GPU's compute capability, the speculative-decoding config and the KV-cache algorithm the engine will run @@ -29,17 +31,17 @@ target claims a configuration regardless of them; ``require`` promises that the target ran, not that a particular knob setting was measured. -Routing modules are imported lazily: resolving a GptOss config never imports -the DeepSeek tree, and a target's modeling code is imported only on the -``route`` branch that returns it. +Routing modules are imported on the first resolve, all of them, and hold no +modeling code: a target's modeling code is imported only on the ``route`` +branch that returns it, so resolving a GptOss config never imports a DeepSeek +target. """ from __future__ import annotations import enum import os -from importlib import import_module -from typing import TYPE_CHECKING, Optional, Tuple +from typing import TYPE_CHECKING, Callable, Dict, Optional, Tuple from tensorrt_llm.logger import logger @@ -47,17 +49,29 @@ from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm.llmapi.llm_args import TorchLlmArgs -_PACKAGE = "tensorrt_llm._torch._experimental.modeling_v2" +# architectures[0] -> the ``route`` that decides where it lands, filled by ``@routes``. +_ROUTES: Dict[str, Callable[..., type]] = {} -# architectures[0] -> routing module, relative to this package. -MODELING_V2_ROUTERS = { - "GptOssForCausalLM": "models.gpt_oss.routing", - "DeepseekV3ForCausalLM": "models.deepseek_v3.routing", - # Qwen3.8-Flash-Next publishes ``Qwen4ExpForConditionalGeneration``, and the - # config loader's composite branch writes the same name back. - "Qwen4ExpForConditionalGeneration": "models.qwen4_exp.routing", - "GlmMoeDsaForCausalLM": "models.glm_moe_dsa.routing", -} + +def routes(*architectures: str): + """Declare that the decorated ``route`` decides where these ``architectures[0]`` values land.""" + + def register(route): + for arch in architectures: + assert arch not in _ROUTES, ( + f"{arch} routed by both {_ROUTES[arch].__module__} and {route.__module__}" + ) + _ROUTES[arch] = route + return route + + return register + + +def registered_routes() -> Dict[str, Callable[..., type]]: + """architectures[0] -> ``route``, for every family ``models/__init__.py`` imports.""" + from . import models # noqa: F401 + + return _ROUTES #: Turns the step contract check on. Off by default, so a served engine runs @@ -99,7 +113,7 @@ def of(cls, llm_args: "TorchLlmArgs") -> "ModelingV2Mode": class NoTarget(Exception): - """Raised by a routing module's ``route`` when no target fits. + """Raised by a family's ``route`` when no target fits. The message says, in words, which fact ruled the targets out. Under ``auto`` the resolver logs it and falls back to the built-in implementation; @@ -130,14 +144,6 @@ def tp_reaches_across_nodes(mapping) -> bool: return MNNVLAllReduce.is_mnnvl(mapping, torch.bfloat16) -def routing_module(arch: str): - """Import and return the routing module for ``arch``, or None.""" - name = MODELING_V2_ROUTERS.get(arch) - if name is None: - return None - return import_module(f"{_PACKAGE}.{name}") - - def _device_compute_capability() -> Tuple[int, int]: import torch @@ -172,8 +178,8 @@ def modeling_v2_resolve(config: "ModelConfig", llm_args: "TorchLlmArgs") -> Opti """Decide which modeling_v2 target, if any, this deployment builds. Returns the target class, or None for the built-in implementation: - when ``llm_args.modeling_v2`` is ``off``, when no routing module claims the - architecture, or when the routing module's ``route`` raises ``NoTarget``. + when ``llm_args.modeling_v2`` is ``off``, when no ``route`` is registered + for the architecture, or when that ``route`` raises ``NoTarget``. In ``require`` every one of those but ``off`` raises ``ValueError`` instead. """ mode = ModelingV2Mode.of(llm_args) @@ -186,13 +192,11 @@ def modeling_v2_resolve(config: "ModelConfig", llm_args: "TorchLlmArgs") -> Opti arch = pretrained_config.architectures[0] try: - routing = routing_module(arch) - if routing is None: - known = ", ".join(sorted(MODELING_V2_ROUTERS)) or "(none)" - raise NoTarget( - f"no routing module for this architecture; routed architectures: {known}" - ) - return routing.route( + route = registered_routes().get(arch) + if route is None: + known = ", ".join(sorted(registered_routes())) or "(none)" + raise NoTarget(f"no route for this architecture; routed architectures: {known}") + return route( pretrained_config, config.mapping, _device_compute_capability(), diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/_op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/_op.py index 421e226b26bf..11f411f438eb 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/_op.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/_op.py @@ -24,10 +24,12 @@ business, not the entry's: nothing on the way to the kernel needs it. `CELLS` is also the answer to "can my model use this op?". A target author -reads the cell list and its reasons, not a document -- and what they read is -what CI proved on the last commit. Cells cover what the shipped targets -actually run; an input outside them is not known to work, which for most of -these ops means "nobody measured it" rather than "known broken". +reads the cell lists and their reasons, not a document -- and what they read is +what CI proved on the last commit. An entry's own `CELLS` are the +configurations that belong to the op whoever calls it; the shapes a shipped +target runs are in that target's `cells.py`, keyed by entry, and the entry's +test drives both. An input outside them is not known to work, which for most +of these ops means "nobody measured it" rather than "known broken". `is_valid` is never called on the way to the kernel. It states, in executable form, the inputs the op accepts and quietly answers wrongly -- the class of @@ -99,8 +101,8 @@ class Cell: `spec` is entry-specific: the entry's test knows how to turn it into tensors, and nothing here interprets it. That keeps environment building -- a KV-cache manager, a 4-rank communicator -- in the test where it belongs, - while the *set of configurations* stays with the op, which is what a target - author needs to read. + while the *set of configurations* stays with the op and the targets that + call it, which is what a target author needs to read. `why` is not decoration. A cell that cannot say what it covers beyond the ones already listed is a cell that costs CI time for nothing. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/activation/flashinfer_silu_and_mul.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/activation/flashinfer_silu_and_mul.py index 3ffd35cc5311..3e67f8d6690b 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/activation/flashinfer_silu_and_mul.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/activation/flashinfer_silu_and_mul.py @@ -37,10 +37,6 @@ class FlashinferSiluAndMul(OpWrapper): why="the dense layers' wider intermediate, 2 x 18432", spec=dict(shape=(16, 36864), dtype=torch.bfloat16), ), - Cell( - why="deepseek default at attention TP8: the shared expert's gate+up, 2 x 2048/8", - spec=dict(shape=(16, 512), dtype=torch.bfloat16), - ), ) note = """ diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/fused_qk_norm_rope.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/fused_qk_norm_rope.py index e6dbc41e29ba..d0eb494f57bb 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/fused_qk_norm_rope.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/fused_qk_norm_rope.py @@ -35,7 +35,7 @@ class FusedQkNormRope(OpWrapper): # heads, 8 key and 8 value, head_dim 64, all of which the cells carry # because the op reads them as scalars rather than deriving them. Under # attention TP each rank sees its own share of the heads. - _GPT_OSS = dict( + GPT_OSS = dict( num_heads_q=64, num_heads_k=8, num_heads_v=8, @@ -52,120 +52,12 @@ class FusedQkNormRope(OpWrapper): attention_factor=1.0, ) - # qwen4_exp (Qwen3.8-Flash-Next) runs Gemma QK norm and an unscaled rope - # on the first 64 dims (32 pairs), base 1e7, at two geometries. Its - # mrope_section is [11, 11, 10], passed as sections 11 / 11. Positions are - # drawn uniformly below `max_position` (a test axis, not an op argument): - # the kernel's fast sincos loses phase as `position * freq` grows, so the - # position range is part of what a cell certifies. - _QWEN_ROPE = dict( - rotary_dim=64, - base=1e7, - is_neox=True, - is_qk_norm=True, - use_gemma=True, - max_position=32768, - ) - _QWEN_MROPE = dict(use_mrope=True, mrope_section1=11, mrope_section2=11, **_QWEN_ROPE) - # Main attention (trunk and MTP): 24 query heads, 2 kv, head_dim 256. - _QWEN_MAIN = dict(num_heads_q=24, num_heads_k=2, num_heads_v=2, head_dim=256, **_QWEN_MROPE) - # The QSA index plane's query rope: 4 index heads plus the one k and one v - # slot of its scratch layout, over index_head_dim 128 (== 2 * rotary_dim, - # so the partial factor is not reapplied). - _QWEN_INDEX_HEADS = dict(num_heads_q=4, num_heads_k=1, num_heads_v=1, head_dim=128) - _QWEN_INDEX_MROPE = dict(**_QWEN_INDEX_HEADS, **_QWEN_MROPE) - _QWEN_INDEX_PLAIN = dict(**_QWEN_INDEX_HEADS, **_QWEN_ROPE) - CELLS: tuple[Cell, ...] = ( - Cell(why="decode: one token", spec=dict(tokens=1, **_GPT_OSS)), - Cell(why="prefill-shaped batch", spec=dict(tokens=2048, **_GPT_OSS)), + Cell(why="decode: one token", spec=dict(tokens=1, **GPT_OSS)), + Cell(why="prefill-shaped batch", spec=dict(tokens=2048, **GPT_OSS)), Cell( why="a token count that is not a multiple of any tile", - spec=dict(tokens=37, **_GPT_OSS), - ), - Cell( - why="default target at attention TP8: 8q/1k/1v heads per rank", - spec=dict( - tokens=64, **{**_GPT_OSS, "num_heads_q": 8, "num_heads_k": 1, "num_heads_v": 1} - ), - ), - Cell( - why="qwen4_exp main rope, decode: one token, Gemma norm and interleaved mRoPE " - "with three distinct position rows", - spec=dict(tokens=1, **_QWEN_MAIN), - ), - Cell( - why="qwen4_exp main rope, 19 tokens: the size its wiring and saturation tests run at", - spec=dict(tokens=19, **_QWEN_MAIN), - ), - Cell( - why="qwen4_exp main rope, 8192-token prefill: the Gemma one-ulp tail, under " - "GEMMA_PREFILL_ATOL", - spec=dict(tokens=8192, gemma_prefill_band=True, **_QWEN_MAIN), - ), - # The main rope run in place on qwen4_exp's fused input projection, whose - # row is [q | k | v | gate | index | pad] = 55 heads of 256: everything - # after k is passed as v heads, which only sets the row stride. - Cell( - why="qwen4_exp fused-row main rope, 1 token: 29 v heads only set the row stride", - spec=dict(tokens=1, **{**_QWEN_MAIN, "num_heads_v": 29}), - ), - Cell( - why="qwen4_exp fused-row main rope, 128 tokens: decode at max concurrency", - spec=dict(tokens=128, **{**_QWEN_MAIN, "num_heads_v": 29}), - ), - Cell( - why="qwen4_exp fused-row main rope, 8192 tokens: prefill, under GEMMA_PREFILL_ATOL", - spec=dict(tokens=8192, gemma_prefill_band=True, **{**_QWEN_MAIN, "num_heads_v": 29}), - ), - Cell( - why="qwen4_exp index rope (mRoPE), 1 token: one k and one v head, the " - "narrowest layout; 5 warps of an 8-warp block", - spec=dict(tokens=1, **_QWEN_INDEX_MROPE), - ), - Cell( - why="qwen4_exp index rope (mRoPE), 2 tokens: the second token's heads " - "straddle a block boundary", - spec=dict(tokens=2, **_QWEN_INDEX_MROPE), - ), - Cell( - why="qwen4_exp index rope (mRoPE), 5 tokens: the last block runs one warp", - spec=dict(tokens=5, **_QWEN_INDEX_MROPE), - ), - Cell( - why="qwen4_exp index rope (mRoPE), 17 tokens: one past a power of two", - spec=dict(tokens=17, **_QWEN_INDEX_MROPE), - ), - Cell( - why="qwen4_exp index rope (mRoPE), 19 tokens: the size its wiring and " - "saturation tests run at", - spec=dict(tokens=19, **_QWEN_INDEX_MROPE), - ), - Cell( - why="qwen4_exp index rope (mRoPE), 33 tokens: one past a power of two", - spec=dict(tokens=33, **_QWEN_INDEX_MROPE), - ), - Cell( - why="qwen4_exp index rope (mRoPE), 65 tokens: one past a power of two", - spec=dict(tokens=65, **_QWEN_INDEX_MROPE), - ), - Cell( - why="qwen4_exp index rope (mRoPE), 8192-token prefill: the index site sees the " - "whole packed batch; under GEMMA_PREFILL_ATOL", - spec=dict(tokens=8192, gemma_prefill_band=True, **_QWEN_INDEX_MROPE), - ), - Cell( - why="qwen4_exp index rope (plain), decode: the form the index plane runs " - "today, row 0 alone as a flat [T] position_ids", - spec=dict(tokens=1, **_QWEN_INDEX_PLAIN), - ), - Cell( - why="qwen4_exp index rope (plain), 19 tokens: an odd count across blocks", - spec=dict(tokens=19, **_QWEN_INDEX_PLAIN), - ), - Cell( - why="qwen4_exp index rope (plain), 8192-token prefill: under GEMMA_PREFILL_ATOL", - spec=dict(tokens=8192, gemma_prefill_band=True, **_QWEN_INDEX_PLAIN), + spec=dict(tokens=37, **GPT_OSS), ), ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/mla_rope_append_paged_kv_assign_q.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/mla_rope_append_paged_kv_assign_q.py index 7a88d076bca1..b603b40c4373 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/mla_rope_append_paged_kv_assign_q.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/mla_rope_append_paged_kv_assign_q.py @@ -31,7 +31,7 @@ class MlaRopeAppendPagedKvAssignQ(OpWrapper): # fixed by the checkpoint -- 512-wide compressed_kv, 64-wide k_pe, 128 # heads -- so the cells vary what a step actually varies: how much of each # sequence was already cached, how many tokens are new, and the page size. - _R1 = dict( + R1 = dict( head_num=128, nope_size=128, rope_size=64, @@ -43,24 +43,15 @@ class MlaRopeAppendPagedKvAssignQ(OpWrapper): CELLS: tuple[Cell, ...] = ( Cell( why="production context call: fp8 pool, page 32, mixed cached lengths", - spec=dict(cached=[0, 32, 100], new=[37, 64, 300], tokens_per_block=32, **_R1), + spec=dict(cached=[0, 32, 100], new=[37, 64, 300], tokens_per_block=32, **R1), ), Cell( why="page 64, the size a tuned MLA target opts into", - spec=dict(cached=[0, 32, 100], new=[37, 64, 300], tokens_per_block=64, **_R1), + spec=dict(cached=[0, 32, 100], new=[37, 64, 300], tokens_per_block=64, **R1), ), Cell( why="a single new token on top of a long cached prefix -- the decode-adjacent shape", - spec=dict(cached=[511], new=[1], tokens_per_block=64, **_R1), - ), - Cell( - why="deepseek default at attention TP8: 16 heads per rank, production context", - spec=dict( - cached=[0, 32, 100], - new=[37, 64, 300], - tokens_per_block=32, - **dict(_R1, head_num=16), - ), + spec=dict(cached=[511], new=[1], tokens_per_block=64, **R1), ), ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/mla_rope_generation.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/mla_rope_generation.py index cf497590ae33..b346a31885de 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/mla_rope_generation.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/mla_rope_generation.py @@ -45,7 +45,7 @@ class MlaRopeGeneration(OpWrapper): # single-layer. Fixed by the checkpoint, so the cells vary what a step # varies -- how many new tokens, the page size, and whether the scales are # passed explicitly. - _R1 = dict( + R1 = dict( num_heads=128, kv_lora_rank=512, qk_rope_head_dim=64, @@ -62,7 +62,7 @@ class MlaRopeGeneration(OpWrapper): q_scaling=Q_SCALING_R1, write_scale=None, read_scale=None, - **_R1, + **R1, ), ), Cell( @@ -73,7 +73,7 @@ class MlaRopeGeneration(OpWrapper): q_scaling=Q_SCALING_R1, write_scale=None, read_scale=None, - **_R1, + **R1, ), ), Cell( @@ -84,29 +84,7 @@ class MlaRopeGeneration(OpWrapper): q_scaling=1.0, write_scale=0.5, read_scale=2.0, - **_R1, - ), - ), - Cell( - why="deepseek default at attention TP8: 16 heads per rank, production decode", - spec=dict( - predicted_tokens_per_seq=1, - tokens_per_block=32, - q_scaling=Q_SCALING_R1, - write_scale=None, - read_scale=None, - **dict(_R1, num_heads=16), - ), - ), - Cell( - why="deepseek default MTP3 at attention TP8: P=4 (draft length 3 + 1) at 16 heads", - spec=dict( - predicted_tokens_per_seq=4, - tokens_per_block=32, - q_scaling=Q_SCALING_R1, - write_scale=None, - read_scale=None, - **dict(_R1, num_heads=16), + **R1, ), ), ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/thop_attention.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/thop_attention.py index 80a2ad211f21..d302d359aa94 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/thop_attention.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/thop_attention.py @@ -35,7 +35,7 @@ #: the paged-context path on (the target serves with block reuse and chunked #: prefill). Pool sizes are chosen, not dumped: 128 pages of 32 tokens and 4 #: block-offset rows, larger than any cell's batch. -_QWEN4_DENSE = dict( +QWEN4_DENSE = dict( input_type=STANDARD, num_heads=24, num_kv_heads=2, @@ -54,7 +54,7 @@ #: qwen4_exp's vision tower: the no-KV-cache mode, one context sequence per #: frame, MHA 16/16 at head_size 72 (hidden 1152 over 16 heads), padding mask. #: `max_num_requests` is the segment count, as the target passes it. -_QWEN4_VISION = dict( +QWEN4_VISION = dict( input_type=STANDARD, kv_cache="none", num_heads=16, @@ -67,33 +67,8 @@ mask_type=0, ) -#: The qwen4_exp full-attention call over an fp8-e4m3 pool at the production kv -#: scale 1.0 (fp32 [1] scale tensors, quant_mode 128): what the trunk layers and -#: the MTP draft layer pass when the engine serves an fp8 KV cache. Same chosen -#: pool sizes as `_QWEN4_DENSE`. -_QWEN4_FP8 = dict(_QWEN4_DENSE, quant_mode=QM_FP8_KV_CACHE, kv_scale=1.0) - -#: The same call over an NVFP4 pool: e2m1 data and a second pool of e4m3 block -#: scales, fp32 [3] unit kv scales, `[1, 2, 2]` pool pointers, bf16 output -- -#: what the trunk layers and the MTP draft layer pass when the engine serves an -#: NVFP4 KV cache. Same chosen pool sizes as `_QWEN4_DENSE`. -_QWEN4_NVFP4 = dict(_QWEN4_DENSE, quant_mode=QM_NVFP4_KV_CACHE, kv_scale=1.0) - -#: `_QWEN4_DENSE` and `_QWEN4_FP8` for a verify step of up to 32 requests, the captured-batch -#: ceiling: one block-offset row per request and pages for 32 histories (the largest cell -#: needs 133 of the 144). -_QWEN4_DENSE_G32 = dict(_QWEN4_DENSE, max_num_requests=32, num_blocks=144) -_QWEN4_FP8_G32 = dict(_QWEN4_FP8, max_num_requests=32, num_blocks=144) - - -def _spread(batch: int, longest: int) -> tuple[int, ...]: - """`batch` histories rising evenly from 1 to `longest` tokens: one-page, page-boundary - and multi-page sequences in one verify step.""" - return tuple(1 + (longest - 1) * i // (batch - 1) for i in range(batch)) - - #: A GQA 24/4/256 full-attention layer (Qwen3.8-27B) over an fp8-e4m3 pool at -#: the production kv scale 1.0. Same chosen pool sizes as `_QWEN4_DENSE`. +#: the production kv scale 1.0. Same chosen pool sizes as `QWEN4_DENSE`. _FP8_24_4_256 = dict( input_type=STANDARD, num_heads=24, @@ -130,300 +105,17 @@ class ThopAttention(OpWrapper): ARCHS = frozenset({Arch.SM_100, Arch.SM_103}) - # gpt-oss alternates its two window classes inside one engine, so both of - # its cells are live every step; deepseek runs context and generation - # against the same pool. qwen4_exp's cells give `cached` (tokens already in - # the pool) and `new` (this call's q rows) per sequence, contexts first. + # The shipped targets' cells are in their cells.py. gpt-oss alternates its + # two window classes inside one engine, so both of its cells are live every + # step; deepseek runs context and generation against the same pool. Cells + # with `cached` (tokens already in the pool) and `new` (this call's q rows) + # give them per sequence, contexts first. CELLS: tuple[Cell, ...] = ( - Cell( - why="gpt-oss sliding layer: GQA 64q/8kv d64, sinks, window 128", - spec=dict( - input_type=STANDARD, - num_heads=64, - num_kv_heads=8, - head_size=64, - v_head_dim=64, - sinks=True, - window=128, - quant_mode=0, - ), - ), - Cell( - why="gpt-oss full layer: the same heads with no window", - spec=dict( - input_type=STANDARD, - num_heads=64, - num_kv_heads=8, - head_size=64, - v_head_dim=64, - sinks=True, - window=None, - quant_mode=0, - ), - ), - # The per-rank head counts the gpt-oss default target runs under - # attention TP: the same code path as the two cells above, at the - # head counts a rank holds after the split. - Cell( - why="gpt-oss default at attention TP8: 8q/1kv per rank, 128-token window", - spec=dict( - input_type=STANDARD, - num_heads=8, - num_kv_heads=1, - head_size=64, - v_head_dim=64, - sinks=True, - window=128, - quant_mode=0, - ), - ), - Cell( - why="deepseek MLA context: explicit k/v, qk 576 -> v 128, fp8 pool", - spec=dict( - input_type=MLA_CONTEXT, - num_heads=128, - num_kv_heads=128, - head_size=576, - v_head_dim=128, - sinks=False, - window=None, - quant_mode=QM_FP8_KV_CACHE, - ), - ), - Cell( - why="deepseek MLA generation: one kv head, latent 576 -> 512, fp8 pool", - spec=dict( - input_type=MLA_GENERATION, - num_heads=128, - num_kv_heads=1, - head_size=576, - v_head_dim=512, - sinks=False, - window=None, - quant_mode=QM_FP8_KV_CACHE, - ), - ), - # The per-rank head counts the deepseek default target runs under - # attention TP: the same two MLA code paths as the cells above, at the - # 128 heads split 2, 4 or 8 ways. - Cell( - why="deepseek default at attention TP8: MLA context, 16 heads per rank", - spec=dict( - input_type=MLA_CONTEXT, - num_heads=16, - num_kv_heads=16, - head_size=576, - v_head_dim=128, - sinks=False, - window=None, - quant_mode=QM_FP8_KV_CACHE, - ), - ), - Cell( - why="deepseek default at attention TP2: MLA generation, 64 heads per rank", - spec=dict( - input_type=MLA_GENERATION, - num_heads=64, - num_kv_heads=1, - head_size=576, - v_head_dim=512, - sinks=False, - window=None, - quant_mode=QM_FP8_KV_CACHE, - ), - ), - Cell( - why="deepseek default at attention TP8: MLA generation, 16 heads per rank", - spec=dict( - input_type=MLA_GENERATION, - num_heads=16, - num_kv_heads=1, - head_size=576, - v_head_dim=512, - sinks=False, - window=None, - quant_mode=QM_FP8_KV_CACHE, - ), - ), - # MTP3 under attention TP: the verify step's taller query block (draft - # length 3 + 1 rows per sequence) at a split head count. The P > 1 mask - # itself is pinned at 128 heads by the MTP cases in the test. - Cell( - why="deepseek default MTP3 at attention TP8: MLA generation, 16 heads, P=4", - spec=dict( - input_type=MLA_GENERATION, - num_heads=16, - num_kv_heads=1, - head_size=576, - v_head_dim=512, - sinks=False, - window=None, - quant_mode=QM_FP8_KV_CACHE, - predicted_tokens_per_seq=4, - ), - ), - Cell( - why="qwen4_exp prefill: GQA 24q/2kv d256, bf16 pool, paged context; " - "a fresh 70-row sequence beside 33 rows over a 48-token cached prefix", - spec=dict( - _QWEN4_DENSE, - cached=(0, 48), - new=(70, 33), - num_contexts=2, - predicted_tokens_per_seq=1, - ), - ), - Cell( - why="qwen4_exp decode P=1: one row per generation sequence over " - "page-crossing histories (70, 35)", - spec=dict( - _QWEN4_DENSE, - cached=(70, 35), - new=(1, 1), - num_contexts=0, - predicted_tokens_per_seq=1, - ), - ), - Cell( - why="qwen4_exp MTP verify P=2 (max_draft_len 1): draft row t attends " - "[0, L-P+t], bottom-right causal within the block", - spec=dict( - _QWEN4_DENSE, - cached=(70, 35), - new=(2, 2), - num_contexts=0, - predicted_tokens_per_seq=2, - ), - ), - Cell( - why="qwen4_exp MTP verify P=3 (max_draft_len 2): a separate decode " - "compile, P reaches the kernel as maxSeqLenQ", - spec=dict( - _QWEN4_DENSE, - cached=(70, 35), - new=(3, 3), - num_contexts=0, - predicted_tokens_per_seq=3, - ), - ), - Cell( - why="qwen4_exp MTP verify P=4 (max_draft_len 3)", - spec=dict( - _QWEN4_DENSE, - cached=(70, 35), - new=(4, 4), - num_contexts=0, - predicted_tokens_per_seq=4, - ), - ), - *( - Cell( - why=f"qwen4_exp MTP verify P={p} (max_draft_len {p - 1})" - + (": the target's ceiling" if p == 8 else ""), - spec=dict( - _QWEN4_DENSE, - cached=(70, 35), - new=(p, p), - num_contexts=0, - predicted_tokens_per_seq=p, - ), - ) - for p in (5, 6, 7, 8) - ), - Cell( - why="qwen4_exp P=8 G=1: one draft block from row 95, across the page edge at 96", - spec=dict( - _QWEN4_DENSE, - cached=(95,), - new=(8,), - num_contexts=0, - predicted_tokens_per_seq=8, - ), - ), - # The trtllm-gen launch picks its generation kernel by batch size at P 5..8, so each P - # is also driven at 32 requests and at one batch in 17..20, with histories of every - # page shape in the step. - *( - Cell( - why=f"qwen4_exp bf16 batch P={p} G=32: 32 blocks over histories 1..218, " - f"{32 * p} q rows", - spec=dict( - _QWEN4_DENSE_G32, - cached=_spread(32, 218), - new=(p,) * 32, - num_contexts=0, - predicted_tokens_per_seq=p, - ), - ) - for p in (5, 6, 7, 8) - ), - *( - Cell( - why=f"qwen4_exp bf16 batch P={p} G={g}: {g} blocks over histories 200..1, " - f"{g * p} q rows", - spec=dict( - _QWEN4_DENSE_G32, - cached=tuple(reversed(_spread(g, 200))), - new=(p,) * g, - num_contexts=0, - predicted_tokens_per_seq=p, - ), - ) - for p, g in ((5, 17), (6, 18), (7, 19), (8, 20)) - ), - Cell( - why="qwen4_exp mixed step: a fresh 20-row context ahead of two P=2 " - "draft blocks, generation rows token-major after the context rows", - spec=dict( - _QWEN4_DENSE, - cached=(0, 70, 35), - new=(20, 2, 2), - num_contexts=1, - predicted_tokens_per_seq=2, - ), - ), - Cell( - why="qwen4_exp mixed step P=8: a fresh 20-row context ahead of two P=8 " - "draft blocks over histories 70 and 35", - spec=dict( - _QWEN4_DENSE, - cached=(0, 70, 35), - new=(20, 8, 8), - num_contexts=1, - predicted_tokens_per_seq=8, - ), - ), - Cell( - why="qwen4_exp hybrid pool: row 47 of the 48-row two-pool mapping " - "(pool 0, layer-in-pool 11), P=2 verify", - spec=dict( - _QWEN4_DENSE, - pool_layout="hybrid48", - local_layer_idx=47, - cached=(64, 25), - new=(2, 2), - num_contexts=0, - predicted_tokens_per_seq=2, - ), - ), - Cell( - why="qwen4_exp P=8 hybrid pool: row 47 of the 48-row two-pool mapping " - "(pool 0, layer-in-pool 11), P=8 verify", - spec=dict( - _QWEN4_DENSE, - pool_layout="hybrid48", - local_layer_idx=47, - cached=(64, 25), - new=(8, 8), - num_contexts=0, - predicted_tokens_per_seq=8, - ), - ), Cell( why="48-layer single pool, identity mapping: row 47 = layer-in-pool 47, " "24 rows over a 40-token cached prefix", spec=dict( - _QWEN4_DENSE, + QWEN4_DENSE, pool_layout="identity48", local_layer_idx=47, cached=(40,), @@ -432,71 +124,41 @@ class ThopAttention(OpWrapper): predicted_tokens_per_seq=1, ), ), - Cell( - why="qwen4_exp vision tower: no KV cache, MHA 16/16 d72, padding mask, " - "frames of 27/9/41 rows each attending only to itself", - spec=dict(_QWEN4_VISION, seg_lens=(27, 9, 41)), - ), Cell( why="encoder-only, d72 causal: measured, not what the tower passes", - spec=dict(_QWEN4_VISION, mask_type=1, seg_lens=(27, 9, 41)), + spec=dict(QWEN4_VISION, mask_type=1, seg_lens=(27, 9, 41)), ), Cell( why="encoder-only, d128 padding: the tower's other admitted width", - spec=dict(_QWEN4_VISION, head_size=128, v_head_dim=128, seg_lens=(27, 9, 41)), + spec=dict(QWEN4_VISION, head_size=128, v_head_dim=128, seg_lens=(27, 9, 41)), ), Cell( why="encoder-only, d128 causal", spec=dict( - _QWEN4_VISION, head_size=128, v_head_dim=128, mask_type=1, seg_lens=(27, 9, 41) + QWEN4_VISION, head_size=128, v_head_dim=128, mask_type=1, seg_lens=(27, 9, 41) ), ), Cell( why="encoder-only, a single one-token frame", - spec=dict(_QWEN4_VISION, seg_lens=(1,)), - ), - Cell( - why="mmmu serving shape: one 5580-row frame, the 62x90 patch grid", - spec=dict(_QWEN4_VISION, seg_lens=(5580,)), - ), - Cell( - why="mmmu serving shape: one 8160-row frame, the largest single-segment call", - spec=dict(_QWEN4_VISION, seg_lens=(8160,)), - ), - Cell( - why="mmmu serving shape: 2204 rows in four frames", - spec=dict(_QWEN4_VISION, seg_lens=(728, 728, 460, 288)), - ), - Cell( - why="mmmu serving shape: 1892 rows in five frames, the most segments per call", - spec=dict(_QWEN4_VISION, seg_lens=(480, 480, 400, 352, 180)), - ), - Cell( - why="mmmu serving shape: 9008 rows in two frames", - spec=dict(_QWEN4_VISION, seg_lens=(8000, 1008)), - ), - Cell( - why="mmmu serving shape: 22144 rows in two frames, the largest call; its " - "12544-row frame is the longest segment", - spec=dict(_QWEN4_VISION, seg_lens=(12544, 9600)), + spec=dict(QWEN4_VISION, seg_lens=(1,)), ), # A vision-tower call is one frame per frame group, at most 32768 rows per call. Cell( why="video: 384 frame groups of 64 rows (8x8 patches), the untiled " "kernel's longest frame", - spec=dict(_QWEN4_VISION, seg_lens=(64,) * 384), + spec=dict(QWEN4_VISION, seg_lens=(64,) * 384), ), Cell( why="video: 256 frame groups of 128 rows, the 32768-row ceiling", - spec=dict(_QWEN4_VISION, seg_lens=(128,) * 256), + spec=dict(QWEN4_VISION, seg_lens=(128,) * 256), ), Cell( why="video: 2 frame groups of 16384 rows (128x128), the longest frame at the ceiling", - spec=dict(_QWEN4_VISION, seg_lens=(16384, 16384)), + spec=dict(QWEN4_VISION, seg_lens=(16384, 16384)), ), Cell( why="video: 8192 frame groups of 2x2 patches, the most a 32768-row call can carry", - spec=dict(_QWEN4_VISION, seg_lens=(4,) * 8192), + spec=dict(QWEN4_VISION, seg_lens=(4,) * 8192), ), Cell( why="fp8 pool packed prefill, GQA 24q/4kv d256: the FMHA reads the " @@ -606,165 +268,6 @@ class ThopAttention(OpWrapper): predicted_tokens_per_seq=3, ), ), - # qwen4_exp's own geometry over an fp8 pool: the kv head count changes the compiled - # decode kernel, so the 24/4/256 cells above do not cover it. - Cell( - why="qwen4_exp fp8 pool paged-context prefill, 24/2/256: a fresh 48-row " - "sequence beside 33 rows over a 48-token e4m3 prefix", - spec=dict( - _QWEN4_FP8, - cached=(0, 48), - new=(48, 33), - num_contexts=2, - predicted_tokens_per_seq=1, - ), - ), - Cell( - why="qwen4_exp fp8 pool decode P=1, 24/2/256: histories 200 and 48 span " - "pages and KV tiles", - spec=dict( - _QWEN4_FP8, - cached=(200, 48), - new=(1, 1), - num_contexts=0, - predicted_tokens_per_seq=1, - ), - ), - *( - Cell( - why=f"qwen4_exp fp8 pool MTP verify P={p}, 24/2/256 (max_draft_len {p - 1})", - spec=dict( - _QWEN4_FP8, - cached=(70, 35), - new=(p, p), - num_contexts=0, - predicted_tokens_per_seq=p, - ), - ) - for p in (2, 4, 5, 6, 7, 8) - ), - Cell( - why="qwen4_exp fp8 pool mixed step P=3, 24/2/256: a 20-row context ahead " - "of two P=3 blocks over histories 31 and 64", - spec=dict( - _QWEN4_FP8, - cached=(0, 31, 64), - new=(20, 3, 3), - num_contexts=1, - predicted_tokens_per_seq=3, - ), - ), - Cell( - why="qwen4_exp fp8 pool mixed step P=8, 24/2/256: a 20-row context ahead " - "of two P=8 blocks over histories 31 and 64", - spec=dict( - _QWEN4_FP8, - cached=(0, 31, 64), - new=(20, 8, 8), - num_contexts=1, - predicted_tokens_per_seq=8, - ), - ), - # KV stays within one 128-key tile (histories to 119 plus 8 rows), where the e4m3 - # probability model of `reference` is the kernel's own at any batch. - Cell( - why="qwen4_exp fp8 batch P=8 G=32: 32 blocks over histories 1..119, 256 q rows", - spec=dict( - _QWEN4_FP8_G32, - cached=_spread(32, 119), - new=(8,) * 32, - num_contexts=0, - predicted_tokens_per_seq=8, - ), - ), - Cell( - why="qwen4_exp fp8 hybrid pool P=2: row 47 of the 48-row two-pool mapping " - "(pool 0, layer-in-pool 11) over e4m3 slabs", - spec=dict( - _QWEN4_FP8, - pool_layout="hybrid48", - local_layer_idx=47, - cached=(64, 25), - new=(2, 2), - num_contexts=0, - predicted_tokens_per_seq=2, - ), - ), - Cell( - why="qwen4_exp fp8 hybrid pool P=8: row 47 of the 48-row two-pool mapping " - "(pool 0, layer-in-pool 11) over e4m3 slabs", - spec=dict( - _QWEN4_FP8, - pool_layout="hybrid48", - local_layer_idx=47, - cached=(64, 25), - new=(8, 8), - num_contexts=0, - predicted_tokens_per_seq=8, - ), - ), - Cell( - why="qwen4_exp nvfp4 pool paged-context prefill, 24/2/256: a fresh 48-row " - "sequence beside 33 rows over a 48-token e2m1 prefix", - spec=dict( - _QWEN4_NVFP4, - cached=(0, 48), - new=(48, 33), - num_contexts=2, - predicted_tokens_per_seq=1, - ), - ), - Cell( - why="qwen4_exp nvfp4 pool decode P=1, 24/2/256: histories 200 and 48", - spec=dict( - _QWEN4_NVFP4, - cached=(200, 48), - new=(1, 1), - num_contexts=0, - predicted_tokens_per_seq=1, - ), - ), - Cell( - why="qwen4_exp nvfp4 pool MTP verify P=2, 24/2/256 (max_draft_len 1)", - spec=dict( - _QWEN4_NVFP4, - cached=(70, 35), - new=(2, 2), - num_contexts=0, - predicted_tokens_per_seq=2, - ), - ), - Cell( - why="qwen4_exp nvfp4 pool MTP verify P=4, 24/2/256 (max_draft_len 3)", - spec=dict( - _QWEN4_NVFP4, - cached=(70, 35), - new=(4, 4), - num_contexts=0, - predicted_tokens_per_seq=4, - ), - ), - Cell( - why="qwen4_exp nvfp4 pool mixed step P=1: a 20-row context ahead of two " - "1-row sequences over histories 31 and 64", - spec=dict( - _QWEN4_NVFP4, - cached=(0, 31, 64), - new=(20, 1, 1), - num_contexts=1, - predicted_tokens_per_seq=1, - ), - ), - Cell( - why="qwen4_exp nvfp4 pool mixed step P=4: a 20-row context ahead of two 4-row blocks", - spec=dict( - _QWEN4_NVFP4, - cached=(0, 31, 64), - new=(20, 4, 4), - num_contexts=1, - predicted_tokens_per_seq=4, - ), - ), ) note = """ diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/bmm_out.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/bmm_out.py index 73265f70118e..b7aaea3d066e 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/bmm_out.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/bmm_out.py @@ -48,18 +48,6 @@ class BmmOut(OpWrapper): "transposed-view path the targets happen to use", spec=dict(batch=128, m=64, k=128, n=512, dtype=torch.bfloat16, transposed_out=False), ), - Cell( - why="deepseek default at attention TP8, decode absorb: 16 heads per rank", - spec=dict(batch=16, m=1, k=128, n=512, dtype=torch.bfloat16, transposed_out=True), - ), - Cell( - why="deepseek default at attention TP4, decode expand: 32 heads per rank", - spec=dict(batch=32, m=1, k=512, n=128, dtype=torch.bfloat16, transposed_out=True), - ), - Cell( - why="deepseek default at attention TP2, prefill-height absorb: 64 heads per rank", - spec=dict(batch=64, m=2048, k=128, n=512, dtype=torch.bfloat16, transposed_out=True), - ), ) note = """ diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/cublas_mm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/cublas_mm.py index eccf450ef885..386ef44b20aa 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/cublas_mm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/cublas_mm.py @@ -20,61 +20,16 @@ class CublasMm(OpWrapper): ARCHS = frozenset({Arch.SM_100, Arch.SM_103}) - # Both targets use this for every dense projection, so the cells are the - # three argument shapes that differ in kind rather than in size: fused bias - # (gpt-oss carries one on qkv, o and the router), no bias (deepseek's + # Both targets use this for every dense projection. Their cells.py carry + # the three argument shapes that differ in kind rather than in size: fused + # bias (gpt-oss carries one on qkv, o and the router), no bias (deepseek's # projections), and an out_dtype that is not the input's (deepseek's router # asks for fp32 logits). CELLS: tuple[Cell, ...] = ( - Cell( - why="gpt-oss qkv with fused bias: hidden 2880 -> (64 + 8 + 8) * 64", - spec=dict(m=16, k=2880, n=5120, dtype=torch.bfloat16, bias=True, out_dtype=None), - ), - Cell( - why="deepseek q_a, no bias: hidden 7168 -> q_lora 1536", - spec=dict(m=16, k=7168, n=1536, dtype=torch.bfloat16, bias=False, out_dtype=None), - ), - Cell( - why="deepseek router: bf16 in, fp32 logits out, 256 experts", - spec=dict( - m=16, k=7168, n=256, dtype=torch.bfloat16, bias=False, out_dtype=torch.float32 - ), - ), Cell( why="decode-shaped: one token, the shape every projection sees per step", spec=dict(m=1, k=7168, n=1536, dtype=torch.bfloat16, bias=False, out_dtype=None), ), - Cell( - why="deepseek default q_b at attention TP4: q_lora 1536 -> 32 heads x 192", - spec=dict(m=16, k=1536, n=6144, dtype=torch.bfloat16, bias=False, out_dtype=None), - ), - Cell( - why="deepseek default kv_b at attention TP8: kv_lora 512 -> 16 heads x 256", - spec=dict(m=64, k=512, n=4096, dtype=torch.bfloat16, bias=False, out_dtype=None), - ), - Cell( - why="deepseek default o_proj at attention TP8: 16 heads x 128 -> hidden 7168", - spec=dict(m=16, k=2048, n=7168, dtype=torch.bfloat16, bias=False, out_dtype=None), - ), - # The MTP layer's own GEMMs. Its eh_proj is replicated on every rank; its - # shared expert is bf16 (the NVFP4 export skips layer 61), so it runs here - # rather than in nvfp4_gemm, split over attention TP like the trunk's. - Cell( - why="deepseek MTP eh_proj: [enorm(e) | hnorm(h)] 2 x 7168 -> hidden 7168", - spec=dict(m=16, k=14336, n=7168, dtype=torch.bfloat16, bias=False, out_dtype=None), - ), - Cell( - why="deepseek MTP shared expert gate_up, whole under attention DP: 7168 -> 2 x 2048", - spec=dict(m=16, k=7168, n=4096, dtype=torch.bfloat16, bias=False, out_dtype=None), - ), - Cell( - why="deepseek MTP shared expert gate_up at attention TP8: 7168 -> 2 x 256", - spec=dict(m=16, k=7168, n=512, dtype=torch.bfloat16, bias=False, out_dtype=None), - ), - Cell( - why="deepseek MTP shared expert down at attention TP8: 256 -> hidden 7168", - spec=dict(m=16, k=256, n=7168, dtype=torch.bfloat16, bias=False, out_dtype=None), - ), ) note = """ diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/nvfp4_gemm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/nvfp4_gemm.py index c4453f0b55f3..9cf942740d50 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/nvfp4_gemm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/nvfp4_gemm.py @@ -73,14 +73,6 @@ class Nvfp4Gemm(OpWrapper): why="an M that is not a multiple of the swizzle's 128-row tile", spec=dict(m=37, k=7168, n=4096, bias=False), ), - Cell( - why="deepseek default at attention TP8, shared expert gate+up: 2 x 2048/8 rows", - spec=dict(m=16, k=7168, n=512, bias=False), - ), - Cell( - why="deepseek default at attention TP8, shared expert down: K = 2048/8", - spec=dict(m=16, k=256, n=7168, bias=False), - ), ) note = """ diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/moe/fp4_block_scale_moe_runner.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/moe/fp4_block_scale_moe_runner.py index fa213428b80e..a546e1c787c9 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/moe/fp4_block_scale_moe_runner.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/moe/fp4_block_scale_moe_runner.py @@ -115,92 +115,23 @@ class Fp4BlockScaleMoeRunner(OpWrapper): # DeepSeek-R1-0528's routed geometry. The deepseek default target makes one # call per rank over its moe_ep window at its moe_tp slice of the # intermediate (dep4 with moe ep4: a 64-wide window at the full 2048; tp8: - # all 256 experts at 256), so the window and intermediate-slice cells are - # the shapes it runs and the full stack is the single-rank one. - _R1 = dict(hidden=7168, intermediate=2048, num_experts=256, top_k=8) - - # qwen4_exp's (Qwen3.8-Flash-Next) routed geometry: many small experts, a - # top_k past 8, an intermediate width of 640, all 512 experts local. Its - # router renormalizes the top-10 softmax weights to sum to 1 before the - # call (`routing="renormalized"`, a test axis rather than an op argument; - # the R1 cells draw uniform [0, 1) weights). Nothing is padded: - # H % 256, 2I % 128, H/16 % 4 and I/16 % 4 are all zero. - _QWEN4_EXP = dict( - hidden=2560, - intermediate=640, - num_experts=512, - top_k=10, - local_expert_offset=0, - local_num_experts=512, - routing="renormalized", - ) + # all 256 experts at 256), so the window cells here and the + # intermediate-slice ones in its cells.py are the shapes it runs, and the + # full stack is the single-rank one. + R1 = dict(hidden=7168, intermediate=2048, num_experts=256, top_k=8) CELLS: tuple[Cell, ...] = ( Cell( why="dep4's own call: the 64-wide expert window at offset 0, R1 geometry", - spec=dict(tokens=64, local_expert_offset=0, local_num_experts=64, **_R1), + spec=dict(tokens=64, local_expert_offset=0, local_num_experts=64, **R1), ), Cell( why="a window at a non-zero offset -- the ids this rank must answer for move with it", - spec=dict(tokens=64, local_expert_offset=192, local_num_experts=64, **_R1), + spec=dict(tokens=64, local_expert_offset=192, local_num_experts=64, **R1), ), Cell( why="the single-rank shape: all 256 experts local, and one token, the decode row count", - spec=dict(tokens=1, local_expert_offset=0, local_num_experts=256, **_R1), - ), - Cell( - why="deepseek default at tp2: all 256 experts, intermediate split 2 ways (1024)", - spec=dict( - tokens=64, - local_expert_offset=0, - local_num_experts=256, - **dict(_R1, intermediate=1024), - ), - ), - Cell( - why="deepseek default at tp4: all 256 experts, intermediate split 4 ways (512)", - spec=dict( - tokens=64, - local_expert_offset=0, - local_num_experts=256, - **dict(_R1, intermediate=512), - ), - ), - Cell( - why="deepseek default at tp8: all 256 experts, intermediate split 8 ways (256)", - spec=dict( - tokens=64, - local_expert_offset=0, - local_num_experts=256, - **dict(_R1, intermediate=256), - ), - ), - Cell( - why="deepseek default at ep8 (tp8 + ep8 and dep8): the last 32-wide window", - spec=dict(tokens=64, local_expert_offset=224, local_num_experts=32, **_R1), - ), - Cell( - why="deepseek default at dep16: a 16-wide window", - spec=dict(tokens=64, local_expert_offset=112, local_num_experts=16, **_R1), - ), - Cell( - why="qwen4_exp decode: one token over 512 experts, top-10 renormalized, H 2560 / I 640", - spec=dict(tokens=1, **_QWEN4_EXP), - ), - Cell( - why="qwen4_exp at 4 rows: the last count the pre-routed path's static-block " - "permutation kernel takes", - spec=dict(tokens=4, **_QWEN4_EXP), - ), - Cell( - why="qwen4_exp at 5 rows: the first count on the single-cluster permutation " - "kernel (the dynamic-block one stops at 256 experts)", - spec=dict(tokens=5, **_QWEN4_EXP), - ), - Cell( - why="qwen4_exp prefill: 8192 rows, the single-cluster kernel's cap and the top " - "of the certified column", - spec=dict(tokens=8192, **_QWEN4_EXP), + spec=dict(tokens=1, local_expert_offset=0, local_num_experts=256, **R1), ), ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/moe/mxe4m3_mxe2m1_block_scale_moe_runner.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/moe/mxe4m3_mxe2m1_block_scale_moe_runner.py index 8ddb8aa512e9..d1024dab23ac 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/moe/mxe4m3_mxe2m1_block_scale_moe_runner.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/moe/mxe4m3_mxe2m1_block_scale_moe_runner.py @@ -170,22 +170,12 @@ class Mxe4m3Mxe2m1BlockScaleMoeRunner(OpWrapper): # clamped SwiGlu constants the checkpoint carries (alpha 1.702, beta 1.0, # clamp limit 7.0). The default target at tp1 holds every expert whole; on # more ranks it holds an expert window (moe_ep) and an intermediate slice - # (moe_tp), which is what the cells after the first three cover. `intermediate` is + # (moe_tp), which is what its cells.py covers. `intermediate` is # the slice's real width and `intermediate_pad` the width every rank # declares (see the default target's weights.py). - _GPT_OSS = dict(hidden=2880, intermediate=2880, num_experts=128, top_k=4) + GPT_OSS = dict(hidden=2880, intermediate=2880, num_experts=128, top_k=4) CELLS: tuple[Cell, ...] = ( - Cell( - why="the gpt-oss target's own call: router entry point, all 128 experts local", - spec=dict( - tokens=64, - routed_by_kernel=True, - local_expert_offset=0, - local_num_experts=128, - **_GPT_OSS, - ), - ), Cell( why="one token -- the decode row count, where the routing has a single row to place", spec=dict( @@ -193,7 +183,7 @@ class Mxe4m3Mxe2m1BlockScaleMoeRunner(OpWrapper): routed_by_kernel=True, local_expert_offset=0, local_num_experts=128, - **_GPT_OSS, + **GPT_OSS, ), ), Cell( @@ -203,77 +193,7 @@ class Mxe4m3Mxe2m1BlockScaleMoeRunner(OpWrapper): routed_by_kernel=False, local_expert_offset=32, local_num_experts=32, - **_GPT_OSS, - ), - ), - Cell( - why="default ep4: the kernel routes all 128 ids and runs this rank's window of 32", - spec=dict( - tokens=64, - routed_by_kernel=True, - local_expert_offset=64, - local_num_experts=32, - **_GPT_OSS, - ), - ), - Cell( - why="default ep8: the kernel routes all 128 ids and runs this rank's window of 16", - spec=dict( - tokens=64, - routed_by_kernel=True, - local_expert_offset=112, - local_num_experts=16, - **_GPT_OSS, - ), - ), - Cell( - why="default moe_tp2: 1440 intermediate columns padded to 1536", - spec=dict( - tokens=64, - routed_by_kernel=True, - local_expert_offset=0, - local_num_experts=128, - **{**_GPT_OSS, "intermediate": 1440, "intermediate_pad": 1536}, - ), - ), - Cell( - why="default moe_tp4, ranks 0-2: 768 intermediate columns, no padding", - spec=dict( - tokens=64, - routed_by_kernel=True, - local_expert_offset=0, - local_num_experts=128, - **{**_GPT_OSS, "intermediate": 768, "intermediate_pad": 768}, - ), - ), - Cell( - why="default moe_tp4, rank 3: 576 real columns inside a 768-wide operand", - spec=dict( - tokens=64, - routed_by_kernel=True, - local_expert_offset=0, - local_num_experts=128, - **{**_GPT_OSS, "intermediate": 576, "intermediate_pad": 768}, - ), - ), - Cell( - why="default moe_tp8, ranks 0-6: 384 intermediate columns, no padding", - spec=dict( - tokens=64, - routed_by_kernel=True, - local_expert_offset=0, - local_num_experts=128, - **{**_GPT_OSS, "intermediate": 384, "intermediate_pad": 384}, - ), - ), - Cell( - why="default moe_tp8, rank 7: 192 real columns inside a 384-wide operand", - spec=dict( - tokens=64, - routed_by_kernel=True, - local_expert_offset=0, - local_num_experts=128, - **{**_GPT_OSS, "intermediate": 192, "intermediate_pad": 384}, + **GPT_OSS, ), ), ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/moe/noaux_tc_op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/moe/noaux_tc_op.py index c64436347f59..c8f522acf75f 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/moe/noaux_tc_op.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/moe/noaux_tc_op.py @@ -42,23 +42,6 @@ class NoauxTcOp(OpWrapper): why="a token count that is not a multiple of any tile", spec=dict(tokens=37, experts=256, n_group=8, topk_group=4, topk=8, scale=2.5), ), - Cell( - why="the default target's expert chunk, which its attention-DP gather fills", - spec=dict(tokens=8192, experts=256, n_group=8, topk_group=4, topk=8, scale=2.5), - ), - Cell( - why="deepseek MTP router: fp32 logits, so the weights are the fp32 " - "token_final_scales fused_moe needs", - spec=dict( - tokens=256, - experts=256, - n_group=8, - topk_group=4, - topk=8, - scale=2.5, - logits_dtype=torch.float32, - ), - ), ) note = """ diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/norm/flashinfer_fused_add_rmsnorm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/norm/flashinfer_fused_add_rmsnorm.py index fad05942169f..26dade284447 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/norm/flashinfer_fused_add_rmsnorm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/norm/flashinfer_fused_add_rmsnorm.py @@ -31,19 +31,11 @@ class FlashinferFusedAddRmsnorm(OpWrapper): ARCHS = frozenset({Arch.SM_100, Arch.SM_103}) - # The two hidden sizes the shipped targets normalize over, at a decode - # shape and a prefill shape. Both targets call this op only on the + # A prefill shape here; the shipped targets' decode shapes at their hidden + # sizes are in their cells.py. Both targets call this op only on the # residual stream, so there is no third width to cover and no 3-D case: # unlike plain rmsnorm this one never runs on a per-head view. CELLS: tuple[Cell, ...] = ( - Cell( - why="decode-shaped, gpt-oss-120b hidden", - spec=dict(shape=(1, 2880), dtype=torch.bfloat16, eps=1e-5), - ), - Cell( - why="decode-shaped, deepseek-r1 hidden", - spec=dict(shape=(1, 7168), dtype=torch.bfloat16, eps=1e-6), - ), Cell( why="prefill-shaped: enough rows to cross the kernel's grid tiling", spec=dict(shape=(2048, 7168), dtype=torch.bfloat16, eps=1e-6), diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/norm/flashinfer_rmsnorm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/norm/flashinfer_rmsnorm.py index b6280682aa7a..ded4bd64b33a 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/norm/flashinfer_rmsnorm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/norm/flashinfer_rmsnorm.py @@ -27,32 +27,17 @@ class FlashinferRmsnorm(OpWrapper): ARCHS = frozenset({Arch.SM_100, Arch.SM_103}) - # The widths the shipped targets normalize over. gpt-oss reaches this op - # only on its residual stream; deepseek reaches it there and on the two - # lora ranks its MLA projects through, which is why the narrow cells are - # here at all -- a 512-wide row is one 128-bit vector per lane and takes a - # different path through the kernel than a 7168-wide one. + # The widths the shipped targets normalize over are in their cells.py. + # gpt-oss reaches this op only on its residual stream; deepseek reaches it + # there and on the two lora ranks its MLA projects through, which is why + # narrow rows are certified at all -- a 512-wide row is one 128-bit vector + # per lane and takes a different path through the kernel than a 7168-wide + # one. CELLS: tuple[Cell, ...] = ( - Cell( - why="decode-shaped, gpt-oss-120b hidden 2880", - spec=dict(shape=(1, 2880), dtype=torch.bfloat16, eps=1e-5), - ), - Cell( - why="decode-shaped, deepseek-r1 hidden 7168", - spec=dict(shape=(1, 7168), dtype=torch.bfloat16, eps=1e-6), - ), Cell( why="prefill-shaped: enough rows to cross the kernel's grid tiling", spec=dict(shape=(2048, 7168), dtype=torch.bfloat16, eps=1e-6), ), - Cell( - why="deepseek q_lora rank 1536, the q_norm this op is called on", - spec=dict(shape=(16, 1536), dtype=torch.bfloat16, eps=1e-6), - ), - Cell( - why="deepseek kv_lora rank 512, the narrowest row any target asks for", - spec=dict(shape=(16, 512), dtype=torch.bfloat16, eps=1e-6), - ), ) note = """ diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/quantization/fp4_quantize.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/quantization/fp4_quantize.py index 160e7077453b..cf058a2c25bd 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/quantization/fp4_quantize.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/quantization/fp4_quantize.py @@ -76,22 +76,6 @@ class Fp4Quantize(OpWrapper): why="a row count that is not a multiple of the swizzle's 128-row tile", spec=dict(shape=(37, 7168), dtype=torch.bfloat16, vec=16, swizzled=True), ), - Cell( - why="deepseek default at attention TP8: the shared expert's down input, 2048/8 wide", - spec=dict(shape=(16, 256), dtype=torch.bfloat16, vec=16, swizzled=True), - ), - Cell( - why="qwen3.8 decode: 1 token, hidden 2560, linear -- its routed-expert input", - spec=dict(shape=(1, 2560), dtype=torch.bfloat16, vec=16, swizzled=False), - ), - Cell( - why="qwen3.8, 7 rows at hidden 2560, linear: several rows, no row padding", - spec=dict(shape=(7, 2560), dtype=torch.bfloat16, vec=16, swizzled=False), - ), - Cell( - why="qwen3.8 prefill at max_num_tokens 8192, hidden 2560, linear", - spec=dict(shape=(8192, 2560), dtype=torch.bfloat16, vec=16, swizzled=False), - ), ) note = """ diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/quantization/mxfp8_quantize.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/quantization/mxfp8_quantize.py index 0e825c497a4f..efa652e51fdc 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/quantization/mxfp8_quantize.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/quantization/mxfp8_quantize.py @@ -38,10 +38,6 @@ class Mxfp8Quantize(OpWrapper): # layouts are separate code paths in the kernel, not because a target uses # it today. CELLS: tuple[Cell, ...] = ( - Cell( - why="gpt-oss hidden 2880, linear scales, aligned to 512 -> padded 3072", - spec=dict(shape=(16, 2880), dtype=torch.bfloat16, swizzled=False, alignment=512), - ), Cell( why="decode-shaped: one token through the same padding", spec=dict(shape=(1, 2880), dtype=torch.bfloat16, swizzled=False, alignment=512), diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/__init__.py index 9da2ae758216..69dd53ea5e1b 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/__init__.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/__init__.py @@ -5,4 +5,13 @@ ``models//`` corresponds one-to-one with the built-in zoo's ``modeling_.py``: the same architecture, split into one self-contained codebase per deployment target instead of one class serving every -checkpoint.""" +checkpoint. + +Importing this package imports every family's ``routing`` module, and each one +registers its ``route`` through ``@routes``. A new family adds one line here. +""" + +from .deepseek_v3 import routing as deepseek_v3_routing # noqa: F401 +from .glm_moe_dsa import routing as glm_moe_dsa_routing # noqa: F401 +from .gpt_oss import routing as gpt_oss_routing # noqa: F401 +from .qwen4_exp import routing as qwen4_exp_routing # noqa: F401 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_10x__default/cells.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_10x__default/cells.py new file mode 100644 index 000000000000..aaaf81526be5 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_10x__default/cells.py @@ -0,0 +1,323 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The cells of the catalog entries the DeepSeek-R1-0528 NVFP4 default target calls. + +Each entry's own cells are the op's: tile edges, row counts, the dtypes its +kernel takes. These are the shapes this target calls it with. The entry's +test drives both through `certified_cells`; nothing in the forward reads +them, so a change here reruns the entry tests and nothing of the target's. +""" + +from __future__ import annotations + +import torch + +from ....catalog._op import Cell, OpWrapper +from ....catalog.activation.flashinfer_silu_and_mul import FlashinferSiluAndMul +from ....catalog.attention.mla_rope_append_paged_kv_assign_q import MlaRopeAppendPagedKvAssignQ +from ....catalog.attention.mla_rope_generation import Q_SCALING_R1, MlaRopeGeneration +from ....catalog.attention.thop_attention import ( + MLA_CONTEXT, + MLA_GENERATION, + QM_FP8_KV_CACHE, + ThopAttention, +) +from ....catalog.gemm.bmm_out import BmmOut +from ....catalog.gemm.cublas_mm import CublasMm +from ....catalog.gemm.nvfp4_gemm import Nvfp4Gemm +from ....catalog.moe.fp4_block_scale_moe_runner import Fp4BlockScaleMoeRunner +from ....catalog.moe.noaux_tc_op import NoauxTcOp +from ....catalog.norm.flashinfer_fused_add_rmsnorm import FlashinferFusedAddRmsnorm +from ....catalog.norm.flashinfer_rmsnorm import FlashinferRmsnorm +from ....catalog.quantization.fp4_quantize import Fp4Quantize + +CELLS: dict[type[OpWrapper], tuple[Cell, ...]] = { + FlashinferSiluAndMul: ( + Cell( + why="deepseek default at attention TP8: the shared expert's gate+up, 2 x 2048/8", + spec=dict(shape=(16, 512), dtype=torch.bfloat16), + ), + ), + MlaRopeAppendPagedKvAssignQ: ( + Cell( + why="deepseek default at attention TP8: 16 heads per rank, production context", + spec=dict( + cached=[0, 32, 100], + new=[37, 64, 300], + tokens_per_block=32, + **dict(MlaRopeAppendPagedKvAssignQ.R1, head_num=16), + ), + ), + ), + MlaRopeGeneration: ( + Cell( + why="deepseek default at attention TP8: 16 heads per rank, production decode", + spec=dict( + predicted_tokens_per_seq=1, + tokens_per_block=32, + q_scaling=Q_SCALING_R1, + write_scale=None, + read_scale=None, + **dict(MlaRopeGeneration.R1, num_heads=16), + ), + ), + Cell( + why="deepseek default MTP3 at attention TP8: P=4 (draft length 3 + 1) at 16 heads", + spec=dict( + predicted_tokens_per_seq=4, + tokens_per_block=32, + q_scaling=Q_SCALING_R1, + write_scale=None, + read_scale=None, + **dict(MlaRopeGeneration.R1, num_heads=16), + ), + ), + ), + ThopAttention: ( + Cell( + why="deepseek MLA context: explicit k/v, qk 576 -> v 128, fp8 pool", + spec=dict( + input_type=MLA_CONTEXT, + num_heads=128, + num_kv_heads=128, + head_size=576, + v_head_dim=128, + sinks=False, + window=None, + quant_mode=QM_FP8_KV_CACHE, + ), + ), + Cell( + why="deepseek MLA generation: one kv head, latent 576 -> 512, fp8 pool", + spec=dict( + input_type=MLA_GENERATION, + num_heads=128, + num_kv_heads=1, + head_size=576, + v_head_dim=512, + sinks=False, + window=None, + quant_mode=QM_FP8_KV_CACHE, + ), + ), + # The per-rank head counts the deepseek default target runs under + # attention TP: the same two MLA code paths as the cells above, at the + # 128 heads split 2, 4 or 8 ways. + Cell( + why="deepseek default at attention TP8: MLA context, 16 heads per rank", + spec=dict( + input_type=MLA_CONTEXT, + num_heads=16, + num_kv_heads=16, + head_size=576, + v_head_dim=128, + sinks=False, + window=None, + quant_mode=QM_FP8_KV_CACHE, + ), + ), + Cell( + why="deepseek default at attention TP2: MLA generation, 64 heads per rank", + spec=dict( + input_type=MLA_GENERATION, + num_heads=64, + num_kv_heads=1, + head_size=576, + v_head_dim=512, + sinks=False, + window=None, + quant_mode=QM_FP8_KV_CACHE, + ), + ), + Cell( + why="deepseek default at attention TP8: MLA generation, 16 heads per rank", + spec=dict( + input_type=MLA_GENERATION, + num_heads=16, + num_kv_heads=1, + head_size=576, + v_head_dim=512, + sinks=False, + window=None, + quant_mode=QM_FP8_KV_CACHE, + ), + ), + # MTP3 under attention TP: the verify step's taller query block (draft + # length 3 + 1 rows per sequence) at a split head count. The P > 1 mask + # itself is pinned at 128 heads by the MTP cases in the test. + Cell( + why="deepseek default MTP3 at attention TP8: MLA generation, 16 heads, P=4", + spec=dict( + input_type=MLA_GENERATION, + num_heads=16, + num_kv_heads=1, + head_size=576, + v_head_dim=512, + sinks=False, + window=None, + quant_mode=QM_FP8_KV_CACHE, + predicted_tokens_per_seq=4, + ), + ), + ), + BmmOut: ( + Cell( + why="deepseek default at attention TP8, decode absorb: 16 heads per rank", + spec=dict(batch=16, m=1, k=128, n=512, dtype=torch.bfloat16, transposed_out=True), + ), + Cell( + why="deepseek default at attention TP4, decode expand: 32 heads per rank", + spec=dict(batch=32, m=1, k=512, n=128, dtype=torch.bfloat16, transposed_out=True), + ), + Cell( + why="deepseek default at attention TP2, prefill-height absorb: 64 heads per rank", + spec=dict(batch=64, m=2048, k=128, n=512, dtype=torch.bfloat16, transposed_out=True), + ), + ), + CublasMm: ( + Cell( + why="deepseek q_a, no bias: hidden 7168 -> q_lora 1536", + spec=dict(m=16, k=7168, n=1536, dtype=torch.bfloat16, bias=False, out_dtype=None), + ), + Cell( + why="deepseek router: bf16 in, fp32 logits out, 256 experts", + spec=dict( + m=16, k=7168, n=256, dtype=torch.bfloat16, bias=False, out_dtype=torch.float32 + ), + ), + Cell( + why="deepseek default q_b at attention TP4: q_lora 1536 -> 32 heads x 192", + spec=dict(m=16, k=1536, n=6144, dtype=torch.bfloat16, bias=False, out_dtype=None), + ), + Cell( + why="deepseek default kv_b at attention TP8: kv_lora 512 -> 16 heads x 256", + spec=dict(m=64, k=512, n=4096, dtype=torch.bfloat16, bias=False, out_dtype=None), + ), + Cell( + why="deepseek default o_proj at attention TP8: 16 heads x 128 -> hidden 7168", + spec=dict(m=16, k=2048, n=7168, dtype=torch.bfloat16, bias=False, out_dtype=None), + ), + # The MTP layer's own GEMMs. Its eh_proj is replicated on every rank; its + # shared expert is bf16 (the NVFP4 export skips layer 61), so it runs + # through cublas_mm rather than nvfp4_gemm, split over attention TP like + # the trunk's. + Cell( + why="deepseek MTP eh_proj: [enorm(e) | hnorm(h)] 2 x 7168 -> hidden 7168", + spec=dict(m=16, k=14336, n=7168, dtype=torch.bfloat16, bias=False, out_dtype=None), + ), + Cell( + why="deepseek MTP shared expert gate_up, whole under attention DP: 7168 -> 2 x 2048", + spec=dict(m=16, k=7168, n=4096, dtype=torch.bfloat16, bias=False, out_dtype=None), + ), + Cell( + why="deepseek MTP shared expert gate_up at attention TP8: 7168 -> 2 x 256", + spec=dict(m=16, k=7168, n=512, dtype=torch.bfloat16, bias=False, out_dtype=None), + ), + Cell( + why="deepseek MTP shared expert down at attention TP8: 256 -> hidden 7168", + spec=dict(m=16, k=256, n=7168, dtype=torch.bfloat16, bias=False, out_dtype=None), + ), + ), + Nvfp4Gemm: ( + Cell( + why="deepseek default at attention TP8, shared expert gate+up: 2 x 2048/8 rows", + spec=dict(m=16, k=7168, n=512, bias=False), + ), + Cell( + why="deepseek default at attention TP8, shared expert down: K = 2048/8", + spec=dict(m=16, k=256, n=7168, bias=False), + ), + ), + Fp4BlockScaleMoeRunner: ( + Cell( + why="deepseek default at tp2: all 256 experts, intermediate split 2 ways (1024)", + spec=dict( + tokens=64, + local_expert_offset=0, + local_num_experts=256, + **dict(Fp4BlockScaleMoeRunner.R1, intermediate=1024), + ), + ), + Cell( + why="deepseek default at tp4: all 256 experts, intermediate split 4 ways (512)", + spec=dict( + tokens=64, + local_expert_offset=0, + local_num_experts=256, + **dict(Fp4BlockScaleMoeRunner.R1, intermediate=512), + ), + ), + Cell( + why="deepseek default at tp8: all 256 experts, intermediate split 8 ways (256)", + spec=dict( + tokens=64, + local_expert_offset=0, + local_num_experts=256, + **dict(Fp4BlockScaleMoeRunner.R1, intermediate=256), + ), + ), + Cell( + why="deepseek default at ep8 (tp8 + ep8 and dep8): the last 32-wide window", + spec=dict( + tokens=64, + local_expert_offset=224, + local_num_experts=32, + **Fp4BlockScaleMoeRunner.R1, + ), + ), + Cell( + why="deepseek default at dep16: a 16-wide window", + spec=dict( + tokens=64, + local_expert_offset=112, + local_num_experts=16, + **Fp4BlockScaleMoeRunner.R1, + ), + ), + ), + NoauxTcOp: ( + Cell( + why="the default target's expert chunk, which its attention-DP gather fills", + spec=dict(tokens=8192, experts=256, n_group=8, topk_group=4, topk=8, scale=2.5), + ), + Cell( + why="deepseek MTP router: fp32 logits, so the weights are the fp32 " + "token_final_scales fused_moe needs", + spec=dict( + tokens=256, + experts=256, + n_group=8, + topk_group=4, + topk=8, + scale=2.5, + logits_dtype=torch.float32, + ), + ), + ), + FlashinferFusedAddRmsnorm: ( + Cell( + why="decode-shaped, deepseek-r1 hidden", + spec=dict(shape=(1, 7168), dtype=torch.bfloat16, eps=1e-6), + ), + ), + FlashinferRmsnorm: ( + Cell( + why="decode-shaped, deepseek-r1 hidden 7168", + spec=dict(shape=(1, 7168), dtype=torch.bfloat16, eps=1e-6), + ), + Cell( + why="deepseek q_lora rank 1536, the q_norm this op is called on", + spec=dict(shape=(16, 1536), dtype=torch.bfloat16, eps=1e-6), + ), + Cell( + why="deepseek kv_lora rank 512, the narrowest row any target asks for", + spec=dict(shape=(16, 512), dtype=torch.bfloat16, eps=1e-6), + ), + ), + Fp4Quantize: ( + Cell( + why="deepseek default at attention TP8: the shared expert's down input, 2048/8 wide", + spec=dict(shape=(16, 256), dtype=torch.bfloat16, vec=16, swizzled=True), + ), + ), +} diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/routing.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/routing.py index 6f4dd9fd6c2f..9adab1845fd1 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/routing.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/routing.py @@ -16,7 +16,7 @@ from tensorrt_llm.llmapi.llm_args import MTPDecodingConfig from tensorrt_llm.quantization.mode import QuantAlgo -from ..._router_index import NoTarget, tp_reaches_across_nodes +from ..._router_index import NoTarget, routes, tp_reaches_across_nodes # The GPUs the target is certified on: sm_100 (B200, GB200) and sm_103 (B300, # GB300). They run the same kernels. @@ -74,6 +74,7 @@ from tensorrt_llm.llmapi.llm_args import TorchLlmArgs +@routes("DeepseekV3ForCausalLM") def route( pretrained_config: Any, mapping: Any, diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/routing.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/routing.py index 7feda0cc4e16..e75ac5b62223 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/routing.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/routing.py @@ -11,24 +11,31 @@ from __future__ import annotations -from typing import Any, Optional, Tuple +from typing import TYPE_CHECKING, Any, Optional, Tuple from tensorrt_llm.llmapi.llm_args import MTPDecodingConfig from tensorrt_llm.quantization.mode import QuantAlgo -from ..._router_index import NoTarget +from ..._router_index import NoTarget, routes + +if TYPE_CHECKING: + from tensorrt_llm._torch.model_config import ModelConfig + from tensorrt_llm.llmapi.llm_args import TorchLlmArgs # (layers, hidden, routed experts, q_lora_rank). Shape-sniffing also matches # fine-tunes of the same shape; the accuracy gate pins identity. _SHAPE = (78, 6144, 256, 2048) +@routes("GlmMoeDsaForCausalLM") def route( pretrained_config: Any, mapping: Any, compute_capability: Tuple[int, int], spec_config: Optional[Any], kv_cache_quant_algo: Optional[Any], + model_config: Optional[ModelConfig] = None, + llm_args: Optional[TorchLlmArgs] = None, ) -> type: """Return the target class for this deployment, or raise ``NoTarget`` saying why none fits. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/gpt_oss/gpt_oss_120b__sm_10x__default/cells.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/gpt_oss/gpt_oss_120b__sm_10x__default/cells.py new file mode 100644 index 000000000000..8d166c14dab1 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/gpt_oss/gpt_oss_120b__sm_10x__default/cells.py @@ -0,0 +1,209 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The cells of the catalog entries the gpt-oss-120b default target calls. + +Each entry's own cells are the op's: tile edges, row counts, the dtypes its +kernel takes. These are the shapes this target calls it with. The entry's +test drives both through `certified_cells`; nothing in the forward reads +them, so a change here reruns the entry tests and nothing of the target's. +""" + +from __future__ import annotations + +import torch + +from ....catalog._op import Cell, OpWrapper +from ....catalog.attention.fused_qk_norm_rope import FusedQkNormRope +from ....catalog.attention.thop_attention import STANDARD, ThopAttention +from ....catalog.gemm.cublas_mm import CublasMm +from ....catalog.moe.mxe4m3_mxe2m1_block_scale_moe_runner import Mxe4m3Mxe2m1BlockScaleMoeRunner +from ....catalog.norm.flashinfer_fused_add_rmsnorm import FlashinferFusedAddRmsnorm +from ....catalog.norm.flashinfer_rmsnorm import FlashinferRmsnorm +from ....catalog.quantization.mxfp8_quantize import Mxfp8Quantize + +CELLS: dict[type[OpWrapper], tuple[Cell, ...]] = { + FusedQkNormRope: ( + Cell( + why="default target at attention TP8: 8q/1k/1v heads per rank", + spec=dict( + tokens=64, + **{ + **FusedQkNormRope.GPT_OSS, + "num_heads_q": 8, + "num_heads_k": 1, + "num_heads_v": 1, + }, + ), + ), + ), + ThopAttention: ( + Cell( + why="gpt-oss sliding layer: GQA 64q/8kv d64, sinks, window 128", + spec=dict( + input_type=STANDARD, + num_heads=64, + num_kv_heads=8, + head_size=64, + v_head_dim=64, + sinks=True, + window=128, + quant_mode=0, + ), + ), + Cell( + why="gpt-oss full layer: the same heads with no window", + spec=dict( + input_type=STANDARD, + num_heads=64, + num_kv_heads=8, + head_size=64, + v_head_dim=64, + sinks=True, + window=None, + quant_mode=0, + ), + ), + # The per-rank head counts the gpt-oss default target runs under + # attention TP: the same code path as the two cells above, at the + # head counts a rank holds after the split. + Cell( + why="gpt-oss default at attention TP8: 8q/1kv per rank, 128-token window", + spec=dict( + input_type=STANDARD, + num_heads=8, + num_kv_heads=1, + head_size=64, + v_head_dim=64, + sinks=True, + window=128, + quant_mode=0, + ), + ), + ), + CublasMm: ( + Cell( + why="gpt-oss qkv with fused bias: hidden 2880 -> (64 + 8 + 8) * 64", + spec=dict(m=16, k=2880, n=5120, dtype=torch.bfloat16, bias=True, out_dtype=None), + ), + ), + Mxe4m3Mxe2m1BlockScaleMoeRunner: ( + Cell( + why="the gpt-oss target's own call: router entry point, all 128 experts local", + spec=dict( + tokens=64, + routed_by_kernel=True, + local_expert_offset=0, + local_num_experts=128, + **Mxe4m3Mxe2m1BlockScaleMoeRunner.GPT_OSS, + ), + ), + Cell( + why="default ep4: the kernel routes all 128 ids and runs this rank's window of 32", + spec=dict( + tokens=64, + routed_by_kernel=True, + local_expert_offset=64, + local_num_experts=32, + **Mxe4m3Mxe2m1BlockScaleMoeRunner.GPT_OSS, + ), + ), + Cell( + why="default ep8: the kernel routes all 128 ids and runs this rank's window of 16", + spec=dict( + tokens=64, + routed_by_kernel=True, + local_expert_offset=112, + local_num_experts=16, + **Mxe4m3Mxe2m1BlockScaleMoeRunner.GPT_OSS, + ), + ), + Cell( + why="default moe_tp2: 1440 intermediate columns padded to 1536", + spec=dict( + tokens=64, + routed_by_kernel=True, + local_expert_offset=0, + local_num_experts=128, + **{ + **Mxe4m3Mxe2m1BlockScaleMoeRunner.GPT_OSS, + "intermediate": 1440, + "intermediate_pad": 1536, + }, + ), + ), + Cell( + why="default moe_tp4, ranks 0-2: 768 intermediate columns, no padding", + spec=dict( + tokens=64, + routed_by_kernel=True, + local_expert_offset=0, + local_num_experts=128, + **{ + **Mxe4m3Mxe2m1BlockScaleMoeRunner.GPT_OSS, + "intermediate": 768, + "intermediate_pad": 768, + }, + ), + ), + Cell( + why="default moe_tp4, rank 3: 576 real columns inside a 768-wide operand", + spec=dict( + tokens=64, + routed_by_kernel=True, + local_expert_offset=0, + local_num_experts=128, + **{ + **Mxe4m3Mxe2m1BlockScaleMoeRunner.GPT_OSS, + "intermediate": 576, + "intermediate_pad": 768, + }, + ), + ), + Cell( + why="default moe_tp8, ranks 0-6: 384 intermediate columns, no padding", + spec=dict( + tokens=64, + routed_by_kernel=True, + local_expert_offset=0, + local_num_experts=128, + **{ + **Mxe4m3Mxe2m1BlockScaleMoeRunner.GPT_OSS, + "intermediate": 384, + "intermediate_pad": 384, + }, + ), + ), + Cell( + why="default moe_tp8, rank 7: 192 real columns inside a 384-wide operand", + spec=dict( + tokens=64, + routed_by_kernel=True, + local_expert_offset=0, + local_num_experts=128, + **{ + **Mxe4m3Mxe2m1BlockScaleMoeRunner.GPT_OSS, + "intermediate": 192, + "intermediate_pad": 384, + }, + ), + ), + ), + FlashinferFusedAddRmsnorm: ( + Cell( + why="decode-shaped, gpt-oss-120b hidden", + spec=dict(shape=(1, 2880), dtype=torch.bfloat16, eps=1e-5), + ), + ), + FlashinferRmsnorm: ( + Cell( + why="decode-shaped, gpt-oss-120b hidden 2880", + spec=dict(shape=(1, 2880), dtype=torch.bfloat16, eps=1e-5), + ), + ), + Mxfp8Quantize: ( + Cell( + why="gpt-oss hidden 2880, linear scales, aligned to 512 -> padded 3072", + spec=dict(shape=(16, 2880), dtype=torch.bfloat16, swizzled=False, alignment=512), + ), + ), +} diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/gpt_oss/routing.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/gpt_oss/routing.py index 7bd51d78e37a..20098df47baa 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/gpt_oss/routing.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/gpt_oss/routing.py @@ -12,7 +12,7 @@ from typing import TYPE_CHECKING, Any, Optional, Tuple -from ..._router_index import NoTarget, tp_reaches_across_nodes +from ..._router_index import NoTarget, routes, tp_reaches_across_nodes # The GPUs the target is certified on: sm_100 (B200, GB200) and sm_103 (B300, # GB300). They run the same kernels. @@ -47,6 +47,7 @@ from tensorrt_llm.llmapi.llm_args import TorchLlmArgs +@routes("GptOssForCausalLM") def route( pretrained_config: Any, mapping: Any, diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/cells.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/cells.py new file mode 100644 index 000000000000..b18424613a1e --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/cells.py @@ -0,0 +1,556 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The cells of the catalog entries the Qwen3.8-Flash-Next NVFP4 sm_103 tp1 target calls. + +Each entry's own cells are the op's: tile edges, row counts, the dtypes its +kernel takes. These are the shapes this target calls it with. The entry's +test drives both through `certified_cells`; nothing in the forward reads +them, so a change here reruns the entry tests and nothing of the target's. +""" + +from __future__ import annotations + +import torch + +from ....catalog._op import Cell, OpWrapper +from ....catalog.attention.fused_qk_norm_rope import FusedQkNormRope +from ....catalog.attention.thop_attention import ( + QM_FP8_KV_CACHE, + QM_NVFP4_KV_CACHE, + QWEN4_DENSE, + QWEN4_VISION, + ThopAttention, +) +from ....catalog.moe.fp4_block_scale_moe_runner import Fp4BlockScaleMoeRunner +from ....catalog.quantization.fp4_quantize import Fp4Quantize + +# qwen4_exp (Qwen3.8-Flash-Next) runs Gemma QK norm and an unscaled rope +# on the first 64 dims (32 pairs), base 1e7, at two geometries. Its +# mrope_section is [11, 11, 10], passed as sections 11 / 11. Positions are +# drawn uniformly below `max_position` (a test axis, not an op argument): +# the kernel's fast sincos loses phase as `position * freq` grows, so the +# position range is part of what a cell certifies. +_QWEN_ROPE = dict( + rotary_dim=64, + base=1e7, + is_neox=True, + is_qk_norm=True, + use_gemma=True, + max_position=32768, +) +_QWEN_MROPE = dict(use_mrope=True, mrope_section1=11, mrope_section2=11, **_QWEN_ROPE) +# Main attention (trunk and MTP): 24 query heads, 2 kv, head_dim 256. +_QWEN_MAIN = dict(num_heads_q=24, num_heads_k=2, num_heads_v=2, head_dim=256, **_QWEN_MROPE) +# The QSA index plane's query rope: 4 index heads plus the one k and one v +# slot of its scratch layout, over index_head_dim 128 (== 2 * rotary_dim, +# so the partial factor is not reapplied). +_QWEN_INDEX_HEADS = dict(num_heads_q=4, num_heads_k=1, num_heads_v=1, head_dim=128) +_QWEN_INDEX_MROPE = dict(**_QWEN_INDEX_HEADS, **_QWEN_MROPE) +_QWEN_INDEX_PLAIN = dict(**_QWEN_INDEX_HEADS, **_QWEN_ROPE) + +# qwen4_exp's (Qwen3.8-Flash-Next) routed geometry: many small experts, a +# top_k past 8, an intermediate width of 640, all 512 experts local. Its +# router renormalizes the top-10 softmax weights to sum to 1 before the +# call (`routing="renormalized"`, a test axis rather than an op argument; +# the R1 cells draw uniform [0, 1) weights). Nothing is padded: +# H % 256, 2I % 128, H/16 % 4 and I/16 % 4 are all zero. +_QWEN4_EXP = dict( + hidden=2560, + intermediate=640, + num_experts=512, + top_k=10, + local_expert_offset=0, + local_num_experts=512, + routing="renormalized", +) + +#: The qwen4_exp full-attention call over an fp8-e4m3 pool at the production kv +#: scale 1.0 (fp32 [1] scale tensors, quant_mode 128): what the trunk layers and +#: the MTP draft layer pass when the engine serves an fp8 KV cache. Same chosen +#: pool sizes as `QWEN4_DENSE`. +_QWEN4_FP8 = dict(QWEN4_DENSE, quant_mode=QM_FP8_KV_CACHE, kv_scale=1.0) + +#: The same call over an NVFP4 pool: e2m1 data and a second pool of e4m3 block +#: scales, fp32 [3] unit kv scales, `[1, 2, 2]` pool pointers, bf16 output -- +#: what the trunk layers and the MTP draft layer pass when the engine serves an +#: NVFP4 KV cache. Same chosen pool sizes as `QWEN4_DENSE`. +_QWEN4_NVFP4 = dict(QWEN4_DENSE, quant_mode=QM_NVFP4_KV_CACHE, kv_scale=1.0) + +#: `QWEN4_DENSE` and `_QWEN4_FP8` for a verify step of up to 32 requests, the captured-batch +#: ceiling: one block-offset row per request and pages for 32 histories (the largest cell +#: needs 133 of the 144). +_QWEN4_DENSE_G32 = dict(QWEN4_DENSE, max_num_requests=32, num_blocks=144) +_QWEN4_FP8_G32 = dict(_QWEN4_FP8, max_num_requests=32, num_blocks=144) + + +def _spread(batch: int, longest: int) -> tuple[int, ...]: + """`batch` histories rising evenly from 1 to `longest` tokens: one-page, page-boundary + and multi-page sequences in one verify step.""" + return tuple(1 + (longest - 1) * i // (batch - 1) for i in range(batch)) + + +CELLS: dict[type[OpWrapper], tuple[Cell, ...]] = { + FusedQkNormRope: ( + Cell( + why="qwen4_exp main rope, decode: one token, Gemma norm and interleaved mRoPE " + "with three distinct position rows", + spec=dict(tokens=1, **_QWEN_MAIN), + ), + Cell( + why="qwen4_exp main rope, 19 tokens: the size its wiring and saturation tests run at", + spec=dict(tokens=19, **_QWEN_MAIN), + ), + Cell( + why="qwen4_exp main rope, 8192-token prefill: the Gemma one-ulp tail, under " + "GEMMA_PREFILL_ATOL", + spec=dict(tokens=8192, gemma_prefill_band=True, **_QWEN_MAIN), + ), + # The main rope run in place on qwen4_exp's fused input projection, whose + # row is [q | k | v | gate | index | pad] = 55 heads of 256: everything + # after k is passed as v heads, which only sets the row stride. + Cell( + why="qwen4_exp fused-row main rope, 1 token: 29 v heads only set the row stride", + spec=dict(tokens=1, **{**_QWEN_MAIN, "num_heads_v": 29}), + ), + Cell( + why="qwen4_exp fused-row main rope, 128 tokens: decode at max concurrency", + spec=dict(tokens=128, **{**_QWEN_MAIN, "num_heads_v": 29}), + ), + Cell( + why="qwen4_exp fused-row main rope, 8192 tokens: prefill, under GEMMA_PREFILL_ATOL", + spec=dict(tokens=8192, gemma_prefill_band=True, **{**_QWEN_MAIN, "num_heads_v": 29}), + ), + Cell( + why="qwen4_exp index rope (mRoPE), 1 token: one k and one v head, the " + "narrowest layout; 5 warps of an 8-warp block", + spec=dict(tokens=1, **_QWEN_INDEX_MROPE), + ), + Cell( + why="qwen4_exp index rope (mRoPE), 2 tokens: the second token's heads " + "straddle a block boundary", + spec=dict(tokens=2, **_QWEN_INDEX_MROPE), + ), + Cell( + why="qwen4_exp index rope (mRoPE), 5 tokens: the last block runs one warp", + spec=dict(tokens=5, **_QWEN_INDEX_MROPE), + ), + Cell( + why="qwen4_exp index rope (mRoPE), 17 tokens: one past a power of two", + spec=dict(tokens=17, **_QWEN_INDEX_MROPE), + ), + Cell( + why="qwen4_exp index rope (mRoPE), 19 tokens: the size its wiring and " + "saturation tests run at", + spec=dict(tokens=19, **_QWEN_INDEX_MROPE), + ), + Cell( + why="qwen4_exp index rope (mRoPE), 33 tokens: one past a power of two", + spec=dict(tokens=33, **_QWEN_INDEX_MROPE), + ), + Cell( + why="qwen4_exp index rope (mRoPE), 65 tokens: one past a power of two", + spec=dict(tokens=65, **_QWEN_INDEX_MROPE), + ), + Cell( + why="qwen4_exp index rope (mRoPE), 8192-token prefill: the index site sees the " + "whole packed batch; under GEMMA_PREFILL_ATOL", + spec=dict(tokens=8192, gemma_prefill_band=True, **_QWEN_INDEX_MROPE), + ), + Cell( + why="qwen4_exp index rope (plain), decode: the form the index plane runs " + "today, row 0 alone as a flat [T] position_ids", + spec=dict(tokens=1, **_QWEN_INDEX_PLAIN), + ), + Cell( + why="qwen4_exp index rope (plain), 19 tokens: an odd count across blocks", + spec=dict(tokens=19, **_QWEN_INDEX_PLAIN), + ), + Cell( + why="qwen4_exp index rope (plain), 8192-token prefill: under GEMMA_PREFILL_ATOL", + spec=dict(tokens=8192, gemma_prefill_band=True, **_QWEN_INDEX_PLAIN), + ), + ), + ThopAttention: ( + Cell( + why="qwen4_exp prefill: GQA 24q/2kv d256, bf16 pool, paged context; " + "a fresh 70-row sequence beside 33 rows over a 48-token cached prefix", + spec=dict( + QWEN4_DENSE, + cached=(0, 48), + new=(70, 33), + num_contexts=2, + predicted_tokens_per_seq=1, + ), + ), + Cell( + why="qwen4_exp decode P=1: one row per generation sequence over " + "page-crossing histories (70, 35)", + spec=dict( + QWEN4_DENSE, + cached=(70, 35), + new=(1, 1), + num_contexts=0, + predicted_tokens_per_seq=1, + ), + ), + Cell( + why="qwen4_exp MTP verify P=2 (max_draft_len 1): draft row t attends " + "[0, L-P+t], bottom-right causal within the block", + spec=dict( + QWEN4_DENSE, + cached=(70, 35), + new=(2, 2), + num_contexts=0, + predicted_tokens_per_seq=2, + ), + ), + Cell( + why="qwen4_exp MTP verify P=3 (max_draft_len 2): a separate decode " + "compile, P reaches the kernel as maxSeqLenQ", + spec=dict( + QWEN4_DENSE, + cached=(70, 35), + new=(3, 3), + num_contexts=0, + predicted_tokens_per_seq=3, + ), + ), + Cell( + why="qwen4_exp MTP verify P=4 (max_draft_len 3)", + spec=dict( + QWEN4_DENSE, + cached=(70, 35), + new=(4, 4), + num_contexts=0, + predicted_tokens_per_seq=4, + ), + ), + *( + Cell( + why=f"qwen4_exp MTP verify P={p} (max_draft_len {p - 1})" + + (": the target's ceiling" if p == 8 else ""), + spec=dict( + QWEN4_DENSE, + cached=(70, 35), + new=(p, p), + num_contexts=0, + predicted_tokens_per_seq=p, + ), + ) + for p in (5, 6, 7, 8) + ), + Cell( + why="qwen4_exp P=8 G=1: one draft block from row 95, across the page edge at 96", + spec=dict( + QWEN4_DENSE, + cached=(95,), + new=(8,), + num_contexts=0, + predicted_tokens_per_seq=8, + ), + ), + # The trtllm-gen launch picks its generation kernel by batch size at P 5..8, so each P + # is also driven at 32 requests and at one batch in 17..20, with histories of every + # page shape in the step. + *( + Cell( + why=f"qwen4_exp bf16 batch P={p} G=32: 32 blocks over histories 1..218, " + f"{32 * p} q rows", + spec=dict( + _QWEN4_DENSE_G32, + cached=_spread(32, 218), + new=(p,) * 32, + num_contexts=0, + predicted_tokens_per_seq=p, + ), + ) + for p in (5, 6, 7, 8) + ), + *( + Cell( + why=f"qwen4_exp bf16 batch P={p} G={g}: {g} blocks over histories 200..1, " + f"{g * p} q rows", + spec=dict( + _QWEN4_DENSE_G32, + cached=tuple(reversed(_spread(g, 200))), + new=(p,) * g, + num_contexts=0, + predicted_tokens_per_seq=p, + ), + ) + for p, g in ((5, 17), (6, 18), (7, 19), (8, 20)) + ), + Cell( + why="qwen4_exp mixed step: a fresh 20-row context ahead of two P=2 " + "draft blocks, generation rows token-major after the context rows", + spec=dict( + QWEN4_DENSE, + cached=(0, 70, 35), + new=(20, 2, 2), + num_contexts=1, + predicted_tokens_per_seq=2, + ), + ), + Cell( + why="qwen4_exp mixed step P=8: a fresh 20-row context ahead of two P=8 " + "draft blocks over histories 70 and 35", + spec=dict( + QWEN4_DENSE, + cached=(0, 70, 35), + new=(20, 8, 8), + num_contexts=1, + predicted_tokens_per_seq=8, + ), + ), + Cell( + why="qwen4_exp hybrid pool: row 47 of the 48-row two-pool mapping " + "(pool 0, layer-in-pool 11), P=2 verify", + spec=dict( + QWEN4_DENSE, + pool_layout="hybrid48", + local_layer_idx=47, + cached=(64, 25), + new=(2, 2), + num_contexts=0, + predicted_tokens_per_seq=2, + ), + ), + Cell( + why="qwen4_exp P=8 hybrid pool: row 47 of the 48-row two-pool mapping " + "(pool 0, layer-in-pool 11), P=8 verify", + spec=dict( + QWEN4_DENSE, + pool_layout="hybrid48", + local_layer_idx=47, + cached=(64, 25), + new=(8, 8), + num_contexts=0, + predicted_tokens_per_seq=8, + ), + ), + Cell( + why="qwen4_exp vision tower: no KV cache, MHA 16/16 d72, padding mask, " + "frames of 27/9/41 rows each attending only to itself", + spec=dict(QWEN4_VISION, seg_lens=(27, 9, 41)), + ), + Cell( + why="mmmu serving shape: one 5580-row frame, the 62x90 patch grid", + spec=dict(QWEN4_VISION, seg_lens=(5580,)), + ), + Cell( + why="mmmu serving shape: one 8160-row frame, the largest single-segment call", + spec=dict(QWEN4_VISION, seg_lens=(8160,)), + ), + Cell( + why="mmmu serving shape: 2204 rows in four frames", + spec=dict(QWEN4_VISION, seg_lens=(728, 728, 460, 288)), + ), + Cell( + why="mmmu serving shape: 1892 rows in five frames, the most segments per call", + spec=dict(QWEN4_VISION, seg_lens=(480, 480, 400, 352, 180)), + ), + Cell( + why="mmmu serving shape: 9008 rows in two frames", + spec=dict(QWEN4_VISION, seg_lens=(8000, 1008)), + ), + Cell( + why="mmmu serving shape: 22144 rows in two frames, the largest call; its " + "12544-row frame is the longest segment", + spec=dict(QWEN4_VISION, seg_lens=(12544, 9600)), + ), + # qwen4_exp's own geometry over an fp8 pool: the kv head count changes the compiled + # decode kernel, so the 24/4/256 cells above do not cover it. + Cell( + why="qwen4_exp fp8 pool paged-context prefill, 24/2/256: a fresh 48-row " + "sequence beside 33 rows over a 48-token e4m3 prefix", + spec=dict( + _QWEN4_FP8, + cached=(0, 48), + new=(48, 33), + num_contexts=2, + predicted_tokens_per_seq=1, + ), + ), + Cell( + why="qwen4_exp fp8 pool decode P=1, 24/2/256: histories 200 and 48 span " + "pages and KV tiles", + spec=dict( + _QWEN4_FP8, + cached=(200, 48), + new=(1, 1), + num_contexts=0, + predicted_tokens_per_seq=1, + ), + ), + *( + Cell( + why=f"qwen4_exp fp8 pool MTP verify P={p}, 24/2/256 (max_draft_len {p - 1})", + spec=dict( + _QWEN4_FP8, + cached=(70, 35), + new=(p, p), + num_contexts=0, + predicted_tokens_per_seq=p, + ), + ) + for p in (2, 4, 5, 6, 7, 8) + ), + Cell( + why="qwen4_exp fp8 pool mixed step P=3, 24/2/256: a 20-row context ahead " + "of two P=3 blocks over histories 31 and 64", + spec=dict( + _QWEN4_FP8, + cached=(0, 31, 64), + new=(20, 3, 3), + num_contexts=1, + predicted_tokens_per_seq=3, + ), + ), + Cell( + why="qwen4_exp fp8 pool mixed step P=8, 24/2/256: a 20-row context ahead " + "of two P=8 blocks over histories 31 and 64", + spec=dict( + _QWEN4_FP8, + cached=(0, 31, 64), + new=(20, 8, 8), + num_contexts=1, + predicted_tokens_per_seq=8, + ), + ), + # KV stays within one 128-key tile (histories to 119 plus 8 rows), where the e4m3 + # probability model of `reference` is the kernel's own at any batch. + Cell( + why="qwen4_exp fp8 batch P=8 G=32: 32 blocks over histories 1..119, 256 q rows", + spec=dict( + _QWEN4_FP8_G32, + cached=_spread(32, 119), + new=(8,) * 32, + num_contexts=0, + predicted_tokens_per_seq=8, + ), + ), + Cell( + why="qwen4_exp fp8 hybrid pool P=2: row 47 of the 48-row two-pool mapping " + "(pool 0, layer-in-pool 11) over e4m3 slabs", + spec=dict( + _QWEN4_FP8, + pool_layout="hybrid48", + local_layer_idx=47, + cached=(64, 25), + new=(2, 2), + num_contexts=0, + predicted_tokens_per_seq=2, + ), + ), + Cell( + why="qwen4_exp fp8 hybrid pool P=8: row 47 of the 48-row two-pool mapping " + "(pool 0, layer-in-pool 11) over e4m3 slabs", + spec=dict( + _QWEN4_FP8, + pool_layout="hybrid48", + local_layer_idx=47, + cached=(64, 25), + new=(8, 8), + num_contexts=0, + predicted_tokens_per_seq=8, + ), + ), + Cell( + why="qwen4_exp nvfp4 pool paged-context prefill, 24/2/256: a fresh 48-row " + "sequence beside 33 rows over a 48-token e2m1 prefix", + spec=dict( + _QWEN4_NVFP4, + cached=(0, 48), + new=(48, 33), + num_contexts=2, + predicted_tokens_per_seq=1, + ), + ), + Cell( + why="qwen4_exp nvfp4 pool decode P=1, 24/2/256: histories 200 and 48", + spec=dict( + _QWEN4_NVFP4, + cached=(200, 48), + new=(1, 1), + num_contexts=0, + predicted_tokens_per_seq=1, + ), + ), + Cell( + why="qwen4_exp nvfp4 pool MTP verify P=2, 24/2/256 (max_draft_len 1)", + spec=dict( + _QWEN4_NVFP4, + cached=(70, 35), + new=(2, 2), + num_contexts=0, + predicted_tokens_per_seq=2, + ), + ), + Cell( + why="qwen4_exp nvfp4 pool MTP verify P=4, 24/2/256 (max_draft_len 3)", + spec=dict( + _QWEN4_NVFP4, + cached=(70, 35), + new=(4, 4), + num_contexts=0, + predicted_tokens_per_seq=4, + ), + ), + Cell( + why="qwen4_exp nvfp4 pool mixed step P=1: a 20-row context ahead of two " + "1-row sequences over histories 31 and 64", + spec=dict( + _QWEN4_NVFP4, + cached=(0, 31, 64), + new=(20, 1, 1), + num_contexts=1, + predicted_tokens_per_seq=1, + ), + ), + Cell( + why="qwen4_exp nvfp4 pool mixed step P=4: a 20-row context ahead of two 4-row blocks", + spec=dict( + _QWEN4_NVFP4, + cached=(0, 31, 64), + new=(20, 4, 4), + num_contexts=1, + predicted_tokens_per_seq=4, + ), + ), + ), + Fp4BlockScaleMoeRunner: ( + Cell( + why="qwen4_exp decode: one token over 512 experts, top-10 renormalized, H 2560 / I 640", + spec=dict(tokens=1, **_QWEN4_EXP), + ), + Cell( + why="qwen4_exp at 4 rows: the last count the pre-routed path's static-block " + "permutation kernel takes", + spec=dict(tokens=4, **_QWEN4_EXP), + ), + Cell( + why="qwen4_exp at 5 rows: the first count on the single-cluster permutation " + "kernel (the dynamic-block one stops at 256 experts)", + spec=dict(tokens=5, **_QWEN4_EXP), + ), + Cell( + why="qwen4_exp prefill: 8192 rows, the single-cluster kernel's cap and the top " + "of the certified column", + spec=dict(tokens=8192, **_QWEN4_EXP), + ), + ), + Fp4Quantize: ( + Cell( + why="qwen3.8 decode: 1 token, hidden 2560, linear -- its routed-expert input", + spec=dict(shape=(1, 2560), dtype=torch.bfloat16, vec=16, swizzled=False), + ), + Cell( + why="qwen3.8, 7 rows at hidden 2560, linear: several rows, no row padding", + spec=dict(shape=(7, 2560), dtype=torch.bfloat16, vec=16, swizzled=False), + ), + Cell( + why="qwen3.8 prefill at max_num_tokens 8192, hidden 2560, linear", + spec=dict(shape=(8192, 2560), dtype=torch.bfloat16, vec=16, swizzled=False), + ), + ), +} diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/routing.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/routing.py index ef2f31f01234..922e081840aa 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/routing.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/routing.py @@ -5,7 +5,7 @@ `route` compares the facts that fix which code runs against the one deployment the target was written for, and raises `NoTarget` stating them otherwise (`auto` falls back to the built-in implementation; `require` raises). The checkpoint's published -architecture name is kept by `load_pretrained_config`, so `_router_index` keys on it. +architecture name is kept by `load_pretrained_config`, so `route` registers under it. All criteria are read off `text_config` / `vision_config`; the composite parent's mirrored fields are unreliable. """ @@ -19,7 +19,7 @@ from tensorrt_llm.llmapi.llm_args import MTPDecodingConfig from tensorrt_llm.quantization.mode import QuantAlgo -from ..._router_index import NoTarget +from ..._router_index import NoTarget, routes if TYPE_CHECKING: from tensorrt_llm._torch.model_config import ModelConfig @@ -134,6 +134,9 @@ def _check_engine( ) +# Qwen3.8-Flash-Next publishes ``Qwen4ExpForConditionalGeneration``, and the +# config loader's composite branch writes the same name back. +@routes("Qwen4ExpForConditionalGeneration") def route( pretrained_config: Any, mapping: Any, diff --git a/tests/integration/test_lists/test-db/l0_b300.yml b/tests/integration/test_lists/test-db/l0_b300.yml index 3c10a654ba35..834d3a9076e8 100644 --- a/tests/integration/test_lists/test-db/l0_b300.yml +++ b/tests/integration/test_lists/test-db/l0_b300.yml @@ -135,9 +135,34 @@ l0_b300: - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_chunked_prefill[quant_dtype=fp8-kv_cache_reuse=True-fp8kv=True-overlap_scheduler=True] # Cover nvbugs 6084445 - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_flashinfer[enable_chunked_prefill=False] - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_flashinfer[enable_chunked_prefill=True] - # TIMEOUT measured: 3122 cases in 13m54s on one GB300 (the cpu_only files run on the CPU - # stage instead), so 30 is 2.2x headroom. - - unittest/_torch/modeling_v2 --ignore=unittest/_torch/modeling_v2/comm TIMEOUT (30) + # modeling_v2 unit tests, one entry per catalog op test and one per model + # family that has unit tests, so CBTS can select the ones a change reaches. + # The whole-tree tests that need no GPU are in l0_cpu; the comm tests in + # l0_gb300_multi_gpus. + - unittest/_torch/modeling_v2/activation/test_modeling_v2_flashinfer_silu_and_mul.py + - unittest/_torch/modeling_v2/attention/test_modeling_v2_fused_qk_norm_rope.py + - unittest/_torch/modeling_v2/attention/test_modeling_v2_load_paged_kv_cache_for_mla.py + - unittest/_torch/modeling_v2/attention/test_modeling_v2_mla_rope_append_paged_kv_assign_q.py + - unittest/_torch/modeling_v2/attention/test_modeling_v2_mla_rope_generation.py + - unittest/_torch/modeling_v2/attention/test_modeling_v2_thop_attention.py + - unittest/_torch/modeling_v2/gemm/test_modeling_v2_bmm_out.py + - unittest/_torch/modeling_v2/gemm/test_modeling_v2_cublas_mm.py + - unittest/_torch/modeling_v2/gemm/test_modeling_v2_nvfp4_gemm.py + - unittest/_torch/modeling_v2/graph/test_modeling_v2_eager_on_graph.py + - unittest/_torch/modeling_v2/moe/test_modeling_v2_fp4_block_scale_moe_runner.py + - unittest/_torch/modeling_v2/moe/test_modeling_v2_mxe4m3_mxe2m1_block_scale_moe_runner.py + - unittest/_torch/modeling_v2/moe/test_modeling_v2_noaux_tc_op.py + - unittest/_torch/modeling_v2/norm/test_modeling_v2_flashinfer_fused_add_rmsnorm.py + - unittest/_torch/modeling_v2/norm/test_modeling_v2_flashinfer_rmsnorm.py + - unittest/_torch/modeling_v2/quantization/test_modeling_v2_fp4_quantize.py + - unittest/_torch/modeling_v2/quantization/test_modeling_v2_mxfp8_quantize.py + - unittest/_torch/modeling_v2/models/deepseek_v3 + - unittest/_torch/modeling_v2/models/qwen4_exp + - unittest/_torch/modeling_v2/test_modeling_v2_deepseek_weight_table.py + - unittest/_torch/modeling_v2/test_modeling_v2_mnnvl_workspace.py + - unittest/_torch/modeling_v2/test_modeling_v2_no_stale_claims.py + - unittest/_torch/modeling_v2/test_modeling_v2_target_contract.py + - unittest/_torch/modeling_v2/test_modeling_v2_weight_table.py - accuracy/modeling_v2/test_gpt_oss.py::TestModelingV2GptOss120bSm10xDefault::test_gsm8k[tp1] # qwen4_exp target, pre-merge: gsm8k; one MTP engine for gsm8k with the acceptance-length # floor (the only gate that sees a miscomputed draft layer), mmmu (the vision tower) and diff --git a/tests/integration/test_lists/test-db/l0_gr100.yml b/tests/integration/test_lists/test-db/l0_gr100.yml index 104093738356..ce9bd2949dcd 100644 --- a/tests/integration/test_lists/test-db/l0_gr100.yml +++ b/tests/integration/test_lists/test-db/l0_gr100.yml @@ -15,3 +15,23 @@ l0_gr100: tests: - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_flashinfer[enable_chunked_prefill=True] - unittest/_torch/attention/test_flashinfer_trtllm_gen_sm107_p128.py + # The GLM-5.2 target's unit tests, one entry per op directory; its comm tests + # run 4 ranks and are in the 4-GPU block below. + - unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/attention + - unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa + - unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/embedding + - unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm + - unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/moe + - unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/quantization + - unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec +- condition: + ranges: + system_gpu_count: + gte: 4 + lte: 4 + terms: + compute_capability: '10.7' + stage: post_merge + backend: pytorch + tests: + - unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 749830020ad0..f7b48b4b5b5e 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -205,7 +205,6 @@ unittest/_torch/kernels/parallel/test_fp4_linear.py::test_fp4_gemm_bias_per_back unittest/_torch/kernels/parallel/test_fp4_linear.py::test_fp4_gemm_bias_per_backend[mnk3-cublaslt] SKIP (https://nvbugs/6581067) unittest/_torch/kernels/parallel/test_fp4_linear.py::test_fp4_gemm_bias_per_backend[mnk3-cutlass] SKIP (https://nvbugs/6581067) unittest/_torch/modeling/test_modeling_nemotron_h_multimodal.py::test_nemotron_nano_v2_vl_image_batch_equivalence SKIP (https://nvbugs/6758946) -unittest/_torch/modeling_v2 --ignore=unittest/_torch/modeling_v2/comm SKIP (https://nvbugs/6853741) unittest/_torch/modeling_v2/comm/test_modeling_v2_allgather_op_matrix.py::test_allgather_op_matrix SKIP (https://nvbugs/6843474) unittest/_torch/modeling_v2/comm/test_modeling_v2_reducescatter_op_matrix.py::test_reducescatter_op_matrix SKIP (https://nvbugs/6843474) unittest/_torch/modules/mamba/test_flashinfer_gdn_decode.py SKIP (https://nvbugs/6842831) diff --git a/tests/unittest/_torch/modeling_v2/_cells.py b/tests/unittest/_torch/modeling_v2/_cells.py new file mode 100644 index 000000000000..699c196c86d9 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/_cells.py @@ -0,0 +1,52 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Every cell a catalog entry is certified over: its own, and its targets'. + +An entry's own `CELLS` are the configurations that belong to the op whoever +calls it -- tile edges, the decode and prefill row counts, the dtypes its +kernel takes. The configurations a particular target runs belong to that +target, in a `cells.py` beside its `modeling.py`: + + CELLS = {FlashinferSiluAndMul: (Cell(why=..., spec=...), ...)} + +A target's cells change no forward, so a change to them reruns the entry tests +they feed and nothing of the target's. An entry's test drives the union, and a +spec several targets declare runs once. +""" + +from __future__ import annotations + +import functools +import importlib +from pathlib import Path + +import tensorrt_llm._torch._experimental.modeling_v2 as _modeling_v2 +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Cell, OpWrapper + +_ROOT = Path(_modeling_v2.__file__).resolve().parent + + +@functools.cache +def target_cells() -> dict[str, dict[type[OpWrapper], tuple[Cell, ...]]]: + """Each `cells.py` under `models/`, keyed by its directory relative to `models/`.""" + found = {} + for path in sorted((_ROOT / "models").rglob("cells.py")): + module = ".".join(path.relative_to(_ROOT).with_suffix("").parts) + cells = importlib.import_module(f"{_modeling_v2.__name__}.{module}").CELLS + found[path.parent.relative_to(_ROOT / "models").as_posix()] = cells + return found + + +def certified_cells(op: OpWrapper | type[OpWrapper]) -> tuple[Cell, ...]: + """`op`'s own cells, then every target's, each distinct spec once.""" + entry = op if isinstance(op, type) else type(op) + declared = [*entry.CELLS] + for by_entry in target_cells().values(): + declared.extend(by_entry.get(entry, ())) + seen, cells = set(), [] + for cell in declared: + key = repr(sorted(cell.spec.items())) + if key not in seen: + seen.add(key) + cells.append(cell) + return tuple(cells) diff --git a/tests/unittest/_torch/modeling_v2/activation/test_modeling_v2_flashinfer_silu_and_mul.py b/tests/unittest/_torch/modeling_v2/activation/test_modeling_v2_flashinfer_silu_and_mul.py index 56719bf1eae5..ed4c6c9c1811 100644 --- a/tests/unittest/_torch/modeling_v2/activation/test_modeling_v2_flashinfer_silu_and_mul.py +++ b/tests/unittest/_torch/modeling_v2/activation/test_modeling_v2_flashinfer_silu_and_mul.py @@ -16,6 +16,7 @@ ) __extra_import_path__ = [".."] +from _cells import certified_cells # noqa: E402 — needs the path declared above from _validating import validating # noqa: E402 — needs the path declared above assert torch.cuda.is_available(), "flashinfer_silu_and_mul requires a CUDA device" @@ -26,7 +27,10 @@ def _build(spec, seed): return torch.randn(spec["shape"], generator=g, device="cuda").to(spec["dtype"]) -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:40] for c in op.CELLS]) +CELLS = certified_cells(op) + + +@pytest.mark.parametrize("cell", CELLS, ids=[c.why[:40] for c in CELLS]) def test_certified_cells(cell) -> None: x = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) with validating(op): diff --git a/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_fused_qk_norm_rope.py b/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_fused_qk_norm_rope.py index f7a304db86bd..9fb0020c0071 100644 --- a/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_fused_qk_norm_rope.py +++ b/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_fused_qk_norm_rope.py @@ -12,6 +12,7 @@ ) __extra_import_path__ = [".."] +from _cells import certified_cells # noqa: E402 — needs the path declared above from _validating import ( # noqa: E402,F401 — needs the path declared above ULP, assert_within_ulp, @@ -134,7 +135,10 @@ def _deviation(out: torch.Tensor, ref: torch.Tensor) -> str: return f"max abs {(o - r).abs().max().item():.6f}, elt {elt:.2f} ulp, rms {rms:.3f} ulp" -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +CELLS = certified_cells(op) + + +@pytest.mark.parametrize("cell", CELLS, ids=[c.why[:44] for c in CELLS]) def test_certified_cells(cell, certifying) -> None: # noqa: F811 certifying(op) qkv, kwargs, compare_kwargs = _build(cell.spec, seed=_seed(cell.why)) @@ -157,7 +161,7 @@ def test_the_v_heads_are_left_alone() -> None: right size, and every cell compares against a reference that makes the same claim -- so this pins it against the input instead. """ - spec = dict(op.CELLS[1].spec) + spec = dict(CELLS[1].spec) qkv, kwargs, _ = _build(spec, seed=11) before = qkv.clone() op(qkv, **kwargs) @@ -170,7 +174,7 @@ def _qwen_spec(num_heads_q: int, tokens: int) -> dict: """The certified qwen4_exp mRoPE cell at this geometry, at `tokens` rows.""" matches = [ c.spec - for c in op.CELLS + for c in CELLS if c.spec["num_heads_q"] == num_heads_q and c.spec.get("use_mrope", False) ] assert matches, f"no certified mRoPE cell with {num_heads_q} query heads" @@ -259,7 +263,7 @@ def call(**over): def _ok(**over): - spec = dict(op.CELLS[0].spec) + spec = dict(CELLS[0].spec) spec.update(over) return _build(spec, seed=3)[:2] diff --git a/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_mla_rope_append_paged_kv_assign_q.py b/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_mla_rope_append_paged_kv_assign_q.py index fc6d547a98bf..a859cbf8034a 100644 --- a/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_mla_rope_append_paged_kv_assign_q.py +++ b/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_mla_rope_append_paged_kv_assign_q.py @@ -51,6 +51,7 @@ from tensorrt_llm.mapping import Mapping __extra_import_path__ = [".."] +from _cells import certified_cells # noqa: E402 — needs the path declared above from _validating import validating # noqa: E402 — needs the path declared above assert torch.cuda.is_available(), "mla_rope_append_paged_kv_assign_q requires a CUDA device" @@ -545,7 +546,10 @@ def _fp8_env(max_batch_size: int = 8, orig_quant: Optional[float] = None) -> _Ml ) -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +CELLS = certified_cells(op) + + +@pytest.mark.parametrize("cell", CELLS, ids=[c.why[:44] for c in CELLS]) def test_certified_cells(cell) -> None: """Every configuration the entry claims, through the same driver the hand-written cases use. diff --git a/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_mla_rope_generation.py b/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_mla_rope_generation.py index 636c563a8470..e5ee380a7969 100644 --- a/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_mla_rope_generation.py +++ b/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_mla_rope_generation.py @@ -61,6 +61,7 @@ from tensorrt_llm.mapping import Mapping __extra_import_path__ = [".."] +from _cells import certified_cells # noqa: E402 — needs the path declared above from _validating import validating # noqa: E402 — needs the path declared above assert torch.cuda.is_available(), "mla_rope_generation requires a CUDA device" @@ -1286,7 +1287,10 @@ def test_fp8_kv_mtp_append_race_control() -> None: env.shutdown() -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +CELLS = certified_cells(op) + + +@pytest.mark.parametrize("cell", CELLS, ids=[c.why[:44] for c in CELLS]) def test_certified_cells(cell) -> None: """Drive every configuration the entry claims, through the entry's own reference and its own gate. diff --git a/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_thop_attention.py b/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_thop_attention.py index 5c196b997868..6973d2f3b5c3 100644 --- a/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_thop_attention.py +++ b/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_thop_attention.py @@ -221,6 +221,7 @@ from tensorrt_llm.mapping import Mapping __extra_import_path__ = [".."] +from _cells import certified_cells # noqa: E402 — needs the path declared above from _validating import ( # noqa: E402,F401 — needs the path declared above ULP, assert_within_ulp, @@ -7264,7 +7265,10 @@ def _cell_encoder(cell, arm) -> None: torch.cuda.empty_cache() -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +CELLS = certified_cells(op) + + +@pytest.mark.parametrize("cell", CELLS, ids=[c.why[:44] for c in CELLS]) def test_certified_cells(cell, certifying) -> None: # noqa: F811 """Drive every configuration the entry claims. diff --git a/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_bmm_out.py b/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_bmm_out.py index 3ed6c1a73a43..29f277f69acf 100644 --- a/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_bmm_out.py +++ b/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_bmm_out.py @@ -12,6 +12,7 @@ ) __extra_import_path__ = [".."] +from _cells import certified_cells # noqa: E402 — needs the path declared above from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above assert torch.cuda.is_available(), "bmm_out requires a CUDA device" @@ -60,7 +61,10 @@ def _build(spec, seed): return a, w, out, parent -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +CELLS = certified_cells(op) + + +@pytest.mark.parametrize("cell", CELLS, ids=[c.why[:44] for c in CELLS]) def test_certified_cells(cell) -> None: a, w, out, parent = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) parent_ptr = parent.data_ptr() diff --git a/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_cublas_mm.py b/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_cublas_mm.py index dfe2c66a6d23..28ce3f13935c 100644 --- a/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_cublas_mm.py +++ b/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_cublas_mm.py @@ -10,6 +10,7 @@ from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.cublas_mm import cublas_mm as op __extra_import_path__ = [".."] +from _cells import certified_cells # noqa: E402 — needs the path declared above from _validating import ( # noqa: E402 — needs the path declared above ULP, assert_within_ulp, @@ -59,7 +60,10 @@ def _build(spec, seed): return mat_a, mat_b, bias, out_dtype -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +CELLS = certified_cells(op) + + +@pytest.mark.parametrize("cell", CELLS, ids=[c.why[:44] for c in CELLS]) def test_certified_cells(cell) -> None: mat_a, mat_b, bias, out_dtype = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) with validating(op): diff --git a/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_nvfp4_gemm.py b/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_nvfp4_gemm.py index 8f7d3f6b8ce0..213b6cf53ed6 100644 --- a/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_nvfp4_gemm.py +++ b/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_nvfp4_gemm.py @@ -13,6 +13,7 @@ ) __extra_import_path__ = [".."] +from _cells import certified_cells # noqa: E402 — needs the path declared above from _validating import ( # noqa: E402 — needs the path declared above ULP, assert_within_ulp, @@ -50,7 +51,10 @@ def _operand(rows, k, seed): return packed, sf, gs -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +CELLS = certified_cells(op) + + +@pytest.mark.parametrize("cell", CELLS, ids=[c.why[:44] for c in CELLS]) def test_certified_cells(cell) -> None: spec = cell.spec seed = zlib.crc32(cell.why.encode()) % 2**31 diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_cute_dsl_weight_path.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_cute_dsl_weight_path.py index ddef4b35ea36..6d7963ed32db 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_cute_dsl_weight_path.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_cute_dsl_weight_path.py @@ -71,8 +71,12 @@ #: the bf16 rounding of the output is 1e-3 of it, the rest is the reference's fp32 / tie #: differences. A wrong layout is ~1; a wrong per-expert scalar is its own relative error. REFERENCE_TOLERANCE = 5e-3 -#: The same, between the backends under a uniform `down_proj.input_scale`. -BACKEND_TOLERANCE = 3e-3 +#: The same, between the backends under a uniform `down_proj.input_scale`. Each backend is +#: within `REFERENCE_TOLERANCE` of the reference, so the two are within twice that of each +#: other. Both requantize the SwiGLU output to NVFP4, and an fp32 difference in the last +#: place moves an element at an e2m1 rounding boundary by a whole code: the token holding +#: it differs by ~5e-3 while the mean stays near 5e-6. A wrong layout is still ~1. +BACKEND_TOLERANCE = 2 * REFERENCE_TOLERANCE # ── the checkpoint ──────────────────────────────────────────────────────── diff --git a/tests/unittest/_torch/modeling_v2/moe/test_modeling_v2_fp4_block_scale_moe_runner.py b/tests/unittest/_torch/modeling_v2/moe/test_modeling_v2_fp4_block_scale_moe_runner.py index 09aaf1b4d565..3d06ad6b83a6 100644 --- a/tests/unittest/_torch/modeling_v2/moe/test_modeling_v2_fp4_block_scale_moe_runner.py +++ b/tests/unittest/_torch/modeling_v2/moe/test_modeling_v2_fp4_block_scale_moe_runner.py @@ -13,6 +13,7 @@ from tensorrt_llm._torch.autotuner import AutoTuner, autotune __extra_import_path__ = [".."] +from _cells import certified_cells # noqa: E402 — needs the path declared above from _validating import certifying, validating # noqa: E402,F401 — needs the path declared above assert torch.cuda.is_available(), "fp4_block_scale_moe_runner requires a CUDA device" @@ -2328,7 +2329,10 @@ def _cell_layer(spec: dict): return _dsr1_at(key[2]) -@pytest.mark.parametrize("cell", moe.CELLS, ids=[c.why[:44] for c in moe.CELLS]) +CELLS = certified_cells(moe) + + +@pytest.mark.parametrize("cell", CELLS, ids=[c.why[:44] for c in CELLS]) def test_certified_cells(cell, certifying): # noqa: F811 """Drive every configuration the entry claims, through the entry's own reference and its own gate. diff --git a/tests/unittest/_torch/modeling_v2/moe/test_modeling_v2_mxe4m3_mxe2m1_block_scale_moe_runner.py b/tests/unittest/_torch/modeling_v2/moe/test_modeling_v2_mxe4m3_mxe2m1_block_scale_moe_runner.py index e146823b3071..989da0cb3587 100644 --- a/tests/unittest/_torch/modeling_v2/moe/test_modeling_v2_mxe4m3_mxe2m1_block_scale_moe_runner.py +++ b/tests/unittest/_torch/modeling_v2/moe/test_modeling_v2_mxe4m3_mxe2m1_block_scale_moe_runner.py @@ -13,6 +13,7 @@ ) __extra_import_path__ = [".."] +from _cells import certified_cells # noqa: E402 — needs the path declared above from _validating import validating # noqa: E402 — needs the path declared above assert torch.cuda.is_available(), "mxe4m3_mxe2m1_block_scale_moe_runner requires a CUDA device" @@ -1621,7 +1622,10 @@ def _window(args, ref, offset: int, num_local: int): return sub, sub_ref -@pytest.mark.parametrize("cell", moe.CELLS, ids=[c.why[:44] for c in moe.CELLS]) +CELLS = certified_cells(moe) + + +@pytest.mark.parametrize("cell", CELLS, ids=[c.why[:44] for c in CELLS]) def test_certified_cells(cell): """Drive every configuration the entry claims, through the entry's own reference and its own gate. diff --git a/tests/unittest/_torch/modeling_v2/moe/test_modeling_v2_noaux_tc_op.py b/tests/unittest/_torch/modeling_v2/moe/test_modeling_v2_noaux_tc_op.py index df301eae3f8d..e74413ecf15b 100644 --- a/tests/unittest/_torch/modeling_v2/moe/test_modeling_v2_noaux_tc_op.py +++ b/tests/unittest/_torch/modeling_v2/moe/test_modeling_v2_noaux_tc_op.py @@ -12,6 +12,7 @@ ) __extra_import_path__ = [".."] +from _cells import certified_cells # noqa: E402 — needs the path declared above from _validating import validating # noqa: E402 — needs the path declared above assert torch.cuda.is_available(), "noaux_tc_op requires a CUDA device" @@ -26,7 +27,10 @@ def _build(spec, seed): return logits, bias, spec["n_group"], spec["topk_group"], spec["topk"], spec["scale"] -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +CELLS = certified_cells(op) + + +@pytest.mark.parametrize("cell", CELLS, ids=[c.why[:44] for c in CELLS]) def test_certified_cells(cell) -> None: args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) with validating(op): diff --git a/tests/unittest/_torch/modeling_v2/norm/test_modeling_v2_flashinfer_fused_add_rmsnorm.py b/tests/unittest/_torch/modeling_v2/norm/test_modeling_v2_flashinfer_fused_add_rmsnorm.py index 8072514beda3..8819358ea4d1 100644 --- a/tests/unittest/_torch/modeling_v2/norm/test_modeling_v2_flashinfer_fused_add_rmsnorm.py +++ b/tests/unittest/_torch/modeling_v2/norm/test_modeling_v2_flashinfer_fused_add_rmsnorm.py @@ -18,6 +18,7 @@ ) __extra_import_path__ = [".."] +from _cells import certified_cells # noqa: E402 — needs the path declared above from _validating import validating # noqa: E402 — needs the path declared above assert torch.cuda.is_available(), "flashinfer_fused_add_rmsnorm requires a CUDA device" @@ -32,7 +33,10 @@ def _build(spec, seed): return x, residual, weight, spec["eps"] -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:40] for c in op.CELLS]) +CELLS = certified_cells(op) + + +@pytest.mark.parametrize("cell", CELLS, ids=[c.why[:40] for c in CELLS]) def test_certified_cells(cell) -> None: """Every certified configuration, kernel against reference. diff --git a/tests/unittest/_torch/modeling_v2/norm/test_modeling_v2_flashinfer_rmsnorm.py b/tests/unittest/_torch/modeling_v2/norm/test_modeling_v2_flashinfer_rmsnorm.py index e581267cd4f0..41b0ac790f09 100644 --- a/tests/unittest/_torch/modeling_v2/norm/test_modeling_v2_flashinfer_rmsnorm.py +++ b/tests/unittest/_torch/modeling_v2/norm/test_modeling_v2_flashinfer_rmsnorm.py @@ -18,6 +18,7 @@ ) __extra_import_path__ = [".."] +from _cells import certified_cells # noqa: E402 — needs the path declared above from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name assert torch.cuda.is_available(), "flashinfer_rmsnorm requires a CUDA device" @@ -31,7 +32,10 @@ def _build(spec, seed): return x, weight, spec["eps"] -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +CELLS = certified_cells(op) + + +@pytest.mark.parametrize("cell", CELLS, ids=[c.why[:44] for c in CELLS]) def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name certifying(op) x, weight, eps = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) diff --git a/tests/unittest/_torch/modeling_v2/quantization/test_modeling_v2_fp4_quantize.py b/tests/unittest/_torch/modeling_v2/quantization/test_modeling_v2_fp4_quantize.py index 83f3b95fb046..0e3f5f699647 100644 --- a/tests/unittest/_torch/modeling_v2/quantization/test_modeling_v2_fp4_quantize.py +++ b/tests/unittest/_torch/modeling_v2/quantization/test_modeling_v2_fp4_quantize.py @@ -13,6 +13,7 @@ ) __extra_import_path__ = [".."] +from _cells import certified_cells # noqa: E402 — needs the path declared above from _validating import validating # noqa: E402 — needs the path declared above assert torch.cuda.is_available(), "fp4_quantize requires a CUDA device" @@ -39,7 +40,10 @@ def _unpack(packed): return codes -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +CELLS = certified_cells(op) + + +@pytest.mark.parametrize("cell", CELLS, ids=[c.why[:44] for c in CELLS]) def test_certified_cells(cell) -> None: """Codes bit-exact away from the rounding ties the entry declines to predict. diff --git a/tests/unittest/_torch/modeling_v2/quantization/test_modeling_v2_mxfp8_quantize.py b/tests/unittest/_torch/modeling_v2/quantization/test_modeling_v2_mxfp8_quantize.py index 1d6831905ee7..dfefad7b1f83 100644 --- a/tests/unittest/_torch/modeling_v2/quantization/test_modeling_v2_mxfp8_quantize.py +++ b/tests/unittest/_torch/modeling_v2/quantization/test_modeling_v2_mxfp8_quantize.py @@ -12,6 +12,7 @@ ) __extra_import_path__ = [".."] +from _cells import certified_cells # noqa: E402 — needs the path declared above from _validating import validating # noqa: E402 — needs the path declared above assert torch.cuda.is_available(), "mxfp8_quantize requires a CUDA device" @@ -36,7 +37,10 @@ def _build(spec, seed): return torch.randn(spec["shape"], generator=g, device="cuda").to(spec["dtype"]) -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +CELLS = certified_cells(op) + + +@pytest.mark.parametrize("cell", CELLS, ids=[c.why[:44] for c in CELLS]) def test_certified_cells(cell) -> None: x = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) swizzled, alignment = cell.spec["swizzled"], cell.spec["alignment"] diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_claims.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_claims.py index f52cafa47e16..42f7c524b4ad 100644 --- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_claims.py +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_claims.py @@ -19,27 +19,26 @@ import ast import re +import sys from pathlib import Path import pytest +import yaml import tensorrt_llm._torch._experimental.modeling_v2 as _modeling_v2 -from tensorrt_llm._torch._experimental.modeling_v2._router_index import ( - MODELING_V2_ROUTERS, - routing_module, -) +from tensorrt_llm._torch._experimental.modeling_v2._router_index import registered_routes # The package, not this file: these paths address the tree under test, and this # test lives in tests/ while that tree lives in tensorrt_llm/. pytestmark = pytest.mark.cpu_only _ROOT = Path(_modeling_v2.__file__).resolve().parent -_ARCHS = sorted(MODELING_V2_ROUTERS) +_ARCHS = sorted(registered_routes()) -def _module_path(dotted: str) -> Path: - """Resolve a package-relative dotted module name to its file.""" - return _ROOT / (dotted.replace(".", "/") + ".py") +def _routing_path(arch: str) -> Path: + """The routing.py whose ``route`` is registered for ``arch``.""" + return Path(sys.modules[registered_routes()[arch].__module__].__file__).resolve() # ``route`` hands back the class itself, imported from the target directory @@ -50,23 +49,19 @@ def _module_path(dotted: str) -> Path: def _routed_targets(arch: str): """(target directory, class name) for every target ``route`` can return.""" - source = _module_path(MODELING_V2_ROUTERS[arch]).read_text() - family_dir = _module_path(MODELING_V2_ROUTERS[arch]).parent + source = _routing_path(arch).read_text() + family_dir = _routing_path(arch).parent return [(family_dir / target, name) for target, name in _TARGET_IMPORT.findall(source)] -def test_every_routed_architecture_has_an_importable_routing_module(): - for arch in _ARCHS: - assert routing_module(arch) is not None, arch - - def test_no_routing_module_is_orphaned(): - """A routing.py the index does not name would never be consulted.""" - indexed = {_module_path(m) for m in MODELING_V2_ROUTERS.values()} + """A routing.py that models/__init__.py does not import, or whose ``route`` + carries no ``@routes``, would never be consulted.""" + registered = {_routing_path(arch) for arch in _ARCHS} on_disk = set((_ROOT / "models").glob("*/routing.py")) - assert on_disk == indexed, ( - f"routing modules on disk but not in MODELING_V2_ROUTERS: " - f"{sorted(p.relative_to(_ROOT) for p in on_disk - indexed)}" + assert on_disk == registered, ( + f"routing modules on disk with no registered route: " + f"{sorted(p.relative_to(_ROOT) for p in on_disk - registered)}" ) @@ -74,7 +69,7 @@ def test_no_routing_module_is_orphaned(): def test_no_target_directory_is_orphaned(arch): """A target directory ``route`` never imports from is dead code: nothing else can reach it.""" - family_dir = _module_path(MODELING_V2_ROUTERS[arch]).parent + family_dir = _routing_path(arch).parent on_disk = {p.parent for p in family_dir.glob("*/modeling.py")} routed = {target_dir for target_dir, _ in _routed_targets(arch)} assert on_disk == routed, ( @@ -155,7 +150,8 @@ def test_every_entry_with_cells_has_a_test_that_drives_them(): Checked by reading rather than importing, like the rest of this file: an entry declares `Cell(` in its source, and some test under this tree has to - name both the entry and `.CELLS`. That is weaker than proving the cells were + name both the entry and `.CELLS` -- or `certified_cells(`, which drives the + entry's cells and its targets'. That is weaker than proving the cells were driven, and it is what a check that runs anywhere can say. """ tests = list(Path(__file__).resolve().parent.rglob("*.py")) @@ -165,7 +161,10 @@ def test_every_entry_with_cells_has_a_test_that_drives_them(): for entry in sorted(_entry_files()): if "Cell(" not in entry.read_text(): continue - if not any(entry.stem in src and ".CELLS" in src for src in sources.values()): + if not any( + entry.stem in src and (".CELLS" in src or "certified_cells(" in src) + for src in sources.values() + ): undriven.append(str(entry.relative_to(_ROOT))) assert not undriven, ( @@ -174,6 +173,111 @@ def test_every_entry_with_cells_has_a_test_that_drives_them(): ) +def _cells_files() -> list[Path]: + return sorted((_ROOT / "models").rglob("cells.py")) + + +@pytest.mark.parametrize("cells", _cells_files(), ids=lambda p: str(p.relative_to(_ROOT))) +def test_target_cells_feed_entries_the_target_calls(cells): + """A `cells.py` holds a target's configurations for the catalog entries it calls. + + It sits in a target's directory, imports entries only from + `catalog//.py`, and each of those its target's own code + imports too: a cell for an entry nothing under it calls certifies nothing + it runs. Each entry's test sits at the mirrored path + `/test_modeling_v2_.py` and drives `certified_cells(`, which + is how it reaches these cells, and how CBTS finds the tests a `cells.py` + change has to rerun. + """ + target_dir = cells.parent + assert (target_dir / "modeling.py").is_file(), f"{cells}: not beside a target's modeling.py" + used = "".join( + p.read_text() + for p in target_dir.rglob("*.py") + if p != cells and "__pycache__" not in p.parts + ) + tests_root = Path(__file__).resolve().parent + for node in ast.walk(ast.parse(cells.read_text())): + if not isinstance(node, ast.ImportFrom) or not node.level: + continue + module = (node.module or "").split(".") + if module == ["catalog", "_op"]: + continue + assert len(module) == 3 and module[0] == "catalog", ( + f"{cells}: imports {node.module!r}; cells.py imports catalog entries only" + ) + _, area, entry = module + assert re.search(rf"catalog\.{area}\.{entry} import", used), ( + f"{cells}: declares cells for {area}/{entry}, which its target never imports" + ) + test = tests_root / area / f"test_modeling_v2_{entry}.py" + assert test.is_file() and "certified_cells(" in test.read_text(), ( + f"{cells}: {area}/{entry} has no test at {test.relative_to(tests_root)} " + f"driving certified_cells(" + ) + + +_TESTS_ROOT = Path(__file__).resolve().parents[3] + + +def test_every_target_has_the_accuracy_gate_of_its_name(): + """A target's gate is `Test` in `accuracy/modeling_v2/test_.py`. + + CBTS reaches a target's gate through that name alone, so a gate renamed + away from its target, or a target with none, is one a change to it would + never run. + """ + routed = { + (_routing_path(arch).parent.name, f"Test{name}") + for arch in _ARCHS + for _, name in _routed_targets(arch) + } + gates = set() + for path in (_TESTS_ROOT / "integration/defs/accuracy/modeling_v2").glob("test_*.py"): + family = path.stem.removeprefix("test_") + for node in ast.parse(path.read_text()).body: + if isinstance(node, ast.ClassDef) and node.name.startswith("TestModelingV2"): + gates.add((family, node.name)) + assert routed == gates, ( + f"targets with no gate of their name: {sorted(routed - gates)}; " + f"gates naming no routed target: {sorted(gates - routed)}" + ) + + +def test_every_unit_test_is_in_a_test_list(): + """The modeling_v2 unit tests are listed per file or per family directory. + + Lists that name files do not pick up a new file by themselves, and a test + no list reaches never runs in CI. The other way round, an entry naming a + path with no test at it fails its CI stage: pytest finds nothing to run. + """ + listed = set() + for path in (_TESTS_ROOT / "integration/test_lists/test-db").glob("l0_*.yml"): + for blocks in (yaml.safe_load(path.read_text()) or {}).values(): + for block in blocks if isinstance(blocks, list) else (): + for entry in block.get("tests") or (): + listed.add(entry.split()[0].split("::")[0]) + + unit_root = Path(__file__).resolve().parent + unlisted = [] + for test in sorted(unit_root.rglob("test_*.py")): + rel = test.relative_to(_TESTS_ROOT / "unittest") + candidates = {f"unittest/{p.as_posix()}" for p in (rel, *rel.parents)} + if not candidates & listed: + unlisted.append(rel.as_posix()) + assert not unlisted, f"unit tests no test-db list reaches: {unlisted}" + + # A directory left behind with only `__pycache__` in it exists on disk but + # holds no test, so require a test file at or under each entry. + empty = sorted( + entry + for entry in listed + if entry.startswith("unittest/_torch/modeling_v2/") + and not ((_TESTS_ROOT / entry).is_file() or any((_TESTS_ROOT / entry).rglob("test_*.py"))) + ) + assert not empty, f"test-db entries with no modeling_v2 unit test at them: {empty}" + + # The phase a step runs at is resolved once, at the dispatcher, and a target is # what runs after that resolution. These two gates are what keep that true: # without them a later edit can put `if md.num_contexts:` back inside a decode diff --git a/tests/unittest/scripts/cbts/test_modeling_v2.py b/tests/unittest/scripts/cbts/test_modeling_v2.py new file mode 100644 index 000000000000..82f6a0d58104 --- /dev/null +++ b/tests/unittest/scripts/cbts/test_modeling_v2.py @@ -0,0 +1,115 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""ModelingV2Rule selects the tests of the node a change lands in.""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +__extra_import_path__ = ["~/jenkins/scripts/cbts"] +from cbts.blocks import Block, YAMLIndex +from cbts.rules.base import PRInputs +from cbts.rules.modeling_v2 import ModelingV2Rule + +_SRC = "tensorrt_llm/_torch/_experimental/modeling_v2" +_UNIT = "unittest/_torch/modeling_v2" +_GATE_A = "accuracy/modeling_v2/test_fam.py::TestModelingV2A" +_GATE_B = "accuracy/modeling_v2/test_fam.py::TestModelingV2B" + +_ENTRIES = ( + f"{_UNIT}/test_modeling_v2_claims.py", + f"{_UNIT}/norm/test_modeling_v2_rmsnorm.py", + f"{_UNIT}/gemm/test_modeling_v2_mm.py", + f"{_UNIT}/models/fam", + f"{_UNIT}/models/other", + f"{_UNIT}/comm", + f"{_GATE_A}::test_gsm8k[tp1]", + f"{_GATE_B}::test_gsm8k[tp4]", + "accuracy/modeling_v2/test_other.py::TestModelingV2C::test_gsm8k", +) + + +def _rule(repo: Path) -> ModelingV2Rule: + """A tree with family `fam` (targets a, b, and an ops-only dir) and family `other`.""" + files = { + f"{_SRC}/models/fam/routing.py": ( + "from .a.modeling import ModelingV2A\nfrom .b.modeling import ModelingV2B\n" + ), + f"{_SRC}/models/fam/a/modeling.py": "", + f"{_SRC}/models/fam/a/ops/norm/fused.py": "", + f"{_SRC}/models/fam/a/cells.py": ( + "from ....catalog._op import Cell\n" + "from ....catalog.norm.rmsnorm import Rmsnorm\n" + "CELLS = {Rmsnorm: (Cell(why='x'),)}\n" + ), + f"{_SRC}/models/fam/b/modeling.py": "", + f"{_SRC}/models/fam/b_sm103/ops/gemm.py": "", + f"{_SRC}/models/other/routing.py": "from .c.modeling import ModelingV2C\n", + f"{_SRC}/models/other/c/modeling.py": "", + f"tests/{_UNIT}/norm/test_modeling_v2_rmsnorm.py": "", + } + for path, text in files.items(): + (repo / path).parent.mkdir(parents=True, exist_ok=True) + (repo / path).write_text(text) + index = YAMLIndex() + index.blocks = [Block("l0_b300", 0, {}, list(_ENTRIES))] + return ModelingV2Rule(index, {}, repo_root=repo) + + +@pytest.mark.parametrize( + ("changed", "selected"), + ( + pytest.param( + "models/fam/a/ops/norm/fused.py", + {f"{_UNIT}/models/fam", f"{_GATE_A}::test_gsm8k[tp1]"}, + id="target-runs-its-family-unit-tests-and-only-its-own-gate", + ), + pytest.param( + "models/fam/b_sm103/ops/gemm.py", + {f"{_UNIT}/models/fam", f"{_GATE_A}::test_gsm8k[tp1]", f"{_GATE_B}::test_gsm8k[tp4]"}, + id="ops-only-dir-belongs-to-its-family", + ), + pytest.param( + "models/fam/routing.py", + {f"{_UNIT}/models/fam", f"{_GATE_A}::test_gsm8k[tp1]", f"{_GATE_B}::test_gsm8k[tp4]"}, + id="family-file-runs-every-target-of-the-family", + ), + pytest.param( + "models/fam/a/cells.py", + {f"{_UNIT}/norm/test_modeling_v2_rmsnorm.py"}, + id="cells-run-only-the-ops-they-certify", + ), + pytest.param("catalog/norm/rmsnorm.py", set(_ENTRIES), id="catalog-runs-everything"), + ), +) +def test_change_selects_its_node(tmp_path: Path, changed: str, selected: set[str]) -> None: + result = _rule(tmp_path).apply(PRInputs(changed_files=[f"{_SRC}/{changed}"], diffs={})) + kept = {entry for waives in result.block_filters[("l0_b300", 0)].values() for entry in waives} + assert kept == selected | {f"{_UNIT}/test_modeling_v2_claims.py"} + + +def test_cells_naming_an_op_without_a_mirrored_test_fall_back_to_the_node(tmp_path: Path) -> None: + rule = _rule(tmp_path) + cells = tmp_path / _SRC / "models/fam/a/cells.py" + cells.write_text("from ....catalog.gemm.unmirrored import Op\n") + result = rule.apply(PRInputs(changed_files=[f"{_SRC}/models/fam/a/cells.py"], diffs={})) + kept = {entry for waives in result.block_filters[("l0_b300", 0)].values() for entry in waives} + assert kept == { + f"{_UNIT}/test_modeling_v2_claims.py", + f"{_UNIT}/models/fam", + f"{_GATE_A}::test_gsm8k[tp1]", + }