diff --git a/.claude/skills/ad-sharding-ir-port/SKILL.md b/.claude/skills/ad-sharding-ir-port/SKILL.md index 1eb40bdbb380..3437555a85f8 100644 --- a/.claude/skills/ad-sharding-ir-port/SKILL.md +++ b/.claude/skills/ad-sharding-ir-port/SKILL.md @@ -33,7 +33,7 @@ You MAY introduce ONLY the following changes: - `nn.Linear(...)` / `F.linear(...)` → `torch.ops.auto_deploy.torch_linear_simple(...)` - `tensor.view(...)` / `tensor.reshape(...)` → `torch.ops.auto_deploy.view(...)` (only when the shape contains a TP-scaled dim) - `torch.split(...)` / `torch.split_with_sizes(...)` → `torch.ops.auto_deploy.split_with_sizes(...)` -- **A2. Sharding-hint kwargs added** to call sites of: `torch_moe`, `torch_ssm`, `torch_gated_delta_rule`, `torch_causal_conv1d`, `torch_rmsnorm_gated`, `torch_mla`, `torch_attention`, `torch_linear_simple`, `auto_deploy.split_with_sizes`, `auto_deploy.view`. Allowed kwargs: `tp_mode`, `layer_type`, `output_sizes`, `tp_min_local_shape`, `tp_scaled_dim`, `shardable`, `enable_sharding`. +- **A2. Sharding-hint kwargs added** to call sites of: `torch_moe`, `torch_ssm`, `torch_gated_delta_rule`, `torch_causal_conv1d`, `torch_rmsnorm_gated`, `torch_mla`, `torch_linear_simple`, `auto_deploy.split_with_sizes`, `auto_deploy.view`. Allowed kwargs: `tp_mode`, `layer_type`, `output_sizes`, `tp_min_local_shape`, `tp_scaled_dim`, `shardable`, `enable_sharding`. - **A3. Inserting `torch.ops.auto_deploy.all_reduce(..., layer_type=...)`** after rowwise projections / at MoE merge points (single all_reduce after routed + shared sums). - **A4. Docstring updates:** - Module-level: a single-line header noting the file uses sharding IR, followed by the existing source-of-truth / HF link block. Example: `"""Llama 3 model (sharding IR)."""`. @@ -100,24 +100,33 @@ Pass `layer_type="moe"` into `torch_moe`; `apply_sharding_hints` handles EP/TP. The model's existing registration (`AutoModelForCausalLMFactory.register_custom_model_cls` at the bottom of the file and its import in `__init__.py`) stays unchanged. No new registration is needed — sharding hints do not change the model identity. -### Step 9: YAML — no per-model opt-in needed +### Step 9: YAML — enable hint-driven sharding -No YAML change is required to enable the IR path. The default sharding pipeline (`apply_sharding_hints`) auto-detects the presence of `torch.ops.auto_deploy.all_reduce` markers in the exported FX graph and routes IR-marked models to the IR pipeline; non-marked models fall through to the legacy `detect_sharding` + `sharding_transform_executor` pair. The markers you added in Steps 1–7 are sufficient. +Add `enable_sharder_ir.yaml` to the model's `yaml_extra` list in `examples/auto_deploy/model_registry/models.yaml` (if not already present). This composable fragment disables legacy sharding passes and enables `apply_sharding_hints`. Registry fragments are deep-merged in `yaml_extra` order (see `DynamicYamlMixInForSettings` in `tensorrt_llm/_torch/auto_deploy/utils/_config.py`). -If the model needs a non-default `apply_sharding_hints` config (for example a non-NCCL `allreduce_strategy`, or selective `shard_layers`), add a per-model yaml override under `examples/auto_deploy/model_registry/configs/` that overrides only the keys you need: +Example transform block: ```yaml +# Typical contents for enable_sharder_ir.yaml (registry composable fragment) transforms: - apply_sharding_hints: - allreduce_strategy: SYMM_MEM - # shard_layers: ['mha', 'mlp'] # optional selective sharding export_to_gm: num_moe_experts_for_export: 2 # often required when expert count is large (>64) + detect_sharding: + stage: sharding + enabled: false + sharding_transform_executor: + stage: sharding + enabled: false + apply_sharding_hints: + stage: sharding + enabled: true + run_shape_prop: true + allreduce_strategy: NCCL + # shard_layers: ['mha', 'mlp'] # optional selective sharding + gather_logits_before_lm_head: + enabled: true ``` -To force the legacy pipeline (e.g. while an IR port has a known bug awaiting fix), add `enable_legacy_sharding.yaml` to the model's `yaml_extra` — that override disables `apply_sharding_hints` and re-enables the legacy stages explicitly. - - Set `world_size` once, to the **maximum number of GPUs available on the machine**, auto-detected with `python -c 'import torch; print(torch.cuda.device_count())'` (or `nvidia-smi --list-gpus | wc -l`). Do **not** hardcode `world_size: 8` (or any other literal) — porting agents run on heterogeneous hardware and an 8-GPU literal will simply fail to launch on a 2- or 4-GPU machine. If the model's `num_attention_heads` (and, for GQA, `num_key_value_heads`) does not divide the detected GPU count, fall back to the largest power-of-two divisor that does (e.g. 4 on an 8-GPU machine if `num_attention_heads = 12`). Run the end-to-end command exactly once at that size — there is no value in repeating it at multiple smaller sizes, because the offline sharding equivalence test (Step 10b) already exercises 2- and 4-GPU dist configs cheaply. Optional `shard_layers` limits which `layer_type` hints are processed; unset means shard all shardable nodes. @@ -135,15 +144,15 @@ Do not report success until a run completes successfully. ### Step 10b — Sharding equivalence test (MANDATORY) -Run the offline sharding-IR equivalence test ([`tests/unittest/auto_deploy/multigpu/transformations/library/test_sharding_num_correctness.py`](tests/unittest/auto_deploy/multigpu/transformations/library/test_sharding_num_correctness.py)) against the modeling file you just edited, under **every** parallelism configuration the test exposes. The port is **not** complete until every configuration passes. Skipping this step or treating a partial pass (e.g. only `tep`) as success is not allowed. +Run the offline sharding-IR equivalence test ([`tests/unittest/auto_deploy/multigpu/transformations/library/test_sharding_ir_equivalence.py`](tests/unittest/auto_deploy/multigpu/transformations/library/test_sharding_ir_equivalence.py)) against the modeling file you just edited, under **every** parallelism configuration the test exposes. The port is **not** complete until every configuration passes. Skipping this step or treating a partial pass (e.g. only `tep`) as success is not allowed. -The test compares a sharded prefill against the unsharded eager reference on a tiny (4-layer, hidden_size=64) instance of the model and asserts `rel_rmse < tol`, where `tol` is the test-defined relative-RMSE tolerance (`REL_RMSE_TOL` constant in [`test_sharding_num_correctness.py`](tests/unittest/auto_deploy/multigpu/transformations/library/test_sharding_num_correctness.py); overridable per invocation via the `SHARDING_IR_REL_RMSE_TOL` env var). It uses no PyExecutor / no compile / no checkpoint download, so each cell runs in ~30s on 4xGPU. +The test compares a sharded prefill against the unsharded eager reference on a tiny (4-layer, hidden_size=64) instance of the model and asserts `rel_rmse < tol`, where `tol` is the test-defined relative-RMSE tolerance (`REL_RMSE_TOL` constant in [`test_sharding_ir_equivalence.py`](tests/unittest/auto_deploy/multigpu/transformations/library/test_sharding_ir_equivalence.py); overridable per invocation via the `SHARDING_IR_REL_RMSE_TOL` env var). It uses no PyExecutor / no compile / no checkpoint download, so each cell runs in ~30s on 4xGPU. **Run the matrix:** ```bash MODEL=tensorrt_llm/_torch/auto_deploy/models/custom/modeling_.py -TEST=tests/unittest/auto_deploy/multigpu/transformations/library/test_sharding_num_correctness.py +TEST=tests/unittest/auto_deploy/multigpu/transformations/library/test_sharding_ir_equivalence.py for CFG in tp-only ep-only tep attn-dp; do pytest "$TEST" --sharding-ir-modeling-file "$MODEL" --sharding-ir-dist-config "$CFG" -s -v \ @@ -226,8 +235,6 @@ You are NOT done until every row in the table is a yes-allowed category. **MLA (DeepSeek):** `layer_type="mla"`: keep `torch_mla` intact with `shardable=True`—do **not** decompose into separate linears + `torch_attention` (introduces bad `expand`/`view` with concrete head counts). q_a/kv_a latent: `tp_mode="none"`; q_b colwise; `o_proj` rowwise + `all_reduce`. -**Per-head free Parameters on `torch_attention` (GPT-OSS-style sinks):** when an attention block has a learnable `nn.Parameter` indexed by Q-head count that flows DIRECTLY into `torch_attention` (not through a Linear) — e.g. GPT-OSS's `self.sinks = nn.Parameter(torch.empty(num_heads))` passed as `sinks=self.sinks` — pass `enable_sharding=True` to the `torch_attention(...)` call. The IR's `WeightedParamShardableNode` is registered for `torch_attention` and will slice every direct `get_attr` arg along dim 0 (= head dim) per rank. Q/K/V/O projection weights are unaffected (they belong to the preceding `torch_linear_simple` nodes and are sharded by `LinearShardableNode`). Models with no such head-wise Parameter (qwen3, llama, smollm3, ...) leave `enable_sharding` at its default `False` and the handler no-ops for them. - ## Common pitfalls 1. **Missing `auto_deploy::view` for head reshapes** — concrete shapes from export break after sharding. @@ -238,7 +245,7 @@ You are NOT done until every row in the table is a yes-allowed category. 6. **Decomposing ops that absorb weights** (e.g. `torch_mla`) — use `shardable` + handler instead of splitting into plain linears. 7. **Interleaved vs contiguous fused weights** — interleaved per-head groups: colwise only; contiguous Q|K|V blocks: require `output_sizes`. 8. **Omitting `layer_type` when using `shard_layers`** — `"unknown"` nodes are skipped; set hints explicitly on sharding-aware ops. -9. **`layer_type` on non-hint ops** — do **not** pass `layer_type` to ops that are not designed for sharding hints (e.g. `torch_l2norm`, `torch_rope_*`); extra positional args break calls. Note: `torch_attention` DOES accept `layer_type` (and `enable_sharding`) — see the per-head Parameters paragraph in "Layer-specific sharding patterns" above. Confirm in `custom_ops/` docstrings which ops accept hints. +9. **`layer_type` on non-hint ops** — do **not** pass `layer_type` to ops that are not designed for sharding hints (e.g. `torch_attention`, `torch_l2norm`, `torch_rope_*`); extra positional args break calls. Confirm in `custom_ops/` docstrings which ops accept hints. 10. **Conditional hint values** — no `if _s else "none"`; use unconditional hints and rely on `shard_layers` / transform config. 11. **Replacing `torch.ops.trtllm.*` ops** — `noaux_tc_op`, `dsv3_router_gemm_op`, fused norm/MLP kernels are TP-replicated and must be kept verbatim (rule F1). AD has no fusion pass to recover them from vanilla PyTorch. diff --git a/.claude/skills/trtllm-model-onboard-multimodal/SKILL.md b/.claude/skills/trtllm-model-onboard-multimodal/SKILL.md index cc96eaaefa02..bef01b2e670b 100644 --- a/.claude/skills/trtllm-model-onboard-multimodal/SKILL.md +++ b/.claude/skills/trtllm-model-onboard-multimodal/SKILL.md @@ -304,18 +304,7 @@ class {Name}Model(PreTrainedModel): ### Phase 3 — Input processor + dummy builder -Subclass **both** `BaseMultimodalInputProcessor` (drives every real request) and `BaseMultimodalDummyInputsBuilder` (drives engine warmup / KV-cache profiling). Colocate in the modeling file. References: `Qwen3VLInputProcessorBase` (image+video), `Mistral3InputProcessor` (Pixtral). - -**Encoder KV-cache memory profiling (deterministic dummy sizing).** The KV-cache profiler sizes the encoder's memory contribution by running the encoder **once** on a worst-case dummy (held resident through the peak measurement), **decoupled** from the text-only LLM dummy. To opt in, the model exposes `encode_multimodal_inputs` (via `MultimodalModelMixin`) and the input processor implements the modality-agnostic dummy contract on `BaseMultimodalDummyInputsBuilder`: - -- `get_mm_max_tokens_per_item() -> {modality: tokens}` — per-modality worst-case single-item encoder-attention tokens. The keys enumerate the modalities the model encodes; the profiler splits the shared `encoder_max_num_tokens` across them in proportion to these (so they share one microbatch cap, not each the whole budget). Default `{}` → no direct encoder profiling. -- `get_dummy_mm_data_for_tokens(*, max_tokens_per_modality, dtype) -> multimodal_data` — materialize the processed encoder tensors **directly** (zeros of the exact shape the processor would emit; no PIL image + HF-processor round-trip), merged into one `multimodal_data` dict so a single `encode_multimodal_inputs` profiles the combined peak. Default raises `NotImplementedError`. - -Vision models implement these via the size trio: `get_num_mm_tokens(*, width, height, num_frames)` (size → **pre-merger encoder-attention tokens**; the single source of truth shared with the hashing path `get_num_tokens_per_image`/`_video`), its inverse `get_size_for_max_tokens(max_tokens)` (largest aspect-bounded size whose token count ≤ budget, capped at `max_pixels`), and `get_dummy_mm_data_for_size(...)`. Qwen builds `pixel_values`/`image_grid_thw`; Mistral builds `pixel_values`/`image_sizes` and keeps its ViT patch count off `get_num_mm_tokens` (a private `_vit_tokens` helper) so the LLM-side Pixtral hashing count is unchanged. A model with neither contract falls back to a text-only dummy (encoder memory unaccounted). Don't hardcode the encoder attention workspace (`max_num_*=8192`): inherit `MultimodalEncoderMixin` and let the engine size it via `setup_attn_metadata` at load. - -The workspace dimensions come from `TorchLlmArgs.get_encoder_runtime_sizes()` → `(encoder_max_batch_size, encoder_max_num_tokens)` — two prototype knobs that size the encoder's `AttentionMetadata` independently of the LLM batch and fall back to the LLM-side `max_batch_size` / `max_num_tokens` when unset. `encoder_max_num_tokens` is exactly the per-iteration encoder microbatch token cap that `get_dummy_mm_data_for_tokens` saturates (and that the profiler splits across modalities), so an encoder microbatch can be sized larger than the LLM `max_num_tokens` without inflating the KV-cache budget. Read them via `get_encoder_runtime_sizes()` rather than the raw fields so the fallback is applied. - -> Mixed image+video+audio models (nemotron-nano / phi4mm) compose multiple modality dummies through the same `get_dummy_mm_data_for_tokens` (return `{"image": ..., "audio": ...}`); a per-modality `ModalityDummySizer` composition is the planned home for the shared orchestration. +Subclass **both** `BaseMultimodalInputProcessor` (drives every real request) and `BaseMultimodalDummyInputsBuilder` (drives engine warmup / profiling — the base shrinks dummy image resolution until the synthetic prompt fits `input_seq_len`). Colocate in the modeling file. Reference: `Qwen3VLInputProcessorBase`. Implement `call_with_text_prompt(inputs, sampling_params)` — the per-model text-prompt path. **Don't override `__call__`**: the base class's concrete `__call__` dispatches here for text prompts, and also detokenizes `prompt_token_ids → prompt` and falls through to here for non-fast-path VLMs. `call_with_text_prompt` does: @@ -444,7 +433,6 @@ Follow `CONTRIBUTING.md`. Title `[JIRA/NVBUG/None][type] description`, `git comm **Input processor** - [ ] Subclasses both `BaseMultimodalInputProcessor` and `BaseMultimodalDummyInputsBuilder`. -- [ ] Encoder KV-cache profiling: implements the deterministic dummy contract (`get_mm_max_tokens_per_item` + `get_dummy_mm_data_for_tokens`, vision via the `get_num_mm_tokens` / `get_size_for_max_tokens` / `get_dummy_mm_data_for_size` trio) and the model exposes `encode_multimodal_inputs`; encoder inherits `MultimodalEncoderMixin` (no hardcoded `max_num_*=8192` — sized by `setup_attn_metadata`). Skipping these = text-only dummy, encoder memory unaccounted. - [ ] `call_with_text_prompt` (not `__call__` — that's the base-class dispatcher) runs HF AutoProcessor + tokenizer, builds `multimodal_data` by modality, computes `mrope_config` on CPU, `_postprocess`-rewrites mm token ids to the OOV sentinel. - [ ] `mm_processor_kwargs` flow-through preserved. (Tokenized fast path is optional: set `supports_token_id_mm_expansion = True` + implement `get_text_with_mm_placeholders` / `expand_prompt_token_ids_for_mm`; otherwise the base class detokenizes token-ID inputs automatically.) - [ ] `_attach_multimodal_embeddings_impl` implemented (not the `attach_multimodal_embeddings` wrapper) if `@support_multimodal_disaggregated`. diff --git a/.claude/skills/trtllm-moe-develop/SKILL.md b/.claude/skills/trtllm-moe-develop/SKILL.md index ab7b282bc415..81698a55b3a5 100644 --- a/.claude/skills/trtllm-moe-develop/SKILL.md +++ b/.claude/skills/trtllm-moe-develop/SKILL.md @@ -268,26 +268,6 @@ Checklist: - Existing legacy `forward` methods can be read for compatibility context, but they are not the default pattern for new backend work. -### Imported Kernel ABI Checklist - -When importing or wrapping an upstream kernel, derive the TRT-LLM adapter -contract from the lowest-level kernel consumer. Comments, docs, design notes, -and parameter names are useful hints, but they are not proof of the runtime ABI. - -- Derive weight shape and layout from the kernel entrypoint, `make_layout`, TMA, - MMA/GEMM transforms, and stride usage. Record required tensor shape, stride, - physical storage layout, and boundary view layout. -- Derive alpha and scale semantics from kernel consumption points. Trace where - alpha, norm constants, block scales, activation scales, and weight scales are - loaded and multiplied before deciding how upper layers compute or pack them. - Treat weight bytes, block scales/SF, and global alpha/norm constants as - separate contracts. -- Design the upper-layer adapter from the kernel ABI upward. Map each kernel - input/output to an adapter responsibility: storage tensor, view/transposition, - dtype reinterpretation, padding, scale packing, workspace ownership, - synchronization, and output reduction. Validate parity with upstream - invocation dumps, not just final output. - ### Quantization And Weights Role: diff --git a/.gitattributes b/.gitattributes index 15bc65ba47ac..177818296355 100644 --- a/.gitattributes +++ b/.gitattributes @@ -10,12 +10,10 @@ triton_backend/tools/gpt/input_data.json filter=lfs diff=lfs merge=lfs -text docs/source/blogs/media/tech_blog3_mla_absorb.png filter=lfs diff=lfs merge=lfs -text tests/integration/test_input_files/*.png filter=lfs diff=lfs merge=lfs -text tests/integration/test_input_files/*.jpg filter=lfs diff=lfs merge=lfs -text -tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/*.zip filter=lfs diff=lfs merge=lfs -text +tests/integration/defs/examples/golden/visual_gen_lpips/*.zip filter=lfs diff=lfs merge=lfs -text docs/source/blogs/media/tech_blog10_baseline_performance_detail.png filter=lfs diff=lfs merge=lfs -text docs/source/blogs/media/tech_blog10_full_strategy_performance.png filter=lfs diff=lfs merge=lfs -text docs/source/blogs/media/tech_blog10_context_wait_performance.png filter=lfs diff=lfs merge=lfs -text cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/cubin/kernelMetaInfo_cubin.cpp filter=lfs diff=lfs merge=lfs -text cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/cubin/xqa_kernel_cubin.cpp filter=lfs diff=lfs merge=lfs -text tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/*/*/*.so filter=lfs diff=lfs merge=lfs -text -docs/source/blogs/media/tech_blog26_deepseek_v4_hybrid_attention.png filter=lfs diff=lfs merge=lfs -text -docs/source/blogs/media/tech_blog26_deepseek_v4_mhc_moe.png filter=lfs diff=lfs merge=lfs -text diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 83e0275ad77c..dc88f8a256d0 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -1,418 +1,217 @@ -# ============================================================================ -# PRECEDENCE: last matching pattern wins — fallback first, module rules next, governance last. -# -# Domain (in file order) Handle -# Engineering baseline/ fallback / untriaged @NVIDIA/trt-llm-devs (global `*`, declared FIRST; incl. deprecating TRT backend) -# Multi-GPU CI gate @NVIDIA/trt-llm-multi-gpu-ci-review -# Infra @NVIDIA/trt-llm-infra-devs -# Agent config @NVIDIA/trt-llm-agent-devs -# Docs / Examples @NVIDIA/trt-llm-doc-owners -# QA @NVIDIA/trt-llm-qa -# Runtime @NVIDIA/trt-llm-runtime-devs -# Kernels - Misc @NVIDIA/trt-llm-kernels-devs -# Models @NVIDIA/trt-llm-models-devs -# General Perf @NVIDIA/trt-llm-perf-devs -# KV Cache Manager @NVIDIA/trt-llm-kv-cache-manager-devs -# Disaggregated Serving @NVIDIA/trt-llm-disagg-devs -# Attention @NVIDIA/trt-llm-torch-attention-devs -# MoE @NVIDIA/trt-llm-moe-devs -# VisualGen (AIGV) @NVIDIA/trt-llm-torch-visual-gen-devs -# Scaffolding @WeiHaocheng -# Self-managed: AutoDeploy, Triton backend, Dynamo. Governance/legal: unchanged. -# -# New handles (trt-llm-devs / doc / agent / runtime / models / perf / moe / kernels / multi-gpu-ci-review) -# must be created & populated before activation (unknown owner blocks merge). -# ============================================================================ - -# ===== FALLBACK (trt-llm-devs) ===== -# Global net; every rule below overrides it. The explicit entries are visible parking (redundant with `*`). -* @NVIDIA/trt-llm-devs - -/tensorrt_llm/commands/eval.py @NVIDIA/trt-llm-devs -/tensorrt_llm/evaluate @NVIDIA/trt-llm-devs -/tensorrt_llm/tools @NVIDIA/trt-llm-devs -/tests/integration/test_lists/test-db @NVIDIA/trt-llm-devs @NVIDIA/trt-llm-qa @NVIDIA/trt-llm-infra-devs -/tests/integration/test_lists/waives.txt @NVIDIA/trt-llm-devs @NVIDIA/trt-llm-qa @NVIDIA/trt-llm-infra-devs -/tests/test_common @NVIDIA/trt-llm-devs -/tests/unittest @NVIDIA/trt-llm-devs - -# ===== TensorRT backend (will be deprecated soon) — also on the trt-llm-devs fallback ===== -/cpp/include/tensorrt_llm/plugins @NVIDIA/trt-llm-devs -/cpp/tensorrt_llm/plugins @NVIDIA/trt-llm-devs -/tensorrt_llm/builder.py @NVIDIA/trt-llm-devs -/tensorrt_llm/commands/build.py @NVIDIA/trt-llm-devs -/tensorrt_llm/commands/prune.py @NVIDIA/trt-llm-devs -/tensorrt_llm/commands/refit.py @NVIDIA/trt-llm-devs -/tensorrt_llm/functional.py @NVIDIA/trt-llm-devs -/tensorrt_llm/graph_rewriting.py @NVIDIA/trt-llm-devs -/tensorrt_llm/layers @NVIDIA/trt-llm-devs -/tensorrt_llm/models @NVIDIA/trt-llm-devs -/tensorrt_llm/module.py @NVIDIA/trt-llm-devs -/tensorrt_llm/network.py @NVIDIA/trt-llm-devs -/tensorrt_llm/parameter.py @NVIDIA/trt-llm-devs -/tensorrt_llm/plugin @NVIDIA/trt-llm-devs -/tensorrt_llm/python_plugin.py @NVIDIA/trt-llm-devs -/tensorrt_llm/runtime @NVIDIA/trt-llm-devs -/tensorrt_llm/tools/plugin_gen @NVIDIA/trt-llm-devs -/tensorrt_llm/top_model_mixin.py @NVIDIA/trt-llm-devs - -# ===== MULTI-GPU / MULTI-NODE CI TEST-LIST GATE (review group, NEW) ===== -# A PR adding/expanding these expensive cases must justify WHAT is tested, WHY existing coverage is -# insufficient, and WHY a unit / single-GPU test can't do it. (Handle is a placeholder.) -/tests/integration/test_lists/test-db/*multi_gpu* @NVIDIA/trt-llm-multi-gpu-ci-review -/tests/integration/test_lists/test-db/*multi_node* @NVIDIA/trt-llm-multi-gpu-ci-review -/tests/integration/test_lists/test-db/l0_dgx_* @NVIDIA/trt-llm-multi-gpu-ci-review - -# ===== INFRA ===== -/.coderabbit.yaml @NVIDIA/trt-llm-infra-devs -/.github @NVIDIA/trt-llm-infra-devs -/.pre-commit-config.yaml @NVIDIA/trt-llm-infra-devs -/docker @NVIDIA/trt-llm-infra-devs -/enroot @NVIDIA/trt-llm-infra-devs -/jenkins @NVIDIA/trt-llm-infra-devs -/ruff-legacy-baseline.json @NVIDIA/trt-llm-infra-devs -/ruff-legacy.toml @NVIDIA/trt-llm-infra-devs -/scripts @NVIDIA/trt-llm-infra-devs -/security_scanning @NVIDIA/trt-llm-infra-devs - -# ===== AGENT ===== -/.claude @NVIDIA/trt-llm-agent-devs -/.codex @NVIDIA/trt-llm-agent-devs -/AGENTS.md @NVIDIA/trt-llm-agent-devs -/CLAUDE.md @NVIDIA/trt-llm-agent-devs -/scripts/check_skill_naming_convention.py @NVIDIA/trt-llm-agent-devs - -# ===== DOCS / EXAMPLES ===== -/CODE_OF_CONDUCT.md @NVIDIA/trt-llm-doc-owners +# This file defines code ownership rules for the repository. + +## TensorRT-LLM QA +### Integration Tests +/tests/integration/test_lists/qa @NVIDIA/trt-llm-qa +/tests/integration/defs/examples/test_ray.py @NVIDIA/trt-llm-qa-function +/tests/integration/defs/examples/test_redrafter.py @NVIDIA/trt-llm-qa-function +/tests/integration/defs/accuracy @NVIDIA/trt-llm-qa-function +/tests/integration/defs/stress_test @NVIDIA/trt-llm-qa-function +/tests/integration/defs/triton_server @NVIDIA/trt-llm-qa-function +/tests/integration/defs/test_e2e.py @NVIDIA/trt-llm-qa-function +/tests/integration/defs/disaggregated @NVIDIA/trt-llm-qa-serving +/tests/integration/defs/sysinfo @NVIDIA/trt-llm-qa-perf +/tests/integration/defs/perf @NVIDIA/trt-llm-qa-perf +/tests/integration/defs/perf/disagg @NVIDIA/trt-llm-qa-serving + +## TensorRT-LLM Infra +### CI +/jenkins @NVIDIA/trt-llm-ci-infra-devs @NVIDIA/trt-llm-infra-devs +### Setup +/docker @NVIDIA/trt-llm-setup-infra-devs @NVIDIA/trt-llm-infra-devs +/.pre-commit-config.yaml @NVIDIA/trt-llm-setup-infra-devs @NVIDIA/trt-llm-infra-devs +### Github workflows +/.github @NVIDIA/trt-llm-gh-workflows-infra-devs @NVIDIA/trt-llm-infra-devs +/.coderabbit.yaml @NVIDIA/trt-llm-gh-workflows-infra-devs @NVIDIA/trt-llm-infra-devs + +## TensorRT-LLM - Docs +/docs @NVIDIA/trt-llm-doc-owners /CODING_GUIDELINES.md @NVIDIA/trt-llm-doc-owners +/CODE_OF_CONDUCT.md @NVIDIA/trt-llm-doc-owners /CONTAINER_SOURCE.md @NVIDIA/trt-llm-doc-owners /CONTRIBUTING.md @NVIDIA/trt-llm-doc-owners /README.md @NVIDIA/trt-llm-doc-owners -/SECURITY.md @NVIDIA/trt-llm-doc-owners -/docs @NVIDIA/trt-llm-doc-owners +/CLAUDE.md @NVIDIA/trt-llm-doc-owners +/AGENTS.md @NVIDIA/trt-llm-doc-owners + +## Examples /examples @NVIDIA/trt-llm-doc-owners -# ===== QA ===== -/tests/integration/defs @NVIDIA/trt-llm-devs @NVIDIA/trt-llm-qa @NVIDIA/trt-llm-infra-devs -/tests/integration/test_lists/qa @NVIDIA/trt-llm-qa +## TensorRT-LLM - Triton backend +/triton_backend @NVIDIA/trt-llm-triton-backend-devs -# ===== RUNTIME ===== -/cpp/include/tensorrt_llm/batch_manager @NVIDIA/trt-llm-runtime-devs -/cpp/include/tensorrt_llm/common @NVIDIA/trt-llm-runtime-devs -/cpp/include/tensorrt_llm/executor @NVIDIA/trt-llm-runtime-devs -/cpp/include/tensorrt_llm/layers @NVIDIA/trt-llm-runtime-devs -/cpp/include/tensorrt_llm/runtime @NVIDIA/trt-llm-runtime-devs -/cpp/tensorrt_llm/batch_manager @NVIDIA/trt-llm-runtime-devs -/cpp/tensorrt_llm/common @NVIDIA/trt-llm-runtime-devs -/cpp/tensorrt_llm/executor @NVIDIA/trt-llm-runtime-devs -/cpp/tensorrt_llm/executor_worker @NVIDIA/trt-llm-runtime-devs -/cpp/tensorrt_llm/layers @NVIDIA/trt-llm-runtime-devs -/cpp/tensorrt_llm/nanobind @NVIDIA/trt-llm-runtime-devs -/cpp/tensorrt_llm/runtime @NVIDIA/trt-llm-runtime-devs -/cpp/tensorrt_llm/testing @NVIDIA/trt-llm-runtime-devs -/cpp/tests @NVIDIA/trt-llm-runtime-devs -/tensorrt_llm/_tensorrt_engine @NVIDIA/trt-llm-runtime-devs -/tensorrt_llm/_torch @NVIDIA/trt-llm-runtime-devs -/tensorrt_llm/commands/__init__.py @NVIDIA/trt-llm-runtime-devs -/tensorrt_llm/commands/serve.py @NVIDIA/trt-llm-runtime-devs -/tensorrt_llm/commands/utils.py @NVIDIA/trt-llm-runtime-devs -/tensorrt_llm/executor @NVIDIA/trt-llm-runtime-devs -/tensorrt_llm/grpc @NVIDIA/trt-llm-runtime-devs -/tensorrt_llm/inputs @NVIDIA/trt-llm-runtime-devs -/tensorrt_llm/llmapi @NVIDIA/trt-llm-runtime-devs -/tensorrt_llm/llmapi/mm_encoder.py @NVIDIA/trt-llm-runtime-devs @NVIDIA/trt-llm-models-devs -/tensorrt_llm/lora_helper.py @NVIDIA/trt-llm-runtime-devs -/tensorrt_llm/lora_manager.py @NVIDIA/trt-llm-runtime-devs -/tensorrt_llm/mapping.py @NVIDIA/trt-llm-runtime-devs -/tensorrt_llm/metrics @NVIDIA/trt-llm-runtime-devs -/tensorrt_llm/prompt_adapter_manager.py @NVIDIA/trt-llm-runtime-devs -/tensorrt_llm/sampling_params.py @NVIDIA/trt-llm-runtime-devs -/tensorrt_llm/scheduling_params.py @NVIDIA/trt-llm-runtime-devs -/tensorrt_llm/serve @NVIDIA/trt-llm-runtime-devs -/tensorrt_llm/tokenizer @NVIDIA/trt-llm-runtime-devs -/tensorrt_llm/usage @NVIDIA/trt-llm-runtime-devs -/tests/torch @NVIDIA/trt-llm-runtime-devs -/tests/unittest/_torch/compilation @NVIDIA/trt-llm-runtime-devs -/tests/unittest/_torch/custom_ops @NVIDIA/trt-llm-runtime-devs -/tests/unittest/_torch/distributed @NVIDIA/trt-llm-runtime-devs -/tests/unittest/_torch/executor @NVIDIA/trt-llm-runtime-devs -/tests/unittest/_torch/lora @NVIDIA/trt-llm-runtime-devs -/tests/unittest/_torch/misc/test_autotuner.py @NVIDIA/trt-llm-runtime-devs -/tests/unittest/_torch/modules @NVIDIA/trt-llm-runtime-devs -/tests/unittest/_torch/multi_gpu @NVIDIA/trt-llm-runtime-devs -/tests/unittest/_torch/pyexecutor @NVIDIA/trt-llm-runtime-devs -/tests/unittest/_torch/ray_orchestrator @NVIDIA/trt-llm-runtime-devs -/tests/unittest/_torch/sampler @NVIDIA/trt-llm-runtime-devs -/tests/unittest/_torch/speculative @NVIDIA/trt-llm-runtime-devs -/tests/unittest/_torch/speculative_hw_agnostic @NVIDIA/trt-llm-runtime-devs -/tests/unittest/executor @NVIDIA/trt-llm-runtime-devs -/tests/unittest/inputs @NVIDIA/trt-llm-runtime-devs -/tests/unittest/llmapi @NVIDIA/trt-llm-runtime-devs -/tests/unittest/llmapi/apps/*multimodal* @NVIDIA/trt-llm-runtime-devs @NVIDIA/trt-llm-models-devs -/tests/unittest/pyexecutor @NVIDIA/trt-llm-runtime-devs - -# ===== KERNELS - MISC ===== -/cpp/include/tensorrt_llm/deep_gemm @NVIDIA/trt-llm-kernels-devs -/cpp/include/tensorrt_llm/kernels @NVIDIA/trt-llm-kernels-devs -/cpp/tensorrt_llm/cutlass_extensions @NVIDIA/trt-llm-kernels-devs -/cpp/tensorrt_llm/deep_gemm @NVIDIA/trt-llm-kernels-devs -/cpp/tensorrt_llm/kernels @NVIDIA/trt-llm-kernels-devs -/cpp/tensorrt_llm/thop @NVIDIA/trt-llm-kernels-devs -/cpp/tests/unit_tests/kernels @NVIDIA/trt-llm-kernels-devs -/tensorrt_llm/_torch/cuda_tile_kernels @NVIDIA/trt-llm-kernels-devs -/tensorrt_llm/_torch/cute_dsl_kernels @NVIDIA/trt-llm-kernels-devs -/tests/scripts/cute_dsl_kernels @NVIDIA/trt-llm-kernels-devs -/tests/unittest/_torch/thop @NVIDIA/trt-llm-kernels-devs - -# ===== MODELS ===== -/docs/source/features/multi-modality.md @NVIDIA/trt-llm-models-devs @NVIDIA/trt-llm-doc-owners -/examples/llm-api/quickstart_multimodal.py @NVIDIA/trt-llm-models-devs @NVIDIA/trt-llm-doc-owners -/examples/models @NVIDIA/trt-llm-models-devs @NVIDIA/trt-llm-doc-owners -/examples/serve/*multimodal* @NVIDIA/trt-llm-models-devs @NVIDIA/trt-llm-doc-owners -/scripts/build_cpp_examples.py @NVIDIA/trt-llm-models-devs -/scripts/generate_config_database_tests.py @NVIDIA/trt-llm-models-devs @NVIDIA/trt-llm-doc-owners -/scripts/generate_config_table.py @NVIDIA/trt-llm-models-devs @NVIDIA/trt-llm-doc-owners -/tensorrt_llm/_torch/models @NVIDIA/trt-llm-models-devs -/tensorrt_llm/_torch/modules/mamba @NVIDIA/trt-llm-models-devs -/tensorrt_llm/quantization @NVIDIA/trt-llm-models-devs -/tests/integration/defs/accuracy/test_llm_api_pytorch_multimodal.py @NVIDIA/trt-llm-models-devs @NVIDIA/trt-llm-qa -/tests/unittest/_torch/modeling @NVIDIA/trt-llm-models-devs -/tests/unittest/_torch/models @NVIDIA/trt-llm-models-devs -/tests/unittest/_torch/modules/mamba @NVIDIA/trt-llm-models-devs -/tests/unittest/_torch/multi_gpu_modeling @NVIDIA/trt-llm-models-devs -/tests/unittest/_torch/multimodal @NVIDIA/trt-llm-models-devs -/tests/unittest/models @NVIDIA/trt-llm-models-devs -/tests/unittest/others/test_multimodal_registry.py @NVIDIA/trt-llm-models-devs - -# ===== GENERAL PERF ===== -/benchmarks @NVIDIA/trt-llm-perf-devs -/cpp/micro_benchmarks @NVIDIA/trt-llm-perf-devs -/docs/source/performance/perf-benchmarking.md @NVIDIA/trt-llm-perf-devs -/scripts/check_pinned_memory_usage.py @NVIDIA/trt-llm-perf-devs -/tensorrt_llm/bench @NVIDIA/trt-llm-perf-devs -/tensorrt_llm/commands/bench.py @NVIDIA/trt-llm-perf-devs -/tensorrt_llm/tools/layer_wise_benchmarks @NVIDIA/trt-llm-perf-devs -/tensorrt_llm/tools/profiler @NVIDIA/trt-llm-perf-devs -/tests/microbenchmarks @NVIDIA/trt-llm-perf-devs -/tests/scripts/allreduce_perf @NVIDIA/trt-llm-perf-devs -/tests/scripts/iteration_log_parser.py @NVIDIA/trt-llm-perf-devs -/tests/scripts/perf @NVIDIA/trt-llm-perf-devs -/tests/scripts/perf-sanity @NVIDIA/trt-llm-perf-devs - -# ===== KV CACHE MANAGER ===== -/cpp/include/tensorrt_llm/batch_manager/allocateKvCache* @NVIDIA/trt-llm-kv-cache-manager-devs -/cpp/include/tensorrt_llm/batch_manager/blockKey* @NVIDIA/trt-llm-kv-cache-manager-devs -/cpp/include/tensorrt_llm/batch_manager/evictionPolicy* @NVIDIA/trt-llm-kv-cache-manager-devs -/cpp/include/tensorrt_llm/batch_manager/kvCache* @NVIDIA/trt-llm-kv-cache-manager-devs -/cpp/include/tensorrt_llm/batch_manager/radixBlockTree* @NVIDIA/trt-llm-kv-cache-manager-devs -/cpp/include/tensorrt_llm/batch_manager/reorderPolicy* @NVIDIA/trt-llm-kv-cache-manager-devs -/cpp/include/tensorrt_llm/batch_manager/stringSetTrie* @NVIDIA/trt-llm-kv-cache-manager-devs -/cpp/include/tensorrt_llm/batch_manager/templatedTrie* @NVIDIA/trt-llm-kv-cache-manager-devs -/cpp/tensorrt_llm/batch_manager/allocateKvCache* @NVIDIA/trt-llm-kv-cache-manager-devs -/cpp/tensorrt_llm/batch_manager/blockKey* @NVIDIA/trt-llm-kv-cache-manager-devs -/cpp/tensorrt_llm/batch_manager/evictionPolicy* @NVIDIA/trt-llm-kv-cache-manager-devs -/cpp/tensorrt_llm/batch_manager/kvCache* @NVIDIA/trt-llm-kv-cache-manager-devs -/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager* @NVIDIA/trt-llm-kv-cache-manager-devs -/cpp/tests/unit_tests/batch_manager/blockKey* @NVIDIA/trt-llm-kv-cache-manager-devs -/cpp/tests/unit_tests/batch_manager/evictionPolicy* @NVIDIA/trt-llm-kv-cache-manager-devs -/cpp/tests/unit_tests/batch_manager/kvCache* @NVIDIA/trt-llm-kv-cache-manager-devs -/cpp/tests/unit_tests/batch_manager/radixBlockTree* @NVIDIA/trt-llm-kv-cache-manager-devs -/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py @NVIDIA/trt-llm-kv-cache-manager-devs -/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py @NVIDIA/trt-llm-kv-cache-manager-devs -/tensorrt_llm/_torch/pyexecutor/resource_manager.py @NVIDIA/trt-llm-kv-cache-manager-devs -/tensorrt_llm/runtime/kv_cache_manager_v2 @NVIDIA/trt-llm-kv-cache-manager-devs -/tests/unittest/_torch/executor/test_dual_pool_kv_cache.py @NVIDIA/trt-llm-kv-cache-manager-devs -/tests/unittest/_torch/executor/test_kv_cache* @NVIDIA/trt-llm-kv-cache-manager-devs -/tests/unittest/_torch/executor/test_kv_pool_rebalance.py @NVIDIA/trt-llm-kv-cache-manager-devs -/tests/unittest/_torch/executor/test_kvcache_aware_router.py @NVIDIA/trt-llm-kv-cache-manager-devs -/tests/unittest/_torch/executor/test_mamba_cache_manager.py @NVIDIA/trt-llm-kv-cache-manager-devs -/tests/unittest/kv_cache_manager_v2_tests @NVIDIA/trt-llm-kv-cache-manager-devs - -# ===== DISAGGREGATED SERVING ===== -/cpp/include/tensorrt_llm/batch_manager/*Formatter* @NVIDIA/trt-llm-disagg-devs -/cpp/include/tensorrt_llm/batch_manager/*TransBuffer* @NVIDIA/trt-llm-disagg-devs -/cpp/include/tensorrt_llm/batch_manager/*Transceiver* @NVIDIA/trt-llm-disagg-devs -/cpp/include/tensorrt_llm/batch_manager/cacheTransferLayer* @NVIDIA/trt-llm-disagg-devs -/cpp/include/tensorrt_llm/batch_manager/disagg* @NVIDIA/trt-llm-disagg-devs -/cpp/tensorrt_llm/batch_manager/*Formatter* @NVIDIA/trt-llm-disagg-devs -/cpp/tensorrt_llm/batch_manager/*TransBuffer* @NVIDIA/trt-llm-disagg-devs -/cpp/tensorrt_llm/batch_manager/*Transceiver* @NVIDIA/trt-llm-disagg-devs -/cpp/tensorrt_llm/batch_manager/cacheTransferLayer* @NVIDIA/trt-llm-disagg-devs -/cpp/tensorrt_llm/batch_manager/disagg* @NVIDIA/trt-llm-disagg-devs -/examples/disaggregated @NVIDIA/trt-llm-disagg-devs @NVIDIA/trt-llm-doc-owners -/examples/disaggregated/slurm/benchmark @NVIDIA/trt-llm-disagg-devs @NVIDIA/trt-llm-perf-devs -/tensorrt_llm/_torch/disaggregation @NVIDIA/trt-llm-disagg-devs -/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py @NVIDIA/trt-llm-disagg-devs -/tensorrt_llm/disaggregated_params.py @NVIDIA/trt-llm-disagg-devs -/tensorrt_llm/serve/openai_disagg_server.py @NVIDIA/trt-llm-disagg-devs -# Disagg tests: co-own with the owning team so disagg-devs review disagg-test changes. -/tests/integration/defs/accuracy/*disagg* @NVIDIA/trt-llm-disagg-devs @NVIDIA/trt-llm-qa -/tests/integration/defs/disaggregated @NVIDIA/trt-llm-disagg-devs @NVIDIA/trt-llm-qa -/tests/integration/defs/stress_test/disagg_cancel @NVIDIA/trt-llm-disagg-devs @NVIDIA/trt-llm-qa -/tests/scripts/perf-sanity/disaggregated @NVIDIA/trt-llm-perf-devs @NVIDIA/trt-llm-disagg-devs -/tests/scripts/perf/disaggregated @NVIDIA/trt-llm-perf-devs @NVIDIA/trt-llm-disagg-devs -/tests/unittest/_torch/executor/*disagg* @NVIDIA/trt-llm-runtime-devs @NVIDIA/trt-llm-disagg-devs -/tests/unittest/_torch/multimodal/*disagg* @NVIDIA/trt-llm-models-devs @NVIDIA/trt-llm-disagg-devs -/tests/unittest/disaggregated @NVIDIA/trt-llm-disagg-devs -/tests/unittest/llmapi/*disagg* @NVIDIA/trt-llm-runtime-devs @NVIDIA/trt-llm-disagg-devs -/tests/unittest/llmapi/apps/*disagg* @NVIDIA/trt-llm-runtime-devs @NVIDIA/trt-llm-disagg-devs - -# ===== ATTENTION ===== -# Where a kernel is both a dir and sibling .cu/.h, keep the bare dir (subtree) AND a trailing-* (siblings). -/cpp/kernels @NVIDIA/trt-llm-torch-attention-devs -/cpp/tensorrt_llm/common/attention* @NVIDIA/trt-llm-torch-attention-devs -/cpp/tensorrt_llm/flash_mla @NVIDIA/trt-llm-torch-attention-devs -/cpp/tensorrt_llm/kernels/IndexerKCache* @NVIDIA/trt-llm-torch-attention-devs -/cpp/tensorrt_llm/kernels/IndexerTopK* @NVIDIA/trt-llm-torch-attention-devs -/cpp/tensorrt_llm/kernels/attentionMask* @NVIDIA/trt-llm-torch-attention-devs -/cpp/tensorrt_llm/kernels/buildRelativeAttentionBiasKernel* @NVIDIA/trt-llm-torch-attention-devs -/cpp/tensorrt_llm/kernels/compressorKernels @NVIDIA/trt-llm-torch-attention-devs -/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention @NVIDIA/trt-llm-torch-attention-devs -/cpp/tensorrt_llm/kernels/convertReqIndexToGlobal* @NVIDIA/trt-llm-torch-attention-devs -/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention @NVIDIA/trt-llm-torch-attention-devs -/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention* @NVIDIA/trt-llm-torch-attention-devs -/cpp/tensorrt_llm/kernels/deepseekV4QNormKernel* @NVIDIA/trt-llm-torch-attention-devs -/cpp/tensorrt_llm/kernels/flashMLA @NVIDIA/trt-llm-torch-attention-devs -/cpp/tensorrt_llm/kernels/fmhaDispatcher* @NVIDIA/trt-llm-torch-attention-devs -/cpp/tensorrt_llm/kernels/fusedCatFp4* @NVIDIA/trt-llm-torch-attention-devs -/cpp/tensorrt_llm/kernels/fusedCatFp8* @NVIDIA/trt-llm-torch-attention-devs -/cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel* @NVIDIA/trt-llm-torch-attention-devs -/cpp/tensorrt_llm/kernels/gptKernels* @NVIDIA/trt-llm-torch-attention-devs -/cpp/tensorrt_llm/kernels/helix* @NVIDIA/trt-llm-torch-attention-devs -/cpp/tensorrt_llm/kernels/indexerKCache* @NVIDIA/trt-llm-torch-attention-devs -/cpp/tensorrt_llm/kernels/indexerTopK* @NVIDIA/trt-llm-torch-attention-devs -/cpp/tensorrt_llm/kernels/inverseRopeFp8QuantKernel* @NVIDIA/trt-llm-torch-attention-devs -/cpp/tensorrt_llm/kernels/mla* @NVIDIA/trt-llm-torch-attention-devs -/cpp/tensorrt_llm/kernels/multiHeadAttentionCommon.h @NVIDIA/trt-llm-torch-attention-devs -/cpp/tensorrt_llm/kernels/sparseAttentionKernels* @NVIDIA/trt-llm-torch-attention-devs -/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha @NVIDIA/trt-llm-torch-attention-devs -/cpp/tensorrt_llm/kernels/unfusedAttentionKernels @NVIDIA/trt-llm-torch-attention-devs -/cpp/tensorrt_llm/kernels/unfusedAttentionKernels* @NVIDIA/trt-llm-torch-attention-devs -/cpp/tensorrt_llm/kernels/xqaDispatcher* @NVIDIA/trt-llm-torch-attention-devs -/cpp/tensorrt_llm/nanobind/thop @NVIDIA/trt-llm-torch-attention-devs -/cpp/tensorrt_llm/thop/IndexerKCache* @NVIDIA/trt-llm-torch-attention-devs -/cpp/tensorrt_llm/thop/IndexerTopK* @NVIDIA/trt-llm-torch-attention-devs -/cpp/tensorrt_llm/thop/attention* @NVIDIA/trt-llm-torch-attention-devs -/cpp/tensorrt_llm/thop/compressorOp.cpp @NVIDIA/trt-llm-torch-attention-devs -/cpp/tensorrt_llm/thop/convertReqIndexToGlobalOp.cpp @NVIDIA/trt-llm-torch-attention-devs -/cpp/tensorrt_llm/thop/deepseekV4QNormOp.cpp @NVIDIA/trt-llm-torch-attention-devs -/cpp/tensorrt_llm/thop/dsv3RopeOp.cpp @NVIDIA/trt-llm-torch-attention-devs -/cpp/tensorrt_llm/thop/fmhaPackMaskOp.cpp @NVIDIA/trt-llm-torch-attention-devs -/cpp/tensorrt_llm/thop/fusedCatFp4Op.cpp @NVIDIA/trt-llm-torch-attention-devs -/cpp/tensorrt_llm/thop/fusedCatFp8Op.cpp @NVIDIA/trt-llm-torch-attention-devs -/cpp/tensorrt_llm/thop/fusedQKNormRopeOp.cpp @NVIDIA/trt-llm-torch-attention-devs -/cpp/tensorrt_llm/thop/helixPostProcessOp.cpp @NVIDIA/trt-llm-torch-attention-devs -/cpp/tensorrt_llm/thop/inverseRopeFp8QuantOp.cpp @NVIDIA/trt-llm-torch-attention-devs -/cpp/tensorrt_llm/thop/mla* @NVIDIA/trt-llm-torch-attention-devs -/cpp/tensorrt_llm/thop/relativeAttentionBiasOp.cpp @NVIDIA/trt-llm-torch-attention-devs -/cpp/tensorrt_llm/thop/trtllmGenQKVProcessOp.cpp @NVIDIA/trt-llm-torch-attention-devs -/cpp/tests/unit_tests/common/attentionWorkspaceTest.cpp @NVIDIA/trt-llm-torch-attention-devs -/cpp/tests/unit_tests/kernels/cascadeAttention* @NVIDIA/trt-llm-torch-attention-devs -/cpp/tests/unit_tests/kernels/mla* @NVIDIA/trt-llm-torch-attention-devs -/cpp/tests/unit_tests/kernels/ropeTest.cu @NVIDIA/trt-llm-torch-attention-devs -/cpp/tests/unit_tests/kernels/sparseAttentionKernelsTest.cpp @NVIDIA/trt-llm-torch-attention-devs -/tensorrt_llm/_torch/attention_backend @NVIDIA/trt-llm-torch-attention-devs -/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits @NVIDIA/trt-llm-torch-attention-devs -/tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md @NVIDIA/trt-llm-torch-attention-devs -/tensorrt_llm/_torch/modules/attention.py @NVIDIA/trt-llm-torch-attention-devs -/tensorrt_llm/_torch/modules/cross_attention.py @NVIDIA/trt-llm-torch-attention-devs -/tensorrt_llm/_torch/modules/mla.py @NVIDIA/trt-llm-torch-attention-devs -/tensorrt_llm/_torch/modules/qk_norm_attention.py @NVIDIA/trt-llm-torch-attention-devs -/tensorrt_llm/_torch/modules/rotary_embedding.py @NVIDIA/trt-llm-torch-attention-devs -/tests/scripts/cute_dsl_kernels/paged_mqa_logits @NVIDIA/trt-llm-torch-attention-devs -/tests/unittest/_torch/attention @NVIDIA/trt-llm-torch-attention-devs -/tests/unittest/_torch/custom_ops/test_deepseek_v4_q_norm.py @NVIDIA/trt-llm-torch-attention-devs -/tests/unittest/_torch/custom_ops/test_fused_inv_rope_fp8_quant.py @NVIDIA/trt-llm-torch-attention-devs -/tests/unittest/_torch/modules/helix_test_utils.py @NVIDIA/trt-llm-torch-attention-devs -/tests/unittest/_torch/modules/test_mha_helix.py @NVIDIA/trt-llm-torch-attention-devs -/tests/unittest/_torch/modules/test_mla_helix.py @NVIDIA/trt-llm-torch-attention-devs -/tests/unittest/_torch/modules/test_rotary_embedding.py @NVIDIA/trt-llm-torch-attention-devs -/tests/unittest/_torch/thop/parallel/test_indexer_topk.py @NVIDIA/trt-llm-torch-attention-devs -/tests/unittest/_torch/thop/parallel_hw_agnostic/test_fused_qk_norm_rope.py @NVIDIA/trt-llm-torch-attention-devs -/tests/unittest/_torch/thop/parallel_hw_agnostic/test_helix_postprocess.py @NVIDIA/trt-llm-torch-attention-devs -/tests/unittest/_torch/thop/serial/test_fused_cat_fp8.py @NVIDIA/trt-llm-torch-attention-devs - -# ===== MoE ===== -/cpp/tensorrt_llm/deep_ep @NVIDIA/trt-llm-moe-devs -/cpp/tensorrt_llm/kernels/*Moe* @NVIDIA/trt-llm-moe-devs -/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm @NVIDIA/trt-llm-moe-devs -/cpp/tensorrt_llm/kernels/moe* @NVIDIA/trt-llm-moe-devs -/cpp/tensorrt_llm/kernels/moeLoadBalance @NVIDIA/trt-llm-moe-devs -/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe @NVIDIA/trt-llm-moe-devs -/cpp/tensorrt_llm/thop/*Moe* @NVIDIA/trt-llm-moe-devs -/cpp/tensorrt_llm/thop/moe* @NVIDIA/trt-llm-moe-devs -/cpp/tensorrt_llm/cutlass_extensions/include/cutlass_extensions/gemm/kernel/*moe* @NVIDIA/trt-llm-moe-devs -/cpp/tensorrt_llm/kernels/communicationKernels/moe* @NVIDIA/trt-llm-moe-devs -/cpp/tensorrt_llm/kernels/cuteDslKernels/moe* @NVIDIA/trt-llm-moe-devs -/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_gemm_kernels.h @NVIDIA/trt-llm-moe-devs -/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_kernels.h @NVIDIA/trt-llm-moe-devs -/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_util_kernels.h @NVIDIA/trt-llm-moe-devs -/cpp/tensorrt_llm/kernels/internal_cutlass_kernels/include/moe_gemm_kernels.h @NVIDIA/trt-llm-moe-devs -/cpp/tensorrt_llm/kernels/internal_cutlass_kernels/include/moe_kernels.h @NVIDIA/trt-llm-moe-devs -/cpp/tensorrt_llm/kernels/llama4MinLatencyKernels/*Moe* @NVIDIA/trt-llm-moe-devs -/cpp/tensorrt_llm/kernels/marlin/*moe* @NVIDIA/trt-llm-moe-devs -/cpp/tensorrt_llm/kernels/trtllmGenKernels/batchedGemm @NVIDIA/trt-llm-kernels-devs @NVIDIA/trt-llm-moe-devs -/cpp/tensorrt_llm/nanobind/runtime/moe* @NVIDIA/trt-llm-moe-devs -/cpp/tensorrt_llm/runtime/moeLoadBalancer @NVIDIA/trt-llm-moe-devs -/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/moe_as_dense_gemm @NVIDIA/trt-llm-moe-devs -/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4 @NVIDIA/trt-llm-moe-devs -/tensorrt_llm/_torch/custom_ops/cute_dsl_megamoe_custom_op.py @NVIDIA/trt-llm-moe-devs -/tensorrt_llm/_torch/distributed/moe_alltoall.py @NVIDIA/trt-llm-moe-devs -/tensorrt_llm/_torch/expert_statistic.py @NVIDIA/trt-llm-moe-devs -/tensorrt_llm/_torch/modules/fused_moe @NVIDIA/trt-llm-moe-devs -/tensorrt_llm/_torch/modules/fused_shared_expert.py @NVIDIA/trt-llm-moe-devs -/tensorrt_llm/deep_ep @NVIDIA/trt-llm-moe-devs -/tests/microbenchmarks/bench_moe @NVIDIA/trt-llm-moe-devs -/tests/microbenchmarks/bench_moe_comm.py @NVIDIA/trt-llm-moe-devs -/tests/microbenchmarks/compare_moe_comm.py @NVIDIA/trt-llm-moe-devs -/tests/scripts/cute_dsl_kernels/moe_as_dense_gemm @NVIDIA/trt-llm-moe-devs -/tests/scripts/cute_dsl_kernels/moe_workload_generator.py @NVIDIA/trt-llm-moe-devs -/tests/unittest/_torch/modules/dwdp/test_dwdp_fixup_moe_backends.py @NVIDIA/trt-llm-moe-devs -/tests/unittest/_torch/modules/fused_moe @NVIDIA/trt-llm-moe-devs -/tests/unittest/_torch/modules/moe @NVIDIA/trt-llm-moe-devs -/tests/unittest/_torch/modules/test_fused_moe.py @NVIDIA/trt-llm-moe-devs -/tests/unittest/_torch/modules/test_fused_shared_expert.py @NVIDIA/trt-llm-moe-devs -/tests/unittest/_torch/modules/test_moe_*.py @NVIDIA/trt-llm-moe-devs -/tests/unittest/_torch/multi_gpu/test_moe_a2a.py @NVIDIA/trt-llm-moe-devs -/tests/unittest/_torch/thop/parallel/*moe* @NVIDIA/trt-llm-moe-devs -/tests/unittest/_torch/thop/serial/*moe* @NVIDIA/trt-llm-moe-devs -/tests/unittest/bindings/test_bindings_moe.py @NVIDIA/trt-llm-moe-devs - -# ===== VisualGen / AIGV ===== -/cpp/tensorrt_llm/thop/fusedDiT* @NVIDIA/trt-llm-torch-visual-gen-devs -/examples/visual_gen @NVIDIA/trt-llm-torch-visual-gen-devs -/scripts/visualgen_eval @NVIDIA/trt-llm-torch-visual-gen-devs +# TensorRT-LLM Pytorch backend +/tensorrt_llm/_torch @NVIDIA/trt-llm-torch-devs + +## TensorRT-LLM Pytorch - VisualGen /tensorrt_llm/_torch/visual_gen @NVIDIA/trt-llm-torch-visual-gen-devs /tensorrt_llm/_torch/visual_gen/attention_backend @NVIDIA/trt-llm-torch-attention-devs @NVIDIA/trt-llm-torch-visual-gen-devs /tensorrt_llm/_torch/visual_gen/modules/attention.py @NVIDIA/trt-llm-torch-attention-devs @NVIDIA/trt-llm-torch-visual-gen-devs -/tensorrt_llm/media @NVIDIA/trt-llm-torch-visual-gen-devs -/tensorrt_llm/visual_gen @NVIDIA/trt-llm-torch-visual-gen-devs @NVIDIA/trt-llm-runtime-devs +/tensorrt_llm/visual_gen @NVIDIA/trt-llm-llmapi-devs /tests/integration/defs/examples/test_visual_gen.py @NVIDIA/trt-llm-torch-visual-gen-devs /tests/integration/defs/visual_gen @NVIDIA/trt-llm-torch-visual-gen-devs /tests/scripts/perf-sanity/visual_gen @NVIDIA/trt-llm-torch-visual-gen-devs -/tests/unittest/_torch/thop/parallel_hw_agnostic/test_fused_dit* @NVIDIA/trt-llm-torch-visual-gen-devs /tests/unittest/_torch/visual_gen @NVIDIA/trt-llm-torch-visual-gen-devs -/tests/unittest/visual_gen @NVIDIA/trt-llm-torch-visual-gen-devs -# ===== SCAFFOLDING ===== -/tensorrt_llm/scaffolding @WeiHaocheng -/tests/unittest/scaffolding @WeiHaocheng +## TensorRT-LLM Pytorch - Modules +/tensorrt_llm/_torch/modules @NVIDIA/trt-llm-torch-modules -# ===== SELF-MANAGED ===== -/docs/source/features/auto_deploy @NVIDIA/trt-llm-torch-autodeploy-devs @NVIDIA/trt-llm-doc-owners -/examples/auto_deploy @NVIDIA/trt-llm-torch-autodeploy-devs -/scripts/check_auto_deploy_imports.py @NVIDIA/trt-llm-torch-autodeploy-devs -/scripts/check_model_registry.py @NVIDIA/trt-llm-torch-autodeploy-devs +## TensorRT-LLM Pytorch Models +/tensorrt_llm/_torch/models @NVIDIA/trt-llm-torch-models-devs +/examples/models @NVIDIA/trt-llm-torch-models-devs @NVIDIA/trt-llm-doc-owners + +## TensorRT-LLM Pytorch backend - runtime +/tensorrt_llm/_torch/pyexecutor @NVIDIA/trt-llm-torch-runtime-devs +## TensorRT-LLM Pytorch backend - AutoDeploy flow /tensorrt_llm/_torch/auto_deploy @NVIDIA/trt-llm-torch-autodeploy-devs -/tests/integration/defs/accuracy/test_llm_api_autodeploy.py @NVIDIA/trt-llm-torch-autodeploy-devs @NVIDIA/trt-llm-qa -/tests/unittest/_torch/auto_deploy @NVIDIA/trt-llm-torch-autodeploy-devs +/examples/auto_deploy @NVIDIA/trt-llm-torch-autodeploy-devs +/docs/source/features/auto_deploy @NVIDIA/trt-llm-torch-autodeploy-devs @NVIDIA/trt-llm-doc-owners /tests/unittest/auto_deploy @NVIDIA/trt-llm-torch-autodeploy-devs /tests/integration/defs/accuracy/test_llm_api_autodeploy.py @NVIDIA/trt-llm-torch-autodeploy-devs @NVIDIA/trt-llm-qa-function +## TensorRT-LLM Pytorch - Speculative Decoding +/tensorrt_llm/_torch/speculative @NVIDIA/trt-llm-torch-spec-decoding + +## TensorRT-LLM Pytorch - Graph Compiler +/tensorrt_llm/_torch/compilation @NVIDIA/trt-llm-torch-graph-compiler +/tensorrt_llm/_torch/custom_ops @NVIDIA/trt-llm-torch-graph-compiler +/tensorrt_llm/_torch/autotuner.py @NVIDIA/trt-llm-torch-graph-compiler +/tests/unittest/_torch/compilation @NVIDIA/trt-llm-torch-graph-compiler +/tests/unittest/_torch/multi_gpu/test_ar_residual_norm.py @NVIDIA/trt-llm-torch-graph-compiler +/tests/unittest/_torch/multi_gpu/test_user_buffers.py @NVIDIA/trt-llm-torch-graph-compiler +/tests/unittest/_torch/thop/test_custom_ops.py @NVIDIA/trt-llm-torch-graph-compiler +/tests/unittest/_torch/misc/test_autotuner.py @NVIDIA/trt-llm-torch-graph-compiler + ## TensorRT-LLM Pytorch - Attention /tensorrt_llm/_torch/attention_backend @NVIDIA/trt-llm-torch-attention-devs /tensorrt_llm/_torch/modules/attention.py @NVIDIA/trt-llm-torch-attention-devs + +### TensorRT-LLM Pytorch - Models - Gemma +/tensorrt_llm/_torch/models/modeling_gemma3.py @NVIDIA/trt-llm-torch-models-gemma-devs @NVIDIA/trt-llm-torch-models-devs +/tensorrt_llm/_torch/models/modeling_gemma3vl.py @NVIDIA/trt-llm-torch-models-gemma-devs @NVIDIA/trt-llm-torch-models-devs @NVIDIA/trt-llm-multimodal-devs +/tests/unittest/_torch/modeling/test_modeling_gemma3.py @NVIDIA/trt-llm-torch-models-gemma-devs @NVIDIA/trt-llm-torch-models-devs + +### TensorRT-LLM Pytorch - Models - Mistral & Mixtral +/tensorrt_llm/_torch/models/modeling_mistral.py @NVIDIA/trt-llm-torch-models-mistral-devs @NVIDIA/trt-llm-torch-models-devs +/tests/unittest/_torch/modeling/test_modeling_mistral.py @NVIDIA/trt-llm-torch-models-mistral-devs @NVIDIA/trt-llm-torch-models-devs +/tests/unittest/_torch/modeling/test_modeling_mixtral.py @NVIDIA/trt-llm-torch-models-mistral-devs @NVIDIA/trt-llm-torch-models-devs + +### TensorRT-LLM Pytorch - Models - CLIP +/tensorrt_llm/_torch/models/modeling_clip.py @NVIDIA/trt-llm-torch-models-clip-devs @NVIDIA/trt-llm-torch-models-devs @NVIDIA/trt-llm-multimodal-devs +/tests/unittest/_torch/modeling/test_modeling_clip.py @NVIDIA/trt-llm-torch-models-clip-devs @NVIDIA/trt-llm-torch-models-devs @NVIDIA/trt-llm-multimodal-devs + +### TensorRT-LLM Pytorch - Models - Phi +/tensorrt_llm/_torch/models/modeling_phi3.py @NVIDIA/trt-llm-torch-models-phi-devs @NVIDIA/trt-llm-torch-models-devs +/tensorrt_llm/_torch/models/modeling_phi4mm.py @NVIDIA/trt-llm-torch-models-phi-devs @NVIDIA/trt-llm-torch-models-devs @NVIDIA/trt-llm-multimodal-devs +/tests/unittest/_torch/modeling/test_modeling_phi3.py @NVIDIA/trt-llm-torch-models-phi-devs @NVIDIA/trt-llm-torch-models-devs + +### TensorRT-LLM Pytorch - Models - Deepseek +/tensorrt_llm/_torch/models/modeling_deepseekv3.py @NVIDIA/trt-llm-torch-models-deepseek-devs @NVIDIA/trt-llm-torch-models-devs +/tests/unittest/_torch/modeling/test_modeling_deepseek.py @NVIDIA/trt-llm-torch-models-deepseek-devs @NVIDIA/trt-llm-torch-models-devs + +### TensorRT-LLM Pytorch - Models - Llama +/tensorrt_llm/_torch/models/modeling_mllama.py @NVIDIA/trt-llm-torch-models-llama-devs @NVIDIA/trt-llm-torch-models-devs @NVIDIA/trt-llm-multimodal-devs +/tensorrt_llm/_torch/models/modeling_llama.py @NVIDIA/trt-llm-torch-models-llama-devs @NVIDIA/trt-llm-torch-models-devs +/tensorrt_llm/_torch/models/modeling_llama_min_latency.py @NVIDIA/trt-llm-torch-models-llama-devs @NVIDIA/trt-llm-torch-models-devs +/tests/unittest/_torch/modeling/test_modeling_llama.py @NVIDIA/trt-llm-torch-models-llama-devs @NVIDIA/trt-llm-torch-models-devs +/tests/unittest/_torch/modeling/test_modeling_llama_min_latency.py @NVIDIA/trt-llm-torch-models-llama-devs @NVIDIA/trt-llm-torch-models-devs + +### TensorRT-LLM Pytorch - Models - Qwen +/tensorrt_llm/_torch/models/modeling_qwen3_moe.py @NVIDIA/trt-llm-torch-models-qwen-devs @NVIDIA/trt-llm-torch-models-devs +/tensorrt_llm/_torch/models/modeling_qwen3.py @NVIDIA/trt-llm-torch-models-qwen-devs @NVIDIA/trt-llm-torch-models-devs +/tensorrt_llm/_torch/models/modeling_qwen.py @NVIDIA/trt-llm-torch-models-qwen-devs @NVIDIA/trt-llm-torch-models-devs +/tensorrt_llm/_torch/models/modeling_qwen_moe.py @NVIDIA/trt-llm-torch-models-qwen-devs @NVIDIA/trt-llm-torch-models-devs + +### TensorRT-LLM Pytorch - Models - VLMs +/tensorrt_llm/_torch/models/modeling_qwen2vl.py @NVIDIA/trt-llm-torch-models-vlm-devs @NVIDIA/trt-llm-torch-models-devs @NVIDIA/trt-llm-multimodal-devs +/tensorrt_llm/_torch/models/modeling_vila.py @NVIDIA/trt-llm-torch-models-vlm-devs @NVIDIA/trt-llm-torch-models-devs @NVIDIA/trt-llm-multimodal-devs +/tests/unittest/_torch/modeling/test_modeling_vila.py @NVIDIA/trt-llm-torch-models-vlm-devs @NVIDIA/trt-llm-torch-models-devs @NVIDIA/trt-llm-multimodal-devs +/tensorrt_llm/_torch/models/modeling_pixtral.py @NVIDIA/trt-llm-torch-models-vlm-devs @NVIDIA/trt-llm-torch-models-devs @NVIDIA/trt-llm-multimodal-devs +/tests/unittest/_torch/modeling/test_modeling_pixtral.py @NVIDIA/trt-llm-torch-models-vlm-devs @NVIDIA/trt-llm-torch-models-devs @NVIDIA/trt-llm-multimodal-devs + +### TensorRT-LLM Pytorch - Models - Nemotron +/tensorrt_llm/_torch/models/modeling_nanov2vlm.py @NVIDIA/trt-llm-torch-models-nemotron-devs @NVIDIA/trt-llm-torch-models-vlm-devs @NVIDIA/trt-llm-torch-models-devs @NVIDIA/trt-llm-multimodal-devs +/tensorrt_llm/_torch/models/modeling_radio.py @NVIDIA/trt-llm-torch-models-nemotron-devs @NVIDIA/trt-llm-torch-models-vlm-devs @NVIDIA/trt-llm-torch-models-devs @NVIDIA/trt-llm-multimodal-devs +/tensorrt_llm/_torch/models/modeling_nemotron.py @NVIDIA/trt-llm-torch-models-nemotron-devs @NVIDIA/trt-llm-torch-models-devs +/tensorrt_llm/_torch/models/modeling_nemotron_nas.py @NVIDIA/trt-llm-torch-models-nemotron-devs @NVIDIA/trt-llm-torch-models-devs +/tensorrt_llm/_torch/models/modeling_nemotron_h.py @NVIDIA/trt-llm-torch-models-nemotron-devs @NVIDIA/trt-llm-torch-models-devs +/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py @NVIDIA/trt-llm-torch-models-nemotron-devs @NVIDIA/trt-llm-torch-models-devs +/tensorrt_llm/_torch/modules/mamba @NVIDIA/trt-llm-torch-models-nemotron-devs @NVIDIA/trt-llm-torch-models-devs +/tensorrt_llm/_torch/models/checkpoints/hf/nemotron_h_weight_mapper.py @NVIDIA/trt-llm-torch-models-nemotron-devs @NVIDIA/trt-llm-torch-models-devs +/tests/unittest/_torch/modeling/test_modeling_nemotron.py @NVIDIA/trt-llm-torch-models-nemotron-devs @NVIDIA/trt-llm-torch-models-devs +/tests/unittest/_torch/modeling/test_modeling_nemotron_h.py @NVIDIA/trt-llm-torch-models-nemotron-devs @NVIDIA/trt-llm-torch-models-devs +/tests/unittest/_torch/modeling/test_modeling_nemotron_nas.py @NVIDIA/trt-llm-torch-models-nemotron-devs @NVIDIA/trt-llm-torch-models-devs + +## TensorRT-LLM Multimodal - Shared Infrastructure +/tensorrt_llm/inputs/multimodal.py @NVIDIA/trt-llm-multimodal-devs +/tensorrt_llm/inputs/registry.py @NVIDIA/trt-llm-multimodal-devs +/tensorrt_llm/inputs/utils.py @NVIDIA/trt-llm-multimodal-devs +/tensorrt_llm/_torch/models/modeling_multimodal_encoder.py @NVIDIA/trt-llm-multimodal-devs @NVIDIA/trt-llm-torch-models-devs +/tensorrt_llm/_torch/models/modeling_multimodal_utils.py @NVIDIA/trt-llm-multimodal-devs @NVIDIA/trt-llm-torch-models-devs +/tensorrt_llm/quantization/image_processing.py @NVIDIA/trt-llm-multimodal-devs +/cpp/tensorrt_llm/executor/multimodalInput.cpp @NVIDIA/trt-llm-multimodal-devs + +### TensorRT-LLM Multimodal - VLM Models (multimodal-primary ownership) +/tensorrt_llm/_torch/models/modeling_siglip.py @NVIDIA/trt-llm-multimodal-devs @NVIDIA/trt-llm-torch-models-devs +/tensorrt_llm/_torch/models/modeling_llava_next.py @NVIDIA/trt-llm-multimodal-devs @NVIDIA/trt-llm-torch-models-devs +/tensorrt_llm/_torch/models/modeling_hyperclovax.py @NVIDIA/trt-llm-multimodal-devs @NVIDIA/trt-llm-torch-models-devs +/tensorrt_llm/_torch/models/modeling_parakeet.py @NVIDIA/trt-llm-multimodal-devs @NVIDIA/trt-llm-torch-models-devs +/tensorrt_llm/_torch/models/modeling_nemotron_nano.py @NVIDIA/trt-llm-multimodal-devs @NVIDIA/trt-llm-torch-models-devs +/tensorrt_llm/_torch/models/modeling_qwen3vl.py @NVIDIA/trt-llm-multimodal-devs @NVIDIA/trt-llm-torch-models-devs +/tensorrt_llm/_torch/models/modeling_qwen3vl_moe.py @NVIDIA/trt-llm-multimodal-devs @NVIDIA/trt-llm-torch-models-devs +/tensorrt_llm/_torch/models/modeling_qwen3_5.py @NVIDIA/trt-llm-multimodal-devs @NVIDIA/trt-llm-torch-models-devs +/tensorrt_llm/_torch/models/modeling_hunyuan_dense.py @NVIDIA/trt-llm-multimodal-devs @NVIDIA/trt-llm-torch-models-devs +/tensorrt_llm/_torch/models/checkpoints/hf/qwen3vl_weight_mapper.py @NVIDIA/trt-llm-multimodal-devs @NVIDIA/trt-llm-torch-models-devs +/tensorrt_llm/_torch/models/checkpoints/hf/qwen3vl_moe_weight_mapper.py @NVIDIA/trt-llm-multimodal-devs @NVIDIA/trt-llm-torch-models-devs +/tensorrt_llm/_torch/models/checkpoints/hf/qwen3_5_weight_mapper.py @NVIDIA/trt-llm-multimodal-devs @NVIDIA/trt-llm-torch-models-devs +/tensorrt_llm/_torch/models/checkpoints/hf/llava_next_weight_mapper.py @NVIDIA/trt-llm-multimodal-devs @NVIDIA/trt-llm-torch-models-devs + +### TensorRT-LLM Multimodal - Tests +/tests/unittest/_torch/multimodal/ @NVIDIA/trt-llm-multimodal-devs +/tests/unittest/_torch/modeling/test_modeling_multimodal.py @NVIDIA/trt-llm-multimodal-devs +/tests/unittest/_torch/modeling/test_modeling_siglip.py @NVIDIA/trt-llm-multimodal-devs @NVIDIA/trt-llm-torch-models-devs +/tests/unittest/_torch/modeling/test_modeling_parakeet.py @NVIDIA/trt-llm-multimodal-devs @NVIDIA/trt-llm-torch-models-devs +/tests/unittest/others/test_multimodal_registry.py @NVIDIA/trt-llm-multimodal-devs +/tests/unittest/llmapi/apps/*multimodal* @NVIDIA/trt-llm-multimodal-devs @NVIDIA/trt-llm-llmapi-devs + +## TensorRT-LLM - PEFT +/tensorrt_llm/_torch/peft @NVIDIA/trt-llm-torch-peft +/tensorrt_llm/lora_manager.py @NVIDIA/trt-llm-torch-peft +/cpp/tensorrt_llm/batch_manager/peftCacheManager.cpp @NVIDIA/trt-llm-torch-peft +/cpp/include/tensorrt_llm/batch_manager/peftCacheManager.h @NVIDIA/trt-llm-torch-peft +/cpp/tensorrt_llm/runtime/loraCache.cpp @NVIDIA/trt-llm-torch-peft +/cpp/include/tensorrt_llm/runtime/loraCache.h @NVIDIA/trt-llm-torch-peft +/cpp/tensorrt_llm/runtime/loraModule.cpp @NVIDIA/trt-llm-torch-peft +/cpp/include/tensorrt_llm/runtime/loraModule.h @NVIDIA/trt-llm-torch-peft +/cpp/tensorrt_llm/runtime/loraManager.cpp @NVIDIA/trt-llm-torch-peft +/cpp/tensorrt_llm/runtime/loraManager.h @NVIDIA/trt-llm-torch-peft +/cpp/tensorrt_llm/runtime/loraUtils.cpp @NVIDIA/trt-llm-torch-peft +/cpp/tensorrt_llm/runtime/loraUtils.h @NVIDIA/trt-llm-torch-peft + + +## TensorRT-LLM trtllm-bench Reviewers +/tensorrt_llm/bench @NVIDIA/trtllm-bench-reviewers +/tensorrt_llm/commands/bench.py @NVIDIA/trtllm-bench-reviewers +docs/source/performance/perf-benchmarking.md @NVIDIA/trtllm-bench-reviewers + +## TensorRT-LLM LLM API +/tensorrt_llm/llmapi @NVIDIA/trt-llm-llmapi-devs +/tensorrt_llm/executor @NVIDIA/trt-llm-llmapi-devs +/tensorrt_llm/serve @NVIDIA/trt-llm-llmapi-devs +/tensorrt_llm/commands @NVIDIA/trt-llm-llmapi-devs + +## TensorRT-LLM Multimodal - LLM API & Serving +/tensorrt_llm/llmapi/mm_encoder.py @NVIDIA/trt-llm-multimodal-devs @NVIDIA/trt-llm-llmapi-devs + +## TensorRT-LLM Multimodal - Integration Tests +/tests/integration/defs/accuracy/test_llm_api_pytorch_multimodal.py @NVIDIA/trt-llm-multimodal-devs @NVIDIA/trt-llm-qa-function + +## TensorRT-LLM Multimodal - Examples & Docs +/examples/llm-api/quickstart_multimodal.py @NVIDIA/trt-llm-multimodal-devs @NVIDIA/trt-llm-doc-owners +/examples/serve/*multimodal* @NVIDIA/trt-llm-multimodal-devs @NVIDIA/trt-llm-doc-owners +/docs/source/features/multi-modality.md @NVIDIA/trt-llm-multimodal-devs @NVIDIA/trt-llm-doc-owners + ## TensorRT-LLM LLM Disaggregated /examples/disaggregated @NVIDIA/trt-llm-disagg-devs @NVIDIA/trt-llm-doc-owners /examples/disaggregated/slurm/benchmark @NVIDIA/trt-llm-disagg-devs @NVIDIA/trtllm-bench-reviewers @@ -453,41 +252,49 @@ /tests/unittest/api_stability/ @NVIDIA/trt-llm-noncommitted-api-review-committee /tests/unittest/api_stability/references_committed/ @NVIDIA/trt-llm-committed-api-review-committee /tests/unittest/dynamo/ @NVIDIA/trt-llm-dynamo-devs -/triton_backend @NVIDIA/trt-llm-triton-backend-devs - -# ===== DISAGG BLAST-RADIUS CO-OWNS (Tier-1) ===== -# Last-match cross-cutting: adds disagg-devs to shared files that can silently break disagg e2e. -/tensorrt_llm/_torch/pyexecutor/resource_manager.py @NVIDIA/trt-llm-kv-cache-manager-devs @NVIDIA/trt-llm-disagg-devs -/tensorrt_llm/_torch/pyexecutor/py_executor.py @NVIDIA/trt-llm-runtime-devs @NVIDIA/trt-llm-disagg-devs -/tensorrt_llm/_torch/pyexecutor/model_engine.py @NVIDIA/trt-llm-runtime-devs @NVIDIA/trt-llm-disagg-devs -# ===== GOVERNANCE / LEGAL GATES ===== -/.github/CODEOWNERS @NVIDIA/trt-llm-infra-devs @NVIDIA/trt-llm-oss-compliance -/.github/tava_architecture_diagram.md @NVIDIA/trt-llm-TAVA-design-change -/3rdparty/** @NVIDIA/trt-llm-oss-compliance -/ATTRIBUTIONS-*.md @NVIDIA/trt-llm-oss-compliance +## OSS Compliance & Legal - License/Attribution Protection +## IMPORTANT: Changes to any files below may impact legal compliance, attributions, and third-party licenses. +## These files require review from the TRTLLM OSS compliance team before merging to ensure proper attribution +## and license compliance when adding, removing, or changing versions of dependencies. +### License Files /LICENSE @NVIDIA/trt-llm-oss-compliance -/constraints.txt @NVIDIA/trt-llm-oss-compliance -/cpp/CMakeLists.txt @NVIDIA/trt-llm-oss-compliance -/cpp/cmake/** @NVIDIA/trt-llm-oss-compliance -/cpp/conan.lock @NVIDIA/trt-llm-oss-compliance -/cpp/conandata.yml @NVIDIA/trt-llm-oss-compliance -/cpp/conanfile.py @NVIDIA/trt-llm-oss-compliance -/cpp/libnuma_conan.py @NVIDIA/trt-llm-oss-compliance -/docker/common/** @NVIDIA/trt-llm-infra-devs @NVIDIA/trt-llm-oss-compliance -/jenkins/license_cpp.json @NVIDIA/trt-llm-infra-devs @NVIDIA/trt-llm-oss-compliance +/ATTRIBUTIONS-*.md @NVIDIA/trt-llm-oss-compliance +/jenkins/license_cpp.json @NVIDIA/trt-llm-ci-infra-devs @NVIDIA/trt-llm-infra-devs @NVIDIA/trt-llm-oss-compliance + +### Python Dependency Management +/setup.py @NVIDIA/trt-llm-oss-compliance /pyproject.toml @NVIDIA/trt-llm-oss-compliance -/requirements-dev.txt @NVIDIA/trt-llm-oss-compliance /requirements.txt @NVIDIA/trt-llm-oss-compliance -/setup.py @NVIDIA/trt-llm-oss-compliance -/tests/unittest/api_stability/ @NVIDIA/trt-llm-noncommitted-api-review-committee -/tests/unittest/api_stability/references_committed/ @NVIDIA/trt-llm-committed-api-review-committee +/requirements-dev.txt @NVIDIA/trt-llm-oss-compliance + +### C++ Build & Dependency Management +/cpp/CMakeLists.txt @NVIDIA/trt-llm-oss-compliance +/cpp/conanfile.py @NVIDIA/trt-llm-oss-compliance +/cpp/cmake/** @NVIDIA/trt-llm-oss-compliance + +### Third-Party Dependencies +## Any changes to versions, additions, or removals of third-party libraries +/3rdparty/** @NVIDIA/trt-llm-oss-compliance + +### Vendored Third-Party Code (triton-kernels) +## This is a temporary vendored copy of triton-kernels from the Triton project (MIT License). +## Do not accept contributions to this directory - it should only be updated via scripts/vendor_triton_kernels.py +## This can be removed if and when triton-kernels is published as a separate wheel. /triton_kernels/** @NVIDIA/trt-llm-oss-compliance -### Usage telemetry / privacy review -/tensorrt_llm/usage/ @NVIDIA/trt-llm-usage-telemetry-devs -/tests/unittest/usage/ @NVIDIA/trt-llm-usage-telemetry-devs -/tensorrt_llm/usage/llm_args_golden_manifest.json @NVIDIA/trt-llm-usage-telemetry-devs @NVIDIA/trt-llm-oss-compliance @NVIDIA/trt-llm-noncommitted-api-review-committee +### Docker & Installation Scripts +## These scripts install and pin dependency versions +/docker/common/** @NVIDIA/trt-llm-setup-infra-devs @NVIDIA/trt-llm-infra-devs @NVIDIA/trt-llm-oss-compliance + +### TAVA Architecture Diagram +/.github/tava_architecture_diagram.md @NVIDIA/trt-llm-TAVA-design-change + +### CODEOWNERS file itself +/.github/CODEOWNERS @NVIDIA/trt-llm-gh-workflows-infra-devs @NVIDIA/trt-llm-infra-devs @NVIDIA/trt-llm-oss-compliance -# Uncomment only on release branches (e.g. release/0.19): +# The following rule should only be uncommented on release branches (e.g., release/0.19). +# The rule below requires that any PR to release/**/* branches must be approved by at least one member +# of the NVIDIA/trt-llm-release-branch-approval team, regardless of who else approves the PR. +# Without approval from a member of this team, PRs cannot be merged to release branches. # * @NVIDIA/trt-llm-release-branch-approval diff --git a/.github/scripts/label_component.py b/.github/scripts/label_component.py deleted file mode 100644 index 1fa64cf7351b..000000000000 --- a/.github/scripts/label_component.py +++ /dev/null @@ -1,251 +0,0 @@ -#!/usr/bin/env python3 -"""Label pull requests by component, driven by .github/CODEOWNERS. - -For each configured (CODEOWNERS team handle -> label) mapping, this resolves the -effective owners of every file a PR changes using CODEOWNERS last-match-wins -semantics, and applies the mapped label when any changed file is owned by that -team. Labels are only added, never removed. - -Usage: - # single PR (what the GitHub Action runs) - label_component.py --pr 16143 --codeowners .github/CODEOWNERS - # a few PRs - label_component.py --pr 16143 16142 - # sweep every open PR (preview first with --dry-run) - label_component.py --all-open --dry-run - -The repo defaults to the upstream NVIDIA/TensorRT-LLM. CODEOWNERS is read from ---codeowners when given, otherwise fetched from the repo's default branch. The -token comes from GITHUB_TOKEN / GH_TOKEN, falling back to `gh auth token`. -""" - -import argparse -import os -import re -import subprocess -import sys - -import requests - -GITHUB_API_URL = "https://api.github.com" -DEFAULT_REPO = "NVIDIA/TensorRT-LLM" - -# CODEOWNERS team handle (lower-cased) -> label to apply. -# Extend this dict to cover more components. -COMPONENT_LABELS = { - "@nvidia/trt-llm-torch-visual-gen-devs": "VisualGen", -} - - -# --- CODEOWNERS parsing / matching --------------------------------------- - - -def parse_codeowners(text): - """Parse CODEOWNERS into an ordered list of (compiled_regex, owners).""" - rules = [] - for raw in text.splitlines(): - line = raw.split("#", 1)[0].strip() - if not line: - continue - parts = line.split() - pattern, owners = parts[0], [o.lower() for o in parts[1:]] - rules.append((_pattern_to_regex(pattern), owners)) - return rules - - -def _pattern_to_regex(pattern): - """Translate a CODEOWNERS (gitignore-style) pattern to a regex. - - '*' matches within a path segment, '**' crosses segments, and a directory - pattern matches everything beneath it. All CODEOWNERS patterns here are - root-anchored. - """ - body = re.escape(pattern.strip("/")) - body = body.replace(r"\*\*", ".*").replace(r"\*", "[^/]*") - return re.compile(rf"(?:{body})(?:/.*)?$") - - -def owners_for_path(path, rules): - """Effective CODEOWNERS owners for a path (last matching rule wins).""" - owners = [] - for regex, rule_owners in rules: - if regex.match(path): - owners = rule_owners - return owners - - -def labels_for_files(files, rules, component_labels=COMPONENT_LABELS): - labels = set() - for path in files: - owners = owners_for_path(path, rules) - for team, label in component_labels.items(): - if team in owners: - labels.add(label) - return labels - - -# --- GitHub access ------------------------------------------------------- - - -def resolve_token(): - for var in ("GITHUB_TOKEN", "GH_TOKEN"): - if os.environ.get(var): - return os.environ[var] - try: - return subprocess.check_output(["gh", "auth", "token"], text=True).strip() - except (OSError, subprocess.CalledProcessError): - raise SystemExit("No token found: set GITHUB_TOKEN / GH_TOKEN, or run `gh auth login`.") - - -def make_session(token): - session = requests.Session() - session.headers.update( - { - "Accept": "application/vnd.github+json", - "X-GitHub-Api-Version": "2022-11-28", - "User-Agent": "trtllm-label-component/1.0", - "Authorization": f"token {token}", - } - ) - return session - - -def load_codeowners(session, repo, path): - if path: - with open(path, encoding="utf-8") as fh: - return parse_codeowners(fh.read()) - r = session.get( - f"{GITHUB_API_URL}/repos/{repo}/contents/.github/CODEOWNERS", - headers={"Accept": "application/vnd.github.raw"}, - timeout=30, - ) - r.raise_for_status() - return parse_codeowners(r.text) - - -def iter_open_prs(session, repo, limit=None): - """Yield (number, existing_labels) for open PRs; labels come free here.""" - page, seen = 1, 0 - while True: - r = session.get( - f"{GITHUB_API_URL}/repos/{repo}/pulls", - params={ - "state": "open", - "per_page": 100, - "page": page, - "sort": "created", - "direction": "desc", - }, - timeout=30, - ) - r.raise_for_status() - batch = r.json() - if not batch: - return - for pr in batch: - yield pr["number"], {lbl["name"] for lbl in pr.get("labels", [])} - seen += 1 - if limit and seen >= limit: - return - page += 1 - - -def get_changed_files(session, repo, pr_number): - files, page = [], 1 - while True: - r = session.get( - f"{GITHUB_API_URL}/repos/{repo}/pulls/{pr_number}/files", - params={"per_page": 100, "page": page}, - timeout=30, - ) - r.raise_for_status() - batch = r.json() - if not batch: - break - files.extend(f["filename"] for f in batch) - page += 1 - return files - - -def get_pr_labels(session, repo, pr_number): - r = session.get(f"{GITHUB_API_URL}/repos/{repo}/issues/{pr_number}", timeout=30) - r.raise_for_status() - return {lbl["name"] for lbl in r.json().get("labels", [])} - - -def add_labels(session, repo, pr_number, labels): - r = session.post( - f"{GITHUB_API_URL}/repos/{repo}/issues/{pr_number}/labels", - json={"labels": labels}, - timeout=30, - ) - r.raise_for_status() - - -def process_pr(session, repo, pr_number, rules, existing_labels, dry_run): - """Return the labels added (or that would be added). Empty if none.""" - files = get_changed_files(session, repo, pr_number) - wanted = labels_for_files(files, rules) - to_add = sorted(wanted - existing_labels) - if not to_add: - return [] - if dry_run: - print(f"PR #{pr_number}: would add {to_add}") - else: - add_labels(session, repo, pr_number, to_add) - print(f"PR #{pr_number}: added {to_add}") - return to_add - - -# --- CLI ----------------------------------------------------------------- - - -def parse_args(argv): - ap = argparse.ArgumentParser(description=__doc__.splitlines()[0]) - ap.add_argument("--repo", default=DEFAULT_REPO, help=f"owner/name (default: {DEFAULT_REPO})") - target = ap.add_mutually_exclusive_group(required=True) - target.add_argument("--pr", type=int, nargs="+", metavar="N", help="label these PR number(s)") - target.add_argument("--all-open", action="store_true", help="label every open PR in the repo") - ap.add_argument( - "--codeowners", - metavar="PATH", - help="local CODEOWNERS file; if omitted, fetched from the repo's default branch", - ) - ap.add_argument( - "--limit", type=int, metavar="N", help="with --all-open, cap the number of PRs scanned" - ) - ap.add_argument( - "--dry-run", action="store_true", help="report what would change without adding labels" - ) - return ap.parse_args(argv) - - -def main(argv=None): - args = parse_args(argv) - session = make_session(resolve_token()) - rules = load_codeowners(session, args.repo, args.codeowners) - - if args.pr: - targets = [(n, get_pr_labels(session, args.repo, n)) for n in args.pr] - else: - targets = iter_open_prs(session, args.repo, args.limit) - - scanned, labeled, failed = 0, 0, 0 - for number, existing in targets: - scanned += 1 - # Isolate per-PR failures so one bad PR (transient 5xx, missing label) - # does not abort an --all-open sweep. - try: - if process_pr(session, args.repo, number, rules, existing, args.dry_run): - labeled += 1 - except requests.HTTPError as exc: - failed += 1 - print(f"PR #{number}: failed ({exc})", file=sys.stderr) - - verb = "would be labeled" if args.dry_run else "labeled" - print(f"Scanned {scanned} PR(s); {labeled} {verb}; {failed} failed.") - return 0 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/.github/workflows/blossom-ci.yml b/.github/workflows/blossom-ci.yml index f5a60658d8dd..3158cfbe0e15 100644 --- a/.github/workflows/blossom-ci.yml +++ b/.github/workflows/blossom-ci.yml @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -57,7 +57,6 @@ jobs: "anish-shanbhag", "arekay", "arysef", - "asfiyab-nvidia", "aswinvisva", "athena-nv", "atrifex", @@ -76,14 +75,11 @@ jobs: "bo-nv", "bobboli", "Boreas618", - "BowenFu", - "BoyueZ", "brb-nv", "brnguyen2", "byshiue", "CarstyYou", "cascade812", - "caitlinw", "chang-l", "chenfeiz0326", "cherichy", @@ -117,7 +113,6 @@ jobs: "elvischenv", "EmmaQiaoCh", "eopXD", - "erictsai-nv", "esha-nvidia", "etz-lmn", "evezhier", @@ -137,11 +132,9 @@ jobs: "govind-ramnarayan", "greg-kwasniewski1", "guangyunh-nv", - "GuanhuaWang2001", "guqiqi", "h-guo18", "HandongLi-01", - "haow-nv", "hchings", "hello-11", "heyuhhh", @@ -155,7 +148,6 @@ jobs: "indrajit96", "inocsin", "ISEEKYAN", - "ishovkun", "ixlmar", "IzzyPutterman", "Jackch-NV", @@ -174,7 +166,6 @@ jobs: "Jie-Fang", "jiefangz-nv", "jieli-matrix", - "jingyu-ml", "JintaoPengCS", "jinyangyuan-nvidia", "jinzh-nvidia", @@ -187,14 +178,12 @@ jobs: "JunyiXu-nv", "JyChang012", "kaiyux", - "Kambili", "kanghui0204", "karljang", "karthikvetrivel", "katec846", "Kefeng-Duan", "KingsleyLiu-NV", - "KleinBlueC", "kris1025", "KrishnanPrash", "kunlunl", @@ -204,7 +193,6 @@ jobs: "lancelly", "LarryXFly", "latency1024", - "leo0519", "leslie-fang25", "lfr-0531", "liji-nv", @@ -219,7 +207,6 @@ jobs: "lkomali", "longcheng-nv", "longlee0622", - "lori-ren", "lowsfer", "lucaslie", "lucifer1004", @@ -229,7 +216,6 @@ jobs: "MatthiasKohl", "mayani-nv", "meenchen", - "mgluhovskoi", "mikeiovine", "milesial", "MinaHuai", @@ -276,7 +262,6 @@ jobs: "pengbowang-nv", "PerkzZheng", "poweiw", - "pranav-nvidia", "qiangxu1996", "qiaoxj07", "QiJune", @@ -310,7 +295,6 @@ jobs: "shuyixiong", "shyeh25", "SimengLiu-nv", - "siyidNV", "sklevtsov-nvidia", "StanleySun639", "stnie", @@ -357,7 +341,6 @@ jobs: "viraatc", "Wanli-Jiang", "WeiHaocheng", - "weiminwang-nv", "weireweire", "wenmingw", "wili-65535", @@ -367,7 +350,6 @@ jobs: "wyw1267", "xavier-nvidia", "xd-nv", - "xguannv", "xiaoweiw-nv", "xinhe-nv", "xmchen1987", diff --git a/.github/workflows/label_component_pr.yml b/.github/workflows/label_component_pr.yml deleted file mode 100644 index 5d584eb237cd..000000000000 --- a/.github/workflows/label_component_pr.yml +++ /dev/null @@ -1,42 +0,0 @@ -name: Label Component for PR - -on: - pull_request_target: - types: [opened, reopened] - -permissions: - contents: read - pull-requests: write - -jobs: - label-component: - runs-on: ubuntu-latest - if: github.repository == 'NVIDIA/TensorRT-LLM' - # This workflow is advisory: it must never turn a PR check red. Every step - # is continue-on-error, so the check is always green even if setup or - # labeling fails. - steps: - - name: Checkout base repository - continue-on-error: true - uses: actions/checkout@v6 - with: - persist-credentials: false - - - name: Set up Python - continue-on-error: true - uses: actions/setup-python@v6 - with: - python-version: '3.x' - - - name: Install dependencies - continue-on-error: true - run: pip install requests - - - name: Label PR by component - continue-on-error: true - env: - GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} - run: >- - python .github/scripts/label_component.py - --pr ${{ github.event.pull_request.number }} - --codeowners .github/CODEOWNERS diff --git a/.github/workflows/llm-api-compatibility.yml b/.github/workflows/llm-api-compatibility.yml index e93fea709ef9..5c727b659e7c 100644 --- a/.github/workflows/llm-api-compatibility.yml +++ b/.github/workflows/llm-api-compatibility.yml @@ -33,8 +33,6 @@ jobs: const referencePrefixes = [ 'tests/unittest/api_stability/references/', 'tests/unittest/api_stability/references_committed/', - 'tensorrt_llm/llmapi/llm_args.py', - 'tensorrt_llm/usage/llm_args_golden_manifest.json', ]; // Keep these names in sync with the PR template and API-change guide. const compatibleApiLabel = 'api-compatible'; diff --git a/.gitignore b/.gitignore index 47e39a1b2d71..8d39054480f5 100644 --- a/.gitignore +++ b/.gitignore @@ -114,9 +114,6 @@ enroot/tensorrt_llm.devel.sqsh # MacOSX Files .DS_Store -# stress test aiperf output artifacts -tests/integration/defs/stress_test/artifacts/ - # Agent related files .claude/agent-memory/ .claude/agent-tests/perf-test-sync/report.html diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 452c3be57ea1..5b0072c49fad 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -302,25 +302,6 @@ common-files: &common_files | tensorrt_llm/_torch/cute_dsl_kernels/blackwell/custom_pipeline.py | tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dense_blockscaled_gemm_persistent.py | tensorrt_llm/_torch/cute_dsl_kernels/blackwell/utils.py | - tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/__init__.py | - tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/blocked_scale.py | - tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/contract.py | - tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/custom_ext.py | - tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/dynamic_mainloop.py | - tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/epilogue_refactor.py | - tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/fc1_fc2_fuse_sched.py | - tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/grid_sync.py | - tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/iket_compat.py | - tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/kernel_fc12.py | - tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/megamoe_constants.py | - tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/megamoe_kernel.py | - tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/moe_persistent_scheduler.py | - tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/moe_utils.py | - tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/ptx_helpers.py | - tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/sf_swizzle.py | - tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/sym_buffer.py | - tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/token_comm.py | - tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/topk_reduce.py | tensorrt_llm/_torch/cute_dsl_utils.py | tensorrt_llm/_torch/debug/__init__.py | tensorrt_llm/_torch/debug/debug_hook.py | @@ -1031,6 +1012,7 @@ common-files: &common_files | tests/unittest/_torch/ray_orchestrator/single_gpu/test_cache_transceiver_comm.py | tests/unittest/_torch/sampler/test_beam_search.py | tests/unittest/_torch/sampler/test_best_of_n.py | + tests/unittest/_torch/sampler/test_torch_multi_arange.py | tests/unittest/_torch/sampler/test_trtllm_sampler.py | tests/unittest/_torch/speculative/test_draft_target.py | tests/unittest/_torch/speculative/test_draft_token_tree_sampling.py | @@ -1045,7 +1027,6 @@ common-files: &common_files | tests/unittest/_torch/speculative/test_torch_rejection_sampling.py | tests/unittest/_torch/speculative/test_user_provided.py | tests/unittest/_torch/test_connector.py | - tests/unittest/_torch/test_torch_multi_arange.py | tests/unittest/_torch/thop/parallel/deep_gemm_tests.py | tests/unittest/_torch/thop/parallel/test_causal_conv1d_op.py | tests/unittest/_torch/thop/parallel/test_cublas_mm.py | @@ -1296,6 +1277,7 @@ common-files: &common_files | tests/unittest/trt/model/test_nemotron_nas.py | tests/unittest/trt/model/test_phi.py | tests/unittest/trt/model/test_unet.py | + tests/unittest/trt/model_api/profile_utils.py | tests/unittest/trt/model_api/test_model_api_multi_gpu.py | tests/unittest/trt/model_api/test_model_level_api.py | tests/unittest/trt/model_api/test_model_quantization.py | @@ -1325,22 +1307,55 @@ common-files: &common_files | tests/unittest/utils/test_util.py | tests/unittest/utils/torch_ref.py | tests/unittest/utils/util.py | + triton_backend/all_models/disaggregated_serving/disaggregated_serving_bls/1/model.py | + triton_backend/all_models/gpt/postprocessing/1/model.py | + triton_backend/all_models/gpt/preprocessing/1/model.py | + triton_backend/all_models/gpt/tensorrt_llm/1/model.py | + triton_backend/all_models/inflight_batcher_llm/postprocessing/1/model.py | + triton_backend/all_models/inflight_batcher_llm/preprocessing/1/model.py | + triton_backend/all_models/inflight_batcher_llm/tensorrt_llm/1/model.py | + triton_backend/all_models/inflight_batcher_llm/tensorrt_llm_bls/1/lib/decode.py | + triton_backend/all_models/inflight_batcher_llm/tensorrt_llm_bls/1/lib/triton_decoder.py | + triton_backend/all_models/inflight_batcher_llm/tensorrt_llm_bls/1/model.py | triton_backend/all_models/llmapi/tensorrt_llm/1/helpers.py | triton_backend/all_models/llmapi/tensorrt_llm/1/model.py | + triton_backend/all_models/multimodal/multimodal_encoders/1/model.py | + triton_backend/all_models/multimodal/multimodal_encoders/1/multimodal_utils.py | + triton_backend/all_models/tests/test_decode.py | triton_backend/all_models/tests/test_llmapi_python_backend.py | + triton_backend/all_models/tests/test_multi_image_preprocess.py | + triton_backend/all_models/tests/test_multimodal_encoders.py | + triton_backend/all_models/tests/test_python_backend.py | + triton_backend/all_models/tests/test_triton_decoder.py | + triton_backend/all_models/whisper/whisper_bls/1/fbank.py | + triton_backend/all_models/whisper/whisper_bls/1/model.py | + triton_backend/all_models/whisper/whisper_bls/1/tokenizer.py | + triton_backend/ci/L0_backend_trtllm/base_metrics_verification_tests.py | + triton_backend/ci/L0_backend_trtllm/custom_metrics_verification_tests.py | + triton_backend/inflight_batcher_llm/client/__init__.py | + triton_backend/inflight_batcher_llm/client/e2e_grpc_speculative_decoding_client.py | + triton_backend/inflight_batcher_llm/client/end_to_end_grpc_client.py | + triton_backend/inflight_batcher_llm/client/inflight_batcher_llm_client.py | triton_backend/scripts/launch_triton_server.py | triton_backend/tools/__init__.py | triton_backend/tools/fill_template.py | + triton_backend/tools/gpt/benchmark_core_model.py | + triton_backend/tools/gpt/client.py | + triton_backend/tools/gpt/client_async.py | + triton_backend/tools/gpt/end_to_end_test.py | + triton_backend/tools/gpt/gen_input_data.py | triton_backend/tools/inflight_batcher_llm/benchmark_core_model.py | triton_backend/tools/inflight_batcher_llm/end_to_end_test.py | triton_backend/tools/inflight_batcher_llm/speculative_decoding_test.py | triton_backend/tools/inflight_batcher_llm/test_max_queue_size.py | triton_backend/tools/llmapi_client.py | + triton_backend/tools/multimodal/client.py | triton_backend/tools/tests/__init__.py | triton_backend/tools/tests/test_fill_template.py | triton_backend/tools/tests/test_llmapi_cancel.py | triton_backend/tools/utils/__init__.py | - triton_backend/tools/utils/utils.py + triton_backend/tools/utils/utils.py | + triton_backend/tools/whisper/client.py )$ # Used by ruff hooks: main ruff (exclude: *legacy_files) and @@ -1643,25 +1658,6 @@ legacy-files: &legacy_files | tensorrt_llm/_torch/cute_dsl_kernels/blackwell/custom_pipeline.py | tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dense_blockscaled_gemm_persistent.py | tensorrt_llm/_torch/cute_dsl_kernels/blackwell/utils.py | - tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/__init__.py | - tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/blocked_scale.py | - tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/contract.py | - tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/custom_ext.py | - tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/dynamic_mainloop.py | - tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/epilogue_refactor.py | - tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/fc1_fc2_fuse_sched.py | - tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/grid_sync.py | - tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/iket_compat.py | - tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/kernel_fc12.py | - tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/megamoe_constants.py | - tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/megamoe_kernel.py | - tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/moe_persistent_scheduler.py | - tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/moe_utils.py | - tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/ptx_helpers.py | - tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/sf_swizzle.py | - tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/sym_buffer.py | - tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/token_comm.py | - tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/topk_reduce.py | tensorrt_llm/_torch/cute_dsl_utils.py | tensorrt_llm/_torch/debug/__init__.py | tensorrt_llm/_torch/debug/debug_hook.py | @@ -2372,6 +2368,7 @@ legacy-files: &legacy_files | tests/unittest/_torch/ray_orchestrator/single_gpu/test_cache_transceiver_comm.py | tests/unittest/_torch/sampler/test_beam_search.py | tests/unittest/_torch/sampler/test_best_of_n.py | + tests/unittest/_torch/sampler/test_torch_multi_arange.py | tests/unittest/_torch/sampler/test_trtllm_sampler.py | tests/unittest/_torch/speculative/test_draft_target.py | tests/unittest/_torch/speculative/test_draft_token_tree_sampling.py | @@ -2386,7 +2383,6 @@ legacy-files: &legacy_files | tests/unittest/_torch/speculative/test_torch_rejection_sampling.py | tests/unittest/_torch/speculative/test_user_provided.py | tests/unittest/_torch/test_connector.py | - tests/unittest/_torch/test_torch_multi_arange.py | tests/unittest/_torch/thop/parallel/deep_gemm_tests.py | tests/unittest/_torch/thop/parallel/test_causal_conv1d_op.py | tests/unittest/_torch/thop/parallel/test_cublas_mm.py | @@ -2637,6 +2633,7 @@ legacy-files: &legacy_files | tests/unittest/trt/model/test_nemotron_nas.py | tests/unittest/trt/model/test_phi.py | tests/unittest/trt/model/test_unet.py | + tests/unittest/trt/model_api/profile_utils.py | tests/unittest/trt/model_api/test_model_api_multi_gpu.py | tests/unittest/trt/model_api/test_model_level_api.py | tests/unittest/trt/model_api/test_model_quantization.py | @@ -2666,33 +2663,64 @@ legacy-files: &legacy_files | tests/unittest/utils/test_util.py | tests/unittest/utils/torch_ref.py | tests/unittest/utils/util.py | + triton_backend/all_models/disaggregated_serving/disaggregated_serving_bls/1/model.py | + triton_backend/all_models/gpt/postprocessing/1/model.py | + triton_backend/all_models/gpt/preprocessing/1/model.py | + triton_backend/all_models/gpt/tensorrt_llm/1/model.py | + triton_backend/all_models/inflight_batcher_llm/postprocessing/1/model.py | + triton_backend/all_models/inflight_batcher_llm/preprocessing/1/model.py | + triton_backend/all_models/inflight_batcher_llm/tensorrt_llm/1/model.py | + triton_backend/all_models/inflight_batcher_llm/tensorrt_llm_bls/1/lib/decode.py | + triton_backend/all_models/inflight_batcher_llm/tensorrt_llm_bls/1/lib/triton_decoder.py | + triton_backend/all_models/inflight_batcher_llm/tensorrt_llm_bls/1/model.py | triton_backend/all_models/llmapi/tensorrt_llm/1/helpers.py | triton_backend/all_models/llmapi/tensorrt_llm/1/model.py | + triton_backend/all_models/multimodal/multimodal_encoders/1/model.py | + triton_backend/all_models/multimodal/multimodal_encoders/1/multimodal_utils.py | + triton_backend/all_models/tests/test_decode.py | triton_backend/all_models/tests/test_llmapi_python_backend.py | + triton_backend/all_models/tests/test_multi_image_preprocess.py | + triton_backend/all_models/tests/test_multimodal_encoders.py | + triton_backend/all_models/tests/test_python_backend.py | + triton_backend/all_models/tests/test_triton_decoder.py | + triton_backend/all_models/whisper/whisper_bls/1/fbank.py | + triton_backend/all_models/whisper/whisper_bls/1/model.py | + triton_backend/all_models/whisper/whisper_bls/1/tokenizer.py | + triton_backend/ci/L0_backend_trtllm/base_metrics_verification_tests.py | + triton_backend/ci/L0_backend_trtllm/custom_metrics_verification_tests.py | + triton_backend/inflight_batcher_llm/client/__init__.py | + triton_backend/inflight_batcher_llm/client/e2e_grpc_speculative_decoding_client.py | + triton_backend/inflight_batcher_llm/client/end_to_end_grpc_client.py | + triton_backend/inflight_batcher_llm/client/inflight_batcher_llm_client.py | triton_backend/scripts/launch_triton_server.py | triton_backend/tools/__init__.py | triton_backend/tools/fill_template.py | + triton_backend/tools/gpt/benchmark_core_model.py | + triton_backend/tools/gpt/client.py | + triton_backend/tools/gpt/client_async.py | + triton_backend/tools/gpt/end_to_end_test.py | + triton_backend/tools/gpt/gen_input_data.py | triton_backend/tools/inflight_batcher_llm/benchmark_core_model.py | triton_backend/tools/inflight_batcher_llm/end_to_end_test.py | triton_backend/tools/inflight_batcher_llm/speculative_decoding_test.py | triton_backend/tools/inflight_batcher_llm/test_max_queue_size.py | triton_backend/tools/llmapi_client.py | + triton_backend/tools/multimodal/client.py | triton_backend/tools/tests/__init__.py | triton_backend/tools/tests/test_fill_template.py | triton_backend/tools/tests/test_llmapi_cancel.py | triton_backend/tools/utils/__init__.py | - triton_backend/tools/utils/utils.py + triton_backend/tools/utils/utils.py | + triton_backend/tools/whisper/client.py )$ # <<<< END AUTO-GENERATED >>>> # Files to be subjected to static analysis static-analysis-files: &static_analysis_files | (?x)^( - tensorrt_llm/_torch/pyexecutor/sampler/sampler.py | - tensorrt_llm/_torch/pyexecutor/sampler/sampling_utils.py | - tensorrt_llm/_torch/pyexecutor/sampler/ops/interface.py | - tensorrt_llm/_torch/pyexecutor/sampler/ops/flashinfer.py | - tensorrt_llm/_torch/pyexecutor/sampler/ops/vanilla.py | + tensorrt_llm/_torch/pyexecutor/sampler.py | + tensorrt_llm/_torch/pyexecutor/sampling_utils.py | + tensorrt_llm/_torch/pyexecutor/sampling_utils_flashinfer.py | tests/unittest/_torch/sampler/test_torch_sampler.py | tests/unittest/_torch/sampler/test_beam_search.py | tests/unittest/_torch/sampler/test_beam_search_util.py | diff --git a/3rdparty/fetch_content.json b/3rdparty/fetch_content.json index d86ebb57fc46..6a4679db5262 100644 --- a/3rdparty/fetch_content.json +++ b/3rdparty/fetch_content.json @@ -32,7 +32,7 @@ { "name": "deepgemm", "git_repository": "https://github.com/deepseek-ai/DeepGEMM", - "git_tag": "245dc5d6a5fe344c61505fe71011d203141d4479", + "git_tag": "c491439ed5966833d56883ca302b6f72e74f8105", "git_submodules_recurse": true, "source_subdir": "dont-add-this-project-with-add-subdirectory" }, diff --git a/AGENTS.md b/AGENTS.md index 20fa71f70f7f..5fe19e95be6a 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -13,8 +13,6 @@ Python and C++ codebase supporting TensorRT engine-based and PyTorch-based execu - `git commit -s` (DCO sign-off required). Never attribute AI tools in sign-off line. Always rely on `git` to do the sign off instead of directly adding sign off in commit message. - Do not add co-authors to the git commit message unless explicitly instructed to do so by the user. - `pre-commit` hooks run on commit — if files are modified by hooks, re-stage and commit again -- LLM args or nested-config changes must run `python3 scripts/generate_llm_args_golden_manifest.py` and commit - `tensorrt_llm/usage/llm_args_golden_manifest.json`; new fields require telemetry/privacy CODEOWNER approval - PR title format: `[JIRA/NVBUG/None][type] description` (e.g., `[TRTLLM-5516][perf] optimize cuda graph padding`) - Set `LLM_MODELS_ROOT` env var when running tests that need model weights @@ -85,7 +83,7 @@ HuggingFace Model → LLM API → Executor (PyTorch/AutoDeploy/TensorRT) | `tensorrt_llm/executor/executor.py` | Execution abstraction (`GenerationExecutor`) | | `tensorrt_llm/models/automodel.py` | Auto-discovery and model registry | | `tensorrt_llm/_torch/models/` | PyTorch backend model implementations (distinct from `models/` used by TensorRT backend) | -| `tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md` | Attention, MLA, backend families, sparse backends, metadata contracts, and KV-cache behavior - **read before modifying `tensorrt_llm/_torch/modules/attention.py`, `tensorrt_llm/_torch/modules/mla.py`, or `tensorrt_llm/_torch/attention_backend/`** | +| `tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md` | Attention, MLA, backend families, sparse backends, metadata contracts, and KV-cache behavior - **read before modifying `tensorrt_llm/_torch/modules/attention.py` or `tensorrt_llm/_torch/attention_backend/`** | | `tensorrt_llm/_torch/modules/fused_moe/MOE_DEVELOPER_GUIDE.md` | MoE architecture, backends, communication, development patterns — **read before modifying MoE code** | | `CODING_GUIDELINES.md` | C++ and Python coding standards (referenced throughout, must read before contributing) | @@ -112,7 +110,6 @@ Key entry points: - Serving CLI: `trtllm-serve --model --visual_gen_args `. Key files: -- `tensorrt_llm/_torch/visual_gen/ENGINEERING_CRITERIA.md`: **Engineering criteria for any change under `tensorrt_llm/visual_gen/` or `tensorrt_llm/_torch/visual_gen/`** — API discipline, feature/test/lossy-vs-lossless requirements, examples & docs rules. Read before modifying anything in those trees. - `tensorrt_llm/visual_gen/`: VisualGen public Python API. **User-facing surface — before modifying anything here, pause and confirm with the user that a public API change is actually intended; do not infer it from the surrounding task.** - `tensorrt_llm/_torch/visual_gen/`: VisualGen internal implementation. All non-user-facing code belongs here. @@ -147,9 +144,8 @@ Key files: The `gh` CLI uses `~/.config/gh` by default for authentication. Different GitHub hosts or forks may require a different config directory. **Before running any `gh` command** (e.g., `gh pr create`, `gh api`, `gh pr comment`): 1. Check if the user has specified a custom `GH_CONFIG_DIR` (e.g., in `CLAUDE.local.md` or environment). If so, use it. -2. If not explicitly set, default to `~/.config/gh`; do not ask for confirmation. +2. If not explicitly set, **ask the user** whether the default `~/.config/gh` is correct or if a different directory should be used. This is especially relevant when the PR target is a fork (e.g., `nv-auto-deploy/TensorRT-LLM`) rather than `NVIDIA/TensorRT-LLM`. 3. Prefix all `gh` commands with the resolved config dir: `GH_CONFIG_DIR= gh ...` -4. If the command fails due to missing authentication or the wrong GitHub host/account, report the failure and ask for the correct `GH_CONFIG_DIR`. ## CI / Testing diff --git a/ATTRIBUTIONS-Python.md b/ATTRIBUTIONS-Python.md index e43a78cafda0..0ba92e7c6bc8 100644 --- a/ATTRIBUTIONS-Python.md +++ b/ATTRIBUTIONS-Python.md @@ -5261,7 +5261,7 @@ For more information, please refer to - `Tracker`: https://github.com/tox-dev/py-filelock/issues -## flashinfer-python (0.6.14) +## flashinfer-python (0.6.12) ### Licenses License: `Apache-2.0` diff --git a/CODING_GUIDELINES.md b/CODING_GUIDELINES.md index 8e4d0b5be7bf..33031ece5247 100644 --- a/CODING_GUIDELINES.md +++ b/CODING_GUIDELINES.md @@ -551,9 +551,6 @@ When defining any user-facing configuration classes (particularly `LlmArgs` or a - Prefer `PositiveInt`, `NonNegativeInt`, `NonNegativeFloat`, `PositiveFloat`, `Field(gt=0)`, `Field(ge=0)`, etc. for numeric constraints instead of defining custom validators - Use `Field(min_length=1)` to enforce minimum length of a list -- After changing LLM args or nested configs, run `python3 scripts/generate_llm_args_golden_manifest.py` and commit - `tensorrt_llm/usage/llm_args_golden_manifest.json`; new fields require telemetry/privacy CODEOWNER approval. - **Validation:** - Use `@field_validator` and `@model_validator` instead of manual `validate()` or `is_valid()` methods - Raise `ValueError` instead of using assertions diff --git a/README.md b/README.md index 6ba10b8a5adb..1f9e7fa262b7 100644 --- a/README.md +++ b/README.md @@ -8,9 +8,9 @@ TensorRT LLM [![Ask DeepWiki](https://deepwiki.com/badge.svg)](https://deepwiki.com/NVIDIA/TensorRT-LLM) [![python](https://img.shields.io/badge/python-3.12-green)](https://www.python.org/downloads/release/python-3123/) [![python](https://img.shields.io/badge/python-3.10-green)](https://www.python.org/downloads/release/python-31012/) -[![cuda](https://img.shields.io/badge/cuda-13.2.1-green)](https://developer.nvidia.com/cuda-downloads) -[![torch](https://img.shields.io/badge/torch-2.11.0-green)](https://pytorch.org) -[![version](https://img.shields.io/badge/release-1.3.0rc21-green)](https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/version.py) +[![cuda](https://img.shields.io/badge/cuda-13.1.1-green)](https://developer.nvidia.com/cuda-downloads) +[![torch](https://img.shields.io/badge/torch-2.10.0-green)](https://pytorch.org) +[![version](https://img.shields.io/badge/release-1.3.0rc18-green)](https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/version.py) [![license](https://img.shields.io/badge/license-Apache%202-blue)](https://github.com/NVIDIA/TensorRT-LLM/blob/main/LICENSE) [Architecture](https://nvidia.github.io/TensorRT-LLM/developer-guide/overview.html)   |   [Performance](https://nvidia.github.io/TensorRT-LLM/developer-guide/perf-overview.html)   |   [Examples](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html)   |   [Documentation](https://nvidia.github.io/TensorRT-LLM/)   |   [Roadmap](https://github.com/NVIDIA/TensorRT-LLM/issues?q=is%3Aissue%20state%3Aopen%20label%3Aroadmap) @@ -22,12 +22,6 @@ TensorRT LLM -* [07/17] DeepSeek-V4 on NVIDIA Blackwell: Model-Specific and Agentic-Workload Optimizations in TensorRT LLM -✨ [➡️ link](https://github.com/NVIDIA/TensorRT-LLM/blob/main/docs/source/blogs/tech_blog/blog26_DeepSeek_V4_on_NVIDIA_Blackwell_Model_Specific_and_Agentic_Workload_Optimizations_in_TensorRT-LLM.md) - -* [07/01] Scaling Video Generation Across NVL72 Rack with TensorRT-LLM -✨ [➡️ link](https://github.com/NVIDIA/TensorRT-LLM/blob/main/docs/source/blogs/tech_blog/blog25_Scaling_Video_Generation_Across_NVL72_Rack_with_TensorRT-LLM.md) - * [05/15] Joint Optimization of Agent Applications and TensorRT-LLM ✨ [➡️ link](https://github.com/NVIDIA/TensorRT-LLM/blob/main/docs/source/blogs/tech_blog/blog23_Joint_Optimization_of_Agent_Applications_and_TensorRT-LLM.md) @@ -304,10 +298,9 @@ Deprecation is used to inform developers that some APIs and tools are no longer TensorRT-LLM collects anonymous telemetry data by default. This data is used in aggregate to understand usage patterns and prioritize engineering efforts. **This data cannot be traced back to any individual user.** No prompts, -outputs, model weights, model paths, tokenizer paths, user-identifying -information, raw free-form configuration strings, or persistent identifiers are -collected. Any deployment identifiers are ephemeral, randomly generated per -deployment, and not linked to users. The data we collect includes: +user-identifying information, or persistent identifiers are collected. Any +deployment identifiers are ephemeral, randomly generated per deployment, and +not linked to users. The data we collect includes: - Ingress point (e.g., LLM API, CLI, serve command) - Deployment duration (via periodic heartbeats) @@ -316,10 +309,8 @@ deployment, and not linked to users. The data we collect includes: - Parallelism configuration (TP/PP/CP/MoE-EP/MoE-TP sizes), quantization algorithm, dtype, KV cache dtype - System information (OS platform, Python version, CPU architecture, CPU count) - TRT-LLM version and backend -- Feature summary flags (LoRA, speculative decoding, prefix caching, CUDA graphs, chunked context, data parallelism) +- Feature flags (LoRA, speculative decoding, prefix caching, CUDA graphs, chunked context, data parallelism) - Disaggregated serving metadata (role and deployment ID) -- Selected LLM API configuration values: parallelism, dtype, KV cache, scheduler, CUDA graph, and compile settings -- Capture diagnostics for that payload: a schema checksum (for provenance), the count of captured fields, and whether any free-form value was skipped Telemetry is automatically disabled in CI and test environments. diff --git a/constraints.txt b/constraints.txt index 09a933566229..519a9a29f9fc 100644 --- a/constraints.txt +++ b/constraints.txt @@ -1,5 +1,11 @@ # These vulnerabilities were inherited from the base image (pytorch:25.12-py3) and should be removed when the base image # is updated. +# WAR against https://github.com/advisories/GHSA-8rrh-rw8j-w5fx +wheel>=0.46.2 +# WAR against https://github.com/advisories/GHSA-qjxf-f2mg-c6mc +tornado>=6.5.5 +# WAR against https://github.com/advisories/GHSA-3936-cmfr-pm3m +black>=26.3.1 # Upgrade base image nvidia-cutlass-dsl 4.3.5 to 4.4.2 nvidia-cutlass-dsl>=4.4.2 # The `nvidia-cutlass-dsl` package does not pin numpy at all, which can be problematic in certain CI diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index a323b32b82b5..7a46725c80c7 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -406,13 +406,8 @@ endif() option(ENABLE_BOLT_COMPATIBLE "Enable BOLT-compatible build flags" OFF) if(ENABLE_BOLT_COMPATIBLE AND NOT WIN32) message(STATUS "BOLT compatible flags enabled") - # Compiler flags for C/C++ -fno-reorder-blocks-and-partition is GCC-only - # (required for GCC 8+ per BOLT docs). Clang does not support this flag and - # does not enable the corresponding optimization by default. - if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU") - add_compile_options(-fno-reorder-blocks-and-partition) - endif() - add_compile_options(-fno-plt) + # Compiler flags for C/C++ + add_compile_options(-fno-reorder-blocks-and-partition -fno-plt) # Linker flags - applies to shared, module, and executable targets add_link_options(-Wl,--emit-relocs) # Disable stripping - required for BOLT (affects pybind11 POST_BUILD strip) diff --git a/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h b/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h index 0fcee1435005..f79d51ae59ae 100644 --- a/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h +++ b/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h @@ -26,6 +26,7 @@ #include "tensorrt_llm/executor/dataTransceiverState.h" #include "tensorrt_llm/runtime/utils/mpiUtils.h" #include "tensorrt_llm/runtime/utils/pgUtils.h" +#include #include #include #include @@ -226,11 +227,6 @@ class BaseCacheTransceiver [[nodiscard]] virtual bool checkGenTransferComplete() const = 0; virtual bool cancelRequest(std::shared_ptr llmRequest) = 0; - - [[nodiscard]] virtual bool hasPoisonedTransferBuffer() const - { - return false; - } }; class CacheTransceiver : public BaseCacheTransceiver @@ -242,6 +238,7 @@ class CacheTransceiver : public BaseCacheTransceiver executor::kv_cache::CacheState::AttentionType attentionType = executor::kv_cache::CacheState::AttentionType::kDEFAULT, std::optional cacheTransceiverConfig = std::nullopt, + rnn_state_manager::RnnStateManager* rnnStateManager = nullptr, std::vector const& rnnLayerNumPerPP = {}); CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheManager, std::vector numKvHeadsPerLayer, @@ -250,10 +247,11 @@ class CacheTransceiver : public BaseCacheTransceiver executor::kv_cache::CacheState::AttentionType attentionType = executor::kv_cache::CacheState::AttentionType::kDEFAULT, std::optional cacheTransceiverConfig = std::nullopt, + rnn_state_manager::RnnStateManager* rnnStateManager = nullptr, std::vector const& rnnLayerNumPerPP = {}) : CacheTransceiver(cacheManager, executor::kv_cache::CacheState::ModelConfig{numKvHeadsPerLayer, sizePerHead, tokensPerBlock}, worldConfig, - attentionLayerNumPerPP, dataType, attentionType, cacheTransceiverConfig, rnnLayerNumPerPP) + attentionLayerNumPerPP, dataType, attentionType, cacheTransceiverConfig, rnnStateManager, rnnLayerNumPerPP) { } @@ -276,13 +274,21 @@ class CacheTransceiver : public BaseCacheTransceiver virtual bool cancelRequest(std::shared_ptr llmRequest) override; - [[nodiscard]] bool hasPoisonedTransferBuffer() const override; - private: void initializeCommState(); void setContextState(LlmRequest* llmRequest); + // Default-off, transition-only lifecycle tracing for NVBUG#6448152. + int mWorldRank; + bool mNvbug6448152TraceEnabled; + std::uint64_t mNvbug6448152ContextCheckSequence{0}; + std::uint64_t mNvbug6448152ContextTpConsensusSequence{0}; + std::uint64_t mNvbug6448152ContextPpConsensusSequence{0}; + std::uint64_t mNvbug6448152GenerationCheckSequence{0}; + std::uint64_t mNvbug6448152ContextEnqueuedTransitionCount{0}; + std::unordered_set mNvbug6448152SenderWaitTimeoutIds; + std::unique_ptr mCacheSender; std::unique_ptr mCacheReceiver; // shared_ptr (not raw LlmRequest*) so the futures hold a strong reference for @@ -290,12 +296,9 @@ class CacheTransceiver : public BaseCacheTransceiver // request while a C++ status check still dereferences it. std::vector, std::future>> mSenderFutures; std::vector, std::future>> mRequesterFutures; - // Dedup timeout logs separately from accepted cancellation requests so a - // backend that initially declines cancellation is retried on later polls. + // Dedup sets so observe-only timeout WARN logs fire at most once per stuck request. std::unordered_set mTimedOutSenderIds; std::unordered_set mTimedOutRequesterIds; - std::unordered_set mCancelRequestedSenderIds; - std::unordered_set mCancelRequestedRequesterIds; std::unordered_set mCompletedSenderRequestIds; std::unordered_set mFailedSenderRequestIds; std::unordered_map> mSenderRequestsAwaitingConsensus; @@ -314,6 +317,7 @@ class CacheTransceiver : public BaseCacheTransceiver std::vector> mCacheTransBufferManagers; std::vector mCacheTransBufferManagerPtrs; + rnn_state_manager::RnnStateManager* mRnnStateManager{nullptr}; // TODO(shreyasm): update this to use same container as kv by using base trans buffers instead std::unique_ptr mRnnCacheTransBufferManager{nullptr}; diff --git a/cpp/include/tensorrt_llm/batch_manager/capacityScheduler.h b/cpp/include/tensorrt_llm/batch_manager/capacityScheduler.h index d4faa6c5120d..2ea4f47ce4bc 100644 --- a/cpp/include/tensorrt_llm/batch_manager/capacityScheduler.h +++ b/cpp/include/tensorrt_llm/batch_manager/capacityScheduler.h @@ -1,5 +1,5 @@ /* - * Copyright (c) 2023-2026, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2023-2024, NVIDIA CORPORATION. All rights reserved. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -87,44 +87,31 @@ class MaxRequestsScheduler : public BaseCapacityScheduler /// @brief Schedule requests using the MAX_UTILIZATION policy /// @details Try reserving resources to advance requests by one step, -/// may pause previously started requests. When a -/// ``crossKvCacheManager`` is supplied, requests in the -/// ``ENCODER_INIT`` state may be admitted for encoder compute -/// without consuming self- or cross-KV blocks; the later -/// ``CONTEXT_INIT`` decoder admission owns cross-pool budgeting. +/// may pause previously started requests. class MaxUtilizationScheduler : public BaseCapacityScheduler { public: MaxUtilizationScheduler(SizeType32 maxNumRequests, bool twoStepsLookAhead, LlmRequestState noScheduleUntilState = LlmRequestState::kCONTEXT_INIT, - LlmRequestState noScheduleAfterState = LlmRequestState::kGENERATION_COMPLETE, - bool enablePrefixAwareScheduling = true); + LlmRequestState noScheduleAfterState = LlmRequestState::kGENERATION_COMPLETE); [[nodiscard]] std::tuple operator()( - kv_cache_manager::BaseKVCacheManager& kvCacheManager, - OptionalRef crossKvCacheManager, - OptionalRef peftCacheManager, RequestList const& activeRequests) const; + kv_cache_manager::BaseKVCacheManager& kvCacheManager, OptionalRef peftCacheManager, + RequestList const& activeRequests) const; private: SizeType32 mMaxNumRequests; /// @brief Boolean that indicates if two step lookahead is enabled bool mTwoStepsLookAhead; - /// @brief Whether to use KV prefix-reuse estimates in scheduling decisions. - bool mEnablePrefixAwareScheduling; }; /// @brief Schedule requests using the GUARANTEED_NO_EVICT policy -/// @details When a ``crossKvCacheManager`` is supplied, requests in the -/// ``ENCODER_INIT`` state may be admitted for encoder compute -/// without consuming self- or cross-KV blocks. The later -/// ``CONTEXT_INIT`` decoder admission owns cross-pool budgeting. class GuaranteedNoEvictScheduler : public BaseCapacityScheduler { public: GuaranteedNoEvictScheduler(SizeType32 maxNumRequests, LlmRequestState noScheduleUntilState = LlmRequestState::kCONTEXT_INIT, - LlmRequestState noScheduleAfterState = LlmRequestState::kGENERATION_COMPLETE, - bool enablePrefixAwareScheduling = true); + LlmRequestState noScheduleAfterState = LlmRequestState::kGENERATION_COMPLETE); [[nodiscard]] std::tuple operator()( kv_cache_manager::BaseKVCacheManager const& kvCacheManager, @@ -140,8 +127,6 @@ class GuaranteedNoEvictScheduler : public BaseCapacityScheduler private: SizeType32 mMaxNumRequests; - /// @brief Whether to use KV prefix-reuse estimates in scheduling decisions. - bool mEnablePrefixAwareScheduling; }; /// @brief Schedule requests using the STATIC_BATCH policy @@ -150,8 +135,7 @@ class StaticBatchScheduler : public GuaranteedNoEvictScheduler public: StaticBatchScheduler(SizeType32 maxNumRequests, LlmRequestState noScheduleUntilState = LlmRequestState::kCONTEXT_INIT, - LlmRequestState noScheduleAfterState = LlmRequestState::kGENERATION_COMPLETE, - bool enablePrefixAwareScheduling = true); + LlmRequestState noScheduleAfterState = LlmRequestState::kGENERATION_COMPLETE); [[nodiscard]] std::tuple operator()( kv_cache_manager::BaseKVCacheManager const& kvCacheManager, @@ -167,19 +151,14 @@ class CapacityScheduler : public Algorithm explicit CapacityScheduler(SizeType32 maxNumRequests, executor::CapacitySchedulerPolicy capacitySchedulerPolicy, bool hasKvCacheManager, bool twoStepsLookAhead = false, LlmRequestState noScheduleUntilState = LlmRequestState::kCONTEXT_INIT, - LlmRequestState noScheduleAfterState = LlmRequestState::kGENERATION_COMPLETE, - bool enablePrefixAwareScheduling = true); + LlmRequestState noScheduleAfterState = LlmRequestState::kGENERATION_COMPLETE); /** * @brief Schedules requests following the selected policy. * * @param kvCacheManager Required in MaxUtilizationScheduler (as a ref) and in GuaranteedNoEvictScheduler and * StaticBatchScheduler (as a const ref). - * @param crossKvCacheManager Optional cross-attention KV cache manager. Used by - * MaxUtilizationScheduler (mutates: ``startScheduling`` / ``schedulingRemoveSequence``) - * and GuaranteedNoEvictScheduler / StaticBatchScheduler (read-only). Required for - * encoder-decoder admission. Encoder-init requests only require this pool - * to be configured; decoder context admission budgets blocks from it. + * @param crossKvCacheManager Optional used in GuaranteedNoEvictScheduler and StaticBatchScheduler. * @param peftCacheManager Optional used in MaxUtilizationScheduler, GuaranteedNoEvictScheduler and * StaticBatchScheduler. * @param activeRequests @@ -189,7 +168,7 @@ class CapacityScheduler : public Algorithm [[nodiscard]] std::tuple operator()(RequestList const& activeRequests, OptionalRef kvCacheManager = std::nullopt, OptionalRef peftCacheManager = std::nullopt, - OptionalRef crossKvCacheManager = std::nullopt) const; + OptionalRef crossKvCacheManager = std::nullopt) const; /// @brief Sets the reorder policy to use AgentTreePolicy with the given configuration. /// @param agentPercentage The ratio of agent requests to schedule (0.0-1.0, -1.0 for random). diff --git a/cpp/include/tensorrt_llm/batch_manager/disaggTransferAdmissionController.h b/cpp/include/tensorrt_llm/batch_manager/disaggTransferAdmissionController.h deleted file mode 100644 index 271a969052fb..000000000000 --- a/cpp/include/tensorrt_llm/batch_manager/disaggTransferAdmissionController.h +++ /dev/null @@ -1,165 +0,0 @@ -/* - * 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. - */ - -#pragma once - -#include "tensorrt_llm/batch_manager/common.h" -#include "tensorrt_llm/batch_manager/llmRequest.h" - -#include -#include - -namespace tensorrt_llm::batch_manager -{ - -class DisaggTransferAdmissionController -{ -public: - enum class Policy - { - kFcfsEstimatedBlockBudget - }; - - struct Result - { - RequestVector admittedRequests; - std::size_t activeTransferBlocks{}; - std::size_t admittedTransferBlocks{}; - std::size_t deferredRequestCount{}; - bool limitedByBudget{}; - - [[nodiscard]] bool isBlockedByActiveTransfers() const - { - return limitedByBudget && admittedRequests.empty() && activeTransferBlocks > 0; - } - }; - - explicit DisaggTransferAdmissionController(std::optional maxTokensInBuffer, SizeType32 tokensPerBlock, - Policy policy = Policy::kFcfsEstimatedBlockBudget) - : mMaxTransferBlocks(toBlockBudget(maxTokensInBuffer, tokensPerBlock)) - , mTokensPerBlock(tokensPerBlock) - , mPolicy(policy) - { - } - - [[nodiscard]] bool enabled() const - { - return mMaxTransferBlocks.has_value(); - } - - [[nodiscard]] std::optional getMaxTransferBlocks() const - { - return mMaxTransferBlocks; - } - - [[nodiscard]] Result select(RequestList const& activeRequests, RequestVector const& candidates) const - { - if (!enabled()) - { - return Result{ - candidates, estimateActiveTransferBlocks(activeRequests), estimateRequestsBlocks(candidates), 0, false}; - } - - switch (mPolicy) - { - case Policy::kFcfsEstimatedBlockBudget: return selectFcfsEstimatedBlockBudget(activeRequests, candidates); - } - - return Result{}; - } - -private: - [[nodiscard]] static std::optional toBlockBudget( - std::optional maxTokensInBuffer, SizeType32 tokensPerBlock) - { - if (!maxTokensInBuffer.has_value() || maxTokensInBuffer.value() == 0 || tokensPerBlock <= 0) - { - return std::nullopt; - } - auto const blockSize = static_cast(tokensPerBlock); - return (maxTokensInBuffer.value() + blockSize - 1) / blockSize; - } - - [[nodiscard]] std::size_t estimateRequestBlocks(LlmRequest const& request) const - { - if (mTokensPerBlock <= 0) - { - return 0; - } - auto const promptLen = static_cast(request.getPromptLen()); - auto const blockSize = static_cast(mTokensPerBlock); - return (promptLen + blockSize - 1) / blockSize; - } - - [[nodiscard]] std::size_t estimateRequestsBlocks(RequestVector const& requests) const - { - std::size_t blocks{}; - for (auto const& request : requests) - { - blocks += estimateRequestBlocks(*request); - } - return blocks; - } - - [[nodiscard]] std::size_t estimateActiveTransferBlocks(RequestList const& activeRequests) const - { - std::size_t blocks{}; - for (auto const& request : activeRequests) - { - if (request->isDisaggGenerationTransmissionInProgress()) - { - blocks += estimateRequestBlocks(*request); - } - } - return blocks; - } - - [[nodiscard]] Result selectFcfsEstimatedBlockBudget( - RequestList const& activeRequests, RequestVector const& candidates) const - { - Result result; - result.activeTransferBlocks = estimateActiveTransferBlocks(activeRequests); - - auto const maxTransferBlocks = mMaxTransferBlocks.value(); - auto usedBlocks = result.activeTransferBlocks; - for (auto const& request : candidates) - { - auto const requestBlocks = estimateRequestBlocks(*request); - bool const fitsBudget = usedBlocks + requestBlocks <= maxTransferBlocks; - bool const admitOversizedHead = result.admittedRequests.empty() && result.activeTransferBlocks == 0 - && requestBlocks > maxTransferBlocks; - if (!fitsBudget && !admitOversizedHead) - { - result.limitedByBudget = true; - break; - } - - result.admittedRequests.push_back(request); - usedBlocks += requestBlocks; - result.admittedTransferBlocks += requestBlocks; - } - - result.deferredRequestCount = candidates.size() - result.admittedRequests.size(); - return result; - } - - std::optional mMaxTransferBlocks; - SizeType32 mTokensPerBlock; - Policy mPolicy; -}; - -} // namespace tensorrt_llm::batch_manager diff --git a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h index e821cafd5c3f..c665f7a8df95 100644 --- a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h +++ b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h @@ -51,11 +51,6 @@ namespace kvc = tensorrt_llm::executor::kv_cache; -namespace tensorrt_llm::batch_manager::kv_cache_manager -{ -class FabricMemory; -} // namespace tensorrt_llm::batch_manager::kv_cache_manager - namespace tensorrt_llm::batch_manager::eviction_policy { class BaseEvictionPolicy; @@ -103,13 +98,6 @@ template std::list> chopVectorIntoBlocks( std::vector const& vec, SizeType32 usableSize, SizeType32 elementsPerBlock, bool allowPartial) { - // No usable elements yields no blocks. Guard non-positive usableSize explicitly: callers may pass - // usableSize = inputLength - 1, which is -1 for a Helix CP "empty" rank with 0 input tokens. With a - // negative usableSize, downstream usage of "vec.begin() + usableSize" is undefined behavior. - if (usableSize <= 0) - { - return {}; - } TLLM_CHECK_WITH_INFO( usableSize <= static_cast(vec.size()), "usableSize=%d > %ld=vec.size()", usableSize, vec.size()); std::list> blockedVectors; @@ -327,8 +315,8 @@ struct KvCacheStats std::size_t allocatedBytes{}; }; -/// @brief Per-iteration KV cache statistics. All delta counters and peak gauges represent values since the last call -/// to getIterationStats(). Snapshot gauges are instantaneous. +/// @brief Per-iteration KV cache statistics. All delta counters represent changes since the last call to +/// getIterationStats(). Gauges are instantaneous snapshots. struct KvCacheIterationStats { // --- Instantaneous gauges --- @@ -336,23 +324,10 @@ struct KvCacheIterationStats SizeType32 primaryMaxNumBlocks{0}; SizeType32 primaryFreeNumBlocks{0}; SizeType32 primaryUsedNumBlocks{0}; - // Cached-but-unpinned blocks in the primary pool. Distinct from primaryUsedNumBlocks, - // which also counts blocks pinned during onboard memcpy windows. - SizeType32 primaryEvictableNumBlocks{0}; - SizeType32 primaryPeakFreeNumBlocks{0}; - SizeType32 primaryPeakUsedNumBlocks{0}; - SizeType32 primaryPeakEvictableNumBlocks{0}; // Secondary (host) pool SizeType32 secondaryMaxNumBlocks{0}; SizeType32 secondaryFreeNumBlocks{0}; SizeType32 secondaryUsedNumBlocks{0}; - // Cached-but-unpinned blocks in the secondary pool. Useful to gauge "how full is the - // host cache"; secondaryUsedNumBlocks only counts pinned blocks during the sub-ms - // onboard memcpy window so it cannot answer that question on its own. - SizeType32 secondaryEvictableNumBlocks{0}; - SizeType32 secondaryPeakFreeNumBlocks{0}; - SizeType32 secondaryPeakUsedNumBlocks{0}; - SizeType32 secondaryPeakEvictableNumBlocks{0}; // --- Per-iteration deltas (reset on each read) --- // Context phase: block allocation and reuse @@ -374,11 +349,6 @@ struct KvCacheIterationStats // Intra-device (GPU → GPU) block copies (e.g. partial reuse when source block has refs) SizeType32 iterIntraDeviceCopyBlocks{0}; std::size_t iterIntraDeviceCopyBytes{0}; - - // Pages released by LRU from the last cache tier without ever being onboarded back - // to GPU during their stay at that tier (i.e. fully dropped from the hierarchy). - SizeType32 iterHostDroppedBlocks{0}; - std::size_t iterHostDroppedBytes{0}; }; // Basic building block of a paged KV cache - a single @@ -517,10 +487,6 @@ class KVCacheBlock : public std::enable_shared_from_this [[nodiscard]] bool isLeaf() const; - //! \brief Test if block is detached from radix search tree. - //! \return True if block is detached from search tree. - [[nodiscard]] bool isDetached() const; - void setPriority(executor::RetentionPriority priority); [[nodiscard]] executor::RetentionPriority getPriority() const; @@ -1374,8 +1340,6 @@ class WindowBlockManager // Buffer manager runtime::BufferManager mBufferManager; - // Fabric memory backing for primary pools (MNNVL-capable allocation) - std::vector> mFabricMemoryPools; // Used to keep track of number of free blocks during scheduling SizeType32 mSchedulingNumFreeBlocks; @@ -2578,24 +2542,6 @@ class KVCacheManager : public BaseKVCacheManager [[nodiscard]] std::vector commitAndGetBlockHashesForRequest( LlmRequest const& llmRequest, SizeType32 windowSize) override; - //! @brief Translate logical block IDs into primary-pool block indices. - //! @details A block ID is stable for the lifetime of a block, but its position inside the - //! memory pool can change after offload/onboard cycles. This function performs - //! that translation. The returned index is the value of - //! `KVCacheBlock::getMemoryPoolBlockIndex()` for each input, with the pool flag - //! stripped (see `kernels::KVCacheIndex::get()`), so it is only meaningful for - //! blocks resident in the primary pool. Every referenced block must therefore be - //! primary; this is asserted. Callers (e.g. the disaggregation cache transceiver - //! on the Python side) cannot check residency themselves, and the invariant holds - //! because allocation onboards offloaded blocks and offload only ever selects free - //! blocks — a violation indicates a block-lifetime bug. - //! @param blockIds IDs to translate. - //! @param windowSize Attention window the IDs belong to (selects the WindowBlockManager). - //! @throws Aborts via TLLM_CHECK_WITH_INFO if any referenced block is not found or is not - //! currently in the primary pool. - [[nodiscard]] std::vector getMemoryPoolBlockIndicesByBlockIds( - std::vector const& blockIds, SizeType32 windowSize) const; - std::optional getLastBlockId(LlmRequest::RequestIdType requestId) const override; /// @brief Calculates the number of kv-cache blocks that a sequence will require, for a single beam. diff --git a/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h b/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h index 9e9a07b19d8e..57f8f928a8be 100644 --- a/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h +++ b/cpp/include/tensorrt_llm/batch_manager/kvCacheTransferManager.h @@ -24,11 +24,6 @@ namespace kvc = tensorrt_llm::executor::kv_cache; #pragma once -namespace tensorrt_llm::testing -{ -class KVCacheTransferManagerTestAccess; -} // namespace tensorrt_llm::testing - namespace tensorrt_llm::batch_manager::kv_cache_manager { @@ -81,15 +76,10 @@ class KVCacheTransferManager [[nodiscard]] KvCacheTransferStats getAndResetTransferStats(); private: - friend class ::tensorrt_llm::testing::KVCacheTransferManagerTestAccess; - //! \brief Get pointer to pool specified by cache block. static tr::ITensor::SharedPtr computeBlockPointer( BlockPtr const& block, std::vector const& pools, size_t poolIdx); - //! \brief Get pool-qualified index for pending transfer tracking. - [[nodiscard]] static kernels::KVCacheIndex::UnderlyingType getPendingTransferIndex(BlockPtr const& block); - /*! * \brief The key method that copies the src block to the dst block. * @@ -117,8 +107,8 @@ class KVCacheTransferManager runtime::BufferManager mOnboardManager; runtime::BufferManager mOffloadManager; - // Track reads and writes for blocks. Note that it is the pool-qualified memory pool index - // that identifies the raw memory blocks involved in I/O, not the block Id. + // Track reads and writes for blocks. Note that it is the memory pool index that + // identifies the raw memory blocks involved in I/O, not the block Id. std::unordered_map mPendingReads; std::unordered_map mPendingWrites; // Reference to parent loopback agent diff --git a/cpp/include/tensorrt_llm/batch_manager/llmRequest.h b/cpp/include/tensorrt_llm/batch_manager/llmRequest.h index bc1ca3e6d012..886147a09c73 100644 --- a/cpp/include/tensorrt_llm/batch_manager/llmRequest.h +++ b/cpp/include/tensorrt_llm/batch_manager/llmRequest.h @@ -665,9 +665,9 @@ class GenericLlmRequest return mEncoderUniqueTokens; } - /// @brief Get length of encoder input when present, without throwing for decoder-only requests. - /// @return Encoder input length, or nullopt when this request has no encoder side. - [[nodiscard]] std::optional tryGetEncoderInputLen() const + /// @brief Get length of encoder input (could be tokens or features length) + /// @return An integer. + [[nodiscard]] SizeType32 getEncoderInputLen() const { if (mEncoderInputFeatures.has_value()) { @@ -678,45 +678,19 @@ class GenericLlmRequest return getEncoderTokens().value()->size(); } - return std::nullopt; - } - - /// @brief Get length of encoder input (could be tokens or features length) - /// @return An integer. - [[nodiscard]] SizeType32 getEncoderInputLen() const - { - auto const encoderInputLen = tryGetEncoderInputLen(); - if (encoderInputLen.has_value()) - { - return encoderInputLen.value(); - } - TLLM_THROW("GenericLlmRequest::getEncoderInputLen - Do not have encoder length!"); } - /// @brief Get length of encoder output when present, without throwing for decoder-only requests. - /// @return Encoder output length, or nullopt when this request has no encoder side. - [[nodiscard]] std::optional tryGetEncoderOutputLen() const + /// @brief Get length of encoder output. Fall back to encoder input length if not present + /// @return An integer. + [[nodiscard]] SizeType32 getEncoderOutputLen() const { if (mEncoderOutputLength.has_value()) { return mEncoderOutputLength.value(); } - return tryGetEncoderInputLen(); - } - - /// @brief Get length of encoder output, or throw if the request has no encoder side. - /// @return Explicit encoder output length, or encoder input length when the output length is not present. - [[nodiscard]] SizeType32 getEncoderOutputLen() const - { - auto const encoderOutputLen = tryGetEncoderOutputLen(); - if (encoderOutputLen.has_value()) - { - return encoderOutputLen.value(); - } - - TLLM_THROW("GenericLlmRequest::getEncoderInputLen - Do not have encoder length!"); + return getEncoderInputLen(); } [[nodiscard]] std::optional>> getPositionIds() const @@ -1908,15 +1882,6 @@ class GenericLlmRequest return mPerfMetrics.kvCacheMetrics.numNewAllocatedBlocks; } - void updateKvCachePerfMetrics( - SizeType32 allocTotalBlocks, SizeType32 allocNewBlocks, SizeType32 reusedBlocks, SizeType32 missedBlocks) - { - updateAllocTotalBlocksPerRequest(allocTotalBlocks); - updateAllocNewBlocksPerRequest(allocNewBlocks); - updateReusedBlocksPerRequest(reusedBlocks); - updateMissedBlocksPerRequest(missedBlocks); - } - void updateReusedBlocksPerRequest(SizeType32 reusedBlocksPerRequest) { mPerfMetrics.kvCacheMetrics.numReusedBlocks += reusedBlocksPerRequest; @@ -2247,10 +2212,7 @@ class GenericLlmRequest // Scatter the input tokens to other beam mTokens = BeamTokens(mSamplingConfig.beamWidth, inputTokens); - // A request may legitimately have no input tokens on this rank (e.g. an "empty" Helix CP rank that owns zero KV - // blocks for the sequence). Guard against calling .back() on an empty vector (undefined behavior). - mLastTokens = inputTokens.empty() ? VecTokens(mSamplingConfig.beamWidth) - : VecTokens(mSamplingConfig.beamWidth, inputTokens.back()); + mLastTokens = VecTokens(mSamplingConfig.beamWidth, inputTokens.back()); // Init mUniqueTokens VecUniqueTokens uniqueTokens{inputTokens.size()}; diff --git a/cpp/include/tensorrt_llm/batch_manager/rnnCacheFormatter.h b/cpp/include/tensorrt_llm/batch_manager/rnnCacheFormatter.h index b57d08f2c12b..1bc3fadf62f4 100644 --- a/cpp/include/tensorrt_llm/batch_manager/rnnCacheFormatter.h +++ b/cpp/include/tensorrt_llm/batch_manager/rnnCacheFormatter.h @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -34,11 +34,14 @@ class TransferSession; namespace rnn_state_manager { +class RnnStateManager; class RnnCacheTransBufferManager; } // namespace rnn_state_manager /// @brief RNN Cache Formatter for formatting/unformatting RNN states during transfer. -/// Uses unified pool mode via BaseKVCacheManager (for CppMambaHybridCacheManager, block-indexed pool). +/// Supports two operating modes: +/// - Slot mode: uses RnnStateManager (for CppMambaCacheManager, separate tensor storage) +/// - Unified pool mode: uses BaseKVCacheManager (for CppMambaHybridCacheManager, block-indexed pool) class RnnCacheFormatter : public kv_cache_manager::BaseCacheFormatter { public: @@ -46,6 +49,12 @@ class RnnCacheFormatter : public kv_cache_manager::BaseCacheFormatter using CacheState = executor::kv_cache::CacheState; using RequestIdType = tensorrt_llm::batch_manager::RequestIdType; + /// @brief Constructor for slot-based mode (CppMambaCacheManager with RnnStateManager). + /// @param rnnStateManager Pointer to the RNN state manager. + /// @param rnnCacheTransBufferManager Pointer to the RNN cache transfer buffer manager. + RnnCacheFormatter(rnn_state_manager::RnnStateManager* rnnStateManager, + rnn_state_manager::RnnCacheTransBufferManager* rnnCacheTransBufferManager); + /// @brief Constructor for unified pool mode (CppMambaHybridCacheManager). /// @param kvCacheManager Pointer to the KV cache manager with unified pool. /// @param rnnCacheTransBufferManager Pointer to the RNN cache transfer buffer manager. @@ -72,13 +81,39 @@ class RnnCacheFormatter : public kv_cache_manager::BaseCacheFormatter CacheState const& selfConfig, SizeType32 selfIdx, CacheState const& destConfig, std::vector const& counterPartRanks) const override; - /// @brief Returns the KV cache manager. + /// @brief Returns the KV cache manager (non-null in unified pool mode). [[nodiscard]] kv_cache_manager::BaseKVCacheManager* getCacheManager() const noexcept override { return mKvCacheManager; } + /// @brief Get the RNN state manager (non-null in slot mode). + /// @return Pointer to the RNN state manager. + [[nodiscard]] rnn_state_manager::RnnStateManager* getRnnStateManager() const noexcept + { + return mRnnStateManager; + } + + /// @brief Check if operating in unified pool mode. + [[nodiscard]] bool isUnifiedPoolMode() const noexcept + { + return mKvCacheManager != nullptr; + } + private: + /// @brief Format logic for slot-based path (RnnStateManager). + void formatSlotMode(TransferSession& session); + + /// @brief Unformat logic for slot-based path (RnnStateManager). + void unformatSlotMode(TransferSession& session); + + /// @brief Format logic for unified pool path (BaseKVCacheManager). + void formatUnifiedPoolMode(TransferSession& session); + + /// @brief Unformat logic for unified pool path (BaseKVCacheManager). + void unformatUnifiedPoolMode(TransferSession& session); + + rnn_state_manager::RnnStateManager* mRnnStateManager{nullptr}; rnn_state_manager::RnnCacheTransBufferManager* mRnnCacheTransBufferManager; kv_cache_manager::BaseKVCacheManager* mKvCacheManager{nullptr}; }; diff --git a/cpp/include/tensorrt_llm/batch_manager/runtimeBuffers.h b/cpp/include/tensorrt_llm/batch_manager/runtimeBuffers.h index 97a4ae67acdd..13bde6d07a5e 100644 --- a/cpp/include/tensorrt_llm/batch_manager/runtimeBuffers.h +++ b/cpp/include/tensorrt_llm/batch_manager/runtimeBuffers.h @@ -1,5 +1,5 @@ /* - * Copyright (c) 2023-2026, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2023-2024, NVIDIA CORPORATION. All rights reserved. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -208,9 +208,7 @@ class RuntimeBuffers //! Temporarily store the transposed results of multiple fragment logits, [maxBeamWidth, kCACHE_LENGTH] TensorPtr transposedLogits; - //! Temporarily store logits buffer address during the transposing, [maxBatchSize, kCACHE_LENGTH] - //! One row per batch slot (same layout as fragmentPointerHost) so concurrent flushes for - //! different requests in the same batch never clobber each other's pointer arrays. + //! Temporarily store logits buffer address during the transposing, [kCACHE_LENGTH] TensorPtr fragmentPointerDevice; //! Temporarily store logits buffer address during the transposing, [maxBatchSize, kCACHE_LENGTH] @@ -224,14 +222,11 @@ class RuntimeBuffers workIdx = (workIdx + 1) % (fragmentPointerHost->getShape().d[0]); } - //! Returns matching host and device pointer rows for the current workIdx, then advances - //! workIdx. Always call this instead of the individual getters to avoid ordering bugs. - [[nodiscard]] std::pair getFragmentPointerSlot() + [[nodiscard]] TensorPtr getFragmentPointerHost() { - TensorPtr host = runtime::ITensor::slice(fragmentPointerHost, workIdx, 1); - TensorPtr device = runtime::ITensor::slice(fragmentPointerDevice, workIdx, 1); + TensorPtr slice = runtime::ITensor::slice(fragmentPointerHost, workIdx, 1); cycleWorkIdx(); - return {std::move(host), std::move(device)}; + return slice; }; }; diff --git a/cpp/include/tensorrt_llm/common/optionalRef.h b/cpp/include/tensorrt_llm/common/optionalRef.h index 46723f1c697c..f55b377981d2 100644 --- a/cpp/include/tensorrt_llm/common/optionalRef.h +++ b/cpp/include/tensorrt_llm/common/optionalRef.h @@ -78,13 +78,6 @@ class OptionalRef { } - // Implicit conversion from OptionalRef to OptionalRef - template >> - OptionalRef(OptionalRef> const& other) - : opt(other ? std::optional>(std::ref(*other)) : std::nullopt) - { - } - T* operator->() const { return opt ? &(opt->get()) : nullptr; diff --git a/cpp/include/tensorrt_llm/common/quantization.h b/cpp/include/tensorrt_llm/common/quantization.h index e53a982316c5..df13a674d688 100644 --- a/cpp/include/tensorrt_llm/common/quantization.h +++ b/cpp/include/tensorrt_llm/common/quantization.h @@ -134,11 +134,6 @@ class QuantMode return QuantMode(BaseType(1u) << 16); } - static constexpr QuantMode mxfp8() noexcept - { - return QuantMode(BaseType(1u) << 17); - } - constexpr BaseType value() const noexcept { return mValue; @@ -229,11 +224,6 @@ class QuantMode return isSet(w4a16Mxfp4()); } - constexpr bool hasMxfp8() const noexcept - { - return isSet(mxfp8()); - } - constexpr bool hasKvCacheQuant() const noexcept { return hasInt8KvCache() || hasFp8KvCache() || hasFp4KvCache(); diff --git a/cpp/include/tensorrt_llm/executor/executor.h b/cpp/include/tensorrt_llm/executor/executor.h index 825b8ad75959..f716bef6e3cc 100644 --- a/cpp/include/tensorrt_llm/executor/executor.h +++ b/cpp/include/tensorrt_llm/executor/executor.h @@ -758,7 +758,6 @@ class Request ~Request(); [[nodiscard]] VecTokens getInputTokenIds() const; - [[nodiscard]] SizeType32 getNumInputTokens() const; [[nodiscard]] SizeType32 getMaxTokens() const; [[nodiscard]] bool getStreaming() const; [[nodiscard]] SamplingConfig getSamplingConfig() const; @@ -990,8 +989,6 @@ class DynamicBatchConfig [[nodiscard]] std::vector> getBatchSizeTable() const; - bool operator==(DynamicBatchConfig const& other) const; - /// @brief The default value of batch size table static std::vector> const kDefaultBatchSizeTable; @@ -1022,7 +1019,7 @@ class SchedulerConfig explicit SchedulerConfig( CapacitySchedulerPolicy capacitySchedulerPolicy = CapacitySchedulerPolicy::kGUARANTEED_NO_EVICT, std::optional contextChunkingPolicy = std::nullopt, - std::optional dynamicBatchConfig = std::nullopt, bool enablePrefixAwareScheduling = true); + std::optional dynamicBatchConfig = std::nullopt); bool operator==(SchedulerConfig const& other) const; @@ -1032,8 +1029,6 @@ class SchedulerConfig [[nodiscard]] std::optional getDynamicBatchConfig() const; - [[nodiscard]] bool getEnablePrefixAwareScheduling() const; - private: friend class Serialization; @@ -1045,9 +1040,6 @@ class SchedulerConfig /// @brief The config for tuning batch size dynamically. See DynamicBatchSizeConfig. std::optional mDynamicBatchConfig; - - /// @brief Whether schedulers use KV prefix-reuse estimates for admission and token-budget decisions. - bool mEnablePrefixAwareScheduling; }; /// @brief Configuration class for the KV cache @@ -1509,25 +1501,20 @@ class CacheTransceiverConfig NIXL = 3, MOONCAKE = 4 }; - static constexpr int kDefaultKvTransferPollIntervalMs = 5000; - explicit CacheTransceiverConfig(std::optional backendType = std::nullopt, std::optional maxNumTokens = std::nullopt, std::optional kvTransferTimeoutMs = std::nullopt, - std::optional kvTransferSenderFutureTimeoutMs = std::nullopt, - std::optional kvTransferPollIntervalMs = kDefaultKvTransferPollIntervalMs); + std::optional kvTransferSenderFutureTimeoutMs = std::nullopt); bool operator==(CacheTransceiverConfig const& other) const; void setBackendType(std::optional backendType); void setMaxTokensInBuffer(std::optional maxTokensInBuffer); void setKvTransferTimeoutMs(std::optional kvTransferTimeoutMs); void setKvTransferSenderFutureTimeoutMs(std::optional kvTransferSenderFutureTimeoutMs); - void setKvTransferPollIntervalMs(std::optional kvTransferPollIntervalMs); [[nodiscard]] std::optional getMaxTokensInBuffer() const; [[nodiscard]] std::optional getBackendType() const; [[nodiscard]] std::optional getKvTransferTimeoutMs() const; [[nodiscard]] std::optional getKvTransferSenderFutureTimeoutMs() const; - [[nodiscard]] std::optional getKvTransferPollIntervalMs() const; private: std::optional mBackendType; @@ -1539,9 +1526,6 @@ class CacheTransceiverConfig // @brief Timeout in milliseconds to wait for the sender future to be ready when scheduled batch size is 0. This // allows the request to be eventually cancelled by the user or because of kv_transfer_timeout_ms std::optional mKvTransferSenderFutureTimeoutMs; - // @brief Bounded wait interval in milliseconds for polling KV transfer progress when active transfers block - // disaggregated admission. - std::optional mKvTransferPollIntervalMs; }; /// @brief Configuration class for the model executor @@ -1556,9 +1540,6 @@ class ExecutorConfig // Per request stats may have additional overhead due to going through all requests. Turned off by default. static constexpr SizeType32 kDefaultRequestStatsMaxIterations = 0; - // A value of -1 keeps all iteration/request stats until they are fetched. - static constexpr SizeType32 kUnlimitedStatsMaxIterations = -1; - explicit ExecutorConfig(SizeType32 maxBeamWidth = 1, SchedulerConfig schedulerConfig = SchedulerConfig(), KvCacheConfig kvCacheConfig = KvCacheConfig(), bool enableChunkedContext = true, bool normalizeLogProbs = false, SizeType32 iterStatsMaxIterations = kDefaultIterStatsMaxIterations, @@ -1664,11 +1645,9 @@ class ExecutorConfig bool mNormalizeLogProbs; /// @brief Controls the maximum number of iterations for which to keep statistics. - /// Set to -1 to keep all iteration statistics. Set to 0 to disable iteration statistics. SizeType32 mIterStatsMaxIterations; /// @brief Controls the maximum number of iterations for which to keep per-request statistics. - /// Set to -1 to keep all per-request statistics. Set to 0 to disable per-request statistics. SizeType32 mRequestStatsMaxIterations; /// @brief The type of batching strategy to use. See BatchingType. @@ -1949,12 +1928,12 @@ class Executor void shutdown(); /// @brief Returns the per-iterations statistics computed since last call to getLatestIterationStats. - /// Contains at most iterStatsMaxIterations iterations, or all iterations when set to -1. + /// Contains at most iterStatsMaxIterations iterations. /// @return Iteration stats std::deque getLatestIterationStats(); /// @brief Returns the request stats of each iteration computed since last call to getLatestRequestStats. - /// Contains at most requestStatsMaxIterations iterations, or all iterations when set to -1. + /// Contains at most requestStatsMaxIterations iterations. /// @return Request stats grouped by iterations std::deque getLatestRequestStats(); diff --git a/cpp/include/tensorrt_llm/executor/transferAgent.h b/cpp/include/tensorrt_llm/executor/transferAgent.h index e1685c7c4ba5..532f0ae70c44 100644 --- a/cpp/include/tensorrt_llm/executor/transferAgent.h +++ b/cpp/include/tensorrt_llm/executor/transferAgent.h @@ -1,5 +1,5 @@ /* - * Copyright (c) 2025-2026, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -356,13 +356,6 @@ class TransferStatus virtual ~TransferStatus() = default; [[nodiscard]] virtual bool isCompleted() const = 0; virtual TransferState wait(int64_t timeout_ms = -1) const = 0; - - /// Release the backend transfer request handle. A true return means the backend accepted the handle release; it - /// does not prove remote memory quiescence. - [[nodiscard]] virtual bool release() - { - return false; - } }; struct BaseAgentConfig diff --git a/cpp/include/tensorrt_llm/runtime/ipcNvlsMemory.h b/cpp/include/tensorrt_llm/runtime/ipcNvlsMemory.h index 3f0f7216b8e6..8cf270694988 100644 --- a/cpp/include/tensorrt_llm/runtime/ipcNvlsMemory.h +++ b/cpp/include/tensorrt_llm/runtime/ipcNvlsMemory.h @@ -1,5 +1,5 @@ /* - * Copyright (c) 2022-2026, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -44,23 +44,8 @@ struct IpcNvlsHandle void MPI_group_barrier(std::set ranks); -//! \brief Whether NVLS (NVLink SHARP) multicast memory can be allocated on this -//! node. Checks only the static capability (CUDA driver >= 12010 and -//! CU_DEVICE_ATTRIBUTE_MULTICAST_SUPPORTED). This is the precondition for -//! ipcNvlsAllocate(): the allocator itself selects a fabric or POSIX-FD handle -//! via getMemHandleType(), so single-node NVLS works over POSIX-FD even when the -//! fabric/IMEX plane is not provisioned. The result is cached. bool ipcNvlsSupported(); -//! \brief Whether the NVLink fabric/IMEX plane is provisioned so that NVLS -//! multicast memory can actually be *bound* (not merely statically supported). -//! Extends ipcNvlsSupported() with a live fabric probe (getMemHandleType() must -//! resolve to CU_MEM_HANDLE_TYPE_FABRIC). Use this to decide whether NCCL may -//! attempt NVLS: on an unprovisioned node the static multicast attribute is a -//! false positive and NCCL aborts during init, so callers should disable -//! NCCL_NVLS when this returns false. The (heavy) result is cached. -bool ipcNvlsFabricUsable(); - IpcNvlsHandle* ipcNvlsAllocate(size_t size, std::set ranks); void ipcNvlsFree(IpcNvlsHandle* handle); diff --git a/cpp/include/tensorrt_llm/runtime/utils/mpiUtils.h b/cpp/include/tensorrt_llm/runtime/utils/mpiUtils.h index 75ec7a534815..be12b743db1d 100644 --- a/cpp/include/tensorrt_llm/runtime/utils/mpiUtils.h +++ b/cpp/include/tensorrt_llm/runtime/utils/mpiUtils.h @@ -467,39 +467,6 @@ int getNumNodes(); void initialize(MpiThreadSupport threadMode = MpiThreadSupport::THREAD_MULTIPLE, bool forwardAbortToParent = false); -//! \brief Returns true iff the WideEP fault-tolerance MPI mode is enabled. -//! -//! Reads the `TLLM_FAULT_TOLERANCE_MODE` environment variable and returns -//! `true` exactly when its value is the string `"1"`. Any other value -//! (including unset, empty, `"0"`, `"true"`, `"yes"`, etc.) returns `false`. -//! -//! This is an internal/dev knob for the WideEP fault-tolerance MVP (PR 1d.0). -//! It is intentionally not surfaced as a user-facing CLI/config option; -//! a proper `LLMArgs` field will replace it in PR 1d.1. -bool isFaultToleranceModeEnabled(); - -//! \brief Install WideEP fault-tolerance signal handlers on `SIGABRT` and `SIGSEGV`. -//! -//! Replaces the default handlers (which call `MPI_Abort(MPI_COMM_WORLD)` -//! and optionally `kill(parent, SIGKILL)`) with handlers that call -//! `_exit(137)` instead. This lets surviving ranks outlive a peer's -//! crash so the higher-level FT layers (kernel rank-mask, EPLB remap, -//! AlltoAll watchdog, FT subcomm) get a chance to run. -//! -//! `_exit` (not `exit`) is used because it is async-signal-safe and -//! skips `atexit` / Python finalizers / `MPI_Finalize`, all of which -//! can deadlock on a poisoned state. -//! -//! \note This source-level change is necessary but not sufficient on -//! its own: OpenMPI's `mpirun` terminates the world on any abnormal -//! child exit (any code, any signal) unless launched with -//! `--mca orte_enable_recovery 1`. See the WideEP FT design §5.4 -//! and audit-1a-findings.md Day 2 for the empirical results that motivate -//! this combination. -//! -//! Idempotent: safe to call multiple times. -void installFaultToleranceSignalHandlers(); - class MpiWaitThread { public: diff --git a/cpp/include/tensorrt_llm/runtime/utils/pgUtils.h b/cpp/include/tensorrt_llm/runtime/utils/pgUtils.h index 049de46cb26f..32a3f1c86c9d 100644 --- a/cpp/include/tensorrt_llm/runtime/utils/pgUtils.h +++ b/cpp/include/tensorrt_llm/runtime/utils/pgUtils.h @@ -86,8 +86,6 @@ c10::intrusive_ptr get_local_pg(); void init_pg(c10::intrusive_ptr const& process_group_world, c10::intrusive_ptr const& process_group_local); -void shutdown_pg(); - // Tensor wrapping utilities for ProcessGroup operations inline torch::Tensor wrap_tensor(torch::Tensor data) { diff --git a/cpp/include/tensorrt_llm/runtime/virtualMemory.h b/cpp/include/tensorrt_llm/runtime/virtualMemory.h index d74673f37e01..a7e95b42d707 100644 --- a/cpp/include/tensorrt_llm/runtime/virtualMemory.h +++ b/cpp/include/tensorrt_llm/runtime/virtualMemory.h @@ -505,7 +505,10 @@ class CudaVirtualMemoryAllocator { std::size_t gpuAlignment = 1; CUmemAllocationProp const prop{CU_MEM_ALLOCATION_TYPE_PINNED, CU_MEM_HANDLE_TYPE_NONE, - CUmemLocation{CU_MEM_LOCATION_TYPE_DEVICE, {device}}}; + { + CU_MEM_LOCATION_TYPE_DEVICE, + device, + }}; TLLM_CU_CHECK( cuMemGetAllocationGranularity(&gpuAlignment, &prop, CU_MEM_ALLOC_GRANULARITY_RECOMMENDED)); alignment = std::lcm(getpagesize(), gpuAlignment); diff --git a/cpp/kernels/fmha_v2/setup.py b/cpp/kernels/fmha_v2/setup.py index cb854eab6205..0d505c233bc6 100644 --- a/cpp/kernels/fmha_v2/setup.py +++ b/cpp/kernels/fmha_v2/setup.py @@ -6810,7 +6810,7 @@ def enumerate_kernels(): enumerate_qmma_flash_kernels(specs, sm=120, dtype='e4m3_fp32', - head_sizes=[64, 128, 192, 576], + head_sizes=[128, 192, 576], output_dtype="bf16") if 'ENABLE_HMMA_FP32' in os.environ: diff --git a/cpp/kernels/fmha_v2/src/fmha/kernel_traits.h b/cpp/kernels/fmha_v2/src/fmha/kernel_traits.h index daebcbaaf677..e4a54252bf52 100644 --- a/cpp/kernels/fmha_v2/src/fmha/kernel_traits.h +++ b/cpp/kernels/fmha_v2/src/fmha/kernel_traits.h @@ -145,9 +145,7 @@ template < // Do we use half epilogue for the 2nd GEMM (hmma_fp32) bool BMM2_FP16_EPILOGUE = true, // non-positive means disabled - int SAGE_BLOCK_SIZE_Q_ = 0, int SAGE_BLOCK_SIZE_K_ = 0, int SAGE_BLOCK_SIZE_V_ = 0, - // Enable skip softmax attention feature. - bool ENABLE_SKIP_SOFTMAX_ = false> + int SAGE_BLOCK_SIZE_Q_ = 0, int SAGE_BLOCK_SIZE_K_ = 0, int SAGE_BLOCK_SIZE_V_ = 0> struct Kernel_traits_ { @@ -199,9 +197,6 @@ struct Kernel_traits_ SAGE_BLOCK_SIZE_V = SAGE_BLOCK_SIZE_V_ }; - // Are we enabling skip softmax attention feature? - static constexpr bool ENABLE_SKIP_SOFTMAX = ENABLE_SKIP_SOFTMAX_; - // TODO: expose these tiling params to the interface enum { @@ -1010,13 +1005,10 @@ template < // The output type. typename OutputType = typename Traits::A_type, // The sage attention block size for Q, K and V - int SAGE_BLOCK_SIZE_Q = 0, int SAGE_BLOCK_SIZE_K = 0, int SAGE_BLOCK_SIZE_V = 0, - // Enable skip softmax attention feature. - bool ENABLE_SKIP_SOFTMAX = false> + int SAGE_BLOCK_SIZE_Q = 0, int SAGE_BLOCK_SIZE_K = 0, int SAGE_BLOCK_SIZE_V = 0> using Kernel_traits_v2 = Kernel_traits_::Gmem_tile_o, S, D, DV, STEP, WARPS_M, WARPS_N, - CTAS_PER_HEAD, FLAGS, 2, MASK_VERSION, BMM2_FP16_EPILOGUE, SAGE_BLOCK_SIZE_Q, SAGE_BLOCK_SIZE_K, SAGE_BLOCK_SIZE_V, - ENABLE_SKIP_SOFTMAX>; + CTAS_PER_HEAD, FLAGS, 2, MASK_VERSION, BMM2_FP16_EPILOGUE, SAGE_BLOCK_SIZE_Q, SAGE_BLOCK_SIZE_K, SAGE_BLOCK_SIZE_V>; //////////////////////////////////////////////////////////////////////////////////////////////////// @@ -1046,13 +1038,11 @@ template < // The output type. typename OutputType = typename Traits::A_type, // The sage attention block size for Q, K and V - int SAGE_BLOCK_SIZE_Q = 0, int SAGE_BLOCK_SIZE_K = 0, int SAGE_BLOCK_SIZE_V = 0, - // Enable skip softmax attention feature. - bool ENABLE_SKIP_SOFTMAX = false> -using Kernel_traits_v2_q_k_v = Kernel_traits_::Gmem_tile_o, S, D, DV, STEP, WARPS_M, - WARPS_N, CTAS_PER_HEAD, FLAGS, 2, MASK_VERSION, BMM2_FP16_EPILOGUE, SAGE_BLOCK_SIZE_Q, SAGE_BLOCK_SIZE_K, - SAGE_BLOCK_SIZE_V, ENABLE_SKIP_SOFTMAX>; + int SAGE_BLOCK_SIZE_Q = 0, int SAGE_BLOCK_SIZE_K = 0, int SAGE_BLOCK_SIZE_V = 0> +using Kernel_traits_v2_q_k_v + = Kernel_traits_::Gmem_tile_o, S, D, DV, STEP, WARPS_M, WARPS_N, CTAS_PER_HEAD, FLAGS, + 2, MASK_VERSION, BMM2_FP16_EPILOGUE, SAGE_BLOCK_SIZE_Q, SAGE_BLOCK_SIZE_K, SAGE_BLOCK_SIZE_V>; //////////////////////////////////////////////////////////////////////////////////////////////////// @@ -1082,13 +1072,11 @@ template < // The output type. typename OutputType = typename Traits::A_type, // The sage attention block size for Q, K and V - int SAGE_BLOCK_SIZE_Q = 0, int SAGE_BLOCK_SIZE_K = 0, int SAGE_BLOCK_SIZE_V = 0, - // Enable skip softmax attention feature. - bool ENABLE_SKIP_SOFTMAX = false> -using Kernel_traits_v2_paged_kv_cache = Kernel_traits_::Gmem_tile_o, S, D, DV, STEP, WARPS_M, - WARPS_N, CTAS_PER_HEAD, FLAGS, 2, MASK_VERSION, BMM2_FP16_EPILOGUE, SAGE_BLOCK_SIZE_Q, SAGE_BLOCK_SIZE_K, - SAGE_BLOCK_SIZE_V, ENABLE_SKIP_SOFTMAX>; + int SAGE_BLOCK_SIZE_Q = 0, int SAGE_BLOCK_SIZE_K = 0, int SAGE_BLOCK_SIZE_V = 0> +using Kernel_traits_v2_paged_kv_cache + = Kernel_traits_::Gmem_tile_o, S, D, DV, STEP, WARPS_M, WARPS_N, CTAS_PER_HEAD, FLAGS, + 2, MASK_VERSION, BMM2_FP16_EPILOGUE, SAGE_BLOCK_SIZE_Q, SAGE_BLOCK_SIZE_K, SAGE_BLOCK_SIZE_V>; //////////////////////////////////////////////////////////////////////////////////////////////////// @@ -1118,13 +1106,11 @@ template < // The output type. typename OutputType = typename Traits::A_type, // The sage attention block size for Q, K and V - int SAGE_BLOCK_SIZE_Q = 0, int SAGE_BLOCK_SIZE_K = 0, int SAGE_BLOCK_SIZE_V = 0, - // Enable skip softmax attention feature. - bool ENABLE_SKIP_SOFTMAX = false> + int SAGE_BLOCK_SIZE_Q = 0, int SAGE_BLOCK_SIZE_K = 0, int SAGE_BLOCK_SIZE_V = 0> using Kernel_traits_v2_contiguous_kv_cache = Kernel_traits_::Gmem_tile_o, S, D, 0, STEP, WARPS_M, WARPS_N, CTAS_PER_HEAD, FLAGS, 2, - MASK_VERSION, BMM2_FP16_EPILOGUE, SAGE_BLOCK_SIZE_Q, SAGE_BLOCK_SIZE_K, SAGE_BLOCK_SIZE_V, ENABLE_SKIP_SOFTMAX>; + MASK_VERSION, BMM2_FP16_EPILOGUE, SAGE_BLOCK_SIZE_Q, SAGE_BLOCK_SIZE_K, SAGE_BLOCK_SIZE_V>; //////////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/cpp/kernels/fmha_v2/src/fmha/warpspec_sm120/README.md b/cpp/kernels/fmha_v2/src/fmha/warpspec_sm120/README.md deleted file mode 100644 index 5f5f396b8178..000000000000 --- a/cpp/kernels/fmha_v2/src/fmha/warpspec_sm120/README.md +++ /dev/null @@ -1,124 +0,0 @@ -# skip_softmax — TMA-load + sync-MMA warp-specialized FMHA for sm_120 / sm_121 - -> This is the sm_120 / sm_121 warp-specialized context FMHA that carries the -> per-warp **skip-softmax** optimization (hence the name). Only half of the -> Hopper warp-specialization recipe ports to consumer Blackwell: TMA-driven -> async loads survive, but async MMA does not (sm_120 / sm_121 have no -> `wgmma.async` equivalent), so the compute warps stay on `mma.sync` while a -> dedicated producer warp drives the loads with TMA. - -This directory implements a warp-specialized context FMHA for the sm_120 -family (sm_120 / sm_121). It targets BF16, causal mask, `head_dim == -head_dim_v` in `{128, 256}`, and the PACKED_QKV layout. The kernel carries the -per-warp skip-softmax optimization into the warp-specialized design. - -## Files - -| File | Role | -|------|------| -| `kernel_traits.h` | `Kernel_traits_skip_softmax_sm120`: wraps `fmha::Kernel_traits_v2` for the LDGSTS-friendly `Smem_tile_*` types, then layers on the producer/consumer warp roles, the granular smem buffers, the circular-buffer barriers, and the V re-tile (see below). | -| `dma_sync_mma.h` | Producer (`DMA::run`). Issues `cp.async.bulk.tensor.3d.shared::cta.global.tile` for Q / K / V into the granular buffers. `DMA::Host::init_params` builds the three `CUtensorMap` descriptors with the driver-API `cuTensorMapEncodeTiled`. | -| `compute_sync_mma.h` | Consumer (`Compute::run`). The kv-loop body — BMM1 (`fmha::gemm`) + softmax + causal mask + per-warp skip-softmax vote + BMM2 + epilogue — reading the granular `Smem_tile_q/k/v` per ring slot. | - -The translation unit and the in-engine dispatch bridges live in -`cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/skip_softmax_sm120/fused_multihead_flash_attention_ws_sm120.cu`, -and the entry kernel in -`cpp/kernels/fmha_v2/src/fused_multihead_flash_attention_kernel_ws_sm120.h`. - -## How the runner reaches this kernel - -This is the **default** sm_120 / sm_121 context FMHA — there is no opt-in flag. -`FusedMultiHeadAttentionXMMAKernelV2::run` dispatches every prefill whose config -matches (sm_120 / sm_121, BF16 in/out, causal, `head_dim == head_dim_v` in -`{128, 256}`, PACKED_QKV) **and** that carries no feature the kernel does not -implement (alibi, logit softcapping, sage attention, sliding-window / custom mask, -returning softmax stats, interleaved) to the `run_skip_softmax_*` bridges; every -other config falls through to the cubin/launcher path. - -The per-tile skip-softmax optimization is selected by -`Launch_params::enableSkipSoftmax` (set when a skip-softmax threshold `> 0` is -configured): the bridges instantiate the `ENABLE_SKIP_SOFTMAX = true` kernel -variant when skipping is requested, and the `false` variant — a plain -full-softmax prefill with no skip-check overhead — otherwise. - -The translation unit is compiled only into the `_context_attention_kernels_120` -CMake target (sm_120 family). The all-architecture dispatch TU references the -bridge symbols under `TLLM_ENABLE_SKIP_SOFTMAX_SM120`, which CMake defines only when -sm_120 is built, so builds that exclude sm_120 neither reference nor link the -(then-absent) symbols. - -## Design rationale - -### Why TMA loads, not "just split the warps" - -In the non-warp-specialized tiled kernel, the Q / K / V loads are *multi-thread* -LDGSTS operations: each of the 128 threads issues several `LDGSTS` instructions -to cover `(tile rows × D bytes)`. There is no way to "have warp 0 do the load" -without rewriting the gmem/smem tile load helpers — the partition is baked into -them. TMA fixes exactly this: a single descriptor + a single -`cp.async.bulk.tensor` from one thread issues an entire tile load, and the -consumers wait on an `mbarrier`. So the producer warp uses TMA, not LDGSTS. - -### TMA descriptor format - -Blackwell's TMA engine requires the driver-API `cuTensorMapEncodeTiled` -(128-byte `CUtensorMap`) descriptor — the same form the shipping -trtllmGenKernels FMHA uses. The fmha_v2 hand-rolled 64-byte `fmha::cudaTmaDesc` -(Hopper-era bit layout) is rejected and faults at `UTMALDG`. The descriptors -are built host-side in `DMA::Host::init_params` and passed to the kernel as -`__grid_constant__` params. - -### Why the LDGSTS smem tiles can be filled by TMA - -The make-or-break question for reusing the existing consumer `Smem_tile_*` is -whether their LDGSTS XOR swizzle equals a TMA hardware swizzle mode. It does: -the Q and K granular tiles use `BYTES_PER_ROW = 128`, `BYTES_PER_STS = 16`, -`ROWS_PER_XOR_PATTERN = 8`, i.e. a physical 16-byte chunk index of -`(col / 8) ^ (row % 8)` — byte-identical to the TMA 128B hardware swizzle. So a -chunked 128B-swizzle TMA load fills `Smem_tile_q/k` directly and the consumer's -`ldmatrix` reads correct data. - -### V is re-tiled to 64-wide DV chunks - -The natural `Smem_tile_v` packs the full `DV` (256) into the lead dim, giving -512-byte smem rows that no TMA swizzle mode can reproduce (`cuTensorMapEncodeTiled` -caps the leading box dim at the 128-byte swizzle width; a 512-byte leading dim -only encodes with `SWIZZLE_NONE`, which is plain row-major and does not match -the consumer's XOR-swizzled read). Instead, V is tiled into `BMM2_DV_CHUNK = 64` -wide groups so the V smem tile has `LEAD_DIM = 64` → 128-byte rows — the same -proven layout as K — and the existing `N == 64` `ldsmt` read path applies -unchanged. The producer streams `DV / 64` dv-chunks per kv-tile; the consumer -BMM2 contracts per dv-chunk into the corresponding `acc_o` sub-range. - -### `setmaxnreg` is unavailable here - -`setmaxnreg.{dec,inc}` is a Hopper / datacenter-Blackwell instruction -(sm_90 / 100 / 103); ptxas hard-errors on sm_120 / sm_121. The producer/consumer -register-budget split therefore does not exist on this hardware and is guarded -off (no-op on sm_120 / sm_121). - -## What the port wins, and what it does not - -Wins on sm_120 / sm_121: - -- **Fewer load instructions** — one `cp.async.bulk.tensor` per tile replaces - the many per-thread `LDGSTS` of the tiled kernel. -- **Per-buffer-slot waits** (`mbarrier`) instead of CTA-wide `__syncthreads()` - between load and compute: a consumer warp unblocks as soon as its tile lands. - -Does not win: - -- **MMA / softmax overlap** — there is no `wgmma.async` on sm_120, so a consumer - warp's `mma.sync` blocks its issuing thread until result registers commit. The - Hopper warpspec hides BMM1/BMM2 MMA latency behind softmax/`frag_p` work; that - is not achievable with sync MMA only. -- **Register-budget split** — `setmaxnreg` is unavailable (see above). - -## Relationship to a CuTe-DSL kernel - -CUTLASS 4.x has Blackwell sm_120 FMHA examples implementing the TMA-load + -sync-MMA pattern in CuTe DSL. A longer-term direction is to route the sm_120 / -sm_121 dispatch into a CuTe-DSL kernel. This fmha_v2 implementation maps the -relationship between the existing fmha_v2 infrastructure and that design and is -self-contained: the dispatch is gated, and the directory plus the entry-kernel -header are isolated (no other code includes them). diff --git a/cpp/kernels/fmha_v2/src/fmha/warpspec_sm120/compute_sync_mma.h b/cpp/kernels/fmha_v2/src/fmha/warpspec_sm120/compute_sync_mma.h deleted file mode 100644 index 8ea45b14f9e0..000000000000 --- a/cpp/kernels/fmha_v2/src/fmha/warpspec_sm120/compute_sync_mma.h +++ /dev/null @@ -1,436 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2011-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. - */ - -#pragma once - -// skip_softmax consumer: BMM1 + softmax + skip-softmax + BMM2 body. -// -// This is a port of fused_multihead_flash_attention_kernel_noloop_tiled.h -// where: -// -// * `gmem_q.load(smem_q)` / `gmem_k.load(smem_k)` / `gmem_v.load(smem_v)` -// are removed -- the producer warp (in dma_sync_mma.h) issues TMA loads -// into the ring instead. -// * `fmha::ldgdepbar()` + `__syncthreads()` between load and -// compute are replaced by `cbr_*.wait()` against the entry-produced -// mbarrier for the slot we're about to read. -// * After consumer is done with a slot, `cbr_*.complete(tidx == 0, slot)` -// arrives on the entry-consumed mbarrier so the producer can recycle. -// -// EVERYTHING ELSE -- BMM1 inner loop via fmha::gemm(), softmax, mask, the -// per-warp skip-softmax vote with log-threshold, the BMM2 split between -// skip-path and no-skip-path -- mirrors the non-warp-specialized tiled sm_120 -// kernel so the attention math is numerically identical. - -#include -#include -#include -#include -#include -#include // Single_cta, Block_info_padded - -namespace fmha -{ -namespace ws_sm120 -{ - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -template -struct Compute -{ - using Shared = typename Kernel_traits::Shared; - - using Cbr_q = typename Kernel_traits::Circular_buffer_q_reader; - using Cbr_k = typename Kernel_traits::Circular_buffer_k_reader; - using Cbr_v = typename Kernel_traits::Circular_buffer_v_reader; - - using Traits_p = typename Kernel_traits::Traits_p; - using Traits_o = typename Kernel_traits::Traits_o; - using Cta_tile_p = typename Kernel_traits::Cta_tile_p; - using Cta_tile_o = typename Kernel_traits::Cta_tile_o; - using Mma_tile_p = typename Kernel_traits::Mma_tile_p; - using Mma_tile_o = typename Kernel_traits::Mma_tile_o; - // MMA tile for the dv-chunked V read (MMAS_K kv-steps x MMAS_N=64/16 dv tiles). - using Mma_tile_v = typename Kernel_traits::Mma_tile_v; - - using Smem_tile_q = typename Kernel_traits::Smem_tile_q; - using Smem_tile_k = typename Kernel_traits::Smem_tile_k; - using Smem_tile_v = typename Kernel_traits::Smem_tile_v; - using Smem_tile_o = typename Kernel_traits::Smem_tile_o; - - using Gmem_tile_o = typename Kernel_traits::Gmem_tile_o; - - using Softmax = fmha::Softmax; - - enum - { - STEP_Q = Kernel_traits::STEP_Q - }; - - enum - { - STEP_KV = Kernel_traits::STEP_KV - }; - - enum - { - CAUSAL_MASK = Kernel_traits::CAUSAL_MASK - }; - - enum - { - ENABLE_SKIP_SOFTMAX = Kernel_traits::ENABLE_SKIP_SOFTMAX - }; - - enum - { - CHECK_NEG_INF = Kernel_traits::SLIDING_WINDOW_ATTENTION || Kernel_traits::CUSTOM_MASK - }; - - inline __device__ Compute() {} - - // Run on the consumer warps. tidx is the thread index within the - // consumer group (0 .. NUM_CONSUMER_WARPS*32 - 1). - template - inline __device__ void run(int tidx, Shared* shared, Params const& params) - { - // Block / head / batch indexing -- same as noloop_tiled.h. - int const bidb = blockIdx.z; - int const bidh = blockIdx.y; - int const q_loop = blockIdx.x; // 1 CTA per (B, H, Q-tile) - - fused_multihead_attention::Single_cta const binfo(params, bidb, bidh, 0, tidx); - - int const q_sequence_start = q_loop * STEP_Q + (binfo.actual_kv_seqlen - binfo.actual_q_seqlen); - if (binfo.stop_early(q_loop * STEP_Q)) - { - return; - } - - // Mask + softmax setup. - fmha::Mask_dispatcher mask( - params, binfo, tidx); - // Initialize the mask's query-row offset for this Q-tile. Without this, - // the causal diagonal defaults to q_sequence_start=0 and every Q-tile - // after the first masks against the wrong row range. (noloop_tiled.h - // does the same via mask.load() before the kv loop.) - mask.load(q_sequence_start); - // softmax tail buffer is in shared->smem_v's tail; noloop_tiled - // gives softmax `smem_[Smem_tile_q::BYTES_PER_TILE]` which is the - // K/V smem region. On skip_softmax we don't share -- softmax does not - // touch the K/V smem ring. If Softmax::USE_SHARED_MEMORY is needed, - // we'd allocate a separate softmax_scratch buffer in Shared (TODO). - Softmax softmax(params, /*smem_scratch=*/nullptr, bidb, tidx); - static_assert(!Softmax::USE_SHARED_MEMORY, - "skip_softmax consumer needs Softmax::USE_SHARED_MEMORY = false; if your " - "kernel_traits enables it, add a softmax_scratch buffer to Shared."); - - // Per-granular-buffer ring readers (Q/K stream the head dim, V streams - // kv-positions; each cycles GRANULAR_DEPTH buffers). - Cbr_q cbr_q(&shared->q_barriers); - Cbr_k cbr_k(&shared->k_barriers); - Cbr_v cbr_v(&shared->v_barriers); - - // Smem tiles constructed at the granular tile base (buffer 0). Each - // tile advances its internal read buffer via move_to_next_read_buffer - // in lockstep with the per-chunk barrier handshake below. The producer - // (dma_sync_mma.h) TMAs chunk c into buffer (c % GRANULAR_DEPTH). - Smem_tile_q smem_q(shared->q_buf(0), tidx); - Smem_tile_k smem_k(shared->k_buf(0), tidx); - Smem_tile_v smem_v(shared->v_buf(0), tidx); - - // ----- Per-row state shared by the whole KV loop -------------------- - fmha::Fragment_accumulator acc_o[Mma_tile_o::MMAS_M][Mma_tile_o::VALID_MMAS_N]; - using Acc_type_o = typename Traits_o::Accumulator_type; - fmha::Clear_accumulator::apply(acc_o); - - fmha::Tile_o_normalizer acc_o_normalizer(params, binfo); - float global_max[Softmax::ROWS_PER_THREAD]; - float global_sum[Softmax::ROWS_PER_THREAD]; - - // Tail-chunk MMA validity bound for BMM1 (head-dim chunking). When - // VALID_K is a multiple of the chunk K (the common case, e.g. D=256, - // chunk=64), all MMAs in every chunk are valid; otherwise the last - // chunk only runs its first BMM1_TAIL_MMAS_K_BOUND MMAs. - constexpr int BMM1_VALID_MMAS_K = Mma_tile_p::VALID_MMAS_K; - constexpr int BMM1_TAIL_MMAS_K_BOUND - = BMM1_VALID_MMAS_K % Mma_tile_p::MMAS_K ? BMM1_VALID_MMAS_K % Mma_tile_p::MMAS_K : Mma_tile_p::MMAS_K; - - // Skip-softmax: precompute log(threshold/L) once. - float const skip_softmax_log_threshold = ENABLE_SKIP_SOFTMAX - ? __logf(params.skip_softmax_threshold_scale_factor / static_cast(binfo.actual_kv_seqlen)) - : 0.0f; - - int const valid_seqlen - = CAUSAL_MASK ? min(q_sequence_start + Cta_tile_p::M, binfo.actual_kv_seqlen) : binfo.actual_kv_seqlen; - int const kv_loop_start = 0; - int const kv_loop_end = fmha::div_up(valid_seqlen, int(Cta_tile_p::N)) * int(Cta_tile_p::N); - int const kv_mask_loop_start = int(q_sequence_start / Cta_tile_p::N) * Cta_tile_p::N; - -#ifdef SKIP_SOFTMAX_STAT - // Skip-softmax block counters (compiled only in a -DSKIP_SOFTMAX_STAT - // build). tile_negligible is per-warp and uniform within a warp, so - // every thread keeps an identical local tally; only the elected thread - // (tidx == 0) flushes to global below, reporting the elected consumer - // warp's rate as the CTA proxy (same convention as noloop_tiled.h). - [[maybe_unused]] uint32_t skip_softmax_total = 0; - [[maybe_unused]] uint32_t skip_softmax_skipped = 0; -#endif - - // ----- KV loop ------------------------------------------------------ - for (int kv_loop = kv_loop_start; kv_loop < kv_loop_end; kv_loop += Cta_tile_p::N) - { - bool const first_step = (kv_loop == kv_loop_start); - bool tile_negligible = false; - bool const apply_mask = params.has_alibi || (kv_loop >= kv_mask_loop_start); - - fmha::Fragment_accumulator acc_p[Mma_tile_p::MMAS_M][Mma_tile_p::MMAS_N]; - using Acc_type_p = typename Traits_p::Accumulator_type; - fmha::Clear_accumulator::apply(acc_p); - - mask.move_to_offset(kv_loop); - - typename Smem_tile_q::Fragment frag_q[Mma_tile_p::MMAS_K][Mma_tile_p::MMAS_M]; - typename Smem_tile_k::Fragment frag_k[Mma_tile_p::MMAS_K][Mma_tile_p::MMAS_N]; - - // ---- BMM1: stream the head dim in NUM_BMM1_CHUNKS granular chunks. - // Each chunk c is a separate TMA-filled buffer (chunk % GRANULAR_DEPTH); - // we wait on its produced barrier, do MMAS_K MMAs, then signal - // consumed (all consumer threads arrive -> doubles as the - // pre-recycle sync) and advance to the next granular read buffer. - // This replaces noloop_tiled.h's per-chunk LDGSTS reload + - // ldgdepbar + __syncthreads. -#pragma unroll - for (int chunk = 0; chunk < Kernel_traits::NUM_BMM1_CHUNKS; ++chunk) - { - bool const is_tail = (chunk == Kernel_traits::NUM_BMM1_CHUNKS - 1); - int const k_slot = cbr_k.wait(); - int const q_slot = cbr_q.wait(); -#pragma unroll - for (int ki = 0; ki < Mma_tile_p::MMAS_K; ++ki) - { - smem_q.load(frag_q[ki], ki); - smem_k.load(frag_k[ki], ki); - if (!is_tail || Cta_tile_p::VALID_K % Cta_tile_p::K == 0 || ki < BMM1_TAIL_MMAS_K_BOUND) - { - fmha::gemm(acc_p, frag_q[ki], frag_k[ki]); - } - } - // Every consumer thread arrives on the consumed barrier (count - // == CONSUMER_THREADS), so the producer cannot overwrite the - // buffer until all reads are done. - cbr_k.complete(/*arrive=*/1, k_slot); - cbr_k.advance(); - cbr_q.complete(/*arrive=*/1, q_slot); - cbr_q.advance(); - smem_k.move_to_next_read_buffer(); - smem_q.move_to_next_read_buffer(); - } - - // ---- Softmax ---- - softmax.unpack(acc_p); - if (apply_mask) - { - if (params.has_alibi) - { - softmax.apply_mask_alibi(mask, bidh, params.alibi_params); - } - else - { - softmax.apply_mask(mask); - } - } - - // Hoist frag_p (lifted from noloop_tiled.h; skip-softmax pack lives in tail gate). - fmha::Fragment_a frag_p[Kernel_traits::TOTAL_BMM2_MMAS_K][Mma_tile_o::MMAS_M]; - - if (first_step) - { - softmax.template reduce(global_max); - softmax.template apply_exp_with_mask(global_max); - softmax.template reduce(global_sum); - if constexpr (ENABLE_SKIP_SOFTMAX) - { - softmax.pack(frag_p); - } - } - else - { - float tmp[Softmax::ROWS_PER_THREAD]; -#pragma unroll - for (int i = 0; i < Softmax::ROWS_PER_THREAD; i++) - { - tmp[i] = global_max[i]; - } - softmax.template reduce(global_max); - - // Per-warp skip-softmax vote. - if constexpr (ENABLE_SKIP_SOFTMAX) - { -#ifdef SKIP_SOFTMAX_STAT - ++skip_softmax_total; -#endif - bool skip = ((global_max[0] - tmp[0]) < skip_softmax_log_threshold); -#pragma unroll - for (int i = 1; i < Softmax::ROWS_PER_THREAD; i++) - { - skip = skip & ((global_max[i] - tmp[i]) < skip_softmax_log_threshold); - } - tile_negligible = __all_sync(0xffffffffu, skip); - if (tile_negligible) - { -#ifdef SKIP_SOFTMAX_STAT - ++skip_softmax_skipped; -#endif -#pragma unroll - for (int i = 0; i < Softmax::ROWS_PER_THREAD; i++) - { - global_max[i] = tmp[i]; - } - } - } - - if (!tile_negligible) - { - acc_o_normalizer.update(acc_o, global_max, tmp, global_sum); - softmax.template apply_exp_with_mask(global_max); -#pragma unroll - for (int i = 0; i < Softmax::ROWS_PER_THREAD; i++) - { - tmp[i] = global_sum[i]; - global_sum[i] = 0.f; - } - softmax.template reduce(global_sum); -#pragma unroll - for (int i = 0; i < Softmax::ROWS_PER_THREAD; i++) - { - global_sum[i] += tmp[i]; - } - if constexpr (ENABLE_SKIP_SOFTMAX) - { - softmax.pack(frag_p); - } - } - } - - // Baseline (non-skip-softmax) pack: same location as noloop_tiled.h. - if constexpr (!ENABLE_SKIP_SOFTMAX) - { - softmax.pack(frag_p); - } - - // ---- BMM2: V is tiled in DV (outer) x kv-positions (inner). - // Each sub-tile is a [kv-chunk, dv-chunk=64] 128-byte-row granular - // buffer. For dv-chunk dvc we contract all kv (across the kv-chunks) - // into the acc_o columns [dvc*MMAS_N .. +MMAS_N). frag_p (the - // softmax probs over all kv) is indexed by the global k-step - // kvc*MMAS_K + ki. We ALWAYS consume the V sub-tiles (to drain the - // producer pipeline) but skip the HMMAs on a negligible tile. - typename Smem_tile_v::Fragment frag_v[Mma_tile_v::MMAS_K][Mma_tile_v::VALID_MMAS_N]; - - bool const do_bmm2 = !(ENABLE_SKIP_SOFTMAX && tile_negligible); -#pragma unroll - for (int dvc = 0; dvc < Kernel_traits::NUM_BMM2_DV_CHUNKS; ++dvc) - { -#pragma unroll - for (int kvc = 0; kvc < Kernel_traits::NUM_BMM2_KV_CHUNKS; ++kvc) - { - int const v_slot = cbr_v.wait(); - if (do_bmm2) - { -#pragma unroll - for (int ki = 0; ki < Mma_tile_v::MMAS_K; ++ki) - { - int const p_ki = kvc * Mma_tile_v::MMAS_K + ki; // global frag_p k-step - smem_v.load(frag_v[ki], ki); -#pragma unroll - for (int ni = 0; ni < Mma_tile_v::VALID_MMAS_N; ++ni) - { - int const acc_ni = dvc * Mma_tile_v::VALID_MMAS_N + ni; -#pragma unroll - for (int mi = 0; mi < Mma_tile_o::MMAS_M; ++mi) - { - acc_o[mi][acc_ni].mma(frag_p[p_ki][mi], frag_v[ki][ni]); - } - } - } - } - cbr_v.complete(/*arrive=*/1, v_slot); - cbr_v.advance(); - smem_v.move_to_next_read_buffer(); - } - } - } - -#ifdef SKIP_SOFTMAX_STAT - // Flush this CTA's skip tally to the global counters. Only the elected - // thread (tidx == 0 of the consumer group) writes, so we record the - // first consumer warp's tally as the CTA proxy (noloop_tiled.h does the - // same). trtllm.py reads + prints skipped/total per layer when - // TRTLLM_PRINT_SKIP_SOFTMAX_STAT=1. - if constexpr (ENABLE_SKIP_SOFTMAX) - { - if (tidx == 0 && params.skip_softmax_total_blocks != nullptr) - { - atomicAdd(params.skip_softmax_total_blocks, skip_softmax_total); - atomicAdd(params.skip_softmax_skipped_blocks, skip_softmax_skipped); - } - } -#endif - - // ---- Epilogue: normalize acc_o by global_sum, store O ---- - // Ported from noloop_tiled.h, with two skip_softmax adaptations: - // * __syncthreads() -> named_barrier over the CONSUMER_THREADS group - // only (the producer warp is not part of the epilogue and must not - // be caught in a CTA-wide barrier). - // * Smem_tile_o aliases the start of the smem region (q/k/v buffers - // are free once the kv-loop is done); it needs - // Smem_tile_o::BYTES_PER_TILE bytes, which fits within the - // q+k+v span ahead of the barrier arrays. - acc_o_normalizer.update_sum(global_max, global_sum); - acc_o_normalizer.final_update(acc_o, global_sum); - - Gmem_tile_o gmem_o(params, binfo, tidx, q_loop * Gmem_tile_o::ROWS); - Smem_tile_o smem_o(&shared->smem_q[0], tidx); - -#pragma unroll - for (int ii = 0; ii < Gmem_tile_o::LOOPS; ++ii) - { - // Swizzle the elements and do the final cross-warp reduction. - smem_o.store(acc_o, ii); - // Make sure the data is in shared memory (consumer group only). - fmha::named_barrier_wait(Kernel_traits::CONSUMER_SYNC_BARRIER_ID, Kernel_traits::CONSUMER_THREADS); - - uint4 out[Gmem_tile_o::STGS_PER_LOOP]; - smem_o.load(out); - - // Make sure the data was read from shared memory before reuse. - if (ii < Gmem_tile_o::LOOPS - 1) - { - fmha::named_barrier_wait(Kernel_traits::CONSUMER_SYNC_BARRIER_ID, Kernel_traits::CONSUMER_THREADS); - } - - gmem_o.store(out, ii); - } - } -}; - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -} // namespace ws_sm120 -} // namespace fmha diff --git a/cpp/kernels/fmha_v2/src/fmha/warpspec_sm120/dma_sync_mma.h b/cpp/kernels/fmha_v2/src/fmha/warpspec_sm120/dma_sync_mma.h deleted file mode 100644 index 7a5cd777bd1c..000000000000 --- a/cpp/kernels/fmha_v2/src/fmha/warpspec_sm120/dma_sync_mma.h +++ /dev/null @@ -1,406 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2011-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. - */ - -#pragma once - -// Producer (TMA-loader) half of the sm_120 / sm_121 warp-specialized FMHA. -// -// Single dedicated warp issues `cp.async.bulk.tensor.2d` for Q (once) and -// K / V (every kv iter), arrives on the entry-produced mbarriers, and waits -// on the entry-consumed mbarriers before recycling a ring slot. -// -// Reuses the existing TMA descriptor + utmaldg helpers from -// fmha/hopper/utils_tma.h (which compile under __CUDA_ARCH__ >= 900, inclusive -// of sm_120 / sm_121) and the CircularBufferWriter from -// fmha/warpspec/circular_buffer.h (Arrive_wait-based, also CC >= 9.0). - -#include -#include // abort - -#include // CUtensorMap + cuTensorMapEncodeTiled (driver API) - -#include -#include -#include -#include -#include -#include -#include // Block_info_padded - -namespace fmha -{ -namespace ws_sm120 -{ - -//////////////////////////////////////////////////////////////////////////////////////////////////// -// -// Kernel-local TMA load helper using `shared::cta` (single-CTA) variant. -// -// The Hopper fmha::utmaldg<...> uses `shared::cluster` qualifier, which -// requires cluster launch (cluster_dim > 0) to be valid PTX. sm_120 / sm_121 -// kernels are launched without an explicit cluster attribute, so the -// shared::cluster variant emits an Illegal Instruction at runtime even -// though it assembles successfully. The `shared::cta` variant is the -// single-CTA-no-cluster form and is what we need on consumer Blackwell. -// -// CRITICAL: the descriptor passed here MUST be a driver -// API `CUtensorMap` built by `cuTensorMapEncodeTiled` (128 bytes). The -// fmha_v2 hand-rolled `fmha::cudaTmaDesc` (64 bytes, Hopper-era bit layout) -// is REJECTED by Blackwell's TMA engine -- UTMALDG.3D faults with an -// "Illegal Instruction" at runtime even though the PTX assembles. A minimal -// reproducer confirmed: hand-rolled desc -> illegal instruction; encode-tiled -// CUtensorMap -> loads correct data on GB10 (sm_121). -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline __device__ void utmaldg_3d_cta( - void const* p_desc, uint32_t smem_ptr, uint32_t smem_barrier, int32_t const (&coord)[3], uint32_t elect_one) -{ - if (elect_one) - { -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 - // Note the `.tile` qualifier: required on sm_120 / sm_121 PTX (omitting - // it emits an illegal-instruction at runtime). The trtllmGenKernels FMHA - // shipping today uses the same `.shared::cta.global.tile` variant - // (cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/cuda_ptx/cuda_ptx.h:2073). - asm volatile( - "cp.async.bulk.tensor.3d.shared::cta.global.tile.mbarrier::complete_tx::bytes " - "[%0], [%1, {%2, %3, %4}], [%5];\n" - : - : "r"(smem_ptr), "l"(reinterpret_cast(p_desc)), "r"(coord[0]), "r"(coord[1]), "r"(coord[2]), - "r"(smem_barrier) - : "memory"); -#endif - } -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -template -struct DMA -{ - using Shared = typename Kernel_traits::Shared; - - using Cbw_q = typename Kernel_traits::Circular_buffer_q_writer; - using Cbw_k = typename Kernel_traits::Circular_buffer_k_writer; - using Cbw_v = typename Kernel_traits::Circular_buffer_v_writer; - - enum - { - STEP_Q = Kernel_traits::STEP_Q - }; - - enum - { - STEP_KV = Kernel_traits::STEP_KV - }; - - enum - { - D = Kernel_traits::D - }; - - enum - { - DV = Kernel_traits::DV - }; - - enum - { - ELEMENT_BYTES = Kernel_traits::ELEMENT_BYTES - }; - - using Cta_tile_p = typename Kernel_traits::Cta_tile_p; - - enum - { - CAUSAL_MASK = Kernel_traits::CAUSAL_MASK - }; - - // Per-granular-chunk transaction byte counts (one chunk == one granular - // smem buffer). These MUST equal the smem buffer sizes so the TMA fills - // exactly one buffer (no over/underrun). - enum - { - TX_BYTES_Q = Kernel_traits::BYTES_PER_BUFFER_Q - }; - - enum - { - TX_BYTES_K = Kernel_traits::BYTES_PER_BUFFER_K - }; - - enum - { - TX_BYTES_V = Kernel_traits::BYTES_PER_BUFFER_V - }; - - explicit inline __device__ DMA(uint32_t elect_one) - : elect_one_(elect_one) - { - } - - // Runs on a single warp (32 threads); only thread 0 (`elect_one`) issues - // the cp.async.bulk.tensor.* instructions. - // - // The three TMA descriptors are CUtensorMaps built host-side by - // Host::init_params (cuTensorMapEncodeTiled) and passed in as - // __grid_constant__ kernel params. Each is a chunk descriptor: - // Q/K box = (BMM1_CHUNK_ELTS=64 head-dim, 1, STEP_Q / STEP_KV), 128B swiz - // V box = (DV, 1, BMM2_CHUNK_ELTS=32 kv-positions) - // - // Per kv-tile the producer streams the head dim of Q/K in NUM_BMM1_CHUNKS - // chunks (selected by coord[0]=c*64) and the kv-positions of V in - // NUM_BMM2_CHUNKS chunks (coord[2]=kv_loop+c*32), each into granular buffer - // (chunk % GRANULAR_DEPTH). The CircularBuffer consumed-barrier throttles - // the producer to <= GRANULAR_DEPTH chunks ahead of the consumer. - template - inline __device__ void run(Params const& params, Shared* shared, CUtensorMap const* desc_q, - CUtensorMap const* desc_k, CUtensorMap const* desc_v) - { - int const tidx = static_cast(threadIdx.x) & 31; - fused_multihead_attention::Single_cta const binfo( - params, blockIdx.z, blockIdx.y, 0, tidx); - if (binfo.stop_early(blockIdx.x * STEP_Q)) - { - return; - } - - int const bidh = static_cast(blockIdx.y); - int const bidh_kv = bidh / static_cast(params.h_q_per_kv); - int const q_row = static_cast(blockIdx.x) * STEP_Q; - - // The Q/K/V TMA descriptors span the whole packed [total_tokens, H, D] - // buffer, so the seq coordinate is the global row = this request's - // cumulative token offset (binfo.sum_s == cu_q_seqlens[bidb]) plus the - // request-local position; without it every batch element re-reads - // request 0. KV reuses sum_s (PACKED_QKV: K/V share Q's token range) -- - // sum_s_kv would deref the null cu_kv_seqlens on the self-attention path. - int const q_seq_offset = binfo.sum_s; - int const kv_seq_offset = binfo.sum_s; - - // kv-loop range -- MUST match the consumer's exactly (else the - // consumed-barrier handshake deadlocks). Mirror compute_sync_mma.h. - int const q_sequence_start = q_row + (binfo.actual_kv_seqlen - binfo.actual_q_seqlen); - int const valid_seqlen - = CAUSAL_MASK ? min(q_sequence_start + int(Cta_tile_p::M), binfo.actual_kv_seqlen) : binfo.actual_kv_seqlen; - int const kv_loop_end = fmha::div_up(valid_seqlen, int(Cta_tile_p::N)) * int(Cta_tile_p::N); - - Cbw_q cbw_q(&shared->q_barriers); - Cbw_k cbw_k(&shared->k_barriers); - Cbw_v cbw_v(&shared->v_barriers); - - for (int kv_loop = 0; kv_loop < kv_loop_end; kv_loop += STEP_KV) - { - // ---- BMM1: head-dim chunks of K and Q ---- -#pragma unroll - for (int c = 0; c < Kernel_traits::NUM_BMM1_CHUNKS; ++c) - { - int const head_off = c * Kernel_traits::BMM1_CHUNK_ELTS; - - int const k_slot = cbw_k.tmaReserve(elect_one_, TX_BYTES_K); - uint32_t const k_smem = __nvvm_get_smem_pointer(shared->k_buf(k_slot)); - uint32_t const k_bar = __nvvm_get_smem_pointer(cbw_k.barrier_ptr(k_slot)); - int32_t const k_coord[3] = {head_off, bidh_kv, kv_seq_offset + kv_loop}; - utmaldg_3d_cta(desc_k, k_smem, k_bar, k_coord, elect_one_); - - int const q_slot = cbw_q.tmaReserve(elect_one_, TX_BYTES_Q); - uint32_t const q_smem = __nvvm_get_smem_pointer(shared->q_buf(q_slot)); - uint32_t const q_bar = __nvvm_get_smem_pointer(cbw_q.barrier_ptr(q_slot)); - int32_t const q_coord[3] = {head_off, bidh, q_seq_offset + q_row}; - utmaldg_3d_cta(desc_q, q_smem, q_bar, q_coord, elect_one_); - } - - // ---- BMM2: V sub-tiles, tiled in DV (outer) x kv-positions (inner). - // Each sub-tile is a [BMM2_KV_CHUNK_ELTS kv, BMM2_DV_CHUNK dv] box - // (128-byte rows, 128B swizzle -- same layout as K). The consumer - // reads them in the same (dv, kv) order. -#pragma unroll - for (int dvc = 0; dvc < Kernel_traits::NUM_BMM2_DV_CHUNKS; ++dvc) - { - int const dv_off = dvc * Kernel_traits::BMM2_DV_CHUNK; -#pragma unroll - for (int kvc = 0; kvc < Kernel_traits::NUM_BMM2_KV_CHUNKS; ++kvc) - { - int const kv_off = kv_loop + kvc * Kernel_traits::BMM2_KV_CHUNK_ELTS; - - int const v_slot = cbw_v.tmaReserve(elect_one_, TX_BYTES_V); - uint32_t const v_smem = __nvvm_get_smem_pointer(shared->v_buf(v_slot)); - uint32_t const v_bar = __nvvm_get_smem_pointer(cbw_v.barrier_ptr(v_slot)); - int32_t const v_coord[3] = {dv_off, bidh_kv, kv_seq_offset + kv_off}; - utmaldg_3d_cta(desc_v, v_smem, v_bar, v_coord, elect_one_); - } - } - } - } - - uint32_t elect_one_; - - //////////////////////////////////////////////////////////////////////////////////////////////// - // - // Host-side TMA descriptor setup. - // - // Called once per LLM forward (or once per layer if descriptors are layer- - // varying) before the kernel launch. Builds three driver-API CUtensorMaps - // (cuTensorMapEncodeTiled) for Q / K / V. These are passed into the kernel - // as __grid_constant__ params (see fused_multihead_flash_attention_ws_sm120.cu). - // - // Why the driver API (not the fmha_v2 hand-rolled fmha::cudaTmaDesc): - // The 64-byte hand-rolled descriptor uses a Hopper-era bit layout that - // Blackwell's TMA engine rejects -- UTMALDG faults at runtime. The - // 128-byte CUtensorMap from cuTensorMapEncodeTiled is the only portable, - // Blackwell-valid form (it is what the shipping trtllmGenKernels FMHA - // uses). Proven via reproducer on GB10 (sm_121). - // - // Each descriptor: - // * 3D tensor (D, H_or_HKV, total_seqlen), fastest-varying axis = head - // dim (contiguous), and a 3D box of (D, 1, STEP_Q / STEP_KV). - // * Element format BFloat16 / Float16 (2-byte). - // - // Swizzle note: cuTensorMapEncodeTiled requires the leading box dim *in - // bytes* (boxDim[0] * ELEMENT_BYTES) to be <= the swizzle width. For - // head_dim=256 BF16 that is 512 bytes, which exceeds the 128B max, so the - // whole-head-dim box only encodes with SWIZZLE_NONE. Matching a 128B - // swizzle (what the consumer Smem_tile ldmatrix wants) requires splitting - // the head-dim load into <=64-element chunks -- that is the remaining - // producer/consumer layout-matching work (see README phase plan). For now - // we pick the widest swizzle the leading dim permits. - // - // Scope of v0: - // * BF16 / FP16 only (2-byte). FP8 needs U8 format + V-transpose. - // * PACKED_QKV input layout only. - // * No TMA store -- epilogue still uses the scalar STG Gmem_tile_o path. - //////////////////////////////////////////////////////////////////////////////////////////////// - - struct Host - { - Host() = default; - - // Pick the widest swizzle the leading box dim (bytes) permits. - static CUtensorMapSwizzle pick_swizzle(uint32_t lead_bytes) - { - if (lead_bytes % 128 == 0 && lead_bytes <= 128) - return CU_TENSOR_MAP_SWIZZLE_128B; - if (lead_bytes % 64 == 0 && lead_bytes <= 64) - return CU_TENSOR_MAP_SWIZZLE_64B; - if (lead_bytes % 32 == 0 && lead_bytes <= 32) - return CU_TENSOR_MAP_SWIZZLE_32B; - return CU_TENSOR_MAP_SWIZZLE_NONE; - } - - // Encode one 3D tiled descriptor. tensor_size / box_size are in - // elements (fastest-varying first); the global stride of dim>=1 is in - // bytes (dim 0 is implicitly contiguous at element size). - static void encode(CUtensorMap& out, void* gmem_ptr, uint32_t const (&tensor_size)[3], - uint64_t seq_stride_bytes, uint32_t const (&box_size)[3]) - { - // This kernel is BF16-only (Ampere_hmma_bf16_traits). The descriptor - // data type only drives out-of-bounds fill, which is disabled below - // (CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE), so TMA moves raw bytes regardless - // of fmt -- a 2-byte element therefore maps to BFLOAT16 safely. If FP16 - // support is ever added, distinguish it here with a real data-type knob. - CUtensorMapDataType fmt = (Kernel_traits::ELEMENT_BYTES == 2) ? CU_TENSOR_MAP_DATA_TYPE_BFLOAT16 - : CU_TENSOR_MAP_DATA_TYPE_FLOAT32; - - // globalDim fastest-first: (D, H, seq). - uint64_t global_dim[3] = {tensor_size[0], tensor_size[1], tensor_size[2]}; - // globalStrides are for dims 1.. (dim 0 implicit). Bytes. - // dim1 (head) stride = D * ELEMENT_BYTES (next head) - // dim2 (seq) stride = seq_stride_bytes - uint64_t global_stride[2] - = {static_cast(tensor_size[0]) * Kernel_traits::ELEMENT_BYTES, seq_stride_bytes}; - uint32_t box_dim[3] = {box_size[0], box_size[1], box_size[2]}; - uint32_t elem_stride[3] = {1, 1, 1}; - - uint32_t const lead_bytes = box_size[0] * Kernel_traits::ELEMENT_BYTES; - CUtensorMapSwizzle swizzle = pick_swizzle(lead_bytes); - - CUresult res = cuTensorMapEncodeTiled(&out, fmt, /*rank=*/3, gmem_ptr, global_dim, global_stride, box_dim, - elem_stride, CU_TENSOR_MAP_INTERLEAVE_NONE, swizzle, CU_TENSOR_MAP_L2_PROMOTION_L2_128B, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - if (res != CUDA_SUCCESS) - { - char const* err = nullptr; - cuGetErrorString(res, &err); - fprintf(stderr, - "[skip_softmax] cuTensorMapEncodeTiled failed: %s " - "(dim=%u,%u,%u box=%u,%u,%u lead_bytes=%u swizzle=%d)\n", - err, tensor_size[0], tensor_size[1], tensor_size[2], box_size[0], box_size[1], box_size[2], - lead_bytes, static_cast(swizzle)); - // Abort rather than continue: `out` is left unencoded, so launching - // with it would be undefined behavior (a UTMALDG on a garbage - // descriptor). A failure here is a host-side programming error - // (invalid shape/swizzle), not a recoverable runtime condition. - abort(); - } - } - - template - void init_params(Params& params, Launch_params const& launch_params, CUtensorMap& tma_q, CUtensorMap& tma_k, - CUtensorMap& tma_v) const - { - uint32_t const d = params.d; - uint32_t const dv = params.dv; - uint32_t const h = params.h; - uint32_t const h_kv = params.h_kv; - - uint32_t const total_seqlen = params.is_s_padded ? static_cast(params.b * params.s) - : static_cast(launch_params.total_q_seqlen); - - static_assert( - Kernel_traits::ELEMENT_BYTES == 2, "skip_softmax v0 only supports BF16 / FP16 (2-byte elements)."); - static_assert(STEP_Q <= 256 && STEP_KV <= 256, "TMA box dimensions are capped at 256 elements per axis."); - - char* const q_ptr = reinterpret_cast(params.qkv_ptr); - // PACKED_QKV (H_q + H_kv + H_kv heads of D elements): - char* const k_ptr = q_ptr + h * d * Kernel_traits::ELEMENT_BYTES; - char* const v_ptr = k_ptr + h_kv * d * Kernel_traits::ELEMENT_BYTES; - - // Chunk widths: BMM1 streams the head dim in BMM1_CHUNK_ELTS-wide - // (=64, 128B) chunks -> box leading dim 64 -> 128B swizzle (matches - // the consumer Smem_tile_q/k granular buffer layout, verified). BMM2 - // streams kv-positions in BMM2_CHUNK_ELTS-wide (=32) chunks -> box - // seq dim 32, leading = full DV (SWIZZLE_NONE for now -- Smem_tile_v - // layout match is a follow-up). - constexpr uint32_t Q_CHUNK = Kernel_traits::BMM1_CHUNK_ELTS; - constexpr uint32_t K_CHUNK = Kernel_traits::BMM1_CHUNK_ELTS; - constexpr uint32_t V_DV_CHUNK = Kernel_traits::BMM2_DV_CHUNK; // 64 (128-byte leading) - constexpr uint32_t V_KV_CHUNK = Kernel_traits::BMM2_KV_CHUNK_ELTS; // 32 kv-positions - - // ---- Q ---- tensor (D, H, seq); box (chunk, 1, STEP_Q) - uint32_t const tensor_size_q[3] = {d, h, total_seqlen}; - uint32_t const box_size_q[3] = {Q_CHUNK, 1, STEP_Q}; - encode(tma_q, q_ptr, tensor_size_q, static_cast(params.q_stride_in_bytes), box_size_q); - - // ---- K ---- tensor (D, H_kv, seq); box (chunk, 1, STEP_KV) - uint32_t const tensor_size_k[3] = {d, h_kv, total_seqlen}; - uint32_t const box_size_k[3] = {K_CHUNK, 1, STEP_KV}; - encode(tma_k, k_ptr, tensor_size_k, static_cast(params.k_stride_in_bytes), box_size_k); - - // ---- V ---- tensor (DV, H_kv, seq); box (dv-chunk=64, 1, kv-chunk=32) - // -> 128-byte leading dim -> 128B swizzle (matches the re-tiled - // Smem_tile_v with LEAD_DIM=64). - uint32_t const tensor_size_v[3] = {dv, h_kv, total_seqlen}; - uint32_t const box_size_v[3] = {V_DV_CHUNK, 1, V_KV_CHUNK}; - encode(tma_v, v_ptr, tensor_size_v, static_cast(params.v_stride_in_bytes), box_size_v); - } - }; -}; - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -} // namespace ws_sm120 -} // namespace fmha diff --git a/cpp/kernels/fmha_v2/src/fmha/warpspec_sm120/kernel_traits.h b/cpp/kernels/fmha_v2/src/fmha/warpspec_sm120/kernel_traits.h deleted file mode 100644 index 5b1c89109ba5..000000000000 --- a/cpp/kernels/fmha_v2/src/fmha/warpspec_sm120/kernel_traits.h +++ /dev/null @@ -1,366 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2011-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. - */ - -#pragma once - -// Kernel_traits for the skip_softmax sm_120 / sm_121 warp-specialized FMHA. -// -// Wraps the existing fmha::Kernel_traits_ template (which already provides -// LDGSTS-friendly Smem_tile_a/b/v with ldmatrix swizzle and the right -// Cta_tile / Mma_tile / fragment shapes) and layers on the warp-spec -// pieces: -// * Shared struct with smem tiles + entry-produced / entry-consumed -// mbarrier arrays. -// * Circular_buffer_{q,k,v}_{reader,writer} type aliases against the -// existing fmha::ws::CircularBuffer infrastructure. -// * Named-barrier ids (collision-safe with the existing skip-softmax -// barrier ids 0x3 / 0x4 used on the non-warpspec tiled kernel). -// -// What this header does NOT include: -// * Host-side TMA descriptor setup (phase 3). -// * Persistence of which slot has which kv_loop offset (handled by -// ring writer state, not traits). - -#include -#include -#include - -namespace fmha -{ -namespace ws_sm120 -{ - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -template < - // The instruction traits (e.g. Ampere_hmma_bf16_traits) -- shared with the - // non-warpspec tiled kernel for sm_120. - typename Traits_, - // Sequence length upper bound (e.g. 8192 for typical chunked prefill). - int S, - // Hidden head dim (e.g. 128, 192, 256). - int VALID_D_, - // Hidden head dim of V (= D in standard MHA). - int VALID_DV_, - // The iteration step of the outer Q loop. - int STEP_Q_, - // The number of vertical warps in the compute group (consumer-side). - int WARPS_M_, - // The number of horizontal warps in the compute group. - int WARPS_N_, - // The version of the kernel (passes through to Kernel_traits_). - int VERSION_, - // The mask version of the kernel. - int MASK_VERSION_, - // Skip-softmax knob. - bool ENABLE_SKIP_SOFTMAX_ = false, - // Producer warp count -- single 32-thread warp by default. - int NUM_PRODUCER_WARPS_ = 1> -struct Kernel_traits_skip_softmax_sm120 -{ - - // Compose the existing PACKED_QKV tiled kernel traits with our skip_softmax - // overrides. Kernel_traits_v2 (in fmha/kernel_traits.h) bakes in - // fmha::v2::Gmem_tile_qkv for Q / K / V, which is the right gmem tile - // for the PACKED_QKV input layout TRT-LLM passes for Qwen3.5 prefill. - // - // FLAGS: bit 0x1 = USE_LDGSTS_Q - // bit 0x2 = USE_LDGSTS_K - // bit 0x4 = USE_LDGSTS_V - // Skip_softmax leaves all three LDGSTS bits OFF; the - // producer warp issues TMA, not LDGSTS. The smem - // layout is governed by BYTES_PER_LDG (=16, - // independent of USE_LDGSTS) and the buffer count - // (controlled by USE_GRANULAR_TILING), so turning - // LDGSTS off does NOT change the layout the - // consumer reads via ldmatrix. It also satisfies - // gmem_tile_qkv_packed.h's static assertion that - // USE_LDGSTS=>(PRED_REGS==1 || IS_HOPPER), which - // we'd fail on sm_120 with non-trivial Q tiles. - // bit 0x200 = NO_LOOP (we are a no-loop kernel) - // bit 0x1000= USE_GRANULAR_TILING (matches noloop_tiled.h) - static constexpr uint32_t TILED_FLAGS = 0x200u // NO_LOOP - | 0x1000u // USE_GRANULAR_TILING - ; - - using Base = fmha::Kernel_traits_v2; - - // Carry through the math types -- these are what compute_sync_mma.h needs. - using Traits_p = typename Base::Traits_p; - using Traits_o = typename Base::Traits_o; - using Traits_e = typename Base::Traits_e; - using Cta_tile_p = typename Base::Cta_tile_p; - using Cta_tile_o = typename Base::Cta_tile_o; - using Mma_tile_p = typename Base::Mma_tile_p; - using Mma_tile_o = typename Base::Mma_tile_o; - - // Smem tiles: reuse the existing tiled-kernel types and their buffer counts - // (Base::BUFFERS_PER_TILE_SMEM_*); the swizzle and the consumer's ldmatrix - // access patterns are independent of the buffer count. - using Smem_tile_q = typename Base::Smem_tile_q; - using Smem_tile_k = typename Base::Smem_tile_k; - using Smem_tile_o = typename Base::Smem_tile_o; - - // ----- V re-tiled to 64-wide DV chunks (6c.3) ----------------------------- - // - // The Base Smem_tile_v packs the full DV (256) into the lead dim, giving - // 512-byte smem rows that no TMA swizzle mode can reproduce (the encode caps - // the leading box dim at the 128-byte swizzle width). We instead tile DV - // into BMM2_DV_CHUNK=64-wide groups so the V smem tile has LEAD_DIM=64 -> - // 128-byte rows == the same proven layout as K (TMA-128B-swizzle fillable), - // and the existing Smem_tile_v `N==64` ldsmt read path applies unchanged. - static constexpr int BMM2_DV_CHUNK = 64; - using Cta_tile_v = typename Traits_o::template Cta_tile_extd; - using Smem_tile_v = fmha::Smem_tile_v; - // MMA tile for the dv-chunk V read (MMAS_K kv-steps, MMAS_N = 64/16 dv tiles). - using Mma_tile_v = typename Traits_o::template Mma_tile; - - using Gmem_tile_o = typename Base::Gmem_tile_o; - - enum - { - VALID_D = Base::VALID_D - }; - - enum - { - D = Base::D - }; - - enum - { - VALID_DV = Base::VALID_DV - }; - - enum - { - DV = Base::DV - }; - - enum - { - STEP_Q = STEP_Q_ - }; - - enum - { - STEP_KV = Cta_tile_p::N - }; - - enum - { - VERSION = VERSION_ - }; - - enum - { - MASK_VERSION = MASK_VERSION_ - }; - - enum - { - CAUSAL_MASK = Base::CAUSAL_MASK - }; - - enum - { - SLIDING_WINDOW_ATTENTION = Base::SLIDING_WINDOW_ATTENTION - }; - - enum - { - BIDIRECTIONAL_SLIDING_WINDOW_ATTENTION = Base::BIDIRECTIONAL_SLIDING_WINDOW_ATTENTION - }; - - enum - { - CUSTOM_MASK = Base::CUSTOM_MASK - }; - - enum - { - ELEMENT_BYTES = sizeof(typename Traits_p::A_type) - }; - - enum - { - TOTAL_BMM2_MMAS_K = Base::TOTAL_BMM2_MMAS_K - }; - - enum - { - ENABLE_BMM1_SOFTCAPPING_SCALE = Base::ENABLE_BMM1_SOFTCAPPING_SCALE - }; - - enum - { - IS_MTP = Base::IS_MTP - }; - - // Skip-softmax knob. - static constexpr bool ENABLE_SKIP_SOFTMAX = ENABLE_SKIP_SOFTMAX_; - - // Producer + consumer warp layout. - enum - { - NUM_PRODUCER_WARPS = NUM_PRODUCER_WARPS_ - }; - - enum - { - NUM_CONSUMER_WARPS = WARPS_M_ * WARPS_N_ - }; - - enum - { - THREADS = (NUM_PRODUCER_WARPS + NUM_CONSUMER_WARPS) * 32 - }; - - // Named-barrier ids. Collision-safe with the existing skip-softmax - // barriers (0x3, 0x4 on the non-warpspec path) since we don't run both - // kernels in the same launch. - static constexpr int DMA_SYNC_BARRIER_ID = 0x1; - static constexpr int MMA_SYNC_BARRIER_ID = 0x2; - - // Named-barrier id for the consumer group's pre-recycle sync (all consumer - // warps must finish reading a granular buffer before tid0 lets the producer - // overwrite it). - static constexpr int CONSUMER_SYNC_BARRIER_ID = 0x5; - - // Single CTA cluster (no DSMEM on consumer Blackwell -- CTAS_PER_CGA=1). - static constexpr int CTAS_PER_CGA = 1; - - enum - { - CONSUMER_THREADS = NUM_CONSUMER_WARPS * 32 - }; - - // ----- Granular head-dim / kv-position chunking --------------------------- - // - // The consumer reuses the existing LDGSTS Smem_tile_q/k/v which, with - // USE_GRANULAR_TILING, stream the contraction dim in chunks through a - // BUFFERS_PER_TILE-deep ping-pong (== GRANULAR_DEPTH below). The skip_softmax - // TMA producer fills those exact granular buffers chunk by chunk -- it does - // NOT add a ring on top. So the barrier depth == the granular buffer count. - static constexpr int GRANULAR_DEPTH = Smem_tile_k::BUFFERS_PER_TILE; - static_assert(GRANULAR_DEPTH == Smem_tile_q::BUFFERS_PER_TILE && GRANULAR_DEPTH == Smem_tile_v::BUFFERS_PER_TILE, - "skip_softmax assumes Q/K/V share the same granular buffer depth."); - - // Bytes of one granular buffer (one chunk) for Q / K / V. - static constexpr int BYTES_PER_BUFFER_Q = Smem_tile_q::BYTES_PER_BUFFER; - static constexpr int BYTES_PER_BUFFER_K = Smem_tile_k::BYTES_PER_BUFFER; - static constexpr int BYTES_PER_BUFFER_V = Smem_tile_v::BYTES_PER_BUFFER; - - // BMM1 streams the head dim in chunks of Cta_tile_p::K (=64) elements. - static constexpr int BMM1_CHUNK_ELTS = Cta_tile_p::K; // head-dim elements / chunk - static constexpr int NUM_BMM1_CHUNKS = Mma_tile_p::VALID_MMAS_K / Mma_tile_p::MMAS_K; - - // BMM2 V is tiled in BOTH kv-positions (Cta_tile_o::K = 32 each) and DV - // (BMM2_DV_CHUNK = 64 each). Each V sub-tile is one [kv-chunk, dv-chunk] - // granular buffer ([32, 64] -> 128-byte rows). The producer streams - // NUM_BMM2_DV_CHUNKS * NUM_BMM2_KV_CHUNKS sub-tiles per kv-tile. - static constexpr int BMM2_KV_CHUNK_ELTS = Cta_tile_o::K; // 32 - static constexpr int NUM_BMM2_KV_CHUNKS = TOTAL_BMM2_MMAS_K / Mma_tile_o::MMAS_K; // 4 - static constexpr int NUM_BMM2_DV_CHUNKS = VALID_DV / BMM2_DV_CHUNK; // 4 - static constexpr int NUM_BMM2_CHUNKS = NUM_BMM2_KV_CHUNKS * NUM_BMM2_DV_CHUNKS; // 16 - - using Circular_buffer_q_reader = typename fmha::ws::CircularBuffer::Reader; - using Circular_buffer_q_writer = typename fmha::ws::CircularBuffer::Writer; - using Circular_buffer_k_reader = typename fmha::ws::CircularBuffer::Reader; - using Circular_buffer_k_writer = typename fmha::ws::CircularBuffer::Writer; - using Circular_buffer_v_reader = typename fmha::ws::CircularBuffer::Reader; - using Circular_buffer_v_writer = typename fmha::ws::CircularBuffer::Writer; - - // Shared struct: the granular smem tiles (flat, == GRANULAR_DEPTH buffers - // laid out contiguously) + barrier arrays. The TMA producer writes chunk c - // into buffer (c % GRANULAR_DEPTH); the consumer's Smem_tile cycles its - // internal read buffer via move_to_next_read_buffer in lockstep. - struct __align__(128) Shared - { - uint8_t smem_q[Smem_tile_q::BYTES_PER_TILE]; - uint8_t smem_k[Smem_tile_k::BYTES_PER_TILE]; - uint8_t smem_v[Smem_tile_v::BYTES_PER_TILE]; - - fmha::ws::CircularBufferBarriers q_barriers; - fmha::ws::CircularBufferBarriers k_barriers; - fmha::ws::CircularBufferBarriers v_barriers; - - // Smem address of granular buffer `slot` for each tensor (producer TMA - // destination / consumer tile base + slot stride). - inline __device__ uint8_t* q_buf(int slot) - { - return &smem_q[slot * BYTES_PER_BUFFER_Q]; - } - inline __device__ uint8_t* k_buf(int slot) - { - return &smem_k[slot * BYTES_PER_BUFFER_K]; - } - inline __device__ uint8_t* v_buf(int slot) - { - return &smem_v[slot * BYTES_PER_BUFFER_V]; - } - - // Initialize all mbarriers. Called by thread 0 of the CTA at startup. - // entryProducedBarriers: count 1 -- the elect-one DMA thread arms the - // transaction count via tmaReserve; the TMA completion's tx-bytes - // arrival flips the barrier. - // entryConsumedBarriers: count CONSUMER_THREADS -- every consumer - // thread arrives once it has finished reading the buffer, which also - // serves as the pre-recycle sync (no separate __syncthreads needed). - inline __device__ void init(bool tid0) - { - if (tid0) - { -#pragma unroll - for (int i = 0; i < GRANULAR_DEPTH; i++) - { - fmha::bar_create(&q_barriers.entryProducedBarriers[i], 1); - fmha::bar_create(&q_barriers.entryConsumedBarriers[i], CONSUMER_THREADS); - fmha::bar_create(&k_barriers.entryProducedBarriers[i], 1); - fmha::bar_create(&k_barriers.entryConsumedBarriers[i], CONSUMER_THREADS); - fmha::bar_create(&v_barriers.entryProducedBarriers[i], 1); - fmha::bar_create(&v_barriers.entryConsumedBarriers[i], CONSUMER_THREADS); - } - } - } - }; - - // Pad to align. The non-Hopper kernel allocates BYTES_PER_SMEM in the - // extern __shared__ block; the skip_softmax version uses sizeof(Shared). - enum - { - BYTES_PER_SMEM = sizeof(Shared) - }; -}; - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -} // namespace ws_sm120 -} // namespace fmha diff --git a/cpp/kernels/fmha_v2/src/fused_multihead_flash_attention_kernel_ws_sm120.h b/cpp/kernels/fmha_v2/src/fused_multihead_flash_attention_kernel_ws_sm120.h deleted file mode 100644 index 040b43908ca6..000000000000 --- a/cpp/kernels/fmha_v2/src/fused_multihead_flash_attention_kernel_ws_sm120.h +++ /dev/null @@ -1,143 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2011-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. - */ - -#pragma once - -// ===================================================================== -// Skip_softmax: TMA-load + sync-MMA warp-specialized flash-attention prefill -// kernel for sm_120 / sm_121. -// ===================================================================== -// -// Named for the per-warp skip-softmax optimization it carries. Only half of the -// Hopper warp-specialization recipe ports to sm_120: TMA-driven async loads -// survive, but async MMA does not (sm_120 has no wgmma.async equivalent), so -// the compute warps stay on mma.sync. -// -// Differences from the non-warp-specialized tiled sm_120 path: -// -// 1. Producer / consumer warp split. -// - 1 producer warp (32 threads) issues cp.async.bulk.tensor for Q, K and -// V into the granular smem buffers via host-built CUtensorMap -// descriptors (cuTensorMapEncodeTiled). -// - The remaining warps consume those buffers and run a BMM1 + softmax + -// skip-softmax + BMM2 body. -// -// 2. mbarrier producer/consumer handshake instead of CTA-wide -// __syncthreads(): consumers unblock as soon as their tile's -// cp.async.bulk.tensor completes. -// -// Compute math: -// - sync mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 for both BMM1 -// and BMM2 (via fmha::gemm + Fragment_accumulator::mma in fmha/fragment.h). -// - Per-warp skip-softmax vote with a log-threshold predicate and a BMM2 -// split (skip vs no-skip). -// -// Note: setmaxnreg (the Hopper register-budget split) is NOT available on -// sm_120 / sm_121 and is intentionally not used here. -// -// See fmha/warpspec_sm120/README.md for the full design rationale. - -#include // CUtensorMap - -#include -#include -#include -#include - -namespace fused_multihead_attention -{ - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -namespace -{ - -// sm_120 register-budget split. Numbers below are starting points; -// final values depend on the launched warps and the per-tile working set. -constexpr int DMA_NREG = 40; // producer: TMA issue + coord math only -constexpr int COMPUTE_NREG = 232; // consumers: acc_o + softmax + frag_p live - -// NB: `setmaxnreg.{dec,inc}` is a Hopper / datacenter-Blackwell feature -// (sm_90, sm_100, sm_103). It is NOT supported on *consumer* Blackwell -// (sm_120 / sm_121) -- ptxas rejects it with a hard error there. So the -// producer/consumer register-budget split simply does not exist on the -// skip_softmax target hardware; these helpers compile to a no-op for sm_120/121. -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 && __CUDA_ARCH__ < 1200 -#define FMHA_HAS_SETMAXNREG 1 -#else -#define FMHA_HAS_SETMAXNREG 0 -#endif - -inline __device__ void setmaxnreg_dma() -{ -#if FMHA_HAS_SETMAXNREG - asm volatile("setmaxnreg.dec.sync.aligned.u32 %0;\n" ::"n"(DMA_NREG)); -#endif -} - -inline __device__ void setmaxnreg_compute() -{ -#if FMHA_HAS_SETMAXNREG - asm volatile("setmaxnreg.inc.sync.aligned.u32 %0;\n" ::"n"(COMPUTE_NREG)); -#endif -} - -} // namespace - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -template -inline __device__ void device_flash_attention_ws_sm120( - Params const& params, CUtensorMap const* desc_q, CUtensorMap const* desc_k, CUtensorMap const* desc_v) -{ - using Shared = typename Kernel_traits::Shared; - - // The shared struct contains: - // - smem_q[CIRCULAR_BUFFER_DEPTH] aligned tiles for Q - // - smem_k[CIRCULAR_BUFFER_DEPTH] aligned tiles for K - // - smem_v[CIRCULAR_BUFFER_DEPTH] aligned tiles for V - // - q_barriers / k_barriers / v_barriers (entry-produced / - // entry-consumed mbarrier pairs) - extern __shared__ char smem_[]; - char* smem_aligned = fmha::align_1024(smem_); - Shared* shared = reinterpret_cast(&smem_aligned[0]); - shared->init(threadIdx.x == 0); - __syncthreads(); - - // 32-thread warps. Warp 0 = TMA producer. Warps 1+ = sync-MMA consumers. - int const warp_id = threadIdx.x / 32; - int const lane = threadIdx.x % 32; - int const tidx_in_compute_group = threadIdx.x - 32; - - if (warp_id == 0) - { - setmaxnreg_dma(); - uint32_t const elect_one = (lane == 0) ? 1u : 0u; - fmha::ws_sm120::DMA dma(elect_one); - dma.run(params, shared, desc_q, desc_k, desc_v); - } - else - { - setmaxnreg_compute(); - fmha::ws_sm120::Compute compute; - compute.run(tidx_in_compute_group, shared, params); - } -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -} // namespace fused_multihead_attention diff --git a/cpp/kernels/xqa/mha.cu b/cpp/kernels/xqa/mha.cu index bc70bebe862a..593e27a62f95 100644 --- a/cpp/kernels/xqa/mha.cu +++ b/cpp/kernels/xqa/mha.cu @@ -1475,19 +1475,10 @@ __device__ inline void addAttentionSinks( { for (uint32_t i = 0; i < globalRowSum.size; i++) { - uint32_t const rowOffset = warp_size * i + laneId(); - if constexpr (SPEC_DEC) + uint32_t srcOffset = warp_size * i + laneId(); + if (srcOffset < headGrpSize) { - // Spec-dec rows flatten [query token, head], so repeat the sink indices for every token. - if (rowOffset < warpTile.y) - { - uint32_t const srcOffset = rowOffset % headGrpSize; - globalRowSum[i] += expf(attentionSinks[srcOffset] - globalRowMax[i]); - } - } - else if (rowOffset < headGrpSize) - { - globalRowSum[i] += expf(attentionSinks[rowOffset] - globalRowMax[i]); + globalRowSum[i] += expf(attentionSinks[srcOffset] - globalRowMax[i]); } } } @@ -1707,7 +1698,7 @@ CUBIN_EXPORT __global__ uint32_t const cacheSeqLen = getCacheSeqLen(cacheList, idxReq); #if SLIDING_WINDOW && SPEC_DEC && !IS_SPEC_DEC_TREE - uint32_t const tok0SeqLen = cacheSeqLen - actualQSeqLen + 1; + uint32_t const tok0SeqLen = cacheSeqLen - actualQSeqLen + 1 + idxHeadTokenInGrp; // ctaTokOffset; int32_t const tok0WinBeg = int32_t(tok0SeqLen) - int32_t(slidingWinSize); uint32_t const nbTotalSkipTokens = mha::max(0, tok0WinBeg); bool const rtIsReallySliding = (cacheSeqLen + actualQSeqLen > slidingWinSize); diff --git a/cpp/kernels/xqa/test/refAttention.cpp b/cpp/kernels/xqa/test/refAttention.cpp index 5208a31edf65..cc218f4cbd3c 100644 --- a/cpp/kernels/xqa/test/refAttention.cpp +++ b/cpp/kernels/xqa/test/refAttention.cpp @@ -180,8 +180,7 @@ template #if SPEC_DEC Eigen::Matrix refAttention(IOHead const* q, CacheSeq const& k, CacheSeq const& v, uint32_t seqLen, float qScale, - float kvScale, float xScale, uint32_t slidingWinSize, float* attentionSinks, bool* hostMask, const uint32_t qSeqLen, - const uint32_t q_len) + float kvScale, float xScale, uint32_t slidingWinSize, bool* hostMask, const uint32_t qSeqLen, const uint32_t q_len) { #else Eigen::Matrix refAttention(IOHead const* q, @@ -245,6 +244,7 @@ Eigen::Matrix refAttenti } // Add the attention sinks. +#if !SPEC_DEC if (attentionSinks != nullptr) { for (uint32_t i = 0; i < headGrpSize; i++) @@ -252,6 +252,7 @@ Eigen::Matrix refAttenti rowSum[i] += expf(attentionSinks[i] - rowMax[i]); } } +#endif Eigen::Matrix out = gemm1Acc.array().colwise() * (xScale * kvScale / rowSum.array()); @@ -264,7 +265,7 @@ Eigen::Matrix refAttenti template Eigen::Matrix \ refAttention(IOHead const* q, CacheSeq const& k, \ CacheSeq const& v, uint32_t seqLen, float qScale, float kvScale, float xScale, \ - uint32_t slidingWinSize, float* attentionSinks, bool* hostMask, const uint32_t qSeqLen, const uint32_t q_len) + uint32_t slidingWinSize, bool* hostMask, const uint32_t qSeqLen, const uint32_t q_len) #else #define INSTANTIATE_refAttention(prec, isPaged, useBeamSearch) \ template Eigen::Matrix \ diff --git a/cpp/kernels/xqa/test/refAttention.h b/cpp/kernels/xqa/test/refAttention.h index 8a9ecedb04de..8a3f67b53866 100644 --- a/cpp/kernels/xqa/test/refAttention.h +++ b/cpp/kernels/xqa/test/refAttention.h @@ -95,8 +95,7 @@ template #if SPEC_DEC Eigen::Matrix refAttention(IOHead const* q, CacheSeq const& k, CacheSeq const& v, uint32_t seqLen, float qScale, - float kvScale, float xScale, uint32_t slidingWinSize, float* attentionSinks, bool* hostMask, const uint32_t qSeqLen, - const uint32_t q_len); + float kvScale, float xScale, uint32_t slidingWinSize, bool* hostMask, const uint32_t qSeqLen, const uint32_t q_len); #else Eigen::Matrix refAttention(IOHead const* q, CacheSeq const& k, CacheSeq const& v, uint32_t seqLen, float qScale, diff --git a/cpp/kernels/xqa/test/test.cpp b/cpp/kernels/xqa/test/test.cpp index 8e54ea1c8495..934b3d62d7ec 100644 --- a/cpp/kernels/xqa/test/test.cpp +++ b/cpp/kernels/xqa/test/test.cpp @@ -1225,15 +1225,15 @@ void runTest(uint32_t batchSize, uint32_t seqLen, bool testPerf, bool refCheck, #endif #endif - auto const refAttentionSinks - = hasAttentionSinks ? attentionSinksPtr + headGrpSize * idxKHead : nullptr; #if SPEC_DEC Eigen::Matrix refOutput; refOutput = refAttention(&qHeads[req][b][q_len][runtimeHeadGrpSize * idxKHead], kCacheSeq, vCacheSeq, seqLen, qScaleForRef, kvCacheScale[0], xScale, slidingWinSize, - refAttentionSinks, hostMask, qSeqLen, q_len); + hostMask, qSeqLen, q_len); #else Eigen::Matrix refOutput; + auto const refAttentionSinks + = hasAttentionSinks ? attentionSinksPtr + headGrpSize * idxKHead : nullptr; if (useQGMMA) { refOutput = refFlashAttention(&qHeads[req][b][headGrpSize * idxKHead], kCacheSeq, @@ -1361,20 +1361,6 @@ TEST(RefCheck, llama_V2_70b_3) #endif } -#if SLIDING_WINDOW && !IS_SPEC_DEC_TREE -TEST(RefCheck, gpt_oss_spec_swa_rows) -{ - runTest<8, HEAD_GROUP_SIZE, 2>(1, 258, false, true, true, false, true, ~0U, 128); - runTest<8, HEAD_GROUP_SIZE, 4>(1, 260, false, true, true, false, true, ~0U, 128); - runTest<8, HEAD_GROUP_SIZE, 2>(4, 258, false, true, true, false, true, ~0U, 128); - runTest<8, HEAD_GROUP_SIZE, 4>(4, 260, false, true, true, false, true, ~0U, 128); - runTest<2, HEAD_GROUP_SIZE, 2>(1, 258, false, true, true, false, true, ~0U, 128); - runTest<2, HEAD_GROUP_SIZE, 4>(1, 260, false, true, true, false, true, ~0U, 128); - runTest<2, HEAD_GROUP_SIZE, 2>(4, 258, false, true, true, false, true, ~0U, 128); - runTest<2, HEAD_GROUP_SIZE, 4>(4, 260, false, true, true, false, true, ~0U, 128); -} -#endif - #endif #else diff --git a/cpp/tensorrt_llm/CMakeLists.txt b/cpp/tensorrt_llm/CMakeLists.txt index afd2d3a1f415..06aefb3886ba 100644 --- a/cpp/tensorrt_llm/CMakeLists.txt +++ b/cpp/tensorrt_llm/CMakeLists.txt @@ -200,8 +200,6 @@ set(TRTLLM_LINK_LIBS layers_src runtime_src testing_src - compressorKernels_src - mhcKernels_src userbuffers_src ${DECODER_SHARED_TARGET_0} ${DECODER_SHARED_TARGET_1}) diff --git a/cpp/tensorrt_llm/batch_manager/baseTransBuffer.cpp b/cpp/tensorrt_llm/batch_manager/baseTransBuffer.cpp index cbb8159bce44..3be8ee22bc1f 100644 --- a/cpp/tensorrt_llm/batch_manager/baseTransBuffer.cpp +++ b/cpp/tensorrt_llm/batch_manager/baseTransBuffer.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -21,60 +21,14 @@ #include "tensorrt_llm/common/logger.h" #include "tensorrt_llm/common/opUtils.h" -#include -#include #include namespace tensorrt_llm::batch_manager { -namespace -{ - -char const* bufferKindName(BufferKind kind) -{ - switch (kind) - { - case BufferKind::kKV: return "kv"; - case BufferKind::kKV_INDEXER: return "kv_indexer"; - case BufferKind::kRNN: return "rnn"; - } - return "unknown"; -} - -} // namespace - void BufferIndexHolder::release() noexcept { - if (mMgr == nullptr) - { - return; - } - if (mHeld) - { - try - { - if (mIsRecv) - { - mMgr->freeBufferIndexForRecv(mIndex); - } - else - { - mMgr->freeBufferIndexForSend(mIndex); - } - } - catch (...) - { - // noexcept: swallow so the destructor can never throw. - } - } - mHeld = false; - mMgr = nullptr; -} - -void BufferIndexHolder::poison() noexcept -{ - if (mMgr == nullptr) + if (!mHeld || mMgr == nullptr) { return; } @@ -82,19 +36,18 @@ void BufferIndexHolder::poison() noexcept { if (mIsRecv) { - mMgr->poisonBufferIndexForRecv(mIndex); + mMgr->freeBufferIndexForRecv(mIndex); } else { - mMgr->poisonBufferIndexForSend(mIndex); + mMgr->freeBufferIndexForSend(mIndex); } } catch (...) { - // noexcept: poison is a fail-closed best effort from exception paths. + // noexcept: swallow so the destructor can never throw. } mHeld = false; - mMgr = nullptr; } BaseTransBufferManager::BaseTransBufferManager( @@ -108,7 +61,7 @@ BaseTransBufferManager::BaseTransBufferManager( mRecvBufferCount = common::getEnvRequestKVCacheConcurrent() ? common::getEnvKVCacheRecvBufferCount() : 1; mSendBufferCount = common::getEnvKVCacheSendMaxConcurrenceNum(); mUseFabricMemory = !(common::getEnvKVCacheTransferUseSyncBuffer() || common::getEnvKVCacheTransferUseAsyncBuffer()) - && kv_cache_manager::FabricMemory::supportFabricMemory(); + && kv_cache_manager::FabricMemory::supportFbaricMemory(); if (mUseFabricMemory) { mTransferBufferSize = kv_cache_manager::FabricMemory::getAlignedSize(mTransferBufferSize); @@ -125,11 +78,9 @@ BaseTransBufferManager::BaseTransBufferManager( allocateBuffer(); } -std::optional BaseTransBufferManager::assignBufferIndexForSend( - std::atomic const* perRequestCancel, int64_t waitSliceMs) +std::optional BaseTransBufferManager::assignBufferIndexForSend() { - return assignBufferIndex( - mConcurrenceSendResource, mSendBufferCount, mOnlyUseDynamicBuffer, perRequestCancel, waitSliceMs); + return assignBufferIndex(mConcurrenceSendResource, mSendBufferCount, mOnlyUseDynamicBuffer); } void BaseTransBufferManager::freeBufferIndexForSend(std::optional bufferId) @@ -137,16 +88,9 @@ void BaseTransBufferManager::freeBufferIndexForSend(std::optional bufferId) freeBufferIndex(mConcurrenceSendResource, bufferId, mSendBufferCount, mOnlyUseDynamicBuffer); } -void BaseTransBufferManager::poisonBufferIndexForSend(std::optional bufferId) noexcept +std::optional BaseTransBufferManager::assignBufferIndexForRecv() { - poisonBufferIndex(mConcurrenceSendResource, bufferId, mSendBufferCount, mOnlyUseDynamicBuffer, "send"); -} - -std::optional BaseTransBufferManager::assignBufferIndexForRecv( - std::atomic const* perRequestCancel, int64_t waitSliceMs) -{ - return assignBufferIndex( - mConcurrenceRecvResource, mRecvBufferCount, mOnlyUseDynamicBuffer, perRequestCancel, waitSliceMs); + return assignBufferIndex(mConcurrenceRecvResource, mRecvBufferCount, mOnlyUseDynamicBuffer); } void BaseTransBufferManager::freeBufferIndexForRecv(std::optional bufferId) @@ -154,11 +98,6 @@ void BaseTransBufferManager::freeBufferIndexForRecv(std::optional bufferId) freeBufferIndex(mConcurrenceRecvResource, bufferId, mRecvBufferCount, mOnlyUseDynamicBuffer); } -void BaseTransBufferManager::poisonBufferIndexForRecv(std::optional bufferId) noexcept -{ - poisonBufferIndex(mConcurrenceRecvResource, bufferId, mRecvBufferCount, mOnlyUseDynamicBuffer, "recv"); -} - std::tuple, size_t, bool> BaseTransBufferManager::getOrAllocateSendBuffers( std::optional bufferId, int targetNum, std::vector const& requestedNumberOfElements, runtime::BufferManager const& bufferManagerToUse) @@ -310,51 +249,16 @@ void BaseTransBufferManager::allocateBuffer() } } -std::optional BaseTransBufferManager::assignBufferIndex(ConcurrenceResource& resource, size_t bufferCount, - bool onlyUseDynamicBuffer, std::atomic const* perRequestCancel, int64_t waitSliceMs) +std::optional BaseTransBufferManager::assignBufferIndex( + ConcurrenceResource& resource, size_t bufferCount, bool onlyUseDynamicBuffer) { - auto const isCancelled = [perRequestCancel]() - { return perRequestCancel != nullptr && perRequestCancel->load(std::memory_order_relaxed); }; - if (isCancelled()) - { - TLLM_THROW("Cache transfer buffer acquisition cancelled"); - } if (onlyUseDynamicBuffer) { - TLLM_CHECK_WITH_INFO(!resource.mPoisoned.load(std::memory_order_relaxed), - "Cannot assign dynamic cache transfer buffer kind=%s because the transfer buffer pool is poisoned", - bufferKindName(getBufferKind())); return std::nullopt; } std::unique_lock lk(resource.mBuffersMutex); - auto const predicate = [&resource, bufferCount]() - { - return resource.mPoisoned.load(std::memory_order_relaxed) - || static_cast(resource.mConcurrence) < bufferCount; - }; - if (perRequestCancel == nullptr) - { - resource.mBuffersCV.wait(lk, predicate); - } - else - { - auto const slice = std::chrono::milliseconds{waitSliceMs}; - while (!predicate()) - { - resource.mBuffersCV.wait_for(lk, slice); - if (isCancelled()) - { - TLLM_THROW("Cache transfer buffer acquisition cancelled"); - } - } - } - if (isCancelled()) - { - TLLM_THROW("Cache transfer buffer acquisition cancelled"); - } - TLLM_CHECK_WITH_INFO(!resource.mPoisoned.load(std::memory_order_relaxed), - "Cannot assign cache transfer buffer kind=%s because the transfer buffer pool is poisoned", - bufferKindName(getBufferKind())); + resource.mBuffersCV.wait( + lk, [&resource, bufferCount]() { return static_cast(resource.mConcurrence) < bufferCount; }); int bufferId = -1; for (size_t i = 0; i < bufferCount; i++) { @@ -384,12 +288,6 @@ void BaseTransBufferManager::freeBufferIndex( TLLM_CHECK(static_cast(bufferId.value()) < bufferCount); { std::scoped_lock lk(resource.mBuffersMutex); - if (resource.mBufferIndexFlag[bufferId.value()] == 2) - { - TLLM_LOG_ERROR("Refusing to free poisoned cache transfer buffer kind=%s index=%d", - bufferKindName(getBufferKind()), bufferId.value()); - return; - } resource.mBufferIndexFlag[bufferId.value()] = 0; } resource.mConcurrence--; @@ -397,48 +295,6 @@ void BaseTransBufferManager::freeBufferIndex( } } -void BaseTransBufferManager::poisonBufferIndex(ConcurrenceResource& resource, std::optional bufferId, - size_t bufferCount, bool onlyUseDynamicBuffer, char const* direction) noexcept -{ - resource.mPoisoned.store(true, std::memory_order_relaxed); - - if (onlyUseDynamicBuffer) - { - TLLM_LOG_ERROR("Poisoned dynamic %s cache transfer buffer kind=%s; process restart is required", direction, - bufferKindName(getBufferKind())); - resource.mBuffersCV.notify_all(); - return; - } - - try - { - if (bufferId.has_value()) - { - TLLM_CHECK(static_cast(bufferId.value()) < bufferCount); - { - std::scoped_lock lk(resource.mBuffersMutex); - if (resource.mBufferIndexFlag[bufferId.value()] == 1) - { - resource.mBufferIndexFlag[bufferId.value()] = 2; - } - } - } - TLLM_LOG_ERROR("Poisoned %s cache transfer buffer kind=%s index=%d; process restart is required", direction, - bufferKindName(getBufferKind()), bufferId.value_or(-1)); - } - catch (std::exception const& e) - { - TLLM_LOG_ERROR("Exception while poisoning %s cache transfer buffer kind=%s index=%d: %s", direction, - bufferKindName(getBufferKind()), bufferId.value_or(-1), e.what()); - } - catch (...) - { - TLLM_LOG_ERROR("Unknown exception while poisoning %s cache transfer buffer kind=%s index=%d", direction, - bufferKindName(getBufferKind()), bufferId.value_or(-1)); - } - resource.mBuffersCV.notify_all(); -} - size_t BaseTransBufferManager::getRecvBufferCount() { return mRecvBufferCount; diff --git a/cpp/tensorrt_llm/batch_manager/baseTransBuffer.h b/cpp/tensorrt_llm/batch_manager/baseTransBuffer.h index 8c1fd313f5a6..2cbf9f514bd7 100644 --- a/cpp/tensorrt_llm/batch_manager/baseTransBuffer.h +++ b/cpp/tensorrt_llm/batch_manager/baseTransBuffer.h @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -49,9 +49,8 @@ enum class BufferKind : uint8_t class BaseTransBufferManager; /// @brief RAII holder for an index from BaseTransBufferManager::assignBufferIndexFor{Send,Recv}. -/// Releases on destruction (incl. exception unwind). A dynamic-buffer holder has no -/// concrete index, but retains its manager binding so poison() can fail closed. -/// Move-only; call release() on the happy path or detach() when ownership is handed off downstream. +/// Releases on destruction (incl. exception unwind). Move-only; call release() on +/// the happy path or detach() when ownership is handed off downstream. class BufferIndexHolder { public: @@ -83,7 +82,6 @@ class BufferIndexHolder , mIsRecv(other.mIsRecv) { other.mHeld = false; - other.mMgr = nullptr; } BufferIndexHolder& operator=(BufferIndexHolder&& other) noexcept @@ -96,7 +94,6 @@ class BufferIndexHolder mHeld = other.mHeld; mIsRecv = other.mIsRecv; other.mHeld = false; - other.mMgr = nullptr; } return *this; } @@ -111,27 +108,18 @@ class BufferIndexHolder return mHeld; } - [[nodiscard]] bool isBoundTo(BaseTransBufferManager const& manager) const noexcept - { - return mMgr == &manager; - } - /// @brief Relinquish ownership without releasing. Use when a downstream /// owner (e.g. the formatter inside receiveSync) takes over the /// release responsibility on the happy path. std::optional detach() noexcept { mHeld = false; - mMgr = nullptr; return mIndex; } /// @brief Release the slot now and disarm the destructor. Safe to call multiple times. void release() noexcept; - /// @brief Fail-closed release for an exit path where transfer-buffer quiescence is unknown. - void poison() noexcept; - private: BaseTransBufferManager* mMgr{nullptr}; std::optional mIndex{}; @@ -139,8 +127,6 @@ class BufferIndexHolder bool mIsRecv{true}; }; -inline constexpr int64_t kBufferAcquireSliceMs = 100; - /// @brief Base class for cache transfer buffer management. /// Handles buffer pool allocation, index assignment, and slicing. /// Derived classes provide cache-specific size calculations. @@ -153,28 +139,20 @@ class BaseTransBufferManager /// @brief Assign a buffer index for sending. /// @return Assigned buffer index, or nullopt if using dynamic buffers. - std::optional assignBufferIndexForSend( - std::atomic const* perRequestCancel = nullptr, int64_t waitSliceMs = kBufferAcquireSliceMs); + std::optional assignBufferIndexForSend(); /// @brief Free a buffer index used for sending. /// @param bufferId The buffer index to free. void freeBufferIndexForSend(std::optional bufferId); - /// @brief Poison a send buffer index after an unquiesced transfer exit. - void poisonBufferIndexForSend(std::optional bufferId) noexcept; - /// @brief Assign a buffer index for receiving. /// @return Assigned buffer index, or nullopt if using dynamic buffers. - std::optional assignBufferIndexForRecv( - std::atomic const* perRequestCancel = nullptr, int64_t waitSliceMs = kBufferAcquireSliceMs); + std::optional assignBufferIndexForRecv(); /// @brief Free a buffer index used for receiving. /// @param bufferId The buffer index to free. void freeBufferIndexForRecv(std::optional bufferId); - /// @brief Poison a receive buffer index after an unquiesced transfer exit. - void poisonBufferIndexForRecv(std::optional bufferId) noexcept; - /// @brief Get or allocate send buffers for cache transfer. /// @param bufferId The assigned buffer ID. /// @param targetNum Number of target sequences. @@ -213,12 +191,6 @@ class BaseTransBufferManager return mMaxNumTokens; } - [[nodiscard]] bool hasPoisonedBuffer() const noexcept - { - return mConcurrenceSendResource.mPoisoned.load(std::memory_order_relaxed) - || mConcurrenceRecvResource.mPoisoned.load(std::memory_order_relaxed); - } - protected: /// @brief Constructor - derived classes call this after computing buffer sizes. /// @param transferBufferSize Size of each transfer buffer in bytes. @@ -234,7 +206,6 @@ class BaseTransBufferManager std::mutex mBuffersMutex; std::condition_variable mBuffersCV; std::atomic mConcurrence{0}; - std::atomic mPoisoned{false}; }; std::tuple, size_t, bool> getOrAllocateBuffers(std::optional bufferId, @@ -242,12 +213,9 @@ class BaseTransBufferManager runtime::BufferManager const& bufferManagerToUse, ConcurrenceResource& concurrenceResource); void allocateBuffer(); - std::optional assignBufferIndex(ConcurrenceResource& resource, size_t bufferCount, bool onlyUseDynamicBuffer, - std::atomic const* perRequestCancel = nullptr, int64_t waitSliceMs = kBufferAcquireSliceMs); + std::optional assignBufferIndex(ConcurrenceResource& resource, size_t bufferCount, bool onlyUseDynamicBuffer); void freeBufferIndex( ConcurrenceResource& resource, std::optional bufferId, size_t bufferCount, bool onlyUseDynamicBuffer); - void poisonBufferIndex(ConcurrenceResource& resource, std::optional bufferId, size_t bufferCount, - bool onlyUseDynamicBuffer, char const* direction) noexcept; size_t mPreAllocBufferSize; size_t mRecvBufferCount; diff --git a/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp b/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp index 9dc44531d409..68a398f1b52d 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp @@ -74,13 +74,6 @@ void sendBuffer(TransferSession& session, int deviceId, size_t localIdx, size_t bufferIdx = computeBufferIdx(localIdx, targetInfo); size_t size = outputBuffers[bufferIdx]->getSizeInBytes(); - // Skip Helix CP ranks that own no blocks for this sequence (num_total_blocks < cp_size). - // The matching gen rank skips its receive, so no 0-byte transfer is posted on either side. - if (size == 0) - { - return; - } - if (bufferIdx < bufferCoverTargetNum) { TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), " send connIdx: %ld bufferIdx: %ld size:%ld", connIdx, @@ -529,9 +522,7 @@ void CacheFormatter::format(tensorrt_llm::batch_manager::TransferSession& sessio // cache blocks to the corresponding buffer. // 5. send the buffer to the corresponding target. Ideally, we send only once (one buffer) for each target. - auto const* sendCancelFlag - = common::getEnvDisaggEnableInflightCancel() ? &session.getDataContext().getTransferTerminate() : nullptr; - auto cacheBufferId = mCacheTransBufferManager->assignBufferIndexForSend(sendCancelFlag); + auto cacheBufferId = mCacheTransBufferManager->assignBufferIndexForSend(); BufferIndexHolder sendHolder(*mCacheTransBufferManager, cacheBufferId, /*isRecv=*/false); int peerDuplicateHeadFactor = targetInfo.mPeerDupHeadFactor; auto bufferTargetNum = targetNum / peerDuplicateHeadFactor; @@ -611,24 +602,8 @@ void CacheFormatter::format(tensorrt_llm::batch_manager::TransferSession& sessio == inputKvCacheBlocksPerWindow.begin()->second.front()->getDataType()); } - if (sendCancelFlag != nullptr && sendCancelFlag->load(std::memory_order_relaxed)) - { - TLLM_THROW("KV cache transfer cancelled before NIXL submission"); - } - - try - { - sendAllBuffers(session, deviceId, outputSplitCaches, bufferCoverTargetNum, preAllocSendBuffer, - bufferManager, targetInfo, pickUpConnections); - } - catch (...) - { - if (agentConnection != nullptr && common::getEnvDisaggEnableInflightCancel()) - { - sendHolder.poison(); - } - throw; - } + sendAllBuffers(session, deviceId, outputSplitCaches, bufferCoverTargetNum, preAllocSendBuffer, bufferManager, + targetInfo, pickUpConnections); session.setTime(TransferSession::kTimeTransmissions); @@ -711,13 +686,6 @@ void CacheFormatter::unformat(tensorrt_llm::batch_manager::TransferSession& sess "outputBuffersPerWindow size: %ld,blockNum: %d , kvWindowSizes: %ld", outputBuffersPerWindow.size(), blockNum, kvWindowSizes.size()); TLLM_CHECK(!outputBuffersPerWindow.empty()); - - // An "empty" Helix CP rank owns no KV blocks for this sequence (num_total_blocks < cp_size). - // There is nothing to receive; the sender (context, CP=1) skips the matching 0-byte transfer. - if (blockNum == 0) - { - return; - } if (outputBuffersPerWindow.size() > 1) { // We only support limited case for VSWA. @@ -893,16 +861,12 @@ void CacheFormatter::unformat(tensorrt_llm::batch_manager::TransferSession& sess if (preAssignedKvId.has_value()) { cacheBufferId = static_cast(*preAssignedKvId); - if (!session.hasReservedRecvBuffer(*mCacheTransBufferManager)) - { - recvHolder = BufferIndexHolder(*mCacheTransBufferManager, cacheBufferId, /*isRecv=*/true); - } } else { cacheBufferId = mCacheTransBufferManager->assignBufferIndexForRecv(); - recvHolder = BufferIndexHolder(*mCacheTransBufferManager, cacheBufferId, /*isRecv=*/true); } + recvHolder = BufferIndexHolder(*mCacheTransBufferManager, cacheBufferId, /*isRecv=*/true); auto [recvSplitCachestmp, bufferCoverTargetNumtmp, onlyUseDynamicBuffer] = mCacheTransBufferManager->getOrAllocateRecvBuffers( cacheBufferId, static_cast(targetNum), bufferEleSizes, bufferManager); @@ -1036,7 +1000,6 @@ void CacheFormatter::unformat(tensorrt_llm::batch_manager::TransferSession& sess recvSplitCaches, outputBuffersPerWindow, destConfig, selfConfig, selfIdx, bufferManager); bufferManager.getStream().synchronize(); - (void) session.releaseReservedRecvBuffer(*mCacheTransBufferManager); recvHolder.release(); } session.setTime(TransferSession::kTimePostprocess); diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransBuffer.cpp b/cpp/tensorrt_llm/batch_manager/cacheTransBuffer.cpp index 772c9555f0f3..b6f1e4e8df15 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheTransBuffer.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheTransBuffer.cpp @@ -128,7 +128,7 @@ size_t FabricMemory::getAlignedSize(size_t size) return (size + granularity - 1) / granularity * granularity; } -bool FabricMemory::supportFabricMemory() +bool FabricMemory::supportFbaricMemory() { #ifdef __aarch64__ auto support_fun = []() @@ -309,7 +309,7 @@ size_t CacheTransBufferManager::preAllocBufferSize( transferBufferSize += validTokenNum * cacheSizeBytesPerToken; } } - bool useFabricMemory = FabricMemory::supportFabricMemory() + bool useFabricMemory = FabricMemory::supportFbaricMemory() && (!(common::getEnvKVCacheTransferUseSyncBuffer() || common::getEnvKVCacheTransferUseAsyncBuffer())); if (useFabricMemory) { diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransBuffer.h b/cpp/tensorrt_llm/batch_manager/cacheTransBuffer.h index 1635c11bc673..b63f18ab797f 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheTransBuffer.h +++ b/cpp/tensorrt_llm/batch_manager/cacheTransBuffer.h @@ -50,7 +50,7 @@ class FabricMemory size_t getSize() const; static size_t getAlignedSize(size_t size); - static bool supportFabricMemory(); + static bool supportFbaricMemory(); private: class Impl; diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp index aa05cd033892..6d3bfa658fdf 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp @@ -56,7 +56,6 @@ #include #include #include -#include #include #include @@ -70,78 +69,31 @@ namespace using RequestIdType = LlmRequest::RequestIdType; -constexpr int kTransferFuturePollIntervalMs = 10; +constexpr long long kNvbug6448152SlowConsensusThresholdUs = 1'000'000; -// Finite status checks are scheduler polls, not terminal deadlines. Pure polls -// use short slices; calls that ask for at least one completion keep bounded -// backpressure by waiting up to the configured future timeout. -std::chrono::milliseconds getTransferFutureWaitInterval( - std::optional const& configuredTimeoutMs, bool const needsProgress) +bool isNvbug6448152TraceEnabled() { - auto waitMs = kTransferFuturePollIntervalMs; - if (configuredTimeoutMs.has_value()) - { - waitMs = needsProgress ? configuredTimeoutMs.value() - : std::min(configuredTimeoutMs.value(), kTransferFuturePollIntervalMs); - } - return std::chrono::milliseconds(std::max(1, waitMs)); + return common::getBoolEnv("TRTLLM_NVBUG_6448152_TRACE"); } enum class TransferConsensusState : std::uint64_t { kCompleted = 1, kFailed = 2, - kTimedOut = 3, }; struct TransferStateCounts { int completedCount{0}; int failedCount{0}; - int timedOutCount{0}; }; struct TransferConsensusOutcome { std::unordered_set completedRequestIds; std::unordered_set failedRequestIds; - std::unordered_set timedOutRequestIds; }; -template -bool requestCancellationNoThrow(RequestIdType requestId, char const* transferKind, CancelFn&& cancelFn) noexcept -{ - try - { - return cancelFn(); - } - catch (std::exception const& error) - { - TLLM_LOG_ERROR( - "%s cancellation for request %ld failed and will be retried: %s", transferKind, requestId, error.what()); - } - catch (...) - { - TLLM_LOG_ERROR("%s cancellation for request %ld failed with an unknown error and will be retried", transferKind, - requestId); - } - return false; -} - -long getTransferElapsedMs(std::shared_ptr const& request, LlmRequest::TimePoint end) -{ - auto const elapsed - = std::chrono::duration_cast(end - request->getKvCacheTransferStart()); - return static_cast(elapsed.count()); -} - -std::vector sortedRequestIds(std::unordered_set const& requestIds) -{ - std::vector result(requestIds.begin(), requestIds.end()); - std::sort(result.begin(), result.end()); - return result; -} - void appendPackedTransferState( std::vector& packedStates, RequestIdType requestId, TransferConsensusState state) { @@ -181,11 +133,10 @@ std::vector gatherPackedTransferStates( TransferConsensusOutcome reduceTransferStates(std::shared_ptr const& comm, std::unordered_set const& completedRequestIds, - std::unordered_set const& failedRequestIds, - std::unordered_set const& timedOutRequestIds) + std::unordered_set const& failedRequestIds) { std::vector localStates; - localStates.reserve((completedRequestIds.size() + failedRequestIds.size() + timedOutRequestIds.size()) * 2); + localStates.reserve((completedRequestIds.size() + failedRequestIds.size()) * 2); for (auto const requestId : completedRequestIds) { if (failedRequestIds.find(requestId) == failedRequestIds.end()) @@ -197,10 +148,6 @@ TransferConsensusOutcome reduceTransferStates(std::shared_ptrgetSize() : 1; auto const gatheredStates @@ -220,7 +167,6 @@ TransferConsensusOutcome reduceTransferStates(std::shared_ptr 0) - { - outcome.timedOutRequestIds.insert(requestId); - } - if (terminalCount == syncSize && (counts.failedCount > 0 || counts.timedOutCount > 0)) + if (terminalCount == syncSize && counts.failedCount > 0) { outcome.failedRequestIds.insert(requestId); } @@ -247,13 +189,10 @@ TransferConsensusOutcome reduceTransferStates(std::shared_ptr const& firstComm, std::shared_ptr const& secondComm, std::unordered_set const& completedRequestIds, - std::unordered_set const& failedRequestIds, - std::unordered_set const& timedOutRequestIds) + std::unordered_set const& failedRequestIds) { - auto const firstOutcome - = reduceTransferStates(firstComm, completedRequestIds, failedRequestIds, timedOutRequestIds); - return reduceTransferStates( - secondComm, firstOutcome.completedRequestIds, firstOutcome.failedRequestIds, firstOutcome.timedOutRequestIds); + auto const firstOutcome = reduceTransferStates(firstComm, completedRequestIds, failedRequestIds); + return reduceTransferStates(secondComm, firstOutcome.completedRequestIds, firstOutcome.failedRequestIds); } void recordLocalTransferOutcome(RequestIdType requestId, std::shared_ptr request, bool failed, @@ -293,9 +232,6 @@ std::unique_ptr CacheTransceiverFactory::createCacheTransc TLLM_LOG_INFO("CacheTransceiver is disabled."); return nullptr; } - TLLM_CHECK_WITH_INFO(!common::getEnvDisaggEnableInflightCancel(), - "TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL=1 is supported only by the PyExecutor C++ NIXL transceiver path; " - "the legacy C++ executor does not provide the required deferred cleanup and poison escalation."); auto backendType = cacheTransceiverConfig.value().getBackendType(); if (backendType.value() == executor::CacheTransceiverConfig::BackendType::DEFAULT) { @@ -347,34 +283,13 @@ CacheTransceiver::CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheMa std::vector const& attentionLayerNumPerPP, nvinfer1::DataType dataType, executor::kv_cache::CacheState::AttentionType attentionType, std::optional cacheTransceiverConfig, - std::vector const& rnnLayerNumPerPP) - : mCacheTransceiverConfig{cacheTransceiverConfig} + rnn_state_manager::RnnStateManager* rnnStateManager, std::vector const& rnnLayerNumPerPP) + : mWorldRank{worldConfig.getRank()} + , mNvbug6448152TraceEnabled{isNvbug6448152TraceEnabled()} + , mCacheTransceiverConfig{cacheTransceiverConfig} + , mRnnStateManager{rnnStateManager} { using tensorrt_llm::batch_manager::kv_cache_manager::CacheFormatter; - TLLM_CHECK_WITH_INFO(mCacheTransceiverConfig.has_value(), "CacheTransceiverConfig is not set."); - auto const backendType = mCacheTransceiverConfig.value().getBackendType(); - TLLM_CHECK_WITH_INFO( - backendType.has_value() && (backendType.value() != executor::CacheTransceiverConfig::BackendType::DEFAULT), - " CacheTransceiverConfig::BackendType is not set."); - if (common::getEnvDisaggEnableInflightCancel()) - { - auto const nixlBackend = common::getEnvNixlBackend(); - TLLM_CHECK_WITH_INFO( - backendType.value() == executor::CacheTransceiverConfig::BackendType::NIXL && nixlBackend == "UCX", - "TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL=1 is experimental and currently supported only with the " - "NIXL cache transceiver and the UCX NIXL backend."); - TLLM_CHECK_WITH_INFO(mCacheTransceiverConfig->getKvTransferTimeoutMs().has_value(), - "TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL=1 requires kv_transfer_timeout_ms to enforce a finite deadline."); - TLLM_CHECK_WITH_INFO(!common::getEnvDisableKVCacheTransferOverlap(), - "TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL=1 requires asynchronous KV cache transfer; " - "TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP=1 is not supported."); - TLLM_CHECK_WITH_INFO(!common::getEnvDisaggLayerwise(), - "TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL=1 does not support layer-wise KV cache transfer."); - TLLM_CHECK_WITH_INFO(!common::getEnvTryZCopyForKVCacheTransfer(), - "TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL=1 does not support zero-copy KV cache transfer because request " - "blocks cannot be quarantined after an unquiesced cancellation."); - } - if (useMPI()) { mGroupComm = std::make_shared(std::addressof(tensorrt_llm::mpi::MpiComm::session())); @@ -425,6 +340,12 @@ CacheTransceiver::CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheMa } } bool isMLA = attentionType == executor::kv_cache::CacheState::AttentionType::kMLA; + TLLM_CHECK_WITH_INFO(mCacheTransceiverConfig.has_value(), "CacheTransceiverConfig is not set."); + auto backendType = mCacheTransceiverConfig.value().getBackendType(); + TLLM_CHECK_WITH_INFO( + backendType.has_value() && (backendType.value() != executor::CacheTransceiverConfig::BackendType::DEFAULT), + " CacheTransceiverConfig::BackendType is not set."); + std::optional maxNumTokens = mCacheTransceiverConfig.value().getMaxTokensInBuffer(); mCacheTransBufferManagers.push_back( @@ -435,9 +356,28 @@ CacheTransceiver::CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheMa std::make_unique(cacheManager, maxNumTokens, true)); } + // RNN specific setup + if (mRnnStateManager != nullptr) + { + TLLM_LOG_DEBUG("Setting up RNN cache transfer components."); + TLLM_CHECK(!rnnLayerNumPerPP.empty()); + + mRnnCacheTransBufferManager + = std::make_unique(mRnnStateManager, maxNumTokens); + + auto rnnModelCfg = mRnnStateManager->getRnnCacheStateModelConfig(); + + auto const convStateDataType = mRnnStateManager->getConvStateDataType(); + auto const ssmStateDataType = mRnnStateManager->getSsmStateDataType(); + + mCacheState->setRnnConfig(rnnModelCfg, rnnLayerNumPerPP, convStateDataType, ssmStateDataType); + + TLLM_LOG_INFO("RNN cache transfer components initialized."); + } + // Unified pool path (CppMambaHybridCacheManager): build RnnModelConfig from - // LinearAttentionMetadata. Detected by rnnLayerNumPerPP being non-empty. - if (!rnnLayerNumPerPP.empty()) + // LinearAttentionMetadata. Detected by rnnLayerNumPerPP set but no RnnStateManager. + if (mRnnStateManager == nullptr && !rnnLayerNumPerPP.empty()) { auto const& blockManager = cacheManager->getBlockManager(); auto const& linearMeta = blockManager.getLinearAttentionMetadata(); @@ -554,6 +494,11 @@ CacheTransceiver::CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheMa auto makeRnnFormatter = [this, cacheManager]() -> std::unique_ptr { + if (mRnnStateManager != nullptr && mRnnCacheTransBufferManager != nullptr) + { + // Slot-based path (CppMambaCacheManager) + return std::make_unique(mRnnStateManager, mRnnCacheTransBufferManager.get()); + } // Unified pool path (CppMambaHybridCacheManager) if (mCacheState->hasRnnConfig() && mRnnCacheTransBufferManager != nullptr) { @@ -573,11 +518,6 @@ CacheTransceiver::CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheMa CacheTransceiver::~CacheTransceiver() { - // Stop sender/receiver workers while the connection manager and transfer - // plugin are still alive. The workers can access both during termination. - mCacheSender.reset(); - mCacheReceiver.reset(); - if (mWrapperLibHandle) { std::lock_guard lock(mDllMutex); @@ -625,6 +565,17 @@ void CacheTransceiver::respondAndSendAsync(std::shared_ptr llmReques setContextState(llmRequest.get()); auto future = mCacheSender->sendAsync(llmRequest); mSenderFutures.emplace_back(std::move(llmRequest), std::move(future)); + if (mNvbug6448152TraceEnabled) + { + mNvbug6448152ContextEnqueuedTransitionCount++; + auto const requestId = mSenderFutures.back().first->mRequestId; + TLLM_LOG_INFO( + "NVBUG6448152_DIAG event=transfer_start side=ctx rank=%d request_id=%llu check_seq=%llu " + "future_count=%zu awaiting_count=%zu timeout_count=%zu", + mWorldRank, static_cast(requestId), + static_cast(mNvbug6448152ContextCheckSequence), mSenderFutures.size(), + mSenderRequestsAwaitingConsensus.size(), mNvbug6448152SenderWaitTimeoutIds.size()); + } } void CacheTransceiver::respondAndSendLayerWise( @@ -641,43 +592,26 @@ void CacheTransceiver::respondAndSendLayerWise( setContextState(llmRequest.get()); auto future = mCacheSender->sendAsync(llmRequest); mSenderFutures.emplace_back(llmRequest, std::move(future)); + if (mNvbug6448152TraceEnabled) + { + mNvbug6448152ContextEnqueuedTransitionCount++; + TLLM_LOG_INFO( + "NVBUG6448152_DIAG event=transfer_start side=ctx_layerwise rank=%d request_id=%llu " + "check_seq=%llu future_count=%zu awaiting_count=%zu timeout_count=%zu", + mWorldRank, static_cast(llmRequest->mRequestId), + static_cast(mNvbug6448152ContextCheckSequence), mSenderFutures.size(), + mSenderRequestsAwaitingConsensus.size(), mNvbug6448152SenderWaitTimeoutIds.size()); + } } } void CacheTransceiver::requestAndReceiveSync(std::shared_ptr llmRequest) { TLLM_CHECK(llmRequest && llmRequest->isGenerationOnlyRequest()); - auto const requestId = llmRequest->mRequestId; - auto const contextRequestId = llmRequest->getContextPhaseParams().value().getReqId(); - TLLM_LOG_DEBUG("Synchronous KV cache receive request %zu, context request %zu waiting for native completion.", - requestId, contextRequestId); - try { auto future = mCacheReceiver->receiveAsync(llmRequest); future.get(); } - catch (std::exception const& err) - { - llmRequest->setState(LlmRequestState::kDISAGG_TRANS_ERROR); - llmRequest->setKvCacheTransferEnd(LlmRequest::getSteadyClockNow()); - TLLM_LOG_ERROR("Synchronous KV cache receive request %zu, context request %zu failed: %s", requestId, - contextRequestId, err.what()); - return; - } - catch (...) - { - llmRequest->setState(LlmRequestState::kDISAGG_TRANS_ERROR); - llmRequest->setKvCacheTransferEnd(LlmRequest::getSteadyClockNow()); - TLLM_LOG_ERROR("Synchronous KV cache receive request %zu, context request %zu failed with an unknown error", - requestId, contextRequestId); - return; - } - if (llmRequest->getState() == LlmRequestState::kDISAGG_TRANS_ERROR) - { - TLLM_LOG_ERROR("Synchronous KV cache receive request %zu, context request %zu completed with an error state.", - requestId, contextRequestId); - return; - } llmRequest->setState(LlmRequestState::kDISAGG_GENERATION_TRANS_COMPLETE); } @@ -694,11 +628,19 @@ void CacheTransceiver::requestAndReceiveAsync(std::shared_ptr llmReq return; } - llmRequest->setKvCacheTransferStart(LlmRequest::getSteadyClockNow()); auto future = mCacheReceiver->receiveAsync(llmRequest); auto* requestPtr = llmRequest.get(); mRequesterFutures.emplace_back(std::move(llmRequest), std::move(future)); requestPtr->setState(LlmRequestState::kDISAGG_GENERATION_TRANS_IN_PROGRESS); + if (mNvbug6448152TraceEnabled) + { + TLLM_LOG_INFO( + "NVBUG6448152_DIAG event=transfer_start side=gen rank=%d request_id=%llu check_seq=%llu " + "future_count=%zu awaiting_count=%zu timeout_count=%zu", + mWorldRank, static_cast(requestPtr->mRequestId), + static_cast(mNvbug6448152GenerationCheckSequence), mRequesterFutures.size(), + mRequesterRequestsAwaitingConsensus.size(), mTimedOutRequesterIds.size()); + } } std::vector gatherRequestIds( @@ -798,20 +740,31 @@ void updateKVCacheTransferBW(std::shared_ptr const& mComm, RequestStatuses CacheTransceiver::checkContextTransferStatus( std::optional const& atLeastRequestNum, bool markComplete) { - bool const blockAll = !atLeastRequestNum.has_value(); - bool const inflightCancelEnabled = common::getEnvDisaggEnableInflightCancel(); - TLLM_CHECK_WITH_INFO(!inflightCancelEnabled || !blockAll, - "In-flight cancellation requires a finite context-transfer status poll; pass 0 for a nonblocking poll."); + std::uint64_t contextCheckSequence = 0; + bool traceConsensusTransition = false; + bool periodicConsensusCheckpoint = false; + std::uint64_t enqueuedTransitionCount = 0; + size_t localTerminalTransitionCount = 0; + size_t firstWaitTimeoutTransitionCount = 0; + if (mNvbug6448152TraceEnabled) + { + contextCheckSequence = ++mNvbug6448152ContextCheckSequence; + enqueuedTransitionCount = mNvbug6448152ContextEnqueuedTransitionCount; + mNvbug6448152ContextEnqueuedTransitionCount = 0; + periodicConsensusCheckpoint + = contextCheckSequence % 512 == 0 && (!mSenderFutures.empty() || !mSenderRequestsAwaitingConsensus.empty()); + traceConsensusTransition = enqueuedTransitionCount > 0 || periodicConsensusCheckpoint; + } + + bool blockAll = !atLeastRequestNum.has_value(); std::optional senderFutureTimeoutMs = std::nullopt; - if (mCacheTransceiverConfig.has_value()) + // If blockAll is true, we want to block and not use a timeout + if (!blockAll && mCacheTransceiverConfig.has_value()) { senderFutureTimeoutMs = mCacheTransceiverConfig->getKvTransferSenderFutureTimeoutMs(); } - bool const needsProgress = atLeastRequestNum.value_or(0) > 0; - auto const futureWaitInterval = getTransferFutureWaitInterval(senderFutureTimeoutMs, needsProgress); - // Without the opt-in flag, deadline checks remain observe-only. With the - // flag, timeout IDs participate in the same topology consensus as terminal - // outcomes and request cancellation is requested on every nonterminal rank. + // Observe-only: WARN per-request when the wall-clock transfer time exceeds + // kvTransferTimeoutMs. No cancellation, eviction, or state transition. std::optional kvTransferTimeoutMs = std::nullopt; if (mCacheTransceiverConfig.has_value()) { @@ -873,53 +826,71 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( for (auto it = mSenderFutures.begin(); it != mSenderFutures.end();) { auto& [request, future] = *it; - auto const requestId = request->mRequestId; - if (kvTransferTimeoutMs.has_value() - && future.wait_for(std::chrono::milliseconds(0)) != std::future_status::ready) + if (kvTransferTimeoutMs.has_value()) { - auto const elapsedMs = getTransferElapsedMs(request, LlmRequest::getSteadyClockNow()); - if (elapsedMs > kvTransferTimeoutMs.value()) + auto elapsed = std::chrono::duration_cast( + LlmRequest::getSteadyClockNow() - request->getKvCacheTransferStart()); + auto elapsedMs = static_cast(elapsed.count()); + if (elapsedMs > kvTransferTimeoutMs.value() && mTimedOutSenderIds.insert(request->mRequestId).second) { - if (mTimedOutSenderIds.insert(requestId).second) - { - TLLM_LOG_WARNING( - "Context KV cache transfer for request %ld exceeded configured timeout: " - "elapsed %ld ms > limit %d ms (%s).", - requestId, elapsedMs, kvTransferTimeoutMs.value(), - inflightCancelEnabled ? "requesting cancellation" : "observe-only"); - } + TLLM_LOG_WARNING( + "Context KV cache transfer for request %ld exceeded configured timeout: " + "elapsed %ld ms > limit %d ms (observe-only).", + request->mRequestId, elapsedMs, kvTransferTimeoutMs.value()); } } - if (blockAll || (toCompleteIdSet.find(requestId) != toCompleteIdSet.end())) + if (blockAll || (toCompleteIdSet.find(request->mRequestId) != toCompleteIdSet.end())) { + auto const requestId = request->mRequestId; try { - auto const status = blockAll ? std::future_status::ready : future.wait_for(futureWaitInterval); - if (status == std::future_status::ready) + // Wait for up to a specified timeout + auto status = future.wait_for(std::chrono::milliseconds(senderFutureTimeoutMs.value_or(0))); + if (status == std::future_status::ready || !senderFutureTimeoutMs.has_value()) { future.get(); - if (kvTransferTimeoutMs.has_value()) + bool const failed = request->getState() == LlmRequestState::kDISAGG_TRANS_ERROR; + long elapsedMs = 0; + if (mNvbug6448152TraceEnabled) { - auto const elapsedMs = getTransferElapsedMs(request, request->getKvCacheTransferEnd()); - if (elapsedMs > kvTransferTimeoutMs.value() && mTimedOutSenderIds.insert(requestId).second) - { - TLLM_LOG_WARNING( - "Context KV cache transfer for request %ld completed after its deadline: " - "elapsed %ld ms > limit %d ms (%s).", - requestId, elapsedMs, kvTransferTimeoutMs.value(), - inflightCancelEnabled ? "failing request" : "observe-only"); - } + auto const elapsed = std::chrono::duration_cast( + LlmRequest::getSteadyClockNow() - request->getKvCacheTransferStart()); + elapsedMs = static_cast(elapsed.count()); } - recordLocalTransferOutcome(requestId, request, /*failed=*/false, mCompletedSenderRequestIds, + recordLocalTransferOutcome(requestId, request, failed, mCompletedSenderRequestIds, mFailedSenderRequestIds, mSenderRequestsAwaitingConsensus); it = mSenderFutures.erase(it); + if (mNvbug6448152TraceEnabled) + { + traceConsensusTransition = true; + localTerminalTransitionCount++; + TLLM_LOG_INFO( + "NVBUG6448152_DIAG event=local_future_terminal side=ctx rank=%d request_id=%llu " + "outcome=%s elapsed_ms=%ld check_seq=%llu future_count=%zu awaiting_count=%zu " + "timeout_count=%zu", + mWorldRank, static_cast(requestId), failed ? "failed" : "completed", + elapsedMs, static_cast(contextCheckSequence), mSenderFutures.size(), + mSenderRequestsAwaitingConsensus.size(), mNvbug6448152SenderWaitTimeoutIds.size()); + } } else if (status == std::future_status::timeout) { - TLLM_LOG_DEBUG( - "Context KV cache transfer for request %ld is not ready after %ld ms wait slice; keeping it " - "in progress.", - requestId, static_cast(futureWaitInterval.count())); + if (!mNvbug6448152TraceEnabled) + { + TLLM_LOG_WARNING("Timed out waiting for context KV cache transfer after %d milliseconds.", + senderFutureTimeoutMs.value()); + } + else if (mNvbug6448152SenderWaitTimeoutIds.insert(requestId).second) + { + traceConsensusTransition = true; + firstWaitTimeoutTransitionCount++; + TLLM_LOG_WARNING( + "NVBUG6448152_DIAG event=first_future_wait_timeout side=ctx rank=%d request_id=%llu " + "wait_ms=%d check_seq=%llu future_count=%zu awaiting_count=%zu timeout_count=%zu", + mWorldRank, static_cast(requestId), senderFutureTimeoutMs.value(), + static_cast(contextCheckSequence), mSenderFutures.size(), + mSenderRequestsAwaitingConsensus.size(), mNvbug6448152SenderWaitTimeoutIds.size()); + } ++it; } else @@ -930,6 +901,18 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( recordLocalTransferOutcome(requestId, request, /*failed=*/true, mCompletedSenderRequestIds, mFailedSenderRequestIds, mSenderRequestsAwaitingConsensus); it = mSenderFutures.erase(it); + if (mNvbug6448152TraceEnabled) + { + traceConsensusTransition = true; + localTerminalTransitionCount++; + TLLM_LOG_INFO( + "NVBUG6448152_DIAG event=local_future_terminal side=ctx rank=%d request_id=%llu " + "outcome=unexpected_future_status check_seq=%llu future_count=%zu awaiting_count=%zu " + "timeout_count=%zu", + mWorldRank, static_cast(requestId), + static_cast(contextCheckSequence), mSenderFutures.size(), + mSenderRequestsAwaitingConsensus.size(), mNvbug6448152SenderWaitTimeoutIds.size()); + } } } catch (std::exception const& e) @@ -938,13 +921,17 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( recordLocalTransferOutcome(requestId, request, /*failed=*/true, mCompletedSenderRequestIds, mFailedSenderRequestIds, mSenderRequestsAwaitingConsensus); it = mSenderFutures.erase(it); - } - catch (...) - { - TLLM_LOG_ERROR("Unknown error occurred during context transfer for request %ld", requestId); - recordLocalTransferOutcome(requestId, request, /*failed=*/true, mCompletedSenderRequestIds, - mFailedSenderRequestIds, mSenderRequestsAwaitingConsensus); - it = mSenderFutures.erase(it); + if (mNvbug6448152TraceEnabled) + { + traceConsensusTransition = true; + localTerminalTransitionCount++; + TLLM_LOG_INFO( + "NVBUG6448152_DIAG event=local_future_terminal side=ctx rank=%d request_id=%llu " + "outcome=exception check_seq=%llu future_count=%zu awaiting_count=%zu timeout_count=%zu", + mWorldRank, static_cast(requestId), + static_cast(contextCheckSequence), mSenderFutures.size(), + mSenderRequestsAwaitingConsensus.size(), mNvbug6448152SenderWaitTimeoutIds.size()); + } } } else @@ -954,33 +941,118 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( } RequestStatuses requestsStatus{}; - auto const consensusOutcome = reduceTransferStates(syncComm, mGroupPipeParaComm, mCompletedSenderRequestIds, - mFailedSenderRequestIds, inflightCancelEnabled ? mTimedOutSenderIds : std::unordered_set{}); - if (inflightCancelEnabled) - { - for (auto const requestId : consensusOutcome.timedOutRequestIds) + TransferConsensusOutcome consensusOutcome; + std::uint64_t tpConsensusSequence = 0; + std::uint64_t ppConsensusSequence = 0; + long long tpConsensusDurationUs = 0; + long long ppConsensusDurationUs = 0; + if (mNvbug6448152TraceEnabled) + { + tpConsensusSequence = ++mNvbug6448152ContextTpConsensusSequence; + ppConsensusSequence = ++mNvbug6448152ContextPpConsensusSequence; + + int const tpRank = syncComm != nullptr ? syncComm->getRank() : 0; + int const tpSize = syncComm != nullptr ? syncComm->getSize() : 1; + if (traceConsensusTransition) { - auto const futureIt = std::find_if(mSenderFutures.begin(), mSenderFutures.end(), - [requestId](auto const& entry) { return entry.first->mRequestId == requestId; }); - if (futureIt == mSenderFutures.end() - || futureIt->second.wait_for(std::chrono::milliseconds(0)) == std::future_status::ready - || mCancelRequestedSenderIds.find(requestId) != mCancelRequestedSenderIds.end()) - { - continue; - } - mTimedOutSenderIds.insert(requestId); - if (requestCancellationNoThrow( - requestId, "Context", [&]() { return mCacheSender->cancelRequest(*futureIt->first); })) - { - mCancelRequestedSenderIds.insert(requestId); - } - else - { - TLLM_LOG_DEBUG("Context cancellation for request %ld was not accepted; will retry", requestId); - } + TLLM_LOG_INFO( + "NVBUG6448152_DIAG event=consensus_enter side=ctx stage=tp rank=%d comm_rank=%d comm_size=%d " + "check_seq=%llu consensus_seq=%llu local_terminal_transitions=%zu " + "first_wait_timeout_transitions=%zu enqueued_transitions=%llu periodic_checkpoint=%d " + "completed_count=%zu failed_count=%zu future_count=%zu awaiting_count=%zu timeout_count=%zu", + mWorldRank, tpRank, tpSize, static_cast(contextCheckSequence), + static_cast(tpConsensusSequence), localTerminalTransitionCount, + firstWaitTimeoutTransitionCount, static_cast(enqueuedTransitionCount), + static_cast(periodicConsensusCheckpoint), mCompletedSenderRequestIds.size(), + mFailedSenderRequestIds.size(), mSenderFutures.size(), mSenderRequestsAwaitingConsensus.size(), + mNvbug6448152SenderWaitTimeoutIds.size()); + } + auto const tpConsensusStart = std::chrono::steady_clock::now(); + auto const tpOutcome = reduceTransferStates(syncComm, mCompletedSenderRequestIds, mFailedSenderRequestIds); + tpConsensusDurationUs + = std::chrono::duration_cast(std::chrono::steady_clock::now() - tpConsensusStart) + .count(); + if (traceConsensusTransition) + { + TLLM_LOG_INFO( + "NVBUG6448152_DIAG event=consensus_exit side=ctx stage=tp rank=%d comm_rank=%d comm_size=%d " + "check_seq=%llu consensus_seq=%llu duration_us=%lld local_terminal_transitions=%zu " + "first_wait_timeout_transitions=%zu enqueued_transitions=%llu completed_count=%zu failed_count=%zu " + "future_count=%zu awaiting_count=%zu timeout_count=%zu", + mWorldRank, tpRank, tpSize, static_cast(contextCheckSequence), + static_cast(tpConsensusSequence), tpConsensusDurationUs, + localTerminalTransitionCount, firstWaitTimeoutTransitionCount, + static_cast(enqueuedTransitionCount), tpOutcome.completedRequestIds.size(), + tpOutcome.failedRequestIds.size(), mSenderFutures.size(), mSenderRequestsAwaitingConsensus.size(), + mNvbug6448152SenderWaitTimeoutIds.size()); + } + else if (tpConsensusDurationUs >= kNvbug6448152SlowConsensusThresholdUs) + { + TLLM_LOG_WARNING( + "NVBUG6448152_DIAG event=consensus_slow side=ctx stage=tp rank=%d comm_rank=%d comm_size=%d " + "check_seq=%llu consensus_seq=%llu duration_us=%lld completed_count=%zu failed_count=%zu " + "future_count=%zu awaiting_count=%zu timeout_count=%zu", + mWorldRank, tpRank, tpSize, static_cast(contextCheckSequence), + static_cast(tpConsensusSequence), tpConsensusDurationUs, + tpOutcome.completedRequestIds.size(), tpOutcome.failedRequestIds.size(), mSenderFutures.size(), + mSenderRequestsAwaitingConsensus.size(), mNvbug6448152SenderWaitTimeoutIds.size()); + } + + int const ppRank = mGroupPipeParaComm != nullptr ? mGroupPipeParaComm->getRank() : 0; + int const ppSize = mGroupPipeParaComm != nullptr ? mGroupPipeParaComm->getSize() : 1; + if (traceConsensusTransition) + { + TLLM_LOG_INFO( + "NVBUG6448152_DIAG event=consensus_enter side=ctx stage=pp rank=%d comm_rank=%d comm_size=%d " + "check_seq=%llu consensus_seq=%llu local_terminal_transitions=%zu " + "first_wait_timeout_transitions=%zu enqueued_transitions=%llu periodic_checkpoint=%d " + "completed_count=%zu failed_count=%zu future_count=%zu awaiting_count=%zu timeout_count=%zu", + mWorldRank, ppRank, ppSize, static_cast(contextCheckSequence), + static_cast(ppConsensusSequence), localTerminalTransitionCount, + firstWaitTimeoutTransitionCount, static_cast(enqueuedTransitionCount), + static_cast(periodicConsensusCheckpoint), tpOutcome.completedRequestIds.size(), + tpOutcome.failedRequestIds.size(), mSenderFutures.size(), mSenderRequestsAwaitingConsensus.size(), + mNvbug6448152SenderWaitTimeoutIds.size()); + } + auto const ppConsensusStart = std::chrono::steady_clock::now(); + consensusOutcome + = reduceTransferStates(mGroupPipeParaComm, tpOutcome.completedRequestIds, tpOutcome.failedRequestIds); + ppConsensusDurationUs + = std::chrono::duration_cast(std::chrono::steady_clock::now() - ppConsensusStart) + .count(); + if (traceConsensusTransition) + { + TLLM_LOG_INFO( + "NVBUG6448152_DIAG event=consensus_exit side=ctx stage=pp rank=%d comm_rank=%d comm_size=%d " + "check_seq=%llu consensus_seq=%llu duration_us=%lld local_terminal_transitions=%zu " + "first_wait_timeout_transitions=%zu enqueued_transitions=%llu completed_count=%zu failed_count=%zu " + "future_count=%zu awaiting_count=%zu timeout_count=%zu", + mWorldRank, ppRank, ppSize, static_cast(contextCheckSequence), + static_cast(ppConsensusSequence), ppConsensusDurationUs, + localTerminalTransitionCount, firstWaitTimeoutTransitionCount, + static_cast(enqueuedTransitionCount), consensusOutcome.completedRequestIds.size(), + consensusOutcome.failedRequestIds.size(), mSenderFutures.size(), + mSenderRequestsAwaitingConsensus.size(), mNvbug6448152SenderWaitTimeoutIds.size()); + } + else if (ppConsensusDurationUs >= kNvbug6448152SlowConsensusThresholdUs) + { + TLLM_LOG_WARNING( + "NVBUG6448152_DIAG event=consensus_slow side=ctx stage=pp rank=%d comm_rank=%d comm_size=%d " + "check_seq=%llu consensus_seq=%llu duration_us=%lld completed_count=%zu failed_count=%zu " + "future_count=%zu awaiting_count=%zu timeout_count=%zu", + mWorldRank, ppRank, ppSize, static_cast(contextCheckSequence), + static_cast(ppConsensusSequence), ppConsensusDurationUs, + consensusOutcome.completedRequestIds.size(), consensusOutcome.failedRequestIds.size(), + mSenderFutures.size(), mSenderRequestsAwaitingConsensus.size(), + mNvbug6448152SenderWaitTimeoutIds.size()); } } - for (auto const requestId : sortedRequestIds(consensusOutcome.failedRequestIds)) + else + { + consensusOutcome + = reduceTransferStates(syncComm, mGroupPipeParaComm, mCompletedSenderRequestIds, mFailedSenderRequestIds); + } + for (auto const requestId : consensusOutcome.failedRequestIds) { auto const requestIt = mSenderRequestsAwaitingConsensus.find(requestId); if (requestIt == mSenderRequestsAwaitingConsensus.end()) @@ -990,11 +1062,24 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( requestIt->second->setState(LlmRequestState::kDISAGG_TRANS_ERROR); requestsStatus.errorRequestIds.insert(requestId); mTimedOutSenderIds.erase(requestId); - mCancelRequestedSenderIds.erase(requestId); eraseLocalTransferOutcome( requestId, mCompletedSenderRequestIds, mFailedSenderRequestIds, mSenderRequestsAwaitingConsensus); + if (mNvbug6448152TraceEnabled) + { + bool const hadWaitTimeout = mNvbug6448152SenderWaitTimeoutIds.erase(requestId) > 0; + TLLM_LOG_INFO( + "NVBUG6448152_DIAG event=global_commit side=ctx rank=%d request_id=%llu outcome=failed " + "check_seq=%llu tp_consensus_seq=%llu pp_consensus_seq=%llu tp_duration_us=%lld " + "pp_duration_us=%lld had_wait_timeout=%d future_count=%zu awaiting_count=%zu timeout_count=%zu", + mWorldRank, static_cast(requestId), + static_cast(contextCheckSequence), + static_cast(tpConsensusSequence), + static_cast(ppConsensusSequence), tpConsensusDurationUs, ppConsensusDurationUs, + static_cast(hadWaitTimeout), mSenderFutures.size(), mSenderRequestsAwaitingConsensus.size(), + mNvbug6448152SenderWaitTimeoutIds.size()); + } } - for (auto const requestId : sortedRequestIds(consensusOutcome.completedRequestIds)) + for (auto const requestId : consensusOutcome.completedRequestIds) { auto const requestIt = mSenderRequestsAwaitingConsensus.find(requestId); if (requestIt == mSenderRequestsAwaitingConsensus.end()) @@ -1007,9 +1092,22 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( requestIt->second->setState(LlmRequestState::kDISAGG_CONTEXT_COMPLETE); } mTimedOutSenderIds.erase(requestId); - mCancelRequestedSenderIds.erase(requestId); eraseLocalTransferOutcome( requestId, mCompletedSenderRequestIds, mFailedSenderRequestIds, mSenderRequestsAwaitingConsensus); + if (mNvbug6448152TraceEnabled) + { + bool const hadWaitTimeout = mNvbug6448152SenderWaitTimeoutIds.erase(requestId) > 0; + TLLM_LOG_INFO( + "NVBUG6448152_DIAG event=global_commit side=ctx rank=%d request_id=%llu outcome=completed " + "mark_complete=%d check_seq=%llu tp_consensus_seq=%llu pp_consensus_seq=%llu tp_duration_us=%lld " + "pp_duration_us=%lld had_wait_timeout=%d future_count=%zu awaiting_count=%zu timeout_count=%zu", + mWorldRank, static_cast(requestId), static_cast(markComplete), + static_cast(contextCheckSequence), + static_cast(tpConsensusSequence), + static_cast(ppConsensusSequence), tpConsensusDurationUs, ppConsensusDurationUs, + static_cast(hadWaitTimeout), mSenderFutures.size(), mSenderRequestsAwaitingConsensus.size(), + mNvbug6448152SenderWaitTimeoutIds.size()); + } } return requestsStatus; @@ -1017,41 +1115,19 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( void CacheTransceiver::checkGenTransferStatus(std::optional const& atLeastRequestNum) { - bool const blockAll = !atLeastRequestNum.has_value(); - bool const inflightCancelEnabled = common::getEnvDisaggEnableInflightCancel(); - TLLM_CHECK_WITH_INFO(!inflightCancelEnabled || !blockAll, - "In-flight cancellation requires a finite generation-transfer status poll; pass 0 for a nonblocking poll."); - bool const needsProgress = atLeastRequestNum.value_or(0) > 0; - std::optional genTransferPollIntervalMs = std::nullopt; - if (mCacheTransceiverConfig.has_value()) + std::uint64_t generationCheckSequence = 0; + if (mNvbug6448152TraceEnabled) { - genTransferPollIntervalMs = mCacheTransceiverConfig->getKvTransferPollIntervalMs(); + generationCheckSequence = ++mNvbug6448152GenerationCheckSequence; } - auto const futureWaitInterval = getTransferFutureWaitInterval(genTransferPollIntervalMs, needsProgress); + bool blockAll = !atLeastRequestNum.has_value(); std::vector genTransferReadyRequestIds; - auto collectReadyRequestIds = [&]() - { - genTransferReadyRequestIds.clear(); - for (auto&& [request, future] : mRequesterFutures) - { - if (future.wait_for(std::chrono::milliseconds(0)) == std::future_status::ready) - { - genTransferReadyRequestIds.push_back(request->mRequestId); - } - } - }; - collectReadyRequestIds(); - if (needsProgress) + for (auto&& [request, future] : mRequesterFutures) { - auto const deadline = std::chrono::steady_clock::now() + futureWaitInterval; - while (static_cast(genTransferReadyRequestIds.size()) < atLeastRequestNum.value() - && std::chrono::steady_clock::now() < deadline) + if (future.wait_for(std::chrono::milliseconds(0)) == std::future_status::ready) { - auto const remaining - = std::chrono::duration_cast(deadline - std::chrono::steady_clock::now()); - std::this_thread::sleep_for(std::min(std::chrono::milliseconds(kTransferFuturePollIntervalMs), remaining)); - collectReadyRequestIds(); + genTransferReadyRequestIds.push_back(request->mRequestId); } } std::unordered_map frequencyMap; @@ -1080,6 +1156,53 @@ void CacheTransceiver::checkGenTransferStatus(std::optional const& atLeastR [](std::pair const& left, std::pair const& right) { return left.second > right.second; }); std::unordered_set toCompleteIdSet; + size_t idx = 0; + while (atLeastRequestNum.value_or(0) > static_cast(toCompleteIdSet.size())) + { + if (idx >= freqVec.size()) + { + break; + } + toCompleteIdSet.insert(freqVec.at(idx).first); + if (useMPI()) + { + TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), + " checkGenTransferStatus at least from freqVec requestId: %zu ", freqVec.at(idx).first); + } + else + { + TLLM_LOG_DEBUG(tensorrt_llm::pg_utils::get_world_pg()->getRank(), + " checkGenTransferStatus at least from freqVec requestId: %zu ", freqVec.at(idx).first); + } + idx++; + } + idx = 0; + + // insert order + while (atLeastRequestNum.value_or(0) > static_cast(toCompleteIdSet.size())) + { + if (idx >= mRequesterFutures.size()) + { + break; + } + if (toCompleteIdSet.find(mRequesterFutures.at(idx).first->mRequestId) == toCompleteIdSet.end()) + { + toCompleteIdSet.insert(mRequesterFutures.at(idx).first->mRequestId); + if (useMPI()) + { + TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), + " checkGenTransferStatus at least from RequesterFuture requestId: %zu atLeastRequestNum:%d", + mRequesterFutures.at(idx).first->mRequestId, atLeastRequestNum.value_or(0)); + } + else + { + TLLM_LOG_DEBUG(tensorrt_llm::pg_utils::get_world_pg()->getRank(), + " checkGenTransferStatus at least from RequesterFuture requestId: %zu atLeastRequestNum:%d", + mRequesterFutures.at(idx).first->mRequestId, atLeastRequestNum.value_or(0)); + } + } + idx++; + } for (auto&& [requestId, freq] : freqVec) { if (freq == ((syncComm != nullptr) ? syncComm->getSize() : 1)) @@ -1109,8 +1232,7 @@ void CacheTransceiver::checkGenTransferStatus(std::optional const& atLeastR " checkGenTransferStatus toCompleteIdSet size: %zu, atLeastRequestNum: %d ", toCompleteIdSet.size(), atLeastRequestNum.value_or(0)); } - - // Gen-side mirror of the context deadline/consensus path. + // Observe-only: gen-side mirror of the context-side timeout WARN. std::optional kvTransferTimeoutMs = std::nullopt; if (mCacheTransceiverConfig.has_value()) { @@ -1120,60 +1242,35 @@ void CacheTransceiver::checkGenTransferStatus(std::optional const& atLeastR { auto& request = it->first; auto const requestId = request->mRequestId; - if (kvTransferTimeoutMs.has_value() - && it->second.wait_for(std::chrono::milliseconds(0)) != std::future_status::ready) + if (kvTransferTimeoutMs.has_value()) { - auto const elapsedMs = getTransferElapsedMs(request, LlmRequest::getSteadyClockNow()); - if (elapsedMs > kvTransferTimeoutMs.value()) + auto elapsed = std::chrono::duration_cast( + LlmRequest::getSteadyClockNow() - request->getKvCacheTransferStart()); + auto elapsedMs = static_cast(elapsed.count()); + if (elapsedMs > kvTransferTimeoutMs.value() && mTimedOutRequesterIds.insert(request->mRequestId).second) { - if (mTimedOutRequesterIds.insert(requestId).second) - { - TLLM_LOG_WARNING( - "Generation KV cache transfer for request %ld exceeded configured timeout: " - "elapsed %ld ms > limit %d ms (%s).", - requestId, elapsedMs, kvTransferTimeoutMs.value(), - inflightCancelEnabled ? "requesting cancellation" : "observe-only"); - } + TLLM_LOG_WARNING( + "Generation KV cache transfer for request %ld exceeded configured timeout: " + "elapsed %ld ms > limit %d ms (observe-only).", + request->mRequestId, elapsedMs, kvTransferTimeoutMs.value()); } } if (blockAll || toCompleteIdSet.find(requestId) != toCompleteIdSet.end()) { + bool localFailed = true; try { - auto const status = blockAll ? std::future_status::ready : it->second.wait_for(futureWaitInterval); - if (status == std::future_status::ready) + it->second.get(); + bool const failed = request->getState() == LlmRequestState::kDISAGG_TRANS_ERROR; + localFailed = failed; + if (failed) { - it->second.get(); - if (kvTransferTimeoutMs.has_value()) - { - auto const elapsedMs = getTransferElapsedMs(request, request->getKvCacheTransferEnd()); - if (elapsedMs > kvTransferTimeoutMs.value() && mTimedOutRequesterIds.insert(requestId).second) - { - TLLM_LOG_WARNING( - "Generation KV cache transfer for request %ld completed after its deadline: " - "elapsed %ld ms > limit %d ms (%s).", - requestId, elapsedMs, kvTransferTimeoutMs.value(), - inflightCancelEnabled ? "failing request" : "observe-only"); - } - } - recordLocalTransferOutcome(requestId, request, /*failed=*/false, mCompletedRequesterRequestIds, - mFailedRequesterRequestIds, mRequesterRequestsAwaitingConsensus); - } - else if (status == std::future_status::timeout) - { - TLLM_LOG_DEBUG( - "Generation KV cache transfer for request %ld is not ready after %ld ms wait slice; keeping " - "it in progress.", - requestId, static_cast(futureWaitInterval.count())); - ++it; - continue; - } - else - { - TLLM_LOG_ERROR("Future returned unexpected status for request %ld. Marking as error.", requestId); - recordLocalTransferOutcome(requestId, request, /*failed=*/true, mCompletedRequesterRequestIds, - mFailedRequesterRequestIds, mRequesterRequestsAwaitingConsensus); + // The receiver uses the error state as a local transfer-failed signal. + // Keep that signal local until the consensus outcome commits it globally. + request->setState(LlmRequestState::kDISAGG_GENERATION_TRANS_IN_PROGRESS); } + recordLocalTransferOutcome(requestId, request, failed, mCompletedRequesterRequestIds, + mFailedRequesterRequestIds, mRequesterRequestsAwaitingConsensus); } catch (std::exception const& e) { @@ -1181,12 +1278,6 @@ void CacheTransceiver::checkGenTransferStatus(std::optional const& atLeastR recordLocalTransferOutcome(requestId, request, /*failed=*/true, mCompletedRequesterRequestIds, mFailedRequesterRequestIds, mRequesterRequestsAwaitingConsensus); } - catch (...) - { - TLLM_LOG_ERROR("Unknown error occurred during generation transfer for request %ld", requestId); - recordLocalTransferOutcome(requestId, request, /*failed=*/true, mCompletedRequesterRequestIds, - mFailedRequesterRequestIds, mRequesterRequestsAwaitingConsensus); - } if (useMPI()) { TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), @@ -1200,6 +1291,15 @@ void CacheTransceiver::checkGenTransferStatus(std::optional const& atLeastR request->getContextPhaseParams().value().getReqId()); } it = mRequesterFutures.erase(it); + if (mNvbug6448152TraceEnabled) + { + TLLM_LOG_INFO( + "NVBUG6448152_DIAG event=local_future_terminal side=gen rank=%d request_id=%llu outcome=%s " + "check_seq=%llu future_count=%zu awaiting_count=%zu timeout_count=%zu", + mWorldRank, static_cast(requestId), localFailed ? "failed" : "completed", + static_cast(generationCheckSequence), mRequesterFutures.size(), + mRequesterRequestsAwaitingConsensus.size(), mTimedOutRequesterIds.size()); + } } else { @@ -1208,33 +1308,8 @@ void CacheTransceiver::checkGenTransferStatus(std::optional const& atLeastR } auto const consensusOutcome - = reduceTransferStates(syncComm, mCompletedRequesterRequestIds, mFailedRequesterRequestIds, - inflightCancelEnabled ? mTimedOutRequesterIds : std::unordered_set{}); - if (inflightCancelEnabled) - { - for (auto const requestId : consensusOutcome.timedOutRequestIds) - { - auto const futureIt = std::find_if(mRequesterFutures.begin(), mRequesterFutures.end(), - [requestId](auto const& entry) { return entry.first->mRequestId == requestId; }); - if (futureIt == mRequesterFutures.end() - || futureIt->second.wait_for(std::chrono::milliseconds(0)) == std::future_status::ready - || mCancelRequestedRequesterIds.find(requestId) != mCancelRequestedRequesterIds.end()) - { - continue; - } - mTimedOutRequesterIds.insert(requestId); - if (requestCancellationNoThrow( - requestId, "Generation", [&]() { return mCacheReceiver->cancelRequest(*futureIt->first); })) - { - mCancelRequestedRequesterIds.insert(requestId); - } - else - { - TLLM_LOG_DEBUG("Generation cancellation for request %ld was not accepted; will retry", requestId); - } - } - } - for (auto const requestId : sortedRequestIds(consensusOutcome.failedRequestIds)) + = reduceTransferStates(syncComm, mCompletedRequesterRequestIds, mFailedRequesterRequestIds); + for (auto const requestId : consensusOutcome.failedRequestIds) { auto const requestIt = mRequesterRequestsAwaitingConsensus.find(requestId); if (requestIt == mRequesterRequestsAwaitingConsensus.end()) @@ -1243,11 +1318,19 @@ void CacheTransceiver::checkGenTransferStatus(std::optional const& atLeastR } requestIt->second->setState(LlmRequestState::kDISAGG_TRANS_ERROR); mTimedOutRequesterIds.erase(requestId); - mCancelRequestedRequesterIds.erase(requestId); eraseLocalTransferOutcome( requestId, mCompletedRequesterRequestIds, mFailedRequesterRequestIds, mRequesterRequestsAwaitingConsensus); + if (mNvbug6448152TraceEnabled) + { + TLLM_LOG_INFO( + "NVBUG6448152_DIAG event=global_commit side=gen rank=%d request_id=%llu outcome=failed " + "check_seq=%llu future_count=%zu awaiting_count=%zu timeout_count=%zu", + mWorldRank, static_cast(requestId), + static_cast(generationCheckSequence), mRequesterFutures.size(), + mRequesterRequestsAwaitingConsensus.size(), mTimedOutRequesterIds.size()); + } } - for (auto const requestId : sortedRequestIds(consensusOutcome.completedRequestIds)) + for (auto const requestId : consensusOutcome.completedRequestIds) { auto const requestIt = mRequesterRequestsAwaitingConsensus.find(requestId); if (requestIt == mRequesterRequestsAwaitingConsensus.end()) @@ -1262,9 +1345,17 @@ void CacheTransceiver::checkGenTransferStatus(std::optional const& atLeastR updateKVCacheTransferBW(syncComm, requestIt->second.get()); } mTimedOutRequesterIds.erase(requestId); - mCancelRequestedRequesterIds.erase(requestId); eraseLocalTransferOutcome( requestId, mCompletedRequesterRequestIds, mFailedRequesterRequestIds, mRequesterRequestsAwaitingConsensus); + if (mNvbug6448152TraceEnabled) + { + TLLM_LOG_INFO( + "NVBUG6448152_DIAG event=global_commit side=gen rank=%d request_id=%llu outcome=completed " + "check_seq=%llu future_count=%zu awaiting_count=%zu timeout_count=%zu", + mWorldRank, static_cast(requestId), + static_cast(generationCheckSequence), mRequesterFutures.size(), + mRequesterRequestsAwaitingConsensus.size(), mTimedOutRequesterIds.size()); + } } } @@ -1273,19 +1364,8 @@ bool CacheTransceiver::checkGenTransferComplete() const return mRequesterFutures.empty() && mCompletedRequesterRequestIds.empty() && mFailedRequesterRequestIds.empty(); } -bool CacheTransceiver::hasPoisonedTransferBuffer() const -{ - return std::any_of(mCacheTransBufferManagerPtrs.begin(), mCacheTransBufferManagerPtrs.end(), - [](BaseTransBufferManager const* manager) { return manager != nullptr && manager->hasPoisonedBuffer(); }); -} - bool CacheTransceiver::cancelRequest(std::shared_ptr llmRequest) { - if (llmRequest == nullptr) - { - TLLM_LOG_WARNING("Cannot cancel a null KV cache transfer request"); - return false; - } if (llmRequest->isContextOnlyRequest()) { return mCacheSender->cancelRequest(*llmRequest); diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransferLayer.cpp b/cpp/tensorrt_llm/batch_manager/cacheTransferLayer.cpp index d0a54dbb7d3c..c013fd75c6e4 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheTransferLayer.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheTransferLayer.cpp @@ -54,7 +54,8 @@ void CacheTransferLayer::validateSupport(executor::DataTransceiverState const& p if (mRnnFormatter && selfHasRnn) { - // Unified pool path (CppMambaHybridCacheManager) uses RnnCacheFormatter. + // Both slot-based (CppMambaCacheManager) and unified pool (CppMambaHybridCacheManager) + // paths now use RnnCacheFormatter. if (peerHasRnn) { TLLM_CHECK_WITH_INFO(mRnnFormatter->inquireSupport(mCacheState, peerState.getCacheState().value()), diff --git a/cpp/tensorrt_llm/batch_manager/capacityScheduler.cpp b/cpp/tensorrt_llm/batch_manager/capacityScheduler.cpp index 335e30403d0f..60b6cc14ebee 100644 --- a/cpp/tensorrt_llm/batch_manager/capacityScheduler.cpp +++ b/cpp/tensorrt_llm/batch_manager/capacityScheduler.cpp @@ -20,7 +20,6 @@ #include "tensorrt_llm/batch_manager/kvCacheManager.h" #include "tensorrt_llm/batch_manager/peftCacheManager.h" #include "tensorrt_llm/batch_manager/scheduledBlocksManager.h" -#include "tensorrt_llm/common/assert.h" #include "tensorrt_llm/common/logger.h" #include "tensorrt_llm/common/nvtxUtils.h" @@ -115,32 +114,6 @@ bool beneficialToSkip(std::optional const& return false; } -template -void checkRequiredCrossKvCacheManager( - LlmRequestState noScheduleUntilState, OptionalRef crossKvCacheManager) -{ - if (noScheduleUntilState != LlmRequestState::kENCODER_INIT) - { - return; - } - - TLLM_CHECK_WITH_INFO( - static_cast(crossKvCacheManager), "Encoder-decoder scheduling requires a cross_kv_cache_manager."); -} - -void claimPeftPagesForRequest(std::shared_ptr const& req, - OptionalRef peftCacheManager, SizeType32& claimedPeftPages, - std::unordered_set& uniqTaskIds) -{ - bool const reqHasLora = req->getLoraTaskId().has_value(); - bool const isNewTask = reqHasLora && !uniqTaskIds.count(req->getLoraTaskId().value()); - if (isNewTask) - { - claimedPeftPages += peftCacheManager ? peftCacheManager->determineNumPages(req) : 0; - uniqTaskIds.insert(req->getLoraTaskId().value()); - } -} - } // namespace MaxRequestsScheduler::MaxRequestsScheduler( @@ -151,26 +124,23 @@ MaxRequestsScheduler::MaxRequestsScheduler( } MaxUtilizationScheduler::MaxUtilizationScheduler(SizeType32 maxNumRequests, bool twoStepsLookAhead, - LlmRequestState noScheduleUntilState, LlmRequestState noScheduleAfterState, bool enablePrefixAwareScheduling) + LlmRequestState noScheduleUntilState, LlmRequestState noScheduleAfterState) : BaseCapacityScheduler(noScheduleUntilState, noScheduleAfterState) , mMaxNumRequests(maxNumRequests) , mTwoStepsLookAhead{twoStepsLookAhead} - , mEnablePrefixAwareScheduling{enablePrefixAwareScheduling} { } -GuaranteedNoEvictScheduler::GuaranteedNoEvictScheduler(SizeType32 maxNumRequests, LlmRequestState noScheduleUntilState, - LlmRequestState noScheduleAfterState, bool enablePrefixAwareScheduling) +GuaranteedNoEvictScheduler::GuaranteedNoEvictScheduler( + SizeType32 maxNumRequests, LlmRequestState noScheduleUntilState, LlmRequestState noScheduleAfterState) : BaseCapacityScheduler(noScheduleUntilState, noScheduleAfterState) , mMaxNumRequests(maxNumRequests) - , mEnablePrefixAwareScheduling{enablePrefixAwareScheduling} { } -StaticBatchScheduler::StaticBatchScheduler(SizeType32 maxNumRequests, LlmRequestState noScheduleUntilState, - LlmRequestState noScheduleAfterState, bool enablePrefixAwareScheduling) - : GuaranteedNoEvictScheduler( - maxNumRequests, noScheduleUntilState, noScheduleAfterState, enablePrefixAwareScheduling) +StaticBatchScheduler::StaticBatchScheduler( + SizeType32 maxNumRequests, LlmRequestState noScheduleUntilState, LlmRequestState noScheduleAfterState) + : GuaranteedNoEvictScheduler(maxNumRequests, noScheduleUntilState, noScheduleAfterState) { } @@ -222,14 +192,12 @@ std::tuple GuaranteedNoEvictScheduler::impl( { RequestVector scheduledRequests; - checkRequiredCrossKvCacheManager(getNoScheduleUntilState(), crossKvCacheManager); - // Now check if we can add pending requests auto const maxPeftCachePages = peftCacheManager ? peftCacheManager->getMaxDevicePages() : std::numeric_limits::max(); // The optimization of delaying requests won't work for variable window attention - bool skippingIsRelevant = mEnablePrefixAwareScheduling && (!kvCacheManager.getBlockManager().isVariableWindow()) + bool skippingIsRelevant = (!kvCacheManager.getBlockManager().isVariableWindow()) && (!crossKvCacheManager || !crossKvCacheManager->getBlockManager().isVariableWindow()); // Keep track of blocks contributed by requests in context phase @@ -253,7 +221,6 @@ std::tuple GuaranteedNoEvictScheduler::impl( : std::nullopt; SizeType32 claimedPeftPages{0}; std::unordered_set uniqTaskIds{}; - std::size_t numAdmittedRequests{0}; RequestVector pendingRequests; RequestVector pendingDisGenInitRequests; pendingRequests.reserve(activeRequests.size()); @@ -263,24 +230,20 @@ std::tuple GuaranteedNoEvictScheduler::impl( // if request cannot be scheduled yet or request should no longer be scheduled, skip if ( // Allow disagg_generation_init requests to be scheduled, so that we'll allocate their KV cache - !req->isDisaggGenerationInitState() && !req->isDisaggGenerationTransmissionInProgress() + !req->isDisaggGenerationInitState() && (!req->hasReachedState(getNoScheduleUntilState()) || req->hasReachedState(getNoScheduleAfterState()))) { continue; } - if (numAdmittedRequests >= static_cast(mMaxNumRequests)) + if (scheduledRequests.size() >= static_cast(mMaxNumRequests)) { break; } - if (req->isDisaggGenerationTransmissionInProgress() || req->isGenerationInProgressState()) + if (req->isGenerationInProgressState()) { - ++numAdmittedRequests; - if (req->isGenerationInProgressState()) - { - scheduledRequests.emplace_back(req); - } + scheduledRequests.emplace_back(req); reservedBlocks.enoughAvailableBlocks(*req); reservedBlocks.commitBlocks(); if (reservedCrossBlocks) @@ -288,7 +251,13 @@ std::tuple GuaranteedNoEvictScheduler::impl( reservedCrossBlocks->enoughAvailableBlocks(*req); reservedCrossBlocks->commitBlocks(); } - claimPeftPagesForRequest(req, peftCacheManager, claimedPeftPages, uniqTaskIds); + bool const reqHasLora = req->getLoraTaskId().has_value(); + bool const isNewTask = reqHasLora && !uniqTaskIds.count(req->getLoraTaskId().value()); + if (isNewTask) + { + claimedPeftPages += peftCacheManager ? peftCacheManager->determineNumPages(req) : 0; + uniqTaskIds.insert(req->getLoraTaskId().value()); + } } else if (req->isDisaggGenerationInitState()) { @@ -302,7 +271,7 @@ std::tuple GuaranteedNoEvictScheduler::impl( // If StaticBatchScheduling == true check if we can add pending requests only when no requests are active. // Otherwise, add just check that we can add pending requests. - if (!StaticBatchScheduling || numAdmittedRequests == 0) + if (!StaticBatchScheduling || scheduledRequests.size() == 0) { auto availablePeftPages = maxPeftCachePages - claimedPeftPages; @@ -318,79 +287,39 @@ std::tuple GuaranteedNoEvictScheduler::impl( // eliminating 2 redundant walks per request. bool const isFirstChunkContext = req->isContextInitState() && req->isFirstContextChunk() && !req->isDisaggGenerationInitState(); - // Encoder-init requests do not consume self- or cross-KV - // blocks. We still keep the cross reuse summary available for - // beneficial-to-skip so duplicate encoder inputs can be ordered - // consistently before their decoder-context admission budgets - // the cross pool. - bool const isEncoderInit = req->isEncoderInitState(); std::optional summary; std::optional crossSummary; - if (mEnablePrefixAwareScheduling) + if (isFirstChunkContext) { - if (isFirstChunkContext) + // analyzePrefixReuse asserts on variable-window managers; skip the walk there + // and let downstream callers fall back to their fresh tree-walk path. + if (kvCacheManager.isEnableBlockReuse() && !kvCacheManager.getBlockManager().isVariableWindow()) { - // analyzePrefixReuse asserts on variable-window managers; skip the walk there - // and let downstream callers fall back to their fresh tree-walk path. - if (kvCacheManager.isEnableBlockReuse() && !kvCacheManager.getBlockManager().isVariableWindow()) - { - auto uniqueTokens = req->getUniqueTokens(0); - summary = kvCacheManager.analyzePrefixReuse(uniqueTokens, *req); - } - if (crossKvCacheManager && crossKvCacheManager->isEnableBlockReuse() - && !crossKvCacheManager->getBlockManager().isVariableWindow()) - { - auto uniqueTokens = *(req->getEncoderUniqueTokens().value()); - crossSummary = crossKvCacheManager->analyzePrefixReuse(uniqueTokens, *req); - } + auto uniqueTokens = req->getUniqueTokens(0); + summary = kvCacheManager.analyzePrefixReuse(uniqueTokens, *req); } - else if (isEncoderInit && crossKvCacheManager && crossKvCacheManager->isEnableBlockReuse() + if (crossKvCacheManager && crossKvCacheManager->isEnableBlockReuse() && !crossKvCacheManager->getBlockManager().isVariableWindow()) { - // Encoder admission only needs the cross summary for reuse ordering. auto uniqueTokens = *(req->getEncoderUniqueTokens().value()); crossSummary = crossKvCacheManager->analyzePrefixReuse(uniqueTokens, *req); } } - else if (isFirstChunkContext) - { - summary = kv_cache_manager::PrefixReuseSummary{}; - if (crossKvCacheManager) - { - crossSummary = kv_cache_manager::PrefixReuseSummary{}; - } - } + // Beneficial-to-skip check using the cached summary - if (!StaticBatchScheduling && skippingIsRelevant && (isFirstChunkContext || isEncoderInit) + if (!StaticBatchScheduling && skippingIsRelevant && isFirstChunkContext && beneficialToSkip( summary, crossSummary, newlyContributedContextBlocks, newlyContributedCrossContextBlocks)) { continue; } - if (numAdmittedRequests >= static_cast(mMaxNumRequests)) + if (scheduledRequests.size() >= static_cast(mMaxNumRequests)) { break; } - if (isEncoderInit) - { - bool reqHasLora = req->getLoraTaskId().has_value(); - bool isNewTask = reqHasLora && !uniqTaskIds.count(req->getLoraTaskId().value()); - auto neededPeftPages = isNewTask && peftCacheManager ? peftCacheManager->determineNumPages(req) : 0; - - if (neededPeftPages <= availablePeftPages) - { - scheduledRequests.emplace_back(req); - ++numAdmittedRequests; - availablePeftPages -= neededPeftPages; - if (isNewTask) - { - uniqTaskIds.insert(req->getLoraTaskId().value()); - } - } - } - else if (req->isContextInitState() || req->isDisaggGenerationInitState()) + if (req->isContextInitState() || req->isDisaggGenerationInitState()) { // Check block availability using the cached summary when available. // enoughAvailableBlocks is check-only (no decrement) — safe if cross check fails. @@ -407,7 +336,6 @@ std::tuple GuaranteedNoEvictScheduler::impl( if (enoughBlocks && enoughCrossBlocks && neededPeftPages <= availablePeftPages) { scheduledRequests.emplace_back(req); - ++numAdmittedRequests; // Decrement using the cached values computed by enoughAvailableBlocks. reservedBlocks.commitBlocks(); if (reservedCrossBlocks) @@ -435,55 +363,31 @@ std::tuple GuaranteedNoEvictScheduler::impl( // TODO(nhaber): remove forward declare and just keep the function here, right before the merge. I put it below just so // the remote diff is easier to look at/rebase conflicts bool trySchedulingRequestMaxUtilization(std::shared_ptr const& req, SizeType32 maxNumRequests, - std::size_t& numAdmittedRequests, RequestVector& scheduledRequests, - kv_cache_manager::MaxUtilizationScheduledBlocksManager& blocksManager, - std::optional& crossBlocksManager, + RequestVector& scheduledRequests, kv_cache_manager::MaxUtilizationScheduledBlocksManager& blocksManager, OptionalRef peftCacheManager, SizeType32& numScheduledPeftPages, std::unordered_set& seenTaskIds, std::optional const& cachedSummary); std::tuple MaxUtilizationScheduler::operator()( - kv_cache_manager::BaseKVCacheManager& kvCacheManager, - OptionalRef crossKvCacheManager, - OptionalRef peftCacheManager, RequestList const& activeRequests) const + kv_cache_manager::BaseKVCacheManager& kvCacheManager, OptionalRef peftCacheManager, + RequestList const& activeRequests) const { - checkRequiredCrossKvCacheManager(getNoScheduleUntilState(), crossKvCacheManager); - kvCacheManager.startScheduling(); - if (crossKvCacheManager) - { - crossKvCacheManager->startScheduling(); - } // The optimization of delaying requests won't work for variable window attention - bool skippingIsRelevant = mEnablePrefixAwareScheduling && !kvCacheManager.getBlockManager().isVariableWindow(); + bool skippingIsRelevant = !kvCacheManager.getBlockManager().isVariableWindow(); // Keep track of number of requests and block needed for the scheduled requests auto scheduledBlocksManager = kv_cache_manager::MaxUtilizationScheduledBlocksManager(kvCacheManager, mTwoStepsLookAhead); - // Mirror the budget tracker for the cross pool when present. - // Encoder-init requests do not consume either tracker; decoder - // context/generation requests update both trackers in lockstep. - std::optional scheduledCrossBlocksManager; - if (crossKvCacheManager) - { - scheduledCrossBlocksManager.emplace(*crossKvCacheManager, mTwoStepsLookAhead); - } SizeType32 numScheduledPeftPages{0}; std::unordered_set seenTaskIds; // Keep track of blocks contributed by requests in context phase - std::unordered_set newlyContributedContextBlocks; - std::unordered_set newlyContributedCrossContextBlocks; - if (skippingIsRelevant) - { - std::tie(newlyContributedContextBlocks, newlyContributedCrossContextBlocks) - = prefillWithChunkedContextsAlreadyExecuting(activeRequests, kvCacheManager); - } + auto [newlyContributedContextBlocks, newlyContributedCrossContextBlocks] + = prefillWithChunkedContextsAlreadyExecuting(activeRequests, kvCacheManager); - // Find last active in case we need to evict. Encoder-init requests are - // intentionally excluded here: they hold no started self- or cross-pool - // blocks, so pausing them would not free any KV budget. + // Find last active in case we need to evict auto startedReqLambda = [this](std::shared_ptr const& req) { return (req->hasReachedState(getNoScheduleUntilState()) && !req->hasReachedState(getNoScheduleAfterState()) @@ -492,7 +396,6 @@ std::tuple MaxUtilizationScheduler::operator()( RequestVector scheduledRequests; RequestVector pausedRequests; - std::size_t numAdmittedRequests{0}; auto reqItEnd = std::end(activeRequests); for (auto reqIt = std::begin(activeRequests); reqIt != reqItEnd;) { @@ -502,7 +405,7 @@ std::tuple MaxUtilizationScheduler::operator()( // if request cannot be scheduled yet or request should no longer be scheduled, skip if ( // Allow disagg_generation_init requests to be scheduled, so that we'll allocate their KV cache - !req->isDisaggGenerationInitState() && !req->isDisaggGenerationTransmissionInProgress() + !req->isDisaggGenerationInitState() && (!req->hasReachedState(getNoScheduleUntilState()) || req->hasReachedState(getNoScheduleAfterState()))) { TLLM_LOG_DEBUG("MaxUtilizationScheduler: request ID %lu cannot / should not be scheduled", req->mRequestId); @@ -510,18 +413,6 @@ std::tuple MaxUtilizationScheduler::operator()( continue; } - if (req->isDisaggGenerationTransmissionInProgress()) - { - if (numAdmittedRequests >= static_cast(mMaxNumRequests)) - { - break; - } - claimPeftPagesForRequest(req, peftCacheManager, numScheduledPeftPages, seenTaskIds); - ++numAdmittedRequests; - reqIt++; - continue; - } - // For first-chunk context requests with block reuse, compute the prefix reuse // summary once. This single radix tree walk serves both the beneficial-to-skip // check and the block budget estimation in getNeededBlocksOneStep. @@ -530,11 +421,7 @@ std::tuple MaxUtilizationScheduler::operator()( std::optional summary; // analyzePrefixReuse asserts on variable-window managers; skip the walk there // and let downstream callers fall back to their fresh tree-walk path. - if (isFirstChunkContext && !mEnablePrefixAwareScheduling) - { - summary = kv_cache_manager::PrefixReuseSummary{}; - } - else if (isFirstChunkContext && kvCacheManager.isEnableBlockReuse() + if (isFirstChunkContext && kvCacheManager.isEnableBlockReuse() && !kvCacheManager.getBlockManager().isVariableWindow()) { auto uniqueTokens = req->getUniqueTokens(0); @@ -550,9 +437,8 @@ std::tuple MaxUtilizationScheduler::operator()( continue; } - bool const wasScheduled = trySchedulingRequestMaxUtilization(req, mMaxNumRequests, numAdmittedRequests, - scheduledRequests, scheduledBlocksManager, scheduledCrossBlocksManager, peftCacheManager, - numScheduledPeftPages, seenTaskIds, summary); + bool const wasScheduled = trySchedulingRequestMaxUtilization(req, mMaxNumRequests, scheduledRequests, + scheduledBlocksManager, peftCacheManager, numScheduledPeftPages, seenTaskIds, summary); if (wasScheduled) { TLLM_LOG_DEBUG("MaxUtilizationScheduler: request ID %lu -> start", req->mRequestId); @@ -569,13 +455,6 @@ std::tuple MaxUtilizationScheduler::operator()( // from the end of the vector and try again // Here we simulate freeing the kvCache blocks associated with that sequence kvCacheManager.schedulingRemoveSequence((*lastStartedReqIt)->mRequestId); - if (crossKvCacheManager) - { - // Mirror self-pool eviction on the cross pool so any cross - // blocks held by the paused request are released for reuse - // by other admissions in this iteration. - crossKvCacheManager->schedulingRemoveSequence((*lastStartedReqIt)->mRequestId); - } pausedRequests.emplace_back(*lastStartedReqIt); TLLM_LOG_INFO("MaxUtilizationScheduler: request ID %lu -> pause", (*lastStartedReqIt)->mRequestId); reqItEnd = std::next(lastStartedReqIt).base(); @@ -591,13 +470,11 @@ std::tuple MaxUtilizationScheduler::operator()( } bool trySchedulingRequestMaxUtilization(std::shared_ptr const& req, SizeType32 maxNumRequests, - std::size_t& numAdmittedRequests, RequestVector& scheduledRequests, - kv_cache_manager::MaxUtilizationScheduledBlocksManager& blocksManager, - std::optional& crossBlocksManager, + RequestVector& scheduledRequests, kv_cache_manager::MaxUtilizationScheduledBlocksManager& blocksManager, OptionalRef peftCacheManager, SizeType32& numScheduledPeftPages, std::unordered_set& seenTaskIds, std::optional const& cachedSummary) { - if (numAdmittedRequests < static_cast(maxNumRequests)) + if (scheduledRequests.size() < static_cast(maxNumRequests)) { bool reqHasLora = req->getLoraTaskId().has_value(); bool isNewTask = reqHasLora && !seenTaskIds.count(req->getLoraTaskId().value()); @@ -605,56 +482,19 @@ bool trySchedulingRequestMaxUtilization(std::shared_ptr const& req, = (isNewTask && peftCacheManager) ? peftCacheManager->determineNumPages(req) : 0; TLLM_LOG_DEBUG( "MaxUtilizationScheduler: request ID %lu required peft pages: %i", req->mRequestId, numRequiredPeftPages); - bool fitsPeft - = (peftCacheManager ? numRequiredPeftPages + numScheduledPeftPages <= peftCacheManager->getMaxDevicePages() - : true); - - if (req->isEncoderInitState()) - { - // Encoder admission does not reserve KV blocks. The scheduler - // entry point verifies the cross manager globally before encoder - // work can be admitted. - if (fitsPeft) - { - numScheduledPeftPages += numRequiredPeftPages; - scheduledRequests.emplace_back(req); - ++numAdmittedRequests; - if (isNewTask) - { - seenTaskIds.insert(req->getLoraTaskId().value()); - } - return true; - } - return false; - } - // Use the cached summary when available to avoid a redundant tree walk auto const scheduledBlocksIfFitsKvCache = blocksManager.prepareNewNumberOfBlocksIfWeEndUpScheduling(*req, cachedSummary); - // Context/generation requests must fit in both pools when a cross - // manager is present. Self-pool fit is checked first so that the - // budget probe is cheap when self is already saturated. - std::optional> crossScheduledIfFits; - if (crossBlocksManager) - { - crossScheduledIfFits = crossBlocksManager->prepareNewNumberOfBlocksIfWeEndUpScheduling(*req); - if (!crossScheduledIfFits) - { - return false; - } - } + bool fitsPeft + = (peftCacheManager ? numRequiredPeftPages + numScheduledPeftPages <= peftCacheManager->getMaxDevicePages() + : true); if (scheduledBlocksIfFitsKvCache && fitsPeft) { blocksManager.updateScheduledBlocks(scheduledBlocksIfFitsKvCache.value()); - if (crossScheduledIfFits) - { - crossBlocksManager->updateScheduledBlocks(crossScheduledIfFits.value()); - } numScheduledPeftPages += numRequiredPeftPages; TLLM_LOG_DEBUG("MaxUtilizationScheduler: scheduled peft pages: %i", numRequiredPeftPages); scheduledRequests.emplace_back(req); - ++numAdmittedRequests; if (isNewTask) { seenTaskIds.insert(req->getLoraTaskId().value()); @@ -667,7 +507,7 @@ bool trySchedulingRequestMaxUtilization(std::shared_ptr const& req, CapacityScheduler::CapacityScheduler(SizeType32 maxNumRequests, executor::CapacitySchedulerPolicy capacitySchedulerPolicy, bool hasKvCacheManager, bool twoStepsLookAhead, - LlmRequestState noScheduleUntilState, LlmRequestState noScheduleAfterState, bool enablePrefixAwareScheduling) + LlmRequestState noScheduleUntilState, LlmRequestState noScheduleAfterState) { if (!hasKvCacheManager) { @@ -675,18 +515,16 @@ CapacityScheduler::CapacityScheduler(SizeType32 maxNumRequests, } else if (capacitySchedulerPolicy == executor::CapacitySchedulerPolicy::kMAX_UTILIZATION) { - mScheduler = MaxUtilizationScheduler{ - maxNumRequests, twoStepsLookAhead, noScheduleUntilState, noScheduleAfterState, enablePrefixAwareScheduling}; + mScheduler + = MaxUtilizationScheduler{maxNumRequests, twoStepsLookAhead, noScheduleUntilState, noScheduleAfterState}; } else if (capacitySchedulerPolicy == executor::CapacitySchedulerPolicy::kGUARANTEED_NO_EVICT) { - mScheduler = GuaranteedNoEvictScheduler{ - maxNumRequests, noScheduleUntilState, noScheduleAfterState, enablePrefixAwareScheduling}; + mScheduler = GuaranteedNoEvictScheduler{maxNumRequests, noScheduleUntilState, noScheduleAfterState}; } else if (capacitySchedulerPolicy == executor::CapacitySchedulerPolicy::kSTATIC_BATCH) { - mScheduler = StaticBatchScheduler{ - maxNumRequests, noScheduleUntilState, noScheduleAfterState, enablePrefixAwareScheduling}; + mScheduler = StaticBatchScheduler{maxNumRequests, noScheduleUntilState, noScheduleAfterState}; } else { @@ -708,7 +546,7 @@ void CapacityScheduler::setAgentTreeReorderPolicy( std::tuple CapacityScheduler::operator()(RequestList const& activeRequests, OptionalRef kvCacheManager, OptionalRef peftCacheManager, - OptionalRef crossKvCacheManager) const + OptionalRef crossKvCacheManager) const { NVTX3_SCOPED_RANGE(capacitySchedulerScheduling); @@ -728,7 +566,7 @@ std::tuple CapacityScheduler::opera else if constexpr (std::is_same_v, MaxUtilizationScheduler>) { std::tie(tmpFittingRequests, pausedRequests) - = scheduler(*kvCacheManager, crossKvCacheManager, peftCacheManager, requestsToSchedule); + = scheduler(*kvCacheManager, peftCacheManager, requestsToSchedule); } else if constexpr (std::is_same_v, GuaranteedNoEvictScheduler> || std::is_same_v, StaticBatchScheduler>) diff --git a/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp b/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp index 109417965a75..b3eb1c0f9843 100644 --- a/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp +++ b/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp @@ -29,7 +29,6 @@ #include "tensorrt_llm/executor/cache_transmission/cacheSplitConcat.h" #include "tensorrt_llm/runtime/common.h" #include "tensorrt_llm/runtime/utils/mpiUtils.h" -#include #include #include #include @@ -126,49 +125,6 @@ void TransferSession::appendMeasure(LlmRequest::TimePoint start, LlmRequest::Tim } } -void TransferSession::setReservedRecvBuffers(std::vector holders) -{ - TLLM_CHECK(mReservedRecvBuffers.empty()); - mReservedRecvBuffers = std::move(holders); -} - -bool TransferSession::hasReservedRecvBuffer(BaseTransBufferManager const& manager) const noexcept -{ - return std::any_of(mReservedRecvBuffers.begin(), mReservedRecvBuffers.end(), - [&manager](BufferIndexHolder const& holder) { return holder.isBoundTo(manager); }); -} - -bool TransferSession::releaseReservedRecvBuffer(BaseTransBufferManager const& manager) noexcept -{ - auto const holderIt = std::find_if(mReservedRecvBuffers.begin(), mReservedRecvBuffers.end(), - [&manager](BufferIndexHolder const& holder) { return holder.isBoundTo(manager); }); - if (holderIt == mReservedRecvBuffers.end()) - { - return false; - } - holderIt->release(); - mReservedRecvBuffers.erase(holderIt); - return true; -} - -void TransferSession::releaseReservedRecvBuffers() noexcept -{ - for (auto& holder : mReservedRecvBuffers) - { - holder.release(); - } - mReservedRecvBuffers.clear(); -} - -void TransferSession::poisonReservedRecvBuffers() noexcept -{ - for (auto& holder : mReservedRecvBuffers) - { - holder.poison(); - } - mReservedRecvBuffers.clear(); -} - void TransferSession::exportMeasure(std::ofstream& outFile, bool isContext) const { if (!mTimes || mTimes->measures.empty()) @@ -347,10 +303,6 @@ class CacheSender::Impl std::promise promise; auto future = promise.get_future(); llmRequest->setKvCacheTransferStart(LlmRequest::getSteadyClockNow()); - if (common::getEnvDisaggEnableInflightCancel()) - { - (void) getOrCreateInFlightCancelFlag(llmRequest->mRequestId); - } { std::scoped_lock lock(mSenderMutex); TLLM_CHECK_WITH_INFO( @@ -364,19 +316,6 @@ class CacheSender::Impl return future; } - std::shared_ptr> getOrCreateInFlightCancelFlag(RequestIdType requestId) - { - std::lock_guard lg(mInFlightCancelMutex); - auto it = mInFlightCancelFlags.find(requestId); - if (it != mInFlightCancelFlags.end()) - { - return it->second; - } - auto flag = std::make_shared>(false); - mInFlightCancelFlags.emplace(requestId, flag); - return flag; - } - [[nodiscard]] executor::kv_cache::CommState const& getCommState() const { return mSelfState.getCommState().value(); @@ -397,65 +336,24 @@ class CacheSender::Impl void release(LlmRequest::RequestIdType requestId) { + std::unique_lock lk(mMtxForMap); + auto it = mRequestToSession.find(requestId); + TLLM_CHECK(it != mRequestToSession.end()); + if (!common::getEnvKVCacheTimeOutputPath().empty()) { - std::unique_lock lk(mMtxForMap); - auto it = mRequestToSession.find(requestId); - TLLM_CHECK(it != mRequestToSession.end()); - if (!common::getEnvKVCacheTimeOutputPath().empty()) + if (!mMeasuresFile.is_open()) { - if (!mMeasuresFile.is_open()) - { - auto outputPath = getTransferOutputPath("send"); - mMeasuresFile.open(outputPath); - TLLM_CHECK_WITH_INFO(mMeasuresFile.is_open(), "Failed to open transfer output file: %s", - outputPath.string().c_str()); - } - it->second.exportMeasure(mMeasuresFile, true); + auto outputPath = getTransferOutputPath("send"); + mMeasuresFile.open(outputPath); + TLLM_CHECK_WITH_INFO( + mMeasuresFile.is_open(), "Failed to open transfer output file: %s", outputPath.string().c_str()); } - mRequestToSession.erase(it); - } - if (common::getEnvDisaggEnableInflightCancel()) - { - std::lock_guard lg(mInFlightCancelMutex); - mInFlightCancelFlags.erase(requestId); - } - } - - void discardTransferState(LlmRequest::RequestIdType requestId) noexcept - { - try - { - std::unique_lock lk(mMtxForMap); - mRequestToSession.erase(requestId); - } - catch (std::exception const& e) - { - TLLM_LOG_WARNING("Failed to discard sender session for request %ld: %s", requestId, e.what()); - } - catch (...) - { - TLLM_LOG_WARNING("Failed to discard sender session for request %ld: unknown exception", requestId); - } - if (!common::getEnvDisaggEnableInflightCancel()) - { - return; - } - try - { - std::lock_guard lg(mInFlightCancelMutex); - mInFlightCancelFlags.erase(requestId); - } - catch (std::exception const& e) - { - TLLM_LOG_WARNING("Failed to discard in-flight cancel flag for request %ld: %s", requestId, e.what()); - } - catch (...) - { - TLLM_LOG_WARNING("Failed to discard in-flight cancel flag for request %ld: unknown exception", requestId); + it->second.exportMeasure(mMeasuresFile, true); } + mRequestToSession.erase(it); } - [[nodiscard]] std::optional recvRequestInfo() + [[nodiscard]] RequestInfo recvRequestInfo() { auto* agentConnectionManager = dynamic_cast(mManager); bool isAgent = agentConnectionManager != nullptr; @@ -465,10 +363,10 @@ class CacheSender::Impl auto const* connection = isAgent ? agentConnectionManager->recvConnectionAndRequestInfo(info, mTerminate) : mManager->recvConnect(DataContext{TransceiverTag::kID_TAG, mTerminate}, &id, sizeof(id)); - if (connection == nullptr) + if (connection == nullptr && !mManager->isRunning()) { - TLLM_LOG_WARNING("recvRequestInfo connection is nullptr, maybe the server is terminating"); - return std::nullopt; + TLLM_LOG_WARNING(" recvRequestInfo connection is nullptr, maybe the server is terminating"); + return info; } if (!isAgent) @@ -495,25 +393,15 @@ class CacheSender::Impl TLLM_CHECK_WITH_INFO(peerIdx < static_cast(allCounterparts.size()), "Peer rank %d not found in expected counterparts", peerSelfIdx); - std::shared_ptr> cancelFlag; - if (common::getEnvDisaggEnableInflightCancel()) - { - cancelFlag = getOrCreateInFlightCancelFlag(requestId); - } { std::unique_lock lk(mMtxForMap); auto it = mRequestToSession.find(requestId); if (it == mRequestToSession.end()) { - auto session = cancelFlag != nullptr - ? TransferSession(std::vector(allCounterparts.size(), nullptr), - DataContext{tagFromRequestId(requestId), *cancelFlag}, allCounterparts, mSelfState, - info.getTransState(), mBufferManager, info.getIndexFromEnd(), info.getLastBlockKey(), nullptr, - !common::getEnvKVCacheTimeOutputPath().empty()) - : TransferSession(std::vector(allCounterparts.size(), nullptr), - DataContext{tagFromRequestId(requestId), mTerminate}, allCounterparts, mSelfState, - info.getTransState(), mBufferManager, info.getIndexFromEnd(), info.getLastBlockKey(), nullptr, - !common::getEnvKVCacheTimeOutputPath().empty()); + auto session = TransferSession(std::vector(allCounterparts.size(), nullptr), + DataContext{tagFromRequestId(requestId), mTerminate}, allCounterparts, mSelfState, + info.getTransState(), mBufferManager, info.getIndexFromEnd(), info.getLastBlockKey(), nullptr, + !common::getEnvKVCacheTimeOutputPath().empty()); session.setTime(TransferSession::kTimeRequestInfo); it = mRequestToSession.emplace(requestId, std::move(session)).first; } @@ -532,60 +420,25 @@ class CacheSender::Impl session = std::addressof(it->second); } session->setLlmRequest(llmRequest); - TLLM_LOG_DEBUG("KV cache transfer request %zu phase=transfer-submit begin.", llmRequest.mRequestId); mCacheTransferLayer.format(*session); - TLLM_LOG_DEBUG("KV cache transfer request %zu phase=transfer-complete end.", llmRequest.mRequestId); llmRequest.setKvCacheTransferEnd(LlmRequest::getSteadyClockNow()); } bool cancelRequest(LlmRequest const& llmRequest) { - bool const inflightCancelEnabled = common::getEnvDisaggEnableInflightCancel(); bool isCancelled = false; - bool isCurrentRequest = false; + std::scoped_lock lock(mSenderMutex); + auto it = mReadyResponses.find(llmRequest.mRequestId); + // If the request is not the current request and already in the ready queue, we can cancel it. + if (it != mReadyResponses.end() + && (!mCurrentRequest.has_value() || mCurrentRequest.value() != llmRequest.mRequestId)) { - std::scoped_lock lock(mSenderMutex); - auto it = mReadyResponses.find(llmRequest.mRequestId); - if (it != mReadyResponses.end()) - { - isCurrentRequest = mCurrentRequest.has_value() && mCurrentRequest.value() == llmRequest.mRequestId; - // The legacy path cannot interrupt a ready/active transfer, so - // preserve its false return until the opt-in is enabled. - if (!isCurrentRequest || inflightCancelEnabled) - { - mCancelledRequests.insert(llmRequest.mRequestId); - isCancelled = true; - if (inflightCancelEnabled && !isCurrentRequest) - { - // Keep only the request ID as a tombstone so a late peer - // receives ready=false without retaining the request. - failResponse(it->second, - std::make_exception_ptr( - TLLM_REQUEST_EXCEPTION(llmRequest.mRequestId, common::RequestErrorCode::kNETWORK_ERROR, - "Context KV cache request cancelled before a peer was ready for request %zu", - llmRequest.mRequestId))); - mReadyResponses.erase(it); - } - } - } - } - if (inflightCancelEnabled && (!isCancelled || isCurrentRequest)) - { - std::lock_guard lg(mInFlightCancelMutex); - auto flagIt = mInFlightCancelFlags.find(llmRequest.mRequestId); - if (flagIt != mInFlightCancelFlags.end()) - { - flagIt->second->store(true, std::memory_order_relaxed); - isCancelled = true; - } - } - if (!isCancelled) - { - TLLM_LOG_WARNING("Cannot cancel request %zu", llmRequest.mRequestId); + mCancelledRequests.insert(llmRequest.mRequestId); + isCancelled = true; } else { - mSenderCv.notify_all(); + TLLM_LOG_WARNING("Cannot cancel request %zu", llmRequest.mRequestId); } return isCancelled; } @@ -607,7 +460,8 @@ class CacheSender::Impl { auto* agentConnection = dynamic_cast(connections.at(i)); TLLM_CHECK(agentConnection); - agentConnection->sendReadySignal(session->getDataContext(), isReady); + agentConnection->sendReadySignal( + executor::kv_cache::DataContext{TransceiverTag::kREADY_SIGNAL_TAG}, isReady); } else { @@ -682,23 +536,13 @@ class CacheSender::Impl catch (tensorrt_llm::common::RequestSpecificException const& e) { TLLM_LOG_ERROR("Exception in sendAndRemoveResponse: %s ", e.what()); - discardTransferState(id); auto new_exception = TLLM_REQUEST_EXCEPTION(id, e.getErrorCode(), "%s", e.what()); - failResponse(resp, std::make_exception_ptr(new_exception)); + resp.mPromise.set_exception(std::make_exception_ptr(new_exception)); } catch (std::exception const& e) { - auto const exception = std::current_exception(); TLLM_LOG_ERROR("Exception in sendAndRemoveResponse: %s request id: %ld", e.what(), id); - discardTransferState(id); - failResponse(resp, exception); - } - catch (...) - { - auto const exception = std::current_exception(); - TLLM_LOG_ERROR("Unknown exception in sendAndRemoveResponse for request id: %ld", id); - discardTransferState(id); - failResponse(resp, exception); + resp.mPromise.set_exception(std::current_exception()); } } @@ -713,13 +557,6 @@ class CacheSender::Impl catch (std::exception const& err) { TLLM_LOG_ERROR("Failed to queue asynchronous KV cache send for request %zu: %s", id, err.what()); - discardTransferState(id); - failResponse(resp, std::current_exception()); - } - catch (...) - { - TLLM_LOG_ERROR("Unknown error while queueing asynchronous KV cache send for request %zu", id); - discardTransferState(id); failResponse(resp, std::current_exception()); } } @@ -727,44 +564,21 @@ class CacheSender::Impl void sendResponse(RequestIdType reqId) { bool isReady = true; - bool allCounterpartsReady = false; - std::optional cancelledResponse; { std::scoped_lock lock(mSenderMutex); TLLM_CHECK(mCurrentRequest.has_value() && mCurrentRequest.value() == reqId); - auto responseIt = mReadyResponses.find(reqId); - bool const isCancelled = mCancelledRequests.find(reqId) != mCancelledRequests.end(); - TLLM_CHECK(responseIt != mReadyResponses.end() || isCancelled); + TLLM_CHECK(mReadyResponses.find(reqId) != mReadyResponses.end()); auto countIt = mRemainSendCount.find(reqId); TLLM_CHECK(countIt != mRemainSendCount.end()); auto const count = --countIt->second; TLLM_CHECK(count >= 0); - if (isCancelled && responseIt != mReadyResponses.end()) - { - cancelledResponse.emplace(std::move(responseIt->second)); - mReadyResponses.erase(responseIt); - } if (count > 0) { mCurrentRequest = std::nullopt; + return; } - else - { - mRemainSendCount.erase(countIt); - isReady = !isCancelled; - allCounterpartsReady = true; - } - } - - if (cancelledResponse.has_value()) - { - failResponse(*cancelledResponse, - std::make_exception_ptr(TLLM_REQUEST_EXCEPTION(reqId, common::RequestErrorCode::kNETWORK_ERROR, - "KV cache transfer for request %zu was cancelled", reqId))); - } - if (!allCounterpartsReady) - { - return; + mRemainSendCount.erase(countIt); + isReady = mCancelledRequests.find(reqId) == mCancelledRequests.end(); } // Keep mCurrentRequest set while notifying the peer so cancellation cannot change the decision after it has @@ -775,12 +589,9 @@ class CacheSender::Impl { std::scoped_lock lock(mSenderMutex); auto it = mReadyResponses.find(reqId); - if (isReady) - { - TLLM_CHECK(it != mReadyResponses.end()); - response = std::move(it->second); - mReadyResponses.erase(it); - } + TLLM_CHECK(it != mReadyResponses.end()); + response = std::move(it->second); + mReadyResponses.erase(it); mCancelledRequests.erase(reqId); mCurrentRequest = std::nullopt; } @@ -802,7 +613,8 @@ class CacheSender::Impl } else { - discardTransferState(reqId); + response.mPromise.set_exception(std::make_exception_ptr(TLLM_REQUEST_EXCEPTION(reqId, + common::RequestErrorCode::kNETWORK_ERROR, "KV cache transfer for request %zu was cancelled", reqId))); } } @@ -817,20 +629,19 @@ class CacheSender::Impl { { std::unique_lock lock(mSenderMutex); - mSenderCv.wait(lock, - [this]() { return mTerminate || !mReadyResponses.empty() || !mCancelledRequests.empty(); }); + mSenderCv.wait(lock, [this]() { return mTerminate || !mReadyResponses.empty(); }); if (mTerminate) { break; } } - auto requestInfo = recvRequestInfo(); - if (!requestInfo.has_value() || mTerminate || !mManager->isRunning()) + auto const requestInfo = recvRequestInfo(); + if (mTerminate || !mManager->isRunning()) { break; } - auto const reqId = requestInfo->getRequestId(); + auto const reqId = requestInfo.getRequestId(); if (mRemainSendCount.find(reqId) == mRemainSendCount.end()) { @@ -841,11 +652,7 @@ class CacheSender::Impl std::unique_lock lock(mSenderMutex); mCurrentRequest = reqId; mSenderCv.wait(lock, - [this, reqId]() - { - return mTerminate || mReadyResponses.find(reqId) != mReadyResponses.end() - || mCancelledRequests.find(reqId) != mCancelledRequests.end(); - }); + [this, reqId]() { return mTerminate || mReadyResponses.find(reqId) != mReadyResponses.end(); }); if (mTerminate) { mCurrentRequest = std::nullopt; @@ -860,11 +667,6 @@ class CacheSender::Impl TLLM_LOG_ERROR("Exception in CacheSender response: %s", err.what()); responseException = std::current_exception(); } - catch (...) - { - TLLM_LOG_ERROR("Unknown exception in CacheSender response"); - responseException = std::current_exception(); - } if (!responseException) { @@ -885,16 +687,6 @@ class CacheSender::Impl std::scoped_lock lock(mSenderMutex); mTerminate = true; } - if (common::getEnvDisaggEnableInflightCancel()) - { - std::lock_guard lg(mInFlightCancelMutex); - for (auto& [id, flag] : mInFlightCancelFlags) - { - flag->store(true, std::memory_order_relaxed); - } - } - // Wake the sender loop and make in-flight agent transfers observe termination through - // their per-request cancellation flags. mSenderCv.notify_all(); if (mResponseFuture.valid()) { @@ -977,8 +769,6 @@ class CacheSender::Impl std::mutex mMtxForMap; runtime::BufferManager mBufferManager; std::ofstream mMeasuresFile; - std::mutex mInFlightCancelMutex; - std::unordered_map>> mInFlightCancelFlags; }; class CacheReceiver::Impl @@ -1026,16 +816,9 @@ class CacheReceiver::Impl mRequestFutures.emplace_back(std::move(requestFuture)); } auto& asyncResource = mInstanceToAsyncResource.at(processInfo); - std::shared_ptr> cancelFlag; - if (common::getEnvDisaggEnableInflightCancel()) - { - cancelFlag = std::make_shared>(false); - std::lock_guard lg(mInFlightCancelMutex); - mInFlightCancelFlags[llmRequest->mRequestId] = cancelFlag; - } { std::unique_lock lck(asyncResource->mMtxForQueue); - asyncResource->mRequestsQueue.emplace_back(llmRequest, std::move(promise), cancelFlag); + asyncResource->mRequestsQueue.emplace_back(llmRequest, std::move(promise)); } asyncResource->mCVforQueue.notify_all(); return future; @@ -1048,34 +831,22 @@ class CacheReceiver::Impl void receiveSync(TransferSession& session) { - try - { - mCacheTransferLayer.unformat(session); - if (!common::getEnvKVCacheTimeOutputPath().empty()) - { - std::unique_lock lock(mMeasuresFileMutex); - if (!mMeasuresFile.is_open()) - { - auto outputPath = getTransferOutputPath("recv"); - mMeasuresFile.open(outputPath); - TLLM_CHECK_WITH_INFO(mMeasuresFile.is_open(), "Failed to open transfer output file: %s", - outputPath.string().c_str()); - } - session.exportMeasure(mMeasuresFile, false); - } - session.releaseReservedRecvBuffers(); - } - catch (...) + mCacheTransferLayer.unformat(session); + if (!common::getEnvKVCacheTimeOutputPath().empty()) { - if (common::getEnvDisaggEnableInflightCancel()) + std::unique_lock lock(mMeasuresFileMutex); + if (!mMeasuresFile.is_open()) { - session.poisonReservedRecvBuffers(); + auto outputPath = getTransferOutputPath("recv"); + mMeasuresFile.open(outputPath); + TLLM_CHECK_WITH_INFO( + mMeasuresFile.is_open(), "Failed to open transfer output file: %s", outputPath.string().c_str()); } - throw; + session.exportMeasure(mMeasuresFile, false); } } - TransferSession sendRequestInfo(LlmRequest const& llmRequest, std::atomic const* perRequestCancel = nullptr) + TransferSession sendRequestInfo(LlmRequest const& llmRequest) { uint64_t requestId = llmRequest.getContextPhaseParams().value().getReqId(); auto const& contextState = llmRequest.getDataTransceiverState(); @@ -1088,55 +859,35 @@ class CacheReceiver::Impl if (!mCacheTransferLayer.getCacheManager()->getBlockManager().isVariableWindow()) { auto* cacheManager = mCacheTransferLayer.getCacheManager(); + auto beam = 0; auto const srcPpSize = destCacheState.getParallelConfig().mPipelineParallelism; auto requestedBlockRange = getBlockRangeForReceiving(cacheManager, llmRequest, destCacheState.getEnableBlockReuse(), destCacheState.getEnablePartialReuse(), /*recvSideHasCP=*/false, srcPpSize); + auto const& uniqueTokens = llmRequest.getUniqueTokens(beam); + auto lastBlockKey + = BlockKey(llmRequest.getInputTokensExtraIds().has_value(), llmRequest.getLoraTaskId(), uniqueTokens); + auto tokensPerBlock = cacheManager->getBlockManager().getTokensPerBlock(); + SizeType32 startTokenIdx = static_cast(uniqueTokens.size() / tokensPerBlock) * tokensPerBlock; + SizeType32 endTokenIdx = static_cast(uniqueTokens.size()); + auto extraKeys = kv_cache_manager::generateBlockHashExtraKeys(llmRequest, startTokenIdx, endTokenIdx); + lastBlockKey.extraKeys = std::move(extraKeys); + // Compute indexFromEnd from the number of requested blocks int32_t requestedBlockSize = requestedBlockRange.getBlockIdsPerWindow().begin()->second.size(); - // An empty Helix CP rank owns zero KV blocks for this sequence (fewer blocks than - // cp_size). It still sends a RequestInfo so the context's per-request counterpart count - // is satisfied, but requests zero blocks: the default RequestInfo (indexFromEnd=0, empty - // lastBlockKey) is used and the context transmits nothing to it. - if (requestedBlockSize > 0) - { - auto const beam = 0; - auto const& uniqueTokens = llmRequest.getUniqueTokens(beam); - auto lastBlockKey = BlockKey( - llmRequest.getInputTokensExtraIds().has_value(), llmRequest.getLoraTaskId(), uniqueTokens); - auto tokensPerBlock = cacheManager->getBlockManager().getTokensPerBlock(); - SizeType32 startTokenIdx - = static_cast(uniqueTokens.size() / tokensPerBlock) * tokensPerBlock; - SizeType32 endTokenIdx = static_cast(uniqueTokens.size()); - auto extraKeys = kv_cache_manager::generateBlockHashExtraKeys(llmRequest, startTokenIdx, endTokenIdx); - lastBlockKey.extraKeys = std::move(extraKeys); - int32_t indexFromEnd = requestedBlockSize - 1; - - requestInfo = RequestInfo(requestId, mSelfState, indexFromEnd, lastBlockKey); - } + TLLM_CHECK_WITH_INFO(requestedBlockSize > 0, "requestedBlockSize must be > 0"); + int32_t indexFromEnd = requestedBlockSize - 1; + + requestInfo = RequestInfo(requestId, mSelfState, indexFromEnd, lastBlockKey); } auto* agentConnectionManager = dynamic_cast(mManager); - std::vector recvHolders; std::vector> cacheBufferIds; if (agentConnectionManager) { - auto const* bufferCancel = common::getEnvDisaggEnableInflightCancel() ? perRequestCancel : nullptr; - auto const& managers = agentConnectionManager->getCacheTransBufferManagers(); - recvHolders.reserve(managers.size()); - cacheBufferIds.reserve(managers.size()); - for (auto& cacheTransBufferManager : managers) + for (auto& cacheTransBufferManager : agentConnectionManager->getCacheTransBufferManagers()) { - auto rawIdx = cacheTransBufferManager->assignBufferIndexForRecv(bufferCancel); - recvHolders.emplace_back(*cacheTransBufferManager, rawIdx, /*isRecv=*/true); - if (rawIdx.has_value()) - { - cacheBufferIds.push_back(static_cast(rawIdx.value())); - } - else - { - cacheBufferIds.push_back(std::nullopt); - } + cacheBufferIds.push_back(cacheTransBufferManager->assignBufferIndexForRecv()); } TLLM_CHECK(!cacheBufferIds.empty()); } @@ -1165,102 +916,70 @@ class CacheReceiver::Impl allConnections.emplace_back(connection); } - if (common::getEnvDisaggEnableInflightCancel() && perRequestCancel != nullptr - && perRequestCancel->load(std::memory_order_relaxed)) + for (size_t ci = 0; ci < allCounterparts.size(); ci++) { - TLLM_THROW("KV cache receive request cancelled before publishing receive buffers"); - } + auto rank = allCounterparts[ci]; + auto const* connection = connections.at(rank); - try - { - for (size_t ci = 0; ci < allCounterparts.size(); ci++) - { - auto rank = allCounterparts[ci]; - auto const* connection = connections.at(rank); - - bool isKvCounterpart - = std::find(kvCounterParts.begin(), kvCounterParts.end(), rank) != kvCounterParts.end(); - bool isRnnCounterpart = hasRnn - && std::find(rnnCounterParts.begin(), rnnCounterParts.end(), rank) != rnnCounterParts.end(); + bool isKvCounterpart + = std::find(kvCounterParts.begin(), kvCounterParts.end(), rank) != kvCounterParts.end(); + bool isRnnCounterpart + = hasRnn && std::find(rnnCounterParts.begin(), rnnCounterParts.end(), rank) != rnnCounterParts.end(); - if (agentConnectionManager) + if (agentConnectionManager) + { + auto idsForRank = cacheBufferIds; + auto const& managers = agentConnectionManager->getCacheTransBufferManagers(); + for (size_t i = 0; i < idsForRank.size(); i++) { - auto idsForRank = cacheBufferIds; - auto const& managers = agentConnectionManager->getCacheTransBufferManagers(); - for (size_t i = 0; i < idsForRank.size(); i++) - { - auto kind = managers[i]->getBufferKind(); - bool include = (kind != BufferKind::kRNN) ? isKvCounterpart : isRnnCounterpart; - if (!include) - { - idsForRank[i] = std::nullopt; - } - } - - int validConnectionIdx = 0; - if (isKvCounterpart) + auto kind = managers[i]->getBufferKind(); + bool include = (kind != BufferKind::kRNN) ? isKvCounterpart : isRnnCounterpart; + if (!include) { - auto kvCpIdx - = std::find(kvCounterParts.begin(), kvCounterParts.end(), rank) - kvCounterParts.begin(); - auto [pickUpIdx, localRankIdx] = mCacheTransferLayer.getKvFormatter()->pickRecvConnections( - allCounterparts.size(), mSelfState.getCacheState().value(), - mSelfState.getCommState().value().getSelfIdx(), destCacheState, allCounterparts); - validConnectionIdx - = std::find(localRankIdx.begin(), localRankIdx.end(), kvCpIdx) - localRankIdx.begin(); + idsForRank[i] = std::nullopt; } - else if (isRnnCounterpart) - { - auto rnnTargetInfo = executor::kv_cache::targetIRanksForRnn(destCacheState, - mCacheTransferLayer.getCacheState(), mSelfState.getCommState().value().getSelfIdx()); - auto rnnCpIdx - = std::find(rnnCounterParts.begin(), rnnCounterParts.end(), rank) - rnnCounterParts.begin(); - auto [pickUpIdx, localRankIdx] - = cache_formatter_utils::pickRecvConnections(rnnCounterParts.size(), - mCacheTransferLayer.getCacheState(), mSelfState.getCommState().value().getSelfIdx(), - destCacheState, rnnCounterParts, rnnTargetInfo); - validConnectionIdx - = std::find(localRankIdx.begin(), localRankIdx.end(), rnnCpIdx) - localRankIdx.begin(); - } - - auto* agentConnection = dynamic_cast(connection); - TLLM_CHECK(agentConnection != nullptr); + } - const_cast(agentConnection) - ->sendRequestAndBufferInfo(requestInfo, idsForRank, validConnectionIdx, perRequestCancel); + int validConnectionIdx = 0; + if (isKvCounterpart) + { + auto kvCpIdx + = std::find(kvCounterParts.begin(), kvCounterParts.end(), rank) - kvCounterParts.begin(); + auto [pickUpIdx, localRankIdx] = mCacheTransferLayer.getKvFormatter()->pickRecvConnections( + allCounterparts.size(), mSelfState.getCacheState().value(), + mSelfState.getCommState().value().getSelfIdx(), destCacheState, allCounterparts); + validConnectionIdx + = std::find(localRankIdx.begin(), localRankIdx.end(), kvCpIdx) - localRankIdx.begin(); } - else + else if (isRnnCounterpart) { - sendRequestInfo(connection, requestInfo); + auto rnnTargetInfo = executor::kv_cache::targetIRanksForRnn(destCacheState, + mCacheTransferLayer.getCacheState(), mSelfState.getCommState().value().getSelfIdx()); + auto rnnCpIdx + = std::find(rnnCounterParts.begin(), rnnCounterParts.end(), rank) - rnnCounterParts.begin(); + auto [pickUpIdx, localRankIdx] = cache_formatter_utils::pickRecvConnections(rnnCounterParts.size(), + mCacheTransferLayer.getCacheState(), mSelfState.getCommState().value().getSelfIdx(), + destCacheState, rnnCounterParts, rnnTargetInfo); + validConnectionIdx + = std::find(localRankIdx.begin(), localRankIdx.end(), rnnCpIdx) - localRankIdx.begin(); } - } - auto const& resource = getReceiveCacheResource(llmRequest); - TransferSession session = perRequestCancel != nullptr - ? TransferSession(std::move(allConnections), - DataContext{tagFromRequestId(requestId), *perRequestCancel}, std::move(allCounterparts), mSelfState, - contextState, resource->mBufferManager, requestInfo.getIndexFromEnd(), - requestInfo.getLastBlockKey(), &llmRequest, !common::getEnvKVCacheTimeOutputPath().empty()) - : TransferSession(std::move(allConnections), DataContext{tagFromRequestId(requestId), mTerminate}, - std::move(allCounterparts), mSelfState, contextState, resource->mBufferManager, - requestInfo.getIndexFromEnd(), requestInfo.getLastBlockKey(), &llmRequest, - !common::getEnvKVCacheTimeOutputPath().empty()); - if (!recvHolders.empty()) - { - session.setReservedRecvBuffers(std::move(recvHolders)); + auto* agentConnection = dynamic_cast(connection); + TLLM_CHECK(agentConnection != nullptr); + + const_cast(agentConnection) + ->sendRequestAndBufferInfo(requestInfo, idsForRank, validConnectionIdx); } - return session; - } - catch (...) - { - if (common::getEnvDisaggEnableInflightCancel()) + else { - for (auto& holder : recvHolders) - { - holder.poison(); - } + sendRequestInfo(connection, requestInfo); } - throw; } + auto const& resource = getReceiveCacheResource(llmRequest); + return TransferSession(std::move(allConnections), DataContext{tagFromRequestId(requestId), mTerminate}, + std::move(allCounterparts), mSelfState, contextState, resource->mBufferManager, + requestInfo.getIndexFromEnd(), requestInfo.getLastBlockKey(), &llmRequest, + !common::getEnvKVCacheTimeOutputPath().empty()); } std::unique_ptr const& getReceiveCacheResource(LlmRequest const& llmRequest) @@ -1299,27 +1018,11 @@ class CacheReceiver::Impl std::string processInfo = kDefaultProcessInfo; if (common::getEnvRequestKVCacheConcurrent()) { - auto const& commState = llmRequest.getDataTransceiverState().getCommState(); - if (!commState.has_value()) - { - TLLM_LOG_WARNING("Cannot cancel request %zu: the request has no data-transceiver communication state", - llmRequest.mRequestId); - return false; - } - processInfo = commState->toString(); - } - - auto const resourceIt = mInstanceToAsyncResource.find(processInfo); - if (resourceIt == mInstanceToAsyncResource.end()) - { - TLLM_LOG_WARNING("Cannot cancel request %zu: receive worker %s is not registered", llmRequest.mRequestId, - processInfo.c_str()); - return false; + processInfo = llmRequest.getDataTransceiverState().getCommState()->toString(); } bool isCancelled = false; - auto& asyncResource = resourceIt->second; - std::optional queuedCancelledReqId; + auto& asyncResource = mInstanceToAsyncResource.at(processInfo); { std::unique_lock lck(asyncResource->mMtxForQueue); auto it = std::find_if(asyncResource->mRequestsQueue.begin(), asyncResource->mRequestsQueue.end(), @@ -1346,116 +1049,45 @@ class CacheReceiver::Impl } asyncResource->mRequestsQueue.erase(it); isCancelled = true; - queuedCancelledReqId = llmRequest.mRequestId; } - } - if (common::getEnvDisaggEnableInflightCancel() && queuedCancelledReqId.has_value()) - { - std::lock_guard lg(mInFlightCancelMutex); - mInFlightCancelFlags.erase(*queuedCancelledReqId); - } - if (!isCancelled && common::getEnvDisaggEnableInflightCancel()) - { - std::lock_guard lg(mInFlightCancelMutex); - auto flagIt = mInFlightCancelFlags.find(llmRequest.mRequestId); - if (flagIt != mInFlightCancelFlags.end()) + else { - flagIt->second->store(true, std::memory_order_relaxed); - isCancelled = true; + TLLM_LOG_WARNING("Cannot cancel request %zu", llmRequest.mRequestId); } } - if (!isCancelled) - { - TLLM_LOG_WARNING("Cannot cancel request %zu", llmRequest.mRequestId); - } return isCancelled; } - enum class ReadySignalResult - { - kReady, - kNotReady, - kMixed, - kCancelled, - }; - - ReadySignalResult receiveReadySignalDetailed(TransferSession& session, std::atomic const& perRequestCancel) + bool receiveReadySignal(TransferSession& session) { + bool isReadyFinal = true; bool isReady = false; - bool anyReady = false; - bool anyNotReady = false; auto const& connections = session.getConnections(); + for (size_t i = 0; i < connections.size(); i++) { - if (perRequestCancel.load(std::memory_order_relaxed)) - { - return ReadySignalResult::kCancelled; - } auto* agentConnectionManager = dynamic_cast(mManager); if (agentConnectionManager) { auto* agentConnection = dynamic_cast(connections.at(i)); TLLM_CHECK(agentConnection); - auto ready = agentConnection->recvReadySignalWithStatus( - executor::kv_cache::DataContext{session.getDataContext().getTag(), perRequestCancel}); - if (!ready.has_value()) - { - return ReadySignalResult::kCancelled; - } - isReady = ready.value(); + isReady = agentConnection->recvReadySignal( + executor::kv_cache::DataContext{TransceiverTag::kREADY_SIGNAL_TAG, mTerminate}); } else { connections.at(i)->recv( executor::kv_cache::DataContext{TransceiverTag::kREADY_SIGNAL_TAG}, &isReady, sizeof(isReady)); - if (perRequestCancel.load(std::memory_order_relaxed)) - { - return ReadySignalResult::kCancelled; - } } - anyReady |= isReady; - anyNotReady |= !isReady; + isReadyFinal &= isReady; } - if (anyReady && anyNotReady) - { - return ReadySignalResult::kMixed; - } - return anyReady ? ReadySignalResult::kReady : ReadySignalResult::kNotReady; - } - - bool receiveReadySignal(TransferSession& session) - { - auto const result = receiveReadySignalDetailed(session, mTerminate); - if (result == ReadySignalResult::kNotReady) - { - session.releaseReservedRecvBuffers(); - } - else if (result == ReadySignalResult::kMixed) - { - if (common::getEnvDisaggEnableInflightCancel()) - { - session.poisonReservedRecvBuffers(); - } - else - { - session.releaseReservedRecvBuffers(); - } - } - return result == ReadySignalResult::kReady; + return isReadyFinal; } ~Impl() { mTerminate.store(true); - if (common::getEnvDisaggEnableInflightCancel()) - { - std::lock_guard lg(mInFlightCancelMutex); - for (auto& [id, flag] : mInFlightCancelFlags) - { - flag->store(true, std::memory_order_relaxed); - } - } for (auto&& [processInfo, asyncResource] : mInstanceToAsyncResource) { asyncResource->mTerminate = true; @@ -1468,105 +1100,29 @@ class CacheReceiver::Impl } private: - void requestSync(LlmRequest& llmRequest, std::atomic const& perRequestCancel) + void requestSync(LlmRequest& llmRequest) { - auto const requestId = llmRequest.mRequestId; - auto const contextRequestId = llmRequest.getContextPhaseParams().value().getReqId(); - char const* phase = "request-info"; - TLLM_LOG_DEBUG("KV cache receive request %zu, context request %zu started.", requestId, contextRequestId); - if (llmRequest.getKvCacheTransferStart() == LlmRequest::TimePoint{}) - { - llmRequest.setKvCacheTransferStart(LlmRequest::getSteadyClockNow()); - } - - std::optional session; - try - { - if (perRequestCancel.load(std::memory_order_relaxed) || mTerminate.load(std::memory_order_relaxed)) - { - TLLM_THROW("KV cache receive request %zu cancelled before request-info", requestId); - } - TLLM_CUDA_CHECK(cudaSetDevice(mDeviceId)); - TLLM_LOG_DEBUG("KV cache receive request %zu, context request %zu phase=%s begin.", requestId, - contextRequestId, phase); - auto const* cancelFlag = common::getEnvDisaggEnableInflightCancel() ? &perRequestCancel : nullptr; - session.emplace(sendRequestInfo(llmRequest, cancelFlag)); - session->setTime(TransferSession::kTimeRequestInfo); - TLLM_LOG_DEBUG( - "KV cache receive request %zu, context request %zu phase=%s end.", requestId, contextRequestId, phase); - - phase = "ready-signal"; - TLLM_LOG_DEBUG("KV cache receive request %zu, context request %zu phase=%s begin.", requestId, - contextRequestId, phase); - auto readyResult = receiveReadySignalDetailed(*session, perRequestCancel); - TLLM_LOG_DEBUG("KV cache receive request %zu, context request %zu phase=%s end: result=%d.", requestId, - contextRequestId, phase, static_cast(readyResult)); - if (readyResult == ReadySignalResult::kCancelled) - { - if (common::getEnvDisaggEnableInflightCancel()) - { - session->poisonReservedRecvBuffers(); - } - TLLM_THROW("KV cache receive request %zu cancelled while waiting for the ready signal", requestId); - } - if (readyResult == ReadySignalResult::kNotReady) - { - session->releaseReservedRecvBuffers(); - TLLM_THROW("KV cache receive request %zu was rejected by the context peer", requestId); - } - if (readyResult == ReadySignalResult::kMixed) - { - if (common::getEnvDisaggEnableInflightCancel()) - { - session->poisonReservedRecvBuffers(); - } - else - { - session->releaseReservedRecvBuffers(); - } - TLLM_THROW("KV cache receive request %zu received inconsistent ready signals from its context peers", - requestId); - } - - phase = "transfer-completion-notification"; - TLLM_LOG_DEBUG("KV cache receive request %zu, context request %zu phase=%s begin.", requestId, - contextRequestId, phase); - receiveSync(*session); - TLLM_LOG_DEBUG( - "KV cache receive request %zu, context request %zu phase=%s end.", requestId, contextRequestId, phase); - llmRequest.setKvCacheTransferEnd(LlmRequest::getSteadyClockNow()); - } - catch (std::exception const& err) - { - if (common::getEnvDisaggEnableInflightCancel() && session.has_value()) - { - session->poisonReservedRecvBuffers(); - } - llmRequest.setKvCacheTransferEnd(LlmRequest::getSteadyClockNow()); - TLLM_LOG_ERROR("KV cache receive request %zu, context request %zu failed in phase=%s: %s", requestId, - contextRequestId, phase, err.what()); - throw; - } - catch (...) - { - if (common::getEnvDisaggEnableInflightCancel() && session.has_value()) - { - session->poisonReservedRecvBuffers(); - } - llmRequest.setKvCacheTransferEnd(LlmRequest::getSteadyClockNow()); - TLLM_LOG_ERROR( - "KV cache receive request %zu, context request %zu failed in phase=%s with an unknown " - "exception", - requestId, contextRequestId, phase); - throw; + TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), + "Start calling requestSync for request ID: %zu, context request ID: %zu.", llmRequest.mRequestId, + llmRequest.getContextPhaseParams().value().getReqId()); + llmRequest.setKvCacheTransferStart(std::chrono::steady_clock::now()); + TLLM_CUDA_CHECK(cudaSetDevice(mDeviceId)); + auto session = sendRequestInfo(llmRequest); + session.setTime(TransferSession::kTimeRequestInfo); + bool isReady = receiveReadySignal(session); + if (!isReady) + { + // Reuse the error state for the cancelled request. + llmRequest.setState(LlmRequestState::kDISAGG_TRANS_ERROR); + llmRequest.setKvCacheTransferEnd(std::chrono::steady_clock::now()); + return; } + receiveSync(session); + llmRequest.setKvCacheTransferEnd(std::chrono::steady_clock::now()); - TLLM_LOG_DEBUG("KV cache receive request %zu, context request %zu completed.", requestId, contextRequestId); - } - - void requestSync(LlmRequest& llmRequest) - { - requestSync(llmRequest, mTerminate); + TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), + "End calling requestSync for request ID: %zu, context request ID: %zu.", llmRequest.mRequestId, + llmRequest.getContextPhaseParams().value().getReqId()); } struct RequestAndPromise @@ -1575,20 +1131,16 @@ class CacheReceiver::Impl // protects worker-side dereferences and the promise itself from premature destruction. std::shared_ptr mRequest; std::unique_ptr> mPromise; - std::shared_ptr> mCancelFlag; RequestAndPromise() : mRequest(nullptr) , mPromise(nullptr) - , mCancelFlag(nullptr) { } - RequestAndPromise(std::shared_ptr request, std::unique_ptr>&& promise, - std::shared_ptr> cancelFlag) + RequestAndPromise(std::shared_ptr request, std::unique_ptr>&& promise) : mRequest(std::move(request)) , mPromise(std::move(promise)) - , mCancelFlag(std::move(cancelFlag)) { } @@ -1597,7 +1149,6 @@ class CacheReceiver::Impl RequestAndPromise(RequestAndPromise&& other) noexcept : mRequest(std::move(other.mRequest)) , mPromise(std::move(other.mPromise)) - , mCancelFlag(std::move(other.mCancelFlag)) { } @@ -1613,7 +1164,6 @@ class CacheReceiver::Impl mRequest = std::move(other.mRequest); mPromise = std::move(other.mPromise); - mCancelFlag = std::move(other.mCancelFlag); } return *this; } @@ -1657,9 +1207,7 @@ class CacheReceiver::Impl try { TLLM_CHECK_WITH_INFO(requestAndPromise.mRequest != nullptr, "requestAndPromise.mRequest is null"); - auto const& cancelFlag - = requestAndPromise.mCancelFlag != nullptr ? *requestAndPromise.mCancelFlag : mTerminate; - requestSync(*requestAndPromise.mRequest, cancelFlag); + requestSync(*requestAndPromise.mRequest); requestAndPromise.mPromise->set_value(); } catch (tensorrt_llm::common::RequestSpecificException const& err) @@ -1678,19 +1226,6 @@ class CacheReceiver::Impl requestAndPromise.mRequest->getContextPhaseParams().value().getReqId(), err.what()); requestAndPromise.mPromise->set_exception(std::current_exception()); } - catch (...) - { - TLLM_LOG_ERROR("Unknown exception in CacheReceiver request() loop"); - if (requestAndPromise.mPromise) - { - requestAndPromise.mPromise->set_exception(std::current_exception()); - } - } - if (common::getEnvDisaggEnableInflightCancel() && requestAndPromise.mRequest != nullptr) - { - std::lock_guard lg(mInFlightCancelMutex); - mInFlightCancelFlags.erase(requestAndPromise.mRequest->mRequestId); - } } } } @@ -1718,8 +1253,6 @@ class CacheReceiver::Impl std::ofstream mMeasuresFile; std::mutex mMeasuresFileMutex; std::atomic mTerminate{false}; - std::mutex mInFlightCancelMutex; - std::unordered_map>> mInFlightCancelFlags; }; void CacheSender::ImplDeleter::operator()(Impl* ptr) @@ -1762,9 +1295,7 @@ void CacheSender::sendSync(LlmRequest const& llmRequest) RequestInfo CacheSender::recvRequestInfo() { - auto requestInfo = mImpl->recvRequestInfo(); - TLLM_CHECK(requestInfo.has_value()); - return *requestInfo; + return mImpl->recvRequestInfo(); } bool CacheSender::cancelRequest(LlmRequest const& llmRequest) diff --git a/cpp/tensorrt_llm/batch_manager/dataTransceiver.h b/cpp/tensorrt_llm/batch_manager/dataTransceiver.h index 8e84a71556af..3362574da902 100644 --- a/cpp/tensorrt_llm/batch_manager/dataTransceiver.h +++ b/cpp/tensorrt_llm/batch_manager/dataTransceiver.h @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -22,7 +22,6 @@ #include #include -#include "tensorrt_llm/batch_manager/baseTransBuffer.h" #include "tensorrt_llm/batch_manager/cacheTransceiver.h" #include "tensorrt_llm/batch_manager/cacheTransferLayer.h" #include "tensorrt_llm/batch_manager/llmRequest.h" @@ -129,20 +128,6 @@ class TransferSession void appendMeasure(LlmRequest::TimePoint start, LlmRequest::TimePoint end, size_t size); - /// @brief Transfer ownership of pre-assigned receive buffers to this session. - void setReservedRecvBuffers(std::vector holders); - - [[nodiscard]] bool hasReservedRecvBuffer(BaseTransBufferManager const& manager) const noexcept; - - /// @brief Release one formatter's pre-assigned buffer after its receive and postprocessing complete. - bool releaseReservedRecvBuffer(BaseTransBufferManager const& manager) noexcept; - - /// @brief Release all pre-assigned receive buffers after the full receive pipeline completes. - void releaseReservedRecvBuffers() noexcept; - - /// @brief Fail closed when receive-buffer quiescence cannot be established. - void poisonReservedRecvBuffers() noexcept; - // TODO: 1. use global id instead of context request id; 2. export to llm metrics instead of file void exportMeasure(std::ofstream& outFile, bool isContext) const; @@ -175,7 +160,6 @@ class TransferSession runtime::BufferManager const* mBufferManager; LlmRequest const* mRequest; std::unique_ptr mTimes; - std::vector mReservedRecvBuffers; int32_t mIndexFromEnd{0}; BlockKey mLastBlockKey{}; }; diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp index fd5f8d4ddd75..0fb8af1527ae 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp @@ -17,14 +17,12 @@ #include "tensorrt_llm/batch_manager/kvCacheManager.h" -#include "tensorrt_llm/batch_manager/cacheTransBuffer.h" #include "tensorrt_llm/batch_manager/common.h" #include "tensorrt_llm/batch_manager/evictionPolicy.h" #include "tensorrt_llm/batch_manager/kvCacheTransferManager.h" #include "tensorrt_llm/batch_manager/radixBlockTree.h" #include "tensorrt_llm/common/assert.h" #include "tensorrt_llm/common/cudaUtils.h" -#include "tensorrt_llm/common/envUtils.h" #include "tensorrt_llm/common/logger.h" #include "tensorrt_llm/common/memoryUtils.h" #include "tensorrt_llm/executor/executor.h" @@ -522,14 +520,6 @@ bool KVCacheBlock::isLeaf() const return !mLookupNode || !mLookupNode->hasChildren(); } -bool KVCacheBlock::isDetached() const -{ - // A block is "detached" when it is not registered in the reuse lookup tree - // (mLookupNode == nullptr), i.e. it holds no state that future requests can reuse. - // This mirrors the reuse condition used by isShared(). - return mLookupNode == nullptr; -} - // This function calculates the number of block a layer should have, given // the total free memory and the window size of each layer. // For example, if we have 1 layer of window size 1024, and 2 layer of window @@ -1094,23 +1084,6 @@ void WindowBlockManager::allocatePools(bool useUvm) { constexpr nvinfer1::DataType kScaleDtypeNVFP4 = nvinfer1::DataType::kFP8; - bool const requestFabricMemory = tc::getEnvKVCachePoolUseFabricMemory(); - bool const fabricMemorySupported = FabricMemory::supportFabricMemory(); - if (requestFabricMemory && !fabricMemorySupported) - { - TLLM_LOG_WARNING( - "[%s] TRTLLM_KVCACHE_POOL_USE_FABRIC_MEMORY=1 was set but fabric memory is not supported on this " - "platform (FabricMemory::supportFabricMemory() returned false); falling back to standard GPU " - "allocation.", - mLogPrefix.c_str()); - } - bool const useFabricMemory = requestFabricMemory && fabricMemorySupported; - - if (useFabricMemory) - { - TLLM_LOG_INFO("[%s] KV cache pool using fabric memory for MNNVL support", mLogPrefix.c_str()); - } - // Allocate a memory pool backing the blocks for each numKvHeads // TODO(oargov): allocate pools in a single buffer and split it, to avoid fragmentation for (auto& pool : mPools) @@ -1143,27 +1116,9 @@ void WindowBlockManager::allocatePools(bool useUvm) cacheShape.d[2], cacheShape.d[3], pool.layerFirstLayout ? " (layer-first)" : ""); if (useUvm) - { pool.primaryPtr = BufferManager::managed(cacheShape, poolDtype); - } - else if (useFabricMemory) - { - auto const numElements = ITensor::volume(cacheShape); - auto const elementSize = tc::getDTypeSize(poolDtype); - auto const totalBytes = static_cast(numElements) * elementSize; - - // Record ownership before exposing the raw pointer: if FabricMemory's ctor throws nothing - // is wrapped; if ITensor::wrap throws afterwards, the unique_ptr in mFabricMemoryPools - // still owns and will free the allocation. - mFabricMemoryPools.reserve(mFabricMemoryPools.size() + 1); - mFabricMemoryPools.emplace_back(std::make_unique(totalBytes)); - pool.primaryPtr = ITensor::wrap(mFabricMemoryPools.back()->getPtr(), poolDtype, cacheShape, numElements); - } else - { pool.primaryPtr = mBufferManager.gpuSync(cacheShape, poolDtype); - } - if (mNumSecondaryBlocks > 0) { nvinfer1::Dims cacheShapeOffload = isRecurrentState() @@ -1186,12 +1141,6 @@ void BlockManager::releasePools() void WindowBlockManager::releasePools() { - if (mTransferManager) - { - mTransferManager->syncTransfers(); - } - mBufferManager.getStream().synchronize(); - for (auto& pool : mPools) { if (pool.primaryPtr) @@ -1203,8 +1152,7 @@ void WindowBlockManager::releasePools() pool.secondaryPtr->release(); } } - // Release fabric memory backing (must happen after ITensor release). - mFabricMemoryPools.clear(); + mBufferManager.getStream().synchronize(); mBufferManager.memoryPoolTrimTo(0); } @@ -3182,16 +3130,7 @@ std::optional WindowBlockManager::releaseBlocks( // mRefCount==0 and are silently ignored by EvictionPolicy::releaseBlock(). if (!block->hasRefs()) { - auto const isDetached = block->isDetached(); - if (isDetached) - { - // Detached blocks have no reusable hash chain. Drop the stale owning link before recycling the block; - // otherwise front-queue reuse can join completed request chains into an unbounded ownership chain. - block->setPrevBlockInSeq(nullptr); - } - // Send block to front of free queue if it has no reusable state, - // so detached blocks are evicted before blocks cached for reuse. - mEvictionPolicy->releaseBlock(block, /*toFront=*/isDetached); + mEvictionPolicy->releaseBlock(block); } } // Remove stored block ids in sequence @@ -3993,28 +3932,6 @@ tle::RetentionPriority KVCacheManager::getPriorityByBlockId(KVCacheBlock::IdType } } -std::vector KVCacheManager::getMemoryPoolBlockIndicesByBlockIds( - std::vector const& blockIds, SizeType32 windowSize) const -{ - std::vector indices; - indices.reserve(blockIds.size()); - for (auto const blockId : blockIds) - { - BlockPtr const& block = mBlockManager.getBlockById(blockId, windowSize); - TLLM_CHECK_WITH_INFO(block != nullptr, "Block not found (blockId=%d, windowSize=%d)", blockId, windowSize); - // Invariant, not a recoverable condition: blocks referenced here are held by a live - // request, allocation onboards offloaded blocks, and offload only selects free blocks. - // A secondary block therefore indicates a block-lifetime bug, and the returned index - // (pool flag stripped) would silently alias a primary slot. - TLLM_CHECK_WITH_INFO(block->isPrimary(), - "Block is not in the primary pool (blockId=%d, windowSize=%d); the returned index is only " - "valid for primary-pool blocks", - blockId, windowSize); - indices.push_back(block->getMemoryPoolBlockIndex()); - } - return indices; -} - SizeType32 KVCacheManager::copyBlockOffsets(ITensor& output, SizeType32 outputSlotOffset, RequestIdType requestId) const { auto const& sequence = getSequence(requestId); diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp index c28d1e476137..edf3fe549bf6 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp @@ -97,12 +97,6 @@ tr::ITensor::SharedPtr KVCacheTransferManager::computeBlockPointer( return blockTensor; } -tk::KVCacheIndex::UnderlyingType KVCacheTransferManager::getPendingTransferIndex(BlockPtr const& block) -{ - auto const blockOffset = block->getMemoryPoolBlockIndex(); - return block->isPrimary() ? blockOffset : blockOffset | tk::KVCacheIndex::kSecondaryPoolFlag; -} - void KVCacheTransferManager::copyBlock(BlockPtr const& src, BlockPtr const& dst, std::vector const& pools, bool isOffload, int numTokensToCopy, executor::KvCacheTransferMode mode, std::string const& directory) @@ -258,9 +252,10 @@ void KVCacheTransferManager::copyBlock(BlockPtr const& src, BlockPtr const& dst, } // -// Note about recording events to wait for cudaMemcpyAsync calls between blocks: -// The memory copy involves raw memory blocks, which are identified by the pool-qualified -// memory pool block index. Using getBlockId() when recording events is wrong. +// Note about recording events to wait for cudaMempyAsync calls between blocks: +// The memory copy involves raw memory blocks, which are pointed to by the +// memory pool block index. When recording events, you must use getMemoryPoolBlockIndex() +// as the raw memory block identifier. Using getBlockId() when recording events is wrong. // getBlockId() returns the logical block id, which has nothing to do with the raw memory // block pointers involved in a cudaMemcpy. // @@ -294,25 +289,22 @@ void KVCacheTransferManager::onboard(BlockPtr const& offloadedBlock, BlockPtr co std::vector const& pools, int numTokensToCopy, executor::KvCacheTransferMode mode, std::string const& directory) { - auto const offloadedBlockIndex = getPendingTransferIndex(offloadedBlock); - auto const blockIndex = getPendingTransferIndex(block); - // Wait for any pending writes before reading from offloadedBlock - auto offloadedBlockPendingWriteItr = mPendingWrites.find(offloadedBlockIndex); + auto offloadedBlockPendingWriteItr = mPendingWrites.find(offloadedBlock->getMemoryPoolBlockIndex()); if (offloadedBlockPendingWriteItr != mPendingWrites.end()) { mOnboardManager.getStream().wait(offloadedBlockPendingWriteItr->second); // Don't erase, we are not changing state of offloadedBlock } // Wait for any pending reads before overwriting block - auto blockPendingReadItr = mPendingReads.find(blockIndex); + auto blockPendingReadItr = mPendingReads.find(block->getMemoryPoolBlockIndex()); if (blockPendingReadItr != mPendingReads.end()) { mOnboardManager.getStream().wait(blockPendingReadItr->second); mPendingReads.erase(blockPendingReadItr); } // Wait for any pending writes before overwriting block - auto blockPendingWriteItr = mPendingWrites.find(blockIndex); + auto blockPendingWriteItr = mPendingWrites.find(block->getMemoryPoolBlockIndex()); if (blockPendingWriteItr != mPendingWrites.end()) { mOnboardManager.getStream().wait(blockPendingWriteItr->second); @@ -338,36 +330,33 @@ void KVCacheTransferManager::onboard(BlockPtr const& offloadedBlock, BlockPtr co } // Record new pending read from offloadedBlock - mPendingReads[offloadedBlockIndex] = tr::CudaEvent(); - mOnboardManager.getStream().record(mPendingReads[offloadedBlockIndex]); + mPendingReads[offloadedBlock->getMemoryPoolBlockIndex()] = tr::CudaEvent(); + mOnboardManager.getStream().record(mPendingReads[offloadedBlock->getMemoryPoolBlockIndex()]); // Record new pending write to block - mPendingWrites[blockIndex] = tr::CudaEvent(); - mOnboardManager.getStream().record(mPendingWrites[blockIndex]); + mPendingWrites[block->getMemoryPoolBlockIndex()] = tr::CudaEvent(); + mOnboardManager.getStream().record(mPendingWrites[block->getMemoryPoolBlockIndex()]); } void KVCacheTransferManager::offload(BlockPtr const& block, BlockPtr const& offloadBlock, std::vector const& pools, int numTokensToCopy, executor::KvCacheTransferMode mode, std::string const& directory) { - auto const blockIndex = getPendingTransferIndex(block); - auto const offloadBlockIndex = getPendingTransferIndex(offloadBlock); - // Wait for any pending writes before reading from block - auto blockPendingWriteItr = mPendingWrites.find(blockIndex); + auto blockPendingWriteItr = mPendingWrites.find(block->getMemoryPoolBlockIndex()); if (blockPendingWriteItr != mPendingWrites.end()) { mOffloadManager.getStream().wait(blockPendingWriteItr->second); // Don't erase, we are not changing state of block } // Wait for any pending reads before overwriting offloadBlock - auto offloadBlockPendingReadItr = mPendingReads.find(offloadBlockIndex); + auto offloadBlockPendingReadItr = mPendingReads.find(offloadBlock->getMemoryPoolBlockIndex()); if (offloadBlockPendingReadItr != mPendingReads.end()) { mOffloadManager.getStream().wait(offloadBlockPendingReadItr->second); mPendingReads.erase(offloadBlockPendingReadItr); } // Wait for any pending writes before overwriting offloadBlock - auto offloadBlockPendingWriteItr = mPendingWrites.find(offloadBlockIndex); + auto offloadBlockPendingWriteItr = mPendingWrites.find(offloadBlock->getMemoryPoolBlockIndex()); if (offloadBlockPendingWriteItr != mPendingWrites.end()) { mOffloadManager.getStream().wait(offloadBlockPendingWriteItr->second); @@ -384,11 +373,11 @@ void KVCacheTransferManager::offload(BlockPtr const& block, BlockPtr const& offl } // Record new pending read from block - mPendingReads[blockIndex] = tr::CudaEvent(); - mOffloadManager.getStream().record(mPendingReads[blockIndex]); + mPendingReads[block->getMemoryPoolBlockIndex()] = tr::CudaEvent(); + mOffloadManager.getStream().record(mPendingReads[block->getMemoryPoolBlockIndex()]); // Record new pending write to offloadBlock - mPendingWrites[offloadBlockIndex] = tr::CudaEvent(); - mOffloadManager.getStream().record(mPendingWrites[offloadBlockIndex]); + mPendingWrites[offloadBlock->getMemoryPoolBlockIndex()] = tr::CudaEvent(); + mOffloadManager.getStream().record(mPendingWrites[offloadBlock->getMemoryPoolBlockIndex()]); } void KVCacheTransferManager::syncWithBufferManager() @@ -401,7 +390,7 @@ void KVCacheTransferManager::syncWithBufferManager() mBufferManager.getStream().record(readyForOnboardEvent); mOnboardManager.getStream().wait(readyForOnboardEvent); - // Once we synchronize, clear our list of pending transfers. + // Once we synchronize, clear our list of pending thransfers. mPendingReads.clear(); mPendingWrites.clear(); } @@ -416,7 +405,7 @@ void KVCacheTransferManager::syncTransfers() mOnboardManager.getStream().record(onboardEvent); mBufferManager.getStream().wait(onboardEvent); - // Once we synchronize, clear our list of pending transfers. + // Once we synchronize, clear our list of pending thransfers. mPendingReads.clear(); mPendingWrites.clear(); } diff --git a/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp b/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp index 24b7182fe628..812323e92b81 100644 --- a/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp +++ b/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -253,9 +253,7 @@ void MLACacheFormatter::format(tensorrt_llm::batch_manager::TransferSession& ses return bufferSizeForTarget; }; auto bufferEleSizes = getBufferSizeForTarget(); - auto const* sendCancelFlag - = common::getEnvDisaggEnableInflightCancel() ? &session.getDataContext().getTransferTerminate() : nullptr; - auto cacheBufferId = mCacheTransBufferManagers[transferIndexerKCache]->assignBufferIndexForSend(sendCancelFlag); + auto cacheBufferId = mCacheTransBufferManagers[transferIndexerKCache]->assignBufferIndexForSend(); BufferIndexHolder sendHolder( *mCacheTransBufferManagers[transferIndexerKCache], cacheBufferId, /*isRecv=*/false); auto result = mCacheTransBufferManagers[transferIndexerKCache]->getOrAllocateSendBuffers( @@ -306,13 +304,6 @@ void MLACacheFormatter::format(tensorrt_llm::batch_manager::TransferSession& ses auto cpDomainIdx = processIdx / connectionsPerCPDomain; auto ppDomainIdx = (processIdx % connectionsPerCPDomain) % pPDomainSize; auto cacheIdx = cpDomainIdx * pPDomainSize + ppDomainIdx; - // Helix: skip CP ranks that own no blocks for this sequence (num_total_blocks < cp_size). - // The matching gen rank skips its receive, so no 0-byte transfer is posted on either side. - auto const& splitCache = outputSplitCaches.at(cacheIdx); - if (splitCache == nullptr || splitCache->getSizeInBytes() == 0) - { - return; - } if (cacheIdx < bufferCoverTargetNum) { size_t size = outputSplitCaches.at(cacheIdx)->getSizeInBytes(); @@ -348,64 +339,48 @@ void MLACacheFormatter::format(tensorrt_llm::batch_manager::TransferSession& ses session.appendMeasure(startTime, endTime, outputSplitCaches.at(cacheIdx)->getSizeInBytes()); }; - if (sendCancelFlag != nullptr && sendCancelFlag->load(std::memory_order_relaxed)) - { - TLLM_THROW("MLA cache transfer cancelled before NIXL submission"); - } - - try + if (pickUpConnections.size() > 1) { - if (pickUpConnections.size() > 1) + if (!common::getEnvEnableReceiveKVCacheParallel()) { - if (!common::getEnvEnableReceiveKVCacheParallel()) + TLLM_LOG_DEBUG("Disable parallel receiving of the KV cache."); + for (size_t i = 0; i < pickUpConnections.size(); i++) { - TLLM_LOG_DEBUG("Disable parallel receiving of the KV cache."); - for (size_t i = 0; i < pickUpConnections.size(); i++) - { - sendBufferFun(deviceId, pickUpConnections[i]); - } + sendBufferFun(deviceId, pickUpConnections[i]); } - else - { - // concurrency num - auto concurrencyNum - = std::min(std::max(static_cast(1), bufferCoverTargetNum), pPDomainSize * cPDomainSize); + } + else + { + // concurrency num + auto concurrencyNum + = std::min(std::max(static_cast(1), bufferCoverTargetNum), pPDomainSize * cPDomainSize); - auto remainSendNum = pickUpConnections.size(); + auto remainSendNum = pickUpConnections.size(); - while (remainSendNum > 0) + while (remainSendNum > 0) + { + auto sendConcurrencyNum = std::min(remainSendNum, concurrencyNum); + std::vector> futures; + futures.reserve(sendConcurrencyNum); + for (size_t i = 0; i < sendConcurrencyNum; i++) { - auto sendConcurrencyNum = std::min(remainSendNum, concurrencyNum); - std::vector> futures; - futures.reserve(sendConcurrencyNum); - for (size_t i = 0; i < sendConcurrencyNum; i++) - { - size_t idx = i + (pickUpConnections.size() - remainSendNum); - size_t connIdx = pickUpConnections[idx]; - TLLM_CHECK(idx < pickUpConnections.size()); - TLLM_CHECK(connIdx < session.getConnections().size()); - futures.push_back(std::async(std::launch::async, sendBufferFun, deviceId, connIdx)); - } - for (auto& future : futures) - { - future.get(); - } - remainSendNum -= sendConcurrencyNum; + size_t idx = i + (pickUpConnections.size() - remainSendNum); + size_t connIdx = pickUpConnections[idx]; + TLLM_CHECK(idx < pickUpConnections.size()); + TLLM_CHECK(connIdx < session.getConnections().size()); + futures.push_back(std::async(std::launch::async, sendBufferFun, deviceId, connIdx)); + } + for (auto& future : futures) + { + future.get(); } + remainSendNum -= sendConcurrencyNum; } } - else - { - sendBufferFun(deviceId, pickUpConnections[0]); - } } - catch (...) + else { - if (agentConnection != nullptr && common::getEnvDisaggEnableInflightCancel()) - { - sendHolder.poison(); - } - throw; + sendBufferFun(deviceId, pickUpConnections[0]); } sendHolder.release(); } @@ -483,13 +458,6 @@ void MLACacheFormatter::unformat(tensorrt_llm::batch_manager::TransferSession& s } } - // Helix: an "empty" CP rank owns no KV blocks for this sequence (num_total_blocks < cp_size). - // There is nothing to receive; the sender (context, CP=1) skips the matching 0-byte transfer. - if (blockNum == 0) - { - continue; - } - int deviceId = bufferManager.getStream().getDevice(); std::optional cacheBufferId = std::nullopt; @@ -525,18 +493,13 @@ void MLACacheFormatter::unformat(tensorrt_llm::batch_manager::TransferSession& s if (preAssignedId.has_value()) { cacheBufferId = static_cast(*preAssignedId); - if (!session.hasReservedRecvBuffer(*mCacheTransBufferManagers[transferIndexerKCache])) - { - recvHolder = BufferIndexHolder( - *mCacheTransBufferManagers[transferIndexerKCache], cacheBufferId, /*isRecv=*/true); - } } else { cacheBufferId = mCacheTransBufferManagers[transferIndexerKCache]->assignBufferIndexForRecv(); - recvHolder = BufferIndexHolder( - *mCacheTransBufferManagers[transferIndexerKCache], cacheBufferId, /*isRecv=*/true); } + recvHolder + = BufferIndexHolder(*mCacheTransBufferManagers[transferIndexerKCache], cacheBufferId, /*isRecv=*/true); auto targetNum = pickUpConnections.size(); @@ -684,7 +647,6 @@ void MLACacheFormatter::unformat(tensorrt_llm::batch_manager::TransferSession& s bufferManager.getStream().synchronize(); } - (void) session.releaseReservedRecvBuffer(*mCacheTransBufferManagers[transferIndexerKCache]); recvHolder.release(); } session.setTime(TransferSession::kTimePostprocess); diff --git a/cpp/tensorrt_llm/batch_manager/rnnCacheFormatter.cpp b/cpp/tensorrt_llm/batch_manager/rnnCacheFormatter.cpp index c58733f25885..05742af17a28 100644 --- a/cpp/tensorrt_llm/batch_manager/rnnCacheFormatter.cpp +++ b/cpp/tensorrt_llm/batch_manager/rnnCacheFormatter.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -20,9 +20,9 @@ #include "tensorrt_llm/batch_manager/dataTransceiver.h" #include "tensorrt_llm/batch_manager/kvCacheManager.h" #include "tensorrt_llm/batch_manager/kvCacheUtils.h" +#include "tensorrt_llm/batch_manager/rnnStateManager.h" #include "tensorrt_llm/common/assert.h" #include "tensorrt_llm/common/dataType.h" -#include "tensorrt_llm/common/envUtils.h" #include "tensorrt_llm/common/logger.h" #include "tensorrt_llm/common/nvtxUtils.h" #include "tensorrt_llm/executor/cache_transmission/agent_utils/connection.h" @@ -33,6 +33,15 @@ namespace tensorrt_llm::batch_manager { using CacheState = executor::kv_cache::CacheState; +RnnCacheFormatter::RnnCacheFormatter(rnn_state_manager::RnnStateManager* rnnStateManager, + rnn_state_manager::RnnCacheTransBufferManager* rnnCacheTransBufferManager) + : mRnnStateManager{rnnStateManager} + , mRnnCacheTransBufferManager{rnnCacheTransBufferManager} +{ + TLLM_CHECK(mRnnStateManager != nullptr); + TLLM_CHECK(mRnnCacheTransBufferManager != nullptr); +} + RnnCacheFormatter::RnnCacheFormatter(kv_cache_manager::BaseKVCacheManager* kvCacheManager, rnn_state_manager::RnnCacheTransBufferManager* rnnCacheTransBufferManager) : mRnnCacheTransBufferManager{rnnCacheTransBufferManager} @@ -43,6 +52,470 @@ RnnCacheFormatter::RnnCacheFormatter(kv_cache_manager::BaseKVCacheManager* kvCac } void RnnCacheFormatter::format(TransferSession& session) +{ + if (isUnifiedPoolMode()) + { + formatUnifiedPoolMode(session); + } + else + { + formatSlotMode(session); + } +} + +void RnnCacheFormatter::unformat(TransferSession& session) +{ + if (isUnifiedPoolMode()) + { + unformatUnifiedPoolMode(session); + } + else + { + unformatSlotMode(session); + } +} + +void RnnCacheFormatter::formatSlotMode(TransferSession& session) +{ + NVTX3_SCOPED_RANGE(RnnCacheFormatter_format); + session.setTime(TransferSession::kTimeFormatter); + + auto const& llmRequest = session.getLlmRequest(); + TLLM_LOG_DEBUG( + mpi::MpiComm::world().getRank(), "Start sending RNN state for request ID: %ld.", llmRequest.mRequestId); + TLLM_CHECK_WITH_INFO(llmRequest.mSamplingConfig.beamWidth == 1, "Currently, only beam width 1 is supported."); + + auto const& connections = session.getConnections(); + auto const& selfConfig = session.getSelfState().getCacheState().value(); + auto const& destConfig = session.getOtherState().getCacheState().value(); + auto const selfIdx = session.getSelfState().getCommState().value().getSelfIdx(); + auto& bufferManager = session.getBufferManager(); + + auto targetInfo = executor::kv_cache::targetIRanksForRnn(destConfig, selfConfig, selfIdx); + if (!cache_formatter_utils::needSendCache(selfConfig, destConfig, selfIdx, targetInfo)) + { + return; + } + + auto pickUpConnections = cache_formatter_utils::pickSendConnections( + connections.size(), selfConfig, selfIdx, destConfig, session.getCounterPartRanks(), targetInfo); + auto const targetNum = pickUpConnections.size(); + if (targetNum == 0) + { + TLLM_LOG_DEBUG("No targets to send RNN state to for request ID: %ld", llmRequest.mRequestId); + return; + } + + auto const slotIdx = mRnnStateManager->getCacheIndex(llmRequest.mRequestId); + int deviceId; + TLLM_CUDA_CHECK(cudaGetDevice(&deviceId)); + + auto const& selfParallel = selfConfig.getParallelConfig(); + auto const selfTPNum = selfParallel.mTensorParallelism; + auto const selfPPRank = selfIdx / selfTPNum; + auto const& selfLayersPerPP = selfConfig.getRnnCacheState().mLayerNumPerPP; + SizeType32 const numLocalLayers = selfLayersPerPP[selfPPRank]; + + if (common::getEnvTryZCopyForKVCacheTransfer() && destConfig == selfConfig) + { + TLLM_LOG_DEBUG("Try using zero-copy for the RNN cache."); + NVTX3_SCOPED_RANGE(RnnZeroCopySend); + + TLLM_CHECK(pickUpConnections.size() == 1); + + TLLM_CUDA_CHECK(cudaSetDevice(deviceId)); + for (size_t i = 0; i < pickUpConnections.size(); i++) + { + for (SizeType32 layer = 0; layer < numLocalLayers; layer++) + { + + // Get conv state for this layer: shape is [maxBatchSize, convDim, dConv-1] + auto convState = mRnnStateManager->getConvStates(mRnnStateManager->getGlobalLayerNum(layer)); + // Slice out the specific slot: shape becomes [convDim, dConv-1] + auto slotConv = runtime::ITensor::slice(convState, slotIdx, 1); + slotConv->squeeze(0); + + // Receive conv state + // llmRequest.updateKvCacheSize(slotConv->getSizeInBytes()); + session.send(pickUpConnections[i], slotConv->data(), slotConv->getSizeInBytes()); + + // Get SSM state for this layer: shape is [maxBatchSize, numHeads, headDim, dState] + auto ssmState = mRnnStateManager->getSsmStates(mRnnStateManager->getGlobalLayerNum(layer)); + // Slice out the specific slot: shape becomes [numHeads, headDim, dState] + auto slotSsm = runtime::ITensor::slice(ssmState, slotIdx, 1); + slotSsm->squeeze(0); + + // Receive SSM state + // llmRequest.updateKvCacheSize(slotSsm->getSizeInBytes()); + session.send(pickUpConnections[i], slotSsm->data(), slotSsm->getSizeInBytes()); + } + } + TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), "End the sending of RNN cache for the request ID: %ld.", + llmRequest.mRequestId); + + return; + } + + // Calculate buffer sizes for each target + // Each target gets: conv states + ssm states for overlapping layers + auto const& modelConfig = selfConfig.getRnnModelConfig(); + auto const maxBatchSize = mRnnStateManager->getMaxBatchSize(); + int const selfTPSizePerDPGroup = selfConfig.getParallelConfig().mEnableAttentionDP + ? selfTPNum / selfConfig.getParallelConfig().mDPsize + : selfTPNum; + SizeType32 convDimLocal = modelConfig.mConvDimSize / selfTPSizePerDPGroup; + SizeType32 numHeadsLocal = modelConfig.mNumHeads / selfTPSizePerDPGroup; + + size_t convBytesPerLayer + = convDimLocal * (modelConfig.mDConv - 1) * common::getDTypeSize(selfConfig.getConvStateDataType()); + convBytesPerLayer = (convBytesPerLayer + 15) & ~static_cast(15); + size_t ssmBytesPerLayer = numHeadsLocal * modelConfig.mHeadDim * modelConfig.mDState + * common::getDTypeSize(selfConfig.getSsmStateDataType()); + + int peerDuplicateHeadFactor = targetInfo.mPeerDupHeadFactor; + auto bufferTargetNum = targetNum / peerDuplicateHeadFactor; + + std::vector bufferSizesPerTarget(targetNum, 0); + + for (size_t i = 0; i < targetNum; i++) + { + SizeType32 layersForTarget = targetInfo.getPeerPPDomainLayerNum(static_cast(i)); + bufferSizesPerTarget[i] = layersForTarget * (convBytesPerLayer + ssmBytesPerLayer) * peerDuplicateHeadFactor + / targetInfo.mDomainTPSize; + } + + auto cacheBufferId = mRnnCacheTransBufferManager->assignBufferIndexForSend(); + auto allocationResult = mRnnCacheTransBufferManager->getOrAllocateSendBuffers( + cacheBufferId, static_cast(bufferTargetNum), bufferSizesPerTarget, bufferManager); + auto& outputBuffers = std::get<0>(allocationResult); + auto& bufferCoverTargetNum = std::get<1>(allocationResult); + auto& onlyUseDynamicBuffer = std::get<2>(allocationResult); + + TLLM_CHECK(cacheBufferId.has_value() || onlyUseDynamicBuffer); + + auto const* agentConnection + = dynamic_cast(connections[pickUpConnections[0]]); + if (agentConnection != nullptr) + { + TLLM_CHECK_WITH_INFO(bufferCoverTargetNum == bufferTargetNum, "Agent needs all RNN send buffers pre-allocated"); + TLLM_CHECK(onlyUseDynamicBuffer == false); + } + + std::vector inputConvBlocks; + std::vector inputSsmBlocks; + + auto convStates = mRnnStateManager->getConvStates(); // [numLocalLayers, maxBatchSize, convDim, dConv-1] + auto ssmStates = mRnnStateManager->getSsmStates(); // [numLocalLayers, maxBatchSize, numHeads, headDim, dState] + + inputConvBlocks.push_back(convStates); + inputSsmBlocks.push_back(ssmStates); + + tensorrt_llm::executor::rnn_cache::splitRnnConvStateDispatch( + inputConvBlocks, outputBuffers, slotIdx, maxBatchSize, destConfig, selfConfig, selfIdx, bufferManager); + + // Conv and SSM use same output buffer. So need to track convBytesPerLayer to compute offset. + tensorrt_llm::executor::rnn_cache::splitRnnSsmStateDispatch(inputSsmBlocks, outputBuffers, slotIdx, maxBatchSize, + convBytesPerLayer, destConfig, selfConfig, selfIdx, bufferManager); + + bufferManager.getStream().synchronize(); + session.setTime(TransferSession::kTimePreprocess); + + auto preAllocSendBuffer = mRnnCacheTransBufferManager->getSendBuffer(cacheBufferId); + + sendAllBuffers(session, deviceId, outputBuffers, bufferCoverTargetNum, preAllocSendBuffer, bufferManager, + targetInfo, pickUpConnections); + + session.setTime(TransferSession::kTimeTransmissions); + + mRnnCacheTransBufferManager->freeBufferIndexForSend(cacheBufferId); + session.setTime(TransferSession::kTimePostprocess); + + TLLM_LOG_DEBUG( + mpi::MpiComm::world().getRank(), "End sending RNN state for request ID: %ld.", llmRequest.mRequestId); +} + +void RnnCacheFormatter::unformatSlotMode(TransferSession& session) +{ + NVTX3_SCOPED_RANGE(RnnCacheFormatter_unformat); + session.setTime(TransferSession::kTimeFormatter); + + auto& llmRequest = session.getLlmRequest(); + TLLM_LOG_DEBUG( + mpi::MpiComm::world().getRank(), "Start receiving RNN state for request ID: %ld.", llmRequest.mRequestId); + TLLM_CHECK_WITH_INFO(llmRequest.mSamplingConfig.beamWidth == 1, "Currently, only beam width 1 is supported."); + + auto const& connections = session.getConnections(); + auto const& selfConfig = session.getSelfState().getCacheState().value(); + auto const& destConfig = session.getOtherState().getCacheState().value(); + auto const selfIdx = session.getSelfState().getCommState().value().getSelfIdx(); + auto& bufferManager = session.getBufferManager(); + + auto sourceInfo = executor::kv_cache::targetIRanksForRnn(destConfig, selfConfig, selfIdx); + int deviceId; + TLLM_CUDA_CHECK(cudaGetDevice(&deviceId)); + + auto pickRecvConnResult = cache_formatter_utils::pickRecvConnections( + connections.size(), selfConfig, selfIdx, destConfig, session.getCounterPartRanks(), sourceInfo); + auto pickUpConnections = std::get<0>(pickRecvConnResult); + auto localRankIndices = std::get<1>(pickRecvConnResult); + auto const sourceNum = pickUpConnections.size(); + + if (sourceNum == 0) + { + TLLM_LOG_DEBUG("No sources to receive RNN state from for request ID: %ld", llmRequest.mRequestId); + return; + } + + if (common::getEnvDisaggLayerwise()) + { + TLLM_LOG_ERROR("Layer-wise RNN cache transfer is not supported yet"); + return; + } + + // Since allocation happens earlier + auto const slotIdx = mRnnStateManager->getCacheIndex(llmRequest.mRequestId); + + auto const& selfParallel = selfConfig.getParallelConfig(); + auto const selfTPNum = selfParallel.mTensorParallelism; + auto const selfPPRank = selfIdx / selfTPNum; + auto const& selfLayersPerPP = selfConfig.getRnnCacheState().mLayerNumPerPP; + SizeType32 const numLocalLayers = selfLayersPerPP[selfPPRank]; + + if (common::getEnvTryZCopyForKVCacheTransfer() && destConfig == selfConfig) + { + TLLM_LOG_DEBUG("try zcopy for RNN cache"); + NVTX3_SCOPED_RANGE(RnnZeroCopyRecv); + + TLLM_CHECK(sourceNum == 1); + + TLLM_CUDA_CHECK(cudaSetDevice(deviceId)); + for (size_t i = 0; i < sourceNum; i++) + { + for (SizeType32 layer = 0; layer < numLocalLayers; layer++) + { + + // Get conv state for this layer: shape is [maxBatchSize, convDim, dConv-1] + auto convState = mRnnStateManager->getConvStates(mRnnStateManager->getGlobalLayerNum(layer)); + // Slice out the specific slot: shape becomes [convDim, dConv-1] + auto slotConv = runtime::ITensor::slice(convState, slotIdx, 1); + slotConv->squeeze(0); + + // Send conv state + // llmRequest.updateKvCacheSize(slotConv->getSizeInBytes()); + session.recv(pickUpConnections[i], slotConv->data(), slotConv->getSizeInBytes()); + + // Get SSM state for this layer: shape is [maxBatchSize, numHeads, headDim, dState] + auto ssmState = mRnnStateManager->getSsmStates(mRnnStateManager->getGlobalLayerNum(layer)); + // Slice out the specific slot: shape becomes [numHeads, headDim, dState] + auto slotSsm = runtime::ITensor::slice(ssmState, slotIdx, 1); + slotSsm->squeeze(0); + + // Send SSM state + // llmRequest.updateKvCacheSize(slotSsm->getSizeInBytes()); + session.recv(pickUpConnections[i], slotSsm->data(), slotSsm->getSizeInBytes()); + } + } + TLLM_LOG_DEBUG( + mpi::MpiComm::world().getRank(), "End receiving RNN cache for request ID: %ld.", llmRequest.mRequestId); + return; + } + + // Calculate buffer sizes + auto const& modelConfig = selfConfig.getRnnModelConfig(); + int const selfTPSizePerDPGroup = selfParallel.mEnableAttentionDP ? selfTPNum / selfParallel.mDPsize : selfTPNum; + SizeType32 selfConvDimLocal = modelConfig.mConvDimSize / selfTPSizePerDPGroup; + int const selfNumHeadsLocal = modelConfig.mNumHeads / selfTPSizePerDPGroup; + + size_t convBytesPerLayer + = selfConvDimLocal * (modelConfig.mDConv - 1) * common::getDTypeSize(selfConfig.getConvStateDataType()); + convBytesPerLayer = (convBytesPerLayer + 15) & ~static_cast(15); + size_t ssmBytesPerLayer = selfNumHeadsLocal * modelConfig.mHeadDim * modelConfig.mDState + * common::getDTypeSize(selfConfig.getSsmStateDataType()); + + std::vector bufferSizesPerSource(sourceNum, 0); + size_t validTpSources = sourceNum / sourceInfo.mDomainPPSize; + + // Compute source conv bytes for SSM offset + size_t sourceConvBytesPerLayer = convBytesPerLayer / validTpSources; + + for (size_t i = 0; i < sourceNum; i++) + { + SizeType32 layersFromSource = sourceInfo.getPeerPPDomainLayerNum(static_cast(localRankIndices[i])); + bufferSizesPerSource[i] = layersFromSource * (convBytesPerLayer + ssmBytesPerLayer) / validTpSources; + } + + // Allocate receive buffers + size_t remainNoCoverSourceNum = 0; + size_t bufferCoverSourceNum = 0; + std::optional cacheBufferId = std::nullopt; + + auto preAssignedRnnId + = connections[pickUpConnections[0]]->getPreAssignedBufferId(static_cast(BufferKind::kRNN)); + if (preAssignedRnnId.has_value()) + { + cacheBufferId = static_cast(*preAssignedRnnId); + } + else + { + cacheBufferId = mRnnCacheTransBufferManager->assignBufferIndexForRecv(); + } + + auto allocationResult = mRnnCacheTransBufferManager->getOrAllocateRecvBuffers( + cacheBufferId, static_cast(sourceNum), bufferSizesPerSource, bufferManager); + auto& recvBuffers = std::get<0>(allocationResult); + auto& bufferCoverSourceNumTmp = std::get<1>(allocationResult); + auto& onlyUseDynamicBuffer = std::get<2>(allocationResult); + + TLLM_CHECK(cacheBufferId.has_value() || onlyUseDynamicBuffer); + + if (preAssignedRnnId.has_value()) + { + TLLM_CHECK_WITH_INFO(bufferCoverSourceNumTmp == sourceNum, "Agent needs all RNN recv buffers pre-allocated"); + TLLM_CHECK(onlyUseDynamicBuffer == false); + } + + bufferCoverSourceNum = bufferCoverSourceNumTmp; + remainNoCoverSourceNum = sourceNum > bufferCoverSourceNum ? sourceNum - bufferCoverSourceNum : 0; + + bufferManager.getStream().synchronize(); + session.setTime(TransferSession::kTimePreprocess); + + // Get pre-allocated buffer for chunked receive + runtime::ITensor::SharedPtr preAllocRecvBuffer = nullptr; + if (cacheBufferId.has_value()) + { + preAllocRecvBuffer = mRnnCacheTransBufferManager->getRecvBuffer(cacheBufferId); + TLLM_CHECK(preAllocRecvBuffer != nullptr); + } + + auto recvBufferFun = [&](int devId, size_t srcIdx) + { + NVTX3_SCOPED_RANGE(recvBufferFun); + TLLM_CUDA_CHECK(cudaSetDevice(devId)); + TLLM_CHECK(recvBuffers.size() > srcIdx); + auto startTime = LlmRequest::getSteadyClockNow(); + size_t size = 0; + + if (srcIdx >= remainNoCoverSourceNum) + { + // Fast path: buffer is pre-allocated, receive directly + auto& buffer = recvBuffers[srcIdx]; + size = buffer->getSizeInBytes(); + TLLM_LOG_DEBUG( + mpi::MpiComm::world().getRank(), " start recv srcIdx: %lu size:%lu", srcIdx, buffer->getSizeInBytes()); + session.recv(pickUpConnections[srcIdx], buffer->data(), buffer->getSizeInBytes()); + TLLM_LOG_DEBUG( + mpi::MpiComm::world().getRank(), " recv srcIdx: %lu size:%lu", srcIdx, buffer->getSizeInBytes()); + } + else + { + // Slow path: chunked receive for buffers that couldn't be pre-allocated + auto recvBufferIdx = bufferCoverSourceNum == 0 ? 0 : srcIdx % bufferCoverSourceNum + remainNoCoverSourceNum; + auto recvBufferUsed = bufferCoverSourceNum == 0 ? preAllocRecvBuffer : recvBuffers[recvBufferIdx]; + + size_t remainRecvSize = recvBuffers[srcIdx]->getSize(); + size_t needRecvSize = recvBuffers[srcIdx]->getSize(); + + while (remainRecvSize > 0) + { + TLLM_CHECK(recvBufferUsed != nullptr); + auto recvBufferEleSize = recvBufferUsed->getSize(); + auto recvSize = std::min(remainRecvSize, recvBufferEleSize); + auto recvSlice = runtime::ITensor::slice(recvBufferUsed, 0, recvSize); + auto copySlice = runtime::ITensor::slice(recvBuffers[srcIdx], needRecvSize - remainRecvSize, recvSize); + size += recvSlice->getSizeInBytes(); + session.recv(pickUpConnections[srcIdx], recvSlice->data(), recvSlice->getSizeInBytes()); + // Use cudaMemcpyAsync since we're copying bytes + TLLM_CUDA_CHECK(cudaMemcpyAsync(copySlice->data(), recvSlice->data(), recvSlice->getSizeInBytes(), + cudaMemcpyDeviceToDevice, bufferManager.getStream().get())); + bufferManager.getStream().synchronize(); + remainRecvSize -= recvSize; + } + } + + auto endTime = LlmRequest::getSteadyClockNow(); + session.appendMeasure(startTime, endTime, size); + }; + + // Dispatch receives (sequential or parallel based on env var) + if (sourceNum > 1) + { + if (!common::getEnvEnableReceiveKVCacheParallel()) + { + TLLM_LOG_DEBUG("Sequential receive for RNN cache."); + for (size_t i = 0; i < sourceNum; i++) + { + recvBufferFun(deviceId, i); + } + } + else + { + // Parallel receive with controlled concurrency + auto concurrencyNum = std::min(std::max(static_cast(1), bufferCoverSourceNum), sourceNum); + auto remainRecvNum = sourceNum; + + while (remainRecvNum > 0) + { + auto recvConcurrencyNum = std::min(remainRecvNum, concurrencyNum); + + // Avoid leaving a tiny remainder + if (remainRecvNum > concurrencyNum && remainRecvNum < (2 * concurrencyNum)) + { + recvConcurrencyNum = remainRecvNum - concurrencyNum; + } + + std::vector> futures; + futures.reserve(recvConcurrencyNum); + for (size_t i = 0; i < recvConcurrencyNum; i++) + { + size_t idx = i + (sourceNum - remainRecvNum); + TLLM_CHECK(idx < sourceNum); + futures.push_back(std::async(std::launch::async, recvBufferFun, deviceId, idx)); + } + for (auto& future : futures) + { + future.get(); + } + remainRecvNum -= recvConcurrencyNum; + } + } + } + else + { + recvBufferFun(deviceId, 0); + } + session.setTime(TransferSession::kTimeTransmissions); + + // Unpack received buffers into RNN states + std::vector outputConvBlocks; + std::vector outputSsmBlocks; + + auto const maxBatchSize = mRnnStateManager->getMaxBatchSize(); + auto convStates = mRnnStateManager->getConvStates(); // [numLocalLayers, maxBatchSize, convDim, dConv-1] + auto ssmStates = mRnnStateManager->getSsmStates(); // [numLocalLayers, maxBatchSize, numHeads, headDim, dState] + + outputConvBlocks.push_back(convStates); + outputSsmBlocks.push_back(ssmStates); + + tensorrt_llm::executor::rnn_cache::concatRnnConvStateDispatch( + recvBuffers, outputConvBlocks, slotIdx, maxBatchSize, destConfig, selfConfig, selfIdx, bufferManager); + + tensorrt_llm::executor::rnn_cache::concatRnnSsmStateDispatch(recvBuffers, outputSsmBlocks, slotIdx, maxBatchSize, + sourceConvBytesPerLayer, destConfig, selfConfig, selfIdx, bufferManager); + + bufferManager.getStream().synchronize(); + + if (cacheBufferId.has_value()) + { + mRnnCacheTransBufferManager->freeBufferIndexForRecv(cacheBufferId); + } + session.setTime(TransferSession::kTimePostprocess); + + TLLM_LOG_DEBUG( + mpi::MpiComm::world().getRank(), "End receiving RNN state for request ID: %ld.", llmRequest.mRequestId); +} + +void RnnCacheFormatter::formatUnifiedPoolMode(TransferSession& session) { NVTX3_SCOPED_RANGE(RnnCacheFormatter_formatUnifiedPool); session.setTime(TransferSession::kTimeFormatter); @@ -170,26 +643,11 @@ void RnnCacheFormatter::format(TransferSession& session) bufferSizesPerTarget[t] = ssmBufBytes + convBufBytes; } - auto const* sendCancelFlag = common::getEnvDisaggEnableInflightCancel() - ? &session.getDataContext().getTransferTerminate() - : nullptr; - auto cacheBufferId = mRnnCacheTransBufferManager->assignBufferIndexForSend(sendCancelFlag); - BufferIndexHolder sendHolder(*mRnnCacheTransBufferManager, cacheBufferId, /*isRecv=*/false); + auto cacheBufferId = mRnnCacheTransBufferManager->assignBufferIndexForSend(); auto allocationResult = mRnnCacheTransBufferManager->getOrAllocateSendBuffers( cacheBufferId, static_cast(bufferTargetNum), bufferSizesPerTarget, bufferManager); auto& outputBuffers = std::get<0>(allocationResult); auto& bufferCoverTargetNum = std::get<1>(allocationResult); - auto& onlyUseDynamicBuffer = std::get<2>(allocationResult); - TLLM_CHECK(cacheBufferId.has_value() || onlyUseDynamicBuffer); - auto const* agentConnection = rnnSendConns.empty() - ? nullptr - : dynamic_cast(connections[rnnSendConns[0]]); - if (agentConnection != nullptr) - { - TLLM_CHECK_WITH_INFO(bufferCoverTargetNum == bufferTargetNum, - "Agent needs all unified-pool RNN send buffers pre-allocated"); - TLLM_CHECK(onlyUseDynamicBuffer == false); - } // Split each outputBuffer into SSM and conv portions. std::vector ssmOutputBuffers(numTargets); @@ -215,30 +673,15 @@ void RnnCacheFormatter::format(TransferSession& session) bufferManager.getStream().synchronize(); session.setTime(TransferSession::kTimePreprocess); - // Send buffers to targets. + // Send (same protocol as slot mode). int deviceId; TLLM_CUDA_CHECK(cudaGetDevice(&deviceId)); auto preAllocSendBuffer = mRnnCacheTransBufferManager->getSendBuffer(cacheBufferId); - if (sendCancelFlag != nullptr && sendCancelFlag->load(std::memory_order_relaxed)) - { - TLLM_THROW("Unified-pool RNN cache transfer cancelled before NIXL submission"); - } - try - { - sendAllBuffers(session, deviceId, outputBuffers, bufferCoverTargetNum, preAllocSendBuffer, - bufferManager, targetInfo, rnnSendConns); - } - catch (...) - { - if (agentConnection != nullptr && common::getEnvDisaggEnableInflightCancel()) - { - sendHolder.poison(); - } - throw; - } + sendAllBuffers(session, deviceId, outputBuffers, bufferCoverTargetNum, preAllocSendBuffer, bufferManager, + targetInfo, rnnSendConns); session.setTime(TransferSession::kTimeTransmissions); - sendHolder.release(); + mRnnCacheTransBufferManager->freeBufferIndexForSend(cacheBufferId); } } @@ -246,7 +689,7 @@ void RnnCacheFormatter::format(TransferSession& session) llmRequest.mRequestId); } -void RnnCacheFormatter::unformat(TransferSession& session) +void RnnCacheFormatter::unformatUnifiedPoolMode(TransferSession& session) { NVTX3_SCOPED_RANGE(RnnCacheFormatter_unformatUnifiedPool); session.setTime(TransferSession::kTimeFormatter); @@ -376,23 +819,17 @@ void RnnCacheFormatter::unformat(TransferSession& session) bufferSizesPerSource[t] = ssmBufBytes + convBufBytes; } - // Use pre-assigned buffer ID from NIXL connection if available. + // Use pre-assigned buffer ID from NIXL connection if available (same as slot mode). std::optional cacheBufferId = std::nullopt; - BufferIndexHolder recvHolder; auto preAssignedRnnId = connections[rnnRecvConns[0]]->getPreAssignedBufferId(static_cast(BufferKind::kRNN)); if (preAssignedRnnId.has_value()) { cacheBufferId = static_cast(*preAssignedRnnId); - if (!session.hasReservedRecvBuffer(*mRnnCacheTransBufferManager)) - { - recvHolder = BufferIndexHolder(*mRnnCacheTransBufferManager, cacheBufferId, /*isRecv=*/true); - } } else { cacheBufferId = mRnnCacheTransBufferManager->assignBufferIndexForRecv(); - recvHolder = BufferIndexHolder(*mRnnCacheTransBufferManager, cacheBufferId, /*isRecv=*/true); } auto allocationResult = mRnnCacheTransBufferManager->getOrAllocateRecvBuffers( @@ -439,12 +876,10 @@ void RnnCacheFormatter::unformat(TransferSession& session) bufferManager.getStream().synchronize(); - recvHolder.release(); + mRnnCacheTransBufferManager->freeBufferIndexForRecv(cacheBufferId); } } - (void) session.releaseReservedRecvBuffer(*mRnnCacheTransBufferManager); - TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), "End receiving unified pool RNN state for request ID: %ld.", llmRequest.mRequestId); } diff --git a/cpp/tensorrt_llm/batch_manager/rnnCacheTransBuffer.cpp b/cpp/tensorrt_llm/batch_manager/rnnCacheTransBuffer.cpp index 37af8e31baf9..1e52a04b9580 100644 --- a/cpp/tensorrt_llm/batch_manager/rnnCacheTransBuffer.cpp +++ b/cpp/tensorrt_llm/batch_manager/rnnCacheTransBuffer.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -27,6 +27,52 @@ namespace tensorrt_llm::batch_manager::rnn_state_manager { +size_t RnnCacheTransBufferManager::computeTransferBufferSize( + RnnStateManager* rnnStateManager, std::optional maxNumTokens) +{ + SizeType32 numLocalLayers = rnnStateManager->getNumLocalLayers(); + + // Get the tensor for one layer to determine per-slot dimensions + // Conv state shape per layer: [maxBatchSize, convDim_local, dConv-1] + // SSM state shape per layer: [maxBatchSize, numHeads_local, headDim, dState] + // The tensors are shaped [maxBatchSize, ...], so one slot = total_size / maxBatchSize + auto convState = rnnStateManager->getConvStates(rnnStateManager->getGlobalLayerNum(0)); // Get first layer's tensor + auto ssmState = rnnStateManager->getSsmStates(rnnStateManager->getGlobalLayerNum(0)); + + auto convShape = convState->getShape(); + auto ssmShape = ssmState->getShape(); + + // Compute elements per slot per layer (divide total volume by batch size) + size_t convElemsPerSlotPerLayer = runtime::ITensor::volume(convShape) / convShape.d[0]; + size_t ssmElemsPerSlotPerLayer = runtime::ITensor::volume(ssmShape) / ssmShape.d[0]; + + size_t convDtypeSize = common::getDTypeSize(rnnStateManager->getConvStateDataType()); + size_t ssmDtypeSize = common::getDTypeSize(rnnStateManager->getSsmStateDataType()); + + size_t convBytesPerSlotPerLayer = convElemsPerSlotPerLayer * convDtypeSize; + size_t ssmBytesPerSlotPerLayer = ssmElemsPerSlotPerLayer * ssmDtypeSize; + + size_t bufferSizePerSlot = numLocalLayers * (convBytesPerSlotPerLayer + ssmBytesPerSlotPerLayer); + + TLLM_LOG_DEBUG( + "RNN computeTransferBufferSize: numLocalLayers=%d, convBytesPerLayer=%lu, ssmBytesPerLayer=%lu, " + "totalPerSlot=%lu", + numLocalLayers, convBytesPerSlotPerLayer, ssmBytesPerSlotPerLayer, bufferSizePerSlot); + + return bufferSizePerSlot > 0 ? bufferSizePerSlot : common::getEnvMemSizeForKVCacheTransferBuffer(); +} + +RnnCacheTransBufferManager::RnnCacheTransBufferManager( + RnnStateManager* rnnStateManager, std::optional maxNumTokens) + : BaseTransBufferManager(computeTransferBufferSize(rnnStateManager, maxNumTokens), + nvinfer1::DataType::kUINT8, // Use byte buffer for mixed dtypes + maxNumTokens) + , mRnnStateManager{rnnStateManager} +{ + TLLM_CHECK(mRnnStateManager != nullptr); + TLLM_LOG_INFO("RnnCacheTransBufferManager created for RNN cache"); +} + size_t RnnCacheTransBufferManager::computeTransferBufferSizeFromPool( kv_cache_manager::BaseKVCacheManager* kvCacheManager, executor::kv_cache::CacheState const& cacheState, std::optional maxNumTokens) @@ -96,6 +142,7 @@ RnnCacheTransBufferManager::RnnCacheTransBufferManager(kv_cache_manager::BaseKVC executor::kv_cache::CacheState const& cacheState, std::optional maxNumTokens) : BaseTransBufferManager(computeTransferBufferSizeFromPool(kvCacheManager, cacheState, maxNumTokens), nvinfer1::DataType::kUINT8, maxNumTokens) + , mRnnStateManager{nullptr} { TLLM_CHECK(kvCacheManager != nullptr); TLLM_LOG_INFO("RnnCacheTransBufferManager created for unified pool RNN cache"); @@ -116,7 +163,7 @@ size_t RnnCacheTransBufferManager::preAllocBufferSize( size_t transferBufferSize = rnnStateSizeBytes > 0 ? rnnStateSizeBytes : common::getEnvMemSizeForKVCacheTransferBuffer(); - bool useFabricMemory = kv_cache_manager::FabricMemory::supportFabricMemory() + bool useFabricMemory = kv_cache_manager::FabricMemory::supportFbaricMemory() && (!(common::getEnvKVCacheTransferUseSyncBuffer() || common::getEnvKVCacheTransferUseAsyncBuffer())); if (useFabricMemory) diff --git a/cpp/tensorrt_llm/batch_manager/rnnCacheTransBuffer.h b/cpp/tensorrt_llm/batch_manager/rnnCacheTransBuffer.h index 124525184815..e6ffa06db994 100644 --- a/cpp/tensorrt_llm/batch_manager/rnnCacheTransBuffer.h +++ b/cpp/tensorrt_llm/batch_manager/rnnCacheTransBuffer.h @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -18,6 +18,7 @@ #pragma once #include "tensorrt_llm/batch_manager/baseTransBuffer.h" +#include "tensorrt_llm/batch_manager/rnnStateManager.h" #include "tensorrt_llm/executor/dataTransceiverState.h" #include "tensorrt_llm/executor/executor.h" #include "tensorrt_llm/runtime/bufferManager.h" @@ -42,6 +43,11 @@ class RnnCacheTransBufferManager : public BaseTransBufferManager using SizeType32 = tensorrt_llm::runtime::SizeType32; using CacheState = executor::kv_cache::CacheState; + /// @brief Constructor for slot-based path (CppMambaCacheManager with RnnStateManager). + /// @param rnnStateManager Pointer to the RNN state manager. + /// @param maxNumTokens Optional maximum number of tokens for buffer sizing. + RnnCacheTransBufferManager(RnnStateManager* rnnStateManager, std::optional maxNumTokens = std::nullopt); + /// @brief Constructor for unified pool path (CppMambaHybridCacheManager). /// Computes buffer sizes from the KV cache manager's recurrent state pool metadata. /// @param kvCacheManager Pointer to the KV cache manager with unified pool. @@ -57,15 +63,26 @@ class RnnCacheTransBufferManager : public BaseTransBufferManager static size_t preAllocBufferSize( size_t rnnStateSizeBytes, std::optional const& cacheTransceiverConfig); + /// @brief Get the RNN state manager. + [[nodiscard]] RnnStateManager* getRnnStateManager() const noexcept + { + return mRnnStateManager; + } + [[nodiscard]] BufferKind getBufferKind() const override { return BufferKind::kRNN; } private: + /// @brief Compute transfer buffer size from RNN state configuration. + static size_t computeTransferBufferSize(RnnStateManager* rnnStateManager, std::optional maxNumTokens); + /// @brief Compute transfer buffer size from unified pool metadata. static size_t computeTransferBufferSizeFromPool(kv_cache_manager::BaseKVCacheManager* kvCacheManager, executor::kv_cache::CacheState const& cacheState, std::optional maxNumTokens); + + RnnStateManager* mRnnStateManager{nullptr}; }; } // namespace tensorrt_llm::batch_manager::rnn_state_manager diff --git a/cpp/tensorrt_llm/batch_manager/runtimeBuffers.cpp b/cpp/tensorrt_llm/batch_manager/runtimeBuffers.cpp index ea5b9b06a96e..691fb9c7efda 100644 --- a/cpp/tensorrt_llm/batch_manager/runtimeBuffers.cpp +++ b/cpp/tensorrt_llm/batch_manager/runtimeBuffers.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -152,8 +152,8 @@ void RuntimeBuffers::create(SizeType32 maxBatchSize, SizeType32 maxBeamWidth, ITensor::makeShape({GenerationLogitsCache::kCACHE_LENGTH, maxBatchSize * maxBeamWidth, vocabSizePadded}), logitsType); - generationLogitsCache.fragmentPointerDevice = manager.gpu( - ITensor::makeShape({maxBatchSize, GenerationLogitsCache::kCACHE_LENGTH}), nvinfer1::DataType::kINT64); + generationLogitsCache.fragmentPointerDevice + = manager.gpu(ITensor::makeShape({GenerationLogitsCache::kCACHE_LENGTH}), nvinfer1::DataType::kINT64); generationLogitsCache.fragmentPointerHost = tensorrt_llm::runtime::BufferManager::pinnedPool( ITensor::makeShape({maxBatchSize, GenerationLogitsCache::kCACHE_LENGTH}), nvinfer1::DataType::kINT64); } diff --git a/cpp/tensorrt_llm/batch_manager/trtEncoderModel.cpp b/cpp/tensorrt_llm/batch_manager/trtEncoderModel.cpp index 0d7dbfde42e6..de1525b07730 100644 --- a/cpp/tensorrt_llm/batch_manager/trtEncoderModel.cpp +++ b/cpp/tensorrt_llm/batch_manager/trtEncoderModel.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -76,11 +76,8 @@ TrtEncoderModel::TrtEncoderModel(runtime::ModelConfig const& modelConfig, WorldC // handling of maximizing utilization or pause/evict // TODO: finer control on encoder requests scheduling mCapacityScheduler = std::make_unique( - getMaxBatchSize() * mNumMicroBatches, executorConfig.getSchedulerConfig().getCapacitySchedulerPolicy(), - /*hasKvCacheManager=*/false, /*twoStepsLookAhead=*/false, - /*noScheduleUntilState=*/LlmRequestState::kENCODER_INIT, - /*noScheduleAfterState=*/LlmRequestState::kCONTEXT_INIT, - /*enablePrefixAwareScheduling=*/executorConfig.getSchedulerConfig().getEnablePrefixAwareScheduling()); + getMaxBatchSize() * mNumMicroBatches, executorConfig.getSchedulerConfig().getCapacitySchedulerPolicy(), false, + false, LlmRequestState::kENCODER_INIT, LlmRequestState::kCONTEXT_INIT); mMicroBatchScheduler = std::make_unique( std::nullopt, mModelConfig.getMaxInputLen(), LlmRequestState::kENCODER_INIT, LlmRequestState::kCONTEXT_INIT); diff --git a/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp b/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp index 7a0d78beb8a0..05ed827a9511 100644 --- a/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp +++ b/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp @@ -25,7 +25,6 @@ #include "tensorrt_llm/batch_manager/contextProgress.h" #include "tensorrt_llm/batch_manager/createNewDecoderRequests.h" #include "tensorrt_llm/batch_manager/decoderBuffers.h" -#include "tensorrt_llm/batch_manager/disaggTransferAdmissionController.h" #include "tensorrt_llm/batch_manager/guidedDecoder.h" #include "tensorrt_llm/batch_manager/handleContextLogits.h" #include "tensorrt_llm/batch_manager/handleGenerationLogits.h" @@ -334,8 +333,6 @@ TrtGptModelInflightBatching::TrtGptModelInflightBatching(std::shared_ptr( - cacheTransceiverConfig.getMaxTokensInBuffer(), mModelConfig.getTokensPerBlock()); } if (mModelConfig.getSpeculativeDecodingMode().needsKVCacheRewind()) @@ -445,10 +442,7 @@ TrtGptModelInflightBatching::TrtGptModelInflightBatching(std::shared_ptr(getMaxNumSequences(), executorConfig.getSchedulerConfig().getCapacitySchedulerPolicy(), mKvCacheManager != nullptr, - /*twoStepsLookAhead=*/mWorldConfig.isPipelineParallel(), - /*noScheduleUntilState=*/LlmRequestState::kCONTEXT_INIT, - /*noScheduleAfterState=*/LlmRequestState::kGENERATION_COMPLETE, - /*enablePrefixAwareScheduling=*/executorConfig.getSchedulerConfig().getEnablePrefixAwareScheduling()); + mWorldConfig.isPipelineParallel()); mMicroBatchScheduler = std::make_unique(ctxChunkConfig, maxContextLength); @@ -1024,26 +1018,8 @@ void TrtGptModelInflightBatching::forwardAsync(RequestList const& activeRequests auto [fittingRequests, fittingDisaggGenInitRequests, requestsToPause] = (*mCapacityScheduler)(activeRequests, mKvCacheManager, mPeftCacheManager, mCrossKvCacheManager); // Remove from fitting requests the requests that cannot be scheduled due to disagg KV cache transfer - bool waitForDisaggGenTransferProgress = false; if (mModelConfig.isTransformerBased() && getKVCacheManager() && mCacheTransceiver) { - if (mDisaggTransferAdmissionController && mDisaggTransferAdmissionController->enabled() - && !fittingDisaggGenInitRequests.empty()) - { - auto admissionResult - = mDisaggTransferAdmissionController->select(activeRequests, fittingDisaggGenInitRequests); - waitForDisaggGenTransferProgress = admissionResult.isBlockedByActiveTransfers(); - if (admissionResult.deferredRequestCount > 0) - { - TLLM_LOG_DEBUG( - "Disagg transfer admission deferred %zu requests; active transfer blocks=%zu, admitted " - "transfer blocks=%zu, budget=%zu", - admissionResult.deferredRequestCount, admissionResult.activeTransferBlocks, - admissionResult.admittedTransferBlocks, - mDisaggTransferAdmissionController->getMaxTransferBlocks().value_or(0)); - } - fittingDisaggGenInitRequests = std::move(admissionResult.admittedRequests); - } prepareDisaggGenInitRequests(activeRequests, fittingDisaggGenInitRequests); } if (fittingRequests.empty() && fittingDisaggGenInitRequests.empty()) @@ -1055,16 +1031,8 @@ void TrtGptModelInflightBatching::forwardAsync(RequestList const& activeRequests mIterCounter); if (mCacheTransceiver) { - if (waitForDisaggGenTransferProgress) - { - TLLM_LOG_DEBUG("Waiting for generation KV cache transfer progress to free disagg admission budget"); - mCacheTransceiver->checkGenTransferStatus(1); - } - else - { - mCacheTransceiver->checkContextTransferStatus(1, true); - // will free kvCache in next iteration. - } + mCacheTransceiver->checkContextTransferStatus(1, true); + // will free kvCache in next iteration. } } std::tie(currRequests.contextRequests, currRequests.generationRequests) @@ -1248,21 +1216,8 @@ void TrtGptModelInflightBatching::forwardAsync(RequestList const& activeRequests { for (auto const& llmReq : activeRequests) { - // Remove from mInflightReqIds so changeBeamWidth can proceed on the next iteration. - // terminateRequest frees seqSlot/KV cache but does not clean up mInflightReqIds. - mInflightReqIds.erase(llmReq->mRequestId); terminateRequest(llmReq); } - // Force buffer/decoder reset to clean up any partial state from the aborted batch - // (e.g. partially-filled cross-KV block offsets from mid-context-chunk processing). - // Guard on mInflightReqIds.empty(): in pipeline-parallel multi-micro-batch mode, - // other micro-batches may still have requests tracked here; changeBeamWidth asserts - // emptiness so we skip the reset and let the next successful forwardAsync iteration - // perform it when the set is clear. - if (mWorldConfig.isLastPipelineParallelRank() && mInflightReqIds.empty()) - { - changeBeamWidth(mOperatingBeamWidth); - } } catch (std::exception const& e) { @@ -1661,10 +1616,10 @@ void TrtGptModelInflightBatching::prepareDisaggGenInitRequests( auto const blockTransfer = std::all_of(activeRequests.begin(), activeRequests.end(), [](auto const& req) { return req->isDisaggGenerationTransmissionInProgress(); }); TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), - "newGenReqs.size():%ld requests, activeRequests.size():%ld allTransferInProgress:%d original " + "newGenReqs.size():%ld requests, activeRequests.size():%ld checkGenTransferStatus :%d original " "gen_only_requests_num:%ld", newGenReqs.size(), activeRequests.size(), blockTransfer, genInitReqNum); - mCacheTransceiver->checkGenTransferStatus(0); + mCacheTransceiver->checkGenTransferStatus(blockTransfer ? 1 : 0); auto timeEnd = std::chrono::steady_clock::now(); auto duration = std::chrono::duration(timeEnd - timeStart).count(); TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), @@ -1693,14 +1648,21 @@ void TrtGptModelInflightBatching::checkDisaggGenTransferStatus(RequestList const if (needCheck) { - mCacheTransceiver->checkGenTransferStatus(0); + auto const needCheckOne = std::all_of(activeRequests.begin(), activeRequests.end(), + [](auto const& req) { return req->isDisaggGenerationTransmissionInProgress(); }); + + int atLeastNum = needCheckOne ? 1 : 0; + TLLM_LOG_DEBUG( + mpi::MpiComm::world().getRank(), "noPreppared requests, checkGenTransferStatus atLeastNum:%d", atLeastNum); + + mCacheTransceiver->checkGenTransferStatus(atLeastNum); auto timeEnd = std::chrono::steady_clock::now(); auto duration = std::chrono::duration(timeEnd - timeStart).count(); TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), "no Prepare checkDisaggGenTransferStatus time:%f ms, " - "needCheck:%d,activeRequests.size():%ld", - duration, needCheck, activeRequests.size()); + "needCheckOne:%d,needCheck:%ld,activeRequests.size():%ld", + duration, needCheckOne, needCheck, activeRequests.size()); } } diff --git a/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.h b/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.h index d6550281a758..ca9f7c7f4a7f 100644 --- a/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.h +++ b/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.h @@ -88,7 +88,6 @@ class TrtGptModelTest; // Algorithms class CapacityScheduler; -class DisaggTransferAdmissionController; class MicroBatchScheduler; class PauseRequests; class AssignReqSeqSlots; @@ -601,7 +600,6 @@ class TrtGptModelInflightBatching : public TrtGptModel /******************** Cache transceiver ********************/ std::unique_ptr mCacheTransceiver; - std::unique_ptr mDisaggTransferAdmissionController; /******************** Spec dec ***********************/ std::unique_ptr mDraftModelSendLogitsThread; diff --git a/cpp/tensorrt_llm/batch_manager/utils/inflightBatchingUtils.cpp b/cpp/tensorrt_llm/batch_manager/utils/inflightBatchingUtils.cpp index 416235f347b8..bdb12886337c 100644 --- a/cpp/tensorrt_llm/batch_manager/utils/inflightBatchingUtils.cpp +++ b/cpp/tensorrt_llm/batch_manager/utils/inflightBatchingUtils.cpp @@ -103,11 +103,10 @@ void copyGenerationLogits(RuntimeBuffers::GenerationLogitsCache& generationLogit auto const fragmentSize = llmReq.getGenerationLogitsFragmentsSize(); - // Merge logits fragments on device. getFragmentPointerSlot() returns the matching host and - // device rows for the current workIdx and advances the index atomically, so concurrent flushes - // for different requests in the same batch never clobber each other's pointer arrays. + // Merge logits fragments on device auto const& transposeBufferPtr = generationLogitsCache.transposedLogits; - auto [cachePointerHost, cachePointerDevice] = generationLogitsCache.getFragmentPointerSlot(); + auto const& cachePointerDevice = generationLogitsCache.fragmentPointerDevice; + auto const& cachePointerHost = generationLogitsCache.getFragmentPointerHost(); tensorrt_llm::runtime::kernels::mergeLogitsFragments(bufferManager, *transposeBufferPtr, llmReq.getGenerationLogitsFragments(), *cachePointerDevice, *cachePointerHost, 0, 1, reqBeamWidth, bufferManager.getStream(), 0); diff --git a/cpp/tensorrt_llm/common/attentionOp.cpp b/cpp/tensorrt_llm/common/attentionOp.cpp index b91c0ef98df6..1c838bbcf12c 100644 --- a/cpp/tensorrt_llm/common/attentionOp.cpp +++ b/cpp/tensorrt_llm/common/attentionOp.cpp @@ -23,7 +23,6 @@ #include "tensorrt_llm/common/memoryUtils.h" #include "tensorrt_llm/common/sageQuant.h" #include "tensorrt_llm/kernels/decoderMaskedMultiheadAttention.h" -#include "tensorrt_llm/kernels/decoderMaskedMultiheadAttention/cascadeAttentionKernel.h" #include "tensorrt_llm/kernels/flashMLA/flash_mla.h" #include "tensorrt_llm/kernels/gptKernels.h" #include "tensorrt_llm/kernels/kvCacheUtils.h" @@ -118,12 +117,6 @@ struct FusedQKVMaskedAttentionDispatchParams float* partial_sum; float* partial_max; int* block_counter; - // Cascade attention prefix-side workspace (fp32). Sliced from the same - // generation workspace and forwarded into Multihead_attention_params so - // the cascade fast-path no longer needs its own cudaMalloc. - float* cascade_partial_out{}; - float* cascade_partial_max{}; - float* cascade_partial_sum{}; float const* kv_scale_orig_quant; float const* kv_scale_quant_orig; tc::QuantMode kv_cache_quant_mode; @@ -700,10 +693,6 @@ void fusedQKV_masked_attention_dispatch(Multihead_attention_params> max_num_tokens). - // enqueueContext sizes these buffers by total_kv_len, so workspace must match. - size_t const kv_buf_tokens = std::max( - static_cast(total_kv_len), static_cast(mChunkPrefillBufferBatchSize) * max_num_tokens); - fp8_k_buf_size = kv_buf_tokens * static_cast(total_k_dim_all_heads); - fp8_v_buf_size = kv_buf_tokens * static_cast(total_v_dim_all_heads); + fp8_k_buf_size = mChunkPrefillBufferBatchSize * max_num_tokens * static_cast(total_k_dim_all_heads); + fp8_v_buf_size = mChunkPrefillBufferBatchSize * max_num_tokens * static_cast(total_v_dim_all_heads); } } else if (useSageAttnSeparateQkv) @@ -989,13 +974,6 @@ size_t AttentionOp::getWorkspaceSizeForGeneration(nvinfer1::DataType type, int32 generationWorkspaceSizes.partialSum = partial_sum_size; generationWorkspaceSizes.partialMax = partial_max_size; generationWorkspaceSizes.shiftKCache = shift_k_cache_size; - { - auto const cascadeSizes - = tensorrt_llm::kernels::mmha::cascade::getCascadeWorkspaceSizes(batch_beam, mNumHeads, mHeadSize); - generationWorkspaceSizes.cascadeOut = cascadeSizes.out; - generationWorkspaceSizes.cascadeMax = cascadeSizes.mMax; - generationWorkspaceSizes.cascadeSum = cascadeSizes.lSum; - } generation_workspace_size = AttentionWorkspaceManager::buildGenerationLayout(generationWorkspaceSizes).totalSize; size_t xqa_workspace_size = 0; @@ -1130,12 +1108,6 @@ int AttentionOp::mlaGeneration( tllmRunnerParams.oPtr = reinterpret_cast(params.context_buf); tllmRunnerParams.oSfPtr = generation_params.context_buf_sf; - if (params.dsv4_epilogue_fusion.enabled) - { - tllmRunnerParams.mDsv4EpilogueFusion.enabled = true; - tllmRunnerParams.mDsv4EpilogueFusion.cosSinCache = params.dsv4_epilogue_fusion.cos_sin_cache; - tllmRunnerParams.mDsv4EpilogueFusion.scaleBufM = params.dsv4_epilogue_fusion.scale_buf_m; - } // softmax stats if needed tllmRunnerParams.softmaxStatsPtr = generation_params.softmax_stats; @@ -1607,29 +1579,13 @@ int AttentionOp::enqueueContext(EnqueueContextParams const& params, cudaStrea auto const workspaceViews = AttentionWorkspaceManager::materializeContext( params.workspace, workspaceLayout, cpMaxPadedSequenceLength, getHeadSize(), mNumHeads, mNumKVHeads); - auto* fp8QBuf = workspaceViews.fp8QBuf; - // Fused FP8-Q path: caller pre-fills the nope segment of `quant_q_buf`; - // route the context-MLA Q pointer to it so the fused RoPE kernel appends - // rope FP8 in place and the FMHA Q load reads the merged [nope|rope] buffer. - if (mIsMLAEnabled && params.mla_param != nullptr && params.mla_param->fuse_q_fp8_in_rope - && params.mla_param->quant_q_buf != nullptr) - { - fp8QBuf = reinterpret_cast<__nv_fp8_e4m3*>(params.mla_param->quant_q_buf); - } - // build attention mask, cu_seqlens, and padding offset tensors // Note: self attn and cross attn should use different params // cross attn's seqlen info is from encoder input lengths, not decoder input lengths! // moreover, attn mask for cross attn should be set separately (see below) BuildDecoderInfoParams decoder_params{}; - int32_t const* precomputedCuQSeqlens = params.cu_q_seqlens; - int32_t const* precomputedCuKvSeqlens = params.cu_kv_seqlens != nullptr ? params.cu_kv_seqlens - : params.cu_q_seqlens != nullptr ? params.cu_q_seqlens - : nullptr; decoder_params.seqQOffsets = workspaceViews.cuQSeqlens; decoder_params.seqKVOffsets = workspaceViews.cuKvSeqlens; - decoder_params.precomputedSeqQOffsets = precomputedCuQSeqlens; - decoder_params.precomputedSeqKVOffsets = precomputedCuKvSeqlens; decoder_params.seqCpPartialOffsets = workspaceViews.cuCpPartialSeqlens; decoder_params.cpSize = mCpSize; decoder_params.packedMaskRowOffsets = workspaceViews.cuMaskRows; @@ -1673,11 +1629,6 @@ int AttentionOp::enqueueContext(EnqueueContextParams const& params, cudaStrea invokeBuildDecoderInfo(decoder_params, stream); sync_check_cuda_error(stream); - int32_t const* contextCuQSeqlens - = precomputedCuQSeqlens != nullptr ? precomputedCuQSeqlens : workspaceViews.cuQSeqlens; - int32_t const* contextCuKvSeqlens - = precomputedCuKvSeqlens != nullptr ? precomputedCuKvSeqlens : workspaceViews.cuKvSeqlens; - // In cross attention context phase, the attention mask should be a matrix of all ones. // Override the attention mask produced by invokeBuildDecoderInfo(). // also, invokeBuildDecoderInfo can only handle square mask, not cross B x q_len x kv_len mask @@ -1742,7 +1693,8 @@ int AttentionOp::enqueueContext(EnqueueContextParams const& params, cudaStrea if (mCpSize > 1 && mAttnTpSize > 1 && mAttnCpSize == 1) { this->template ulyssesContextPreprocess(attention_input, workspaceViews.gatherInBuffer, - workspaceViews.gatherOutBuffer, params, contextCuQSeqlens, workspaceViews.cuCpPartialSeqlens, stream); + workspaceViews.gatherOutBuffer, params, workspaceViews.cuQSeqlens, workspaceViews.cuCpPartialSeqlens, + stream); attention_input = workspaceViews.gatherInBuffer; sync_check_cuda_error(stream); } @@ -1769,16 +1721,12 @@ int AttentionOp::enqueueContext(EnqueueContextParams const& params, cudaStrea preprocessingParams.qkv_bias = params.qkv_bias; preprocessingParams.tokens_info = decoder_params.tokensInfo; preprocessingParams.seq_lens = params.context_lengths; - // For self-attention, cache_seq_lens indicates whether chunked context is used - // (i.e. cache_seq_len > seq_len). - // For cross-attention, callers do not consistently use sequence_lengths as decoder length; use decoder - // context lengths so the encoder KV-cache write gate opens. - preprocessingParams.cache_seq_lens = isCrossAttention() ? params.context_lengths : params.sequence_lengths; - + // Indicate if chunked-context is used (i.e. q_seqlen > kv_seqlen). + preprocessingParams.cache_seq_lens = params.sequence_lengths; preprocessingParams.encoder_seq_lens = params.encoder_input_lengths; - preprocessingParams.cu_seq_lens = contextCuQSeqlens; + preprocessingParams.cu_seq_lens = workspaceViews.cuQSeqlens; // Cross-attention only. - preprocessingParams.cu_kv_seq_lens = contextCuKvSeqlens; + preprocessingParams.cu_kv_seq_lens = workspaceViews.cuKvSeqlens; preprocessingParams.rotary_embedding_inv_freq = workspaceViews.rotaryInvFreq; preprocessingParams.rotary_coef_cache_buffer = params.rotary_cos_sin; preprocessingParams.mrope_rotary_cos_sin = params.mrope_rotary_cos_sin; @@ -1837,13 +1785,12 @@ int AttentionOp::enqueueContext(EnqueueContextParams const& params, cudaStrea { TLLM_CHECK_WITH_INFO(params.mla_param != nullptr, "MLA param is nullptr"); params.mla_param->cache_type = cache_type; - params.mla_param->cu_q_seqlens = const_cast(contextCuQSeqlens); - params.mla_param->cu_kv_seqlens = const_cast(contextCuKvSeqlens); + params.mla_param->cu_q_seqlens = workspaceViews.cuQSeqlens; params.mla_param->quant_scale_kv = params.kv_scale_orig_quant; // Set BMM scales for FP8 context computation params.mla_param->bmm1_scale = workspaceViews.fmhaBmm1Scale; params.mla_param->bmm2_scale = workspaceViews.fmhaBmm2Scale; - params.mla_param->quant_q_buf = mFP8ContextMLA ? fp8QBuf : nullptr; + params.mla_param->quant_q_buf = mFP8ContextMLA ? workspaceViews.fp8QBuf : nullptr; params.mla_param->quant_k_buf = mFP8ContextMLA ? workspaceViews.fp8KBuf : nullptr; params.mla_param->quant_v_buf = mFP8ContextMLA ? workspaceViews.fp8VBuf : nullptr; // Set additional scales for context phase @@ -1856,16 +1803,11 @@ int AttentionOp::enqueueContext(EnqueueContextParams const& params, cudaStrea = 1 / (mQScaling * sqrt((float) (mMLAParams.qk_nope_head_dim + mMLAParams.qk_rope_head_dim))); // The sparse MLA is in the absorption mode for the context phase. params.mla_param->absorption_mode = useSparseMLA(); - // Fused FP8-Q-quant: RoPE kernel writes FP8 rope into `quant_q_buf`, - // so we skip the standalone invokeMLAContextFp8Quantize call below. - bool const useFusedQFp8 = params.mla_param->fuse_q_fp8_in_rope && mFP8ContextMLA - && params.mla_param->absorption_mode && cache_type == KvCacheDataType::FP8 - && params.mla_param->quant_q_buf != nullptr && params.mla_param->quant_scale_qkv != nullptr; if (params.mla_param->latent_cache != nullptr) { invokeMLARopeContext(*params.mla_param, kv_cache_buffer, stream); } - if (mFP8ContextMLA && !useFusedQFp8) + if (mFP8ContextMLA) { invokeMLAContextFp8Quantize(*params.mla_param, params.total_kv_len, stream); } @@ -1977,14 +1919,14 @@ int AttentionOp::enqueueContext(EnqueueContextParams const& params, cudaStrea { TLLM_CHECK_WITH_INFO( mFmhaDispatcher->isSeparateQAndKvInput(), "Separate QKV input is required for fp8 context MLA"); - TLLM_CHECK_WITH_INFO(fp8QBuf != nullptr, "FP8 q buffer is required for fp8 context MLA"); + TLLM_CHECK_WITH_INFO(workspaceViews.fp8QBuf != nullptr, "FP8 q buffer is required for fp8 context MLA"); // In sparse MLA (absorption mode), K and V are stored in KV cache, not as separate FP8 buffers TLLM_CHECK_WITH_INFO(useSparseMLA() || workspaceViews.fp8KBuf != nullptr, "FP8 k buffer is required for fp8 context MLA in non-sparse mode"); TLLM_CHECK_WITH_INFO(useSparseMLA() || workspaceViews.fp8VBuf != nullptr, "FP8 v buffer is required for fp8 context MLA in non-sparse mode"); - fmhaParams.qPtr = reinterpret_cast(fp8QBuf); + fmhaParams.qPtr = reinterpret_cast(workspaceViews.fp8QBuf); fmhaParams.kPtr = useSparseMLA() ? nullptr : reinterpret_cast(workspaceViews.fp8KBuf); fmhaParams.vPtr = useSparseMLA() ? nullptr : reinterpret_cast(workspaceViews.fp8VBuf); } @@ -2024,12 +1966,6 @@ int AttentionOp::enqueueContext(EnqueueContextParams const& params, cudaStrea // Only use [totalLength, h / cpSize, Dh]. fmhaParams.outputPtr = mCpSize > 1 ? workspaceViews.gatherOutBuffer : params.context_buf; fmhaParams.outputSfPtr = params.context_buf_sf; - if (params.mla_param != nullptr && params.mla_param->dsv4_epilogue_fusion.enabled) - { - fmhaParams.dsv4EpilogueFusion.enabled = true; - fmhaParams.dsv4EpilogueFusion.cosSinCache = params.mla_param->dsv4_epilogue_fusion.cos_sin_cache; - fmhaParams.dsv4EpilogueFusion.scaleBufM = params.mla_param->dsv4_epilogue_fusion.scale_buf_m; - } fmhaParams.attentionSinksPtr = params.attention_sinks; fmhaParams.packedMaskPtr = params.attention_packed_mask; if constexpr (std::is_same_v) @@ -2037,9 +1973,9 @@ int AttentionOp::enqueueContext(EnqueueContextParams const& params, cudaStrea fmhaParams.pagedKvCache = kv_cache_buffer; fmhaParams.pagedKvSfCache = kv_scale_cache_buffer; } - fmhaParams.cuQSeqLenPtr = contextCuQSeqlens; + fmhaParams.cuQSeqLenPtr = workspaceViews.cuQSeqlens; fmhaParams.kvSeqLenPtr = decoder_params.seqKVLengths; - fmhaParams.cuKvSeqLenPtr = contextCuKvSeqlens; + fmhaParams.cuKvSeqLenPtr = workspaceViews.cuKvSeqlens; fmhaParams.cuMaskRowsPtr = workspaceViews.cuMaskRows; fmhaParams.tileCounterPtr = workspaceViews.fmhaTileCounter; fmhaParams.scaleBmm1Ptr = workspaceViews.fmhaBmm1Scale; @@ -2087,8 +2023,8 @@ int AttentionOp::enqueueContext(EnqueueContextParams const& params, cudaStrea if (mCpSize > 1 && mAttnTpSize > 1 && mAttnCpSize == 1) { this->template ulyssesContextPostprocess(workspaceViews.gatherOutBuffer, - reinterpret_cast(params.context_buf), workspaceViews.gatherInBuffer, params, contextCuQSeqlens, - workspaceViews.cuCpPartialSeqlens, stream); + reinterpret_cast(params.context_buf), workspaceViews.gatherInBuffer, params, + workspaceViews.cuQSeqlens, workspaceViews.cuCpPartialSeqlens, stream); sync_check_cuda_error(stream); } @@ -2575,13 +2511,6 @@ int AttentionOp::enqueueGeneration(EnqueueGenerationParams const& params, cud workspaceSizes.partialSum = partial_sum_size; workspaceSizes.partialMax = partial_max_size; workspaceSizes.shiftKCache = shift_k_cache_size; - { - auto const cascadeSizes - = tensorrt_llm::kernels::mmha::cascade::getCascadeWorkspaceSizes(batch_beam, mNumHeads, mHeadSize); - workspaceSizes.cascadeOut = cascadeSizes.out; - workspaceSizes.cascadeMax = cascadeSizes.mMax; - workspaceSizes.cascadeSum = cascadeSizes.lSum; - } auto const workspaceLayout = AttentionWorkspaceManager::buildGenerationLayout(workspaceSizes); auto const workspaceViews = AttentionWorkspaceManager::materializeGeneration( params.workspace, workspaceLayout, cpMaxPaddedSequenceLength, mNumHeads, mNumKVHeads, mHeadSize); @@ -2650,9 +2579,6 @@ int AttentionOp::enqueueGeneration(EnqueueGenerationParams const& params, cud dispatch_params.partial_out = workspaceViews.partialOut; dispatch_params.partial_sum = workspaceViews.partialSum; dispatch_params.partial_max = workspaceViews.partialMax; - dispatch_params.cascade_partial_out = workspaceViews.cascadeOut; - dispatch_params.cascade_partial_max = workspaceViews.cascadeMax; - dispatch_params.cascade_partial_sum = workspaceViews.cascadeSum; dispatch_params.block_counter = mMultiBlockSemaphores.get(); dispatch_params.kv_cache_quant_mode = mKVCacheQuantMode; dispatch_params.kv_scale_orig_quant = params.kv_scale_orig_quant; @@ -2978,10 +2904,6 @@ int AttentionOp::initialize() noexcept fmhaParams.dataTypeKv = DATA_TYPE_E4M3; fmhaParams.dataTypeOut = DATA_TYPE_BF16; } - if (mFusesDsv4InvRopeFp8Quant) - { - fmhaParams.dataTypeOut = DATA_TYPE_E4M3; - } // TODO: remove forceFp32Acc from MHARunnerFixedParams after adding host_runtime_perf_knobs to // bertAttentionPlugin input tensors, so that we can change mLaunchParams.force_fp32_acc value in runtime. fmhaParams.forceFp32Acc = false; @@ -3056,7 +2978,6 @@ int AttentionOp::initialize() noexcept fmhaParams.scaleAlibi = isAliBiWithScale(); fmhaParams.useSparseMLA = useSparseMLA(); fmhaParams.useTllmGenSparseAttention = useTllmGenSparseAttention(); - fmhaParams.fusesDsv4InvRopeFp8Quant = mFusesDsv4InvRopeFp8Quant; // SageAttention: set block sizes for sage quantization. if (useSageAttn) @@ -3101,14 +3022,9 @@ int AttentionOp::initialize() noexcept qDataType = DATA_TYPE_E4M3; kvDataType = DATA_TYPE_E4M3; } - if (mFusesDsv4InvRopeFp8Quant) - { - outputDataType = DATA_TYPE_E4M3; - } // Instantiate the mTllmGenFMHARunner used for MLA - mTllmGenFMHARunner.reset(new TllmGenFmhaRunner( - qDataType, kvDataType, kvDataType, outputDataType, 0, 0, 0, 0, mFusesDsv4InvRopeFp8Quant)); + mTllmGenFMHARunner.reset(new TllmGenFmhaRunner(qDataType, kvDataType, kvDataType, outputDataType)); } else if (mIsGenerationMLA && !mUseGenFlashMLA) { diff --git a/cpp/tensorrt_llm/common/attentionOp.h b/cpp/tensorrt_llm/common/attentionOp.h index f7337c9c9cb2..c8f0d33e4705 100644 --- a/cpp/tensorrt_llm/common/attentionOp.h +++ b/cpp/tensorrt_llm/common/attentionOp.h @@ -57,7 +57,7 @@ class AttentionOp [[nodiscard]] int getHeadSize(bool checkInit = true) const; [[nodiscard]] int getMaxNumSeqLenTile(int batch_beam_size = 1) const; [[nodiscard]] size_t getWorkspaceSizeForContext(nvinfer1::DataType type, int32_t nbReq, int32_t max_input_length, - int32_t cross_kv_length = 0, int32_t max_num_tokens = 0, int32_t total_kv_len = 0) const noexcept; + int32_t cross_kv_length = 0, int32_t max_num_tokens = 0) const noexcept; // total_num_seq is the sum of beam_width for multiple requests [[nodiscard]] size_t getWorkspaceSizeForGeneration(nvinfer1::DataType type, int32_t total_num_seq, int32_t max_attention_window_size, int32_t max_num_tokens, int32_t max_blocks_per_sequence) const noexcept; @@ -153,12 +153,6 @@ class AttentionOp int32_t const* helix_position_offsets = nullptr; bool const* helix_is_inactive_rank = nullptr; - // Optional packed-varlen boundaries for context attention. When set, - // these describe attention sequences/segments and are used directly by - // FMHA instead of the workspace boundaries rebuilt from context_lengths. - int32_t const* cu_q_seqlens = nullptr; - int32_t const* cu_kv_seqlens = nullptr; - std::string enqueueContextParamsToString() const { // variables from the params coming from the runtime @@ -217,8 +211,6 @@ class AttentionOp ss << "softmaxStatsPtr: " << this->softmax_stats << std::endl; ss << "k_ptr: " << this->k_ptr << std::endl; ss << "v_ptr: " << this->v_ptr << std::endl; - ss << "cu_q_seqlens: " << this->cu_q_seqlens << std::endl; - ss << "cu_kv_seqlens: " << this->cu_kv_seqlens << std::endl; return ss.str(); } }; @@ -533,8 +525,6 @@ class AttentionOp // Whether to fuse FP4 quant into attention kernel. bool mFuseFp4Quant = false; - // Whether to fuse DSv4 inverse-RoPE + FP8 output quant into trtllm-gen FMHA. - bool mFusesDsv4InvRopeFp8Quant = false; kernels::SparseAttentionParams mRuntimeSparseAttentionParams; @@ -576,9 +566,9 @@ class AttentionOp mMLAParams.data(), mCpSize, mCpRank, mCpGroup, mNumAttnHeads, mNumAttnKVHeads, mNumKVHeadsOrigin, mAttnTpSize, mAttnTpRank, mAttnCpSize, mAttnCpRank, mUlyssesMQABroadcast, mEnableContextFMHA, mFMHAForceFP32Acc, mMultiBlockMode, mEnableXQA, mUseKVCache, mSkipAttn, mFuseFp4Quant, - mFusesDsv4InvRopeFp8Quant, mNbMultiBlockSemaphores, mAttentionChunkSize.value_or(-1), - mSkipSoftmaxThresholdScaleFactorPrefill, mSkipSoftmaxThresholdScaleFactorDecode, mSageAttnNumEltsPerBlkQ, - mSageAttnNumEltsPerBlkK, mSageAttnNumEltsPerBlkV, mSageAttnQkInt8); + mNbMultiBlockSemaphores, mAttentionChunkSize.value_or(-1), mSkipSoftmaxThresholdScaleFactorPrefill, + mSkipSoftmaxThresholdScaleFactorDecode, mSageAttnNumEltsPerBlkQ, mSageAttnNumEltsPerBlkK, + mSageAttnNumEltsPerBlkV, mSageAttnQkInt8); }; private: diff --git a/cpp/tensorrt_llm/common/attentionWorkspace.h b/cpp/tensorrt_llm/common/attentionWorkspace.h index 3ba53da3ff34..a9c69b693590 100644 --- a/cpp/tensorrt_llm/common/attentionWorkspace.h +++ b/cpp/tensorrt_llm/common/attentionWorkspace.h @@ -142,12 +142,6 @@ struct AttentionGenerationWorkspaceSizes size_t partialSum{}; size_t partialMax{}; size_t shiftKCache{}; - // Cascade-attention prefix-side partials (fp32), populated only when the - // owning AttentionOp's generation path can dispatch the cascade fast-path. - // See cascade::getCascadeWorkspaceSizes for the field layout. - size_t cascadeOut{}; - size_t cascadeMax{}; - size_t cascadeSum{}; }; struct AttentionGenerationWorkspaceLayout @@ -157,9 +151,6 @@ struct AttentionGenerationWorkspaceLayout WorkspaceSlice partialSum{}; WorkspaceSlice partialMax{}; WorkspaceSlice shiftKCache{}; - WorkspaceSlice cascadeOut{}; - WorkspaceSlice cascadeMax{}; - WorkspaceSlice cascadeSum{}; size_t totalSize{}; }; @@ -172,9 +163,6 @@ struct AttentionGenerationWorkspaceViews float* partialSum{}; float* partialMax{}; T* shiftKCache{}; - float* cascadeOut{}; - float* cascadeMax{}; - float* cascadeSum{}; }; struct AttentionFlashMlaWorkspaceSizes @@ -314,9 +302,6 @@ class AttentionWorkspaceManager layout.partialSum = nextSlice(offset, sizes.partialSum, alignment); layout.partialMax = nextSlice(offset, sizes.partialMax, alignment); layout.shiftKCache = nextSlice(offset, sizes.shiftKCache, alignment); - layout.cascadeOut = nextSlice(offset, sizes.cascadeOut, alignment); - layout.cascadeMax = nextSlice(offset, sizes.cascadeMax, alignment); - layout.cascadeSum = nextSlice(offset, sizes.cascadeSum, alignment); layout.totalSize = offset; return layout; } @@ -338,9 +323,6 @@ class AttentionWorkspaceManager views.partialSum = ptr(workspace, layout.partialSum); views.partialMax = ptr(workspace, layout.partialMax); views.shiftKCache = ptr(workspace, layout.shiftKCache); - views.cascadeOut = ptr(workspace, layout.cascadeOut); - views.cascadeMax = ptr(workspace, layout.cascadeMax); - views.cascadeSum = ptr(workspace, layout.cascadeSum); return views; } diff --git a/cpp/tensorrt_llm/common/envUtils.cpp b/cpp/tensorrt_llm/common/envUtils.cpp index 14c637699584..61e81c9135af 100644 --- a/cpp/tensorrt_llm/common/envUtils.cpp +++ b/cpp/tensorrt_llm/common/envUtils.cpp @@ -259,12 +259,6 @@ bool getEnvEnablePDL() return enablePDL; } -bool getEnvEnableCascadeMmha() -{ - static bool const enable = getBoolEnv("TRTLLM_ENABLE_CASCADE_MMHA"); - return enable; -} - bool getEnvEnableTrtllmgenMoeRoutingRenormPDL() { static std::once_flag flag; @@ -401,12 +395,6 @@ bool getEnvTryZCopyForKVCacheTransfer() return zcopyForSysmmetricKVCache; } -bool getEnvDisaggEnableInflightCancel() -{ - static bool const enabled = getBoolEnv("TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL"); - return enabled; -} - bool getEnvForceDeterministic() { static bool const forceDeterministic = getBoolEnv("FORCE_DETERMINISTIC"); @@ -505,12 +493,6 @@ bool getEnvKVCacheTransferAllBlocksForWindow() return allBlocksForWindow; } -bool getEnvKVCachePoolUseFabricMemory() -{ - static bool const useFabricMemory = getBoolEnv("TRTLLM_KVCACHE_POOL_USE_FABRIC_MEMORY"); - return useFabricMemory; -} - uint16_t getEnvNixlPort() { static uint16_t const nixlPort = getUInt64Env("TRTLLM_NIXL_PORT").value_or(0); diff --git a/cpp/tensorrt_llm/common/envUtils.h b/cpp/tensorrt_llm/common/envUtils.h index 37b5934b4515..196c15aadbf8 100644 --- a/cpp/tensorrt_llm/common/envUtils.h +++ b/cpp/tensorrt_llm/common/envUtils.h @@ -55,11 +55,6 @@ int getEnvMmhaKernelBlockSize(); // Whether PDL is enabled. bool getEnvEnablePDL(); -// Whether the experimental cascade attention kernel is enabled (replaces -// masked_multihead_attention_kernel for beam-search decoding). -// Controlled by env var TRTLLM_ENABLE_CASCADE_MMHA (default: false). -bool getEnvEnableCascadeMmha(); - // Whether PDL is enabled for MoE Renormalize routing kernel. // Disabled by default to avoid NaN corruption (https://nvbugs/5955170). // Set TRTLLM_ENABLE_TRTLLMGEN_MOE_ROUTING_RENORM_PDL=1 to re-enable. @@ -115,9 +110,6 @@ std::string const& getEnvKVCacheTimeOutputPath(); bool getEnvTryZCopyForKVCacheTransfer(); -// Opt-in for disaggregated KV transfer in-flight cancellation and fail-closed transfer-buffer quarantine. -bool getEnvDisaggEnableInflightCancel(); - // Force deterministic behavior for all kernels. bool getEnvForceDeterministic(); @@ -147,8 +139,6 @@ size_t getEnvKVCacheSendMaxConcurrenceNum(); size_t getEnvMemSizeForKVCacheTransferBuffer(); -bool getEnvKVCachePoolUseFabricMemory(); - uint16_t getEnvNixlPort(); bool getEnvNixlEnableCoalesce(); diff --git a/cpp/tensorrt_llm/common/lamportUtils.cuh b/cpp/tensorrt_llm/common/lamportUtils.cuh index 60639e2f4870..9e2f22d1a182 100644 --- a/cpp/tensorrt_llm/common/lamportUtils.cuh +++ b/cpp/tensorrt_llm/common/lamportUtils.cuh @@ -1,5 +1,5 @@ /* - * Copyright (c) 2025-2026, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -36,8 +36,6 @@ namespace common { constexpr uint16_t kNEGZERO_FP16 = 0x8000U; -constexpr uint32_t kNEGZERO_FP32 = 0x80000000U; -constexpr uint32_t kWARP_SIZE = 32U; template union Fp16BitCast @@ -85,7 +83,7 @@ static inline __device__ bool isNegZero(T val) if constexpr (std::is_same_v) { - return __float_as_uint(val) == kNEGZERO_FP32; + return val == 0.F && signbit(val); } else if constexpr (std::is_same_v || std::is_same_v) { @@ -123,56 +121,6 @@ constexpr __device__ __host__ PackedType getPackedLamportInit() return initValue.mPacked; } -template -union VolatilePackedLoad -{ - PackedType packed; - uint32_t words[sizeof(PackedType) / sizeof(uint32_t)]; -}; - -template -inline __device__ VolatilePackedLoad loadPackedVolatile(void const* ptr) -{ - static_assert(sizeof(PackedType) == 0, "loadPackedVolatile not specialized for this type"); - return {}; -} - -template <> -inline __device__ VolatilePackedLoad loadPackedVolatile(void const* ptr) -{ - VolatilePackedLoad returnValue; - asm volatile( - "ld.volatile.global.v4.u32 {%0, %1, %2, %3}, [%4];\n" - : "=r"(returnValue.words[0]), "=r"(returnValue.words[1]), "=r"(returnValue.words[2]), "=r"(returnValue.words[3]) - : "l"(ptr) - : "memory"); - return returnValue; -} - -template <> -inline __device__ VolatilePackedLoad loadPackedVolatile(void const* ptr) -{ - VolatilePackedLoad returnValue; - asm volatile("ld.volatile.global.v2.u32 {%0, %1}, [%2];\n" - : "=r"(returnValue.words[0]), "=r"(returnValue.words[1]) - : "l"(ptr) - : "memory"); - return returnValue; -} - -template -inline __device__ bool isLamportDirty(VolatilePackedLoad const& value) -{ - // The dirty sentinel is a raw word; typed fp compares can flush nearby bit patterns. - bool dirty = false; -#pragma unroll - for (int i = 0; i < sizeof(PackedType) / sizeof(uint32_t); i++) - { - dirty |= value.words[i] == kNEGZERO_FP32; - } - return dirty; -} - // A helper class to get the correct base pointer for a given layout struct LamportBufferLayout { @@ -206,8 +154,8 @@ struct LamportBufferLayout namespace cg = cooperative_groups; // PackedType is the one used in kernel for Lamport buffer (LDG.128 or LDG.64) -template -struct __attribute__((aligned(32))) LamportFlags +template +__device__ struct __attribute__((aligned(32))) LamportFlags { public: __device__ explicit LamportFlags(uint32_t* bufferFlags, uint32_t numStages = 1) @@ -258,38 +206,25 @@ public: __device__ void ctaArrive() { - if constexpr (UseCGA) - { + int tid{0}; #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - cg::cluster_group cluster = cg::this_cluster(); - __cluster_barrier_arrive(); - if (cluster.block_rank() == 0 && threadIdx.x < kWARP_SIZE) - { - __cluster_barrier_wait(); - arriveCounter(threadIdx.x); - } + + cg::cluster_group cluster = cg::this_cluster(); + // We update the atomic counter per cluster + tid = cluster.thread_rank(); + cluster.sync(); #else - __syncthreads(); - arriveCounter(threadIdx.x); + tid = threadIdx.x; + __syncthreads(); #endif - } - else + if (tid == 0) { -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 700)) - uint32_t const barrierThreads - = ((static_cast(blockDim.x) + kWARP_SIZE - 1U) / kWARP_SIZE) * kWARP_SIZE; - if (threadIdx.x < kWARP_SIZE) - { - asm volatile("barrier.cta.sync 1, %0;" ::"r"(barrierThreads) : "memory"); - arriveCounter(threadIdx.x); - } - else - { - asm volatile("barrier.cta.arrive 1, %0;" ::"r"(barrierThreads) : "memory"); - } +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)) + asm volatile("red.async.release.global.gpu.add.u32 [%0], %1;" ::"l"(mFlagAccessPtr), "r"(1) : "memory"); +#elif (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 700)) + asm volatile("red.release.global.gpu.add.u32 [%0], %1;" ::"l"(mFlagAccessPtr), "r"(1) : "memory"); #else - __syncthreads(); - arriveCounter(threadIdx.x); + atomicAdd(mFlagAccessPtr, 1); #endif } } @@ -298,40 +233,21 @@ public: { bool isLastCtaT0{false}; int targetCount{0}; - if constexpr (UseCGA) - { #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - cg::grid_group grid = cg::this_grid(); - // Use the first thread instead of the last thread as the last thread may exit early. - isLastCtaT0 = grid.thread_rank() == 0; - targetCount = grid.num_clusters(); + cg::grid_group grid = cg::this_grid(); + // Use the first thread instead of the last thread as the last thread may exit early + isLastCtaT0 = grid.thread_rank() == 0; + targetCount = grid.num_clusters(); #else - isLastCtaT0 = threadIdx.x == 0 && blockIdx.x == 0 && blockIdx.y == 0 && blockIdx.z == 0; - targetCount = gridDim.x * gridDim.y * gridDim.z; + isLastCtaT0 = threadIdx.x == 0 && blockIdx.x == 0 && blockIdx.y == 0; + targetCount = gridDim.x * gridDim.y; #endif - } - else - { - isLastCtaT0 = threadIdx.x == 0 && blockIdx.x == 0 && blockIdx.y == 0 && blockIdx.z == 0; - targetCount = gridDim.x * gridDim.y * gridDim.z; - } if (isLastCtaT0) { uint4* flagPtr = reinterpret_cast(mBufferFlagsPtr); -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 700)) - uint32_t arrivedCount; - do - { - asm volatile("ld.acquire.gpu.global.u32 %0, [%1];" - : "=r"(arrivedCount) - : "l"(mFlagAccessPtr) - : "memory"); - } while (arrivedCount < static_cast(targetCount)); -#else - while (*reinterpret_cast(mFlagAccessPtr) < static_cast(targetCount)) + while (*reinterpret_cast(mFlagAccessPtr) < targetCount) { } -#endif // 'Current' becomes 'Dirty' flagPtr[0] = {(mCurrentIndex + 1) % 3, // Current index mCurrentIndex, // Dirty index @@ -343,20 +259,6 @@ public: } private: - __device__ void arriveCounter(int tid) - { - if (tid == 0) - { -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)) - asm volatile("red.async.release.global.gpu.add.u32 [%0], %1;" ::"l"(mFlagAccessPtr), "r"(1) : "memory"); -#elif (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 700)) - asm volatile("red.release.global.gpu.add.u32 [%0], %1;" ::"l"(mFlagAccessPtr), "r"(1) : "memory"); -#else - atomicAdd(mFlagAccessPtr, 1); -#endif - } - } - uint32_t* mBufferFlagsPtr; uint32_t* mFlagAccessPtr; diff --git a/cpp/tensorrt_llm/common/ncclUtils.cpp b/cpp/tensorrt_llm/common/ncclUtils.cpp index e36cfddcd404..6b3d4209f7e9 100644 --- a/cpp/tensorrt_llm/common/ncclUtils.cpp +++ b/cpp/tensorrt_llm/common/ncclUtils.cpp @@ -1,5 +1,5 @@ /* - * Copyright (c) 2025-2026, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -27,7 +27,7 @@ namespace { -// RAII guard for cudaMalloc. Frees the pointer on destruction, logging a warning on failure. +// RAII guard for cudaMalloc — frees the pointer on destruction, logging a warning on failure. struct CudaMallocGuard { void* ptr{nullptr}; @@ -56,7 +56,7 @@ struct CudaMallocGuard CudaMallocGuard& operator=(CudaMallocGuard const&) = delete; }; -// RAII guard for ncclMemAlloc. Frees the pointer on destruction, logging a warning on failure. +// RAII guard for ncclMemAlloc — frees the pointer on destruction, logging a warning on failure. struct NcclMemGuard { void* ptr{nullptr}; @@ -90,121 +90,6 @@ struct NcclMemGuard namespace tensorrt_llm::common::nccl_util { -namespace -{ - -#if NCCL_VERSION_CODE >= NCCL_VERSION(2, 28, 0) -constexpr int kNcclWindowMinRuntimeVersion = NCCL_VERSION(2, 28, 0); -constexpr int kNcclGb10WindowFixedVersion = NCCL_VERSION(2, 30, 4); -constexpr int kGb10RealSmVersion = 121; - -bool isGb10Platform(int realSmVersion, bool isIntegrated) -{ - return realSmVersion == kGb10RealSmVersion && isIntegrated; -} -#endif - -bool queryNcclWindowSupported() -{ -#if NCCL_VERSION_CODE >= NCCL_VERSION(2, 28, 0) - int version = 0; - if (ncclGetVersion(&version) != ncclSuccess) - { - TLLM_LOG_WARNING("[NCCLUtil] Failed to query NCCL runtime version; falling back to regular tensors."); - return false; - } - - if (version < kNcclWindowMinRuntimeVersion) - { - TLLM_LOG_WARNING( - "[NCCLUtil] NCCL runtime version %d.%d.%d does not support window buffers; falling back to regular " - "tensors.", - version / 10000, (version % 10000) / 100, version % 100); - return false; - } - - if (version >= kNcclGb10WindowFixedVersion) - { - return true; - } - - int device = -1; - cudaError_t const deviceErr = cudaGetDevice(&device); - if (deviceErr != cudaSuccess) - { - TLLM_LOG_WARNING( - "[NCCLUtil] Failed to query the current CUDA device while checking NCCL window support: %s; " - "falling back to regular tensors.", - cudaGetErrorString(deviceErr)); - return false; - } - - int isIntegrated = 0; - cudaError_t const integratedErr = cudaDeviceGetAttribute(&isIntegrated, cudaDevAttrIntegrated, device); - if (integratedErr != cudaSuccess) - { - TLLM_LOG_WARNING( - "[NCCLUtil] Failed to query CUDA integrated-device attribute for device %d while checking NCCL window " - "support: %s; falling back to regular tensors.", - device, cudaGetErrorString(integratedErr)); - return false; - } - - int realSmVersion = -1; - try - { - realSmVersion = tensorrt_llm::common::getSMVersion(/*queryRealSmArch=*/true); - } - catch (std::exception const& e) - { - TLLM_LOG_WARNING( - "[NCCLUtil] Failed to query real CUDA SM version while checking NCCL window support: %s; falling back " - "to regular tensors.", - e.what()); - return false; - } - - bool const supported = !isGb10Platform(realSmVersion, isIntegrated != 0); - if (!supported) - { - TLLM_LOG_WARNING( - "[NCCLUtil] Disabling NCCL window buffers on integrated SM %d with NCCL runtime version %d.%d.%d; " - "GB10 requires NCCL 2.30.4 or newer for symmetric window registration.", - realSmVersion, version / 10000, (version % 10000) / 100, version % 100); - } - return supported; -#else - return false; -#endif -} - -} // namespace - -bool isNcclWindowSupportedForPlatform(int realSmVersion, bool isIntegrated, int ncclRuntimeVersion) -{ -#if NCCL_VERSION_CODE >= NCCL_VERSION(2, 28, 0) - if (ncclRuntimeVersion < kNcclWindowMinRuntimeVersion) - { - return false; - } - - return !(ncclRuntimeVersion < kNcclGb10WindowFixedVersion && isGb10Platform(realSmVersion, isIntegrated)); -#else - (void) realSmVersion; - (void) isIntegrated; - (void) ncclRuntimeVersion; - return false; -#endif -} - -bool isNcclWindowSupported() -{ - static std::once_flag supportCheckFlag; - static bool windowBuffersSupported = false; - std::call_once(supportCheckFlag, []() { windowBuffersSupported = queryNcclWindowSupported(); }); - return windowBuffersSupported; -} - //============================================================================== // NcclCommResourceManager Implementation //============================================================================== @@ -402,7 +287,26 @@ NCCLWindowAllocator& NCCLWindowAllocator::getInstance() NCCLWindowBuffer NCCLWindowAllocator::requestBuffer(ncclComm_t comm, size_t size) { - if (!isNcclWindowSupported()) + // One-time runtime version check: the runtime NCCL library must also support window buffers. + static std::once_flag versionCheckFlag; + static bool runtimeVersionOk = false; + std::call_once(versionCheckFlag, + []() + { + int version = 0; + if (ncclGetVersion(&version) == ncclSuccess && version >= NCCL_VERSION(2, 28, 0)) + { + runtimeVersionOk = true; + } + else + { + TLLM_LOG_WARNING( + "[NCCLUtil] NCCL runtime version %d.%d.%d does not support window buffers; " + "falling back to regular tensors.", + version / 10000, (version % 10000) / 100, version % 100); + } + }); + if (!runtimeVersionOk) { return NCCLWindowBuffer(); } @@ -440,17 +344,6 @@ NCCLWindowBuffer NCCLWindowAllocator::requestBuffer(ncclComm_t comm, size_t size return bestFit->buffer; } - // If a previous allocateAndRegisterBuffer call collectively failed for this comm at a size - // no larger than this request, do not retry the known-failing new allocation path. Smaller - // requests and already-pooled buffers can still use NCCL windows. - auto const failureIt = mMinSymmetricFailureSize.find(comm); - if (failureIt != mMinSymmetricFailureSize.end() && size >= failureIt->second) - { - TLLM_LOG_DEBUG("[NCCLUtil] Skipping NCCL window allocation for comm %p, size=%zu; known failure threshold=%zu", - static_cast(comm), size, failureIt->second); - return NCCLWindowBuffer(); - } - // No available buffer found, avoid registration during CUDA graph capture auto stream = at::cuda::getCurrentCUDAStream(); cudaStreamCaptureStatus capture_status = cudaStreamCaptureStatusNone; @@ -471,47 +364,11 @@ NCCLWindowBuffer NCCLWindowAllocator::requestBuffer(ncclComm_t comm, size_t size "[NCCLUtil] Allocating new NCCL window buffer for comm %p, size=%zu", static_cast(comm), size); int handle = static_cast(commBuffers.size()); NCCLWindowBuffer buffer = allocateAndRegisterBuffer(comm, size, handle); - // Only cache valid buffers. allocateAndRegisterBuffer returns an empty buffer when any rank - // failed ncclMemAlloc (collective fallback to plain allreduce); caching it would leak a - // permanently "in use" empty entry per request because releaseBuffer is a no-op for nullptr. - if (buffer.isValid()) - { - commBuffers.push_back({buffer, true}); - } - else - { - // The collective allreduce inside allocateAndRegisterBuffer agreed that this request - // cannot use symmetric memory on at least one rank. Remember the smallest failing - // request size so repeated too-large autotuner probes do not keep stressing this path. - recordSymmetricFailureLocked(comm, size); - } + commBuffers.push_back({buffer, true}); return buffer; } -void NCCLWindowAllocator::recordSymmetricFailureLocked(ncclComm_t comm, size_t size) -{ - auto failureIt = mMinSymmetricFailureSize.find(comm); - if (failureIt == mMinSymmetricFailureSize.end()) - { - mMinSymmetricFailureSize.emplace(comm, size); - } - else if (size < failureIt->second) - { - failureIt->second = size; - } -} - -cudaError_t NCCLWindowAllocator::clearCudaErrorIfSymmetricAllocationFailed( - int localAllocOk, CudaGetLastErrorFunc getLastError) noexcept -{ - if (localAllocOk == 0) - { - return getLastError(); - } - return cudaSuccess; -} - NCCLWindowBuffer NCCLWindowAllocator::searchBuffer(ncclComm_t comm, void* ptr) const { if (!comm || !ptr) @@ -609,37 +466,24 @@ bool NCCLWindowAllocator::isCommValid(ncclComm_t comm) const noexcept NCCLWindowBuffer NCCLWindowAllocator::allocateAndRegisterBuffer(ncclComm_t comm, size_t size, int handle) { - // Step 1: Pre-allocate the rank-sync flag before ncclMemAlloc. ncclMemAlloc can fail - // asymmetrically with ncclUnhandledCudaError on configurations where the symmetric/VMM path - // is unavailable; that failure may leave a sticky CUDA last-error on the device. If we - // deferred this cudaMalloc until after the failure, the sticky error would propagate into - // cudaMalloc, TLLM_CUDA_CHECK would throw, and the failing rank would never reach the - // collective ncclAllReduce(min) below, hanging every other rank that did succeed. - int* rankSyncFlag = nullptr; - TLLM_CUDA_CHECK(cudaMalloc(&rankSyncFlag, sizeof(int))); - CudaMallocGuard flagGuard{rankSyncFlag}; // frees rankSyncFlag on any early return or exception - auto stream = at::cuda::getCurrentCUDAStream().stream(); - TLLM_CUDA_CHECK(cudaMemsetAsync(rankSyncFlag, 0, sizeof(int), stream)); - - // Step 2: Allocate symmetric memory. This per-rank, non-collective call can fail - // asymmetrically. When it fails, NCCL may leave a sticky CUDA error behind; clear it before - // the stream-ordered flag copy and collective fallback so the failing rank still reaches - // ncclAllReduce with the other ranks. + // Step 1: Allocate symmetric memory (per-rank, non-collective — can fail asymmetrically). void* ncclPtr = nullptr; TLLM_NCCL_CHECK_WARN(ncclMemAlloc(&ncclPtr, size)); int const localAllocOk = (ncclPtr != nullptr) ? 1 : 0; NcclMemGuard ncclGuard{ncclPtr}; // frees ncclPtr on any early return or exception - clearCudaErrorIfSymmetricAllocationFailed(localAllocOk); - // Step 3: ncclCommWindowRegister is collective. If any rank skips it, all other ranks hang. - // Populate flag, reduce with min across ranks (0 if any rank failed), then read back. - // The flag is initialized to 0, so H2D failure is non-fatal and conservatively falls back - // to regular NCCL while still reaching the collective. allreduce and D2H failures throw. - if (localAllocOk != 0) - { - TLLM_CUDA_CHECK_WARN( - cudaMemcpyAsync(rankSyncFlag, &localAllocOk, sizeof(localAllocOk), cudaMemcpyHostToDevice, stream)); - } + // Step 2: ncclCommWindowRegister is collective — if any rank skips it, all other ranks hang. + // Synchronize the per-rank alloc status using a small cudaMalloc flag (not ncclMemAlloc, so + // OOM on symmetric memory does not prevent us from allocating the flag). + int* rankSyncFlag = nullptr; + TLLM_CUDA_CHECK(cudaMalloc(&rankSyncFlag, sizeof(int))); + CudaMallocGuard flagGuard{rankSyncFlag}; // frees rankSyncFlag on any early return or exception + + // Step 3: Populate flag, reduce with min across ranks (0 if any rank failed), then read back. + // H2D failure is non-fatal: warn and continue — device flag may be stale but the allreduce + // must still be reached by all ranks. allreduce and D2H failures are catastrophic (throw). + auto stream = at::cuda::getCurrentCUDAStream().stream(); + TLLM_CUDA_CHECK_WARN(cudaMemcpy(rankSyncFlag, &localAllocOk, sizeof(int), cudaMemcpyHostToDevice)); TLLM_NCCL_CHECK(ncclAllReduce(rankSyncFlag, rankSyncFlag, 1, ncclInt32, ncclMin, comm, stream)); TLLM_CUDA_CHECK_WARN(cudaStreamSynchronize(stream)); @@ -659,7 +503,7 @@ NCCLWindowBuffer NCCLWindowAllocator::allocateAndRegisterBuffer(ncclComm_t comm, return NCCLWindowBuffer{}; // ncclGuard frees ncclPtr } - // Step 4: Register with NCCL as a window. This is collective, so all ranks must reach it. + // Step 4: Register with NCCL as a window (collective — all ranks must reach this call). // Failure here is non-fatal: warn and fall back to regular allreduce. // ncclGuard frees ncclPtr on return. ncclWindow_t window = nullptr; @@ -670,7 +514,7 @@ NCCLWindowBuffer NCCLWindowAllocator::allocateAndRegisterBuffer(ncclComm_t comm, return NCCLWindowBuffer{}; } - // Step 5: Success. Transfer ownership to the returned buffer. + // Step 5: Success — transfer ownership to the returned buffer. ncclGuard.release(); NCCLWindowBuffer buffer{ncclPtr, handle, size, window}; TLLM_LOG_TRACE("[NCCLUtil] Allocated and registered NCCL window buffer: handle=%d, ptr=%p, size=%zu, window=%p", @@ -743,7 +587,6 @@ void NCCLWindowAllocator::cleanupBuffersForComm(ncclComm_t comm) noexcept { // No buffers to clean up, but mark as cleaned mRegisteredComms.erase(comm); - mMinSymmetricFailureSize.erase(comm); return; } @@ -819,7 +662,6 @@ void NCCLWindowAllocator::cleanupBuffersForComm(ncclComm_t comm) noexcept mBufferPool.erase(commIt); mRegisteredComms.erase(comm); - mMinSymmetricFailureSize.erase(comm); } #endif // NCCL_VERSION_CODE >= NCCL_VERSION(2, 28, 0) diff --git a/cpp/tensorrt_llm/common/ncclUtils.h b/cpp/tensorrt_llm/common/ncclUtils.h index fdc88f1aaaba..f4699d71c334 100644 --- a/cpp/tensorrt_llm/common/ncclUtils.h +++ b/cpp/tensorrt_llm/common/ncclUtils.h @@ -1,5 +1,5 @@ /* - * Copyright (c) 2025-2026, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -164,14 +164,21 @@ class NcclCommResource // NCCL Version Check //============================================================================== -// Returns true if NCCL window buffers (ncclMemAlloc / ncclCommWindowRegister) -// are supported for the given real SM version, integrated-device flag, and runtime NCCL version. -// Exposed for focused unit testing of platform/version gates. -bool isNcclWindowSupportedForPlatform(int realSmVersion, bool isIntegrated, int ncclRuntimeVersion); - // Returns true if the compile-time and runtime NCCL versions support window buffers -// and the current CUDA device is not in a known-unsupported platform/version set. -bool isNcclWindowSupported(); +// (ncclMemAlloc / ncclCommWindowRegister). +inline bool isNcclWindowSupported() +{ +#if NCCL_VERSION_CODE >= NCCL_VERSION(2, 28, 0) + int version = 0; + if (ncclGetVersion(&version) != ncclSuccess) + { + return false; + } + return version >= NCCL_VERSION(2, 28, 0); +#else + return false; +#endif +} //============================================================================== // NCCL Window Buffer Allocation @@ -258,23 +265,12 @@ class NCCLWindowAllocator NCCLWindowAllocator& operator=(NCCLWindowAllocator&&) = delete; private: - friend class NCCLWindowAllocatorTestAccess; - NCCLWindowAllocator() = default; ~NCCLWindowAllocator() = default; // Allocate a new buffer and register it with NCCL as a window NCCLWindowBuffer allocateAndRegisterBuffer(ncclComm_t comm, size_t size, int handle); - // Record a failed new symmetric allocation (assumes mMutex is already locked). - void recordSymmetricFailureLocked(ncclComm_t comm, size_t size); - - using CudaGetLastErrorFunc = cudaError_t (*)(); - - // Drain the sticky CUDA error left by a failed symmetric allocation. - static cudaError_t clearCudaErrorIfSymmetricAllocationFailed( - int localAllocOk, CudaGetLastErrorFunc getLastError = cudaGetLastError) noexcept; - // Search for a buffer by pointer (assumes mMutex is already locked) NCCLWindowBuffer searchBufferLocked(ncclComm_t comm, void* ptr) const; @@ -293,10 +289,6 @@ class NCCLWindowAllocator mutable std::mutex mMutex; std::unordered_map> mBufferPool; std::unordered_set mRegisteredComms; - // Smallest request size that is known to fail collectively for each communicator. - // Requests below the recorded size may still succeed and already-pooled buffers are always - // reused before consulting this cache. - std::unordered_map mMinSymmetricFailureSize; }; // RAII wrapper for NCCL window buffers diff --git a/cpp/tensorrt_llm/common/opUtils.cpp b/cpp/tensorrt_llm/common/opUtils.cpp index ff9b57cdd099..a738e28377a1 100644 --- a/cpp/tensorrt_llm/common/opUtils.cpp +++ b/cpp/tensorrt_llm/common/opUtils.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -17,7 +17,6 @@ #include "tensorrt_llm/common/opUtils.h" #include "tensorrt_llm/common/ncclUtils.h" -#include "tensorrt_llm/runtime/ipcNvlsMemory.h" #include "tensorrt_llm/runtime/utils/mpiTags.h" #include "tensorrt_llm/runtime/utils/mpiUtils.h" @@ -161,21 +160,8 @@ std::shared_ptr getComm(std::set const& group) #else setenv("NCCL_RUNTIME_CONNECT", "0", 0); setenv("NCCL_GRAPH_REGISTER", "0", 0); - // NCCL aborts during init if it tries NVLS multicast but the fabric/IMEX - // plane can't bind it. Disable NVLS when the fabric is not usable so NCCL - // falls back to NVLink P2P. No-overwrite preserves an explicit user setting. - if (!tensorrt_llm::runtime::ipcNvlsFabricUsable()) - { - setenv("NCCL_NVLS_ENABLE", "0", 0); - } #endif // _WIN32 -#if NCCL_VERSION_CODE >= NCCL_VERSION(2, 29, 0) - ncclConfig_t config = NCCL_CONFIG_INITIALIZER; - config.graphUsageMode = 1; - NCCLCHECK_THROW(ncclCommInitRankConfig(ncclComm.get(), group.size(), id, groupRank, &config)); -#else NCCLCHECK_THROW(ncclCommInitRank(ncclComm.get(), group.size(), id, groupRank)); -#endif // NCCL_VERSION_CODE >= NCCL_VERSION(2, 29, 0) commMap[group] = ncclComm; TLLM_LOG_TRACE("%s stop for rank %d", __PRETTY_FUNCTION__, rank); return ncclComm; diff --git a/cpp/tensorrt_llm/cutlass_extensions/include/cutlass_extensions/gemm_configs.h b/cpp/tensorrt_llm/cutlass_extensions/include/cutlass_extensions/gemm_configs.h index a17b83e81473..0e817644d220 100644 --- a/cpp/tensorrt_llm/cutlass_extensions/include/cutlass_extensions/gemm_configs.h +++ b/cpp/tensorrt_llm/cutlass_extensions/include/cutlass_extensions/gemm_configs.h @@ -361,12 +361,7 @@ struct CutlassGemmConfig GROUPED_GEMM = 1u << 5, FP8_ONLY = 1u << 6, FP4_ONLY = 1u << 7, - FP8FP4_MIXED = 1u << 8, - // MXFP8xMXFP8 block-scaled MoE on SM100/103. Restricts the candidate - // tile shapes to the subset valid for the Mxf8f6f4 tensor-op (TileM=128, - // TileN in {64,128,256}); otherwise autotuning would enumerate FP8 tile - // shapes that the runtime dispatcher rejects. - MXFP8_MXFP8 = 1u << 9 + FP8FP4_MIXED = 1u << 8 }; CutlassTileConfig tile_config_sm80 = CutlassTileConfig::ChooseWithHeuristic; diff --git a/cpp/tensorrt_llm/deep_ep/CMakeLists.txt b/cpp/tensorrt_llm/deep_ep/CMakeLists.txt index e00815eed13e..562c9e7d694c 100644 --- a/cpp/tensorrt_llm/deep_ep/CMakeLists.txt +++ b/cpp/tensorrt_llm/deep_ep/CMakeLists.txt @@ -120,15 +120,6 @@ if(NOT CMAKE_CXX_COMPILER_ID STREQUAL "GNU") set(CMAKE_C_COMPILER gcc) set(CMAKE_CXX_COMPILER g++) set(CMAKE_CUDA_HOST_COMPILER g++) - # PyTorch's cmake/public/cuda.cmake (loaded transitively by - # find_package(Torch)) appends -Xcompiler=-fclang-abi-compat=17 to - # CMAKE_CUDA_FLAGS whenever the parent build is configured with Clang>=18 (see - # pytorch PR #175233). Since this subdirectory falls back to GCC for NVSHMEM - # compatibility, that Clang-only flag would be forwarded to g++ via `nvcc - # -ccbin=g++` and abort the build with: g++: error: unrecognized command-line - # option '-fclang-abi-compat=17' - string(REGEX REPLACE "-Xcompiler=-fclang-abi-compat=[0-9]+" "" - CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS}") endif() # Add nvshmem external project @@ -213,12 +204,6 @@ target_compile_options( target_compile_definitions( deep_ep_cpp_tllm PRIVATE DISABLE_AGGRESSIVE_PTX_INSTRS TORCH_EXTENSION_NAME=deep_ep_cpp_tllm) -# Newer CUDA containers provide NVSHMEM headers in the default CUDA include -# directory. DeepEP must compile against the vendored NVSHMEM headers because it -# links the vendored NVSHMEM static library below. -target_include_directories( - deep_ep_cpp_tllm BEFORE - PRIVATE ${CMAKE_CURRENT_BINARY_DIR}/nvshmem-build/src/include) target_link_libraries( deep_ep_cpp_tllm PRIVATE nvshmem_project::nvshmem ${TORCH_LIBRARIES} ${TORCH_PYTHON_LIB}) diff --git a/cpp/tensorrt_llm/executor/cacheTransceiverConfig.cpp b/cpp/tensorrt_llm/executor/cacheTransceiverConfig.cpp index 45d994a46d21..81f8085a7ae3 100644 --- a/cpp/tensorrt_llm/executor/cacheTransceiverConfig.cpp +++ b/cpp/tensorrt_llm/executor/cacheTransceiverConfig.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -23,21 +23,18 @@ namespace tensorrt_llm::executor CacheTransceiverConfig::CacheTransceiverConfig(std::optional backendType, std::optional maxNumTokens, std::optional kvTransferTimeoutMs, - std::optional kvTransferSenderFutureTimeoutMs, std::optional kvTransferPollIntervalMs) + std::optional kvTransferSenderFutureTimeoutMs) : mBackendType(backendType) , mMaxTokensInBuffer(maxNumTokens) , mKvTransferTimeoutMs(kvTransferTimeoutMs) , mKvTransferSenderFutureTimeoutMs(kvTransferSenderFutureTimeoutMs) { - setKvTransferPollIntervalMs(kvTransferPollIntervalMs); } bool CacheTransceiverConfig::operator==(CacheTransceiverConfig const& other) const { return mMaxTokensInBuffer == other.mMaxTokensInBuffer && mBackendType == other.mBackendType - && mKvTransferTimeoutMs == other.mKvTransferTimeoutMs - && mKvTransferSenderFutureTimeoutMs == other.mKvTransferSenderFutureTimeoutMs - && mKvTransferPollIntervalMs == other.mKvTransferPollIntervalMs; + && mKvTransferTimeoutMs == other.mKvTransferTimeoutMs; } void CacheTransceiverConfig::setBackendType(std::optional backendType) @@ -68,15 +65,6 @@ void CacheTransceiverConfig::setKvTransferSenderFutureTimeoutMs(std::optional kvTransferPollIntervalMs) -{ - if (kvTransferPollIntervalMs.has_value() && kvTransferPollIntervalMs.value() <= 0) - { - TLLM_THROW("kvTransferPollIntervalMs must be positive"); - } - mKvTransferPollIntervalMs = kvTransferPollIntervalMs; -} - std::optional CacheTransceiverConfig::getBackendType() const { return mBackendType; @@ -96,9 +84,4 @@ std::optional CacheTransceiverConfig::getKvTransferSenderFutureTimeoutMs() { return mKvTransferSenderFutureTimeoutMs; } - -std::optional CacheTransceiverConfig::getKvTransferPollIntervalMs() const -{ - return mKvTransferPollIntervalMs; -} } // namespace tensorrt_llm::executor diff --git a/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.cpp b/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.cpp index 2be557ac1f54..a21d470b898a 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.cpp +++ b/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.cpp @@ -189,45 +189,8 @@ void AgentConnection::send(DataContext const& ctx, void const* data, size_t size NotificationInfo notificationInfo{syncInfo}; std::stringstream ss; NotificationInfo::serialize(notificationInfo, ss); - bool const inflightCancelEnabled = common::getEnvDisaggEnableInflightCancel(); - TransferState transferState; - if (!inflightCancelEnabled) - { - transferState = status->wait(); - } - else - { - static constexpr int64_t kCancelPollTimeoutMs = 100; - transferState = TransferState::kIN_PROGRESS; - while (transferState == TransferState::kIN_PROGRESS) - { - transferState = status->wait(kCancelPollTimeoutMs); - if (transferState == TransferState::kIN_PROGRESS - && ctx.getTransferTerminate().load(std::memory_order_relaxed)) - { - bool const released = status->release(); - TLLM_LOG_WARNING( - "AgentConnection::send cancelled while transfer was in progress (ctx tag=%d, remote=%s, " - "releaseAccepted=%d)", - ctx.getTag(), mRemoteAgentName.c_str(), released); - TLLM_CHECK_WITH_INFO( - released, "AgentConnection::send cancel could not release the backend transfer handle"); - TLLM_THROW("AgentConnection::send cancelled mid-transfer"); - } - } - } + TransferState transferState = status->wait(); TLLM_CHECK_WITH_INFO(transferState == TransferState::kSUCCESS, "AgentConnection::send failed"); - if (inflightCancelEnabled && ctx.getTransferTerminate().load(std::memory_order_relaxed)) - { - bool const released = status->release(); - TLLM_LOG_WARNING( - "AgentConnection::send cancelled after transfer completed but before notify (ctx tag=%d, remote=%s, " - "releaseAccepted=%d)", - ctx.getTag(), mRemoteAgentName.c_str(), released); - TLLM_CHECK_WITH_INFO( - released, "AgentConnection::send pre-notify cancel could not release the backend transfer handle"); - TLLM_THROW("AgentConnection::send cancelled pre-notify"); - } // TODO: there is a bug in request_with_notify https://github.com/ai-dynamo/nixl/pull/252 mAgentConnectionManager->getAgent()->notifySyncMessage(mRemoteAgentName, ss.str()); } @@ -236,19 +199,11 @@ void AgentConnection::recv(DataContext const& ctx, void* data, size_t size) cons { NotificationSyncInfo syncInfo{mAgentName, ctx}; - bool const received - = mAgentConnectionManager->waitForSyncInfo(mRemoteAgentName, syncInfo, ctx.getTransferTerminate()); - if (common::getEnvDisaggEnableInflightCancel()) - { - TLLM_CHECK_WITH_INFO(received, - "AgentConnection::recv ended before receiving sync notification (ctx tag=%d, remote=%s)", ctx.getTag(), - mRemoteAgentName.c_str()); - } + mAgentConnectionManager->waitForSyncInfo(mRemoteAgentName, syncInfo, ctx.getTransferTerminate()); } void AgentConnection::sendRequestAndBufferInfo(batch_manager::RequestInfo& requestInfo, - std::vector> const& cacheBufferIds, int connectionIdx, - std::atomic const* perRequestCancel) + std::vector> const& cacheBufferIds, int connectionIdx) { TLLM_CHECK(!common::getEnvTryZCopyForKVCacheTransfer()); @@ -300,11 +255,6 @@ void AgentConnection::sendRequestAndBufferInfo(batch_manager::RequestInfo& reque std::stringstream ss; NotificationInfo notificationInfo{requestAndBufferInfo}; NotificationInfo::serialize(notificationInfo, ss); - if (common::getEnvDisaggEnableInflightCancel() && perRequestCancel != nullptr - && perRequestCancel->load(std::memory_order_relaxed)) - { - TLLM_THROW("sendRequestAndBufferInfo cancelled before notify"); - } mAgentConnectionManager->getAgent()->notifySyncMessage(mRemoteAgentName, ss.str()); } @@ -341,17 +291,9 @@ void AgentConnection::sendReadySignal(DataContext const& ctx, bool isReady) cons } bool AgentConnection::recvReadySignal(DataContext const& ctx) const -{ - return recvReadySignalWithStatus(ctx).value_or(false); -} - -std::optional AgentConnection::recvReadySignalWithStatus(DataContext const& ctx) const { ReadySignalInfo readySignalInfo{mAgentName, ctx, false}; - if (!mAgentConnectionManager->waitForReadySignal(mRemoteAgentName, readySignalInfo, ctx.getTransferTerminate())) - { - return std::nullopt; - } + mAgentConnectionManager->waitForReadySignal(mRemoteAgentName, readySignalInfo, ctx.getTransferTerminate()); return readySignalInfo.mIsReady; } @@ -750,7 +692,7 @@ int AgentConnectionManager::getDeviceId() const } template -bool AgentConnectionManager::waitForNotification( +void AgentConnectionManager::waitForNotification( std::string const& remoteAgentName, NotificationType& expectedInfo, std::atomic const& terminateFlag) { while (!terminateFlag.load()) @@ -758,7 +700,7 @@ bool AgentConnectionManager::waitForNotification( if (!mIsRunning) { - return false; + return; } updateUnhandledNotifications(); std::scoped_lock lock(mNotificationMutex); @@ -791,7 +733,7 @@ bool AgentConnectionManager::waitForNotification( { it = mUnhandledNotifications.erase(it); } - return true; + return; } } } @@ -811,7 +753,7 @@ bool AgentConnectionManager::waitForNotification( { it = mUnhandledNotifications.erase(it); } - return true; + return; } } } @@ -831,25 +773,24 @@ bool AgentConnectionManager::waitForNotification( } } } - return false; } // Explicit template instantiations -template bool AgentConnectionManager::waitForNotification( +template void AgentConnectionManager::waitForNotification( std::string const& remoteAgentName, NotificationSyncInfo& expectedInfo, std::atomic const& terminateFlag); -template bool AgentConnectionManager::waitForNotification( +template void AgentConnectionManager::waitForNotification( std::string const& remoteAgentName, ReadySignalInfo& expectedInfo, std::atomic const& terminateFlag); -bool AgentConnectionManager::waitForSyncInfo( +void AgentConnectionManager::waitForSyncInfo( std::string const& remoteAgentName, NotificationSyncInfo& syncInfo, std::atomic const& terminateFlag) { - return waitForNotification(remoteAgentName, syncInfo, terminateFlag); + waitForNotification(remoteAgentName, syncInfo, terminateFlag); } -bool AgentConnectionManager::waitForReadySignal( +void AgentConnectionManager::waitForReadySignal( std::string const& remoteAgentName, ReadySignalInfo& readySignalInfo, std::atomic const& terminateFlag) { - return waitForNotification(remoteAgentName, readySignalInfo, terminateFlag); + waitForNotification(remoteAgentName, readySignalInfo, terminateFlag); } std::string const& AgentConnectionManager::getAgentName() const diff --git a/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.h b/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.h index 410eff0248c1..25283d1341a1 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.h +++ b/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.h @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -260,15 +260,13 @@ class AgentConnection : public Connection void send(DataContext const& ctx, void const* data, size_t size) const override; void recv(DataContext const& ctx, void* data, size_t size) const override; void sendRequestAndBufferInfo(batch_manager::RequestInfo& requestInfo, - std::vector> const& cacheBufferIds, int validConnectionIdx, - std::atomic const* perRequestCancel = nullptr); - void setSenderState(std::vector cacheReceiverBufferDescs, int validSegmentIdx, + std::vector> const& cacheBufferIds, int validConnectionIdx); + void setSenderState(std::vector cacheReceiverBufferDescs, int valideSegmentIdx, std::vector> offsetRatios, std::vector bufferKinds); void setHasLoadRemoteAgent(bool hasLoadRemoteAgent); [[nodiscard]] bool hasLoadRemoteAgent() const; void sendReadySignal(DataContext const& ctx, bool isReady) const; bool recvReadySignal(DataContext const& ctx) const; - std::optional recvReadySignalWithStatus(DataContext const& ctx) const; void activateBuffer(uint8_t kind) const override; [[nodiscard]] std::optional getPreAssignedBufferId(uint8_t kind) const override; @@ -322,11 +320,11 @@ class AgentConnectionManager : public ConnectionManager [[nodiscard]] std::string const& getAgentName() const; template - bool waitForNotification( + void waitForNotification( std::string const& remoteAgentName, NotificationType& expectedInfo, std::atomic const& terminateFlag); - bool waitForSyncInfo( + void waitForSyncInfo( std::string const& remoteAgentName, NotificationSyncInfo& syncInfo, std::atomic const& terminateFlag); - bool waitForReadySignal( + void waitForReadySignal( std::string const& remoteAgentName, ReadySignalInfo& readySignalInfo, std::atomic const& terminateFlag); [[nodiscard]] bool isRunning() const override; diff --git a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/agentBindings.cpp b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/agentBindings.cpp index 040979e7197d..220868893f08 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/agentBindings.cpp +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/agentBindings.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -214,82 +214,61 @@ NB_MODULE(tensorrt_llm_transfer_agent_binding, m) .def_rw("backend_params", &kvc::BaseAgentConfig::backendParams); // BaseTransferAgent class (abstract base) - // All transfer-engine operations release the GIL: they may block on NIXL / - // UCX / network / GPU memory pinning, and holding the GIL across them - // starves Python listener / progress threads in the same process. nb::class_(m, "BaseTransferAgent") - .def("register_memory", &kvc::BaseTransferAgent::registerMemory, nb::arg("descs"), - nb::call_guard()) - .def("deregister_memory", &kvc::BaseTransferAgent::deregisterMemory, nb::arg("descs"), - nb::call_guard()) + .def("register_memory", &kvc::BaseTransferAgent::registerMemory, nb::arg("descs")) + .def("deregister_memory", &kvc::BaseTransferAgent::deregisterMemory, nb::arg("descs")) .def("load_remote_agent", nb::overload_cast(&kvc::BaseTransferAgent::loadRemoteAgent), - nb::arg("name"), nb::arg("agent_desc"), nb::call_guard()) + nb::arg("name"), nb::arg("agent_desc")) .def("load_remote_agent_by_connection", nb::overload_cast( &kvc::BaseTransferAgent::loadRemoteAgent), - nb::arg("name"), nb::arg("connection_info"), nb::call_guard()) - .def("get_local_agent_desc", &kvc::BaseTransferAgent::getLocalAgentDesc, - nb::call_guard()) - .def("invalidate_remote_agent", &kvc::BaseTransferAgent::invalidateRemoteAgent, nb::arg("name"), - nb::call_guard()) + nb::arg("name"), nb::arg("connection_info")) + .def("get_local_agent_desc", &kvc::BaseTransferAgent::getLocalAgentDesc) + .def("invalidate_remote_agent", &kvc::BaseTransferAgent::invalidateRemoteAgent, nb::arg("name")) .def( "submit_transfer_requests", [](kvc::BaseTransferAgent& self, kvc::TransferRequest const& request) { return self.submitTransferRequests(request).release(); }, - nb::arg("request"), nb::rv_policy::take_ownership, nb::call_guard(), - nb::keep_alive<0, 1>()) - .def("notify_sync_message", &kvc::BaseTransferAgent::notifySyncMessage, nb::arg("name"), - nb::arg("sync_message"), nb::call_guard()) - .def("get_notified_sync_messages", &kvc::BaseTransferAgent::getNotifiedSyncMessages, - nb::call_guard()) - .def("get_local_connection_info", &kvc::BaseTransferAgent::getLocalConnectionInfo, - nb::call_guard()) - .def("check_remote_descs", &kvc::BaseTransferAgent::checkRemoteDescs, nb::arg("name"), nb::arg("memory_descs"), - nb::call_guard()); + nb::arg("request"), nb::rv_policy::take_ownership, nb::keep_alive<0, 1>()) + .def( + "notify_sync_message", &kvc::BaseTransferAgent::notifySyncMessage, nb::arg("name"), nb::arg("sync_message")) + .def("get_notified_sync_messages", &kvc::BaseTransferAgent::getNotifiedSyncMessages) + .def("get_local_connection_info", &kvc::BaseTransferAgent::getLocalConnectionInfo) + .def("check_remote_descs", &kvc::BaseTransferAgent::checkRemoteDescs, nb::arg("name"), nb::arg("memory_descs")); #ifdef ENABLE_NIXL // NixlTransferStatus class - release GIL for blocking operations nb::class_(m, "NixlTransferStatus") .def("is_completed", &kvc::NixlTransferStatus::isCompleted, nb::call_guard()) .def("wait", &kvc::NixlTransferStatus::wait, nb::arg("timeout_ms") = -1, - nb::call_guard()) - .def("get_last_status", &kvc::NixlTransferStatus::getLastStatus) - .def("get_last_status_str", &kvc::NixlTransferStatus::getLastStatusStr); + nb::call_guard()); // NixlTransferAgent class nb::class_(m, "NixlTransferAgent") - .def(nb::init(), nb::arg("config"), nb::call_guard()) - .def("shutdown", &kvc::NixlTransferAgent::shutdown, nb::call_guard()) - .def("register_memory", &kvc::NixlTransferAgent::registerMemory, nb::arg("descs"), - nb::call_guard()) - .def("deregister_memory", &kvc::NixlTransferAgent::deregisterMemory, nb::arg("descs"), - nb::call_guard()) + .def(nb::init(), nb::arg("config")) + .def("register_memory", &kvc::NixlTransferAgent::registerMemory, nb::arg("descs")) + .def("deregister_memory", &kvc::NixlTransferAgent::deregisterMemory, nb::arg("descs")) .def("load_remote_agent", nb::overload_cast(&kvc::NixlTransferAgent::loadRemoteAgent), - nb::arg("name"), nb::arg("agent_desc"), nb::call_guard()) + nb::arg("name"), nb::arg("agent_desc")) .def("load_remote_agent_by_connection", nb::overload_cast( &kvc::NixlTransferAgent::loadRemoteAgent), - nb::arg("name"), nb::arg("connection_info"), nb::call_guard()) - .def("get_local_agent_desc", &kvc::NixlTransferAgent::getLocalAgentDesc, - nb::call_guard()) - .def("get_local_connection_info", &kvc::NixlTransferAgent::getLocalConnectionInfo, - nb::call_guard()) - .def("invalidate_remote_agent", &kvc::NixlTransferAgent::invalidateRemoteAgent, nb::arg("name"), - nb::call_guard()) + nb::arg("name"), nb::arg("connection_info")) + .def("get_local_agent_desc", &kvc::NixlTransferAgent::getLocalAgentDesc) + .def("get_local_connection_info", &kvc::NixlTransferAgent::getLocalConnectionInfo) + .def("invalidate_remote_agent", &kvc::NixlTransferAgent::invalidateRemoteAgent, nb::arg("name")) .def( "submit_transfer_requests", [](kvc::NixlTransferAgent& self, kvc::TransferRequest const& request) { return self.submitTransferRequests(request).release(); }, nb::arg("request"), nb::rv_policy::take_ownership, nb::call_guard(), nb::keep_alive<0, 1>()) - .def("notify_sync_message", &kvc::NixlTransferAgent::notifySyncMessage, nb::arg("name"), - nb::arg("sync_message"), nb::call_guard()) - .def("get_notified_sync_messages", &kvc::NixlTransferAgent::getNotifiedSyncMessages, - nb::call_guard()) - .def("check_remote_descs", &kvc::NixlTransferAgent::checkRemoteDescs, nb::arg("name"), nb::arg("memory_descs"), - nb::call_guard()); + .def( + "notify_sync_message", &kvc::NixlTransferAgent::notifySyncMessage, nb::arg("name"), nb::arg("sync_message")) + .def("get_notified_sync_messages", &kvc::NixlTransferAgent::getNotifiedSyncMessages) + .def("check_remote_descs", &kvc::NixlTransferAgent::checkRemoteDescs, nb::arg("name"), nb::arg("memory_descs")); #endif // NOTE: MooncakeTransferAgent/MooncakeTransferStatus class bindings are intentionally diff --git a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp index 711750f84f54..09dd63b0a299 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -29,13 +29,11 @@ #include #include #include -#include #include #include #include #include #include -#include #include #include #include @@ -318,35 +316,14 @@ void NixlHelper::posixFileToGpuFallback(MemoryDescs const& memoryDescs, FileDesc } } -NixlTransferStatus::NixlTransferStatus(std::weak_ptr agent, nixlXferReqH* handle) - : mWeakAgent{std::move(agent)} +NixlTransferStatus::NixlTransferStatus(nixlAgent* agent, nixlXferReqH* handle) + : mRawAgent{agent} , mHandle{handle} - , mSynchronizeHandleAccess{common::getEnvDisaggEnableInflightCancel()} { - TLLM_CHECK(!mWeakAgent.expired()); + TLLM_CHECK(mRawAgent); TLLM_CHECK(mHandle); } -NixlTransferStatus::~NixlTransferStatus() noexcept -{ - try - { - if (!release()) - { - TLLM_LOG_WARNING( - "NIXL transfer handle release failed during destruction; backend handle may remain active"); - } - } - catch (std::exception const& e) - { - TLLM_LOG_WARNING("~NixlTransferStatus: releaseXferReq threw: %s", e.what()); - } - catch (...) - { - TLLM_LOG_WARNING("~NixlTransferStatus: releaseXferReq threw unknown exception"); - } -} - [[nodiscard]] MemoryDescs NixlHelper::coalesceMemoryDescs(MemoryDescs const& descs) { auto const& descVec = descs.getDescs(); @@ -511,7 +488,7 @@ TransferState NixlTransferStatus::wait(int64_t timeout_ms) const while (true) { - auto const status = queryStatus(); + auto status = mRawAgent->getXferStatus(mHandle); if (status == NIXL_SUCCESS) { return TransferState::kSUCCESS; @@ -541,76 +518,9 @@ TransferState NixlTransferStatus::wait(int64_t timeout_ms) const } } -int NixlTransferStatus::getLastStatus() const noexcept -{ - return mLastStatus.load(std::memory_order_relaxed); -} - -std::string NixlTransferStatus::getLastStatusStr() const -{ - return nixlEnumStrings::statusStr(static_cast(getLastStatus())); -} - [[nodiscard]] bool NixlTransferStatus::isCompleted() const { - return queryStatus() == NIXL_SUCCESS; -} - -nixl_status_t NixlTransferStatus::queryStatus() const -{ - auto const query = [this]() - { - if (mHandle == nullptr) - { - mLastStatus.store(static_cast(NIXL_ERR_INVALID_PARAM), std::memory_order_relaxed); - return NIXL_ERR_INVALID_PARAM; - } - auto agent = mWeakAgent.lock(); - if (!agent) - { - // Owning agent was reset; report failure so callers don't deref a null status. - mLastStatus.store(static_cast(NIXL_ERR_INVALID_PARAM), std::memory_order_relaxed); - return NIXL_ERR_INVALID_PARAM; - } - auto const status = agent->getXferStatus(mHandle); - mLastStatus.store(static_cast(status), std::memory_order_relaxed); - return status; - }; - - if (mSynchronizeHandleAccess) - { - std::lock_guard lock(mHandleMutex); - return query(); - } - return query(); -} - -[[nodiscard]] bool NixlTransferStatus::release() -{ - std::lock_guard lock(mHandleMutex); - if (mHandle == nullptr) - { - return true; - } - - auto agent = mWeakAgent.lock(); - if (!agent) - { - mHandle = nullptr; - mLastStatus.store(static_cast(NIXL_ERR_INVALID_PARAM), std::memory_order_relaxed); - return true; - } - - auto status = agent->releaseXferReq(mHandle); - mLastStatus.store(static_cast(status), std::memory_order_relaxed); - if (status == NIXL_SUCCESS) - { - mHandle = nullptr; - return true; - } - - TLLM_LOG_WARNING("NIXL releaseXferReq failed with status: %s", nixlEnumStrings::statusStr(status).c_str()); - return false; + return mRawAgent->getXferStatus(mHandle) == NIXL_SUCCESS; } NixlTransferAgent::NixlTransferAgent(BaseAgentConfig const& config) @@ -632,7 +542,7 @@ NixlTransferAgent::NixlTransferAgent(BaseAgentConfig const& config) nixlAgentConfig nixlConfig{config.useProgThread, true, port, nixl_thread_sync_t::NIXL_THREAD_SYNC_DEFAULT, numWorker, 0, 10000, config.enableTelemetry}; mAddress = getAvailableIP() + ":" + std::to_string(port); - mRawAgent = std::make_shared(config.mName, std::move(nixlConfig)); + mRawAgent = std::make_unique(config.mName, std::move(nixlConfig)); } else { @@ -642,7 +552,7 @@ NixlTransferAgent::NixlTransferAgent(BaseAgentConfig const& config) mAddress.clear(); nixlAgentConfig nixlConfig{config.useProgThread, false, 0, nixl_thread_sync_t::NIXL_THREAD_SYNC_DEFAULT, numWorker, 0, 10000, config.enableTelemetry}; - mRawAgent = std::make_shared(config.mName, std::move(nixlConfig)); + mRawAgent = std::make_unique(config.mName, std::move(nixlConfig)); } std::string nixlBackend = common::getEnvNixlBackend(); @@ -678,12 +588,14 @@ NixlTransferAgent::NixlTransferAgent(BaseAgentConfig const& config) } mExtraParams.backends.push_back(mRawBackend); TLLM_LOG_INFO("NixlTransferAgent::NixlTransferAgent mAddress: %s", mAddress.c_str()); + mDRamSrcBuffer.resize(16); + mDRamDstBuffer.resize(16); + MemoryDescs descs{MemoryType::kDRAM, {MemoryDesc{mDRamSrcBuffer}, MemoryDesc{mDRamDstBuffer}}}; + registerMemory(descs); } void NixlTransferAgent::registerMemory(RegisterDescs const& descs) { - std::unique_lock lock(mLock); - TLLM_CHECK_WITH_INFO(!mShutdown.load(), "NixlTransferAgent::registerMemory called after shutdown"); // Split VRAM descriptors at VMM chunk boundaries so each sub-descriptor // falls within a single cuMemCreate allocation (required by gdr_copy / cuda_ipc). size_t detectedChunkSize = 0; @@ -708,8 +620,6 @@ void NixlTransferAgent::registerMemory(RegisterDescs const& descs) void NixlTransferAgent::deregisterMemory(RegisterDescs const& descs) { - std::unique_lock lock(mLock); - TLLM_CHECK_WITH_INFO(!mShutdown.load(), "NixlTransferAgent::deregisterMemory called after shutdown"); // Split using per-region registry info to match what was registered auto splitDescs = VmmDescSplitter::splitDescsWithRegionMap(descs, mLocalVramRegionInfo); @@ -733,8 +643,6 @@ void NixlTransferAgent::deregisterMemory(RegisterDescs const& descs) void NixlTransferAgent::loadRemoteAgent(std::string const& name, AgentDesc const& agentDesc) { - std::unique_lock lock(mLock); - TLLM_CHECK_WITH_INFO(!mShutdown.load(), "NixlTransferAgent::loadRemoteAgent called after shutdown"); nixl_status_t status; std::string remoteName; status = mRawAgent->loadRemoteMD(agentDesc.getBackendAgentDesc(), remoteName); @@ -758,8 +666,6 @@ void NixlTransferAgent::loadRemoteAgent(std::string const& name, AgentDesc const AgentDesc NixlTransferAgent::getLocalAgentDesc() { - std::shared_lock lock(mLock); - TLLM_CHECK_WITH_INFO(!mShutdown.load(), "NixlTransferAgent::getLocalAgentDesc called after shutdown"); nixl_blob_t nixlBlob; nixl_status_t status = mRawAgent->getLocalMD(nixlBlob); TLLM_CHECK(status == NIXL_SUCCESS); @@ -779,12 +685,6 @@ AgentDesc NixlTransferAgent::getLocalAgentDesc() void NixlTransferAgent::invalidateRemoteAgent(std::string const& name) { - std::unique_lock lock(mLock); - if (mShutdown.load()) - { - // shutdown() already cleaned everything; treat as no-op for late callers. - return; - } // Clean up remote VMM region info before invalidating the remote agent. mRemoteVramRegionInfo.erase(name); mRawAgent->invalidateRemoteMD(name); @@ -792,22 +692,18 @@ void NixlTransferAgent::invalidateRemoteAgent(std::string const& name) [[nodiscard]] std::unique_ptr NixlTransferAgent::submitTransferRequests(TransferRequest const& request) { - std::shared_lock lock(mLock); - TLLM_CHECK_WITH_INFO(!mShutdown.load(), "NixlTransferAgent::submitTransferRequests called after shutdown"); nixl_status_t status; nixlXferReqH* handle; - // Local per-request copy: hasNotif / notifMsg vary per call; a shared mExtraParams - // would race between concurrent submits even under shared_lock. - nixl_opt_args_t reqParams = mExtraParams; if (request.getSyncMessage().has_value()) { - reqParams.hasNotif = true; - reqParams.notifMsg = request.getSyncMessage().value(); + mExtraParams.hasNotif = true; + + mExtraParams.notifMsg = request.getSyncMessage().value(); } else { - reqParams.hasNotif = false; + mExtraParams.hasNotif = false; } // Split transfer descriptors at VMM chunk boundaries to match registered memory. // Both src and dst are split at chunk boundaries to ensure each descriptor @@ -829,12 +725,12 @@ void NixlTransferAgent::invalidateRemoteAgent(std::string const& name) auto [coalescedSrc, coalescedDst] = NixlHelper::coalesceTransferDescs(splitSrc, splitDst); status = mRawAgent->createXferReq(NixlHelper::convert(request.getOp()), NixlHelper::convertXferDist(coalescedSrc), - NixlHelper::convertXferDist(coalescedDst), request.getRemoteName(), handle, &reqParams); + NixlHelper::convertXferDist(coalescedDst), request.getRemoteName(), handle, &mExtraParams); } else { status = mRawAgent->createXferReq(NixlHelper::convert(request.getOp()), NixlHelper::convertXferDist(splitSrc), - NixlHelper::convertXferDist(splitDst), request.getRemoteName(), handle, &reqParams); + NixlHelper::convertXferDist(splitDst), request.getRemoteName(), handle, &mExtraParams); } TLLM_CHECK_WITH_INFO(status == NIXL_SUCCESS, @@ -843,15 +739,14 @@ void NixlTransferAgent::invalidateRemoteAgent(std::string const& name) request.getRemoteName().c_str()); { NVTX3_SCOPED_RANGE(postXferReq); - status = mRawAgent->postXferReq(handle, &reqParams); + status = mRawAgent->postXferReq(handle, &mExtraParams); } - return std::make_unique(std::weak_ptr(mRawAgent), handle); + return std::make_unique(mRawAgent.get(), handle); } void NixlTransferAgent::notifySyncMessage(std::string const& name, SyncMessage const& syncMessage) { - std::shared_lock lock(mLock); - TLLM_CHECK_WITH_INFO(!mShutdown.load(), "NixlTransferAgent::notifySyncMessage called after shutdown"); + auto status = mRawAgent->genNotif(name, syncMessage); TLLM_CHECK_WITH_INFO( status == NIXL_SUCCESS, "genNotif failed with status: %s", nixlEnumStrings::statusStr(status).c_str()); @@ -859,8 +754,7 @@ void NixlTransferAgent::notifySyncMessage(std::string const& name, SyncMessage c [[nodiscard]] std::unordered_map> NixlTransferAgent::getNotifiedSyncMessages() { - std::shared_lock lock(mLock); - TLLM_CHECK_WITH_INFO(!mShutdown.load(), "NixlTransferAgent::getNotifiedSyncMessages called after shutdown"); + nixl_notifs_t notifs; auto status = mRawAgent->getNotifs(notifs); TLLM_CHECK_WITH_INFO( @@ -871,14 +765,11 @@ void NixlTransferAgent::notifySyncMessage(std::string const& name, SyncMessage c ConnectionInfoType NixlTransferAgent::getLocalConnectionInfo() { - // mAddress is set in ctor and never mutated; no lock needed. return mAddress; } void NixlTransferAgent::loadRemoteAgent(std::string const& name, ConnectionInfoType const& connectionInfo) { - std::unique_lock lock(mLock); - TLLM_CHECK_WITH_INFO(!mShutdown.load(), "NixlTransferAgent::loadRemoteAgent called after shutdown"); std::string ip = connectionInfo.substr(0, connectionInfo.find(":")); std::string port = connectionInfo.substr(connectionInfo.find(":") + 1); TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), @@ -915,67 +806,15 @@ void NixlTransferAgent::loadRemoteAgent(std::string const& name, ConnectionInfoT bool NixlTransferAgent::checkRemoteDescs(std::string const& name, MemoryDescs const& memoryDescs) { - std::shared_lock lock(mLock); - TLLM_CHECK_WITH_INFO(!mShutdown.load(), "NixlTransferAgent::checkRemoteDescs called after shutdown"); auto status = mRawAgent->checkRemoteMD(name, NixlHelper::convertXferDist(memoryDescs)); TLLM_CHECK_WITH_INFO(status == NIXL_SUCCESS || status == NIXL_ERR_NOT_FOUND, "checkRemoteMD failed with status: %s", nixlEnumStrings::statusStr(status).c_str()); return status == NIXL_SUCCESS; } -void NixlTransferAgent::shutdown() noexcept -{ - // unique_lock drains all in-flight shared_lock holders (submit / getDesc / etc.). - // A concurrent second shutdown() blocks here, then sees mShutdown=true and returns. - std::unique_lock lock(mLock); - if (mShutdown.exchange(true)) - { - return; - } - TLLM_LOG_DEBUG("NixlTransferAgent::shutdown"); - - if (mRawAgent) - { - // Inline invalidate: invalidateRemoteAgent() would re-enter the non-recursive lock. - for (auto const& [name, _] : mRemoteVramRegionInfo) - { - try - { - mRawAgent->invalidateRemoteMD(name); - } - catch (std::exception const& e) - { - TLLM_LOG_WARNING( - "NixlTransferAgent::shutdown: invalidateRemoteMD(%s) threw: %s", name.c_str(), e.what()); - } - catch (...) - { - } - } - } - - mExtraParams.backends.clear(); - mRawBackend = nullptr; - try - { - mRawAgent.reset(); - } - catch (std::exception const& e) - { - TLLM_LOG_WARNING("NixlTransferAgent::shutdown: ~nixlAgent threw: %s", e.what()); - } - catch (...) - { - TLLM_LOG_WARNING("NixlTransferAgent::shutdown: ~nixlAgent threw unknown exception"); - } - mLocalVramRegionInfo.clear(); - mRemoteVramRegionInfo.clear(); -} - NixlTransferAgent::~NixlTransferAgent() { TLLM_LOG_DEBUG("NixlTransferAgent::~NixlTransferAgent"); - shutdown(); } NixlLoopbackAgent::NixlLoopbackAgent(BaseAgentConfig const& config) @@ -986,7 +825,7 @@ NixlLoopbackAgent::NixlLoopbackAgent(BaseAgentConfig const& config) nixl_status_t status; nixl_b_params_t init; - mRawAgent = std::make_shared(config.mName, std::move(nixlConfig)); + mRawAgent = std::make_unique(config.mName, std::move(nixlConfig)); init["batch_pool_size"] = std::to_string(8); init["batch_limit"] = std::to_string(128); init["max_request_size"] = std::to_string(16 * 1024 * 1024); @@ -1005,33 +844,6 @@ NixlLoopbackAgent::NixlLoopbackAgent(BaseAgentConfig const& config) } } -void NixlLoopbackAgent::shutdown() noexcept -{ - // unique_lock drains all in-flight shared_lock holders before destroying the agent. - std::unique_lock lock(mLock); - if (mShutdown.exchange(true)) - { - return; - } - try - { - mRawAgent.reset(); - } - catch (std::exception const& e) - { - TLLM_LOG_WARNING("NixlLoopbackAgent::shutdown: ~nixlAgent threw: %s", e.what()); - } - catch (...) - { - TLLM_LOG_WARNING("NixlLoopbackAgent::shutdown: ~nixlAgent threw unknown exception"); - } -} - -NixlLoopbackAgent::~NixlLoopbackAgent() -{ - shutdown(); -} - int NixlLoopbackAgent::registerMemory(MemoryDescs const& descs) { nixl_status_t status = mRawAgent->registerMem(NixlHelper::convertRegDlist(descs)); @@ -1083,14 +895,12 @@ std::unique_ptr NixlLoopbackAgent::submitLoopbackRequests( status = mRawAgent->postXferReq(handle); TLLM_CHECK(status == NIXL_IN_PROG); - return std::make_unique(std::weak_ptr(mRawAgent), handle); + return std::make_unique(mRawAgent.get(), handle); } void NixlLoopbackAgent::executeLoopbackRequest( MemoryDescs const& memoryDescs, FileDescs const& fileDescs, bool isOffload) { - std::shared_lock lock(mLock); - TLLM_CHECK_WITH_INFO(!mShutdown.load(), "NixlLoopbackAgent::executeLoopbackRequest called after shutdown"); bool fallback = false; int ret; diff --git a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.h b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.h index cb371f02439e..50cf672400ce 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.h +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.h @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -20,9 +20,6 @@ #include "nixl.h" #include "tensorrt_llm/executor/transferAgent.h" #include -#include -#include -#include #include namespace tensorrt_llm::executor::kv_cache @@ -59,32 +56,15 @@ struct NixlHelper class NixlTransferStatus final : public TransferStatus { public: - NixlTransferStatus(std::weak_ptr agent, nixlXferReqH* handle); - ~NixlTransferStatus() noexcept override; - - NixlTransferStatus(NixlTransferStatus const&) = delete; - NixlTransferStatus& operator=(NixlTransferStatus const&) = delete; - NixlTransferStatus(NixlTransferStatus&&) = delete; - NixlTransferStatus& operator=(NixlTransferStatus&&) = delete; + NixlTransferStatus(nixlAgent* agent, nixlXferReqH* handle); [[nodiscard]] bool isCompleted() const override; [[nodiscard]] TransferState wait(int64_t timeout_ms = -1) const override; - [[nodiscard]] int getLastStatus() const noexcept; - [[nodiscard]] std::string getLastStatusStr() const; - - [[nodiscard]] bool release() override; - private: - [[nodiscard]] nixl_status_t queryStatus() const; - - // weak_ptr so the status outliving the owning agent is safe (lock() returns null after reset). - std::weak_ptr mWeakAgent; + nixlAgent* mRawAgent{}; nixlXferReqH* mHandle{}; - mutable std::atomic mLastStatus{0}; - bool const mSynchronizeHandleAccess; - mutable std::mutex mHandleMutex; }; class NixlTransferAgent final : public BaseTransferAgent @@ -93,9 +73,6 @@ class NixlTransferAgent final : public BaseTransferAgent NixlTransferAgent(BaseAgentConfig const& config); ~NixlTransferAgent(); - /// Synchronously release NIXL agent / UCX / prog_thread. Idempotent. - void shutdown() noexcept; - void registerMemory(RegisterDescs const& descs) override; void deregisterMemory(RegisterDescs const& descs) override; @@ -129,18 +106,14 @@ class NixlTransferAgent final : public BaseTransferAgent bool checkRemoteDescs(std::string const& name, MemoryDescs const& memoryDescs) override; private: - // shared_ptr so outstanding NixlTransferStatus (via weak_ptr) can detect agent reset. - std::shared_ptr mRawAgent; + std::unique_ptr mRawAgent; nixlBackendH* mRawBackend{}; nixl_opt_args_t mExtraParams; std::string mName; std::string mAddress; - std::atomic mShutdown{false}; - /// Serializes (a) wrapper-map mutations vs reads and (b) drain-on-shutdown. - /// Writers (register/deregister/load/invalidate/shutdown) take unique_lock; - /// readers (submit / getLocalAgentDesc / checkRemoteDescs / etc.) take shared_lock. - mutable std::shared_mutex mLock; + std::vector mDRamSrcBuffer; + std::vector mDRamDstBuffer; /// Local VMM region info (from registerMemory). Keyed by local virtual address. VramRegionMap mLocalVramRegionInfo; @@ -154,10 +127,7 @@ class NixlLoopbackAgent final : public BaseLoopbackAgent { public: NixlLoopbackAgent(BaseAgentConfig const& config); - ~NixlLoopbackAgent() override; - - /// Synchronously release the NIXL agent. Idempotent; drains in-flight requests. - void shutdown() noexcept; + virtual ~NixlLoopbackAgent() = default; virtual void executeLoopbackRequest( MemoryDescs const& memoryDescs, FileDescs const& fileDescs, bool isOffload) override; @@ -171,11 +141,8 @@ class NixlLoopbackAgent final : public BaseLoopbackAgent [[nodiscard]] std::unique_ptr submitLoopbackRequests( MemoryDescs const& memoryDescs, FileDescs const& filedescs, bool isOffload); - std::shared_ptr mRawAgent; + std::unique_ptr mRawAgent; std::string mName; - std::atomic mShutdown{false}; - /// Drain-on-shutdown: executeLoopbackRequest takes shared_lock; shutdown takes unique_lock. - mutable std::shared_mutex mLock; }; #if defined(__clang__) diff --git a/cpp/tensorrt_llm/executor/cache_transmission/rnnCacheSplitConcat.cu b/cpp/tensorrt_llm/executor/cache_transmission/rnnCacheSplitConcat.cu index fb6837e8d188..b697f2eeb3d6 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/rnnCacheSplitConcat.cu +++ b/cpp/tensorrt_llm/executor/cache_transmission/rnnCacheSplitConcat.cu @@ -964,7 +964,8 @@ void concatRnnSsmStateDispatch(std::vector const& i // SSM portion: [numHeads, headDim, dState] — split by heads // Conv portion: [section0_dim, section1_dim, ...] x [dConv-1] — section-aware split // -// The input/output pointer arrays follow the pattern: input pointers → output pointers → prefixLayerNum. +// The input/output pointer arrays follow the same pattern as the existing +// RnnStateManager kernels: input pointers → output pointers → prefixLayerNum. /** * @brief Kernel to split SSM state from unified pool blocks to per-target buffers. diff --git a/cpp/tensorrt_llm/executor/dynamicBatchConfig.cpp b/cpp/tensorrt_llm/executor/dynamicBatchConfig.cpp index 6340d1fa48bd..766a91c7842c 100644 --- a/cpp/tensorrt_llm/executor/dynamicBatchConfig.cpp +++ b/cpp/tensorrt_llm/executor/dynamicBatchConfig.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -49,14 +49,6 @@ std::vector> DynamicBatchConfig::getBatchSizeT return mBatchSizeTable; } -bool DynamicBatchConfig::operator==(DynamicBatchConfig const& other) const -{ - return mEnableBatchSizeTuning == other.mEnableBatchSizeTuning - && mEnableMaxNumTokensTuning == other.mEnableMaxNumTokensTuning - && mDynamicBatchMovingAverageWindow == other.mDynamicBatchMovingAverageWindow - && mBatchSizeTable == other.mBatchSizeTable; -} - std::vector> const DynamicBatchConfig::kDefaultBatchSizeTable{ {144, 128}, {336, 256}, diff --git a/cpp/tensorrt_llm/executor/executorConfig.cpp b/cpp/tensorrt_llm/executor/executorConfig.cpp index 25602a1a83d6..2dff78280f5a 100644 --- a/cpp/tensorrt_llm/executor/executorConfig.cpp +++ b/cpp/tensorrt_llm/executor/executorConfig.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -65,8 +65,8 @@ ExecutorConfig::ExecutorConfig(SizeType32 maxBeamWidth, SchedulerConfig schedule , mEnableTrtOverlap(enableTrtOverlap) , mFailFastOnAttentionWindowTooLarge(failFastOnAttentionWindowTooLarge) { - TLLM_CHECK(iterStatsMaxIterations >= kUnlimitedStatsMaxIterations); - TLLM_CHECK(requestStatsMaxIterations >= kUnlimitedStatsMaxIterations); + TLLM_CHECK(iterStatsMaxIterations >= 0); + TLLM_CHECK(requestStatsMaxIterations >= 0); TLLM_CHECK(mMaxBeamWidth > 0); TLLM_CHECK(maxSeqIdleMicroseconds > 0); } @@ -271,13 +271,13 @@ void ExecutorConfig::setNormalizeLogProbs(bool normalizeLogProbs) void ExecutorConfig::setIterStatsMaxIterations(SizeType32 iterStatsMaxIterations) { mIterStatsMaxIterations = iterStatsMaxIterations; - TLLM_CHECK(mIterStatsMaxIterations >= kUnlimitedStatsMaxIterations); + TLLM_CHECK(mIterStatsMaxIterations >= 0); } void ExecutorConfig::setRequestStatsMaxIterations(SizeType32 requestStatsMaxIterations) { mRequestStatsMaxIterations = requestStatsMaxIterations; - TLLM_CHECK(mRequestStatsMaxIterations >= kUnlimitedStatsMaxIterations); + TLLM_CHECK(mRequestStatsMaxIterations >= 0); } void ExecutorConfig::setBatchingType(BatchingType batchingType) diff --git a/cpp/tensorrt_llm/executor/executorImpl.cpp b/cpp/tensorrt_llm/executor/executorImpl.cpp index 9f7fb654a2d5..2fb20b7ea572 100644 --- a/cpp/tensorrt_llm/executor/executorImpl.cpp +++ b/cpp/tensorrt_llm/executor/executorImpl.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -86,16 +86,6 @@ namespace return fixedExecutorConfig; } -[[nodiscard]] bool statsBufferIsEnabled(SizeType32 maxIterations) -{ - return maxIterations != 0; -} - -[[nodiscard]] bool statsBufferIsBounded(SizeType32 maxIterations) -{ - return maxIterations > 0; -} - SizeType32 getNumChildRequests(Request const& request) { auto samplingConfig = request.getSamplingConfig(); @@ -1918,13 +1908,9 @@ RequestStatsPerIteration Executor::Impl::getCurrentRequestStats( void Executor::Impl::appendCurrentIterStats(IterationStats&& currentIterStats) { std::scoped_lock lck(mIterStatsMtx); - if (statsBufferIsBounded(mIterStatsMaxIterations)) + if (mIterationStats.size() >= mIterStatsMaxIterations) { - auto const maxIterStats = static_cast(mIterStatsMaxIterations); - if (mIterationStats.size() >= maxIterStats) - { - mIterationStats.pop_front(); - } + mIterationStats.pop_front(); } mIterationStats.emplace_back(std::move(currentIterStats)); } @@ -1932,16 +1918,16 @@ void Executor::Impl::appendCurrentIterStats(IterationStats&& currentIterStats) void Executor::Impl::appendMultipleIterStats(std::vector&& currentIterStatsVec) { std::scoped_lock lck(mIterStatsMtx); - mIterationStats.insert(mIterationStats.end(), std::make_move_iterator(currentIterStatsVec.begin()), - std::make_move_iterator(currentIterStatsVec.end())); - if (statsBufferIsBounded(mIterStatsMaxIterations)) + if (mIterationStats.size() + currentIterStatsVec.size() > mIterStatsMaxIterations) { - auto const maxIterStats = static_cast(mIterStatsMaxIterations); - while (mIterationStats.size() > maxIterStats) + size_t removeCount = mIterationStats.size() + currentIterStatsVec.size() - mIterStatsMaxIterations; + for (size_t i = 0; i < removeCount; i++) { mIterationStats.pop_front(); } } + mIterationStats.insert(mIterationStats.end(), std::make_move_iterator(currentIterStatsVec.begin()), + std::make_move_iterator(currentIterStatsVec.end())); } void Executor::Impl::updateIterationStats(RequestList const& activeRequests, double iterLatencyMS, @@ -1949,7 +1935,7 @@ void Executor::Impl::updateIterationStats(RequestList const& activeRequests, dou bool flushToOrchestrator) { NVTX3_SCOPED_RANGE(updateIterationStats); - if (statsBufferIsEnabled(mIterStatsMaxIterations) && mIsLeader) + if (mIterStatsMaxIterations > 0 && mIsLeader) { auto currentIterStats = getCurrentIterationStats( activeRequests, iterLatencyMS, numNewActiveRequests, newActiveRequestsQueueLatencyMS, numCompletedRequests); @@ -1986,13 +1972,9 @@ void Executor::Impl::updateIterationStats(RequestList const& activeRequests, dou void Executor::Impl::appendCurrentRequestStats(RequestStatsPerIteration&& currentRequestStats) { std::scoped_lock lck(mRequestStatsMtx); - if (statsBufferIsBounded(mRequestStatsMaxIterations)) + if (mRequestStats.size() >= mRequestStatsMaxIterations) { - auto const maxRequestStats = static_cast(mRequestStatsMaxIterations); - if (mRequestStats.size() >= maxRequestStats) - { - mRequestStats.pop_front(); - } + mRequestStats.pop_front(); } mRequestStats.emplace_back(std::move(currentRequestStats)); } @@ -2000,23 +1982,23 @@ void Executor::Impl::appendCurrentRequestStats(RequestStatsPerIteration&& curren void Executor::Impl::appendMultipleRequestStats(std::vector&& currentRequestStatsVec) { std::scoped_lock lck(mRequestStatsMtx); - mRequestStats.insert(mRequestStats.end(), std::make_move_iterator(currentRequestStatsVec.begin()), - std::make_move_iterator(currentRequestStatsVec.end())); - if (statsBufferIsBounded(mRequestStatsMaxIterations)) + if (mRequestStats.size() + currentRequestStatsVec.size() > mRequestStatsMaxIterations) { - auto const maxRequestStats = static_cast(mRequestStatsMaxIterations); - while (mRequestStats.size() > maxRequestStats) + size_t removeCount = mRequestStats.size() + currentRequestStatsVec.size() - mRequestStatsMaxIterations; + for (size_t i = 0; i < removeCount; i++) { mRequestStats.pop_front(); } } + mRequestStats.insert(mRequestStats.end(), std::make_move_iterator(currentRequestStatsVec.begin()), + std::make_move_iterator(currentRequestStatsVec.end())); } void Executor::Impl::updateRequestStats( RequestList const& activeRequests, RequestList const& finishedRequests, bool flushToOrchestrator) { NVTX3_SCOPED_RANGE(updateRequestStats); - if (statsBufferIsEnabled(mRequestStatsMaxIterations) && mIsLeader) + if (mRequestStatsMaxIterations > 0 && mIsLeader) { // Add current iteration request stats auto currentRequestStats = getCurrentRequestStats(activeRequests, finishedRequests); diff --git a/cpp/tensorrt_llm/executor/executorImpl.h b/cpp/tensorrt_llm/executor/executorImpl.h index f812b55a3fa0..6e545dbf6def 100644 --- a/cpp/tensorrt_llm/executor/executorImpl.h +++ b/cpp/tensorrt_llm/executor/executorImpl.h @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -310,12 +310,12 @@ class Executor::Impl std::unordered_map> mChildReqIdsMap; // Iteration stats - SizeType32 mIterStatsMaxIterations; + IterationType mIterStatsMaxIterations; std::mutex mIterStatsMtx; std::deque mIterationStats; // Request stats - SizeType32 mRequestStatsMaxIterations; + IterationType mRequestStatsMaxIterations; std::mutex mRequestStatsMtx; std::deque mRequestStats; diff --git a/cpp/tensorrt_llm/executor/request.cpp b/cpp/tensorrt_llm/executor/request.cpp index 4d7b1d909c33..e32045892ba7 100644 --- a/cpp/tensorrt_llm/executor/request.cpp +++ b/cpp/tensorrt_llm/executor/request.cpp @@ -79,11 +79,6 @@ VecTokens Request::getInputTokenIds() const return mImpl->getInputTokenIds(); } -SizeType32 Request::getNumInputTokens() const -{ - return mImpl->getNumInputTokens(); -} - SizeType32 Request::getMaxTokens() const { return mImpl->getMaxNewTokens(); diff --git a/cpp/tensorrt_llm/executor/requestImpl.h b/cpp/tensorrt_llm/executor/requestImpl.h index 145336f77f9a..55610885b1ae 100644 --- a/cpp/tensorrt_llm/executor/requestImpl.h +++ b/cpp/tensorrt_llm/executor/requestImpl.h @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -135,11 +135,6 @@ class Request::Impl return mInputTokenIds; } - [[nodiscard]] SizeType32 getNumInputTokens() const - { - return static_cast(mInputTokenIds.size()); - } - [[nodiscard]] SizeType32 getMaxNewTokens() const { return mMaxNewTokens; diff --git a/cpp/tensorrt_llm/executor/schedulerConfig.cpp b/cpp/tensorrt_llm/executor/schedulerConfig.cpp index c95d543c53a4..20aa5976455a 100644 --- a/cpp/tensorrt_llm/executor/schedulerConfig.cpp +++ b/cpp/tensorrt_llm/executor/schedulerConfig.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -21,20 +21,17 @@ namespace tensorrt_llm::executor { SchedulerConfig::SchedulerConfig(CapacitySchedulerPolicy capacitySchedulerPolicy, - std::optional contextChunkingPolicy, std::optional dynamicBatchConfig, - bool enablePrefixAwareScheduling) + std::optional contextChunkingPolicy, std::optional dynamicBatchConfig) : mCapacitySchedulerPolicy(capacitySchedulerPolicy) , mContextChunkingPolicy(std::move(contextChunkingPolicy)) , mDynamicBatchConfig(std::move(dynamicBatchConfig)) - , mEnablePrefixAwareScheduling(enablePrefixAwareScheduling) { } bool SchedulerConfig::operator==(SchedulerConfig const& other) const { return mCapacitySchedulerPolicy == other.mCapacitySchedulerPolicy - && mContextChunkingPolicy == other.mContextChunkingPolicy && mDynamicBatchConfig == other.mDynamicBatchConfig - && mEnablePrefixAwareScheduling == other.mEnablePrefixAwareScheduling; + && mContextChunkingPolicy == other.mContextChunkingPolicy; } [[nodiscard]] CapacitySchedulerPolicy SchedulerConfig::getCapacitySchedulerPolicy() const @@ -52,9 +49,4 @@ bool SchedulerConfig::operator==(SchedulerConfig const& other) const return mDynamicBatchConfig; } -[[nodiscard]] bool SchedulerConfig::getEnablePrefixAwareScheduling() const -{ - return mEnablePrefixAwareScheduling; -} - } // namespace tensorrt_llm::executor diff --git a/cpp/tensorrt_llm/executor/serialization.cpp b/cpp/tensorrt_llm/executor/serialization.cpp index 020306e03e56..ce081e10c603 100644 --- a/cpp/tensorrt_llm/executor/serialization.cpp +++ b/cpp/tensorrt_llm/executor/serialization.cpp @@ -1427,9 +1427,7 @@ SchedulerConfig Serialization::deserializeSchedulerConfig(std::istream& is) auto capacitySchedulerPolicy = su::deserialize(is); auto contextChunkingPolicy = su::deserialize>(is); auto dynamicBatchConfig = su::deserialize>(is); - auto enablePrefixAwareScheduling = su::deserialize(is); - return SchedulerConfig{ - capacitySchedulerPolicy, contextChunkingPolicy, dynamicBatchConfig, enablePrefixAwareScheduling}; + return SchedulerConfig{capacitySchedulerPolicy, contextChunkingPolicy, dynamicBatchConfig}; } void Serialization::serialize(SchedulerConfig const& schedulerConfig, std::ostream& os) @@ -1437,7 +1435,6 @@ void Serialization::serialize(SchedulerConfig const& schedulerConfig, std::ostre su::serialize(schedulerConfig.getCapacitySchedulerPolicy(), os); su::serialize(schedulerConfig.getContextChunkingPolicy(), os); su::serialize(schedulerConfig.getDynamicBatchConfig(), os); - su::serialize(schedulerConfig.getEnablePrefixAwareScheduling(), os); } size_t Serialization::serializedSize(SchedulerConfig const& schedulerConfig) @@ -1446,7 +1443,6 @@ size_t Serialization::serializedSize(SchedulerConfig const& schedulerConfig) totalSize += su::serializedSize(schedulerConfig.getCapacitySchedulerPolicy()); totalSize += su::serializedSize(schedulerConfig.getContextChunkingPolicy()); totalSize += su::serializedSize(schedulerConfig.getDynamicBatchConfig()); - totalSize += su::serializedSize(schedulerConfig.getEnablePrefixAwareScheduling()); return totalSize; } @@ -1457,9 +1453,7 @@ CacheTransceiverConfig Serialization::deserializeCacheTransceiverConfig(std::ist auto maxTokensInBuffer = su::deserialize>(is); auto kvTransferTimeoutMs = su::deserialize>(is); auto kvTransferSenderFutureTimeoutMs = su::deserialize>(is); - auto kvTransferPollIntervalMs = su::deserialize>(is); - return CacheTransceiverConfig{ - backendType, maxTokensInBuffer, kvTransferTimeoutMs, kvTransferSenderFutureTimeoutMs, kvTransferPollIntervalMs}; + return CacheTransceiverConfig{backendType, maxTokensInBuffer, kvTransferTimeoutMs, kvTransferSenderFutureTimeoutMs}; } void Serialization::serialize(CacheTransceiverConfig const& cacheTransceiverConfig, std::ostream& os) @@ -1468,7 +1462,6 @@ void Serialization::serialize(CacheTransceiverConfig const& cacheTransceiverConf su::serialize(cacheTransceiverConfig.getMaxTokensInBuffer(), os); su::serialize(cacheTransceiverConfig.getKvTransferTimeoutMs(), os); su::serialize(cacheTransceiverConfig.getKvTransferSenderFutureTimeoutMs(), os); - su::serialize(cacheTransceiverConfig.getKvTransferPollIntervalMs(), os); } size_t Serialization::serializedSize(CacheTransceiverConfig const& cacheTransceiverConfig) @@ -1478,7 +1471,6 @@ size_t Serialization::serializedSize(CacheTransceiverConfig const& cacheTranscei totalSize += su::serializedSize(cacheTransceiverConfig.getMaxTokensInBuffer()); totalSize += su::serializedSize(cacheTransceiverConfig.getKvTransferTimeoutMs()); totalSize += su::serializedSize(cacheTransceiverConfig.getKvTransferSenderFutureTimeoutMs()); - totalSize += su::serializedSize(cacheTransceiverConfig.getKvTransferPollIntervalMs()); return totalSize; } diff --git a/cpp/tensorrt_llm/flash_mla/CMakeLists.txt b/cpp/tensorrt_llm/flash_mla/CMakeLists.txt index 024cee521d10..e87f12275f60 100644 --- a/cpp/tensorrt_llm/flash_mla/CMakeLists.txt +++ b/cpp/tensorrt_llm/flash_mla/CMakeLists.txt @@ -44,15 +44,6 @@ if(CMAKE_SYSTEM_PROCESSOR MATCHES "aarch64|arm64" AND CMAKE_CXX_COMPILER_ID set(CMAKE_CUDA_HOST_COMPILER ${GCC_EXECUTABLE}) message( STATUS "FlashMLA: Using GCC at ${GCC_EXECUTABLE} for CUDA compilation") - # PyTorch's cmake/public/cuda.cmake (loaded transitively by - # find_package(Torch)) appends -Xcompiler=-fclang-abi-compat=17 to - # CMAKE_CUDA_FLAGS whenever the parent build is configured with Clang>=18 (see - # pytorch PR #175233). Since CUDA host compilation here falls back to GCC, - # that Clang-only flag would be forwarded to g++ via `nvcc -ccbin=g++` and - # abort the build with: g++: error: unrecognized command-line option - # '-fclang-abi-compat=17' - string(REGEX REPLACE "-Xcompiler=-fclang-abi-compat=[0-9]+" "" - CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS}") endif() # Check CUDA version and architecture support diff --git a/cpp/tensorrt_llm/kernels/CMakeLists.txt b/cpp/tensorrt_llm/kernels/CMakeLists.txt index fcea8829442b..e0e498fa89eb 100644 --- a/cpp/tensorrt_llm/kernels/CMakeLists.txt +++ b/cpp/tensorrt_llm/kernels/CMakeLists.txt @@ -30,8 +30,6 @@ add_subdirectory(dsv3MinLatencyKernels) add_subdirectory(causalConv1d) add_subdirectory(fusedGatedRMSNormQuant) add_subdirectory(mamba2MTPSSMCache) -add_subdirectory(mhcKernels) -add_subdirectory(compressorKernels) file(GLOB_RECURSE SRC_CPP *.cpp) file(GLOB_RECURSE SRC_CU *.cu) @@ -57,13 +55,6 @@ list(FILTER SRC_CU EXCLUDE REGEX "userbuffers/.*") list(FILTER SRC_CU EXCLUDE REGEX "fusedLayernormKernels/.*") list(FILTER SRC_CU EXCLUDE REGEX "fusedGatedRMSNormQuant/.*") list(FILTER SRC_CU EXCLUDE REGEX "mamba2MTPSSMCache/.*") -list(FILTER SRC_CPP EXCLUDE REGEX "mhcKernels/.*") -list(FILTER SRC_CU EXCLUDE REGEX "mhcKernels/.*") -list(FILTER SRC_CPP EXCLUDE REGEX "compressorKernels/.*") -list(FILTER SRC_CU EXCLUDE REGEX "compressorKernels/.*") -# Marlin is built as its own Hopper-only OBJECT library below. -list(FILTER SRC_CPP EXCLUDE REGEX "marlin/.*") -list(FILTER SRC_CU EXCLUDE REGEX "marlin/.*") if(NOT ENABLE_MULTI_DEVICE) list(FILTER SRC_CU EXCLUDE REGEX "customAllReduceKernels*.*cu$") @@ -82,26 +73,7 @@ if(FAST_BUILD) STATUS "FAST_BUILD enabled for kernels: using -O1 for CUDA compilation") endif() -# Marlin NVFP4: Hopper-only OBJECT library. Pinned to sm_90 so the global -# CMAKE_CUDA_ARCHITECTURES doesn't propagate. -file(GLOB_RECURSE MARLIN_SRC "marlin/*.cu" "marlin/*.cpp") -if(MARLIN_SRC) - add_library(marlin_src OBJECT ${MARLIN_SRC}) - set_property(TARGET marlin_src PROPERTY POSITION_INDEPENDENT_CODE ON) - set_property(TARGET marlin_src PROPERTY CUDA_RESOLVE_DEVICE_SYMBOLS ON) - target_include_directories( - marlin_src - PRIVATE - $ - ) - target_link_libraries(marlin_src PRIVATE trtllm_gen_fmha_interface) - set_cuda_architectures(marlin_src 90) -endif() - -add_library( - kernels_src STATIC - ${SRC_CPP} ${SRC_CU} - $<$:$>) +add_library(kernels_src STATIC ${SRC_CPP} ${SRC_CU}) set_property(TARGET kernels_src PROPERTY POSITION_INDEPENDENT_CODE ON) set_property(TARGET kernels_src PROPERTY CUDA_RESOLVE_DEVICE_SYMBOLS ON) target_include_directories( diff --git a/cpp/tensorrt_llm/kernels/IndexerTopK.h b/cpp/tensorrt_llm/kernels/IndexerTopK.h index 6e6e7b29b059..324e597dc779 100644 --- a/cpp/tensorrt_llm/kernels/IndexerTopK.h +++ b/cpp/tensorrt_llm/kernels/IndexerTopK.h @@ -1,5 +1,5 @@ /* - * Copyright (c) 2019-2026, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2019-2025, NVIDIA CORPORATION. All rights reserved. * Copyright (c) 2021, NAVER Corp. Authored by CLOVA. * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -27,63 +27,47 @@ TRTLLM_NAMESPACE_BEGIN namespace kernels { -// Number of blocks-per-row used by the multi-block split + merge dispatch path of -// invokeIndexerTopKDecode. Returns 1 when the single-block path is preferred. -// Callers that allocate aux buffers must use this same helper to size them, and -// must pass the same splitWorkThreshold they will pass to invokeIndexerTopKDecode -// (a value <= 0 selects the internal default). -int computeIndexerTopKDecodeBlocksPerRow(int numRows, int numColumns, int splitWorkThreshold = 0); +/// Indexer TopK decode. Three tiers: +/// - GVR Heuristic (preIdx provided, K in {512,1024,2048}, numColumns in +/// [kSeqSmall, splitWorkThreshold), numRows below the +/// architecture-derived wave/L2 bound). +/// - Single-block (numColumns < split-work threshold) +/// - Multi-pass radix (numColumns >= split-work threshold; requires +/// `scratch` sized via indexerTopKDecodeScratchBytes, +/// zero-init on first call and may be reused). +/// +/// `is_prefill = true` forces single-block (split-work suppressed). +void invokeIndexerTopKDecode(float const* logits, int const* seqLens, int* indices, int const splitWorkThreshold, + int const numRows, int const numColumns, int const stride0, int const stride1, int const next_n, + int const topK = 2048, int const* preIdx = nullptr, int const preIdxStride = 0, int const preIdxCount = 0, + float* heuristicScratch = nullptr, cudaStream_t const stream = 0, void* scratch = nullptr, size_t scratchBytes = 0, + bool is_prefill = false); -/// fp32 indexer TopK decode — L2-aware BS-threshold dispatcher with four -/// fallback tiers: -/// - GVR Heuristic (preIdx provided, kSeqSmall ≤ N < splitWork, BS < kBsLarge, K ∈ {512,1024,2048}) -/// - Insertion sort (N < kSortingAlgorithmThreshold) -/// - Radix sort (kSortingAlgorithmThreshold ≤ N < splitWork) -/// - Radix split-work (N ≥ splitWork — uses outLogitsAux / outIndicesAux) -void invokeIndexerTopKDecode(float const* logits, int const* seqLens, int* indices, float* outLogitsAux, - int* outIndicesAux, int const splitWorkThreshold, int const numRows, int const numColumns, int const stride0, - int const stride1, int const next_n, int const topK = 2048, int const* preIdx = nullptr, int const preIdxStride = 0, - int const preIdxCount = 0, float* heuristicScratch = nullptr, int const compressRatio = 1, - cudaStream_t const stream = 0); +/// Size of the multi-pass radix `scratch` buffer for these shapes. +size_t indexerTopKDecodeScratchBytes(int numRows, int numColumns, int topK); -/// bf16 indexer TopK decode — same dispatch axes as the fp32 entry, except -/// kBsL2 uses sizeof(__nv_bfloat16) bytes/elem (L2 footprint is half) and -/// the split-work tier is unsupported (the bf16/fp16 entry does not expose -/// the float aux buffers required for split-work). Insertion + radix tiers -/// share topKPerRowDecode with fp32 — histogram and sort run on float keys -/// after a static_cast(InputT) at HBM-read sites. -/// -/// Aborts with TLLM_CHECK if numColumns ≥ splitWorkThreshold; callers in -/// that regime must use the fp32 entry. +/// bf16 overload; same contract. void invokeIndexerTopKDecode(__nv_bfloat16 const* logits, int const* seqLens, int* indices, int const splitWorkThreshold, int const numRows, int const numColumns, int const stride0, int const stride1, int const next_n, int const topK = 2048, int const* preIdx = nullptr, int const preIdxStride = 0, - int const preIdxCount = 0, __nv_bfloat16* heuristicScratch = nullptr, int const compressRatio = 1, - cudaStream_t const stream = 0); + int const preIdxCount = 0, __nv_bfloat16* heuristicScratch = nullptr, cudaStream_t const stream = 0, + void* scratch = nullptr, size_t scratchBytes = 0, bool is_prefill = false); -/// fp16 indexer TopK decode — see bf16 overload for dispatcher contract. +/// fp16 overload; same contract. void invokeIndexerTopKDecode(__half const* logits, int const* seqLens, int* indices, int const splitWorkThreshold, int const numRows, int const numColumns, int const stride0, int const stride1, int const next_n, int const topK = 2048, int const* preIdx = nullptr, int const preIdxStride = 0, int const preIdxCount = 0, - __half* heuristicScratch = nullptr, int const compressRatio = 1, cudaStream_t const stream = 0); + __half* heuristicScratch = nullptr, cudaStream_t const stream = 0, void* scratch = nullptr, size_t scratchBytes = 0, + bool is_prefill = false); void invokeIndexerTopKPrefill(float const* logits, int const* rowStarts, int const* rowEnds, int* indices, int const numRows, int const numColumns, int const stride0, int const stride1, int const topK = 2048, cudaStream_t const stream = 0); -/// Returns true iff invokeIndexerTopKDecode would route to the GVR Heuristic -/// kernel for this (numRows, numColumns, topK) triple, assuming valid preIdx -/// is provided and stride1 == 1. Useful for callers that need to provision a -/// preIdx tensor or heuristicScratch buffer only when GVR will be selected. -/// -/// Mirrors the gating logic of the dispatcher: K ∈ {512, 1024, 2048}, -/// numColumns ∈ [kSeqSmall, splitWorkThreshold), numRows < kBsLarge, where -/// kBsLarge = min(kBsWave, kBsL2) and kBsL2 scales with bytesPerElem. -/// -/// @param numRows logits rows (batch · next_n) -/// @param numColumns logits columns (max sequence length) -/// @param topK requested output size -/// @param bytesPerElem element size of logits (4 for fp32, 2 for bf16/fp16) +/// True iff invokeIndexerTopKDecode would pick the GVR tier for this shape: +/// K in {512,1024,2048}, numColumns in [kSeqSmall, splitWorkThreshold), and +/// numRows below the architecture-derived wave/L2 bound. Lets callers +/// provision preIdx / heuristicScratch only when needed. bool canIndexerTopKDecodeUseGvr(int numRows, int numColumns, int topK, int bytesPerElem = 4); } // namespace kernels diff --git a/cpp/tensorrt_llm/kernels/beamSearchKernels.cu b/cpp/tensorrt_llm/kernels/beamSearchKernels.cu index 2dadc9d5b682..005a1539168e 100644 --- a/cpp/tensorrt_llm/kernels/beamSearchKernels.cu +++ b/cpp/tensorrt_llm/kernels/beamSearchKernels.cu @@ -16,7 +16,6 @@ #include "tensorrt_llm/common/config.h" #include "tensorrt_llm/common/cudaUtils.h" -#include "tensorrt_llm/common/reduceKernelUtils.cuh" #include "tensorrt_llm/kernels/beamSearchKernels.h" using namespace tensorrt_llm::common; @@ -136,26 +135,21 @@ void invokeUpdateCacheIndirection(int* tgtCI, int const* srcCI, BeamHypotheses& sync_check_cuda_error(stream); } -__global__ void addCumLogProbs(float* __restrict pStage1LogProbs, int const* __restrict pStage1Ids, - float const* __restrict cumLogProbs, FinishedState const* finished, int const* endIds, float const* diversityRates, +__global__ void addCumLogProbs(float* __restrict pStage1LogProbs, float const* __restrict cumLogProbs, + FinishedState const* finished, int const* endIds, float const* diversityRates, runtime::SizeType32 const* batchSlots, size_t const nBS, size_t const nBMIn, size_t const nBMOut, size_t const nBM) { int const bid = blockIdx.x; // Index of request in batch runtime::SizeType32 const slot = batchSlots[bid]; float const diversityRate{diversityRates[slot]}; float* pLocalLogProbs = pStage1LogProbs + bid * nBMIn * nBMOut * 2; - int const* pLocalIds = pStage1Ids + bid * nBMIn * nBMOut * 2; for (int i = threadIdx.x; i < nBMIn * nBMOut * 2; i += blockDim.x) { int const iBMIn = i / (nBMOut * 2); - if (finished[slot * nBM + iBMIn].isFinished()) + if (finished[slot * nBMIn + iBMIn].isFinished()) { - // In V2 path, i is a candidate-slot index (0..nBMIn*nBMOut*2-1), NOT a vocab token id. - // Use pStage1Ids to look up the actual token id for the EOS comparison. - bool const isEOS = (pLocalIds[i] == endIds[slot]); - // Keep only the EOS candidate with its proper cumulative score; suppress all others. - pLocalLogProbs[i] = isEOS ? (pLocalLogProbs[i] + cumLogProbs[slot * nBM + iBMIn]) : -FLT_MAX; + pLocalLogProbs[i] += (i == endIds[slot]) ? 1.0f : 0.0f; } else { @@ -166,27 +160,21 @@ __global__ void addCumLogProbs(float* __restrict pStage1LogProbs, int const* __r return; } -__global__ void addCumLogProbs(half* __restrict pStage1LogProbs, int const* __restrict pStage1Ids, - float const* __restrict cumLogProbs, FinishedState const* finished, int const* endIds, float const* diversityRates, +__global__ void addCumLogProbs(half* __restrict pStage1LogProbs, float const* __restrict cumLogProbs, + FinishedState const* finished, int const* endIds, float const* diversityRates, runtime::SizeType32 const* batchSlots, size_t const nBS, size_t const nBMIn, size_t const nBMOut, size_t const nBM) { int const bid = blockIdx.x; // Index of request in batch runtime::SizeType32 const slot = batchSlots[bid]; float const diversityRate{diversityRates[slot]}; half* pLocalLogProbs = pStage1LogProbs + bid * nBMIn * nBMOut * 2; - int const* pLocalIds = pStage1Ids + bid * nBMIn * nBMOut * 2; for (int i = threadIdx.x; i < nBMIn * nBMOut * 2; i += blockDim.x) { int const iBMIn = i / (nBMOut * 2); - if (finished[slot * nBM + iBMIn].isFinished()) + if (finished[slot * nBMIn + iBMIn].isFinished()) { - // In V2 path, i is a candidate-slot index (0..nBMIn*nBMOut*2-1), NOT a vocab token id. - // Use pStage1Ids to look up the actual token id for the EOS comparison. - bool const isEOS = (pLocalIds[i] == endIds[slot]); - // Keep only the EOS candidate with its proper cumulative score; suppress all others. - pLocalLogProbs[i] - = isEOS ? (half) (float(pLocalLogProbs[i]) + cumLogProbs[slot * nBM + iBMIn]) : (half) -HALF_FLT_MAX; + pLocalLogProbs[i] += (i == endIds[slot]) ? 1.0f : 0.0f; } else { diff --git a/cpp/tensorrt_llm/kernels/beamSearchKernels.h b/cpp/tensorrt_llm/kernels/beamSearchKernels.h index 345e4659c941..d8a9266e9406 100644 --- a/cpp/tensorrt_llm/kernels/beamSearchKernels.h +++ b/cpp/tensorrt_llm/kernels/beamSearchKernels.h @@ -131,15 +131,13 @@ void invokeTopkBeamSearch(T const* logProbs, T const* bias, void* workspace, Bea void invokeUpdateCacheIndirection(int* tgtCI, int const* srcCI, BeamHypotheses& bh, runtime::SizeType32 const maxAttentionWindow, runtime::SizeType32 sinkTokenLength, cudaStream_t stream); -__global__ void addCumLogProbs(float* __restrict pStage1LogProbs, int const* __restrict pStage1Ids, - float const* __restrict cumLogProbs, ::tensorrt_llm::kernels::FinishedState const* finished, int const* endIds, - float const* diversityRates, runtime::SizeType32 const* batchSlots, size_t const nBS, size_t const nBMIn, - size_t const nBMOut, size_t const nBM); - -__global__ void addCumLogProbs(half* __restrict pStage1LogProbs, int const* __restrict pStage1Ids, - float const* __restrict cumLogProbs, ::tensorrt_llm::kernels::FinishedState const* finished, int const* endIds, - float const* diversityRates, runtime::SizeType32 const* batchSlots, size_t const nBS, size_t const nBMIn, - size_t const nBMOut, size_t const nBM); +__global__ void addCumLogProbs(float* __restrict pStage1LogProbs, float const* __restrict cumLogProbs, + ::tensorrt_llm::kernels::FinishedState const* finished, int const* endIds, float const* diversityRates, + runtime::SizeType32 const* batchSlots, size_t const nBS, size_t const nBMIn, size_t const nBMOut, size_t const nBM); + +__global__ void addCumLogProbs(half* __restrict pStage1LogProbs, float const* __restrict cumLogProbs, + ::tensorrt_llm::kernels::FinishedState const* finished, int const* endIds, float const* diversityRates, + runtime::SizeType32 const* batchSlots, size_t const nBS, size_t const nBMIn, size_t const nBMOut, size_t const nBM); __global__ void gatherId(int const* __restrict pStage1Id, int* __restrict pStage2Id, size_t const nBS, size_t const nBMIn, size_t const nBMOut, size_t const nV); diff --git a/cpp/tensorrt_llm/kernels/beamSearchKernels/beamSearchKernelsTemplate.h b/cpp/tensorrt_llm/kernels/beamSearchKernels/beamSearchKernelsTemplate.h index 09b114ef9752..eb0d9e072997 100644 --- a/cpp/tensorrt_llm/kernels/beamSearchKernels/beamSearchKernelsTemplate.h +++ b/cpp/tensorrt_llm/kernels/beamSearchKernels/beamSearchKernelsTemplate.h @@ -208,7 +208,6 @@ __launch_bounds__(BLOCK_SIZE) __global__ void beamStage3Kernel( __shared__ float smemCumLogProbs[PBM]; __shared__ int smemSeqLen[PBM]; __shared__ KVPair smemTopKV[(IS_V2) ? 1 : PBM * 2]; // Just a placeholder in V2 workflow - __shared__ int smemNBeamForNextStep; if (bh.numBeamsCBA != nullptr) { @@ -218,11 +217,14 @@ __launch_bounds__(BLOCK_SIZE) __global__ void beamStage3Kernel( // Initialize worst score in the first call bh.minNormedScoresCBA[slot] = 0.0f; // logProbs is in range (-inf, 0] } - else if (earlyStopping == 1 && bh.numBeamsCBA[slot] >= nBM || earlyStopping != 1 && bh.batchDones[slot]) + else if (earlyStopping == 1 && bh.numBeamsCBA[slot] == nBM + || earlyStopping != 1 && bh.finished[slot * nBM].isFinished()) { // Condition of early return: // 1. In EarlyStopping mode, and we have got enough beams // 2. In NonEarlyStopping mode, and this batch has been marked as done + // TODO: improve the condition like below + // earlyStopping == 1 && bh.numBeamsCBA[slot] == nBM || earlyStopping != 1 && bh.batchDones[slot] return; } } @@ -294,11 +296,6 @@ __launch_bounds__(BLOCK_SIZE) __global__ void beamStage3Kernel( } __syncthreads(); - // Timestep at which each next-step destination beam's token is stored in the work tree. - // This is the parent beam's sequence length (the true generation step), which may differ - // from the destination slot's own (possibly stale) length when a finished slot is reused. - __shared__ int smemWriteStep[PBM]; - if (tid == 0) { int nBeamForNextStep{0}; @@ -327,14 +324,10 @@ __launch_bounds__(BLOCK_SIZE) __global__ void beamStage3Kernel( { // Condition of this branch: // This token is end-token and belongs to top nBM range in Beam search mode - // Use the actual parent beam index (topId / nV) % nBM, not the candidate rank i, - // to look up the correct sequenceLength and inputLength for length-penalty scoring. - int const parentBeam = (topId / nV) % nBM; - int const nSeqLen - = bh.sequenceLengths[slot * nBM + parentBeam] + 1 - bh.inputLengths[slot * nBM + parentBeam]; + int const nSeqLen = bh.sequenceLengths[slot * nBM + i] + 1 - bh.inputLengths[slot * nBM + i]; float const score = applyLengthPenalty(topLogProb, nSeqLen, lengthPenalty); int nCBA = bh.numBeamsCBA[slot]; - if (nCBA >= nBM) + if (nCBA == nBM) { // There are already nBM beams if (score < bh.minNormedScoresCBA[slot]) @@ -414,18 +407,13 @@ __launch_bounds__(BLOCK_SIZE) __global__ void beamStage3Kernel( // 1. bh.numBeamsCBA == nullptr && i < nBM, i.e., beam search is disable // 2. bh.numBeamsCBA != nullptr && i < nBM && isEndToken == false, i.e., add token at the end // 3. bh.numBeamsCBA != nullptr && i >= nBM && isEndToken == false, i.e., add token at the end - // Write at the parent beam's sequence length (the actual generation step), - // not the destination slot's length, which can be stale if the slot was - // previously finished and is now being reused for a new continuation. - int const parentBeam = topId / nV % nBM; - int const step = bh.sequenceLengths[slot * nBM + parentBeam]; - smemWriteStep[nBeamForNextStep] = step; + int const step = bh.sequenceLengths[slot * nBM + nBeamForNextStep]; // Copy the selected token to work tree bh.outputIdsPtr[slot][nBeamForNextStep * nMSL + step] = topId; if (bh.logProbsTiled != nullptr) { int const index = step * nMBS * nBM + slot * nBM + nBeamForNextStep; - int const indexBeam = parentBeam; + int const indexBeam = topId / nV % nBM; bh.logProbsTiled[index] = (float) topLogProb - smemCumLogProbs[indexBeam]; } bh.cumLogProbs[slot * nBM + nBeamForNextStep] = (float) topLogProb; @@ -449,7 +437,6 @@ __launch_bounds__(BLOCK_SIZE) __global__ void beamStage3Kernel( break; } } - smemNBeamForNextStep = nBeamForNextStep; } // Update bh.batchDones @@ -494,40 +481,23 @@ __launch_bounds__(BLOCK_SIZE) __global__ void beamStage3Kernel( if (tid < nBMOut) { int const indexBatchBeam = slot * nBM + tid; - if (tid < smemNBeamForNextStep) + int const step = smemSeqLen[tid]; + if (!bh.finished[indexBatchBeam].isFinished()) { - // This slot received a valid next-step token from the selection phase. - // Use the timestep recorded by the selection phase (the parent beam's length), - // which matches the position where the encoded token was stored. - int const step = smemWriteStep[tid]; - int const newId = bh.outputIdsPtr[slot][tid * nMSL + step]; - int const newBeamId = (newId / nV) % nBM; - int const newTokenId = newId % nV; - int const indexParentBeam = slot * nBM + newBeamId; - int const parentSeqLen = smemSeqLen[newBeamId]; - bh.sequenceLengths[indexBatchBeam] = parentSeqLen + (!bh.finished[indexParentBeam].isFinished() ? 1 : 0); - if (newTokenId == bh.endIds[slot]) - { - bh.finished[indexBatchBeam].setFinishedEOS(); - } - else - { - // Reset any stale finished state: this slot may have been marked finished in a - // previous step and is now reused for a valid non-EOS beam; otherwise it would be - // wrongly skipped downstream. - bh.finished[indexBatchBeam] = FinishedState::empty(); - } - bh.parentIdsPtr[slot][tid * nMSL + step] = newBeamId; - bh.outputIdsPtr[slot][tid * nMSL + step] = newTokenId; + smemSeqLen[tid]++; } - else + int const newId = bh.outputIdsPtr[slot][tid * nMSL + step]; + int const newBeamId = (newId / nV) % nBM; + int const newTokenId = newId % nV; + bh.sequenceLengths[indexBatchBeam] = smemSeqLen[newBeamId]; + if (newTokenId == bh.endIds[slot]) { - // No valid next-step token for this slot: all top candidates went to CBA. - // Mark as finished so downstream stages (cache indirection, next decode) skip it. - bh.finished[indexBatchBeam].setFinished(); + bh.finished[indexBatchBeam].setFinishedEOS(); } + bh.parentIdsPtr[slot][tid * nMSL + step] = newBeamId; + bh.outputIdsPtr[slot][tid * nMSL + step] = newTokenId; - if ((earlyStopping == 1) && (bh.numBeamsCBA != nullptr && bh.numBeamsCBA[slot] >= nBM) + if ((earlyStopping == 1) && (bh.numBeamsCBA != nullptr && bh.numBeamsCBA[slot] == nBM) || (earlyStopping != 1) && bh.batchDones[slot]) { bh.batchDones[slot] = true; @@ -661,7 +631,7 @@ void beamSearchKernelLauncher( sync_check_cuda_error(stream); int nThread = min(roundUp(nBMIn * nBMOut * 2, 32), 1024); - addCumLogProbs<<>>(pStage1LogProbs, pStage1Ids, bh.cumLogProbs, bh.finished, bh.endIds, + addCumLogProbs<<>>(pStage1LogProbs, bh.cumLogProbs, bh.finished, bh.endIds, bh.diversityRates, bh.batchSlots, nBS, nBMIn, nBMOut, nBM); sync_check_cuda_error(stream); diff --git a/cpp/tensorrt_llm/kernels/communicationKernels/mnnvlAllreduceKernels.cu b/cpp/tensorrt_llm/kernels/communicationKernels/mnnvlAllreduceKernels.cu index eb44f1638a19..8bb09476ba81 100644 --- a/cpp/tensorrt_llm/kernels/communicationKernels/mnnvlAllreduceKernels.cu +++ b/cpp/tensorrt_llm/kernels/communicationKernels/mnnvlAllreduceKernels.cu @@ -1,5 +1,5 @@ /* - * Copyright (c) 2025-2026, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -21,10 +21,8 @@ #include #include #include -#include #include #include -#include #include "tensorrt_llm/common/cudaTypeUtils.cuh" #include "tensorrt_llm/common/cudaUtils.h" @@ -33,7 +31,6 @@ #include "tensorrt_llm/common/lamportUtils.cuh" #include "tensorrt_llm/common/logger.h" #include "tensorrt_llm/common/reduceKernelUtils.cuh" -#include "tensorrt_llm/kernels/quantization.cuh" TRTLLM_NAMESPACE_BEGIN @@ -41,12 +38,10 @@ namespace kernels::mnnvl { using tensorrt_llm::common::isNegZero; -using tensorrt_llm::common::isLamportDirty; using tensorrt_llm::common::LamportFlags; using tensorrt_llm::common::cuda_cast; using tensorrt_llm::common::getMultiProcessorCount; using tensorrt_llm::common::getDTypeSize; -using tensorrt_llm::common::loadPackedVolatile; // Guard the helper function used for this kernel. namespace detail @@ -91,6 +86,31 @@ inline __device__ const PackedType loadPacked(T const* ptr) return *reinterpret_cast(ptr); } +template +inline __device__ PackedType loadPackedVolatile(void const* ptr) +{ + static_assert(sizeof(PackedType) == 0, "Not implemented"); + return PackedType{}; +} + +template <> +inline __device__ float4 loadPackedVolatile(void const* ptr) +{ + float4 returnValue; + asm volatile("ld.volatile.global.v4.f32 {%0, %1, %2, %3}, [%4];\n" + : "=f"(returnValue.x), "=f"(returnValue.y), "=f"(returnValue.z), "=f"(returnValue.w) + : "l"(ptr)); + return returnValue; +} + +template <> +inline __device__ float2 loadPackedVolatile(void const* ptr) +{ + float2 returnValue; + asm volatile("ld.volatile.global.v2.f32 {%0, %1}, [%2];\n" : "=f"(returnValue.x), "=f"(returnValue.y) : "l"(ptr)); + return returnValue; +} + template inline __device__ void copyF4(T_IN* dst, T_IN const* src) { @@ -206,73 +226,44 @@ inline __device__ __host__ T divUp(T m, T n) return (m + n - 1) / n; } -inline bool requiresTwoAccessScaleGroup(ar_fusion::AllReduceFusionPattern pattern) +// A helper function to tune the grid configuration for fused oneshot and rmsnorm kernels +// Return (block_size, cluster_size, loads_per_thread) +std::tuple adjustGridConfig(int numTokens, int dim, int eltsPerThread) { - return pattern == ar_fusion::AllReduceFusionPattern::kARResidualRMSNormFP4Quant - || pattern == ar_fusion::AllReduceFusionPattern::kARResidualRMSNormOutFP4Quant; -} + static int SM = tensorrt_llm::common::getSMVersion(); -// A helper function to tune the grid configuration for fused oneshot and RMSNorm kernels. -// Return (block_size, cluster_size, loads_per_thread). -template -std::tuple adjustGridConfig(int numTokens, int dim, int eltsPerThread, int accessGroupSize = 1) -{ - // Step 1: Start from the widest cluster we are willing to launch. The caller can request a no-CGA - // fallback for multi-load RMSNorm rows, in which case the cluster width stays at one CTA. - int clusterSize = UseCluster ? 8 : 1; + int clusterSize = SM >= 90 ? 8 : 1; int blockSize = 128; + // ========================== Adjust the grid configuration ========================== int threadsNeeded = divUp(dim, eltsPerThread); int loadsPerThread = 1; blockSize = divUp(threadsNeeded, clusterSize); - if constexpr (UseCluster) + if (clusterSize > 1) { - // Step 2: Shrink the cluster until the hidden dimension partitions cleanly across CTAs - // and each CTA contains an integral packed-access group. NVFP4 scale generation uses a - // two-access group, so it must not straddle a CTA boundary. - while ((threadsNeeded % clusterSize != 0 || divUp(threadsNeeded, clusterSize) % accessGroupSize != 0) - && clusterSize > 1) + while (threadsNeeded % clusterSize != 0 && clusterSize > 1) { clusterSize /= 2; } - int const maxDivisibleClusterSize = clusterSize; blockSize = divUp(threadsNeeded, clusterSize); - // Step 3: If divisibility leaves each CTA too small, trade cluster width for at least - // a 128-thread CTA. This improves occupancy and avoids tiny CTAs when the row is narrow. while (blockSize < 128 && clusterSize >= 2) { blockSize *= 2; clusterSize /= 2; } int smCount = getMultiProcessorCount(); - // Step 4: If the token grid already has enough CTAs to cover the GPU, reduce cluster - // width and make CTAs larger. This avoids over-partitioning one token across too many CTAs. while (numTokens * clusterSize > smCount && clusterSize > 1 && blockSize <= 512) { blockSize *= 2; clusterSize /= 2; } - // Step 5: If the token grid still underfills the GPU, restore cluster width up to the - // divisibility limit. We accept 64-thread CTAs here to expose more CTAs per token. - while (clusterSize < maxDivisibleClusterSize) - { - int const candidateClusterSize = clusterSize * 2; - int const candidateBlockSize = divUp(threadsNeeded, candidateClusterSize); - if (candidateBlockSize < 64 || candidateBlockSize % accessGroupSize != 0 - || numTokens * candidateClusterSize > smCount) - { - break; - } - clusterSize = candidateClusterSize; - blockSize = candidateBlockSize; - } } - // Step 6: For very wide rows, first increase cluster width on SM90+ and then increase - // per-thread loads. The goal is to keep block_size within CUDA's 1024-thread limit. + // Trying to scale up use multiple loads or CGA while (blockSize > 1024) { - if constexpr (UseCluster) + // Scale up with CGA if supported + if (SM >= 90) { if (clusterSize < 8) { @@ -280,18 +271,12 @@ std::tuple adjustGridConfig(int numTokens, int dim, int eltsPerTh } else { - if (loadsPerThread < 8) - { - loadsPerThread += 1; - } - else - { - break; - } + break; } } else { + if (loadsPerThread < 8) { loadsPerThread += 1; @@ -303,391 +288,131 @@ std::tuple adjustGridConfig(int numTokens, int dim, int eltsPerTh } blockSize = divUp(threadsNeeded, clusterSize * loadsPerThread); } - - while (blockSize % accessGroupSize != 0 && loadsPerThread < 8) - { - loadsPerThread += 1; - blockSize = divUp(threadsNeeded, clusterSize * loadsPerThread); - } - return {blockSize, clusterSize, loadsPerThread}; } -template -struct MnnvlAllReduceKernelParams -{ - T* outputPtr; - T* residualOutPtr; - T const* shardPtr; - T const* residualInPtr; - T const* gammaPtr; - T** inputPtrs; - T* bufferInputPtr; - T* mcastPtr; - void* quantOutPtr; - void* scaleOutPtr; - float const* scaleFactorPtr; - int numTokens; - int tokenDim; - int nRanks; - int rank; - float epsilon; - uint32_t* bufferFlags; - bool waitForResults; - QuantizationSFLayout layout; -}; - -template -inline __device__ void sanitizeLamportPayload(PackedVec& value) -{ -#pragma unroll - for (int i = 0; i < sizeof(PackedType) / sizeof(T); i++) - { - if (isNegZero(value.elements[i])) - { - value.elements[i] = cuda_cast(0.F); - } - } -} - -template -inline __device__ bool pollOneshotRemoteRank( - PackedVec* remoteValues, T* stagePtrLocal, int token, int tokenDim, int packedIdx) -{ - if constexpr (Rank == LocalRank) - { - return true; - } - else - { - auto loaded = loadPackedVolatile( - &stagePtrLocal[token * tokenDim * WorldSize + Rank * tokenDim + packedIdx * kELTS_PER_THREAD]); - remoteValues[Rank].packed = loaded.packed; - return !isLamportDirty(loaded); - } -} - -template -inline __device__ bool pollOneshotRemoteRanks(PackedVec* remoteValues, T* stagePtrLocal, int token, - int tokenDim, int packedIdx, std::integer_sequence) -{ - bool valid = true; - ((valid &= pollOneshotRemoteRank( - remoteValues, stagePtrLocal, token, tokenDim, packedIdx)), - ...); - return valid; -} - -template -inline __device__ void accumulatePacked(float (&accum)[kELTS_PER_THREAD], PackedVec const& value) -{ -#pragma unroll - for (int i = 0; i < kELTS_PER_THREAD; i++) - { - accum[i] += cuda_cast(value.elements[i]); - } -} - -template -inline __device__ void accumulateLamportRanksChunked( - float (&accum)[kELTS_PER_THREAD], T const* input, int token, int tokenDim, int packedIdx) -{ - static_assert(kRankChunk > 0); - static_assert(WorldSize % kRankChunk == 0); - -#pragma unroll - for (int i = 0; i < kELTS_PER_THREAD; i++) - { - accum[i] = 0.F; - } - -#pragma unroll 1 - for (int rankBase = 0; rankBase < WorldSize; rankBase += kRankChunk) - { - float chunkAccum[kELTS_PER_THREAD]; - while (1) - { - bool valid = true; -#pragma unroll - for (int i = 0; i < kELTS_PER_THREAD; i++) - { - chunkAccum[i] = 0.F; - } -#pragma unroll - for (int rr = 0; rr < kRankChunk; rr++) - { - int const r = rankBase + rr; - auto loaded = loadPackedVolatile( - &input[token * tokenDim * WorldSize + r * tokenDim + packedIdx * kELTS_PER_THREAD]); - PackedVec value; - value.packed = loaded.packed; - valid &= !isLamportDirty(loaded); - accumulatePacked(chunkAccum, value); - } - if (valid) - { - break; - } - } -#pragma unroll - for (int i = 0; i < kELTS_PER_THREAD; i++) - { - accum[i] += chunkAccum[i]; - } - } -} - -template -inline __device__ void accumulateOneshotRank(float (&accum)[kELTS_PER_THREAD], - PackedVec const* remoteValues, PackedVec const& localValue) -{ - if constexpr (Rank == LocalRank) - { - accumulatePacked(accum, localValue); - } - else - { - accumulatePacked(accum, remoteValues[Rank]); - } -} - -template -inline __device__ void accumulateOneshotRanks(float (&accum)[kELTS_PER_THREAD], - PackedVec const* remoteValues, PackedVec const& localValue, - std::integer_sequence) -{ - (accumulateOneshotRank(accum, remoteValues, localValue), ...); -} - -template -inline __device__ void waitOneshotRemoteRanks( - PackedVec* remoteValues, T* stagePtrLocal, int token, int tokenDim, int packedIdx) -{ - static_assert(LocalRank < WorldSize); - while (1) - { - bool const valid = pollOneshotRemoteRanks( - remoteValues, stagePtrLocal, token, tokenDim, packedIdx, std::make_integer_sequence{}); - if (valid) - { - break; - } - } -} - -template -inline __device__ PackedVec reduceOneshotDeterministic( - PackedVec const* remoteValues, PackedVec const& localValue) -{ - static_assert(LocalRank < WorldSize); - float accum[kELTS_PER_THREAD]; -#pragma unroll - for (int i = 0; i < kELTS_PER_THREAD; i++) - { - accum[i] = 0.F; - } - accumulateOneshotRanks( - accum, remoteValues, localValue, std::make_integer_sequence{}); - - PackedVec packedAccum; -#pragma unroll - for (int i = 0; i < kELTS_PER_THREAD; i++) - { - packedAccum.elements[i] = cuda_cast(accum[i]); - } - return packedAccum; -} - -template -inline __device__ PackedVec reduceOneshotDeterministicFastPath( - PackedVec const& localValue, T* stagePtrLocal, int token, int tokenDim, int packedIdx) -{ - PackedVec remoteValues[WorldSize]; - waitOneshotRemoteRanks( - remoteValues, stagePtrLocal, token, tokenDim, packedIdx); - return reduceOneshotDeterministic(remoteValues, localValue); -} - -template -inline __device__ void quantizeEpilogue(PackedVec const& value, - MnnvlAllReduceKernelParams const& params, int packedAccessIdx, int accessIdInToken, int token) -{ - if constexpr (ar_fusion::GetQuantType == ar_fusion::QuantType::kFP4) - { - static_assert(kELTS_PER_THREAD == 8, "NVFP4 quantization expects eight elements per 16-byte access."); - constexpr int kSFVecSize = 16; - using QuantPackedVec = tensorrt_llm::kernels::PackedVec; - QuantPackedVec packVal = *reinterpret_cast(&value.packed); - auto* sfOut = cvt_quant_get_sf_out_offset(std::nullopt /* batchIdx */, token, accessIdInToken, - std::nullopt /* numRows */, params.tokenDim / kSFVecSize, reinterpret_cast(params.scaleOutPtr), - params.layout); - reinterpret_cast(params.quantOutPtr)[packedAccessIdx] - = cvt_warp_fp16_to_fp4(packVal, *params.scaleFactorPtr, sfOut); - } - else if constexpr (ar_fusion::GetQuantType == ar_fusion::QuantType::kFP8) - { - float const scale = 1.F / *params.scaleFactorPtr; - using PackedQuantizedType = std::conditional_t, float, float2>; - PackedQuantizedType quantized; -#pragma unroll - for (int i = 0; i < kELTS_PER_THREAD; ++i) - { - reinterpret_cast<__nv_fp8_e4m3*>(&quantized)[i] - = static_cast<__nv_fp8_e4m3>(static_cast(value.elements[i]) * scale); - } - reinterpret_cast(params.quantOutPtr)[packedAccessIdx] = quantized; - } - else - { - static_assert(ar_fusion::GetQuantType == ar_fusion::QuantType::kNone, "Invalid quant type"); - } -} - -template -inline __device__ void writeEpilogueOutput(PackedVec const& value, - MnnvlAllReduceKernelParams const& params, int threadOffset, int packedAccessIdx, int accessIdInToken, int token) -{ - if constexpr (ar_fusion::HasAllReduceOut || ar_fusion::HasNormOut) - { - if (params.outputPtr != nullptr) - { - reinterpret_cast(¶ms.outputPtr[threadOffset])[0] = value.packed; - } - } - if constexpr (ar_fusion::GetQuantType != ar_fusion::QuantType::kNone) - { - quantizeEpilogue( - value, params, packedAccessIdx, accessIdInToken, token); - } -} - } // namespace detail using detail::PackedVec; using detail::loadPacked; +using detail::loadPackedVolatile; using detail::blockReduceSum; using detail::divUp; using detail::copyF4; -using detail::MnnvlAllReduceKernelParams; -using detail::sanitizeLamportPayload; -using detail::accumulateLamportRanksChunked; -using detail::reduceOneshotDeterministicFastPath; -using detail::writeEpilogueOutput; - -template -__global__ void __launch_bounds__(1024) oneshotAllreduceFusionKernel(MnnvlAllReduceKernelParams params) + +template +__global__ void __launch_bounds__(1024) oneshotAllreduceFusionKernel(T* outputPtr, T* prenormedPtr, T const* shardPtr, + T const* residualInPtr, T const* gammaPtr, T** inputPtrs, T* mcastPtr, int const numTokens, int const tokenDim, + float epsilon, int const rank, uint32_t* bufferFlags) { constexpr int kELTS_PER_THREAD = sizeof(PackedType) / sizeof(T); + constexpr int kLAMPORT_ELTS_PER_PACKED = sizeof(PackedType) / sizeof(float); constexpr uint32_t kELT_SIZE = sizeof(T); #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) namespace cg = cooperative_groups; cg::cluster_group cluster = cg::this_cluster(); int packedIdx = cluster.thread_rank(); int token = blockIdx.x; - int threadOffset = token * params.tokenDim + packedIdx * kELTS_PER_THREAD; + int threadOffset = token * tokenDim + packedIdx * kELTS_PER_THREAD; cudaGridDependencySynchronize(); #else int packedIdx = blockIdx.y * blockDim.x + threadIdx.x; int token = blockIdx.x; // Offset w.r.t. the input shard - int threadOffset = token * params.tokenDim + packedIdx * kELTS_PER_THREAD; + int threadOffset = token * tokenDim + packedIdx * kELTS_PER_THREAD; #endif // We only use 1 stage for the oneshot allreduce - LamportFlags flag(params.bufferFlags, 1); - T* stagePtrMcast = reinterpret_cast(flag.getCurLamportBuf(params.mcastPtr, 0)); - T* stagePtrLocal = reinterpret_cast(flag.getCurLamportBuf(params.inputPtrs[params.rank], 0)); - bool const inBounds = packedIdx * kELTS_PER_THREAD < params.tokenDim; + LamportFlags flag(bufferFlags, 1); + T* stagePtrMcast = reinterpret_cast(flag.getCurLamportBuf(mcastPtr, 0)); + T* stagePtrLocal = reinterpret_cast(flag.getCurLamportBuf(inputPtrs[rank], 0)); + + if (packedIdx * kELTS_PER_THREAD >= tokenDim) + { + flag.clearDirtyLamportBuf(inputPtrs[rank], -1); + return; + } // ==================== Broadcast tokens to each rank ============================= PackedVec val; - if (inBounds) + val.packed = loadPacked(&shardPtr[threadOffset]); +#pragma unroll + for (int i = 0; i < kELTS_PER_THREAD; i++) { - val.packed = loadPacked(¶ms.shardPtr[threadOffset]); - sanitizeLamportPayload(val); - - reinterpret_cast( - &stagePtrMcast[token * params.tokenDim * WorldSize + params.rank * params.tokenDim])[packedIdx] - = val.packed; + if (isNegZero(val.elements[i])) + val.elements[i] = cuda_cast(0.f); } + reinterpret_cast(&stagePtrMcast[token * tokenDim * WorldSize + rank * tokenDim])[packedIdx] + = val.packed; + flag.ctaArrive(); // ======================= Lamport Sync and clear the output buffer from previous iteration // ============================= - flag.clearDirtyLamportBuf(params.inputPtrs[params.rank], -1); + flag.clearDirtyLamportBuf(inputPtrs[rank], -1); - if (!inBounds) + PackedVec valuesLamport[WorldSize]; + while (1) { - return; + bool valid = true; +#pragma unroll + for (int r = 0; r < WorldSize; r++) + { + valuesLamport[r].packed = loadPackedVolatile( + &stagePtrLocal[token * tokenDim * WorldSize + r * tokenDim + packedIdx * kELTS_PER_THREAD]); + +#pragma unroll + for (int i = 0; i < kLAMPORT_ELTS_PER_PACKED; i++) + { + valid &= !isNegZero(valuesLamport[r].elements[i]); + } + } + if (valid) + { + break; + } } + auto values = reinterpret_cast*>(valuesLamport); // ======================= Reduction ============================= - // Fully deterministic: every rank uses the exact same reduction order. For WorldSize <= 8, specialize the local - // slot so the fast path reuses `val` from registers without a dynamic `remoteValues[params.rank]` store. Larger - // world sizes use the compact fallback because the benefit is thin but specializing every rank significantly - // increases compile time. + float accum[kELTS_PER_THREAD]; PackedVec packedAccum; - if constexpr (WorldSize <= 8) + +#pragma unroll + for (int i = 0; i < kELTS_PER_THREAD; i++) { - packedAccum = val; -#define RUN_ONESHOT_LOCAL_RANK(LOCAL_RANK) \ - case LOCAL_RANK: \ - if constexpr (WorldSize > LOCAL_RANK) \ - { \ - packedAccum = reduceOneshotDeterministicFastPath( \ - val, stagePtrLocal, token, params.tokenDim, packedIdx); \ - } \ - break - - switch (params.rank) - { - RUN_ONESHOT_LOCAL_RANK(0); - RUN_ONESHOT_LOCAL_RANK(1); - RUN_ONESHOT_LOCAL_RANK(2); - RUN_ONESHOT_LOCAL_RANK(3); - RUN_ONESHOT_LOCAL_RANK(4); - RUN_ONESHOT_LOCAL_RANK(5); - RUN_ONESHOT_LOCAL_RANK(6); - RUN_ONESHOT_LOCAL_RANK(7); - } -#undef RUN_ONESHOT_LOCAL_RANK + accum[i] = cuda_cast(values[0].elements[i]); } - else + +#pragma unroll + for (int r = 1; r < WorldSize; r++) { - // Chunk Lamport polling so only a bounded rank set is live at once, avoiding register spills for large - // world sizes. - constexpr int kRankChunk = 8; - float accum[kELTS_PER_THREAD]; - accumulateLamportRanksChunked( - accum, stagePtrLocal, token, params.tokenDim, packedIdx); #pragma unroll for (int i = 0; i < kELTS_PER_THREAD; i++) { - packedAccum.elements[i] = cuda_cast(accum[i]); + accum[i] += cuda_cast(values[r].elements[i]); } } + +#pragma unroll + for (int i = 0; i < kELTS_PER_THREAD; i++) + { + packedAccum.elements[i] = cuda_cast(accum[i]); + } #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) cudaTriggerProgrammaticLaunchCompletion(); #endif - if constexpr (ar_fusion::HasResidual) + if constexpr (RMSNormFusion) { // =============================== Residual =============================== PackedVec residualIn; - residualIn.packed = *reinterpret_cast(¶ms.residualInPtr[threadOffset]); + residualIn.packed = *reinterpret_cast(&residualInPtr[threadOffset]); packedAccum += residualIn; - if constexpr (ar_fusion::HasResidualOut) - { - *reinterpret_cast(¶ms.residualOutPtr[threadOffset]) = packedAccum.packed; - } - } - if constexpr (ar_fusion::HasRMSNorm) - { + *reinterpret_cast(&prenormedPtr[threadOffset]) = packedAccum.packed; // =============================== Rmsnorm ================================ PackedVec gamma; - gamma.packed = *reinterpret_cast(¶ms.gammaPtr[packedIdx * kELTS_PER_THREAD]); + gamma.packed = *reinterpret_cast(&gammaPtr[packedIdx * kELTS_PER_THREAD]); float threadSum = 0.F; #pragma unroll @@ -721,7 +446,7 @@ __global__ void __launch_bounds__(1024) oneshotAllreduceFusionKernel(MnnvlAllRed } } #endif - float rcpRms = rsqrtf(fullSum / params.tokenDim + params.epsilon); + float rcpRms = rsqrtf(fullSum / tokenDim + epsilon); #pragma unroll for (int i = 0; i < kELTS_PER_THREAD; i++) { @@ -729,36 +454,32 @@ __global__ void __launch_bounds__(1024) oneshotAllreduceFusionKernel(MnnvlAllRed cuda_cast(packedAccum.elements[i]) * rcpRms * cuda_cast(gamma.elements[i])); } } - int const packedAccessIdx = threadOffset / kELTS_PER_THREAD; - writeEpilogueOutput( - packedAccum, params, threadOffset, packedAccessIdx, packedIdx, token); - flag.waitAndUpdate({static_cast(params.numTokens * params.tokenDim * WorldSize * kELT_SIZE), 0, 0, 0}); + reinterpret_cast(&outputPtr[threadOffset])[0] = packedAccum.packed; + flag.waitAndUpdate({static_cast(numTokens * tokenDim * WorldSize * kELT_SIZE), 0, 0, 0}); } using detail::adjustGridConfig; -using detail::requiresTwoAccessScaleGroup; void oneshotAllreduceFusionOp(AllReduceFusionParams const& params) { static int const kSMVersion = tensorrt_llm::common::getSMVersion(); - TLLM_CHECK_WITH_INFO(kSMVersion >= 90, "[MNNVL AllReduceOneShot] requires SM 90 or newer."); int const numTokens = params.numTokens; int const tokenDim = params.tokenDim; int const eltsPerThread = sizeof(float4) / getDTypeSize(params.dType); - int const accessGroupSize = requiresTwoAccessScaleGroup(params.pattern) ? 2 : 1; - auto [blockSize, clusterSize, loadsPerThread] - = adjustGridConfig(numTokens, tokenDim, eltsPerThread, accessGroupSize); + auto [blockSize, clusterSize, loadsPerThread] = adjustGridConfig(numTokens, tokenDim, eltsPerThread); dim3 grid(numTokens, clusterSize, 1); TLLM_LOG_DEBUG( "[MNNVL AllReduceOneShot] Dispatch: grid size: (%d, %d, 1), block_size: %d, cluster_size: %d, " - "loads_per_thread: %d, deterministic_reduction_order: 1, threads_needed: %d", + "loads_per_thread: %d, " + "threads_needed: %d", numTokens, clusterSize, blockSize, clusterSize, loadsPerThread, divUp(tokenDim, eltsPerThread)); TLLM_CHECK_WITH_INFO(blockSize <= 1024 && loadsPerThread == 1, - "Hidden Dimension %d exceeds the maximum supported hidden dimension (%d)", tokenDim, 1024 * 8 * eltsPerThread); + "Hidden Dimension %d exceeds the maximum supported hidden dimension (%d)", tokenDim, + 1024 * (kSMVersion >= 90 ? 8 : 1) * eltsPerThread); cudaLaunchAttribute attrs[2]; attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; @@ -774,57 +495,21 @@ void oneshotAllreduceFusionOp(AllReduceFusionParams const& params) .dynamicSmemBytes = 0, .stream = params.stream, .attrs = attrs, - .numAttrs = 2U, + .numAttrs = kSMVersion >= 90 ? 2U : 1U, }; -#define LAUNCH_ALLREDUCE_KERNEL(WORLD_SIZE, T, PATTERN) \ - TLLM_CUDA_CHECK(cudaLaunchKernelEx(&config, &oneshotAllreduceFusionKernel, kernelParams)); -#define DISPATCH_ALLREDUCE_PATTERN(WORLD_SIZE, T) \ - if (params.pattern == ar_fusion::AllReduceFusionPattern::kAllReduce) \ - { \ - LAUNCH_ALLREDUCE_KERNEL(WORLD_SIZE, T, ar_fusion::AllReduceFusionPattern::kAllReduce); \ - } \ - else if (params.pattern == ar_fusion::AllReduceFusionPattern::kARResidualRMSNorm) \ - { \ - LAUNCH_ALLREDUCE_KERNEL(WORLD_SIZE, T, ar_fusion::AllReduceFusionPattern::kARResidualRMSNorm); \ - } \ - else if (params.pattern == ar_fusion::AllReduceFusionPattern::kARResidualRMSNormFP8Quant) \ +#define LAUNCH_ALLREDUCE_KERNEL(WORLD_SIZE, T, RMSNORM) \ + TLLM_CUDA_CHECK(cudaLaunchKernelEx(&config, &oneshotAllreduceFusionKernel, output, \ + residualOut, input, residualIn, gamma, ucPtrs, mcPtr, numTokens, tokenDim, static_cast(params.epsilon), \ + params.rank, params.bufferFlags)); +#define DISPATCH_ALLREDUCE_KERNEL(WORLD_SIZE, T) \ + if (params.rmsNormFusion) \ { \ - LAUNCH_ALLREDUCE_KERNEL(WORLD_SIZE, T, ar_fusion::AllReduceFusionPattern::kARResidualRMSNormFP8Quant); \ - } \ - else if (params.pattern == ar_fusion::AllReduceFusionPattern::kARResidualRMSNormOutFP8Quant) \ - { \ - LAUNCH_ALLREDUCE_KERNEL(WORLD_SIZE, T, ar_fusion::AllReduceFusionPattern::kARResidualRMSNormOutFP8Quant); \ - } \ - else if (params.pattern == ar_fusion::AllReduceFusionPattern::kARResidualRMSNormFP4Quant) \ - { \ - if constexpr (!std::is_same_v) \ - { \ - LAUNCH_ALLREDUCE_KERNEL(WORLD_SIZE, T, ar_fusion::AllReduceFusionPattern::kARResidualRMSNormFP4Quant); \ - } \ - else \ - { \ - TLLM_CHECK_WITH_INFO(false, "[MNNVL AllReduceOneShot] NVFP4 quantization does not support FP32 input."); \ - } \ - } \ - else if (params.pattern == ar_fusion::AllReduceFusionPattern::kARResidualRMSNormOutFP4Quant) \ - { \ - if constexpr (!std::is_same_v) \ - { \ - LAUNCH_ALLREDUCE_KERNEL(WORLD_SIZE, T, ar_fusion::AllReduceFusionPattern::kARResidualRMSNormOutFP4Quant); \ - } \ - else \ - { \ - TLLM_CHECK_WITH_INFO(false, "[MNNVL AllReduceOneShot] NVFP4 quantization does not support FP32 input."); \ - } \ - } \ - else if (params.pattern == ar_fusion::AllReduceFusionPattern::kARRMSNorm) \ - { \ - LAUNCH_ALLREDUCE_KERNEL(WORLD_SIZE, T, ar_fusion::AllReduceFusionPattern::kARRMSNorm); \ + LAUNCH_ALLREDUCE_KERNEL(WORLD_SIZE, T, true); \ } \ else \ { \ - TLLM_CHECK_WITH_INFO(false, "[MNNVL AllReduceOneShot] Unsupported fusion pattern."); \ + LAUNCH_ALLREDUCE_KERNEL(WORLD_SIZE, T, false); \ } // C++17 compatible alternative using a template function auto dispatchImpl = [&](auto* type_ptr) -> bool @@ -837,25 +522,21 @@ void oneshotAllreduceFusionOp(AllReduceFusionParams const& params) T const* input = reinterpret_cast(params.input); T const* residualIn = reinterpret_cast(params.residualIn); T const* gamma = reinterpret_cast(params.gamma); - MnnvlAllReduceKernelParams kernelParams{output, residualOut, input, residualIn, gamma, ucPtrs, - reinterpret_cast(params.bufferPtrLocal), mcPtr, params.quantOut, params.scaleOut, params.scaleFactor, - numTokens, tokenDim, params.nRanks, params.rank, static_cast(params.epsilon), params.bufferFlags, - false, params.layout}; switch (params.nRanks) { // FIXME: Do we need other world sizes? - case 2: DISPATCH_ALLREDUCE_PATTERN(2, T); return true; - case 4: DISPATCH_ALLREDUCE_PATTERN(4, T); return true; - case 8: DISPATCH_ALLREDUCE_PATTERN(8, T); return true; - case 16: DISPATCH_ALLREDUCE_PATTERN(16, T); return true; - case 32: DISPATCH_ALLREDUCE_PATTERN(32, T); return true; - case 64: DISPATCH_ALLREDUCE_PATTERN(64, T); return true; + case 2: DISPATCH_ALLREDUCE_KERNEL(2, T); return true; + case 4: DISPATCH_ALLREDUCE_KERNEL(4, T); return true; + case 8: DISPATCH_ALLREDUCE_KERNEL(8, T); return true; + case 16: DISPATCH_ALLREDUCE_KERNEL(16, T); return true; + case 32: DISPATCH_ALLREDUCE_KERNEL(32, T); return true; + case 64: DISPATCH_ALLREDUCE_KERNEL(64, T); return true; } return false; }; #undef LAUNCH_ALLREDUCE_KERNEL -#undef DISPATCH_ALLREDUCE_PATTERN +#undef DISPATCH_ALLREDUCE_KERNEL bool launched = (params.dType == nvinfer1::DataType::kBF16 && dispatchImpl((__nv_bfloat16*) nullptr)) || (params.dType == nvinfer1::DataType::kFLOAT && dispatchImpl((float*) nullptr)) || (params.dType == nvinfer1::DataType::kHALF && dispatchImpl((__nv_half*) nullptr)); @@ -873,101 +554,138 @@ enum MNNVLTwoShotStage : uint8_t }; template -__global__ __launch_bounds__(128) void twoshotAllreduceKernel(MnnvlAllReduceKernelParams params) +__global__ __launch_bounds__(128) void twoshotAllreduceKernel(T* outputPtr, T const* shardPtr, T** inputPtrs, + T* mcastPtr, uint32_t const numTokens, uint32_t const tokenDim, uint32_t const rank, uint32_t* bufferFlags, + bool const wait_for_results) { constexpr int kELTS_PER_THREAD = sizeof(PackedType) / sizeof(T); + constexpr int kLAMPORT_ELTS_PER_PACKED = sizeof(PackedType) / sizeof(float); constexpr uint32_t kELT_SIZE = sizeof(T); int packedIdx = blockIdx.y * blockDim.x + threadIdx.x; int token = blockIdx.x; // Offset w.r.t. the input shard - int threadOffset = token * params.tokenDim + packedIdx * kELTS_PER_THREAD; + int threadOffset = token * tokenDim + packedIdx * kELTS_PER_THREAD; int destRank = token % WorldSize; int destTokenOffset = token / WorldSize; - bool const inBounds = packedIdx * kELTS_PER_THREAD < params.tokenDim; #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) cudaGridDependencySynchronize(); #endif - LamportFlags flag(params.bufferFlags, MNNVLTwoShotStage::NUM_STAGES); + LamportFlags flag(bufferFlags, MNNVLTwoShotStage::NUM_STAGES); - T* scatterBufLocal - = reinterpret_cast(flag.getCurLamportBuf(params.inputPtrs[params.rank], MNNVLTwoShotStage::SCATTER)); - T* scatterBufDest - = reinterpret_cast(flag.getCurLamportBuf(params.inputPtrs[destRank], MNNVLTwoShotStage::SCATTER)); - T* broadcastBufW = reinterpret_cast(flag.getCurLamportBuf(params.mcastPtr, MNNVLTwoShotStage::BROADCAST)); - T* broadcastBufR - = reinterpret_cast(flag.getCurLamportBuf(params.inputPtrs[params.rank], MNNVLTwoShotStage::BROADCAST)); + T* scatterBufLocal = reinterpret_cast(flag.getCurLamportBuf(inputPtrs[rank], MNNVLTwoShotStage::SCATTER)); + T* scatterBufDest = reinterpret_cast(flag.getCurLamportBuf(inputPtrs[destRank], MNNVLTwoShotStage::SCATTER)); + T* broadcastBufW = reinterpret_cast(flag.getCurLamportBuf(mcastPtr, MNNVLTwoShotStage::BROADCAST)); + T* broadcastBufR = reinterpret_cast(flag.getCurLamportBuf(inputPtrs[rank], MNNVLTwoShotStage::BROADCAST)); + +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + cudaTriggerProgrammaticLaunchCompletion(); +#endif + // Make sure the clear function is called before OOB thread exits + if (packedIdx * kELTS_PER_THREAD >= tokenDim) + { + flag.clearDirtyLamportBuf(inputPtrs[rank], -1); + return; + } // =============================== Scatter =============================== // Load vectorized data PackedVec val; - if (inBounds) + val.packed = loadPacked(&shardPtr[threadOffset]); +#pragma unroll + for (int i = 0; i < kELTS_PER_THREAD; i++) { - val.packed = loadPacked(¶ms.shardPtr[threadOffset]); - sanitizeLamportPayload(val); - - // Store vectorized data - reinterpret_cast( - &scatterBufDest[destTokenOffset * params.tokenDim * WorldSize + params.rank * params.tokenDim])[packedIdx] - = val.packed; + if (isNegZero(val.elements[i])) + { + val.elements[i] = cuda_cast(0.F); + } } -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - cudaTriggerProgrammaticLaunchCompletion(); -#endif + // Store vectorized data + reinterpret_cast(&scatterBufDest[destTokenOffset * tokenDim * WorldSize + rank * tokenDim])[packedIdx] + = val.packed; - flag.clearDirtyLamportBuf(params.inputPtrs[params.rank], MNNVLTwoShotStage::SCATTER); + flag.clearDirtyLamportBuf(inputPtrs[rank], MNNVLTwoShotStage::SCATTER); // =============================== Reduction and Broadcast =============================== - if (inBounds && (token % WorldSize) == params.rank) + if ((token % WorldSize) == rank) { int localToken = token / WorldSize; + float accum[kELTS_PER_THREAD] = {0.F}; + + // Use float as we only check each float value for validity + PackedVec valuesLamport[WorldSize]; + while (1) + { + bool valid = true; +#pragma unroll + for (int r = 0; r < WorldSize; r++) + { + valuesLamport[r].packed = loadPackedVolatile( + &scatterBufLocal[localToken * tokenDim * WorldSize + r * tokenDim + packedIdx * kELTS_PER_THREAD]); + + // Check validity across all elements +#pragma unroll + for (int i = 0; i < kLAMPORT_ELTS_PER_PACKED; i++) + { + valid &= !isNegZero(valuesLamport[r].elements[i]); + } + } + if (valid) + { + break; + } + } + + // Now we view it as the value for reduction + auto values = reinterpret_cast*>(valuesLamport); +#pragma unroll + for (int r = 0; r < WorldSize; r++) + { + +#pragma unroll + for (int i = 0; i < kELTS_PER_THREAD; i++) + { + accum[i] += cuda_cast(values[r].elements[i]); + } + } // Store vectorized result PackedVec packedAccum; - // Chunk Lamport polling so only a bounded rank set is live at once, avoiding register spills for large - // world sizes. - constexpr int kRankChunk = WorldSize < 16 ? WorldSize : 16; - float accum[kELTS_PER_THREAD]; - accumulateLamportRanksChunked( - accum, scatterBufLocal, localToken, params.tokenDim, packedIdx); #pragma unroll for (int i = 0; i < kELTS_PER_THREAD; i++) { packedAccum.elements[i] = cuda_cast(accum[i]); } - sanitizeLamportPayload(packedAccum); - reinterpret_cast(&broadcastBufW[token * params.tokenDim])[packedIdx] = packedAccum.packed; + reinterpret_cast(&broadcastBufW[token * tokenDim])[packedIdx] = packedAccum.packed; } - flag.clearDirtyLamportBuf(params.inputPtrs[params.rank], MNNVLTwoShotStage::BROADCAST); + flag.clearDirtyLamportBuf(inputPtrs[rank], MNNVLTwoShotStage::BROADCAST); // Optionally wait for results if the next layer isn't doing the Lamport check - if (params.waitForResults) + if (wait_for_results) { // Update the atomic counter to indicate the block has read the offsets flag.ctaArrive(); - if (inBounds) + PackedVec valLamport; + valLamport.packed = loadPackedVolatile(&broadcastBufR[threadOffset]); + while (isNegZero(valLamport.elements[0])) { - auto loaded = loadPackedVolatile(&broadcastBufR[threadOffset]); - while (isLamportDirty(loaded)) - { - loaded = loadPackedVolatile(&broadcastBufR[threadOffset]); - } - if (params.outputPtr) - { - reinterpret_cast(¶ms.outputPtr[threadOffset])[0] = loaded.packed; - } + valLamport.packed = loadPackedVolatile(&broadcastBufR[threadOffset]); + } + if (outputPtr) + { + reinterpret_cast(&outputPtr[threadOffset])[0] = valLamport.packed; } // Update the buffer flags - flag.waitAndUpdate({static_cast(divUp(params.numTokens, WorldSize) * WorldSize - * params.tokenDim * kELT_SIZE), // Clear Size for scatter stage - static_cast(params.numTokens * params.tokenDim * kELT_SIZE), // Clear Size for broadcast stage + flag.waitAndUpdate({static_cast(divUp(numTokens, WorldSize) * WorldSize * tokenDim + * kELT_SIZE), // Clear Size for scatter stage + static_cast(numTokens * tokenDim * kELT_SIZE), // Clear Size for broadcast stage 0, 0}); // If not wait for results, we will rely on the following kernel to update the buffer } @@ -979,10 +697,13 @@ __global__ __launch_bounds__(128) void twoshotAllreduceKernel(MnnvlAllReduceKern // 1. Use CGA if supported. It expands the hidden dimension to 8k x 8 = 64k. // 2. Set loads_per_thread >1. Which can be used if CGA is not supported. Note that this will be limited by the // shared memory size and register count. -template -__global__ __launch_bounds__(1024) void rmsNormLamport(MnnvlAllReduceKernelParams params) +template +__global__ __launch_bounds__(1024) void rmsNormLamport(T_IN* outputPreNorm, T_OUT* outputNorm, T_IN* bufferInput, + T_IN const* gamma, float epsilon, T_IN const* residual, uint32_t numTokens, uint32_t dim, uint32_t worldSize, + uint32_t* bufferFlags) { - static int const kELTS_PER_LOAD = sizeof(float4) / sizeof(T); + static_assert(std::is_same_v, "T_IN and T_OUT must be the same type"); + static int const kELTS_PER_LOAD = sizeof(float4) / sizeof(T_IN); uint32_t const token = blockIdx.x; uint32_t const blockSize = blockDim.x; @@ -991,18 +712,14 @@ __global__ __launch_bounds__(1024) void rmsNormLamport(MnnvlAllReduceKernelParam uint32_t numThreads = blockSize; uint32_t clusterSize = 1; uint32_t blockOffset = 0; - if constexpr (UseCGA) - { #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - namespace cg = cooperative_groups; - cg::cluster_group cluster = cg::this_cluster(); - numThreads = cluster.num_threads(); - clusterSize = cluster.num_blocks(); - blockOffset = cluster.block_rank(); + namespace cg = cooperative_groups; + cg::cluster_group cluster = cg::this_cluster(); + numThreads = cluster.num_threads(); + clusterSize = cluster.num_blocks(); + blockOffset = cluster.block_rank(); #endif - } - uint32_t const dimPadded - = divUp(static_cast(params.tokenDim), kELTS_PER_LOAD * numThreads) * kELTS_PER_LOAD * numThreads; + uint32_t const dimPadded = divUp(dim, kELTS_PER_LOAD * numThreads) * kELTS_PER_LOAD * numThreads; uint32_t const elemsPerThread = dimPadded / numThreads; uint32_t const loadStride = blockSize; @@ -1010,14 +727,14 @@ __global__ __launch_bounds__(1024) void rmsNormLamport(MnnvlAllReduceKernelParam float rInput[LoadsPerThread * kELTS_PER_LOAD]; uint32_t offsets[LoadsPerThread * kELTS_PER_LOAD]; - uint32_t const smemBufferSize = blockSize * elemsPerThread * sizeof(T); - T* smemInput = reinterpret_cast(&smem[0]); - T* smemResidual = reinterpret_cast(&smem[smemBufferSize]); - T* smemGamma = reinterpret_cast(&smem[2 * smemBufferSize]); + uint32_t const smemBufferSize = blockSize * elemsPerThread * sizeof(T_IN); + T_IN* smemInput = (T_IN*) &smem[0]; + T_IN* smemResidual = (T_IN*) &smem[smemBufferSize]; + T_IN* smemGamma = (T_IN*) &smem[2 * smemBufferSize]; - LamportFlags flag(params.bufferFlags, MNNVLTwoShotStage::NUM_STAGES); - T* input = reinterpret_cast( - flag.getCurLamportBuf(reinterpret_cast(params.bufferInputPtr), MNNVLTwoShotStage::BROADCAST)); + LamportFlags flag(bufferFlags, MNNVLTwoShotStage::NUM_STAGES); + T_IN* input = reinterpret_cast( + flag.getCurLamportBuf(reinterpret_cast(bufferInput), MNNVLTwoShotStage::BROADCAST)); #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) cudaTriggerProgrammaticLaunchCompletion(); @@ -1025,9 +742,8 @@ __global__ __launch_bounds__(1024) void rmsNormLamport(MnnvlAllReduceKernelParam // The offset that current thread should load from. Note that the hidden dimension is split by CGA size and each // block loads a contiguous chunk; // The size of chunk that each block processes - uint32_t const blockChunkSize - = divUp(static_cast(params.tokenDim), clusterSize * kELTS_PER_LOAD) * kELTS_PER_LOAD; - uint32_t const blockLoadOffset = token * params.tokenDim + blockOffset * blockChunkSize; + uint32_t const blockChunkSize = divUp(dim, clusterSize * kELTS_PER_LOAD) * kELTS_PER_LOAD; + uint32_t const blockLoadOffset = token * dim + blockOffset * blockChunkSize; #pragma unroll for (uint32_t i = 0; i < LoadsPerThread; i++) @@ -1037,26 +753,23 @@ __global__ __launch_bounds__(1024) void rmsNormLamport(MnnvlAllReduceKernelParam offsets[i] = blockLoadOffset + threadLoadOffset; } - if constexpr (ar_fusion::HasResidual) - { #pragma unroll - for (uint32_t i = 0; i < LoadsPerThread; i++) + for (uint32_t i = 0; i < LoadsPerThread; i++) + { + uint32_t const threadLoadOffset = (i * loadStride + threadOffset) * kELTS_PER_LOAD; + if (blockOffset * blockChunkSize + threadLoadOffset < dim) { - uint32_t const threadLoadOffset = (i * loadStride + threadOffset) * kELTS_PER_LOAD; - if (blockOffset * blockChunkSize + threadLoadOffset < static_cast(params.tokenDim)) - { - copyF4(&smemResidual[threadLoadOffset], ¶ms.residualInPtr[blockLoadOffset + threadLoadOffset]); - } + copyF4(&smemResidual[threadLoadOffset], &residual[blockLoadOffset + threadLoadOffset]); } - __pipeline_commit(); } + __pipeline_commit(); #pragma unroll for (uint32_t i = 0; i < LoadsPerThread; i++) { uint32_t const threadLoadOffset = (i * loadStride + threadOffset) * kELTS_PER_LOAD; - if (blockOffset * blockChunkSize + threadLoadOffset < static_cast(params.tokenDim)) + if (blockOffset * blockChunkSize + threadLoadOffset < dim) { - copyF4(&smemGamma[threadLoadOffset], ¶ms.gammaPtr[blockOffset * blockChunkSize + threadLoadOffset]); + copyF4(&smemGamma[threadLoadOffset], &gamma[blockOffset * blockChunkSize + threadLoadOffset]); } } __pipeline_commit(); @@ -1072,16 +785,16 @@ __global__ __launch_bounds__(1024) void rmsNormLamport(MnnvlAllReduceKernelParam { uint32_t threadLoadOffset = (i * loadStride + threadOffset) * kELTS_PER_LOAD; - if (blockOffset * blockChunkSize + threadLoadOffset < static_cast(params.tokenDim)) + if (blockOffset * blockChunkSize + threadLoadOffset < dim) { float4* dst4 = reinterpret_cast(&smemInput[threadLoadOffset]); float4 const* src4 = reinterpret_cast(&input[offsets[i]]); - auto loaded = loadPackedVolatile(src4); + float4 value = loadPackedVolatile(src4); // Assume that the 16B were written atomically, so we only need to check one value - valid &= !isLamportDirty(loaded); - *dst4 = loaded.packed; + valid &= !isNegZero(value.x); + *dst4 = value; } } } @@ -1094,27 +807,20 @@ __global__ __launch_bounds__(1024) void rmsNormLamport(MnnvlAllReduceKernelParam for (int i = 0; i < LoadsPerThread; i++) { int threadLoadOffset = (i * loadStride + threadOffset) * kELTS_PER_LOAD; - if (blockOffset * blockChunkSize + threadLoadOffset < static_cast(params.tokenDim)) + if (blockOffset * blockChunkSize + threadLoadOffset < dim) { - PackedVec inp{.packed = loadPacked(&smemInput[threadLoadOffset])}; - PackedVec inpPlusRes = inp; - if constexpr (ar_fusion::HasResidual) - { - PackedVec res{.packed = loadPacked(&smemResidual[threadLoadOffset])}; - inpPlusRes = inp + res; - if constexpr (ar_fusion::HasResidualOut) - { - *reinterpret_cast(¶ms.residualOutPtr[blockLoadOffset + threadLoadOffset]) - = inpPlusRes.packed; - } - } + PackedVec inp{.packed = loadPacked(&smemInput[threadLoadOffset])}; + PackedVec res{.packed = loadPacked(&smemResidual[threadLoadOffset])}; + PackedVec inp_plus_res = inp + res; #pragma unroll for (int j = 0; j < kELTS_PER_LOAD; j++) { - rInput[i * kELTS_PER_LOAD + j] = cuda_cast(inpPlusRes.elements[j]); - threadSum += cuda_cast(inpPlusRes.elements[j] * inpPlusRes.elements[j]); + rInput[i * kELTS_PER_LOAD + j] = cuda_cast(inp_plus_res.elements[j]); + threadSum += cuda_cast(inp_plus_res.elements[j] * inp_plus_res.elements[j]); } + + *reinterpret_cast(&outputPreNorm[blockLoadOffset + threadLoadOffset]) = inp_plus_res.packed; } } @@ -1124,57 +830,49 @@ __global__ __launch_bounds__(1024) void rmsNormLamport(MnnvlAllReduceKernelParam float fullSum = blockSum; // Use CGA Reduction if supported - if constexpr (UseCGA) - { #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - __shared__ float sharedVal[8]; - namespace cg = cooperative_groups; - cg::cluster_group cluster = cg::this_cluster(); - int const numBlocks = cluster.num_blocks(); - if (numBlocks > 1) + __shared__ float sharedVal[8]; + int const numBlocks = cluster.num_blocks(); + if (numBlocks > 1) + { + fullSum = 0.F; + // Need to reduce over the entire cluster + int const blockRank = cluster.block_rank(); + if (threadIdx.x < numBlocks) { - fullSum = 0.F; - // Need to reduce over the entire cluster - int const blockRank = cluster.block_rank(); - if (threadIdx.x < numBlocks) - { - cluster.map_shared_rank(&sharedVal[0], threadIdx.x)[blockRank] = blockSum; - } - // cluster.sync(); - cluster.barrier_wait(cluster.barrier_arrive()); - for (int i = 0; i < numBlocks; ++i) - { - fullSum += sharedVal[i]; - } + cluster.map_shared_rank(&sharedVal[0], threadIdx.x)[blockRank] = blockSum; + } + // cluster.sync(); + cluster.barrier_wait(cluster.barrier_arrive()); + for (int i = 0; i < numBlocks; ++i) + { + fullSum += sharedVal[i]; } -#endif } +#endif - float rcpRms = rsqrtf(fullSum / params.tokenDim + params.epsilon); + float rcpRms = rsqrtf(fullSum / dim + epsilon); #pragma unroll for (int i = 0; i < LoadsPerThread; i++) { - PackedVec r_out; + PackedVec r_out; uint32_t threadLoadOffset = (i * loadStride + threadOffset) * kELTS_PER_LOAD; - if (blockOffset * blockChunkSize + threadLoadOffset < static_cast(params.tokenDim)) + if (blockOffset * blockChunkSize + threadLoadOffset < dim) { - PackedVec gamma = {.packed = loadPacked(&smemGamma[threadLoadOffset])}; + PackedVec gamma = {.packed = loadPacked(&smemGamma[threadLoadOffset])}; #pragma unroll for (uint32_t j = 0; j < kELTS_PER_LOAD; j++) { - r_out.elements[j] = cuda_cast( - cuda_cast(gamma.elements[j]) * rInput[i * kELTS_PER_LOAD + j] * rcpRms); + r_out.elements[j] = cuda_cast( + cuda_cast(gamma.elements[j]) * rInput[i * kELTS_PER_LOAD + j] * rcpRms); } - int const accessIdInToken = (blockOffset * blockChunkSize + threadLoadOffset) / kELTS_PER_LOAD; - int const packedAccessIdx = (blockLoadOffset + threadLoadOffset) / kELTS_PER_LOAD; - writeEpilogueOutput( - r_out, params, blockLoadOffset + threadLoadOffset, packedAccessIdx, accessIdInToken, token); + *reinterpret_cast(&outputNorm[blockLoadOffset + threadLoadOffset]) = r_out.packed; } } - constexpr int kELTS_SIZE = sizeof(T); + constexpr int kELTS_SIZE = sizeof(T_IN); // Issue ACQBLK at the end. Assuming preceding kernel will not modify the buffer_flags. #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) @@ -1182,15 +880,13 @@ __global__ __launch_bounds__(1024) void rmsNormLamport(MnnvlAllReduceKernelParam #endif // Update the buffer pointers - flag.waitAndUpdate({static_cast(divUp(params.numTokens, params.nRanks) * params.nRanks - * params.tokenDim * kELTS_SIZE), - static_cast(params.numTokens * params.tokenDim * kELTS_SIZE), 0, 0}); + flag.waitAndUpdate({static_cast(divUp(numTokens, worldSize) * worldSize * dim * kELTS_SIZE), + static_cast(numTokens * dim * kELTS_SIZE), 0, 0}); } void twoshotAllreduceFusionOp(AllReduceFusionParams const& params) { static int const kSMVersion = tensorrt_llm::common::getSMVersion(); - TLLM_CHECK_WITH_INFO(kSMVersion >= 90, "[MNNVL AllReduceTwoShot] requires SM 90 or newer."); int const numTokens = params.numTokens; int const tokenDim = params.tokenDim; int const numEltsPerThread = sizeof(float4) / getDTypeSize(params.dType); @@ -1219,7 +915,8 @@ void twoshotAllreduceFusionOp(AllReduceFusionParams const& params) "[MNNVL AllReduceTwoShot] Dispatch: grid size: (%d, %d, 1), block_size: 128", numTokens, arNumBlocksPerToken); #define LAUNCH_ALLREDUCE_KERNEL(WORLD_SIZE, T) \ - TLLM_CUDA_CHECK(cudaLaunchKernelEx(&arConfig, &twoshotAllreduceKernel, kernelParams)); + TLLM_CUDA_CHECK(cudaLaunchKernelEx(&arConfig, &twoshotAllreduceKernel, output, input, ucPtrs, \ + mcastPtr, numTokens, tokenDim, params.rank, params.bufferFlags, (!params.rmsNormFusion))); auto dispatchAR = [&](auto* type_ptr) -> bool { using T = std::remove_pointer_t; @@ -1227,11 +924,6 @@ void twoshotAllreduceFusionOp(AllReduceFusionParams const& params) T* mcastPtr = reinterpret_cast(params.multicastPtr); T* output = reinterpret_cast(params.output); T const* input = reinterpret_cast(params.input); - MnnvlAllReduceKernelParams kernelParams{output, reinterpret_cast(params.residualOut), input, - reinterpret_cast(params.residualIn), reinterpret_cast(params.gamma), ucPtrs, - reinterpret_cast(params.bufferPtrLocal), mcastPtr, params.quantOut, params.scaleOut, params.scaleFactor, - numTokens, tokenDim, params.nRanks, params.rank, static_cast(params.epsilon), params.bufferFlags, - !params.rmsNormFusion, params.layout}; switch (params.nRanks) { case 2: LAUNCH_ALLREDUCE_KERNEL(2, T); return true; @@ -1256,19 +948,10 @@ void twoshotAllreduceFusionOp(AllReduceFusionParams const& params) // Launch the rmsnorm lamport kernel if fusion is enabled if (params.rmsNormFusion) { - int const accessGroupSize = requiresTwoAccessScaleGroup(params.pattern) ? 2 : 1; - auto gridConfig = adjustGridConfig(numTokens, tokenDim, numEltsPerThread, accessGroupSize); + auto gridConfig = adjustGridConfig(numTokens, tokenDim, numEltsPerThread); int rnBlockSize = std::get<0>(gridConfig); int rnClusterSize = std::get<1>(gridConfig); int rnLoadsPerThread = std::get<2>(gridConfig); - bool rnUseCGA = rnClusterSize > 1 && rnLoadsPerThread == 1; - if (!rnUseCGA) - { - gridConfig = adjustGridConfig(numTokens, tokenDim, numEltsPerThread, accessGroupSize); - rnBlockSize = std::get<0>(gridConfig); - rnClusterSize = std::get<1>(gridConfig); - rnLoadsPerThread = std::get<2>(gridConfig); - } int rnNumThreads = rnClusterSize * rnBlockSize; dim3 rnGrid(numTokens, rnClusterSize, 1); @@ -1280,15 +963,13 @@ void twoshotAllreduceFusionOp(AllReduceFusionParams const& params) rnConfig.attrs = rnAttrs; rnAttrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; rnAttrs[0].val.programmaticStreamSerializationAllowed = tensorrt_llm::common::getEnvEnablePDL() ? 1 : 0; - rnConfig.numAttrs = 1U; - if (rnUseCGA) - { - rnAttrs[1].id = cudaLaunchAttributeClusterDimension; - rnAttrs[1].val.clusterDim.x = 1; - rnAttrs[1].val.clusterDim.y = rnClusterSize; - rnAttrs[1].val.clusterDim.z = 1; - rnConfig.numAttrs = 2U; - } + rnAttrs[1].id = cudaLaunchAttributeClusterDimension; + rnAttrs[1].val.clusterDim.x = 1; + rnAttrs[1].val.clusterDim.y = rnClusterSize; + rnAttrs[1].val.clusterDim.z = 1; + rnConfig.numAttrs = (kSMVersion >= 90) ? 2U : 1U; + + bool const rnUseCGA = kSMVersion >= 90 && rnClusterSize > 1; int const dimPadded = divUp(tokenDim, numEltsPerThread * rnNumThreads) * numEltsPerThread * rnNumThreads; int const iters = dimPadded / rnNumThreads; @@ -1300,88 +981,40 @@ void twoshotAllreduceFusionOp(AllReduceFusionParams const& params) "threads_needed: %d", numTokens, rnClusterSize, rnBlockSize, rnClusterSize, rnLoadsPerThread, divUp(tokenDim, numEltsPerThread)); -#define RUN_RMSNORM_KERNEL(T, PATTERN, USE_CGA, LOADS_PER_THREAD) \ - TLLM_CUDA_CHECK(cudaFuncSetAttribute(&rmsNormLamport, \ - cudaFuncAttributeMaxDynamicSharedMemorySize, smemSize)); \ +#define RUN_RMSNORM_KERNEL(T_IN, T_OUT, LOADS_PER_THREAD) \ + TLLM_CUDA_CHECK(cudaFuncSetAttribute( \ + &rmsNormLamport, cudaFuncAttributeMaxDynamicSharedMemorySize, smemSize)); \ rnConfig.dynamicSmemBytes = smemSize; \ - TLLM_CUDA_CHECK( \ - cudaLaunchKernelEx(&rnConfig, &rmsNormLamport, kernelParams)); - -#define DISPATCH_RMSNORM_PATTERN(T, USE_CGA, LOADS_PER_THREAD) \ - if (params.pattern == ar_fusion::AllReduceFusionPattern::kARResidualRMSNorm) \ - { \ - RUN_RMSNORM_KERNEL(T, ar_fusion::AllReduceFusionPattern::kARResidualRMSNorm, USE_CGA, LOADS_PER_THREAD); \ - } \ - else if (params.pattern == ar_fusion::AllReduceFusionPattern::kARResidualRMSNormFP8Quant) \ - { \ - RUN_RMSNORM_KERNEL( \ - T, ar_fusion::AllReduceFusionPattern::kARResidualRMSNormFP8Quant, USE_CGA, LOADS_PER_THREAD); \ - } \ - else if (params.pattern == ar_fusion::AllReduceFusionPattern::kARResidualRMSNormOutFP8Quant) \ - { \ - RUN_RMSNORM_KERNEL( \ - T, ar_fusion::AllReduceFusionPattern::kARResidualRMSNormOutFP8Quant, USE_CGA, LOADS_PER_THREAD); \ - } \ - else if (params.pattern == ar_fusion::AllReduceFusionPattern::kARResidualRMSNormFP4Quant) \ - { \ - if constexpr (!std::is_same_v) \ - { \ - RUN_RMSNORM_KERNEL( \ - T, ar_fusion::AllReduceFusionPattern::kARResidualRMSNormFP4Quant, USE_CGA, LOADS_PER_THREAD); \ - } \ - else \ - { \ - TLLM_CHECK_WITH_INFO(false, "[MNNVL AllReduceTwoShot] NVFP4 quantization does not support FP32 input."); \ - } \ - } \ - else if (params.pattern == ar_fusion::AllReduceFusionPattern::kARResidualRMSNormOutFP4Quant) \ - { \ - if constexpr (!std::is_same_v) \ - { \ - RUN_RMSNORM_KERNEL( \ - T, ar_fusion::AllReduceFusionPattern::kARResidualRMSNormOutFP4Quant, USE_CGA, LOADS_PER_THREAD); \ - } \ - else \ - { \ - TLLM_CHECK_WITH_INFO(false, "[MNNVL AllReduceTwoShot] NVFP4 quantization does not support FP32 input."); \ - } \ - } \ - else if (params.pattern == ar_fusion::AllReduceFusionPattern::kARRMSNorm) \ - { \ - RUN_RMSNORM_KERNEL(T, ar_fusion::AllReduceFusionPattern::kARRMSNorm, USE_CGA, LOADS_PER_THREAD); \ - } \ - else \ - { \ - TLLM_CHECK_WITH_INFO(false, "[MNNVL AllReduceTwoShot] Unsupported RMSNorm fusion pattern."); \ - } + TLLM_CUDA_CHECK(cudaLaunchKernelEx(&rnConfig, &rmsNormLamport, residualOut, output, \ + bufferInput, gamma, static_cast(params.epsilon), residualIn, numTokens, tokenDim, params.nRanks, \ + params.bufferFlags)); // C++ 17 does not support capturing structured bindings auto dispatchRN = [&, rnLoadsPerThread](auto* type_ptr) { - using T = std::remove_pointer_t; - MnnvlAllReduceKernelParams kernelParams{reinterpret_cast(params.output), - reinterpret_cast(params.residualOut), reinterpret_cast(params.input), - reinterpret_cast(params.residualIn), reinterpret_cast(params.gamma), - reinterpret_cast(params.bufferPtrsDev), reinterpret_cast(params.bufferPtrLocal), - reinterpret_cast(params.multicastPtr), params.quantOut, params.scaleOut, params.scaleFactor, - numTokens, tokenDim, params.nRanks, params.rank, static_cast(params.epsilon), params.bufferFlags, - false, params.layout}; + using T_IN = std::remove_pointer_t; + using T_OUT = T_IN; + T_OUT* residualOut = reinterpret_cast(params.residualOut); + T_OUT* output = reinterpret_cast(params.output); + T_IN* bufferInput = reinterpret_cast(params.bufferPtrLocal); + T_IN const* gamma = reinterpret_cast(params.gamma); + T_IN const* residualIn = reinterpret_cast(params.residualIn); if (rnUseCGA) { - DISPATCH_RMSNORM_PATTERN(T, true, 1); + RUN_RMSNORM_KERNEL(T_IN, T_OUT, 1); } else { switch (rnLoadsPerThread) { - case 1: DISPATCH_RMSNORM_PATTERN(T, false, 1); break; - case 2: DISPATCH_RMSNORM_PATTERN(T, false, 2); break; - case 3: DISPATCH_RMSNORM_PATTERN(T, false, 3); break; - case 4: DISPATCH_RMSNORM_PATTERN(T, false, 4); break; - case 5: DISPATCH_RMSNORM_PATTERN(T, false, 5); break; - case 6: DISPATCH_RMSNORM_PATTERN(T, false, 6); break; - case 7: DISPATCH_RMSNORM_PATTERN(T, false, 7); break; - case 8: DISPATCH_RMSNORM_PATTERN(T, false, 8); break; + case 1: RUN_RMSNORM_KERNEL(T_IN, T_OUT, 1); break; + case 2: RUN_RMSNORM_KERNEL(T_IN, T_OUT, 2); break; + case 3: RUN_RMSNORM_KERNEL(T_IN, T_OUT, 3); break; + case 4: RUN_RMSNORM_KERNEL(T_IN, T_OUT, 4); break; + case 5: RUN_RMSNORM_KERNEL(T_IN, T_OUT, 5); break; + case 6: RUN_RMSNORM_KERNEL(T_IN, T_OUT, 6); break; + case 7: RUN_RMSNORM_KERNEL(T_IN, T_OUT, 7); break; + case 8: RUN_RMSNORM_KERNEL(T_IN, T_OUT, 8); break; default: return false; } } @@ -1396,7 +1029,6 @@ void twoshotAllreduceFusionOp(AllReduceFusionParams const& params) TLLM_CHECK_WITH_INFO(false, "[MNNVL AllReduceTwoShot] Failed to dispatch rmsnorm lamport kernel."); } #undef RUN_RMSNORM_KERNEL -#undef DISPATCH_RMSNORM_PATTERN } } diff --git a/cpp/tensorrt_llm/kernels/communicationKernels/mnnvlAllreduceKernels.h b/cpp/tensorrt_llm/kernels/communicationKernels/mnnvlAllreduceKernels.h index 2a228e815b8d..5361f50221b4 100644 --- a/cpp/tensorrt_llm/kernels/communicationKernels/mnnvlAllreduceKernels.h +++ b/cpp/tensorrt_llm/kernels/communicationKernels/mnnvlAllreduceKernels.h @@ -1,5 +1,5 @@ /* - * Copyright (c) 2022-2026, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -17,7 +17,6 @@ #define TRTLLM_MNNVL_ALLREDUCE_KERNELS_H #include "tensorrt_llm/common/config.h" -#include "tensorrt_llm/kernels/communicationKernels/allReduceFusionKernels.h" #include #include @@ -49,8 +48,6 @@ struct AllReduceFusionParams void* multicastPtr; //!< Multicast buffer pointer. uint32_t* bufferFlags; //!< Synchronization flags for coordinating communication phases bool rmsNormFusion; //!< Whether to fuse RMS normalization with the AllReduce operation - ar_fusion::AllReduceFusionPattern pattern - = ar_fusion::AllReduceFusionPattern::kAllReduce; //!< Fused epilogue pattern //! @} @@ -62,13 +59,9 @@ struct AllReduceFusionParams void const* gamma; //!< Gamma parameters for RMS normalization (used when rmsnormFusion=true) double epsilon; //!< Epsilon value for RMS normalization numerical stability (used when rmsnormFusion=true) - void* residualOut = nullptr; //!< Output tensor for residual connection result (used when rmsnormFusion=true) - void* output = nullptr; //!< Output tensor containing the AllReduce or RMSNorm result - void* quantOut = nullptr; //!< Quantized RMSNorm output (used by quantized fusion patterns) - void* scaleOut = nullptr; //!< NVFP4 scale-factor output (used by NVFP4 fusion patterns) - float const* scaleFactor = nullptr; //!< Quantization scale factor - QuantizationSFLayout layout = QuantizationSFLayout::SWIZZLED; //!< NVFP4 scale-factor layout - cudaStream_t stream; //!< CUDA stream for asynchronous kernel execution + void* residualOut; //!< Output tensor for residual connection result (used when rmsnormFusion=true) + void* output; //!< Final output tensor containing the AllReduce result + cudaStream_t stream; //!< CUDA stream for asynchronous kernel execution //! @} }; diff --git a/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu b/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu index 472a5877a80d..91cb5725fede 100644 --- a/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu +++ b/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu @@ -210,14 +210,6 @@ __device__ __forceinline__ int compute_target_rank_id(int expert_id, int base, i return remainder + (expert_id - split) / base; } -// Test bit `rank` in a kRankMaskWords-wide little-endian uint64 bitmask. -// Word 0 covers ranks 0..63, word 1 covers ranks 64..127, etc. -// `rank >> 6` and `rank & 63` divide / modulo by 64. -__device__ __forceinline__ bool is_rank_active(uint64_t const* mask, int rank) -{ - return (mask[rank >> 6] >> (rank & 63)) & 1ULL; -} - // ============================================================================ // Helper Functions for Vectorized Memory Operations // ============================================================================ @@ -392,7 +384,7 @@ __global__ void moeA2APrepareDispatchKernel( // Dispatch Kernels // ============================================================================ -template +template __global__ void moeA2ADispatchKernel(int32_t const* token_selected_experts, // [local_num_tokens, TOP_K] const DispatchKernelPointers ptrs, // Struct containing all kernel pointers int num_payloads, // Number of payloads @@ -424,7 +416,7 @@ __global__ void moeA2ADispatchKernel(int32_t const* token_selected_experts, // [ int* smem_topk_target_ranks = smem; int* smem_topk_send_indices = smem + TOP_K; - uint64_t already_copied[kRankMaskWords] = {}; + uint64_t already_copied = 0; // Precompute the ceil/floor partition parameters once per thread, outside the // per-token TOP_K loop. The fast path (remainder == 0) then collapses to a single // integer divide per call, matching the pre-PR uniform-partition cost exactly. @@ -440,17 +432,7 @@ __global__ void moeA2ADispatchKernel(int32_t const* token_selected_experts, // [ // Supports the non-divisible case where num_experts % ep_size != 0. int target_rank = compute_target_rank_id(expert_id, ep_base, ep_remainder); - int const mask_word = target_rank >> 6; - uint64_t const mask_bit = 1ULL << (target_rank & 63); - bool const target_already_copied = (already_copied[mask_word] & mask_bit) != 0; - bool skip_target = target_already_copied; - if constexpr (ENABLE_RANK_MASK) - { - // This is a fail-closed safety guard until post-commit routing is enforced end to end. - // A masked route is not valid model output; the failed execution epoch must be discarded. - skip_target = skip_target || !is_rank_active(ptrs.active_rank_mask, target_rank); - } - if (skip_target) + if (already_copied & (1ULL << target_rank)) { if (thread_idx == 0) { @@ -475,7 +457,7 @@ __global__ void moeA2ADispatchKernel(int32_t const* token_selected_experts, // [ smem_topk_target_ranks[k] = target_rank; smem_topk_send_indices[k] = dst_token_idx; } - already_copied[mask_word] |= mask_bit; + already_copied |= 1ULL << target_rank; } // Sync before dispatching data ThreadingPolicy::sync(); @@ -529,16 +511,10 @@ __global__ void moeA2ADispatchKernel(int32_t const* token_selected_experts, // [ if (is_last_token) { -// Store send_counters to recv_counters. -// Skip masked target ranks: their symmetric memory may be inaccessible. +// Store send_counters to recv_counters #pragma unroll 1 // No unroll as one iter is typically enough for (int target_rank = lane_id; target_rank < ep_size; target_rank += warpSize) { - if constexpr (ENABLE_RANK_MASK) - { - if (!is_rank_active(ptrs.active_rank_mask, target_rank)) - continue; - } int send_count = ptrs.send_counters[target_rank]; ptrs.recv_counters[target_rank][rank_id] = send_count; } @@ -546,15 +522,9 @@ __global__ void moeA2ADispatchKernel(int32_t const* token_selected_experts, // [ if constexpr (ENABLE_EPLB) { // Write local stats into peer buffers before the release fence below. - // Skip masked target ranks for the same reason as above. #pragma unroll 1 for (int target_rank = 0; target_rank < ep_size; ++target_rank) { - if constexpr (ENABLE_RANK_MASK) - { - if (!is_rank_active(ptrs.active_rank_mask, target_rank)) - continue; - } int* target_stats = ptrs.eplb_gathered_stats[target_rank]; for (int expert_id = lane_id; expert_id < eplb_stats_num_experts; expert_id += warpSize) { @@ -573,16 +543,9 @@ __global__ void moeA2ADispatchKernel(int32_t const* token_selected_experts, // [ #else asm volatile("fence.acq_rel.sys;"); #endif - // Signal completion to all active peers; skip dead ranks (their symmetric memory - // is unreachable). #pragma unroll 1 // No unroll as one iter is typically enough for (int target_rank = lane_id; target_rank < ep_size; target_rank += warpSize) { - if constexpr (ENABLE_RANK_MASK) - { - if (!is_rank_active(ptrs.active_rank_mask, target_rank)) - continue; - } uint32_t* flag_addr = &ptrs.completion_flags[target_rank][rank_id]; asm volatile("st.relaxed.sys.u32 [%0], %1;" ::"l"(flag_addr), "r"(expected_value)); @@ -592,16 +555,9 @@ __global__ void moeA2ADispatchKernel(int32_t const* token_selected_experts, // [ #endif } - // Wait for all active peers to signal; skip dead ranks (otherwise we would - // spin forever — this is the bug the rank-mask is here to prevent). #pragma unroll 1 // No unroll for (int peer_rank = lane_id; peer_rank < ep_size; peer_rank += warpSize) { - if constexpr (ENABLE_RANK_MASK) - { - if (!is_rank_active(ptrs.active_rank_mask, peer_rank)) - continue; - } bool flag_set = false; auto s = clock64(); do @@ -647,14 +603,8 @@ void moe_a2a_dispatch_launch(MoeA2ADispatchParams const& params) // Validate parameters TLLM_CHECK(params.top_k > 0 && params.top_k <= kMaxTopK); TLLM_CHECK(params.ep_size > 0 && params.ep_size <= kMaxRanks); - TLLM_CHECK(params.ep_rank >= 0 && params.ep_rank < params.ep_size); TLLM_CHECK(params.local_num_tokens >= 0); TLLM_CHECK(params.num_payloads > 0 && params.num_payloads <= kMaxPayloads); - if (params.enable_rank_mask) - { - TLLM_CHECK_WITH_INFO((params.active_rank_mask[params.ep_rank >> 6] >> (params.ep_rank & 63)) & 1ULL, - "active_rank_mask must mark the local ep_rank (%d) as active", params.ep_rank); - } // Prepare kernel pointers struct DispatchKernelPointers kernel_ptrs = {}; @@ -692,12 +642,6 @@ void moe_a2a_dispatch_launch(MoeA2ADispatchParams const& params) kernel_ptrs.topk_send_indices = params.topk_send_indices; kernel_ptrs.eplb_local_stats = params.eplb_local_stats; - // Copy active-rank bitmask into the kernel pointers struct - for (int w = 0; w < kRankMaskWords; ++w) - { - kernel_ptrs.active_rank_mask[w] = params.active_rank_mask[w]; - } - int const kBlockSize = tensorrt_llm::common::getEnvMoeA2ADispatchBlockSize(); // One block per token: grid_size == local_num_tokens. If 0, launch a single block to @@ -708,15 +652,12 @@ void moe_a2a_dispatch_launch(MoeA2ADispatchParams const& params) grid_size = 1; } int shared_bytes = 2 * params.top_k * (int) sizeof(int); - SWITCH_BOOL(params.enable_rank_mask, ENABLE_RANK_MASK, {SWITCH_BOOL(params.enable_eplb, EPLB_STATS, { - SWITCH_TOP_K(params.top_k, TOP_K, { - auto kernel_fn = moeA2ADispatchKernel; - launchWithPdlWhenEnabled("moeA2ADispatchKernel", kernel_fn, grid_size, kBlockSize, shared_bytes, - params.stream, params.token_selected_experts, kernel_ptrs, params.num_payloads, - params.max_tokens_per_rank, params.local_num_tokens, params.ep_rank, params.ep_size, params.num_experts, - params.eplb_stats_num_experts); - }); - })}) + SWITCH_BOOL(params.enable_eplb, EPLB_STATS, SWITCH_TOP_K(params.top_k, TOP_K, { + auto kernel_fn = moeA2ADispatchKernel; + launchWithPdlWhenEnabled("moeA2ADispatchKernel", kernel_fn, grid_size, kBlockSize, shared_bytes, params.stream, + params.token_selected_experts, kernel_ptrs, params.num_payloads, params.max_tokens_per_rank, + params.local_num_tokens, params.ep_rank, params.ep_size, params.num_experts, params.eplb_stats_num_experts); + })) } // ============================================================================ @@ -1172,7 +1113,7 @@ __global__ void moeA2APrepareCombineKernel(uint8_t* recv_buffer_bytes, void cons // Generic Combine Kernel Implementation (Templated by data type) // ============================================================================ -template +template __global__ void moeA2ACombineKernel( const CombineKernelPointers ptrs, // Combine-specific struct, src_data_ptrs[0] is output int max_tokens_per_rank, int elements_per_token, int local_num_tokens, int rank_id, int ep_size, @@ -1212,16 +1153,9 @@ __global__ void moeA2ACombineKernel( if (blockIdx.x == 0) { - // Signal readiness to all active peers; skip dead ranks (their symmetric memory - // is unreachable). #pragma unroll 1 // No unroll for (int peer_rank = lane_id; peer_rank < ep_size; peer_rank += warpSize) { - if constexpr (ENABLE_RANK_MASK) - { - if (!is_rank_active(ptrs.active_rank_mask, peer_rank)) - continue; - } uint32_t* flag_addr = &ptrs.completion_flags[peer_rank][rank_id]; asm volatile("st.relaxed.sys.u32 [%0], %1;" ::"l"(flag_addr), "r"(expected_value)); #if ENABLE_DEBUG_PRINT @@ -1231,16 +1165,9 @@ __global__ void moeA2ACombineKernel( } } - // Wait for all active peers to signal; skip dead ranks (otherwise we would spin - // forever — this is the bug the rank-mask is here to prevent). #pragma unroll 1 // No unroll for (int peer_rank = lane_id; peer_rank < ep_size; peer_rank += warpSize) { - if constexpr (ENABLE_RANK_MASK) - { - if (!is_rank_active(ptrs.active_rank_mask, peer_rank)) - continue; - } bool flag_set = false; auto s = clock64(); do @@ -1344,14 +1271,8 @@ void moe_a2a_combine_launch(MoeA2ACombineParams const& params) // Validate parameters TLLM_CHECK(params.top_k > 0 && params.top_k <= kMaxTopK); TLLM_CHECK(params.ep_size > 0 && params.ep_size <= kMaxRanks); - TLLM_CHECK(params.ep_rank >= 0 && params.ep_rank < params.ep_size); TLLM_CHECK(params.local_num_tokens >= 0); TLLM_CHECK(params.elements_per_token > 0); - if (params.enable_rank_mask) - { - TLLM_CHECK_WITH_INFO((params.active_rank_mask[params.ep_rank >> 6] >> (params.ep_rank & 63)) & 1ULL, - "active_rank_mask must mark the local ep_rank (%d) as active", params.ep_rank); - } // Configure kernel launch (one block per token). int const kBlockSize = tensorrt_llm::common::getEnvMoeA2ACombineBlockSize(); @@ -1385,12 +1306,6 @@ void moe_a2a_combine_launch(MoeA2ACombineParams const& params) kernel_ptrs.topk_target_ranks = params.topk_target_ranks; kernel_ptrs.topk_send_indices = params.topk_send_indices; - // Copy active-rank bitmask into the kernel pointers struct - for (int w = 0; w < kRankMaskWords; ++w) - { - kernel_ptrs.active_rank_mask[w] = params.active_rank_mask[w]; - } - // stride_per_token: byte distance between tokens in the recv buffer. // FP8 external payload: EPT × 1 (compact FP8 layout) // FP8 in-place / non-FP8: EPT × sizeof(PayloadT) (payload-dtype stride) @@ -1406,16 +1321,14 @@ void moe_a2a_combine_launch(MoeA2ACombineParams const& params) auto const effective_dtype = params.use_low_precision ? nvinfer1::DataType::kFP8 : params.dtype; // Launch appropriate kernel with compact macros - SWITCH_BOOL(params.enable_rank_mask, ENABLE_RANK_MASK, { - SWITCH_DTYPE(effective_dtype, TKernelType, { - SWITCH_TOP_K(params.top_k, TOP_K, { - auto kernel_fn = moeA2ACombineKernel; - launchWithPdlWhenEnabled("moeA2ACombineKernel", kernel_fn, grid, kBlockSize, 0, params.stream, - kernel_ptrs, params.max_tokens_per_rank, params.elements_per_token, params.local_num_tokens, - params.ep_rank, params.ep_size, stride_per_token); - }); + SWITCH_DTYPE(effective_dtype, TKernelType, { + SWITCH_TOP_K(params.top_k, TOP_K, { + auto kernel_fn = moeA2ACombineKernel; + launchWithPdlWhenEnabled("moeA2ACombineKernel", kernel_fn, grid, kBlockSize, 0, params.stream, kernel_ptrs, + params.max_tokens_per_rank, params.elements_per_token, params.local_num_tokens, params.ep_rank, + params.ep_size, stride_per_token); }); - }) + }); } // Kernel to sanitize expert ids for invalid tokens diff --git a/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h b/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h index 177293684874..317ff4d2240c 100644 --- a/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h +++ b/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h @@ -26,12 +26,9 @@ namespace kernels::moe_comm { // Configuration constants -static constexpr int kMaxTopK = 22; // Maximum top-k experts per token -static constexpr int kMaxPayloads = 4; // Maximum number of different payload types -static constexpr int kMaxRanks = 128; // Maximum supported EP size (covers NVL72 with headroom) -static constexpr int kRankMaskWords = 2; // uint64 words to hold the active-rank bitmask - // (kRankMaskWords * 64 must be >= kMaxRanks) -static_assert(kRankMaskWords * 64 >= kMaxRanks, "active_rank_mask too small for kMaxRanks"); +static constexpr int kMaxTopK = 22; // Maximum top-k experts per token +static constexpr int kMaxPayloads = 4; // Maximum number of different payload types +static constexpr int kMaxRanks = 64; // Maximum supported EP size // Describes a single payload type to be communicated struct PayloadDescriptor @@ -62,17 +59,12 @@ struct DispatchKernelPointers int* local_token_counter; // Atomic counter for completed tokens // Top-K compact routing info per local token (size: [local_num_tokens, top_k]) - int* topk_target_ranks; // target rank per k, -1 for invalid or duplicate routes - int* topk_send_indices; // dst index per k, -1 for invalid or duplicate routes + int* topk_target_ranks; // target rank per k, -1 for duplicates + int* topk_send_indices; // dst index per k, -1 for duplicates // Optional: Statistics for EPLB int const* eplb_local_stats; // [eplb_stats_num_experts] int* eplb_gathered_stats[kMaxRanks]; // [ep_size, eplb_stats_num_experts] per rank - - // Active-rank bitmask: bit i set => rank i participates in this collective. - // Word 0 covers ranks 0..63; word 1 covers ranks 64..127. The masked kernel - // rejects inactive route targets and skips their peer counters, stats, and flags. - uint64_t active_rank_mask[kRankMaskWords]; }; // Combine kernel pointers - non-const output in src_data_ptrs[0], const recv buffers @@ -88,12 +80,8 @@ struct CombineKernelPointers uint32_t* flag_val; // The value of the flag for this round (stored on the local rank) // Top-K compact routing info per local token (size: [local_num_tokens, top_k]) - int const* topk_target_ranks; // target rank per k, -1 for invalid or duplicate routes - int const* topk_send_indices; // dst index per k, -1 for invalid or duplicate routes - - // Active-rank bitmask: see DispatchKernelPointers::active_rank_mask. Combine skips - // completion flag writes/waits to/from inactive peers. - uint64_t active_rank_mask[kRankMaskWords]; + int const* topk_target_ranks; // target rank per k, -1 for duplicates + int const* topk_send_indices; // dst index per k, -1 for duplicates }; // Dispatch phase parameters @@ -137,16 +125,6 @@ struct MoeA2ADispatchParams int const* eplb_local_stats; // [eplb_stats_num_experts] int* eplb_gathered_stats[kMaxRanks]; // [ep_size, eplb_stats_num_experts] per rank - // Whether to instantiate a kernel with active-rank checks. - // This is a launch-lifetime mode, independent of future execution-abort handling. - bool enable_rank_mask{false}; - - // Active-rank bitmask: see DispatchKernelPointers::active_rank_mask. Used only when - // enable_rank_mask is true; defaults to all-ones for backwards-compatible behavior. - // The mask is copied by value into kernel arguments. Rank-mask mode must reject - // CUDA graph replay until generation-scoped invalidation and recapture are available. - uint64_t active_rank_mask[kRankMaskWords] = {~uint64_t{0}, ~uint64_t{0}}; - // CUDA stream cudaStream_t stream; }; @@ -192,16 +170,6 @@ struct MoeA2ACombineParams // rank has signaled the target rank void const* recv_buffers[kMaxRanks]; // Per-rank receive buffers (only for single payload) - // Whether to instantiate a kernel with active-rank checks in peer synchronization. - // This is a launch-lifetime mode, independent of future execution-abort handling. - bool enable_rank_mask{false}; - - // Active-rank bitmask: see DispatchKernelPointers::active_rank_mask. Used only when - // enable_rank_mask is true; defaults to all-ones for backwards-compatible behavior. - // The mask is copied by value into kernel arguments. Rank-mask mode must reject - // CUDA graph replay until generation-scoped invalidation and recapture are available. - uint64_t active_rank_mask[kRankMaskWords] = {~uint64_t{0}, ~uint64_t{0}}; - // CUDA stream cudaStream_t stream; }; diff --git a/cpp/tensorrt_llm/kernels/compressorKernels/CMakeLists.txt b/cpp/tensorrt_llm/kernels/compressorKernels/CMakeLists.txt deleted file mode 100644 index 46544a47b11f..000000000000 --- a/cpp/tensorrt_llm/kernels/compressorKernels/CMakeLists.txt +++ /dev/null @@ -1,25 +0,0 @@ -# -# 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. -# - -set(SRC_CU compressorKernels.cu) - -add_library(compressorKernels_src OBJECT ${SRC_CU}) -set_property(TARGET compressorKernels_src PROPERTY POSITION_INDEPENDENT_CODE ON) -set_property(TARGET compressorKernels_src PROPERTY CUDA_RESOLVE_DEVICE_SYMBOLS - ON) -target_compile_options(compressorKernels_src - PRIVATE $<$:--use_fast_math>) diff --git a/cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.cu b/cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.cu deleted file mode 100644 index 1457efeb3138..000000000000 --- a/cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.cu +++ /dev/null @@ -1,1892 +0,0 @@ -/* - * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. - * - * 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. - */ - -// ============================================================================ -// Compressor Kernels — DeepSeek-V4 KV Cache Compression -// ============================================================================ -// -// This file implements CUDA kernels for KV cache compression in the DeepSeek-V4 -// sparse attention system. The compressor reduces sequences of input tokens -// into fewer compressed tokens via learned weighted averaging (online softmax), -// then post-processes and scatters results into a paged KV cache. -// -// Three kernels are provided: -// -// 1. pagedKvCompressKernel — Decode path (single/few new tokens per batch). -// Loads prior compressor state from paged memory, performs online softmax -// with the new token(s), writes updated state back, and emits a compressed -// output token when compress_ratio tokens have been accumulated. -// -// 2. prefillReductionKernel — Prefill path (many tokens per batch). -// Processes full chunks of compress_ratio tokens in one shot via online -// softmax reduction over the input sequence. Also saves compressor state -// for any remainder tokens that don't form a complete chunk. -// -// 3. postProcessScatterKernel — Fused post-processing + paged cache write. -// Takes compressed output tokens and applies: RMSNorm → RoPE → Hadamard -// transform → optional V4-Pro QDQ → scatter to paged KV cache. Supports -// default, FP8, and V4-Pro MXFP8/MXFP4 QDQ cache modes. -// Keeps all intermediate values in float32 registers to avoid extra DRAM -// round-trips. -// -// Vectorization strategy: -// All kernels use 128-bit vectorized loads/stores (float4 / 8×bf16). -// VEC = number of elements per thread, chosen so that NTHRD = HEAD_DIM/VEC >= 32. -// For HEAD_DIM=128, bf16: VEC=4, NTHRD=32. For HEAD_DIM=512, bf16: VEC=8, NTHRD=64. -// -// Overlap mode (compress_ratio=4): -// When enabled, state_dim = 2*head_dim and the compressor uses overlapping -// windows: each compressed output is derived from both the previous and current -// chunk of compress_ratio tokens (previous chunk → first head_dim features, -// current chunk → second head_dim features). This doubles the state stored -// per position but improves compression quality. -// -// Template parameters: -// HEAD_DIM — Head dimension (128 or 512) -// KV_SCORE_ELEM_BYTES — kv_score element size (2=bf16, 4=fp32) -// STATE_ELEM_BYTES — compressor state element size (2=bf16, 4=fp32) -// SCALE_TYPE — Output cache scale/dtype for postProcessScatterKernel -// ============================================================================ - -#include "tensorrt_llm/kernels/compressorKernels/compressorKernels.h" - -#include "tensorrt_llm/common/assert.h" -#include -#include -#include -#include -#include -#include -#include - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels::compressor -{ - -// ============================================================================ -// Helper functions -// ============================================================================ - -// Full-warp butterfly reductions via __shfl_xor_sync (all 32 lanes participate). -__device__ inline float warpReduceSum(float val) -{ - for (int mask = 16; mask > 0; mask >>= 1) - val += __shfl_xor_sync(0xFFFFFFFF, val, mask); - return val; -} - -__device__ inline float warpReduceMax(float val) -{ - for (int mask = 16; mask > 0; mask >>= 1) - val = fmaxf(val, __shfl_xor_sync(0xFFFFFFFF, val, mask)); - return val; -} - -// Runtime-dispatched element load (bf16 or fp32 → float). Used in the decode -// kernel where elem_bytes is a runtime parameter from paged state buffers. -__device__ inline float loadAsFloat(void const* base, int64_t offset, int elem_bytes) -{ - if (elem_bytes == 2) - return __bfloat162float(reinterpret_cast<__nv_bfloat16 const*>(base)[offset]); - else - return reinterpret_cast(base)[offset]; -} - -__device__ inline void storeFromFloat(void* base, int64_t offset, float val, int elem_bytes) -{ - if (elem_bytes == 2) - reinterpret_cast<__nv_bfloat16*>(base)[offset] = __float2bfloat16_rn(val); - else - reinterpret_cast(base)[offset] = val; -} - -// Bit-hack ceil(log2(x)) for x>0: equivalent to V4 reference fast_log2_ceil. -__device__ inline int fastLog2Ceil(float x) -{ - uint32_t const bits = __float_as_uint(x); - int const exp_part = static_cast((bits >> 23) & 0xFFu) - 127; - uint32_t const man_bits = bits & 0x007FFFFFu; - return exp_part + (man_bits != 0u ? 1 : 0); -} - -// Bit-hack 2^n for integer n: equivalent to V4 reference fast_pow2. -__device__ inline float fastPow2(int n) -{ - uint32_t const bits = static_cast(n + 127) << 23; - return __uint_as_float(bits); -} - -// V4-Pro fast_round_scale: 2^ceil(log2(amax * max_value_inv)). Operand order -// (multiplication vs `amax / max_value`) matches V4 byte-for-byte; using fp32 -// bit hacks avoids log2f/exp2f rounding so the resulting power-of-2 is exact. -__device__ inline float roundedPow2Scale(float amax, float max_value_inv, float min_amax) -{ - float const clamped_amax = fmaxf(amax, min_amax); - return fastPow2(fastLog2Ceil(clamped_amax * max_value_inv)); -} - -// Hardware FP4 (e2m1) round-trip via __nv_fp4_e2m1 (cuda_fp4.h). The -// constructor uses the SM100 PTX `cvt.rn.satfinite.e2m1x2.f32` cast -// (round-to-nearest-even with finite saturation), matching V4-Pro's -// Cast(FP4) byte-for-byte. Verified against the V4-Pro reference LUT -// (e2m1 levels {0, 0.5, 1, 1.5, 2, 3, 4, 6} with midpoint thresholds -// {0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0}) over 1.1M sweep inputs covering -// every exact level, every midpoint ±2 ULPs, out-of-range, subnormals, and -// dense random fp32 -- zero mismatches. The Python test reference still -// uses the explicit LUT so the kernel is checked against an independent -// software model rather than the HW cast itself. -__device__ inline uint8_t toUe8m0(float val) -{ - __nv_fp8_e8m0 out; - out.__x = __nv_cvt_float_to_e8m0(val, __NV_SATFINITE, cudaRoundPosInf); - return out.__x; -} - -__device__ inline uint8_t packE2M1x2(float lo, float hi) -{ -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) - uint32_t val; - asm volatile( - "{\n" - ".reg .b8 byte0;\n" - "cvt.rn.satfinite.e2m1x2.f32 byte0, %2, %1;\n" - "mov.b32 %0, {byte0, byte0, byte0, byte0};\n" - "}" - : "=r"(val) - : "f"(lo), "f"(hi)); - return static_cast(val); -#else - return 0; -#endif -} - -// Vectorized load/store types: maps byte-width to CUDA vector type. -template -struct VecType; - -template <> -struct VecType<4> -{ - using type = unsigned int; -}; // 32-bit: 2 bf16 or 1 fp32 - -template <> -struct VecType<8> -{ - using type = uint2; -}; // 64-bit: 4 bf16 or 2 fp32 - -template <> -struct VecType<16> -{ - using type = uint4; -}; // 128-bit: 8 bf16 or 4 fp32 - -// Cache scale / pre-store quantization type for postProcessScatterKernel. -// The store dtype is implied: kNone keeps the input dtype (elem_bytes -// controls bf16 vs fp32), kFP8* writes one byte per element, and -// kMXFP4Blockwise writes packed FP4 (two values per byte) plus per-32 -// UE8M0 scale bytes. -enum class CacheScaleType -{ - kNone = 0, - kFP8PerTensor = 1, // FP8 E4M3 with static scale=1.0 - kFP8Blockwise = 2, // FP8 E4M3 with per-128-element fp32 scales - kMXFP4Blockwise = 3 // packed FP4 E2M1 with per-32 UE8M0 scales -}; - -// ============================================================================ -// Decode Kernel: pagedKvCompressKernel -// -// Template: -// NEXT_N: number of new tokens per sequence in this decode step (1-8) -// -// Grid: (batch_size) — one block per batch element -// Block: (NTHRD) where NTHRD = HEAD_DIM / VEC (>= 32 threads) -// -// Algorithm per batch element: -// For each new token in the decode step: -// 1. Load existing compressor state (partial kv/score) from paged cache -// 2. Perform online softmax: accumulate new token's contribution using -// the numerically stable running max + weighted sum formulation -// 3. Write updated state back to paged cache -// 4. If compress_ratio tokens accumulated → emit compressed output, -// reset state for next compression window -// -// Each thread handles VEC contiguous elements of head_dim. In overlap mode -// (state_dim = 2*head_dim), Phase 1 iterates over 2 column halves. -// -// Memory layout: -// kv_score: [total_tokens, 2 * state_dim] — interleaved KV and score projections -// paged_kv: paged cache for compressor KV state -// paged_score: paged cache for compressor score state (with APE bias) -// output: [total_comp_tokens, head_dim] — compressed output tokens -// ============================================================================ - -// Helper: vectorized online softmax step reading from paged KV/score state. -// Loads one position's KV and score from paged memory and updates the running -// online softmax accumulators (rmax, rsum, rwsum) per element. -// APE is already baked into paged_score (added during Phase 1), so no APE -// addition is performed here. -template -__device__ __forceinline__ void decodeSoftmaxVec(void const* __restrict__ paged_kv_raw, - void const* __restrict__ paged_score_raw, - int64_t page_sd, // page_size * state_dim (in elements) - int state_dim, - int phys_kv, // physical page index for kv - int phys_sc, // physical page index for score - int blk_off, // offset within page - int kv_col_off, // column offset (0 or HEAD_DIM) - int tid, float* __restrict__ rmax, float* __restrict__ rsum, float* __restrict__ rwsum) -{ - using StateElemT = typename std::conditional::type; - using StateVecT = typename VecType::type; - - auto const* kv = reinterpret_cast(paged_kv_raw); - auto const* sc = reinterpret_cast(paged_score_raw); - - int64_t base_kv = static_cast(phys_kv) * page_sd + blk_off * state_dim + kv_col_off; - int64_t base_sc = static_cast(phys_sc) * page_sd + blk_off * state_dim + kv_col_off; - - StateVecT k_raw = reinterpret_cast(&kv[base_kv])[tid]; - StateVecT s_raw = reinterpret_cast(&sc[base_sc])[tid]; - StateElemT const* ke = reinterpret_cast(&k_raw); - StateElemT const* se = reinterpret_cast(&s_raw); - -#pragma unroll - for (int i = 0; i < VEC; i += 4) - { - float kf[4] = {static_cast(ke[i]), static_cast(ke[i + 1]), static_cast(ke[i + 2]), - static_cast(ke[i + 3])}; - // score already includes APE (added during Phase 1 store) - float sf[4] = {static_cast(se[i]), static_cast(se[i + 1]), static_cast(se[i + 2]), - static_cast(se[i + 3])}; - // Online softmax: maintain running (max, sum_exp, weighted_sum) per element. - // nm = new max, sc_f = rescale factor for old accumulators, tm = exp(score - new_max). - // Final output: rwsum / rsum = weighted average of KV values. -#pragma unroll - for (int j = 0; j < 4; j++) - { - float nm = fmaxf(rmax[i + j], sf[j]); - float sc_f = expf(rmax[i + j] - nm); - float tm = expf(sf[j] - nm); - rsum[i + j] = rsum[i + j] * sc_f + tm; - rwsum[i + j] = rwsum[i + j] * sc_f + kf[j] * tm; - rmax[i + j] = nm; - } - } -} - -template -__global__ void pagedKvCompressKernel(void const* __restrict__ kv_score_raw, float const* __restrict__ ape, - void* __restrict__ paged_kv_raw, void* __restrict__ paged_score_raw, int32_t const* __restrict__ block_table_kv, - int32_t const* __restrict__ block_table_score, void* __restrict__ output_raw, int32_t const* __restrict__ kv_lens, - int32_t const* __restrict__ cu_seq_lens, int32_t const* __restrict__ cu_kv_comp, int page_size, int max_blocks, - int out_elem_bytes) -{ - using KvScoreElemT = typename std::conditional::type; - using StateElemT = typename std::conditional::type; - // DeepSeek-V4 model configures compress_ratio to be 4 or 128. - static_assert(COMPRESS_RATIO == 4 || COMPRESS_RATIO == 128, "Unsupported COMPRESS_RATIO"); - constexpr bool IS_OVERLAP = (COMPRESS_RATIO == 4); - constexpr int ELEM_BYTES_FOR_VEC - = (KV_SCORE_ELEM_BYTES > STATE_ELEM_BYTES) ? KV_SCORE_ELEM_BYTES : STATE_ELEM_BYTES; - constexpr int MAX_VEC = 16 / ELEM_BYTES_FOR_VEC; - constexpr int VEC = (HEAD_DIM / MAX_VEC >= 32) ? MAX_VEC : (HEAD_DIM / 32); - using KvScoreVecT = typename VecType::type; - using StateVecT = typename VecType::type; - static_assert(VEC >= 4, "VEC must be >= 4 for float4 ape loads"); - - // HEAD_BLOCKS: split head_dim across blockIdx.y for better SM utilisation. - // For HD=512 and max elem size 2: NTHRD_BASE=64 → HEAD_BLOCKS=2, NTHRD_INNER=32. - // For HD=128 and max elem size 2/4: NTHRD_BASE=32 → HEAD_BLOCKS=1, NTHRD_INNER=32. - // For HD=512 and max elem size 4: NTHRD_BASE=128 → HEAD_BLOCKS=4, NTHRD_INNER=32. - constexpr int NTHRD_BASE = HEAD_DIM / VEC; - constexpr int HEAD_BLOCKS = (NTHRD_BASE > 32) ? (NTHRD_BASE / 32) : 1; - constexpr int NTHRD_INNER = NTHRD_BASE / HEAD_BLOCKS; // always <= 32 - // ELEM_PER_BLOCK: head_dim elements handled by one blockIdx.y block. - // = NTHRD_INNER * VEC = HEAD_DIM / HEAD_BLOCKS. - // Used for the multi-warp shared-memory merge layout so that each block - // only allocates storage for its own head slice, not the full HEAD_DIM. - constexpr int ELEM_PER_BLOCK = NTHRD_INNER * VEC; - - // state_dim is fully determined by template parameters. - constexpr int STATE_DIM = IS_OVERLAP ? 2 * HEAD_DIM : HEAD_DIM; - constexpr int64_t TWO_SD = 2 * STATE_DIM; - constexpr int COFF = IS_OVERLAP ? 2 : 1; - - int const tid = threadIdx.x % NTHRD_INNER; - int const warp_id = threadIdx.x / NTHRD_INNER; // 0..NUM_RED_WARPS-1 - int const batch_idx = blockIdx.x; - int const head_blk = blockIdx.y; - int const eff_tid = head_blk * NTHRD_INNER + tid; - - int const kv_len = kv_lens[batch_idx]; - int const sp = kv_len - NEXT_N; - int const in_off = cu_seq_lens[batch_idx]; - int const out_off = cu_kv_comp[batch_idx]; - int64_t const page_sd = static_cast(page_size) * STATE_DIM; - - auto const* kv_score = reinterpret_cast(kv_score_raw); - auto* paged_kv = reinterpret_cast(paged_kv_raw); - auto* paged_score = reinterpret_cast(paged_score_raw); - - // ================================================================ - // Phase 1: Write NEXT_N new tokens' KV and score state to paged cache. - // - // Only warp 0 participates (all warps share the same eff_tid mapping). - // When NUM_RED_WARPS == 1, the guard compiles away. - // ================================================================ - if (warp_id == 0) - { -#pragma unroll - for (int t = 0; t < NEXT_N; t++) - { - int token_idx = sp + t; - if (token_idx < kv_len) - { - int ape_idx = token_idx % COMPRESS_RATIO; - int log_blk = token_idx / page_size; - int blk_off = token_idx % page_size; - int phys_kv = block_table_kv[batch_idx * max_blocks + log_blk]; - int phys_sc = block_table_score[batch_idx * max_blocks + log_blk]; - - for (int col_idx = 0; col_idx < COFF; col_idx++) - { - int const col = col_idx * HEAD_DIM; - int64_t const src = static_cast(in_off + t) * TWO_SD + col; - int64_t const dkv = static_cast(phys_kv) * page_sd + blk_off * STATE_DIM + col; - int64_t const dsc = static_cast(phys_sc) * page_sd + blk_off * STATE_DIM + col; - - KvScoreVecT kv_raw = reinterpret_cast(&kv_score[src])[eff_tid]; - KvScoreVecT sc_raw = reinterpret_cast(&kv_score[src + STATE_DIM])[eff_tid]; - - KvScoreElemT const* kv_e = reinterpret_cast(&kv_raw); - StateVecT kv_out; - StateElemT* kv_o = reinterpret_cast(&kv_out); -#pragma unroll - for (int i = 0; i < VEC; i++) - { - kv_o[i] = static_cast(static_cast(kv_e[i])); - } - reinterpret_cast(&paged_kv[dkv])[eff_tid] = kv_out; - - KvScoreElemT const* sc_e = reinterpret_cast(&sc_raw); - StateVecT sc_out; - StateElemT* sc_o = reinterpret_cast(&sc_out); -#pragma unroll - for (int i = 0; i < VEC; i += 4) - { - float4 av - = *reinterpret_cast(&ape[ape_idx * STATE_DIM + col + eff_tid * VEC + i]); - sc_o[i] = static_cast(static_cast(sc_e[i]) + av.x); - sc_o[i + 1] = static_cast(static_cast(sc_e[i + 1]) + av.y); - sc_o[i + 2] = static_cast(static_cast(sc_e[i + 2]) + av.z); - sc_o[i + 3] = static_cast(static_cast(sc_e[i + 3]) + av.w); - } - reinterpret_cast(&paged_score[dsc])[eff_tid] = sc_out; - } - } - } - } - - if constexpr (NUM_RED_WARPS > 1) - { - __syncthreads(); - } - - // ================================================================ - // Phase 2: Count how many complete compression windows finished. - // ================================================================ - int last_token_idx = sp + NEXT_N - 1; - int num_compressions = (last_token_idx + 1) / COMPRESS_RATIO - sp / COMPRESS_RATIO; - - // ================================================================ - // Phase 3: Online softmax reduction over each complete chunk. - // - // When NUM_RED_WARPS > 1, the compress_ratio positions are split - // across warps. Each warp reduces its partition independently, then - // partial (rmax, rsum, rwsum) accumulators are merged via shared - // memory using the log-sum-exp identity. - // ================================================================ - for (int c = 0; c < NEXT_N; c++) - { - if (c >= num_compressions) - break; - - int compress_idx = sp / COMPRESS_RATIO + c; - int curr_chunk_start = compress_idx * COMPRESS_RATIO; - - float rmax[VEC], rsum[VEC], rwsum[VEC]; -#pragma unroll - for (int i = 0; i < VEC; i++) - { - rmax[i] = -INFINITY; - rsum[i] = 0.0f; - rwsum[i] = 0.0f; - } - - constexpr int positions_per_warp = COMPRESS_RATIO / NUM_RED_WARPS; - int const my_r_start = warp_id * positions_per_warp; - int const my_r_end = (warp_id == NUM_RED_WARPS - 1) ? COMPRESS_RATIO : (my_r_start + positions_per_warp); - - if constexpr (IS_OVERLAP) - { - int prev_start = curr_chunk_start - COMPRESS_RATIO; - if (prev_start >= 0) - { - if (page_size >= COMPRESS_RATIO) - { - int log_blk_prev = prev_start / page_size; - int phys_kv_prev = block_table_kv[batch_idx * max_blocks + log_blk_prev]; - int phys_sc_prev = block_table_score[batch_idx * max_blocks + log_blk_prev]; - int chunk_off_prev = prev_start % page_size; - for (int r = my_r_start; r < my_r_end; r++) - { - decodeSoftmaxVec(paged_kv_raw, paged_score_raw, page_sd, - STATE_DIM, phys_kv_prev, phys_sc_prev, chunk_off_prev + r, 0, eff_tid, rmax, rsum, rwsum); - } - } - else - { - for (int r = my_r_start; r < my_r_end; r++) - { - int pos = prev_start + r; - int log_blk = pos / page_size; - int blk_off = pos % page_size; - int phys_kv = block_table_kv[batch_idx * max_blocks + log_blk]; - int phys_sc = block_table_score[batch_idx * max_blocks + log_blk]; - decodeSoftmaxVec(paged_kv_raw, paged_score_raw, page_sd, - STATE_DIM, phys_kv, phys_sc, blk_off, 0, eff_tid, rmax, rsum, rwsum); - } - } - } - - if (page_size >= COMPRESS_RATIO) - { - int log_blk_cur = curr_chunk_start / page_size; - int phys_kv_cur = block_table_kv[batch_idx * max_blocks + log_blk_cur]; - int phys_sc_cur = block_table_score[batch_idx * max_blocks + log_blk_cur]; - int chunk_off_cur = curr_chunk_start % page_size; - for (int r = my_r_start; r < my_r_end; r++) - { - decodeSoftmaxVec(paged_kv_raw, paged_score_raw, page_sd, STATE_DIM, - phys_kv_cur, phys_sc_cur, chunk_off_cur + r, HEAD_DIM, eff_tid, rmax, rsum, rwsum); - } - } - else - { - for (int r = my_r_start; r < my_r_end; r++) - { - int pos = curr_chunk_start + r; - int log_blk = pos / page_size; - int blk_off = pos % page_size; - int phys_kv = block_table_kv[batch_idx * max_blocks + log_blk]; - int phys_sc = block_table_score[batch_idx * max_blocks + log_blk]; - decodeSoftmaxVec(paged_kv_raw, paged_score_raw, page_sd, STATE_DIM, - phys_kv, phys_sc, blk_off, HEAD_DIM, eff_tid, rmax, rsum, rwsum); - } - } - } - else - { - if (page_size >= COMPRESS_RATIO) - { - int log_blk = curr_chunk_start / page_size; - int phys_kv = block_table_kv[batch_idx * max_blocks + log_blk]; - int phys_sc = block_table_score[batch_idx * max_blocks + log_blk]; - int chunk_off = curr_chunk_start % page_size; - for (int r = my_r_start; r < my_r_end; r++) - { - decodeSoftmaxVec(paged_kv_raw, paged_score_raw, page_sd, STATE_DIM, - phys_kv, phys_sc, chunk_off + r, 0, eff_tid, rmax, rsum, rwsum); - } - } - else - { - for (int r = my_r_start; r < my_r_end; r++) - { - int pos = curr_chunk_start + r; - int log_blk = pos / page_size; - int blk_off = pos % page_size; - int phys_kv = block_table_kv[batch_idx * max_blocks + log_blk]; - int phys_sc = block_table_score[batch_idx * max_blocks + log_blk]; - decodeSoftmaxVec(paged_kv_raw, paged_score_raw, page_sd, STATE_DIM, - phys_kv, phys_sc, blk_off, 0, eff_tid, rmax, rsum, rwsum); - } - } - } - - // Multi-warp merge epilogue (compiled away when NUM_RED_WARPS == 1). - if constexpr (NUM_RED_WARPS > 1) - { - // Shared-memory layout: [NUM_RED_WARPS * ELEM_PER_BLOCK] per array. - // ELEM_PER_BLOCK = NTHRD_INNER * VEC = HEAD_DIM / HEAD_BLOCKS, i.e. - // the number of head_dim elements covered by this block (blockIdx.y). - // Using ELEM_PER_BLOCK (not HEAD_DIM) avoids 4x over-allocation when - // HEAD_BLOCKS > 1 (e.g. HD=512 fp32 has HEAD_BLOCKS=4). - extern __shared__ float smem[]; - float* s_rmax = smem; - float* s_rsum = s_rmax + NUM_RED_WARPS * ELEM_PER_BLOCK; - float* s_rwsum = s_rsum + NUM_RED_WARPS * ELEM_PER_BLOCK; - - // local_elem: index within this block's head slice [0, ELEM_PER_BLOCK). - // = tid * VEC + i (tid = threadIdx.x % NTHRD_INNER, same as eff_tid - head_blk*NTHRD_INNER) -#pragma unroll - for (int i = 0; i < VEC; i++) - { - int const local_elem = tid * VEC + i; - s_rmax[warp_id * ELEM_PER_BLOCK + local_elem] = rmax[i]; - s_rsum[warp_id * ELEM_PER_BLOCK + local_elem] = rsum[i]; - s_rwsum[warp_id * ELEM_PER_BLOCK + local_elem] = rwsum[i]; - } - __syncthreads(); - - if (warp_id == 0) - { - for (int w = 1; w < NUM_RED_WARPS; w++) - { -#pragma unroll - for (int i = 0; i < VEC; i++) - { - int const local_elem = tid * VEC + i; - float const m2 = s_rmax[w * ELEM_PER_BLOCK + local_elem]; - float const s2 = s_rsum[w * ELEM_PER_BLOCK + local_elem]; - float const ws2 = s_rwsum[w * ELEM_PER_BLOCK + local_elem]; - - float const nm = fmaxf(rmax[i], m2); - float const sc1 = expf(rmax[i] - nm); - float const sc2 = expf(m2 - nm); - rsum[i] = rsum[i] * sc1 + s2 * sc2; - rwsum[i] = rwsum[i] * sc1 + ws2 * sc2; - rmax[i] = nm; - } - } - } - __syncthreads(); - } - - bool const should_write = (NUM_RED_WARPS == 1) || (warp_id == 0); - if (should_write) - { - int64_t const out_base = static_cast(out_off + c) * HEAD_DIM + eff_tid * VEC; - if (out_elem_bytes == 2) - { - __nv_bfloat16 packed[VEC]; -#pragma unroll - for (int i = 0; i < VEC; i++) - packed[i] = __float2bfloat16_rn(rwsum[i] / rsum[i]); - using OutVecT = typename VecType::type; - *reinterpret_cast(&reinterpret_cast<__nv_bfloat16*>(output_raw)[out_base]) - = *reinterpret_cast(packed); - } - else - { - float result[VEC]; -#pragma unroll - for (int i = 0; i < VEC; i++) - result[i] = rwsum[i] / rsum[i]; -#pragma unroll - for (int i = 0; i < VEC; i += 4) - *reinterpret_cast(&reinterpret_cast(output_raw)[out_base + i]) - = *reinterpret_cast(&result[i]); - } - } - } -} - -// ============================================================================ -// Decode kernel configuration matrix — single source of truth. -// -// X-macro listing every supported (HD, KV_EB, STATE_EB, CR, NN, NRW) tuple. -// Both the explicit template instantiations below AND the runtime dispatcher -// in pagedKvCompressLaunch() walk this list, so adding/removing a config is -// a one-line edit. -// -// HD — HEAD_DIM in {128, 512} -// KV_EB — kv_score element bytes in {2 (bf16), 4 (fp32)} -// STATE_EB — paged state element bytes in {2 (bf16), 4 (fp32)} -// CR — COMPRESS_RATIO in {4, 128} -// NN — NEXT_N (new tokens / decode step) in {1..8} -// NRW — NUM_RED_WARPS — 4 when CR=128 and NN<=4 (multi-warp Phase 3 -// reduction hides DRAM latency for the heavier R=128 chunk); -// 1 when CR=4 or when CR=128 and NN>=5. -// -// Multi-warp SMEM budget (per block): 3 * NRW * ELEM_PER_BLOCK * sizeof(float). -// HD=128: ELEM_PER_BLOCK=128 → 6 KB -// HD=512 bf16: ELEM_PER_BLOCK=256 → 12 KB -// HD=512 fp32: ELEM_PER_BLOCK=128 → 6 KB -// ============================================================================ - -// Per-axis fan-outs (used to keep the master list compact). -#define FOREACH_DECODE_NN_1_4(F, HD, KV, ST, CR, NRW) \ - F(HD, KV, ST, CR, 1, NRW) F(HD, KV, ST, CR, 2, NRW) F(HD, KV, ST, CR, 3, NRW) F(HD, KV, ST, CR, 4, NRW) -#define FOREACH_DECODE_NN_5_8(F, HD, KV, ST, CR, NRW) \ - F(HD, KV, ST, CR, 5, NRW) F(HD, KV, ST, CR, 6, NRW) F(HD, KV, ST, CR, 7, NRW) F(HD, KV, ST, CR, 8, NRW) -#define FOREACH_DECODE_DTYPE_1_4(F, HD, CR, NRW) \ - FOREACH_DECODE_NN_1_4(F, HD, 2, 2, CR, NRW) \ - FOREACH_DECODE_NN_1_4(F, HD, 2, 4, CR, NRW) \ - FOREACH_DECODE_NN_1_4(F, HD, 4, 2, CR, NRW) FOREACH_DECODE_NN_1_4(F, HD, 4, 4, CR, NRW) -#define FOREACH_DECODE_DTYPE_5_8(F, HD, CR, NRW) \ - FOREACH_DECODE_NN_5_8(F, HD, 2, 2, CR, NRW) \ - FOREACH_DECODE_NN_5_8(F, HD, 2, 4, CR, NRW) \ - FOREACH_DECODE_NN_5_8(F, HD, 4, 2, CR, NRW) FOREACH_DECODE_NN_5_8(F, HD, 4, 4, CR, NRW) -#define FOREACH_DECODE_DTYPE_1_8(F, HD, CR, NRW) \ - FOREACH_DECODE_DTYPE_1_4(F, HD, CR, NRW) FOREACH_DECODE_DTYPE_5_8(F, HD, CR, NRW) - -// Master list. Order does not matter; the dispatcher walks linearly. -// clang-format off -#define FOREACH_DECODE_CONFIG(F) \ - /* CR=4: single-warp for next_n 1..8 (small reduction; multi-warp would over-subscribe). */ \ - FOREACH_DECODE_DTYPE_1_8(F, 128, 4, 1) FOREACH_DECODE_DTYPE_1_8(F, 512, 4, 1) \ - /* CR=128: single-warp fallback for next_n 5..8. */ \ - FOREACH_DECODE_DTYPE_5_8(F, 128, 128, 1) FOREACH_DECODE_DTYPE_5_8(F, 512, 128, 1) \ - /* CR=128: multi-warp fast path for next_n 1..4. */ \ - FOREACH_DECODE_DTYPE_1_4(F, 128, 128, 4) FOREACH_DECODE_DTYPE_1_4(F, 512, 128, 4) -// clang-format on - -// Generate explicit template instantiations. -#define INST_DECODE(HD, KV_EB, STATE_EB, CR, NN, NRW) \ - template __global__ void pagedKvCompressKernel(void const*, float const*, void*, \ - void*, int32_t const*, int32_t const*, void*, int32_t const*, int32_t const*, int32_t const*, int, int, int); -FOREACH_DECODE_CONFIG(INST_DECODE) -#undef INST_DECODE - -// ============================================================================ -// Decode Launch Wrapper -// -// Dispatches to the correct template instantiation based on head_dim, elem_bytes, -// and next_n (number of new tokens per decode step, in the range 1..8). -// Grid is 2D: (batch_size, head_blocks) where head_blocks = NTHRD_BASE / 32. -// For HD=512 bf16: head_blocks=2; for HD=128 bf16: head_blocks=1. -// ============================================================================ - -void pagedKvCompressLaunch(void const* kv_score, float const* ape, void* paged_kv, void* paged_score, - int32_t const* block_table_kv, int32_t const* block_table_score, void* output, int32_t const* kv_lens, - int32_t const* cu_seq_lens, int32_t const* cu_kv_comp, int batch_size, int page_size, int max_blocks, int head_dim, - int compress_ratio, int next_n, int kv_score_elem_bytes, int state_elem_bytes, int out_elem_bytes, - cudaStream_t stream) -{ - TLLM_CHECK_WITH_INFO( - compress_ratio == 4 || compress_ratio == 128, "pagedKvCompressLaunch only supports compress_ratio 4 or 128"); - TLLM_CHECK_WITH_INFO( - (kv_score_elem_bytes == 2 || kv_score_elem_bytes == 4) && (state_elem_bytes == 2 || state_elem_bytes == 4), - "pagedKvCompressLaunch only supports bf16/fp32 kv_score and paged state"); - constexpr int kMinNextN = 1; - constexpr int kMaxNextN = 8; - TLLM_CHECK_WITH_INFO(next_n >= kMinNextN && next_n <= kMaxNextN, - "pagedKvCompressLaunch only supports next_n in [1, 8], got %d", next_n); - - // Compute HEAD_BLOCKS: mirrors the compile-time constant in the kernel. - // VEC = max_vec if HEAD_DIM/max_vec >= 32, else HEAD_DIM/32. - // NTHRD_BASE = HEAD_DIM / VEC; HEAD_BLOCKS = NTHRD_BASE / 32 (or 1 if <= 32). - // This spreads the head_dim across multiple blocks for better SM utilisation. - int const elem_bytes_for_vec = max(kv_score_elem_bytes, state_elem_bytes); - int const max_vec_elem = 16 / elem_bytes_for_vec; - int const vec = (head_dim / max_vec_elem >= 32) ? max_vec_elem : (head_dim / 32); - int const nthrd_base = head_dim / vec; // mirrors kernel's NTHRD_BASE = HEAD_DIM / VEC - int const head_blocks = (nthrd_base > 32) ? (nthrd_base / 32) : 1; - int const nthreads_inner = nthrd_base / head_blocks; // = min(32, nthrd_base) = always 32 - - // For large compress_ratio, use 4-warp parallel reduction to cut the serial - // softmax loop from COMPRESS_RATIO iterations to COMPRESS_RATIO/4 per warp. - // The multi-warp path supports CR=128, (HD=128 or HD=512), and NEXT_N in - // 1..4. Larger NEXT_N values use a single reduction warp to limit block - // size while still processing every new token exactly. - // - // smem per block = 3 * MULTI_WARP * ELEM_PER_BLOCK * sizeof(float) - // where ELEM_PER_BLOCK = nthreads_inner * vec = HEAD_DIM / HEAD_BLOCKS. - // HD=128: ELEM_PER_BLOCK=128 → 6 KB. - // HD=512 with max elem size 2 (vec=8, HEAD_BLOCKS=2): ELEM_PER_BLOCK=256 → 12 KB. - // HD=512 with max elem size 4 (vec=4, HEAD_BLOCKS=4): ELEM_PER_BLOCK=128 → 6 KB. - constexpr int MULTI_WARP = 4; - bool const use_multi_warp = (compress_ratio == 128 && next_n <= 4); - int const num_red_warps = use_multi_warp ? MULTI_WARP : 1; - int const nthreads = nthreads_inner * num_red_warps; - int const elem_per_block = nthreads_inner * vec; // = HEAD_DIM / HEAD_BLOCKS - int const smem_bytes = use_multi_warp ? (3 * MULTI_WARP * elem_per_block * static_cast(sizeof(float))) : 0; - - dim3 grid(batch_size, head_blocks); - - // Walk FOREACH_DECODE_CONFIG until we find a matching (HD, KV, ST, CR, NN, NRW) - // tuple, then launch that instantiation. Any unsupported tuple bails via TLLM_THROW. -#define TRY_LAUNCH(HD, KV_EB, STATE_EB, CR, NN, NRW) \ - if (head_dim == HD && kv_score_elem_bytes == KV_EB && state_elem_bytes == STATE_EB && compress_ratio == CR \ - && next_n == NN && num_red_warps == NRW) \ - { \ - pagedKvCompressKernel<<>>(kv_score, ape, \ - paged_kv, paged_score, block_table_kv, block_table_score, output, kv_lens, cu_seq_lens, cu_kv_comp, \ - page_size, max_blocks, out_elem_bytes); \ - return; \ - } - FOREACH_DECODE_CONFIG(TRY_LAUNCH) -#undef TRY_LAUNCH - - TLLM_THROW( - "pagedKvCompressLaunch: no matching instantiation for HD=%d, kv_eb=%d, state_eb=%d, CR=%d, NN=%d, NRW=%d", - head_dim, kv_score_elem_bytes, state_elem_bytes, compress_ratio, next_n, num_red_warps); -} - -#undef FOREACH_DECODE_CONFIG -#undef FOREACH_DECODE_DTYPE_1_8 -#undef FOREACH_DECODE_DTYPE_5_8 -#undef FOREACH_DECODE_DTYPE_1_4 -#undef FOREACH_DECODE_NN_5_8 -#undef FOREACH_DECODE_NN_1_4 - -// ============================================================================ -// Prefill Kernel: prefillReductionKernel -// -// Template: -// -// Grid: (batch_size, max_outputs_per_batch, head_blocks) -// Block: (NTHRD_INNER * NUM_RED_WARPS), where NTHRD_INNER covers one head chunk. -// -// Unlike the decode kernel (which operates token-by-token from paged state), -// the prefill kernel processes the full input sequence at once. Each block -// handles one compressed output for one head_dim chunk. For compress_ratio=128, -// four reduction groups split the token dimension for state writes and online -// softmax, then merge per-element partials in shared memory. -// -// The last block (local_output_idx == num_outputs - 1) also handles saving -// compressor state for any remainder tokens that don't form a full chunk. -// All full chunks are also written to paged kv/score caches for block reuse. -// -// Memory layout: -// kv_score: [total_tokens, 2*state_dim] — interleaved KV and score from linear projection -// paged_kv: paged cache for compressor state (remainder) -// paged_score: paged cache for compressor score state (remainder, with APE) -// output: [total_comp_tokens, head_dim] — compressed output tokens -// ============================================================================ - -// Per-element online softmax step on VEC elements via 128-bit vectorized loads. -// Reads directly from the kv_score input buffer (not paged state) since prefill -// has the full sequence available. -template -__device__ __forceinline__ void prefillSoftmaxVec(void const* __restrict__ kv_score_raw, float const* __restrict__ ape, - int64_t row_elem, // (input_offset + row_idx) * two_sd - int kv_col_off, // column offset into kv_score row (0 or HEAD_DIM) - int ape_base, // r * state_dim + ape_col_off - int state_dim, int tid, float* __restrict__ rmax, float* __restrict__ rsum, float* __restrict__ rwsum) -{ - using KvScoreElemT = typename std::conditional::type; - using KvScoreVecT = typename VecType::type; - - auto const* kv = reinterpret_cast(kv_score_raw); - - KvScoreVecT k_raw = reinterpret_cast(&kv[row_elem + kv_col_off])[tid]; - KvScoreVecT s_raw = reinterpret_cast(&kv[row_elem + state_dim + kv_col_off])[tid]; - KvScoreElemT const* ke = reinterpret_cast(&k_raw); - KvScoreElemT const* se = reinterpret_cast(&s_raw); - -#pragma unroll - for (int i = 0; i < VEC; i += 4) - { - float4 av = *reinterpret_cast(&ape[ape_base + tid * VEC + i]); - float kf[4] = {static_cast(ke[i]), static_cast(ke[i + 1]), static_cast(ke[i + 2]), - static_cast(ke[i + 3])}; - float sf[4] = {static_cast(se[i]) + av.x, static_cast(se[i + 1]) + av.y, - static_cast(se[i + 2]) + av.z, static_cast(se[i + 3]) + av.w}; -#pragma unroll - for (int j = 0; j < 4; j++) - { - float nm = fmaxf(rmax[i + j], sf[j]); - float sc = expf(rmax[i + j] - nm); - float tm = expf(sf[j] - nm); - rsum[i + j] = rsum[i + j] * sc + tm; - rwsum[i + j] = rwsum[i + j] * sc + kf[j] * tm; - rmax[i + j] = nm; - } - } -} - -template -__global__ void prefillReductionKernel(void const* __restrict__ kv_score_raw, float const* __restrict__ ape, - void* __restrict__ paged_kv_raw, void* __restrict__ paged_score_raw, int32_t const* __restrict__ block_table_kv, - int32_t const* __restrict__ block_table_score, void* __restrict__ output_raw, int32_t const* __restrict__ kv_lens, - int32_t const* __restrict__ start_pos_arr, int32_t const* __restrict__ cu_seq_lens, - int32_t const* __restrict__ cu_kv_comp, int page_size, int state_dim, int max_blocks, int out_elem_bytes) -{ - using KvScoreElemT = typename std::conditional::type; - using StateElemT = typename std::conditional::type; - static_assert(COMPRESS_RATIO == 4 || COMPRESS_RATIO == 128, "Unsupported COMPRESS_RATIO"); - constexpr bool IS_OVERLAP = (COMPRESS_RATIO == 4); - - constexpr int ELEM_BYTES_FOR_VEC - = (KV_SCORE_ELEM_BYTES > STATE_ELEM_BYTES) ? KV_SCORE_ELEM_BYTES : STATE_ELEM_BYTES; - constexpr int MAX_VEC = 16 / ELEM_BYTES_FOR_VEC; - constexpr int VEC = (HEAD_DIM / MAX_VEC >= 32) ? MAX_VEC : (HEAD_DIM / 32); - using KvScoreVecT = typename VecType::type; - using StateVecT = typename VecType::type; - static_assert(VEC >= 4, "VEC must be >= 4 for float4 ape loads"); - static_assert(NUM_RED_WARPS == 1 || NUM_RED_WARPS == 4, "Unsupported NUM_RED_WARPS"); - - constexpr int NTHRD_BASE = HEAD_DIM / VEC; - constexpr int HEAD_BLOCKS = (COMPRESS_RATIO == 128 && NTHRD_BASE > 32) ? (NTHRD_BASE / 32) : 1; - constexpr int NTHRD_INNER = NTHRD_BASE / HEAD_BLOCKS; - constexpr int ELEM_PER_BLOCK = NTHRD_INNER * VEC; - - int const tid = threadIdx.x % NTHRD_INNER; - int const red_warp = (NUM_RED_WARPS == 1) ? 0 : (threadIdx.x / NTHRD_INNER); - int const batch_idx = blockIdx.x; - int const local_output_idx = blockIdx.y; - int const head_blk = (HEAD_BLOCKS == 1) ? 0 : blockIdx.z; - int const eff_tid = (HEAD_BLOCKS == 1) ? tid : (head_blk * NTHRD_INNER + tid); - - int const sp = start_pos_arr[batch_idx]; - int const kv_len = kv_lens[batch_idx]; - int const input_offset = cu_seq_lens[batch_idx]; - int const output_offset = cu_kv_comp[batch_idx]; - - // Absolute compression index range. Window k covers [k*R, (k+1)*R). - // We emit one output per complete window that gains at least one new token. - int const first_abs_idx = sp / COMPRESS_RATIO; - int const last_abs_idx = kv_len / COMPRESS_RATIO; // exclusive - int const actual_num_outputs = last_abs_idx - first_abs_idx; - // Keep at least one block so the last CTA can still persist remainder state - // even when this chunk has no full compression window. - int const num_outputs = max(actual_num_outputs, 1); - - if (local_output_idx >= num_outputs) - return; - - constexpr int coff = IS_OVERLAP ? 2 : 1; - bool const should_compress = (local_output_idx < actual_num_outputs); - - // Absolute window for this output block. - int const abs_idx = first_abs_idx + local_output_idx; - int const win_start = abs_idx * COMPRESS_RATIO; - - auto const* kv_score = reinterpret_cast(kv_score_raw); - auto* paged_kv = reinterpret_cast(paged_kv_raw); - auto* paged_score = reinterpret_cast(paged_score_raw); - - int64_t const two_sd = 2 * state_dim; - int64_t const page_sd = static_cast(page_size) * state_dim; - - float rmax[VEC], rsum[VEC], rwsum[VEC]; - if constexpr (!IS_OVERLAP) - { -#pragma unroll - for (int i = 0; i < VEC; i++) - { - rmax[i] = -INFINITY; - rsum[i] = 0.0f; - rwsum[i] = 0.0f; - } - } - - // ================================================================ - // Phase 1: State Update (all output blocks, for block reuse support) - // - // 1a runs on every block that has a full chunk (should_compress), writing - // each block's own window to paged cache so any prefix slice is valid for - // block reuse. 1b (remainder) still runs only on the last block. - // ================================================================ - - // Helper: write a contiguous block of positions to paged KV/score cache. - // Only new positions (>= sp) are written; positions already persisted from - // prior calls are skipped by starting the loop at write_r_start. - // APE index is r (position within the compression window), matching the - // index used during the original decode/prefill that first established - // the window alignment. - auto write_to_paged = [&](int range_start, int range_end, int write_r_start, bool accumulateNewTokens) - { - int const write_count = range_end - range_start - write_r_start; - if (write_count <= 0) - { - return; - } - - int const r_begin = write_r_start + write_count * red_warp / NUM_RED_WARPS; - int const r_end = write_r_start + write_count * (red_warp + 1) / NUM_RED_WARPS; - for (int r = r_begin; r < r_end; r++) - { - int const pos = range_start + r; - int const input_row = pos - sp; // >= 0 since pos >= sp (loop starts at write_r_start) - int const log_blk = pos / page_size; - int const blk_off = pos % page_size; - int const phys_kv = block_table_kv[batch_idx * max_blocks + log_blk]; - int const phys_sc = block_table_score[batch_idx * max_blocks + log_blk]; - - for (int col_idx = 0; col_idx < coff; col_idx++) - { - int const col = col_idx * HEAD_DIM; - int64_t const src = static_cast(input_offset + input_row) * two_sd + col; - int64_t const dkv = static_cast(phys_kv) * page_sd + blk_off * state_dim + col; - int64_t const dsc = static_cast(phys_sc) * page_sd + blk_off * state_dim + col; - - KvScoreVecT kv_raw = reinterpret_cast(&kv_score[src])[eff_tid]; - KvScoreVecT sc_raw = reinterpret_cast(&kv_score[src + state_dim])[eff_tid]; - - KvScoreElemT const* kv_e = reinterpret_cast(&kv_raw); - StateVecT kv_out; - StateElemT* kv_o = reinterpret_cast(&kv_out); - float kv_f[VEC]; -#pragma unroll - for (int i = 0; i < VEC; i++) - { - kv_f[i] = static_cast(kv_e[i]); - kv_o[i] = static_cast(kv_f[i]); - } - reinterpret_cast(&paged_kv[dkv])[eff_tid] = kv_out; - - KvScoreElemT const* sc_e = reinterpret_cast(&sc_raw); - StateVecT sc_out; - StateElemT* sc_o = reinterpret_cast(&sc_out); -#pragma unroll - for (int i = 0; i < VEC; i += 4) - { - float4 av = *reinterpret_cast(&ape[r * state_dim + col + eff_tid * VEC + i]); - float const sf0 = static_cast(sc_e[i]) + av.x; - float const sf1 = static_cast(sc_e[i + 1]) + av.y; - float const sf2 = static_cast(sc_e[i + 2]) + av.z; - float const sf3 = static_cast(sc_e[i + 3]) + av.w; - sc_o[i] = static_cast(sf0); - sc_o[i + 1] = static_cast(sf1); - sc_o[i + 2] = static_cast(sf2); - sc_o[i + 3] = static_cast(sf3); - - if constexpr (!IS_OVERLAP) - { - if (accumulateNewTokens) - { - float const sf[4] = {sf0, sf1, sf2, sf3}; -#pragma unroll - for (int j = 0; j < 4; j++) - { - float const nm = fmaxf(rmax[i + j], sf[j]); - float const sc = expf(rmax[i + j] - nm); - float const tm = expf(sf[j] - nm); - rsum[i + j] = rsum[i + j] * sc + tm; - rwsum[i + j] = rwsum[i + j] * sc + kv_f[i + j] * tm; - rmax[i + j] = nm; - } - } - } - } - reinterpret_cast(&paged_score[dsc])[eff_tid] = sc_out; - } - } - }; - - // 1a. Full chunk for this output block. - // Positions [win_start, sp) are already in paged cache from a prior call; - // start the loop at write_r_start to skip them without a per-iteration branch. - if (should_compress) - { - int const write_r_start = (win_start < sp) ? (sp - win_start) : 0; - write_to_paged(win_start, win_start + COMPRESS_RATIO, write_r_start, !IS_OVERLAP); - } - - // 1b. Remainder tokens (last block only). - // Tokens past the last complete window are persisted so a later call can - // continue the same compression window. - // rem_start_pos < sp when the chunk has no full window (actual_num_outputs == 0); - // rem_write_start skips those already-paged positions without a per-iteration branch. - if (local_output_idx == num_outputs - 1) - { - int const rem_start_pos = last_abs_idx * COMPRESS_RATIO; - int const rem_count = kv_len - rem_start_pos; - int const rem_write_start = (rem_start_pos < sp) ? (sp - rem_start_pos) : 0; - write_to_paged(rem_start_pos, rem_start_pos + rem_count, rem_write_start, false); - } - - // ================================================================ - // Phase 2: Online softmax reduction (vectorized) - // - // Each block reduces compress_ratio rows into one output via per-element - // online softmax: output[d] = sum_r(kv[r,d] * softmax(score[r,d] + ape[r,d])) - // In overlap mode, combines previous chunk's first-half and current chunk's - // second-half features (same as decode kernel). - // ================================================================ - if (!should_compress) - return; - - if constexpr (IS_OVERLAP) - { -#pragma unroll - for (int i = 0; i < VEC; i++) - { - rmax[i] = -INFINITY; - rsum[i] = 0.0f; - rwsum[i] = 0.0f; - } - } - - constexpr int positions_per_warp = COMPRESS_RATIO / NUM_RED_WARPS; - int const my_r_start = red_warp * positions_per_warp; - int const my_r_end = (red_warp == NUM_RED_WARPS - 1) ? COMPRESS_RATIO : (my_r_start + positions_per_warp); - - // Helper: online-softmax reduction over a contiguous window of COMPRESS_RATIO positions. - // Positions [range_start, range_start + new_start) come from paged cache (APE already - // fused during the call that wrote them); positions [range_start + new_start, range_start - // + COMPRESS_RATIO) are new tokens read from kv_score with live APE addition. - // Two branch-free loops avoid a per-iteration if/else on pos < sp. - // kv_col_off — column offset into kv_score / paged state (0 or HEAD_DIM for overlap) - // ape_col_off — column offset into APE table for the score column (same as kv_col_off) - auto reduce_window = [&](int range_start, int kv_col_off, int ape_col_off, bool skipNewTokens) - { - // Precompute split once for the whole window. - int const new_start = (range_start < sp) ? min(sp - range_start, COMPRESS_RATIO) : 0; - - // Paged portion — APE already fused when tokens were stored. - int const paged_r_end = min(new_start, my_r_end); - for (int r = my_r_start; r < paged_r_end; r++) - { - int const pos = range_start + r; - int log_blk = pos / page_size; - int blk_off = pos % page_size; - int phys_kv = block_table_kv[batch_idx * max_blocks + log_blk]; - int phys_sc = block_table_score[batch_idx * max_blocks + log_blk]; - decodeSoftmaxVec(paged_kv_raw, paged_score_raw, page_sd, state_dim, - phys_kv, phys_sc, blk_off, kv_col_off, eff_tid, rmax, rsum, rwsum); - } - - // New-token portion — read from kv_score and add APE live. - if (!skipNewTokens) - { - int const new_r_start = max(my_r_start, new_start); - for (int r = new_r_start; r < my_r_end; r++) - { - int const pos = range_start + r; - int const input_row = pos - sp; - int64_t const row = static_cast(input_offset + input_row) * two_sd; - prefillSoftmaxVec(kv_score_raw, ape, row, kv_col_off, - r * state_dim + ape_col_off, state_dim, eff_tid, rmax, rsum, rwsum); - } - } - }; - - if constexpr (IS_OVERLAP) - { - // Overlap mode: each output combines - // prev-segment (first head_dim, kv_col=0) from window (abs_idx-1) - // curr-segment (second head_dim, kv_col=HD) from window abs_idx - if (abs_idx > 0) - reduce_window((abs_idx - 1) * COMPRESS_RATIO, 0, 0, false); // prev window, first half - reduce_window(win_start, HEAD_DIM, HEAD_DIM, false); // curr window, second half - } - else - { - // Non-overlap mode: new tokens were reduced while writing paged state; - // only a reused prefix, if any, still needs to be read from paged state. - reduce_window(win_start, 0, 0, true); - } - - if constexpr (NUM_RED_WARPS > 1) - { - extern __shared__ float smem[]; - float* s_rmax = smem; - float* s_rsum = s_rmax + NUM_RED_WARPS * ELEM_PER_BLOCK; - float* s_rwsum = s_rsum + NUM_RED_WARPS * ELEM_PER_BLOCK; - -#pragma unroll - for (int i = 0; i < VEC; i++) - { - int const local_elem = tid * VEC + i; - s_rmax[red_warp * ELEM_PER_BLOCK + local_elem] = rmax[i]; - s_rsum[red_warp * ELEM_PER_BLOCK + local_elem] = rsum[i]; - s_rwsum[red_warp * ELEM_PER_BLOCK + local_elem] = rwsum[i]; - } - __syncthreads(); - - if (red_warp == 0) - { - for (int w = 1; w < NUM_RED_WARPS; w++) - { -#pragma unroll - for (int i = 0; i < VEC; i++) - { - int const local_elem = tid * VEC + i; - float const m2 = s_rmax[w * ELEM_PER_BLOCK + local_elem]; - float const s2 = s_rsum[w * ELEM_PER_BLOCK + local_elem]; - float const ws2 = s_rwsum[w * ELEM_PER_BLOCK + local_elem]; - - float const nm = fmaxf(rmax[i], m2); - float const sc1 = expf(rmax[i] - nm); - float const sc2 = expf(m2 - nm); - rsum[i] = rsum[i] * sc1 + s2 * sc2; - rwsum[i] = rwsum[i] * sc1 + ws2 * sc2; - rmax[i] = nm; - } - } - } - } - - // ================================================================ - // Store output (vectorized) - // ================================================================ - bool const should_write = (NUM_RED_WARPS == 1) || (red_warp == 0); - if (!should_write) - { - return; - } - - int64_t const out_base = static_cast(output_offset + local_output_idx) * HEAD_DIM + eff_tid * VEC; - - if (out_elem_bytes == 2) - { - __nv_bfloat16 packed[VEC]; -#pragma unroll - for (int i = 0; i < VEC; i++) - packed[i] = __float2bfloat16_rn(rwsum[i] / rsum[i]); - // VEC * 2 bytes: VEC=4 → 8B (uint2), VEC=8 → 16B (uint4) - using OutVecT = typename VecType::type; - *reinterpret_cast(&reinterpret_cast<__nv_bfloat16*>(output_raw)[out_base]) - = *reinterpret_cast(packed); - } - else - { - float result[VEC]; -#pragma unroll - for (int i = 0; i < VEC; i++) - result[i] = rwsum[i] / rsum[i]; - // VEC * 4 bytes: VEC=4 → 16B (uint4/float4), VEC=8 → 32B (2×float4) -#pragma unroll - for (int i = 0; i < VEC; i += 4) - *reinterpret_cast(&reinterpret_cast(output_raw)[out_base + i]) - = *reinterpret_cast(&result[i]); - } -} - -// Explicit instantiations. CR=128 uses four reduction groups to split the long -// token loop; CR=4 stays single-warp because the reduction is too small to amortize -// the merge overhead. -#define INST_PREFILL(HD, KV_EB, STATE_EB, CR, NRW) \ - template __global__ void prefillReductionKernel(void const*, float const*, void*, \ - void*, int32_t const*, int32_t const*, void*, int32_t const*, int32_t const*, int32_t const*, int32_t const*, \ - int, int, int, int); - -#define INST_PREFILL_DTYPES(HD, CR, NRW) \ - INST_PREFILL(HD, 2, 2, CR, NRW) \ - INST_PREFILL(HD, 2, 4, CR, NRW) INST_PREFILL(HD, 4, 2, CR, NRW) INST_PREFILL(HD, 4, 4, CR, NRW) - -INST_PREFILL_DTYPES(128, 4, 1) -INST_PREFILL_DTYPES(128, 128, 4) -INST_PREFILL_DTYPES(512, 4, 1) -INST_PREFILL_DTYPES(512, 128, 4) -#undef INST_PREFILL_DTYPES -#undef INST_PREFILL - -// ============================================================================ -// Prefill Launch Wrapper -// -// Grid is (batch_size, max(max_outputs, 1), head_blocks). Blocks for -// local_output_idx >= num_outputs early-exit inside the kernel. -// ============================================================================ - -// Compute vector width: mirrors compile-time VEC. -static inline int prefillVec(int head_dim, int elem_bytes_for_vec) -{ - int max_vec = 16 / elem_bytes_for_vec; - return (head_dim / max_vec >= 32) ? max_vec : (head_dim / 32); -} - -void prefillReductionLaunch(void const* kv_score, float const* ape, void* paged_kv, void* paged_score, - int32_t const* block_table_kv, int32_t const* block_table_score, void* output, int32_t const* kv_lens, - int32_t const* start_pos, int32_t const* cu_seq_lens, int32_t const* cu_kv_comp, int batch_size, int page_size, - int max_blocks, int head_dim, int compress_ratio, int max_outputs, int kv_score_elem_bytes, int state_elem_bytes, - int out_elem_bytes, cudaStream_t stream) -{ - bool const overlap = (compress_ratio == 4); - TLLM_CHECK_WITH_INFO( - compress_ratio == 4 || compress_ratio == 128, "prefillReductionLaunch only supports compress_ratio 4 or 128"); - TLLM_CHECK_WITH_INFO( - (kv_score_elem_bytes == 2 || kv_score_elem_bytes == 4) && (state_elem_bytes == 2 || state_elem_bytes == 4), - "prefillReductionLaunch only supports bf16/fp32 kv_score and paged state"); - int const elem_bytes_for_vec = max(kv_score_elem_bytes, state_elem_bytes); - int const vec = prefillVec(head_dim, elem_bytes_for_vec); - int const nthrd_base = head_dim / vec; - constexpr int MULTI_WARP = 4; - bool const use_multi_warp = (compress_ratio == 128); - int const head_blocks = (use_multi_warp && nthrd_base > 32) ? (nthrd_base / 32) : 1; - int const nthreads_inner = nthrd_base / head_blocks; - int const num_red_warps = use_multi_warp ? MULTI_WARP : 1; - int const nthreads = nthreads_inner * num_red_warps; - int const elem_per_block = nthreads_inner * vec; - int const smem_bytes = use_multi_warp ? (3 * MULTI_WARP * elem_per_block * static_cast(sizeof(float))) : 0; - int const coff = overlap ? 2 : 1; - int const state_dim = coff * head_dim; - dim3 grid(batch_size, max(max_outputs, 1), head_blocks); - -#define LAUNCH_PREFILL(HD, KV_EB, STATE_EB, CR, NRW) \ - prefillReductionKernel<<>>(kv_score, ape, \ - paged_kv, paged_score, block_table_kv, block_table_score, output, kv_lens, start_pos, cu_seq_lens, cu_kv_comp, \ - page_size, state_dim, max_blocks, out_elem_bytes) - -#define DISPATCH_PREFILL_DTYPE(HD, CR, NRW) \ - do \ - { \ - if (kv_score_elem_bytes == 4 && state_elem_bytes == 4) \ - { \ - LAUNCH_PREFILL(HD, 4, 4, CR, NRW); \ - } \ - else if (kv_score_elem_bytes == 2 && state_elem_bytes == 4) \ - { \ - LAUNCH_PREFILL(HD, 2, 4, CR, NRW); \ - } \ - else if (kv_score_elem_bytes == 4 && state_elem_bytes == 2) \ - { \ - LAUNCH_PREFILL(HD, 4, 2, CR, NRW); \ - } \ - else \ - { \ - LAUNCH_PREFILL(HD, 2, 2, CR, NRW); \ - } \ - } while (false) - - if (head_dim == 512) - { - if (compress_ratio == 4) - DISPATCH_PREFILL_DTYPE(512, 4, 1); - else - DISPATCH_PREFILL_DTYPE(512, 128, 4); - } - else - { - if (compress_ratio == 4) - DISPATCH_PREFILL_DTYPE(128, 4, 1); - else - DISPATCH_PREFILL_DTYPE(128, 128, 4); - } - -#undef DISPATCH_PREFILL_DTYPE -#undef LAUNCH_PREFILL -} - -// ============================================================================ -// Postprocess + Scatter Kernel: postProcessScatterKernel -// -// Template: -// -// Grid: (total_tokens) — one block per compressed token -// Block: (NTHRD = HEAD_DIM / VEC) threads, always >= 32 -// Smem: HEAD_DIM * sizeof(float) — used for cross-warp Hadamard butterfly -// -// This kernel fuses all post-compression processing with the paged cache write -// into a single kernel launch, keeping data in float32 registers throughout. -// This eliminates the DRAM round-trip that a split postprocess+scatter would need. -// -// SCALE_TYPE alone determines the output cache layout: -// - kNone: ELEM_BYTES per value (bf16 / fp32) into kv_cache. -// - kFP8PerTensor: one fp8 byte per value, implicit scale=1.0. -// - kFP8Blockwise: one fp8 byte per value + one fp32 scale per 128 values. -// - kMXFP4Blockwise: packed fp4 (two values per byte) + one ue8m0 byte -// per 32 values. -// -// Pipeline (10 steps, all in fp32 registers): -// 1. Vectorized load compressed token from kv_comp -// 2. RMSNorm: compute sum-of-squares → cross-warp reduce → rsqrt → scale -// 3. Apply RMSNorm weights -// 4. RoPE: interleaved even/odd rotation on rope_dim elements (skip nope_dim) -// 5. Hadamard butterfly transform (3 phases: local → warp shuffle → shared mem) -// 6. Scale by 1/sqrt(HEAD_DIM) (Hadamard normalization) -// 7. Optionally write postprocessed result to kv_out (for callers that need it) -// 8. Binary search cu_kv_comp to find batch_idx for this token -// 9. Compute paged cache destination (logical→physical block via block table) -// 10. Store to cache (layout by SCALE_TYPE). -// ============================================================================ - -template -__global__ void postProcessScatterKernel(void const* __restrict__ kv_comp, // [total_tokens, head_dim] input - void* __restrict__ kv_out, // [total_tokens, head_dim] postprocessed output (may be nullptr) - void const* __restrict__ rms_weight, // [head_dim] - float rms_eps, - float const* __restrict__ cos_sin_table, // [max_pos, 2, rope_dim/2] - int32_t const* __restrict__ position_ids, // [total_tokens] - int nope_dim, int rope_dim, - // scatter params - void* __restrict__ kv_cache, // paged cache buffer - int32_t const* __restrict__ num_outputs_arr, int32_t const* __restrict__ cu_kv_comp, - int32_t const* __restrict__ start_pos_arr, int32_t const* __restrict__ block_offsets, - bool const* __restrict__ compressed_mask, int batch_size, int tokens_per_block, int max_blocks, - int cache_stride_blk_bytes, int total_tokens, int num_scale_blocks, void* __restrict__ quant_output, - void* __restrict__ scale_output) -{ - using ElementT = typename std::conditional::type; - constexpr int MAX_VEC = 16 / ELEM_BYTES; - constexpr int VEC = (HEAD_DIM / MAX_VEC >= 32) ? MAX_VEC : (HEAD_DIM / 32); - constexpr int NTHRD = HEAD_DIM / VEC; - constexpr int VEC_BYTES = VEC * ELEM_BYTES; - using VecT = typename VecType::type; - - int const token_idx = blockIdx.x; - if (token_idx >= total_tokens) - return; - - // Per-token mask: precomputed on host, skips padded generation slots - // and batches that produced no compressed tokens. - if (!compressed_mask[token_idx]) - return; - - // ================================================================ - // Step 0: Find owning batch via binary search on cu_kv_comp. - // ================================================================ - int batch_idx, local_output_idx; - if (batch_size <= 1) - { - batch_idx = 0; - local_output_idx = token_idx; - } - else - { - int lo = 0, hi = batch_size; - while (lo < hi) - { - int mid = (lo + hi) >> 1; - if (cu_kv_comp[mid + 1] <= token_idx) - lo = mid + 1; - else - hi = mid; - } - batch_idx = lo; - if (batch_idx >= batch_size) - return; - local_output_idx = token_idx - cu_kv_comp[batch_idx]; - } - - if (local_output_idx >= num_outputs_arr[batch_idx]) - return; - - int const tid = threadIdx.x; - extern __shared__ float smem[]; - - // ================================================================ - // Step 1: Vectorized load from kv_comp - // ================================================================ - auto const* src = reinterpret_cast( - reinterpret_cast(kv_comp) + static_cast(token_idx) * HEAD_DIM); - VecT raw_in = src[tid]; - ElementT const* in_elems = reinterpret_cast(&raw_in); - - float v[VEC]; -#pragma unroll - for (int i = 0; i < VEC; i++) - v[i] = static_cast(in_elems[i]); - - // ================================================================ - // Step 2: RMSNorm - // ================================================================ - float local_sq = 0.f; -#pragma unroll - for (int i = 0; i < VEC; i++) - local_sq += v[i] * v[i]; - - float warp_sum = warpReduceSum(local_sq); - - constexpr int NUM_WARPS = (NTHRD + 31) / 32; - int const warp_id = tid / 32; - int const lane_id = tid % 32; - - if (lane_id == 0) - smem[warp_id] = warp_sum; - __syncthreads(); - - if (warp_id == 0) - { - float s = (lane_id < NUM_WARPS) ? smem[lane_id] : 0.0f; - for (int offset = 16; offset > 0; offset >>= 1) - s += __shfl_xor_sync(0xFFFFFFFF, s, offset); - if (lane_id == 0) - smem[0] = s; - } - __syncthreads(); - float const rms_scale = rsqrtf(smem[0] / static_cast(HEAD_DIM) + rms_eps); - - // ================================================================ - // Step 3: Load weight, apply RMSNorm - // ================================================================ - auto const* wt_src = reinterpret_cast(reinterpret_cast(rms_weight)); - VecT raw_w = wt_src[tid]; - ElementT const* w_elems = reinterpret_cast(&raw_w); - -#pragma unroll - for (int i = 0; i < VEC; i++) - v[i] = v[i] * rms_scale * static_cast(w_elems[i]); - - // ================================================================ - // Step 4: RoPE (Rotary Positional Embedding) - // - // Applied only to elements in [nope_dim, nope_dim+rope_dim). - // Uses interleaved even/odd pairs: (x_even, x_odd) → rotated by (cos, sin). - // cos_sin_table layout: [max_pos, rope_dim] where first half is cos, second is sin. - // ================================================================ - int const half_rope = rope_dim / 2; - int const pos_id = position_ids[token_idx]; - -#pragma unroll - for (int i = 0; i < VEC; i += 2) - { - int const elem_idx = tid * VEC + i; - if (elem_idx >= nope_dim) - { - int const rope_idx = elem_idx - nope_dim; - int const d = rope_idx >> 1; - float const cos_v = cos_sin_table[pos_id * rope_dim + d]; - float const sin_v = cos_sin_table[pos_id * rope_dim + half_rope + d]; - float const x_even = v[i]; - float const x_odd = v[i + 1]; - v[i] = x_even * cos_v - x_odd * sin_v; - v[i + 1] = x_odd * cos_v + x_even * sin_v; - } - } - - // ================================================================ - // Step 5: Hadamard butterfly transform (rotate activation) - // - // Implements the Walsh-Hadamard transform H_n * v via butterfly network. - // H_n has the recursive structure: H_n = [[H_{n/2}, H_{n/2}], [H_{n/2}, -H_{n/2}]] - // which decomposes into log2(HEAD_DIM) butterfly stages. - // - // Three phases handle increasing stride lengths: - // A) Local: strides < VEC — within each thread's register file - // B) Warp shuffle: strides VEC..32*VEC-1 — via __shfl_xor_sync - // C) Shared memory: strides >= 32*VEC — via XOR-swizzled smem - // - // The XOR swizzle pattern `idx ^ ((idx >> 3) & 0x1F)` ensures bank-conflict- - // free access to shared memory across all butterfly stride patterns. - // - // Skipped entirely when ROTATE_ACTIVATION=false. - // ================================================================ - if constexpr (ROTATE_ACTIVATION) - { - - // Phase A: local butterfly (strides 1..VEC-1, within each thread's VEC registers) -#pragma unroll - for (int stride = 1; stride < VEC; stride <<= 1) - { -#pragma unroll - for (int i = 0; i < VEC; i++) - { - if ((i & stride) == 0) - { - float a = v[i], b = v[i ^ stride]; - v[i] = a + b; - v[i ^ stride] = a - b; - } - } - } - - // Phase B: warp shuffle butterfly (strides VEC..32*VEC-1, within a single warp) - if constexpr (NTHRD > 1) - { - constexpr int SHFL_END = (NTHRD <= 32) ? NTHRD : 32; -#pragma unroll - for (int ts = 1; ts < SHFL_END; ts <<= 1) - { - int const stride = ts * VEC; -#pragma unroll - for (int i = 0; i < VEC; i++) - { - float partner = __shfl_xor_sync(0xFFFFFFFF, v[i], ts); - int const elem_idx = tid * VEC + i; - v[i] = (elem_idx & stride) ? (partner - v[i]) : (v[i] + partner); - } - } - } - - // Phase C: cross-warp butterfly via XOR-swizzled shared memory (strides >= 32*VEC) - // Only needed when NTHRD > 32 (i.e., multiple warps, e.g., HEAD_DIM=512, VEC=8 → 64 threads) - if constexpr (NTHRD > 32) - { -#pragma unroll - for (int i = 0; i < VEC; i++) - { - int const idx = tid * VEC + i; - smem[idx ^ ((idx >> 3) & 0x1F)] = v[i]; - } - __syncthreads(); - - for (int stride = 32 * VEC; stride < HEAD_DIM; stride <<= 1) - { -#pragma unroll - for (int i = 0; i < VEC; i++) - { - int const idx = tid * VEC + i; - int const partner_idx = idx ^ stride; - float const a = smem[idx ^ ((idx >> 3) & 0x1F)]; - float const b = smem[partner_idx ^ ((partner_idx >> 3) & 0x1F)]; - v[i] = (idx & stride) ? (b - a) : (a + b); - } - __syncthreads(); -#pragma unroll - for (int i = 0; i < VEC; i++) - { - int const idx = tid * VEC + i; - smem[idx ^ ((idx >> 3) & 0x1F)] = v[i]; - } - __syncthreads(); - } - } - - // ================================================================ - // Step 6: Scale by Hadamard factor - // ================================================================ - float const had_scale = rsqrtf(static_cast(HEAD_DIM)); - -#pragma unroll - for (int i = 0; i < VEC; i++) - v[i] *= had_scale; - - } // ROTATE_ACTIVATION - - // ================================================================ - // Step 7: Write postprocessed output to kv_out (if requested) - // ================================================================ - if (kv_out != nullptr) - { - VecT raw_out; - ElementT* out_elems = reinterpret_cast(&raw_out); -#pragma unroll - for (int i = 0; i < VEC; i++) - out_elems[i] = static_cast(v[i]); - - auto* dst - = reinterpret_cast(reinterpret_cast(kv_out) + static_cast(token_idx) * HEAD_DIM); - dst[tid] = raw_out; - } - - // ================================================================ - // Step 9: Compute paged cache destination address. - // Map (batch_idx, local_output_idx) → logical block → physical block - // via the block table. block_base points to the start of the physical - // page; token_offset is the slot within that page. - // ================================================================ - int const start_pos = start_pos_arr[batch_idx]; - int const cache_pos = start_pos + local_output_idx; - int const logical_block = cache_pos / tokens_per_block; - int const token_offset = cache_pos % tokens_per_block; - int const phys_block = block_offsets[batch_idx * max_blocks + logical_block]; - - uint8_t* block_base - = reinterpret_cast(kv_cache) + static_cast(phys_block) * cache_stride_blk_bytes; - - // ================================================================ - // Step 11: Store to cache (compile-time dispatch on cache dtype/scale type) - // - // Cache addressing is byte-based: block_base points to the start of - // the physical block, cache_stride_blk_bytes is the total block size. - // ================================================================ - if constexpr (SCALE_TYPE == CacheScaleType::kNone) - { - // Default mode: float→bf16/fp32 pack + vectorized store. - // Cache layout per block: [tokens_per_block * HEAD_DIM] elements of ElementT. - VecT raw_out; - ElementT* out_elems = reinterpret_cast(&raw_out); -#pragma unroll - for (int i = 0; i < VEC; i++) - out_elems[i] = static_cast(v[i]); - - ElementT* row_base = reinterpret_cast(block_base) + token_offset * HEAD_DIM; - reinterpret_cast(row_base)[tid] = raw_out; - } - else if constexpr (SCALE_TYPE == CacheScaleType::kFP8PerTensor) - { - // FP8 per-tensor: direct float→fp8_e4m3fn cast (implicit scale=1.0). - // Cache layout per block: [tokens_per_block * HEAD_DIM] bytes of fp8. - uint8_t fp8_bytes[VEC]; -#pragma unroll - for (int i = 0; i < VEC; i++) - { - __nv_fp8_e4m3 fp8_val(v[i]); - fp8_bytes[i] = *reinterpret_cast(&fp8_val); - } - - using Fp8VecT = typename VecType::type; - uint8_t* fp8_dst = block_base + token_offset * HEAD_DIM + tid * VEC; - *reinterpret_cast(fp8_dst) = *reinterpret_cast(fp8_bytes); - } - else if constexpr (SCALE_TYPE == CacheScaleType::kFP8Blockwise) - { - // FP8 blockwise: per-128-element quantization with explicit scales. - // GROUP_SIZE = number of threads that share one scale factor. - // For HD=512, VEC=8: GROUP_SIZE=16 threads → 128 elements per scale block. - // - // Cache layout per block: - // [fp8_data: tokens_per_block * HEAD_DIM bytes] - // [scales: tokens_per_block * (HEAD_DIM/128) * 4 bytes] - constexpr int GROUP_SIZE = 128 / VEC; - - // Step 11a: Compute per-group amax via warp shuffle reduction. - // GROUP_SIZE <= 16 (< warp), so shuffle is sufficient (no smem needed). - float local_amax = 0.f; -#pragma unroll - for (int i = 0; i < VEC; i++) - local_amax = fmaxf(local_amax, fabsf(v[i])); - -#pragma unroll - for (int offset = GROUP_SIZE / 2; offset > 0; offset >>= 1) - local_amax = fmaxf(local_amax, __shfl_xor_sync(0xFFFFFFFF, local_amax, offset)); - - // Step 11b: Compute scale and inverse scale for quantization. - // 448.0 is the max representable value for fp8_e4m3fn. - float const scale = local_amax / 448.0f; - float const inv_scale = (local_amax > 0.f) ? (448.0f / local_amax) : 1.0f; - - // Step 11c: Quantize to FP8 and store data. - uint8_t fp8_bytes[VEC]; -#pragma unroll - for (int i = 0; i < VEC; i++) - { - __nv_fp8_e4m3 fp8_val(v[i] * inv_scale); - fp8_bytes[i] = *reinterpret_cast(&fp8_val); - } - - using Fp8VecT = typename VecType::type; - uint8_t* fp8_dst = block_base + token_offset * HEAD_DIM + tid * VEC; - *reinterpret_cast(fp8_dst) = *reinterpret_cast(fp8_bytes); - - // Step 11d: Store scale factor (one thread per 128-element group writes it). - if (tid % GROUP_SIZE == 0) - { - int const scale_idx = tid / GROUP_SIZE; - float* scale_dst = reinterpret_cast(block_base + tokens_per_block * HEAD_DIM - + (token_offset * num_scale_blocks + scale_idx) * sizeof(float)); - *scale_dst = scale; - } - - // Step 11e: Optionally write FP8 data and scales to output buffers. - // Used by the indexer compressor which returns (fp8_data, scales) to Python - // for downstream sparse attention indexing. - if (quant_output != nullptr) - { - uint8_t* fp8_out_dst - = reinterpret_cast(quant_output) + static_cast(token_idx) * HEAD_DIM + tid * VEC; - *reinterpret_cast(fp8_out_dst) = *reinterpret_cast(fp8_bytes); - } - if (scale_output != nullptr && tid % GROUP_SIZE == 0) - { - int const scale_idx = tid / GROUP_SIZE; - reinterpret_cast(scale_output)[static_cast(token_idx) * num_scale_blocks + scale_idx] - = scale; - } - } - else if constexpr (SCALE_TYPE == CacheScaleType::kMXFP4Blockwise) - { - constexpr int GROUP_SIZE = 32 / VEC; - constexpr int PACKED_VEC_BYTES = VEC / 2; - constexpr float kFp4Max = 6.0f; - constexpr float kFp4MaxInv = 1.0f / kFp4Max; - constexpr float kFp4MinAmax = kFp4Max * 1.1754943508222875e-38f; - - float local_amax = 0.f; -#pragma unroll - for (int i = 0; i < VEC; i++) - local_amax = fmaxf(local_amax, fabsf(v[i])); - -#pragma unroll - for (int offset = GROUP_SIZE / 2; offset > 0; offset >>= 1) - local_amax = fmaxf(local_amax, __shfl_xor_sync(0xFFFFFFFF, local_amax, offset)); - - float const scale = roundedPow2Scale(local_amax, kFp4MaxInv, kFp4MinAmax); - - uint8_t fp4_bytes[PACKED_VEC_BYTES]; -#pragma unroll - for (int i = 0; i < VEC; i += 2) - { - fp4_bytes[i / 2] = packE2M1x2(v[i] / scale, v[i + 1] / scale); - } - - int const packed_head_dim = HEAD_DIM / 2; - uint8_t* fp4_dst = block_base + token_offset * packed_head_dim + tid * PACKED_VEC_BYTES; -#pragma unroll - for (int i = 0; i < PACKED_VEC_BYTES; ++i) - fp4_dst[i] = fp4_bytes[i]; - - if (tid % GROUP_SIZE == 0) - { - int const scale_idx = tid / GROUP_SIZE; - uint8_t* scale_dst - = block_base + tokens_per_block * packed_head_dim + token_offset * num_scale_blocks + scale_idx; - *scale_dst = toUe8m0(scale); - } - - if (quant_output != nullptr) - { - uint8_t* fp4_out_dst = reinterpret_cast(quant_output) - + static_cast(token_idx) * packed_head_dim + tid * PACKED_VEC_BYTES; -#pragma unroll - for (int i = 0; i < PACKED_VEC_BYTES; ++i) - fp4_out_dst[i] = fp4_bytes[i]; - } - if (scale_output != nullptr && tid % GROUP_SIZE == 0) - { - int const scale_idx = tid / GROUP_SIZE; - reinterpret_cast(scale_output)[static_cast(token_idx) * num_scale_blocks + scale_idx] - = toUe8m0(scale); - } - } -} - -// Explicit instantiations — fused postprocess+scatter. -// kNone supports bf16 (EB=2) and fp32 (EB=4) input types; the quantized -// scale types only support bf16 input since the compressor output is bf16. -// Each combination is instantiated with ROTATE_ACTIVATION=true and false. -#define INST_PPS(HD, EB, CST, AR) \ - template __global__ void postProcessScatterKernel(void const*, void*, void const*, float, \ - float const*, int32_t const*, int, int, void*, int32_t const*, int32_t const*, int32_t const*, int32_t const*, \ - bool const*, int, int, int, int, int, int, void*, void*); - -#define INST_PPS_AR(HD, EB, CST) \ - INST_PPS(HD, EB, CST, true) \ - INST_PPS(HD, EB, CST, false) - -INST_PPS_AR(128, 2, CacheScaleType::kNone) -INST_PPS_AR(128, 4, CacheScaleType::kNone) -INST_PPS_AR(512, 2, CacheScaleType::kNone) -INST_PPS_AR(512, 4, CacheScaleType::kNone) -INST_PPS_AR(128, 2, CacheScaleType::kFP8PerTensor) -INST_PPS_AR(512, 2, CacheScaleType::kFP8PerTensor) -INST_PPS_AR(128, 2, CacheScaleType::kFP8Blockwise) -INST_PPS_AR(512, 2, CacheScaleType::kFP8Blockwise) -INST_PPS_AR(128, 2, CacheScaleType::kMXFP4Blockwise) -INST_PPS_AR(512, 2, CacheScaleType::kMXFP4Blockwise) -#undef INST_PPS_AR -#undef INST_PPS - -// ============================================================================ -// Postprocess + Scatter Launch Wrapper -// -// Derives cache layout parameters (cache_stride_blk_bytes, num_scale_blocks) -// from cache dtype / scale type, then dispatches to the appropriate template instantiation. -// ============================================================================ - -// Compute number of threads per block, mirroring the compile-time VEC/NTHRD logic. -// Ensures NTHRD >= 32 by reducing VEC when HEAD_DIM is small. -static inline int compressorNthreads(int head_dim, int elem_bytes) -{ - int max_vec = 16 / elem_bytes; - int vec = (head_dim / max_vec >= 32) ? max_vec : (head_dim / 32); - return head_dim / vec; -} - -void postProcessScatterLaunch(void const* kv_comp, void* kv_out, void const* rms_weight, float rms_eps, - float const* cos_sin_table, int32_t const* position_ids, int nope_dim, int rope_dim, void* kv_cache, - int32_t const* num_outputs, int32_t const* cu_kv_comp, int32_t const* start_pos, int32_t const* block_offsets, - bool const* compressed_mask, int batch_size, int tokens_per_block, int head_dim, int max_blocks_per_seq, - int elem_bytes, int total_tokens, int cache_scale_type, bool rotate_activation, void* quant_output, - void* scale_output, cudaStream_t stream) -{ - if (total_tokens == 0) - { - return; - } - - TLLM_CHECK_WITH_INFO( - cache_scale_type >= 0 && cache_scale_type <= 3, "Invalid cache_scale_type: %d", cache_scale_type); - auto const cst = static_cast(cache_scale_type); - - bool const is_quantized_store = (cst != CacheScaleType::kNone); - int const nthreads = compressorNthreads(head_dim, elem_bytes); - int const smem_bytes = head_dim * sizeof(float); - TLLM_CHECK_WITH_INFO(cst != CacheScaleType::kMXFP4Blockwise || head_dim % 32 == 0, - "MXFP4 cache requires head_dim divisible by 32, got %d", head_dim); - TLLM_CHECK_WITH_INFO(cst != CacheScaleType::kMXFP4Blockwise || head_dim % 2 == 0, - "FP4 packed cache requires even head_dim, got %d", head_dim); - TLLM_CHECK_WITH_INFO(!is_quantized_store || elem_bytes == 2, - "Quantized cache modes require bf16 compressor output, got elem_bytes=%d", elem_bytes); - - // Derive cache block stride (in bytes) and scale block count from the - // scale type. Each physical cache block stores tokens_per_block tokens: - // none: tpb * HD * elem_bytes - // fp8 pertensor: tpb * HD - // fp8 blockwise: tpb * HD + tpb * (HD/128)*4 - // mxfp4: tpb * (HD/2) + tpb * (HD/32) - int num_scale_blocks = 0; - int cache_stride_blk_bytes = 0; - switch (cst) - { - case CacheScaleType::kFP8PerTensor: cache_stride_blk_bytes = tokens_per_block * head_dim; break; - case CacheScaleType::kFP8Blockwise: - num_scale_blocks = head_dim / 128; - cache_stride_blk_bytes = tokens_per_block * head_dim + tokens_per_block * num_scale_blocks * 4; - break; - case CacheScaleType::kMXFP4Blockwise: - num_scale_blocks = head_dim / 32; - cache_stride_blk_bytes = tokens_per_block * (head_dim / 2) + tokens_per_block * num_scale_blocks; - break; - default: cache_stride_blk_bytes = tokens_per_block * head_dim * elem_bytes; break; - } - -#define LAUNCH_PPS(HD, EB, CST, AR) \ - postProcessScatterKernel<<>>(kv_comp, kv_out, \ - rms_weight, rms_eps, cos_sin_table, position_ids, nope_dim, rope_dim, kv_cache, num_outputs, cu_kv_comp, \ - start_pos, block_offsets, compressed_mask, batch_size, tokens_per_block, max_blocks_per_seq, \ - cache_stride_blk_bytes, total_tokens, num_scale_blocks, quant_output, scale_output) - -#define DISPATCH_ROTATE(HD, EB, CST) \ - if (rotate_activation) \ - { \ - LAUNCH_PPS(HD, EB, CST, true); \ - } \ - else \ - { \ - LAUNCH_PPS(HD, EB, CST, false); \ - } -#define DISPATCH_HD_EB(CST) \ - if (elem_bytes == 4) \ - { \ - switch (head_dim) \ - { \ - case 128: DISPATCH_ROTATE(128, 4, CST); break; \ - default: DISPATCH_ROTATE(512, 4, CST); break; \ - } \ - } \ - else \ - { \ - switch (head_dim) \ - { \ - case 128: DISPATCH_ROTATE(128, 2, CST); break; \ - default: DISPATCH_ROTATE(512, 2, CST); break; \ - } \ - } -#define DISPATCH_HD_BF16(CST) \ - switch (head_dim) \ - { \ - case 128: DISPATCH_ROTATE(128, 2, CST); break; \ - default: DISPATCH_ROTATE(512, 2, CST); break; \ - } - - if (cst == CacheScaleType::kFP8PerTensor) - { - DISPATCH_HD_BF16(CacheScaleType::kFP8PerTensor); - } - else if (cst == CacheScaleType::kFP8Blockwise) - { - DISPATCH_HD_BF16(CacheScaleType::kFP8Blockwise); - } - else if (cst == CacheScaleType::kMXFP4Blockwise) - { - DISPATCH_HD_BF16(CacheScaleType::kMXFP4Blockwise); - } - else - { - DISPATCH_HD_EB(CacheScaleType::kNone); - } - -#undef DISPATCH_HD_BF16 -#undef DISPATCH_HD_EB -#undef DISPATCH_ROTATE -#undef LAUNCH_PPS -} - -} // namespace kernels::compressor - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.h b/cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.h deleted file mode 100644 index 0d83628e1b9b..000000000000 --- a/cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.h +++ /dev/null @@ -1,100 +0,0 @@ -/* - * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. - * - * 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. - */ -#pragma once - -#include -#include -#include -#include - -#include "tensorrt_llm/common/config.h" - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels::compressor -{ - -// Decode kernel: write NEXT_N tokens to paged cache + conditional compression -// via online softmax. Overlap is derived from compress_ratio (ratio=4). -// -// Grid: (batch_size, cdiv(state_dim, block_size)) -// One thread per state_dim element across all phases. -// state_dim is a constexpr derived from COMPRESS_RATIO and HEAD_DIM inside the kernel. -void pagedKvCompressLaunch(void const* kv_score, // [m, 2*state_dim] (bf16 or fp32) - float const* ape, // [compress_ratio, state_dim] - void* paged_kv, // [num_blocks, page_size, state_dim] - void* paged_score, // [num_blocks, page_size, state_dim] - int32_t const* block_table_kv, // [bsz, max_blocks] - int32_t const* block_table_score, // [bsz, max_blocks] - void* output, // [total_outputs, head_dim] - int32_t const* kv_lens, // [bsz] - int32_t const* cu_seq_lens, // [bsz+1] - int32_t const* cu_kv_comp, // [bsz+1] - int batch_size, int page_size, int max_blocks, int head_dim, int compress_ratio, int next_n, - int kv_score_elem_bytes, // bytes per element for kv_score (2=bf16, 4=fp32) - int state_elem_bytes, // bytes per element for paged state (2=bf16, 4=fp32) - int out_elem_bytes, // bytes per element for output - cudaStream_t stream); - -// Prefill kernel: bulk compression with per-token gather/scatter + state update. -// Writes all newly seen token states to paged cache for block reuse, then performs -// online softmax reduction. -// -// Grid: (batch_size, max_outputs_per_batch, num_head_chunks) -// Each block computes one compressed output for one head_dim chunk. -void prefillReductionLaunch(void const* kv_score, // [m, 2*state_dim] (bf16 or fp32) - float const* ape, // [compress_ratio, state_dim] - void* paged_kv, // [num_blocks, page_size, state_dim] - void* paged_score, // [num_blocks, page_size, state_dim] - int32_t const* block_table_kv, // [bsz, max_blocks] - int32_t const* block_table_score, // [bsz, max_blocks] - void* output, // [total_outputs, head_dim] - int32_t const* kv_lens, // [bsz] - int32_t const* start_pos, // [bsz] - int32_t const* cu_seq_lens, // [bsz+1] - int32_t const* cu_kv_comp, // [bsz+1] - int batch_size, int page_size, int max_blocks, int head_dim, int compress_ratio, int max_outputs, - int kv_score_elem_bytes, int state_elem_bytes, int out_elem_bytes, cudaStream_t stream); - -// RMSNorm + RoPE + Hadamard + paged scatter in a single kernel launch. -// Optionally writes postprocessed result to kv_out (nullptr to skip). -// -// Grid: (total_tokens) -- one block per compressed token -// Block: (head_dim / VEC), always >= 32 -void postProcessScatterLaunch(void const* kv_comp, // [total_tokens, head_dim] input - void* kv_out, // [total_tokens, head_dim] postprocessed output (nullptr to skip) - void const* rms_weight, // [head_dim] - float rms_eps, - float const* cos_sin_table, // [max_pos, 2, rope_dim/2] - int32_t const* position_ids, // [total_tokens] - int nope_dim, int rope_dim, - void* kv_cache, // paged cache buffer - int32_t const* num_outputs, // [bsz] - int32_t const* cu_kv_comp, // [bsz+1] - int32_t const* start_pos, // [bsz] - int32_t const* block_offsets, // [bsz, max_blocks] - bool const* compressed_mask, // [total_tokens] — per-token mask, false ⇒ skip - int batch_size, int tokens_per_block, int head_dim, int max_blocks_per_seq, int elem_bytes, int total_tokens, - int cache_scale_type, // 0=none (bf16/fp32 by elem_bytes), 1=fp8_pertensor, - // 2=fp8_blockwise, 3=mxfp4 (packed FP4) - bool rotate_activation, // whether to apply Hadamard transform (false to skip) - void* quant_output, // optional fp8/fp4 packed output (nullptr if unused) - void* scale_output, // optional scale output (float* for fp8, uint8_t* for fp4) - cudaStream_t stream); - -} // namespace kernels::compressor - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/CMakeLists.txt b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/CMakeLists.txt index ce076e893ae0..20363e5f9d08 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/CMakeLists.txt +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/CMakeLists.txt @@ -18,10 +18,6 @@ file(GLOB_RECURSE SRC_CPP *.cpp) file(GLOB_RECURSE SRC_CU *.cu) list(FILTER SRC_CU EXCLUDE REGEX "fmha_v2_cu/.*") -# The skip_softmax sm_120/sm_121 warp-specialized FMHA is a hand-written TU that -# only compiles for the sm_120 family; it is added explicitly to the sm_120 arch -# target below, not to the all-arch source list. -list(FILTER SRC_CU EXCLUDE REGEX "skip_softmax_sm120/.*") add_library(context_attention_src OBJECT) @@ -69,28 +65,11 @@ foreach(arch IN ITEMS 80 86 89 90 100 120) endif() file(GLOB arch_files "fmha_v2_cu/*_sm${arch}.cu") - # Compile the hand-written skip_softmax warp-specialized FMHA into the sm_120 - # family target. It uses sm_120-only TMA + sync-MMA, so it must never be - # compiled for other architectures. - if(${arch} EQUAL 120) - list( - APPEND - arch_files - "${CMAKE_CURRENT_SOURCE_DIR}/skip_softmax_sm120/fused_multihead_flash_attention_ws_sm120.cu" - ) - endif() if(arch_files) set(TARGET_NAME _context_attention_kernels_${arch}) add_library(${TARGET_NAME} OBJECT ${arch_files}) target_compile_definitions(${TARGET_NAME} PRIVATE USE_DEMO_BERT_PARAMS=1 GENERATE_CUBIN=1) - # Let the all-arch dispatch TU (fused_multihead_attention_v2.cpp) know the - # skip_softmax bridge symbols are available so it can reference them without - # an undefined-symbol error on builds that exclude sm_120. - if(${arch} EQUAL 120) - target_compile_definitions(context_attention_src - PRIVATE TLLM_ENABLE_SKIP_SOFTMAX_SM120=1) - endif() set_target_properties( ${TARGET_NAME} PROPERTIES POSITION_INDEPENDENT_CODE ON CUDA_RESOLVE_DEVICE_SYMBOLS ON) diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fmhaRunner.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fmhaRunner.cpp index ca7e920ce7fa..e3bd27bdce2c 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fmhaRunner.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fmhaRunner.cpp @@ -498,16 +498,17 @@ void FusedMHARunnerV2::setupLaunchParams(MHARunnerParams runnerParams) && (mLaunchParams.attention_input_layout == AttentionInputLayout::SEPARATE_Q_K_V)))); } - // Skip-softmax is driven by the threshold alone -- there is no separate enable - // flag. It is realized by two kernels: the Hopper warp-specialized FMHA, which - // is selected through the enableSkipSoftmax cubin-hash bit, and the sm_120 / - // sm_121 warp-specialized context FMHA, which reads the threshold directly and - // therefore does not need the cubin-hash bit (enableSkipSoftmax stays false - // there -- see fused_multihead_attention_v2.cpp). If no skip-capable kernel - // matches the config, skipping is simply not enabled and the request runs full - // softmax. - mLaunchParams.enableSkipSoftmax = runnerParams.skipSoftmaxThresholdScaleFactor > 0 && isSm90 - && mLaunchParams.warp_specialization && mLaunchParams.flash_attention; + // Setup launch params for skip softmax attention + mLaunchParams.enableSkipSoftmax = false; + if (runnerParams.skipSoftmaxThresholdScaleFactor > 0) + { + if (!isSm90 || !mLaunchParams.warp_specialization || !mLaunchParams.flash_attention) + { + TLLM_CHECK_WITH_INFO(false, + "Skip softmax attention is only supported on Hopper with warp specialization and flash attention."); + } + mLaunchParams.enableSkipSoftmax = true; + } } //////////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fused_multihead_attention_common.h b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fused_multihead_attention_common.h index f4ccb43ce96e..68c567105665 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fused_multihead_attention_common.h +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fused_multihead_attention_common.h @@ -147,8 +147,6 @@ struct MHARunnerFixedParams bool useSparseMLA = false; // Use sparse attention in trtllm-gen ? bool useTllmGenSparseAttention = false; - // Fuse DSv4 inverse RoPE and FP8 output quantization in trtllm-gen. - bool fusesDsv4InvRopeFp8Quant = false; // Convert to string for debug. std::string convertToStrOutput() @@ -201,7 +199,6 @@ struct MHARunnerFixedParams output += ", sageBlockSizeV = " + std::to_string(sageBlockSizeV); output += ", useSparseMLA = " + std::string(useSparseMLA ? "true" : "false"); output += ", useTllmGenSparseAttention = " + std::string(useTllmGenSparseAttention ? "true" : "false"); - output += ", fusesDsv4InvRopeFp8Quant = " + std::string(fusesDsv4InvRopeFp8Quant ? "true" : "false"); return output; } @@ -292,21 +289,8 @@ struct MHARunnerParams KVBlockArray pagedKvSfCache; // The output buffer ptr. void* outputPtr; - // The output scaling factor buffer ptr. Used by FP4 output and DSv4 fused epilogue. + // The output scaling factor buffer ptr. (only used for FP4 output) void* outputSfPtr; - - struct Dsv4EpilogueFusionParams - { - // Enable DSv4 inverse-RoPE + FP8 quant epilogue fusion. - bool enabled = false; - // The cos/sin cache used by the fused inverse-RoPE epilogue. - float const* cosSinCache = nullptr; - // The physical token stride of the FP32 output scale tensor. - int32_t scaleBufM = 0; - }; - - // DSv4 fused inverse-RoPE + FP8 quant epilogue parameters. - Dsv4EpilogueFusionParams dsv4EpilogueFusion; // The softmax_status ptr for RingAttention. void* softmaxStatsPtr; // The attention sinks ptr. diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fused_multihead_attention_v2.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fused_multihead_attention_v2.cpp index 90e906848468..b47d4fb3cd8f 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fused_multihead_attention_v2.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fused_multihead_attention_v2.cpp @@ -261,52 +261,9 @@ uint64_t FusedMultiHeadAttentionXMMAKernelV2::hashID(KernelMeta const& kernelMet kernelMeta.mSageBlockSizeV, kernelMeta.mReturnSoftmaxStats, kernelMeta.mEnableSkipSoftmax); } -#if defined(TLLM_ENABLE_SKIP_SOFTMAX_SM120) -// Skip_softmax (TMA-load + sync-MMA warp-specialized FMHA for sm_120 / sm_121) -// launch bridges, defined in the skip_softmax TU -// (skip_softmax_sm120/fused_multihead_flash_attention_ws_sm120.cu). One bridge per -// supported head dim. Only declared when sm_120 is built, so non-sm_120 builds -// neither reference nor link the (then-absent) symbols. -void run_skip_softmax_bf16_d256_causal_sm120( - Fused_multihead_attention_params_v2& params, Launch_params const& launch_params, cudaStream_t stream); -void run_skip_softmax_bf16_d128_causal_sm120( - Fused_multihead_attention_params_v2& params, Launch_params const& launch_params, cudaStream_t stream); -#endif // TLLM_ENABLE_SKIP_SOFTMAX_SM120 - void FusedMultiHeadAttentionXMMAKernelV2::run( Fused_multihead_attention_params_v2& params, Launch_params& launch_params, cudaStream_t stream) const { -#if defined(TLLM_ENABLE_SKIP_SOFTMAX_SM120) - // Default sm_120 / sm_121 context FMHA. Every supported prefill is routed to - // the TMA-load + sync-MMA warp-specialized kernel, which carries the per-tile - // skip-softmax optimization. Skipping is active only when a threshold is set - // (enableSkipSoftmax == threshold > 0); with no threshold the kernel runs a - // plain full-softmax prefill. Shapes / features it does not implement fall - // through to the cubin/launcher path: non-BF16 (incl. fp8), head_dim not in - // {128, 256} or head_dim != head_dim_v, non-causal / sliding-window / custom - // mask, non-PACKED_QKV layout, alibi, logit softcapping, sage attention, - // interleaved, and returning softmax stats. - if ((mSM == kSM_120 || mSM == kSM_121) && launch_params.flash_attention && mInputDataType == DATA_TYPE_BF16 - && mOutputDataType == DATA_TYPE_BF16 && params.d == params.dv && (params.d == 128 || params.d == 256) - && launch_params.attention_mask_type == ContextAttentionMaskType::CAUSAL - && launch_params.attention_input_layout == AttentionInputLayout::PACKED_QKV && !params.has_alibi - && !launch_params.enableAttnLogitSoftcapping && !launch_params.interleaved - && params.softmax_stats_ptr == nullptr && launch_params.sage_block_size_q == 0 - && launch_params.sage_block_size_k == 0 && launch_params.sage_block_size_v == 0) - { - if (params.d == 256) - { - run_skip_softmax_bf16_d256_causal_sm120(params, launch_params, stream); - return; - } - if (params.d == 128) - { - run_skip_softmax_bf16_d128_causal_sm120(params, launch_params, stream); - return; - } - } -#endif // TLLM_ENABLE_SKIP_SOFTMAX_SM120 - bool forceUnroll = useForceUnroll(params, launch_params); auto const findIter = mFunctions.find(hashFromParams(params, launch_params)); diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/skip_softmax_sm120/fused_multihead_flash_attention_ws_sm120.cu b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/skip_softmax_sm120/fused_multihead_flash_attention_ws_sm120.cu deleted file mode 100644 index b99bd4fa6a38..000000000000 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/skip_softmax_sm120/fused_multihead_flash_attention_ws_sm120.cu +++ /dev/null @@ -1,199 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2011-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. - */ - -// Translation unit for the skip_softmax sm_120 / sm_121 warp-specialized FMHA -// (TMA-load + sync-MMA). This kernel is hand-written rather than emitted by -// fmha_v2/setup.py code-gen, so it is compiled directly into the -// _context_attention_kernels_120 target (see the sibling CMakeLists.txt). It -// only ever builds for the sm_120 family (sm_120 / sm_121), which is the only -// hardware that provides the TMA + sync-MMA combination this kernel targets. -// -// Supported shapes: BF16 in/out, head_dim == head_dim_v in {128, 256}, causal -// mask, PACKED_QKV layout. This is the default sm_120 / sm_121 context FMHA: the -// runner dispatches here for every prefill that meets those constraints (and -// carries no unsupported feature), regardless of skip-softmax. The per-tile -// skip-softmax optimization is active only when a prefill threshold is set -// (params.skip_softmax_threshold_scale_factor > 0), which the bridge reads -// directly to pick the kernel variant; see fused_multihead_attention_v2.cpp. -// -// The design rationale lives in -// cpp/kernels/fmha_v2/src/fmha/warpspec_sm120/README.md. - -#include - -#include // CUtensorMap - -#include -#include -#include -#include - -#include "fused_multihead_flash_attention_kernel_ws_sm120.h" - -namespace fmha_skip_softmax -{ - -// NOTE: the `S` template arg is the *kv loop step* (per-iter KV tile size), -// not the runtime maximum sequence length. The runtime kv seqlen is read from -// binfo.actual_kv_seqlen and the kv loop iterates in chunks of S. With the TMA -// box size capped at 256 elements per axis, S must be <= 256 (we load -// STEP_KV = Cta_tile_p::N = S elements per TMA box call). -// -// HEAD_DIM must be a multiple of the 64-element (= 128-byte BF16) TMA chunk -// width so the Q/K head-dim chunks and the 64-wide V dv-chunks keep 128-byte -// smem rows -- the layout that matches the TMA 128B hardware swizzle. head_dim -// 128 and 256 both satisfy this; a non-multiple-of-64 head dim would break the -// swizzle invariant. -template -using Skip_softmax_ktraits = fmha::ws_sm120::Kernel_traits_skip_softmax_sm120< - /*Traits_=*/fmha::Ampere_hmma_bf16_traits, - /*S=*/128, - /*VALID_D_=*/HEAD_DIM, - /*VALID_DV_=*/HEAD_DIM, - /*STEP_Q_=*/64, - /*WARPS_M_=*/4, - /*WARPS_N_=*/1, - /*VERSION_=*/2, - /*MASK_VERSION_=*/3, // 3 = causal - /*ENABLE_SKIP_SOFTMAX_=*/ENABLE_SKIP_SOFTMAX, - /*NUM_PRODUCER_WARPS_=*/1>; - -} // namespace fmha_skip_softmax - -// Templated entry kernel -- one instantiation per head dim (128, 256). THREADS -// is head-dim-independent (1 producer + WARPS_M*WARPS_N consumer warps), so the -// launch_bounds value is identical across instantiations. -template -__global__ __launch_bounds__(Ktraits::THREADS, 1) void skip_softmax_kernel( - bert::Fused_multihead_attention_params_v2 const params, __grid_constant__ const CUtensorMap tma_q, - __grid_constant__ const CUtensorMap tma_k, __grid_constant__ const CUtensorMap tma_v) -{ - // The CUtensorMaps live in const/param space (grid_constant); passing - // their addresses to cp.async.bulk.tensor is a valid tensormap operand. - fused_multihead_attention::device_flash_attention_ws_sm120(params, &tma_q, &tma_k, &tma_v); -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// -// -// Host-side launcher for the skip_softmax kernel. -// -// Sets up the TMA descriptors via DMA::Host::init_params, sizes shared memory, -// configures the launch attribute that lets the kernel use >48KB smem, and -// launches the kernel on the given stream. -// -// Returns the cudaError_t from the launch -- the caller is responsible for -// checking it (and running a separate sync to surface launch-time aborts). -//////////////////////////////////////////////////////////////////////////////////////////////////// - -template -static cudaError_t launch_skip_softmax(bert::Fused_multihead_attention_params_v2 params, - bert::Fused_multihead_attention_launch_params const& launch_params, cudaStream_t stream) -{ - // 1. Build the three TMA descriptors host-side (cuTensorMapEncodeTiled). - // These are passed to the kernel as __grid_constant__ params. - CUtensorMap tma_q{}, tma_k{}, tma_v{}; - typename fmha::ws_sm120::DMA::Host dma_host; - dma_host.init_params(params, launch_params, tma_q, tma_k, tma_v); - - // 2. Size the smem allocation. If it exceeds the 48 KB default, raise the - // cudaFuncAttributeMaxDynamicSharedMemorySize cap before launch. (head_dim - // 128 needs roughly half the head_dim 256 footprint.) - constexpr int smem_bytes = static_cast(Ktraits::BYTES_PER_SMEM); - if (smem_bytes >= 48 * 1024) - { - cudaError_t err = cudaFuncSetAttribute( - skip_softmax_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_bytes); - if (err != cudaSuccess) - { - return err; - } - } - - // 3. Grid = (Q-tiles, H, B). One CTA per (Q-tile, head, batch). Each CTA - // has THREADS threads (1 producer warp + WARPS_M*WARPS_N consumer - // warps = 5 warps total = 160 threads). - int const q_tiles = (params.s + Ktraits::STEP_Q - 1) / Ktraits::STEP_Q; - dim3 const grid(q_tiles, params.h, params.b); - dim3 const block(Ktraits::THREADS); - - skip_softmax_kernel<<>>(params, tma_q, tma_k, tma_v); - return cudaGetLastError(); -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// -// -// In-engine dispatch bridges. -// -// The production fmhaRunner uses `tensorrt_llm::kernels::Fused_multihead_attention_params_v2` -// (defined in contextFusedMultiHeadAttention/fused_multihead_attention_common.h), -// which is a separate but ABI-compatible struct from the fmha_v2 `bert::` one -- -// the generated kernels bridge them with reinterpret_cast, and we do the same. -// One bridge per supported head dim; both share the templated launch_skip_softmax<> -// path. The bert launch_params is NOT ABI-identical to kernels::Launch_params, -// so we copy the fields Host::init_params reads rather than reinterpret_cast it. -//////////////////////////////////////////////////////////////////////////////////////////////////// - -#include "../fused_multihead_attention_common.h" -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/config.h" - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels -{ - -void run_skip_softmax_bf16_d256_causal_sm120( - Fused_multihead_attention_params_v2& params, Launch_params const& launch_params, cudaStream_t stream) -{ - bert::Fused_multihead_attention_launch_params blp{}; - blp.total_q_seqlen = launch_params.total_q_seqlen; - blp.attention_input_layout = fmha::Attention_input_layout::PACKED_QKV; - - // Skip-softmax is enabled by the threshold alone (no separate flag): run the - // skip variant when a prefill threshold is set, otherwise the no-skip variant - // (a plain full-softmax prefill) -- the default sm_120 / sm_121 context path. - auto& bert_params = reinterpret_cast(params); - bool const enable_skip = params.skip_softmax_threshold_scale_factor > 0.f; - auto const err = enable_skip - ? ::launch_skip_softmax<::fmha_skip_softmax::Skip_softmax_ktraits<256, true>>(bert_params, blp, stream) - : ::launch_skip_softmax<::fmha_skip_softmax::Skip_softmax_ktraits<256, false>>(bert_params, blp, stream); - TLLM_CHECK_WITH_INFO( - err == cudaSuccess, "run_skip_softmax_bf16_d256_causal_sm120 launch failed: %s", cudaGetErrorString(err)); -} - -void run_skip_softmax_bf16_d128_causal_sm120( - Fused_multihead_attention_params_v2& params, Launch_params const& launch_params, cudaStream_t stream) -{ - bert::Fused_multihead_attention_launch_params blp{}; - blp.total_q_seqlen = launch_params.total_q_seqlen; - blp.attention_input_layout = fmha::Attention_input_layout::PACKED_QKV; - - // Skip-softmax is enabled by the threshold alone (no separate flag): run the - // skip variant when a prefill threshold is set, otherwise the no-skip variant - // (a plain full-softmax prefill) -- the default sm_120 / sm_121 context path. - auto& bert_params = reinterpret_cast(params); - bool const enable_skip = params.skip_softmax_threshold_scale_factor > 0.f; - auto const err = enable_skip - ? ::launch_skip_softmax<::fmha_skip_softmax::Skip_softmax_ktraits<128, true>>(bert_params, blp, stream) - : ::launch_skip_softmax<::fmha_skip_softmax::Skip_softmax_ktraits<128, false>>(bert_params, blp, stream); - TLLM_CHECK_WITH_INFO( - err == cudaSuccess, "run_skip_softmax_bf16_d128_causal_sm120 launch failed: %s", cudaGetErrorString(err)); -} - -} // namespace kernels - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/customMoeRoutingKernels.cu b/cpp/tensorrt_llm/kernels/customMoeRoutingKernels.cu index f76db261e5a3..58ccb4ac8ea8 100644 --- a/cpp/tensorrt_llm/kernels/customMoeRoutingKernels.cu +++ b/cpp/tensorrt_llm/kernels/customMoeRoutingKernels.cu @@ -15,7 +15,6 @@ */ #include "moeTopKFuncs.cuh" -#include "tensorrt_llm/common/assert.h" #include "tensorrt_llm/common/config.h" #include "tensorrt_llm/common/cudaTypeUtils.cuh" #include "tensorrt_llm/common/envUtils.h" @@ -310,176 +309,6 @@ INSTANTIATE_RENORM_MOE_ROUTING(half, __nv_bfloat16, int32_t, true); INSTANTIATE_RENORM_MOE_ROUTING(__nv_bfloat16, __nv_bfloat16, int32_t, true); #endif -static constexpr int kTOPK = 6; - -// CUDA kernel for gate forward -// Input: pre-computed scores from linear(x, weight) done outside kernel -// Template parameters: -// nExperts: number of experts -// topK: number of top experts to select -// hash: true for hash mode, false for topk mode -// One warp per row (batch element) -template -__global__ void gate_forward_kernel( - float const* __restrict__ scores_in, // [batch_size, nExperts] - pre-computed from linear(x, weight) - float const* __restrict__ bias, // [nExperts] (only used when hash=false) - int const* __restrict__ input_ids, // [batch_size] (only used when hash=true) - int const* __restrict__ tid2eid, // [vocab_size, topK] (only used when hash=true) - float* __restrict__ out_weights, // [batch_size, topK] - int* __restrict__ out_indices, // [batch_size, topK] - int batch_size, float route_scale) -{ - // Compile-time constants - constexpr int kExpertsPerThread = nExperts / WARP_SIZE; - constexpr int kWarpsPerBlock = 4; // Adjust based on occupancy needs - - // Shared memory for original scores (one array per warp in the block) - __shared__ float smem_scores[kWarpsPerBlock][nExperts]; - - // One warp per batch element - int const global_warp_id = (blockIdx.x * blockDim.x + threadIdx.x) / WARP_SIZE; - int const local_warp_id = (threadIdx.x / WARP_SIZE) % kWarpsPerBlock; - int const lane_id = threadIdx.x % WARP_SIZE; - - if (global_warp_id >= batch_size) - return; - - auto warp = cg::tiled_partition(cg::this_thread_block()); - - // Pointer to this warp's shared memory and input scores - float* my_smem = smem_scores[local_warp_id]; - float const* scores_row = scores_in + global_warp_id * nExperts; - -// Load scores, apply score function (softplus + sqrt), and store to shared memory -#pragma unroll - for (int e = 0; e < kExpertsPerThread; ++e) - { - int expert_id = lane_id + e * WARP_SIZE; - float s = scores_row[expert_id]; - float sp = log1pf(expf(s)); - float score = sqrtf(sp); - my_smem[expert_id] = score; // Store original score to shared memory - } - __syncwarp(); // Ensure all scores are written before reading - - // Output: each of first K lanes holds one value - float my_topk_value = 0.0f; - int my_topk_index = 0; - - if constexpr (hash) - { - // Hash mode: directly read from shared memory - int token_id = input_ids[global_warp_id]; - int const* expert_ids = tid2eid + token_id * topK; - - if (lane_id < topK) - { - int expert_id = expert_ids[lane_id]; - my_topk_index = expert_id; - my_topk_value = my_smem[expert_id]; // Direct lookup from shared memory - } - } - else - { - // Topk mode: load from shared memory, add bias to registers for topk - float scores[kExpertsPerThread]; - int indices[kExpertsPerThread]; - -#pragma unroll - for (int e = 0; e < kExpertsPerThread; ++e) - { - int expert_id = lane_id + e * WARP_SIZE; - indices[e] = expert_id; - scores[e] = my_smem[expert_id] + bias[expert_id]; // Add bias for topk selection - } - - // Use reduceTopK to find top-k experts - float topk_values[topK]; - int32_t topk_indices[topK]; - constexpr float minValue = -1e30f; - reduce_topk::reduceTopK( - warp, topk_values, topk_indices, scores, indices, minValue, topK); - - // Gather original weights (without bias) from shared memory - if (lane_id < topK) - { - int expert_id = topk_indices[lane_id]; - my_topk_index = expert_id; - my_topk_value = my_smem[expert_id]; // Read original score (no bias) - } - } - - // Reduce to get sum (first K lanes have values, others have 0) - float weight_sum = cg::reduce(warp, my_topk_value, cg::plus{}); - - // Normalize weights and write output (first K lanes) - if (lane_id < topK) - { - out_weights[global_warp_id * topK + lane_id] = (my_topk_value / weight_sum) * route_scale; - out_indices[global_warp_id * topK + lane_id] = my_topk_index; - } -} - -// C++ wrapper function (output tensors passed as parameters) -// All tensors are float32 -template -void launch_gate_forward_kernel(float* scores_in, float* bias, int* input_ids, int* tid2eid, float* out_weights, - int* out_indices, int batch_size, float route_scale, cudaStream_t stream) -{ - constexpr int warps_per_block = 4; - constexpr int threads_per_block = warps_per_block * WARP_SIZE; - int const blocks = (batch_size + warps_per_block - 1) / warps_per_block; - - gate_forward_kernel<<>>( - scores_in, bias, input_ids, tid2eid, out_weights, out_indices, batch_size, route_scale); -} - -void gate_forward(void* scores_in, // [batch_size, nExperts] - pre-computed from linear(x, weight) - void* bias, // nullptr if hash mode - void* input_ids, // nullptr if non-hash mode - void* tid2eid, // nullptr if non-hash mode - void* out_weights, // [batch_size, topK] - pre-allocated - void* out_indices, // [batch_size, topK] - pre-allocated - int batch_size, int n_experts, float route_scale, bool is_hash, cudaStream_t stream) -{ - auto* scores = static_cast(scores_in); - auto* bias_ptr = static_cast(bias); - auto* input_ids_ptr = static_cast(input_ids); - auto* tid2eid_ptr = static_cast(tid2eid); - auto* weights = static_cast(out_weights); - auto* indices = static_cast(out_indices); - - switch (n_experts) - { - case 256: - if (is_hash) - { - launch_gate_forward_kernel<256, true>( - scores, nullptr, input_ids_ptr, tid2eid_ptr, weights, indices, batch_size, route_scale, stream); - } - else - { - launch_gate_forward_kernel<256, false>( - scores, bias_ptr, nullptr, nullptr, weights, indices, batch_size, route_scale, stream); - } - break; - case 384: - if (is_hash) - { - launch_gate_forward_kernel<384, true>( - scores, nullptr, input_ids_ptr, tid2eid_ptr, weights, indices, batch_size, route_scale, stream); - } - else - { - launch_gate_forward_kernel<384, false>( - scores, bias_ptr, nullptr, nullptr, weights, indices, batch_size, route_scale, stream); - } - break; - default: TLLM_CHECK_WITH_INFO(false, "gate_forward only supports n_experts 256 or 384"); - } - sync_check_cuda_error(stream); -} - } // namespace kernels TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/customMoeRoutingKernels.h b/cpp/tensorrt_llm/kernels/customMoeRoutingKernels.h index 367676fcc334..f8240b436393 100644 --- a/cpp/tensorrt_llm/kernels/customMoeRoutingKernels.h +++ b/cpp/tensorrt_llm/kernels/customMoeRoutingKernels.h @@ -1,5 +1,5 @@ /* - * Copyright (c) 2025-2026, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -29,16 +29,6 @@ namespace kernels template void invokeCustomMoeRouting(InputT* routerLogits, OutputT* topkValues, IdxT* topkIndices, int64_t const numTokens, int64_t const numExperts, int64_t const topK, cudaStream_t const stream); - -// Gate forward function for custom MoE routing -// All tensors are expected to be float32 for scores/weights, int32 for indices -void gate_forward(void* scores_in, // [batch_size, nExperts] - pre-computed from linear(x, weight) - void* bias, // nullptr if hash mode - void* input_ids, // nullptr if non-hash mode - void* tid2eid, // nullptr if non-hash mode - void* out_weights, // [batch_size, topK] - pre-allocated - void* out_indices, // [batch_size, topK] - pre-allocated - int batch_size, int n_experts, float route_scale, bool is_hash, cudaStream_t stream); } // namespace kernels TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/cutlass_heuristic.cpp b/cpp/tensorrt_llm/kernels/cutlass_kernels/cutlass_heuristic.cpp index 7bba57a03d5e..b0ea6333a88c 100644 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/cutlass_heuristic.cpp +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/cutlass_heuristic.cpp @@ -408,65 +408,32 @@ std::vector get_candidate_configs_sm100_dynamic_cluster_shape return candidate_configs; } - std::vector> tile_configs; - if ((config & CutlassGemmConfig::MXFP8_MXFP8) != 0) - { - // MXFP8xMXFP8 always instantiates the Mxf8f6f4 block-scaled tensor-op - // with cutlass::arch::Sm100, even on SM103 (the SM103 dispatch case in - // dispatchMoeGemmSelectTileShapeTmaWarpSpecialized only handles FP4xFP4; - // MXFP8 falls through to the sm_version>=100 && <120 branch which - // instantiates Arch=Sm100). Therefore the TMA-only constraint enforced - // by getDispatchFunctionForSM100 (Arch::kMinComputeCapability==103 is - // false for Sm100) applies on both SM100 and SM103, so we filter out - // non-TMA epilogue candidates unconditionally here. - if (schedule != EpilogueScheduleType::TMA) - return {}; - // MXFP8xMXFP8 uses the Mxf8f6f4 block-scaled tensor-op; only TileM=128 - // and TileN in {64,128,256} are valid (kept in sync with the IsMXFPX - // branch in are_tile_shapes_supported_sm100). Returning the broader FP8 - // tile list would crash autotuning with "Unsupported tile shape" since - // the runtime dispatcher rejects the unsupported combinations. - tile_configs = { - {CutlassTileConfigSM100::CtaShape128x64x128B, cluster1sm}, - {CutlassTileConfigSM100::CtaShape128x128x128B, cluster1sm}, - {CutlassTileConfigSM100::CtaShape128x256x128B, cluster1sm}, - }; - if (supports_2sm) - { - tile_configs.push_back({CutlassTileConfigSM100::CtaShape128x64x128B, cluster2sm}); - tile_configs.push_back({CutlassTileConfigSM100::CtaShape128x128x128B, cluster2sm}); - tile_configs.push_back({CutlassTileConfigSM100::CtaShape128x256x128B, cluster2sm}); - } - } - else - { - tile_configs = { - {CutlassTileConfigSM100::CtaShape64x32x128B, cluster1sm}, - {CutlassTileConfigSM100::CtaShape64x64x128B, cluster1sm}, - {CutlassTileConfigSM100::CtaShape64x128x128B, cluster1sm}, - {CutlassTileConfigSM100::CtaShape64x256x128B, cluster1sm}, - {CutlassTileConfigSM100::CtaShape128x32x128B, cluster1sm}, - {CutlassTileConfigSM100::CtaShape128x64x128B, cluster1sm}, - {CutlassTileConfigSM100::CtaShape128x128x128B, cluster1sm}, - {CutlassTileConfigSM100::CtaShape128x256x128B, cluster1sm}, - }; + std::vector> tile_configs{ + {CutlassTileConfigSM100::CtaShape64x32x128B, cluster1sm}, + {CutlassTileConfigSM100::CtaShape64x64x128B, cluster1sm}, + {CutlassTileConfigSM100::CtaShape64x128x128B, cluster1sm}, + {CutlassTileConfigSM100::CtaShape64x256x128B, cluster1sm}, + {CutlassTileConfigSM100::CtaShape128x32x128B, cluster1sm}, + {CutlassTileConfigSM100::CtaShape128x64x128B, cluster1sm}, + {CutlassTileConfigSM100::CtaShape128x128x128B, cluster1sm}, + {CutlassTileConfigSM100::CtaShape128x256x128B, cluster1sm}, + }; - if (supports_2sm) - { - tile_configs.push_back({CutlassTileConfigSM100::CtaShape64x128x128B, cluster2sm}); - tile_configs.push_back({CutlassTileConfigSM100::CtaShape64x256x128B, cluster2sm}); - tile_configs.push_back({CutlassTileConfigSM100::CtaShape64x64x128B, cluster2sm}); - tile_configs.push_back({CutlassTileConfigSM100::CtaShape128x64x128B, cluster2sm}); - tile_configs.push_back({CutlassTileConfigSM100::CtaShape128x128x128B, cluster2sm}); - tile_configs.push_back({CutlassTileConfigSM100::CtaShape128x256x128B, cluster2sm}); - } + if (supports_2sm) + { + tile_configs.push_back({CutlassTileConfigSM100::CtaShape64x128x128B, cluster2sm}); + tile_configs.push_back({CutlassTileConfigSM100::CtaShape64x256x128B, cluster2sm}); + tile_configs.push_back({CutlassTileConfigSM100::CtaShape64x64x128B, cluster2sm}); + tile_configs.push_back({CutlassTileConfigSM100::CtaShape128x64x128B, cluster2sm}); + tile_configs.push_back({CutlassTileConfigSM100::CtaShape128x128x128B, cluster2sm}); + tile_configs.push_back({CutlassTileConfigSM100::CtaShape128x256x128B, cluster2sm}); + } - if (config & CutlassGemmConfig::FP8_ONLY) - { - tile_configs.push_back({CutlassTileConfigSM100::CtaShape128x16x128B, cluster1sm}); - // TODO: re-enable when handled by the MoE GEMM dispatch - // tile_configs.push_back({ CutlassTileConfigSM100::CtaShape128x8x256B, ClusterShape::ClusterShape_1x1x1 }); - } + if (config & CutlassGemmConfig::FP8_ONLY) + { + tile_configs.push_back({CutlassTileConfigSM100::CtaShape128x16x128B, cluster1sm}); + // TODO: re-enable when handled by the MoE GEMM dispatch + // tile_configs.push_back({ CutlassTileConfigSM100::CtaShape128x8x256B, ClusterShape::ClusterShape_1x1x1 }); } for (auto [tile, cluster] : tile_configs) @@ -482,20 +449,10 @@ std::vector get_candidate_configs_sm100( CutlassGemmConfig::CandidateConfigTypeParam const config, int sm) { #ifdef FAST_BUILD - // Fast build limits the candidate set to a single CTA tile shape but - // keeps both 1SM (cluster 1x1x1) and 2SM (cluster 2x1x1) variants so - // the autotuner can profile both. Block-scaled paths (MXFP8xMXFP8, - // NVFP4) accept both; the 2SM variant is required as a candidate so - // FAST_BUILD doesn't accidentally exclude all 2SM kernels (needed for - // MMA M=256 configurations of the Mxf8f6f4 tensor-op). - return { - CutlassGemmConfig{CutlassTileConfigSM100::CtaShape128x128x128B, MainloopScheduleType::AUTO, - EpilogueScheduleType::TMA, ClusterShape::ClusterShape_1x1x1, ClusterShape::Undefined, - ClusterShape::Undefined, sm}, - CutlassGemmConfig{CutlassTileConfigSM100::CtaShape128x128x128B, MainloopScheduleType::AUTO, - EpilogueScheduleType::TMA, ClusterShape::ClusterShape_2x1x1, ClusterShape::Undefined, - ClusterShape::Undefined, sm}, - }; + // Fast build disables all configs except this one for SM100 + return {CutlassGemmConfig{CutlassTileConfigSM100::CtaShape128x128x128B, MainloopScheduleType::AUTO, + EpilogueScheduleType::TMA, ClusterShape::ClusterShape_1x1x1, ClusterShape::Undefined, ClusterShape::Undefined, + sm}}; #else if (config & CutlassGemmConfig::GROUPED_GEMM) { diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/cutlass_heuristic.h b/cpp/tensorrt_llm/kernels/cutlass_kernels/cutlass_heuristic.h index 28e6b7bbf053..b9a788a4c623 100644 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/cutlass_heuristic.h +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/cutlass_heuristic.h @@ -32,22 +32,11 @@ template (TileShape{}))>; - using SupportedCtaShape2Sm = cute::Shape(TileShape{}))>; - using SupportedCgaShape1Sm = cute::Shape; - using SupportedCgaShape2Sm = cute::Shape; + using SupportedCtaShape = cute::Shape(TileShape{}))>; + using SupportedCgaShape = cute::Shape; - constexpr static bool cta_ok - = cute::is_same_v || cute::is_same_v; - constexpr static bool cga_ok - = cute::is_same_v || cute::is_same_v; - constexpr static bool value = !cta_ok || !cga_ok || DYNAMIC_CGA; + constexpr static bool value = !cute::is_same_v + || !cute::is_same_v || DYNAMIC_CGA; #else constexpr static bool value = false; #endif diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/fp4_gemm/fp4_gemm_bf16.cu b/cpp/tensorrt_llm/kernels/cutlass_kernels/fp4_gemm/fp4_gemm_bf16.cu index 17a197b40c6d..f4f4e40c01f8 100644 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/fp4_gemm/fp4_gemm_bf16.cu +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/fp4_gemm/fp4_gemm_bf16.cu @@ -80,7 +80,6 @@ INSTANTIATE_FP4_GEMM_KERNEL_LAUNCHER_SM120(__nv_bfloat16, 256, 128, 128, 1, 1, 1 template class CutlassFp4GemmRunner<__nv_bfloat16, FP4GemmType::W4A4_NVFP4_NVFP4>; template class CutlassFp4GemmRunner<__nv_bfloat16, FP4GemmType::W4A8_MXFP4_MXFP8>; -template class CutlassFp4GemmRunner<__nv_bfloat16, FP4GemmType::W8A8_MXFP8_MXFP8>; #endif diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/fp4_gemm/fp4_gemm_fp16.cu b/cpp/tensorrt_llm/kernels/cutlass_kernels/fp4_gemm/fp4_gemm_fp16.cu index bb6642a9c186..71453157a51a 100644 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/fp4_gemm/fp4_gemm_fp16.cu +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/fp4_gemm/fp4_gemm_fp16.cu @@ -79,7 +79,6 @@ INSTANTIATE_FP4_GEMM_KERNEL_LAUNCHER_SM120(half, 256, 128, 128, 1, 1, 1) template class CutlassFp4GemmRunner; template class CutlassFp4GemmRunner; -template class CutlassFp4GemmRunner; } // namespace cutlass_kernels } // namespace kernels diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/fp4_gemm/fp4_gemm_fp32.cu b/cpp/tensorrt_llm/kernels/cutlass_kernels/fp4_gemm/fp4_gemm_fp32.cu index acf959ee8777..e1870809388d 100644 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/fp4_gemm/fp4_gemm_fp32.cu +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/fp4_gemm/fp4_gemm_fp32.cu @@ -79,7 +79,6 @@ INSTANTIATE_FP4_GEMM_KERNEL_LAUNCHER_SM120(float, 256, 128, 128, 1, 1, 1) template class CutlassFp4GemmRunner; template class CutlassFp4GemmRunner; -template class CutlassFp4GemmRunner; } // namespace cutlass_kernels } // namespace kernels diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/fp4_gemm/fp4_gemm_template.h b/cpp/tensorrt_llm/kernels/cutlass_kernels/fp4_gemm/fp4_gemm_template.h index e5772b5b39e7..854581349af0 100644 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/fp4_gemm/fp4_gemm_template.h +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/fp4_gemm/fp4_gemm_template.h @@ -37,7 +37,6 @@ #include "../include/fp4_gemm.h" #include "mxfp8_mxfp4_gemm_template_sm100.h" -#include "mxfp8_mxfp8_gemm_template_sm100.h" #include "nvfp4_nvfp4_gemm_template_sm100.h" #include "nvfp4_nvfp4_gemm_template_sm120.h" @@ -324,94 +323,6 @@ size_t dispatchMXFP8xMXFP4GemmCTAShapeSm100(T* D, void const* A, void const* B, } } -template -size_t dispatchMXFP8xMXFP8GemmClusterShapeSm100(T* D, void const* A, void const* B, void const* input_sf, - void const* weight_sf, float const* global_sf, int m, int n, int k, int batch_count, - tkc::CutlassGemmConfig gemmConfig, char* workspace, const size_t workspaceBytes, cudaStream_t stream, - int* occupancy = nullptr) -{ - - TLLM_LOG_DEBUG(__PRETTY_FUNCTION__); - - switch (gemmConfig.cluster_shape) - { - case tkc::ClusterShape::ClusterShape_2x1x1: - return genericMXFP8xMXFP8GemmKernelLauncher, cute::Int<1>, cute::Int<1>, - __2SM>(D, A, B, input_sf, weight_sf, global_sf, m, n, k, batch_count, gemmConfig, workspace, workspaceBytes, - stream, occupancy); - break; - case tkc::ClusterShape::ClusterShape_2x2x1: - return genericMXFP8xMXFP8GemmKernelLauncher, cute::Int<2>, cute::Int<1>, - __2SM>(D, A, B, input_sf, weight_sf, global_sf, m, n, k, batch_count, gemmConfig, workspace, workspaceBytes, - stream, occupancy); - break; - case tkc::ClusterShape::ClusterShape_4x2x1: - return genericMXFP8xMXFP8GemmKernelLauncher, cute::Int<2>, cute::Int<1>, - __2SM>(D, A, B, input_sf, weight_sf, global_sf, m, n, k, batch_count, gemmConfig, workspace, workspaceBytes, - stream, occupancy); - break; - case tkc::ClusterShape::ClusterShape_2x4x1: - return genericMXFP8xMXFP8GemmKernelLauncher, cute::Int<4>, cute::Int<1>, - __2SM>(D, A, B, input_sf, weight_sf, global_sf, m, n, k, batch_count, gemmConfig, workspace, workspaceBytes, - stream, occupancy); - break; - case tkc::ClusterShape::ClusterShape_4x4x1: - return genericMXFP8xMXFP8GemmKernelLauncher, cute::Int<4>, cute::Int<1>, - __2SM>(D, A, B, input_sf, weight_sf, global_sf, m, n, k, batch_count, gemmConfig, workspace, workspaceBytes, - stream, occupancy); - break; - default: - throw std::runtime_error( - "[TensorRT LLM Error][MXFP8][dispatch_gemm_cluster_shape] Config is invalid for MXFP8xMXFP8 GEMM."); - break; - } -} - -template -size_t dispatchMXFP8xMXFP8GemmCTAShapeSm100(T* D, void const* A, void const* B, void const* input_sf, - void const* weight_sf, float const* global_sf, int m, int n, int k, int batch_count, - tkc::CutlassGemmConfig gemmConfig, char* workspace, const size_t workspaceBytes, cudaStream_t stream, - int* occupancy = nullptr) -{ - - TLLM_LOG_DEBUG(__PRETTY_FUNCTION__); - switch (gemmConfig.tile_config_sm100) - { - case tkc::CutlassTileConfigSM100::CtaShape128x64x128B: - return dispatchMXFP8xMXFP8GemmClusterShapeSm100, cute::Int<64>, cute::Int<128>>(D, A, B, - input_sf, weight_sf, global_sf, m, n, k, batch_count, gemmConfig, workspace, workspaceBytes, stream, - occupancy); - break; - case tkc::CutlassTileConfigSM100::CtaShape128x256x128B: - return dispatchMXFP8xMXFP8GemmClusterShapeSm100, cute::Int<256>, cute::Int<128>>(D, A, B, - input_sf, weight_sf, global_sf, m, n, k, batch_count, gemmConfig, workspace, workspaceBytes, stream, - occupancy); - break; - case tkc::CutlassTileConfigSM100::CtaShape128x128x256B: - return dispatchMXFP8xMXFP8GemmClusterShapeSm100, cute::Int<128>, cute::Int<256>>(D, A, B, - input_sf, weight_sf, global_sf, m, n, k, batch_count, gemmConfig, workspace, workspaceBytes, stream, - occupancy); - break; - case tkc::CutlassTileConfigSM100::CtaShape128x256x256B: - return dispatchMXFP8xMXFP8GemmClusterShapeSm100, cute::Int<256>, cute::Int<256>>(D, A, B, - input_sf, weight_sf, global_sf, m, n, k, batch_count, gemmConfig, workspace, workspaceBytes, stream, - occupancy); - break; - case tkc::CutlassTileConfigSM100::Undefined: - throw std::runtime_error("[TensorRT LLM Error][MXFP8][dispatch_gemm_cta_shape] Gemm config undefined."); - break; - case tkc::CutlassTileConfigSM100::ChooseWithHeuristic: - throw std::runtime_error( - "[TensorRT LLM Error][MXFP8][dispatch_gemm_cta_shape] Gemm config should have already been set by " - "heuristic."); - break; - default: - throw std::runtime_error( - "[TensorRT LLM Error][MXFP8][dispatch_gemm_cta_shape] Config is invalid for MXFP8xMXFP8 GEMM."); - break; - } -} - template CutlassFp4GemmRunner::CutlassFp4GemmRunner() { @@ -447,19 +358,6 @@ size_t CutlassFp4GemmRunner::dispatchToArch(T* D, void const* A, "[TensorRT LLM Error][CutlassFp4GemmRunner][GEMM Dispatch] Arch unsupported for CUTLASS FP4 GEMM"); } } - else if constexpr (fp4GemmType == FP4GemmType::W8A8_MXFP8_MXFP8) - { - if (mSm == 100 || mSm == 103) - { - return dispatchMXFP8xMXFP8GemmCTAShapeSm100(D, A, B, input_sf, weight_sf, global_sf, m, n, k, - batch_count, gemmConfig, workspace, workspaceBytes, stream, occupancy); - } - else - { - throw std::runtime_error( - "[TensorRT LLM Error][CutlassFp4GemmRunner][GEMM Dispatch] Arch unsupported for CUTLASS MXFP8 GEMM"); - } - } else if constexpr (fp4GemmType == FP4GemmType::W4A4_NVFP4_NVFP4) { if (mSm == 103) @@ -539,12 +437,9 @@ std::vector CutlassFp4GemmRunner::getCon { for (auto const& cluster_config : clusterShapes) { - if constexpr (fp4GemmType == FP4GemmType::W4A8_MXFP4_MXFP8 - || fp4GemmType == FP4GemmType::W8A8_MXFP8_MXFP8) + if constexpr (fp4GemmType == FP4GemmType::W4A8_MXFP4_MXFP8) { - // Skip for high smem usage (MXFP8xMXFP8 has even higher - // smem pressure than MXFP8xMXFP4 because the B operand is - // 2x wider, so the same skips apply). + // Skip for high smem usage. if (cluster_config == tkc::ClusterShape::ClusterShape_1x1x1 || cluster_config == tkc::ClusterShape::ClusterShape_1x2x1 || cluster_config == tkc::ClusterShape::ClusterShape_1x4x1) diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/fp4_gemm/mxfp8_mxfp4_gemm_template_sm100.h b/cpp/tensorrt_llm/kernels/cutlass_kernels/fp4_gemm/mxfp8_mxfp4_gemm_template_sm100.h index ceff762647e2..276de55c69dc 100644 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/fp4_gemm/mxfp8_mxfp4_gemm_template_sm100.h +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/fp4_gemm/mxfp8_mxfp4_gemm_template_sm100.h @@ -84,19 +84,6 @@ struct MXSMTypeAdapter<__2SM> using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized2SmMxf8f6f4Sm100; }; -namespace detail -{ -template -struct has_bias_ptr : std::false_type -{ -}; - -template -struct has_bias_ptr().bias_ptr)>> : std::true_type -{ -}; -} // namespace detail - #ifdef PLACEHOLDER_KERNELS template (global_sf); - if constexpr (detail::has_bias_ptr>::value) - { - fusion_args.bias_ptr = static_cast(bias); - } + fusion_args.bias_ptr = static_cast(bias); operator_args.problem_shape = cute::make_shape(m, n, k, batch_count); diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/fp4_gemm/mxfp8_mxfp8_gemm_template_sm100.h b/cpp/tensorrt_llm/kernels/cutlass_kernels/fp4_gemm/mxfp8_mxfp8_gemm_template_sm100.h deleted file mode 100644 index 4d473d2c4547..000000000000 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/fp4_gemm/mxfp8_mxfp8_gemm_template_sm100.h +++ /dev/null @@ -1,212 +0,0 @@ -/* - * 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. - */ - -#pragma once - -#ifndef _WIN32 -#pragma GCC diagnostic push -#pragma GCC diagnostic ignored "-Wstrict-aliasing" -#endif // #ifndef _WIN32 - -#include "cutlass/cutlass.h" -#include "cutlass/gemm/device/gemm_universal_adapter.h" - -#include "cutlass/arch/arch.h" -#include "cutlass/epilogue/collective/collective_builder.hpp" -#include "cutlass/gemm/collective/collective_builder.hpp" -#include "cutlass/gemm/gemm.h" - -#include "tensorrt_llm/common/config.h" -#include "tensorrt_llm/common/cudaUtils.h" -#include "tensorrt_llm/common/envUtils.h" -#include "tensorrt_llm/kernels/archCondition.h" -#include "tensorrt_llm/kernels/cutlass_kernels/cutlass_type_conversion.h" -#include "tensorrt_llm/kernels/cutlass_kernels/fp4_gemm/mxfp8_mxfp4_gemm_template_sm100.h" - -#ifndef _WIN32 -#pragma GCC diagnostic pop -#endif // #ifndef _WIN32 - -using namespace cute; -using namespace tensorrt_llm::kernels::cutlass_kernels; - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels -{ -namespace cutlass_kernels -{ - -#ifdef PLACEHOLDER_KERNELS - -template -size_t genericMXFP8xMXFP8GemmKernelLauncher(void* D, void const* A, void const* B, void const* input_sf, - void const* weight_sf, float const* global_sf, int m, int n, int k, int batch_count, - tkc::CutlassGemmConfig gemmConfig, char* workspace, const size_t workspaceBytes, cudaStream_t stream, - int* occupancy) -{ - throw std::runtime_error( - "[TensorRT LLM Error][MXFP8 gemm Runner] TensorRT LLM is not compiled with support for this Architecture."); -} - -#else - -template -struct DeviceGemmMXFP8xMXFP8GemmSm100 -{ - using OutElementType = typename TllmToCutlassTypeAdapter::type; - using ClusterShape = cute::Shape; - using Arch = cutlass::arch::Sm100; - /* // Input A: MXFP8 (e4m3 + UE8M0 block scales) */ - using ElementA = cutlass::mx_float8_t; - using LayoutA = cutlass::layout::RowMajor; - static constexpr int AlignmentA = 16; - /* // Input B: MXFP8 (e4m3 + UE8M0 block scales) -- new vs the MXFP4 template */ - using ElementB = cutlass::mx_float8_t; - using LayoutB = cutlass::layout::ColumnMajor; - static constexpr int AlignmentB = 16; - /* // Input C */ - using ElementC = void; - using LayoutC = cutlass::layout::RowMajor; - static constexpr int AlignmentC = 128 / cutlass::sizeof_bits::value; - - using SFType = cutlass::float_ue8m0_t; - using ElementCompute = float; - using ElementAccumulator = float; - using OperatorClass = cutlass::arch::OpClassBlockScaledTensorOp; - using EpilogueTileType = cutlass::epilogue::collective::EpilogueTileAuto; - using EpilogueSchedule = typename MXSMTypeAdapter::EpilogueSchedule; - using MainloopSchedule = typename MXSMTypeAdapter::MainloopSchedule; - using TileScheduler = cutlass::gemm::PersistentScheduler; - using MmaTileShape = cute::Shape::Scale>, CTA_N, CTA_K>; - using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder>::CollectiveOp; - - using CollectiveMainloop = - typename cutlass::gemm::collective::CollectiveBuilder( - sizeof(typename CollectiveEpilogue::SharedStorage))>, - MainloopSchedule>::CollectiveOp; - - template - struct Sm10xOnly : Base - { - using typename Base::Params; - - CUTLASS_DEVICE - void operator()(Params const& params, char* smem_buf) - { - if constexpr (tensorrt_llm::kernels::arch::is_major_v<10>) - { - this->Base::operator()(params, smem_buf); - } - else - { - if (cute::thread0()) - { - printf("%s : This kernel shall only run on SM10x devices.\n", __PRETTY_FUNCTION__); - __trap(); - } - } - } - }; - - using GemmKernel = Sm10xOnly, - CollectiveMainloop, CollectiveEpilogue, TileScheduler>>; - - using Gemm = typename cutlass::gemm::device::GemmUniversalAdapter; -}; - -template -size_t genericMXFP8xMXFP8GemmKernelLauncher(void* D, void const* A, void const* B, void const* input_sf, - void const* weight_sf, float const* global_sf, int m, int n, int k, int batch_count, - tkc::CutlassGemmConfig gemmConfig, char* workspace, const size_t workspaceBytes, cudaStream_t stream, - int* occupancy) -{ - using ElementOutput__ = - typename cutlass::platform::conditional::value, cutlass::half_t, T>::type; - using ElementOutput_ = - typename cutlass::platform::conditional::value, float, - ElementOutput__>::type; - using ElementOutput = - typename cutlass::platform::conditional::value, - cutlass::bfloat16_t, ElementOutput_>::type; - - using MXFP8xMXFP8GemmOperator = - typename DeviceGemmMXFP8xMXFP8GemmSm100::Gemm; - MXFP8xMXFP8GemmOperator gemm; - // Reuse the MXFP8xMXFP4 argument preparation helper -- the argument layout - // is identical (block-scaled SFA/SFB, UE8M0 scales, same problem shape + - // strides). Only the B element type differs. - auto args = prepareGemmArgsSm100(D, A, B, input_sf, weight_sf, global_sf, m, n, k, - batch_count, dim3(CGA_M{}, CGA_N{}, CGA_K{}), MXSMTypeAdapter::Scale); - /* // Check shared memory size; throw when SMEM exceeds */ - int smem_size = int(sizeof(typename MXFP8xMXFP8GemmOperator::GemmKernel::SharedStorage)); - static int mMaxSmemSize = tk::getMaxSharedMemoryPerBlockOptin(); - if (smem_size > mMaxSmemSize) - { - std::string errMsg = "SMEM size exceeds maximum allowed. Required " + std::to_string(smem_size) + ", got " - + std::to_string(mMaxSmemSize); - throw std::runtime_error("[TensorRT LLM Error][MXFP8 gemm Runner] " + errMsg); - } - /* // Return workspace size */ - if (!A && !B && !D) - { - return gemm.get_workspace_size(args); - } - if (gemm.get_workspace_size(args) > workspaceBytes) - { - std::string errMsg("Requested workspace size insufficient. Required " - + std::to_string(gemm.get_workspace_size(args)) + ", got " + std::to_string(workspaceBytes)); - throw std::runtime_error("[TensorRT LLM Error][MXFP8 gemm Runner] " + errMsg); - } - auto can_implement = gemm.can_implement(args); - if (can_implement != cutlass::Status::kSuccess) - { - std::string errMsg = "MXFP8xMXFP8 Gemm cutlass kernel will fail for params. Error: " - + std::string(cutlassGetStatusString(can_implement)); - throw std::runtime_error("[TensorRT LLM Error][MXFP8 gemm Runner] " + errMsg); - } - auto initStatus = gemm.initialize(args, workspace, stream); - if (initStatus != cutlass::Status::kSuccess) - { - std::string errMsg = "Failed to initialize cutlass MXFP8xMXFP8 gemm. Error: " - + std::string(cutlassGetStatusString(initStatus)); - throw std::runtime_error("[TensorRT LLM Error][MXFP8xMXFP8 gemm Runner] " + errMsg); - } - auto runStatus = gemm.run(args, workspace, stream, nullptr, tensorrt_llm::common::getEnvEnablePDL()); - if (runStatus != cutlass::Status::kSuccess) - { - std::string errMsg - = "Failed to run cutlass MXFP8xMXFP8 gemm. Error: " + std::string(cutlassGetStatusString(runStatus)); - throw std::runtime_error("[TensorRT LLM Error][MXFP8xMXFP8 gemm Runner] " + errMsg); - } - return gemm.get_workspace_size(args); -} - -#endif - -} // namespace cutlass_kernels -} // namespace kernels - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_quant_packed.cu b/cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_quant_packed.cu deleted file mode 100644 index a4923fdbd072..000000000000 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_quant_packed.cu +++ /dev/null @@ -1,224 +0,0 @@ -/* - * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. - * - * 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. - */ - -// Fused 1x128 FP8 quantize + UE8M0 scale packing. -// -// Replaces the (scale_1x128_kernel + pack_fp32_into_ue8m0) two-kernel sequence -// used by SM100 deep_gemm fp8 block-scale GEMMs. Adapted from the SM120 MoE -// in-kernel packing pattern (`scale_1x128_kernel_sm120` in -// sm120_blockwise_gemm/sm120_fp8_moe_gemm_1d1d.cuh), specialised for the -// non-MoE case (single contiguous batch, no token offsets). - -#include "fp8_blockscale_quant_packed.h" - -#include "tensorrt_llm/common/config.h" -#include "tensorrt_llm/common/envUtils.h" - -#include -#include -#include -#include - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels::fp8_blockscale_gemm -{ - -namespace -{ - -__device__ __forceinline__ float reciprocal_approximate_ftz_local(float a) -{ - float b; - asm volatile("rcp.approx.ftz.f32 %0, %1;\n" : "=f"(b) : "f"(a)); - return b; -} - -// Each warp consumes one row × 4 quantization blocks (4 × 128 = 512 K elems). -// 32 lanes split into 4 lane-groups of 8: each group covers 1 quant block -// (8 lanes × 16 BF16 elems = 128 elems). After per-block amax, lanes -// 0/8/16/24 each hold one UE8M0 scale byte; lane 0 packs them into a uint32 -// and stores in the deep_gemm-expected MN-major layout. -template -__global__ void fp8_quantize_1x128_packed_kernel_impl(__nv_fp8_e4m3* __restrict__ fp8_output, - int32_t* __restrict__ packed_scale_output, __nv_bfloat16 const* __restrict__ input, int const m, int const k, - int const scale_leading_dim_uint32) -{ - int const packed_sf_k_idx = static_cast(blockIdx.x); - int const warp_id = static_cast(threadIdx.x) >> 5; - int const lane_id = static_cast(threadIdx.x) & 31; - int const m_idx = static_cast(blockIdx.y) * WarpsPerBlock + warp_id; - -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) - cudaGridDependencySynchronize(); -#endif - - // Padded rows (m_idx >= m) fall through and write packed=0 so the caller can - // skip pre-zeroing the scale output buffer. The kernel-end PDL trigger still - // fires for padded warps. - bool const row_in_range = (m_idx < m); - - uint32_t packed = 0u; - if (row_in_range) - { - int const k_base = packed_sf_k_idx * 512 + lane_id * 16; - - // ---- 1. Load 16 BF16 elements per lane. ---- - auto const* in_ptr = reinterpret_cast(input + static_cast(m_idx) * k + k_base); - constexpr int kLoadNumElems = sizeof(double4) / sizeof(__nv_bfloat16); // 16 - - union LoadTrick - { - double4 pack; - __nv_bfloat16 v[kLoadNumElems]; - }; - - LoadTrick load_trick; - bool const k_in_range = (k_base < k); - load_trick.pack = k_in_range ? in_ptr[0] : double4{}; - - if (k_in_range && k_base + kLoadNumElems > k) - { - int const valid = k - k_base; -#pragma unroll - for (int i = 0; i < kLoadNumElems; ++i) - { - if (i >= valid) - { - load_trick.v[i] = __nv_bfloat16(0.0f); - } - } - } - - // ---- 2. Per-block amax (lanes 0..7 / 8..15 / 16..23 / 24..31 = 4 quant blocks). ---- - __nv_bfloat16 max_elem = __nv_bfloat16(0.0f); -#pragma unroll - for (int i = 0; i < kLoadNumElems; ++i) - { - max_elem = __hmax(max_elem, __habs(load_trick.v[i])); - } - float amax = static_cast(max_elem); - amax = fmaxf(amax, __shfl_xor_sync(0xFFFFFFFFu, amax, 4, 8)); - amax = fmaxf(amax, __shfl_xor_sync(0xFFFFFFFFu, amax, 2, 8)); - amax = fmaxf(amax, __shfl_xor_sync(0xFFFFFFFFu, amax, 1, 8)); - amax = fmaxf(amax, 1e-10f); - - // ---- 3. UE8M0 dequant scale. ---- - float const dequant_scale_raw = amax * reciprocal_approximate_ftz_local(448.0f); - __nv_fp8_e8m0 ue8m0_scale; - ue8m0_scale.__x = __nv_cvt_float_to_e8m0(dequant_scale_raw, __NV_SATFINITE, cudaRoundPosInf); - - // Recover quant_scale = 1 / 2^(exp - 127) for fp8 conversion. - constexpr uint32_t FP32_EXPONENT_BIAS = 127u; - float const quant_scale = (ue8m0_scale.__x == 0) - ? 1.0f - : exp2f(static_cast(FP32_EXPONENT_BIAS) - static_cast(ue8m0_scale.__x)); - - // ---- 4. Quantize and store FP8 output. ---- - constexpr int kStoreNumElems = sizeof(float4) / sizeof(__nv_fp8_e4m3); // 16 - - union StoreTrick - { - float4 pack; - __nv_fp8_e4m3 v[kStoreNumElems]; - }; - - StoreTrick store_trick; - store_trick.pack = float4{}; -#pragma unroll - for (int i = 0; i < kStoreNumElems; ++i) - { - store_trick.v[i] = __nv_fp8_e4m3(static_cast(load_trick.v[i]) * quant_scale); - } - auto* out_ptr = reinterpret_cast(fp8_output + static_cast(m_idx) * k + k_base); - if (k_in_range) - { - if (k_base + kStoreNumElems > k) - { - int const valid = k - k_base; -#pragma unroll - for (int i = 0; i < kStoreNumElems; ++i) - { - if (i >= valid) - { - store_trick.v[i] = __nv_fp8_e4m3(0.0f); - } - } - } - out_ptr[0] = store_trick.pack; - } - - // ---- 5. Pack 4 UE8M0 scales (lanes 0/8/16/24). ---- - uint32_t const s0 = __shfl_sync(0xFFFFFFFFu, static_cast(ue8m0_scale.__x), 0); - uint32_t const s1 = __shfl_sync(0xFFFFFFFFu, static_cast(ue8m0_scale.__x), 8); - uint32_t const s2 = __shfl_sync(0xFFFFFFFFu, static_cast(ue8m0_scale.__x), 16); - uint32_t const s3 = __shfl_sync(0xFFFFFFFFu, static_cast(ue8m0_scale.__x), 24); - if (lane_id == 0) - { - // Mask off scale bytes whose sf_k is past the actual K. - int const num_sf_k = (k + 127) / 128; - int const sf_k_base = packed_sf_k_idx * 4; - if (sf_k_base + 0 < num_sf_k) - packed |= s0; - if (sf_k_base + 1 < num_sf_k) - packed |= (s1 << 8); - if (sf_k_base + 2 < num_sf_k) - packed |= (s2 << 16); - if (sf_k_base + 3 < num_sf_k) - packed |= (s3 << 24); - } - } - - // Always write the packed scale — `packed` is 0 for padded rows. The grid - // covers the full [0, scale_leading_dim_uint32) leading dim (rounded up to - // WarpsPerBlock), and the m_idx guard drops the few rows past the buffer end. - if (lane_id == 0 && m_idx < scale_leading_dim_uint32) - { - packed_scale_output[static_cast(packed_sf_k_idx) * scale_leading_dim_uint32 + m_idx] = packed; - } - -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) - cudaTriggerProgrammaticLaunchCompletion(); -#endif -} - -} // namespace - -void launch_fp8_quantize_1x128_packed_bf16_e4m3(__nv_fp8_e4m3* fp8_output, int32_t* packed_scale_output, - __nv_bfloat16 const* input, int m, int k, int scale_leading_dim_uint32, cudaStream_t stream) -{ - if (m <= 0 || k <= 0) - { - return; - } - - constexpr int kWarpsPerBlock = 4; - int const num_packed_sf_k = (((k + 127) / 128) + 3) / 4; - // Cover the full TMA-padded leading dim so the entire [0, scale_leading_dim_uint32) - // extent of packed_scale_output is written (in-kernel zero for rows past `m`), - // regardless of how the caller padded the input. - int const m_blocks = (scale_leading_dim_uint32 + kWarpsPerBlock - 1) / kWarpsPerBlock; - dim3 const grid(num_packed_sf_k, m_blocks, 1); - dim3 const block(kWarpsPerBlock * 32, 1, 1); - - tensorrt_llm::common::launchWithPdlWhenEnabled("fp8_quantize_1x128_packed_kernel_impl", - fp8_quantize_1x128_packed_kernel_impl, grid, block, 0, stream, fp8_output, packed_scale_output, - input, m, k, scale_leading_dim_uint32); -} - -} // namespace kernels::fp8_blockscale_gemm - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_quant_packed.h b/cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_quant_packed.h deleted file mode 100644 index a4079cd9b054..000000000000 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_quant_packed.h +++ /dev/null @@ -1,52 +0,0 @@ -/* - * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. - * - * 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. - */ - -// Host-callable launcher for the fused FP8 1x128 quantize + UE8M0-pack kernel. -// Kernel implementation lives in fp8_blockscale_quant_packed.cu and is built -// by nvcc; this header is safe to include from .cpp files compiled by g++. - -#pragma once - -#include "tensorrt_llm/common/config.h" - -#include -#include -#include -#include - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels::fp8_blockscale_gemm -{ - -// Launches the fused 1x128 FP8 quant + UE8M0 pack kernel. -// -// Inputs: -// input : BF16 [m, k] row-major contiguous -// Outputs: -// fp8_output : E4M3 [m, k] row-major contiguous -// packed_scale_output : uint32 [packed_sf_k, scale_leading_dim_uint32] -// where packed_sf_k = ceil(ceil(k/128)/4) -// -// `scale_leading_dim_uint32` is the (uint32) stride between consecutive -// packed_sf_k rows of the scale tensor; caller is responsible for choosing -// it (typically aligned to 4 uint32 = 16 bytes for TMA alignment). -void launch_fp8_quantize_1x128_packed_bf16_e4m3(__nv_fp8_e4m3* fp8_output, int32_t* packed_scale_output, - __nv_bfloat16 const* input, int m, int k, int scale_leading_dim_uint32, cudaStream_t stream); - -} // namespace kernels::fp8_blockscale_gemm - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/include/fp4_gemm.h b/cpp/tensorrt_llm/kernels/cutlass_kernels/include/fp4_gemm.h index 8ca27dc55d57..b0e8a307df79 100644 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/include/fp4_gemm.h +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/include/fp4_gemm.h @@ -67,8 +67,6 @@ enum class FP4GemmType { W4A4_NVFP4_NVFP4, W4A8_MXFP4_MXFP8, - // W8A8 MXFP8 weight (e4m3 + UE8M0 1x32 block scales) x dynamic MXFP8 activation. - W8A8_MXFP8_MXFP8, }; template diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_gemm_kernels.h b/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_gemm_kernels.h index dca53bd071d2..a2b7c112bd9e 100644 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_gemm_kernels.h +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_gemm_kernels.h @@ -300,12 +300,10 @@ class MoeGemmRunner void moeGemm(GroupedGemmInput inputs, TmaWarpSpecializedGroupedGemmInput hopper_inputs); - std::vector getConfigs( - bool supports_finalize_fusion, bool use_mxfp8 = false) const; - static std::vector getConfigs( - int sm, bool supports_finalize_fusion, bool use_mxfp8 = false); + std::vector getConfigs(bool supports_finalize_fusion) const; + static std::vector getConfigs(int sm, bool supports_finalize_fusion); static std::vector getTmaWarpSpecializedConfigs( - int sm, bool supports_finalize_fusion, bool use_mxfp8 = false); + int sm, bool supports_finalize_fusion); static std::vector getAmpereConfigs(int sm); [[nodiscard]] bool isTmaWarpSpecialized(cutlass_extensions::CutlassGemmConfig gemm_config) const; @@ -320,7 +318,7 @@ class MoeGemmRunner ActivationType activation_type, int gemm_n, int gemm_k) const; [[nodiscard]] bool supportsFusedGatedActivation(ActivationType activation_type, int gemm_n, int gemm_k) const; - size_t getMaxWorkspaceSize(int num_experts, bool use_mxfp8_weight_scaling = false) const; + size_t getMaxWorkspaceSize(int num_experts) const; [[nodiscard]] int getSM() const; @@ -337,9 +335,8 @@ class MoeGemmRunner int sm_{}; int multi_processor_count_{}; mutable int num_experts_ = 0; - mutable bool use_mxfp8_weight_scaling_ = false; mutable size_t gemm_workspace_size_ = 0; - size_t calcMaxWorkspaceSize(int num_experts, bool use_mxfp8_weight_scaling) const; + size_t calcMaxWorkspaceSize(int num_experts) const; }; } // namespace kernels::cutlass_kernels diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_kernels.h b/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_kernels.h index 24781bec76e7..216877a4ffc7 100644 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_kernels.h +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_kernels.h @@ -21,7 +21,7 @@ #include "tensorrt_llm/common/cudaUtils.h" #include "tensorrt_llm/common/quantization.h" #include "tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_gemm.h" -#include "tensorrt_llm/kernels/cutlass_kernels/include/moe_lora_grouped_gemm.h" +#include "tensorrt_llm/kernels/cutlass_kernels/include/moe_lora_device_path.h" #include #ifdef ENABLE_FP4 #include @@ -69,14 +69,13 @@ struct LoraParams cudaEvent_t* memcpy_event_ptr; - // Capture-safe grouped-GEMM LoRA core scratch. When grouped_gemm.enabled is + // Device-side capture-safe LoRA path scratch. When device_path.enabled is // true, the kernel uses launchMoeLoraPointerExpand, launchMoeLoraProblemBuilder, - // and cudaGraph(SplitK)GroupedGemm instead of the host-pointer LoraImpl::run - // path. The pointers refer to persistent allocations owned by the calling - // FusedMoeRunner, so their addresses are stable across CUDA-graph captures - // and replays. Default-constructed (enabled == false) for the legacy - // TensorRT MoE plugin, which uses its own cuBLAS LoraImpl path. - ::tensorrt_llm::kernels::cutlass_kernels::MoeLoraGroupedGemm grouped_gemm; + // and cudaGraph(SplitK)GroupedGemm instead of the legacy host-pointer + // LoraImpl::run path. The pointers refer to persistent allocations owned by + // the calling FusedMoeRunner, so their addresses are stable across + // CUDA-graph captures and replays. + ::tensorrt_llm::kernels::cutlass_kernels::MoeLoraDevicePath device_path; LoraParams() = default; @@ -302,22 +301,6 @@ struct QuantParams GemmInputs fc2; } mxfp8_mxfp4; - // MXFP8 x MXFP8 quantization params (W8A8 with UE8M0 1x32 block scales on both - // sides). No per-tensor / global alpha (block scales determine output magnitude). - // Kept as a separate slot from mxfp8_mxfp4 so consumers can disambiguate - // B element bitwidth (4-bit vs 8-bit) without relying on naming aliases. - struct MXFP8MXFP8Inputs - { - struct GemmInputs - { - TmaWarpSpecializedGroupedGemmInput::MXFPXElementSF const* weight_block_scale - = nullptr; // (experts, n, k / 32) - }; - - GemmInputs fc1; - GemmInputs fc2; - } mxfp8_mxfp8; - // FP4 quantization params struct FP4Inputs { @@ -421,36 +404,6 @@ struct QuantParams return qp; } - // MXFP8xMXFP8 grouped MoE: e4m3 acts + e4m3 weights, UE8M0 1x32 block - // scales on both sides. No per-tensor / global alpha is required (block - // scales determine output magnitude). Writes to its own dedicated slot - // so consumers can distinguish from MXFP8xMXFP4 (which has 4-bit B). - static QuantParams MXFP8MXFP8(TmaWarpSpecializedGroupedGemmInput::MXFPXElementSF const* fc1_weight_block_scale, - TmaWarpSpecializedGroupedGemmInput::MXFPXElementSF const* fc2_weight_block_scale) - { - QuantParams qp; - qp.mxfp8_mxfp8.fc1 = {fc1_weight_block_scale}; - qp.mxfp8_mxfp8.fc2 = {fc2_weight_block_scale}; - return qp; - } - - // Helpers: return the active MXFPX activation-side block-scale pointer - // regardless of whether B is fp4 (mxfp8_mxfp4) or fp8 (mxfp8_mxfp8). - // Used by consumers that only care "is the activation path block-scaled - // MXFPX" — they don't need to know B's bitwidth, that's already encoded - // in the kernel template instantiation. - TmaWarpSpecializedGroupedGemmInput::MXFPXElementSF const* mxfpxActFc1WeightScale() const - { - return mxfp8_mxfp8.fc1.weight_block_scale ? mxfp8_mxfp8.fc1.weight_block_scale - : mxfp8_mxfp4.fc1.weight_block_scale; - } - - TmaWarpSpecializedGroupedGemmInput::MXFPXElementSF const* mxfpxActFc2WeightScale() const - { - return mxfp8_mxfp8.fc2.weight_block_scale ? mxfp8_mxfp8.fc2.weight_block_scale - : mxfp8_mxfp4.fc2.weight_block_scale; - } - static QuantParams FP4(float const* fc1_act_global_scale, TmaWarpSpecializedGroupedGemmInput::NVFP4ElementSF const* fc1_weight_block_scale, float const* fc1_global_scale, // @@ -598,12 +551,6 @@ class CutlassMoeFCRunnerInterface bool is_profiler = false; bool use_fused_finalize_ = true; - // When the activation/weight pair is , this flag selects - // between the per-tensor FP8 path (false, default) and the MXFP8xMXFP8 - // block-scaled path (true). It is read by getScalingType() in subclasses - // to choose the runtime FpXBlockScalingType. Ignored for other type - // combinations -- their scaling type is fully determined at compile time. - bool use_mxfp8_weight_scaling_ = false; }; // Assumes inputs activations are row major. Weights need to be preprocessed by th_op/weight_quantize.cc . @@ -686,11 +633,7 @@ class CutlassMoeFCRunner : public CutlassMoeFCRunnerInterface std::vector getTactics(MoeGemmId gemm_id) override { - // Pass `use_mxfp8_weight_scaling_` so MXFP8xMXFP8 enumerates only the - // Mxf8f6f4-valid tile shapes; otherwise autotuning would invoke FP8 - // tile shapes that the runtime dispatcher rejects with TLLM_THROW. - return moe_gemm_runner_.getConfigs( - gemm_id == MoeGemmId::GEMM_2 && mayHaveFinalizeFused(), use_mxfp8_weight_scaling_); + return moe_gemm_runner_.getConfigs(gemm_id == MoeGemmId::GEMM_2 && mayHaveFinalizeFused()); } static std::vector getTactics(int sm, MoeGemmId gemm_id) @@ -798,7 +741,7 @@ class CutlassMoeFCRunner : public CutlassMoeFCRunnerInterface virtual size_t getGemmWorkspaceSize(int num_experts_per_node) const override { - return moe_gemm_runner_.getMaxWorkspaceSize(num_experts_per_node, use_mxfp8_weight_scaling_); + return moe_gemm_runner_.getMaxWorkspaceSize(num_experts_per_node); } std::pair @@ -820,7 +763,7 @@ class CutlassMoeFCRunner : public CutlassMoeFCRunnerInterface reinterpret_cast(bias1), reinterpret_cast(bias2), reinterpret_cast(gemm1_output), reinterpret_cast(gemm2_output), router_scales, permuted_row_to_unpermuted_row, - Self::getScalingType(use_mxfp8_weight_scaling_), stream); + stream); } std::pair @@ -863,7 +806,7 @@ class CutlassMoeFCRunner : public CutlassMoeFCRunnerInterface TmaWarpSpecializedGroupedGemmInput::ElementSF const* fp4_act_flat2, QuantParams quant_params, ScaleBiasType const* bias1, ScaleBiasType const* bias2, UnfusedGemmOutputType* gemm1_output, UnfusedGemmOutputType* gemm2_output, float const* router_scales, int const* permuted_row_to_unpermuted_row, - TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType scaling_type, cudaStream_t stream); + cudaStream_t stream); static std::pair computeStridesTmaWarpSpecializedLowLatency(TmaWarpSpecializedGroupedGemmInput layout_info1, TmaWarpSpecializedGroupedGemmInput layout_info2, int64_t num_tokens, int64_t gemm1_n, int64_t gemm1_k, @@ -913,29 +856,12 @@ class CutlassMoeFCRunner : public CutlassMoeFCRunnerInterface return RunnerType::supportsTmaWarpSpecialized(sm) && sm >= 90 && !use_wfp4a16; } - // TODO: This should eventually take the full quant params to give more - // flexibility. For now the only runtime selector is the MXFP8 flag for - // the instantiation (per-tensor FP8 vs MXFP8 block-scaled). - static auto getScalingType(bool use_mxfp8_weight_scaling) + // TODO: This should eventually take the quant params to give more flexibility + static auto getScalingType() { - if constexpr (use_wfp4afp8) - { - return TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::MXFPX; - } - else if constexpr (use_fp4) - { - return TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::NVFP4; - } - else if constexpr (use_fp8 && std::is_same_v) - { - // : per-tensor FP8 (NONE) or MXFP8 block-scaled (MXFPX). - return use_mxfp8_weight_scaling ? TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::MXFPX - : TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::NONE; - } - else - { - return TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::NONE; - } + return use_wfp4afp8 ? TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::MXFPX + : use_fp4 ? TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::NVFP4 + : TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::NONE; } bool setupLoraWorkspace(int64_t expanded_num_rows, int64_t num_rows, int64_t inter_size, int64_t hidden_size, @@ -1036,7 +962,7 @@ struct GemmProfilerBackend nvinfer1::DataType wtype, nvinfer1::DataType otype, int num_experts, int k, int64_t hidden_size, int64_t unpadded_hidden_size, int64_t inter_size, int64_t group_size, ActivationType activation_type, bool bias, bool use_lora, bool min_latency_mode, bool need_weights, MOEParallelismConfig parallelism_config, - bool const enable_alltoall, bool use_mxfp8_weight_scaling = false) + bool const enable_alltoall) { mInterface = &runner; mGemmToProfile = gemm_to_profile; @@ -1057,7 +983,6 @@ struct GemmProfilerBackend mNeedWeights = need_weights; mParallelismConfig = parallelism_config; mEnableAlltoall = enable_alltoall; - mUseMxfp8WeightScaling = use_mxfp8_weight_scaling; mSM = common::getSMVersion(); mScalingType = TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::NONE; @@ -1066,13 +991,6 @@ struct GemmProfilerBackend { mScalingType = TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::MXFPX; } - else if (dtype == nvinfer1::DataType::kFP8 && wtype == nvinfer1::DataType::kFP8 && use_mxfp8_weight_scaling) - { - // MXFP8 W8A8: e4m3 acts × e4m3 weights with UE8M0 1x32 block scales on both sides. - // Profiler must produce MXFPX block-scaled inputs (otherwise the per-expert SF - // pointer arrays stay uninitialized and the kernel reads garbage SF addresses). - mScalingType = TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::MXFPX; - } else if ((dtype == nvinfer1::DataType::kFP4 || dtype == nvinfer1::DataType::kINT64) && (wtype == nvinfer1::DataType::kFP4 || wtype == nvinfer1::DataType::kINT64)) { @@ -1102,7 +1020,6 @@ struct GemmProfilerBackend ActivationType mActivationType{}; MOEParallelismConfig mParallelismConfig{}; bool mEnableAlltoall = false; - bool mUseMxfp8WeightScaling = false; int mSampleIndex = 0; diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_lora_device_path.h b/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_lora_device_path.h new file mode 100644 index 000000000000..8214e84e3b7f --- /dev/null +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_lora_device_path.h @@ -0,0 +1,163 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + * + * 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. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" + +#include + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::cutlass_kernels +{ + +// Forward declaration; the typedef below references it by name. +struct MoeLoraDevicePathModule; + +// Function-pointer dispatch for the libtorch-dependent GEMM stage of the MoE +// LoRA device path. The implementation lives in th_common (moeOp.cpp) because +// the underlying cudaGraph(SplitK)GroupedGemm wrappers allocate workspace via +// at::Tensor, which is not linkable from libmoe_gemm_src.a (that archive ends +// up inside the TensorRT plugin shared object, which deliberately does not +// depend on libtorch). +// +// Contract: +// mod: per-module device-resident scratch produced by +// launchMoeLoraPointerExpand. The implementation repacks +// it into a MoeLoraGemmGroupArrays, calls the problem +// builder, and dispatches the in/out GEMMs. +// num_permuted_tokens: length of the permuted-row tables in mod. +// in_hidden_size: K of the in-GEMM (the input row stride). +// max_lora_rank: workspace ldd and worst-case rank cap. +// dtype_bytes: sizeof(scalar) for the LoRA tensors. +// splitk_slices: split-K factor for the in-GEMM. +// input_base: base of the input matrix (M rows of in_hidden_size). +// output_base: base of the module's output buffer; the implementation +// accumulates into it, so callers must initialize it. +// data_type: scalar dtype expected by the GEMM wrappers +// (fp16/bf16/fp32). +// stream: CUDA stream to launch onto. +using MoeLoraDeviceRunFn = void (*)(MoeLoraDevicePathModule const& mod, int64_t num_permuted_tokens, + int64_t in_hidden_size, int64_t max_lora_rank, int64_t dtype_bytes, int64_t splitk_slices, void const* input_base, + void* output_base, nvinfer1::DataType data_type, cudaStream_t stream); + +// Per-module device-resident scratch for the MoE LoRA capture-safe path. +// Pointers refer to device memory unless noted. +// +// The struct is typed with void* rather than the concrete +// cutlass::gemm::GemmCoord* / int64_t* types so this header can be included +// from moe_kernels.h without dragging in cutlass headers. The concrete types +// are recovered at the call site (matching the contract documented in +// moe_lora_problem_builder.h): +// +// problem_sizes_* -> cutlass::gemm::GemmCoord* (device, [P_max]) +// a_ptrs_*/b/d -> void** (device, [P_max]) +// lda/ldb/ldd_* -> int64_t* (device, [P_max]) +// splitk_offsets -> int64_t* (device, [P_max + 1]) +// lowrank_ws_dev -> void* (device, [P_max, max_lora_rank, dtype_bytes]) +// host_max_* -> cutlass::gemm::GemmCoord* (pinned host, [1]) +// +// The split-K in-GEMM's partial-sum scratch is allocated internally by the +// cuda_graph_split_k_grouped_gemm wrapper (sized from the host max-problem +// hint); only the per-problem splitk_offsets are produced here. +// +// out_hidden_size is the trailing dimension of the module's output buffer; it +// is inter_size for fc1/gated and hidden_size for fc2. The output base address +// itself is passed directly to runMoeLoraDeviceModule at the call site. +struct MoeLoraDevicePathModule +{ + // Per-source-token (rank, A_ptr, B_ptr) device mirrors, staged via a + // pinned-host to device async H2D in FusedMoeRunner::buildMoeLoraParams. + // These feed launchMoeLoraPointerExpand as ranks_src / ptrs_src. + int32_t const* ranks_src_dev = nullptr; + int64_t const* ptrs_src_dev = nullptr; + + // Inner (A) and outer (B) dimensions for this module, fed to the + // pointer-expand kernel as dim_a / dim_b so it can compute the per-expert + // offset weight_index * dim * lora_rank. For fc1/gated this is + // (hidden_size, inter_size); for fc2 it is (inter_size, hidden_size). + int64_t dim_a = 0; + int64_t dim_b = 0; + + // Per-permuted-row (rank, A_ptr + offset, B_ptr + offset). + int32_t* permuted_ranks_dev = nullptr; + int64_t* permuted_ptrs_dev = nullptr; + + // cuda_graph_(split_k_)grouped_gemm-ready bundle. + void* problem_sizes_in_dev = nullptr; + void* problem_sizes_out_dev = nullptr; + void** a_ptrs_in_dev = nullptr; + void** b_ptrs_in_dev = nullptr; + void** d_ptrs_in_dev = nullptr; + void** b_ptrs_out_dev = nullptr; + void** d_ptrs_out_dev = nullptr; + int64_t* lda_in_dev = nullptr; + int64_t* ldb_in_dev = nullptr; + int64_t* ldd_in_dev = nullptr; + int64_t* ldb_out_dev = nullptr; + int64_t* ldd_out_dev = nullptr; + int64_t* splitk_offsets_dev = nullptr; + + // Low-rank intermediate workspace shared between the in- and out-GEMM. The + // split-K partial-sum scratch is owned by the GEMM wrapper, not here. + void* lowrank_workspace_dev = nullptr; + + // Host (pinned) per-call max problem size hints, required by the + // cuda_graph_*_grouped_gemm wrappers for kernel selection. The + // values are upper bounds (max_M, max_N, max_K) safe to fix at + // warmup time. + void* host_max_problem_in_pinned = nullptr; + void* host_max_problem_out_pinned = nullptr; + + // Trailing dimension of the module's output buffer (inter_size for + // fc1/gated, hidden_size for fc2). The output base address is supplied + // directly to runMoeLoraDeviceModule at the call site. + int64_t out_hidden_size = 0; +}; + +// Top-level device-path bundle attached to LoraParams when the device LoRA +// path is active. enabled == false means the FusedMoeRunner runs the legacy +// host path. +struct MoeLoraDevicePath +{ + bool enabled = false; + + // Scalars common to all three modules. Fixed for the lifetime of the + // FusedMoeRunner once the scratch is allocated. + int64_t in_hidden_size = 0; + int64_t max_lora_rank = 0; + int64_t dtype_bytes = 0; + int64_t splitk_slices = 0; + + bool has_gated = false; + + // libtorch-bound GEMM dispatch entry point, populated by moeOp.cpp when the + // device path is enabled. nullptr means the device path is unavailable from + // this consumer (for example, the TensorRT plugin). + MoeLoraDeviceRunFn run = nullptr; + + MoeLoraDevicePathModule fc1; + MoeLoraDevicePathModule fc2; + MoeLoraDevicePathModule gated; +}; + +} // namespace kernels::cutlass_kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_lora_grouped_gemm.h b/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_lora_grouped_gemm.h deleted file mode 100644 index 55ab4e40a3ae..000000000000 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_lora_grouped_gemm.h +++ /dev/null @@ -1,164 +0,0 @@ -/* - * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. - * - * 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. - */ - -#pragma once - -#include "tensorrt_llm/common/config.h" - -#include - -#include -#include - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels::cutlass_kernels -{ - -// Forward declaration; the typedef below references it by name. -struct MoeLoraGroupedGemmModule; - -// Function-pointer dispatch for the libtorch-dependent GEMM stage of the MoE -// LoRA grouped-GEMM core. The implementation lives in th_common (moeOp.cpp) because -// the underlying cudaGraph(SplitK)GroupedGemm wrappers allocate workspace via -// at::Tensor, which is not linkable from libmoe_gemm_src.a (that archive ends -// up inside the TensorRT plugin shared object, which deliberately does not -// depend on libtorch). -// -// Contract: -// mod: per-module device-resident scratch produced by -// launchMoeLoraPointerExpand. The implementation repacks -// it into a MoeLoraGemmGroupArrays, calls the problem -// builder, and dispatches the in/out GEMMs. -// num_permuted_tokens: length of the permuted-row tables in mod. -// in_hidden_size: K of the in-GEMM (the input row stride). -// max_lora_rank: workspace ldd and worst-case rank cap. -// dtype_bytes: sizeof(scalar) for the LoRA tensors. -// splitk_slices: split-K factor for the in-GEMM. -// input_base: base of the input matrix (M rows of in_hidden_size). -// output_base: base of the module's output buffer; the implementation -// accumulates into it, so callers must initialize it. -// data_type: scalar dtype expected by the GEMM wrappers -// (fp16/bf16/fp32). -// stream: CUDA stream to launch onto. -using MoeLoraGroupedGemmRunFn = void (*)(MoeLoraGroupedGemmModule const& mod, int64_t num_permuted_tokens, - int64_t in_hidden_size, int64_t max_lora_rank, int64_t dtype_bytes, int64_t splitk_slices, void const* input_base, - void* output_base, nvinfer1::DataType data_type, cudaStream_t stream); - -// Per-module device-resident scratch for the MoE LoRA capture-safe path. -// Pointers refer to device memory unless noted. -// -// The struct is typed with void* rather than the concrete -// cutlass::gemm::GemmCoord* / int64_t* types so this header can be included -// from moe_kernels.h without dragging in cutlass headers. The concrete types -// are recovered at the call site (matching the contract documented in -// moe_lora_problem_builder.h): -// -// problem_sizes_* -> cutlass::gemm::GemmCoord* (device, [P_max]) -// a_ptrs_*/b/d -> void** (device, [P_max]) -// lda/ldb/ldd_* -> int64_t* (device, [P_max]) -// splitk_offsets -> int64_t* (device, [P_max + 1]) -// lowrank_ws_dev -> void* (device, [P_max, max_lora_rank, dtype_bytes]) -// host_max_* -> cutlass::gemm::GemmCoord* (pinned host, [1]) -// -// The split-K in-GEMM's partial-sum scratch is allocated internally by the -// cuda_graph_split_k_grouped_gemm wrapper (sized from the host max-problem -// hint); only the per-problem splitk_offsets are produced here. -// -// out_hidden_size is the trailing dimension of the module's output buffer; it -// is inter_size for fc1/gated and hidden_size for fc2. The output base address -// itself is passed directly to runMoeLoraGroupedGemmModule at the call site. -struct MoeLoraGroupedGemmModule -{ - // Per-source-token (rank, A_ptr, B_ptr) device mirrors, staged via a - // pinned-host to device async H2D in FusedMoeRunner::buildMoeLoraParams. - // These feed launchMoeLoraPointerExpand as ranks_src / ptrs_src. - int32_t const* ranks_src_dev = nullptr; - int64_t const* ptrs_src_dev = nullptr; - - // Inner (A) and outer (B) dimensions for this module, fed to the - // pointer-expand kernel as dim_a / dim_b so it can compute the per-expert - // offset weight_index * dim * lora_rank. For fc1/gated this is - // (hidden_size, inter_size); for fc2 it is (inter_size, hidden_size). - int64_t dim_a = 0; - int64_t dim_b = 0; - - // Per-permuted-row (rank, A_ptr + offset, B_ptr + offset). - int32_t* permuted_ranks_dev = nullptr; - int64_t* permuted_ptrs_dev = nullptr; - - // cuda_graph_(split_k_)grouped_gemm-ready bundle. - void* problem_sizes_in_dev = nullptr; - void* problem_sizes_out_dev = nullptr; - void** a_ptrs_in_dev = nullptr; - void** b_ptrs_in_dev = nullptr; - void** d_ptrs_in_dev = nullptr; - void** b_ptrs_out_dev = nullptr; - void** d_ptrs_out_dev = nullptr; - int64_t* lda_in_dev = nullptr; - int64_t* ldb_in_dev = nullptr; - int64_t* ldd_in_dev = nullptr; - int64_t* ldb_out_dev = nullptr; - int64_t* ldd_out_dev = nullptr; - int64_t* splitk_offsets_dev = nullptr; - - // Low-rank intermediate workspace shared between the in- and out-GEMM. The - // split-K partial-sum scratch is owned by the GEMM wrapper, not here. - void* lowrank_workspace_dev = nullptr; - - // Host (pinned) per-call max problem size hints, required by the - // cuda_graph_*_grouped_gemm wrappers for kernel selection. The - // values are upper bounds (max_M, max_N, max_K) safe to fix at - // warmup time. - void* host_max_problem_in_pinned = nullptr; - void* host_max_problem_out_pinned = nullptr; - - // Trailing dimension of the module's output buffer (inter_size for - // fc1/gated, hidden_size for fc2). The output base address is supplied - // directly to runMoeLoraGroupedGemmModule at the call site. - int64_t out_hidden_size = 0; -}; - -// Top-level grouped-GEMM bundle attached to LoraParams when the MoE op runs the -// capture-safe grouped-GEMM LoRA core. enabled == false is the default and means -// this consumer does not use the core (the legacy TensorRT MoE plugin leaves it -// false and runs its own cuBLAS LoRA path). -struct MoeLoraGroupedGemm -{ - bool enabled = false; - - // Scalars common to all three modules. Fixed for the lifetime of the - // FusedMoeRunner once the scratch is allocated. - int64_t in_hidden_size = 0; - int64_t max_lora_rank = 0; - int64_t dtype_bytes = 0; - int64_t splitk_slices = 0; - - bool has_gated = false; - - // libtorch-bound GEMM dispatch entry point, populated by moeOp.cpp when the - // grouped-GEMM core is enabled. nullptr means the core is unavailable from - // this consumer (for example, the TensorRT plugin). - MoeLoraGroupedGemmRunFn run = nullptr; - - MoeLoraGroupedGemmModule fc1; - MoeLoraGroupedGemmModule fc2; - MoeLoraGroupedGemmModule gated; -}; - -} // namespace kernels::cutlass_kernels - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/launchers/moe_gemm_tma_ws_launcher.inl b/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/launchers/moe_gemm_tma_ws_launcher.inl index 0044528b4dff..a8f4e71bebad 100644 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/launchers/moe_gemm_tma_ws_launcher.inl +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/launchers/moe_gemm_tma_ws_launcher.inl @@ -264,18 +264,11 @@ using namespace cutlass::epilogue; \ constexpr static bool IsFP8 = cutlass::platform::is_same::value; \ \ - /* MXFP8 weight x MXFP8 activation (W8A8): both A and B are e4m3 with UE8M0 1x32 block scales. */ \ - /* On SM100/103 the existing block-scaled infrastructure uses `cute::tuple` for both A/B, */ \ - /* so once `IsBlockScaled` is true and `IsMXFPX` is true the collective_builder lowers to the same */ \ - /* Mxf8f6f4 tensor-op family as MXFP4xMXFP8 (the B element type just widens from fp4 to fp8). */ \ - constexpr static bool IsWMXFP8AMXFP8 \ - = IsFP8 && cutlass::platform::is_same::value && IsMXFPX; \ - \ /* TODO Update once mixed input support is added */ \ static_assert(cutlass::platform::is_same::value || IsWFP4AFP8, \ "TMA warp specialized MOE implementation does not support mixed input types"); \ \ - constexpr static bool IsBlockScaled = IsFP4 || IsWFP4AFP8 || IsWMXFP8AMXFP8; \ + constexpr static bool IsBlockScaled = IsFP4 || IsWFP4AFP8; \ static_assert(!IsBlockScaled || IsBlackwell, "Block scaled is only implemented for SM100"); \ \ static_assert(FUSION == EpilogueFusion::NONE || FUSION == EpilogueFusion::FINALIZE, \ diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_template_dispatch.h b/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_template_dispatch.h index f65e4cd1911c..2fcd54bd1345 100644 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_template_dispatch.h +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_template_dispatch.h @@ -483,17 +483,17 @@ namespace kernels::cutlass_kernels template std::vector MoeGemmRunner::getConfigs( - bool supports_finalize_fusion, bool use_mxfp8) const + bool supports_finalize_fusion) const { - return getConfigs(sm_, supports_finalize_fusion, use_mxfp8); + return getConfigs(sm_, supports_finalize_fusion); } template std::vector MoeGemmRunner::getConfigs( - int sm, bool supports_finalize_fusion, bool use_mxfp8) + int sm, bool supports_finalize_fusion) { std::vector candidate_configs - = getTmaWarpSpecializedConfigs(sm, supports_finalize_fusion, use_mxfp8); + = getTmaWarpSpecializedConfigs(sm, supports_finalize_fusion); std::vector ampere_configs = getAmpereConfigs(sm); std::copy(ampere_configs.begin(), ampere_configs.end(), std::back_inserter(candidate_configs)); return candidate_configs; @@ -530,7 +530,7 @@ MoeGemmRunner::getAmpereConfigs(int sm template std::vector MoeGemmRunner::getTmaWarpSpecializedConfigs( - int sm, bool supports_finalize_fusion, bool use_mxfp8) + int sm, bool supports_finalize_fusion) { using tensorrt_llm::cutlass_extensions::CutlassGemmConfig; static constexpr auto weight_only_flag @@ -545,16 +545,8 @@ MoeGemmRunner::getTmaWarpSpecializedCo static constexpr auto fp4_only_flag = (use_fp4 || use_wfp4afp8) ? CutlassGemmConfig::FP4_ONLY : CutlassGemmConfig::NONE; static constexpr auto fp8fp4_mixed_flag = use_wfp4afp8 ? CutlassGemmConfig::FP8FP4_MIXED : CutlassGemmConfig::NONE; - // MXFP8xMXFP8 only applies to ; for other type pairs the flag is ignored. -#if defined(ENABLE_FP8) - static constexpr bool is_wfp8afp8 = std::is_same_v && std::is_same_v; -#else - static constexpr bool is_wfp8afp8 = false; -#endif - int const mxfp8_flag = (use_mxfp8 && is_wfp8afp8) ? CutlassGemmConfig::MXFP8_MXFP8 : CutlassGemmConfig::NONE; - auto config_type_param - = static_cast(weight_only_flag | simt_only_flag | grouped_gemm_flag - | enable_blackwell | enable_hopper | fp8_only_flag | fp4_only_flag | fp8fp4_mixed_flag | mxfp8_flag); + auto config_type_param = static_cast(weight_only_flag | simt_only_flag + | grouped_gemm_flag | enable_blackwell | enable_hopper | fp8_only_flag | fp4_only_flag | fp8fp4_mixed_flag); TLLM_CHECK_WITH_INFO(!(enable_blackwell && enable_hopper), "Blackwell and hopper flags are mutually exclusive"); sm = use_wfp4afp8 && sm == 103 ? 100 : sm; @@ -768,48 +760,19 @@ void MoeGemmRunner::dispatchToArch( TLLM_CHECK_WITH_INFO( hopper_inputs.isValid(), "Calling TMA warp specialized configuration with invalid hopper config"); - // Select the appropriate fusion function. For we - // pick IsMXFPX at runtime from fpX_block_scaling_type so the - // dispatcher chain stays a single, type-driven template chain. - constexpr bool is_wfp4afp8 - = std::is_same_v && std::is_same_v; - constexpr bool is_wfp8afp8 - = std::is_same_v && std::is_same_v; - bool const use_mxfp8 = is_wfp8afp8 - && hopper_inputs.fpX_block_scaling_type - == TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::MXFPX; - // Pick the IsMXFPX template parameter for a given FUSION, factoring out the duplicated - // is_wfp4afp8 / is_wfp8afp8 / else chain. C++17-compatible via an integral_constant tag. - auto select_mxfpx_mode = [&](auto fusion_tag) - { - constexpr auto FUSION = decltype(fusion_tag)::value; - if constexpr (is_wfp4afp8) - { - return &cutlass_kernels_oss::dispatchMoeGemmSelectTileShapeTmaWarpSpecialized; - } - else if constexpr (is_wfp8afp8) - { - return use_mxfp8 ? &cutlass_kernels_oss::dispatchMoeGemmSelectTileShapeTmaWarpSpecialized - : &cutlass_kernels_oss::dispatchMoeGemmSelectTileShapeTmaWarpSpecialized; - } - else - { - return &cutlass_kernels_oss::dispatchMoeGemmSelectTileShapeTmaWarpSpecialized; - } - }; + // Select the appropriate fusion function auto select_function = [&]() { - using Fusion = TmaWarpSpecializedGroupedGemmInput::EpilogueFusion; switch (hopper_inputs.fusion) { - case Fusion::FINALIZE: return select_mxfpx_mode(std::integral_constant{}); - case Fusion::NONE: return select_mxfpx_mode(std::integral_constant{}); - case Fusion::ACTIVATION: - case Fusion::GATED_ACTIVATION: + case TmaWarpSpecializedGroupedGemmInput::EpilogueFusion::FINALIZE: + return &cutlass_kernels_oss::dispatchMoeGemmSelectTileShapeTmaWarpSpecialized; + case TmaWarpSpecializedGroupedGemmInput::EpilogueFusion::NONE: + return &cutlass_kernels_oss::dispatchMoeGemmSelectTileShapeTmaWarpSpecialized; + case TmaWarpSpecializedGroupedGemmInput::EpilogueFusion::ACTIVATION: + case TmaWarpSpecializedGroupedGemmInput::EpilogueFusion::GATED_ACTIVATION: default: TLLM_THROW("Unimplemented fusion %d requested", (int) hopper_inputs.fusion); }; }; @@ -910,24 +873,19 @@ void MoeGemmRunner::dispatchToArch( } template -size_t MoeGemmRunner::getMaxWorkspaceSize( - int num_experts, bool use_mxfp8_weight_scaling) const +size_t MoeGemmRunner::getMaxWorkspaceSize(int num_experts) const { - if (num_experts != num_experts_ || use_mxfp8_weight_scaling != use_mxfp8_weight_scaling_) + if (num_experts != num_experts_) { - TLLM_LOG_TRACE( - "Calling getMaxWorkspaceSize() with a new (expert count, use_mxfp8_weight_scaling) (%d, %d) vs (%d, %d)", - num_experts, (int) use_mxfp8_weight_scaling, num_experts_, (int) use_mxfp8_weight_scaling_); + TLLM_LOG_TRACE("Calling getMaxWorkspaceSize() with a new expert count %d vs %d", num_experts, num_experts_); num_experts_ = num_experts; - use_mxfp8_weight_scaling_ = use_mxfp8_weight_scaling; - gemm_workspace_size_ = calcMaxWorkspaceSize(num_experts, use_mxfp8_weight_scaling); + gemm_workspace_size_ = calcMaxWorkspaceSize(num_experts); } return gemm_workspace_size_; } template -size_t MoeGemmRunner::calcMaxWorkspaceSize( - int num_experts, bool use_mxfp8_weight_scaling) const +size_t MoeGemmRunner::calcMaxWorkspaceSize(int num_experts) const { if constexpr (use_w4_groupwise) { @@ -942,15 +900,8 @@ size_t MoeGemmRunner::calcMaxWorkspace && !use_w4afp8 && !use_wfp4a16) { // Finalize fusion may not actually be supported by the kernel, - // if they are not we will catch the error and skip them. Restrict the - // candidate set to MXFP8-valid tiles when the caller is sizing for the - // MXFP8xMXFP8 variant; otherwise the FP8 list would include tiles the - // dispatcher rejects. - auto configs = getTmaWarpSpecializedConfigs(sm_, true, use_mxfp8_weight_scaling); - // For the same template compiles both per-tensor FP8 - // (NONE) and MXFP8 block-scaled (MXFPX) variants; the caller passes - // `use_mxfp8_weight_scaling` so we size workspace for exactly the - // variant that will run. + // if they are not we will catch the error and skip them + auto configs = getTmaWarpSpecializedConfigs(sm_, true); auto fpX_block_scaling_type = TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::NONE; if constexpr (use_wfp4afp8) { @@ -960,12 +911,6 @@ size_t MoeGemmRunner::calcMaxWorkspace { fpX_block_scaling_type = TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::NVFP4; } - else if constexpr (std::is_same_v && std::is_same_v) - { - fpX_block_scaling_type = use_mxfp8_weight_scaling - ? TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::MXFPX - : TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::NONE; - } size_t max_size = 0; bool has_config = false; for (auto conf : configs) diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_template_dispatch_tma_ws.h b/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_template_dispatch_tma_ws.h index 20282a743779..339f95a96df1 100644 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_template_dispatch_tma_ws.h +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_template_dispatch_tma_ws.h @@ -73,12 +73,8 @@ namespace kernels::cutlass_kernels_oss using tensorrt_llm::kernels::cutlass_kernels::TmaWarpSpecializedGroupedGemmInput; using EpilogueFusion = TmaWarpSpecializedGroupedGemmInput::EpilogueFusion; -// `IsMXFPX` selects the Mxf8f6f4 block-scaled tensor-op path: true for -// WFP4AFP8 and WMXFP8AMXFP8, false for per-tensor FP8 / NVFP4 / BF16 / FP16. -// `is_block_scaled` extends the FP4 check to MXFP8xMXFP8 so the TMA-epilogue -// constraint applies to the new path. template + EpilogueFusion FUSION, typename TileShape, typename ClusterShape, bool is_wfp4afp8> auto getDispatchFunctionForSM100( cutlass_extensions::EpilogueScheduleType epilogue_schedule, bool dynamic_cga, bool swap_ab) { @@ -86,22 +82,19 @@ auto getDispatchFunctionForSM100( { auto select_dynamic_cga = [epilogue_schedule](auto dynamic_cga_t) { - constexpr bool is_fp4_block_scaled + constexpr bool is_block_scaled = std::is_same_v || std::is_same_v; - constexpr bool is_mxfp8_mxfp8_block_scaled - = std::is_same_v && std::is_same_v && IsMXFPX; - constexpr bool is_block_scaled = is_fp4_block_scaled || is_mxfp8_mxfp8_block_scaled; if constexpr ((!is_block_scaled || Arch::kMinComputeCapability == 103) && FUSION != EpilogueFusion::FINALIZE) { auto func_map = std::array{ &kernels::cutlass_kernels_oss::tma_warp_specialized_generic_moe_gemm_kernelLauncher, &kernels::cutlass_kernels_oss::tma_warp_specialized_generic_moe_gemm_kernelLauncher }; @@ -116,7 +109,7 @@ auto getDispatchFunctionForSM100( "No Smem epilogue schedule is not supported for block scaled types or finalize fusion"); return &kernels::cutlass_kernels_oss::tma_warp_specialized_generic_moe_gemm_kernelLauncher; } }; @@ -128,7 +121,7 @@ auto getDispatchFunctionForSM100( } template + EpilogueFusion FUSION, typename TileShape, typename ClusterShape> void dispatchMoeGemmFinalDispatchTmaWarpSpecialized(TmaWarpSpecializedGroupedGemmInput hopper_input, int num_experts, cutlass_extensions::CutlassGemmConfig gemm_config, int multi_processor_count, cudaStream_t stream, int* occupancy, size_t* workspace_size) @@ -174,33 +167,12 @@ void dispatchMoeGemmFinalDispatchTmaWarpSpecialized(TmaWarpSpecializedGroupedGem { constexpr static bool is_wfp4afp8 = std::is_same_v && std::is_same_v; - constexpr static bool is_wfp8afp8 - = std::is_same_v && std::is_same_v; - // Compile-time consistency: WFP4AFP8 is always block-scaled (IsMXFPX=true); - // non-(wfp4afp8|wfp8afp8) types must never be instantiated with IsMXFPX=true. - static_assert(!is_wfp4afp8 || IsMXFPX, "WFP4AFP8 must be instantiated with IsMXFPX=true"); - static_assert( - IsMXFPX == false || is_wfp4afp8 || is_wfp8afp8, "IsMXFPX=true is only valid for WFP4AFP8 or WMXFP8AMXFP8"); if constexpr (is_wfp4afp8) { TLLM_CHECK_WITH_INFO( hopper_input.fpX_block_scaling_type == TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::MXFPX, "MXFPX is the only supported scaling type for WFP4AFP8"); } - else if constexpr (is_wfp8afp8 && IsMXFPX) - { - TLLM_CHECK_WITH_INFO( - hopper_input.fpX_block_scaling_type == TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::MXFPX, - "WMXFP8AMXFP8 dispatch requires fpX_block_scaling_type=MXFPX"); - } - else if constexpr (is_wfp8afp8) - { - TLLM_CHECK_WITH_INFO( - hopper_input.fpX_block_scaling_type != TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::MXFPX - && hopper_input.fpX_block_scaling_type - != TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::NVFP4, - "Per-tensor FP8 e4m3xe4m3 dispatch requires fpX_block_scaling_type=NONE"); - } else { TLLM_CHECK_WITH_INFO( @@ -220,7 +192,7 @@ void dispatchMoeGemmFinalDispatchTmaWarpSpecialized(TmaWarpSpecializedGroupedGem std::get<0>(cluster_shape_fallback), std::get<1>(cluster_shape_fallback), cute::_1{}}; auto selected_func = getDispatchFunctionForSM100(gemm_config.epilogue_schedule, dynamic_cga, swap_ab); + TileShape, ClusterShape, is_wfp4afp8>(gemm_config.epilogue_schedule, dynamic_cga, swap_ab); selected_func(hopper_input, num_experts, multi_processor_count, stream, occupancy, workspace_size, cluster_shape_cute, cluster_shape_cute_fallback); } @@ -230,10 +202,10 @@ void dispatchMoeGemmFinalDispatchTmaWarpSpecialized(TmaWarpSpecializedGroupedGem constexpr bool dynamic_cga = false; auto selected_func = hopper_input.swap_ab ? kernels::cutlass_kernels_oss::tma_warp_specialized_generic_moe_gemm_kernelLauncher : kernels::cutlass_kernels_oss::tma_warp_specialized_generic_moe_gemm_kernelLauncher; selected_func(hopper_input, num_experts, multi_processor_count, stream, occupancy, workspace_size, {}, {}); @@ -241,8 +213,7 @@ void dispatchMoeGemmFinalDispatchTmaWarpSpecialized(TmaWarpSpecializedGroupedGem } } -template +template constexpr bool are_tile_shapes_supported_sm100() { // We use a runtime cluster shape for SM100, so we only support 1x1x1 and 2x1x1 cluster shapes. @@ -281,21 +252,6 @@ constexpr bool are_tile_shapes_supported_sm100() } #endif - // MXFP8xMXFP8 uses the Mxf8f6f4 block-scaled tensor-op: TileM=128 and - // TileN in {64,128,192,256}. Keep in sync with is_gemm_op_valid_sm100 in - // generate_kernels.py so the dispatcher and kernel-emission list agree. - if constexpr (IsMXFPX && std::is_same_v && std::is_same_v) - { - if (TileM != 128) - { - return false; - } - if (TileN != 64 && TileN != 128 && TileN != 192 && TileN != 256) - { - return false; - } - } - if constexpr (std::is_same_v) { if constexpr ((TileN == 16 || TileN == 8) && cute::size<0>(ClusterShape{}) == 1 @@ -348,13 +304,12 @@ constexpr bool are_tile_shapes_supported_sm120() We make the above restrictions are to improve compilation speed in TRT-LLM by pruning kernels that may not be very useful in practice. */ -template +template constexpr bool are_tile_shapes_supported() { if constexpr (Arch::kMinComputeCapability >= 100 && Arch::kMinComputeCapability < 120) { - return are_tile_shapes_supported_sm100(); + return are_tile_shapes_supported_sm100(); } else if constexpr (Arch::kMinComputeCapability == 120 || Arch::kMinComputeCapability == 121) { @@ -390,7 +345,7 @@ constexpr bool are_tile_shapes_supported() } template + EpilogueFusion FUSION, typename TileShape> void dispatchMoeGemmSelectClusterShapeTmaWarpSpecialized(TmaWarpSpecializedGroupedGemmInput hopper_input, int num_experts, cutlass_extensions::CutlassGemmConfig gemm_config, int multi_processor_count, cudaStream_t stream, int* occupancy, size_t* workspace_size) @@ -403,10 +358,10 @@ void dispatchMoeGemmSelectClusterShapeTmaWarpSpecialized(TmaWarpSpecializedGroup case cutlass_extensions::ClusterShape::ClusterShape_##M##x##N##x##K: \ { \ using ClusterShape = Shape<_##M, _##N, _##K>; \ - if constexpr (are_tile_shapes_supported()) \ + if constexpr (are_tile_shapes_supported()) \ { \ dispatchMoeGemmFinalDispatchTmaWarpSpecialized( \ + TileShape, ClusterShape>( \ hopper_input, num_experts, gemm_config, multi_processor_count, stream, occupancy, workspace_size); \ break; \ } \ @@ -432,8 +387,7 @@ void dispatchMoeGemmSelectClusterShapeTmaWarpSpecialized(TmaWarpSpecializedGroup } } -template +template void dispatchMoeGemmSelectTileShapeTmaWarpSpecialized(TmaWarpSpecializedGroupedGemmInput hopper_input, int num_experts, cutlass_extensions::CutlassGemmConfig gemm_config, int multi_processor_count, cudaStream_t stream, int* occupancy, size_t* workspace_size) @@ -448,7 +402,7 @@ void dispatchMoeGemmSelectTileShapeTmaWarpSpecialized(TmaWarpSpecializedGroupedG using KTileDim = Int; \ using TileShape = Shape<_##M, _##N, KTileDim>; \ dispatchMoeGemmSelectClusterShapeTmaWarpSpecialized( \ + EpilogueTag, FUSION, TileShape>( \ hopper_input, num_experts, gemm_config, multi_processor_count, stream, occupancy, workspace_size); \ break; \ } @@ -463,13 +417,7 @@ void dispatchMoeGemmSelectTileShapeTmaWarpSpecialized(TmaWarpSpecializedGroupedG if (gemm_config.sm_version == 90) { - // Block-scaled MXFP8xMXFP8 (IsMXFPX=true) is Blackwell-only; the SM90 launcher - // has no `is_mx_fpx=True` explicit instantiation in generate_kernels.py. Gate - // the SM90 dispatch on `!IsMXFPX` so the IsMXFPX=true template is never - // instantiated for Sm90 (otherwise the link of libth_common.so fails with - // undefined references when SM90 is included in CMAKE_CUDA_ARCHITECTURES). - if constexpr (!IsMXFPX - && kernels::cutlass_kernels::isValidHopperMOESpecialisation()) + if constexpr (kernels::cutlass_kernels::isValidHopperMOESpecialisation()) { switch (gemm_config.tile_config_sm90) { @@ -560,34 +508,9 @@ size_t calcMaxWorkspaceSizeTmaWarpSpecialized(int num_experts, cutlass_extension size_t count = 0; TmaWarpSpecializedGroupedGemmInput input{}; input.fpX_block_scaling_type = fpX_block_scaling_type; - // Most of the values are ignored for WS size calculation. We reuse the function to reduce the template bloat. - // needs the IsMXFPX template to match what the runtime dispatch will pick. - constexpr bool is_wfp4afp8 = std::is_same_v && std::is_same_v; - constexpr bool is_wfp8afp8 = std::is_same_v && std::is_same_v; - auto pick_kernel = [&]() - { - if constexpr (is_wfp4afp8) - { - return &dispatchMoeGemmSelectTileShapeTmaWarpSpecialized; - } - else if constexpr (is_wfp8afp8) - { - bool const use_mxfp8 - = fpX_block_scaling_type == TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::MXFPX; - return use_mxfp8 ? &dispatchMoeGemmSelectTileShapeTmaWarpSpecialized - : &dispatchMoeGemmSelectTileShapeTmaWarpSpecialized; - } - else - { - return &dispatchMoeGemmSelectTileShapeTmaWarpSpecialized; - } - }; - auto selected_kernel = pick_kernel(); - selected_kernel(input, num_experts, gemm_config, multi_processor_count, cudaStream_t{0}, nullptr, &count); + // Most of the values are ignored for WS size calculation. We reuse the function to reduce the template bloat + dispatchMoeGemmSelectTileShapeTmaWarpSpecialized(input, num_experts, gemm_config, multi_processor_count, cudaStream_t{0}, nullptr, &count); return count; } diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu b/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu index b7a32be2e285..a99f42003e47 100644 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu @@ -60,11 +60,11 @@ #include "tensorrt_llm/kernels/cutlass_kernels/include/moe_lora_pointer_expand.h" #include "tensorrt_llm/kernels/cutlass_kernels/include/moe_util_kernels.h" -// NOTE: the grouped-GEMM dispatch (cudaGraph(SplitK)GroupedGemm, +// NOTE: the device-path GEMM dispatch (cudaGraph(SplitK)GroupedGemm, // launchMoeLoraProblemBuilder) is not called here. Those wrappers pull in // libtorch via at::Tensor, and this file is archived into libmoe_gemm_src.a, // which the TensorRT plugin also links and must keep libtorch-free. The -// dispatch is reached through the LoraParams::grouped_gemm.run function pointer, +// dispatch is reached through the LoraParams::device_path.run function pointer, // populated in moeOp.cpp. #ifndef CUDART_VERSION @@ -1128,7 +1128,7 @@ float const** computeFP8DequantScale( } template -__device__ void setupBlockScalingFactors(TmaWarpSpecializedGroupedGemmInput& layout_info, int expert, int gemm_m, +__device__ void setupFP4BlockScalingFactors(TmaWarpSpecializedGroupedGemmInput& layout_info, int expert, int gemm_m, int gemm_n, int gemm_k, TmaWarpSpecializedGroupedGemmInput::ElementSF const* fp4_act_flat, TmaWarpSpecializedGroupedGemmInput::ElementSF const* weight_block_scale, int64_t num_tokens_before_expert) { @@ -1325,20 +1325,19 @@ __global__ void computeStridesTmaWarpSpecializedKernel(int64_t const* expert_fir { if (quant_type.fc1.weight_block_scale) { - setupBlockScalingFactors(layout_info1, expert, gemm_m, gemm1_n, gemm1_k, fp4_act_flat1, - quant_type.fc1.weight_block_scale, num_tokens_before_expert); + setupFP4BlockScalingFactors(layout_info1, expert, gemm_m, gemm1_n, gemm1_k, + fp4_act_flat1, quant_type.fc1.weight_block_scale, num_tokens_before_expert); } if (quant_type.fc2.weight_block_scale) { - setupBlockScalingFactors(layout_info2, expert, gemm_m, gemm2_n, gemm2_k, fp4_act_flat2, - quant_type.fc2.weight_block_scale, num_tokens_before_expert); + setupFP4BlockScalingFactors(layout_info2, expert, gemm_m, gemm2_n, gemm2_k, + fp4_act_flat2, quant_type.fc2.weight_block_scale, num_tokens_before_expert); } }; setupIfSelected(TmaWarpSpecializedGroupedGemmInput::NVFP4BlockScaledConfig{}, quant_params.fp4); setupIfSelected(TmaWarpSpecializedGroupedGemmInput::MXFPXBlockScaledConfig{}, quant_params.fp8_mxfp4); setupIfSelected(TmaWarpSpecializedGroupedGemmInput::MXFPXBlockScaledConfig{}, quant_params.mxfp8_mxfp4); - setupIfSelected(TmaWarpSpecializedGroupedGemmInput::MXFPXBlockScaledConfig{}, quant_params.mxfp8_mxfp8); assert(gemm_m <= INT32_MAX); assert(gemm1_n > 0 && gemm1_n <= INT32_MAX); @@ -1635,9 +1634,7 @@ void expandInputRowsKernelLauncher(InputActivationsType const* unpermuted_input, && std::is_same_v) { TLLM_CHECK_WITH_INFO(!prequant_scales, "FP8 is not supported for AWQ"); - // Either MXFP8xMXFP4 (B=fp4) or MXFP8xMXFP8 (B=fp8); both use MXFPX - // activation block-scaling, so we accept either slot. - return quant_params.mxfpxActFc1WeightScale() + return quant_params.mxfp8_mxfp4.fc1.weight_block_scale ? &expandInputRowsKernel : &expandInputRowsKernel) { - // Accept either MXFP8xMXFP4 (B=fp4) or MXFP8xMXFP8 (B=fp8); - // both use MXFPX activation block-scaling. - auto const* mxfpx_fc2_sf = quant_params.mxfpxActFc2WeightScale(); - num_padding_tokens = mxfpx_fc2_sf + num_padding_tokens = quant_params.mxfp8_mxfp4.fc2.weight_block_scale ? TmaWarpSpecializedGroupedGemmInput::MinNDimAlignmentMXFPX * num_experts_per_node : 0; - return mxfpx_fc2_sf ? fn(MXFPX) : fn(NONE); + return quant_params.mxfp8_mxfp4.fc2.weight_block_scale ? fn(MXFPX) : fn(NONE); } else #endif @@ -2996,22 +2990,20 @@ CutlassMoeFCRunner:: auto act_sf_rows = min_latency_mode ? num_moe_inputs : std::min(num_moe_inputs, static_cast(num_rows * num_experts_per_node)); - auto const scaling_type = getScalingType(use_mxfp8_weight_scaling_); - size_t const sf_size = scaling_type == TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::MXFPX + size_t const sf_size = getScalingType() == TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::MXFPX ? sizeof(TmaWarpSpecializedGroupedGemmInput::MXFPXElementSF) : sizeof(TmaWarpSpecializedGroupedGemmInput::NVFP4ElementSF); size_t const fc1_fp4_act_scale_size - = getOffsetActivationSF(num_experts_per_node, act_sf_rows, hidden_size, scaling_type) * sf_size; + = getOffsetActivationSF(num_experts_per_node, act_sf_rows, hidden_size, getScalingType()) * sf_size; size_t const fc2_fp4_act_scale_size - = getOffsetActivationSF(num_experts_per_node, act_sf_rows, inter_size, scaling_type) * sf_size; + = getOffsetActivationSF(num_experts_per_node, act_sf_rows, inter_size, getScalingType()) * sf_size; size_t const fp4_act_scale_size = std::max(fc1_fp4_act_scale_size, fc2_fp4_act_scale_size); size_t const tma_ws_size - = using_tma_ws ? TmaWarpSpecializedGroupedGemmInput::workspaceSize(num_experts_per_node, scaling_type) : 0; + = using_tma_ws ? TmaWarpSpecializedGroupedGemmInput::workspaceSize(num_experts_per_node, getScalingType()) : 0; - size_t const gemm_workspace_size - = moe_gemm_runner_.getMaxWorkspaceSize(num_experts_per_node, use_mxfp8_weight_scaling_); + size_t const gemm_workspace_size = moe_gemm_runner_.getMaxWorkspaceSize(num_experts_per_node); // lora related size_t const lora_input_size @@ -3187,14 +3179,7 @@ void CutlassMoeFCRunner MoE serves both per-tensor FP8 (no SF buffer needed) and - // MXFP8xMXFP8 (needs the per-expert SF workspace). The constexpr - // `use_block_scaling = use_fp4 || use_wfp4afp8` is false for our - // template instantiation, so without the runtime check below the SF - // buffer pointer is left nullptr and the kernel reads garbage SF - // descriptors -> cudaErrorIllegalInstruction. Confirmed via a device - // printf in setupBlockScalingFactors that showed sf_act_ptr=(nil). - if (use_block_scaling || use_mxfp8_weight_scaling_) + if (use_block_scaling) { fc1_fp4_act_scale_ = getWsPtr(TmaWarpSpecializedGroupedGemmInput::ElementSF{}, "fp4_act_scale"); fc2_fp4_act_scale_ = getWsPtr(TmaWarpSpecializedGroupedGemmInput::ElementSF{}, "fp4_act_scale"); @@ -3206,13 +3191,12 @@ void CutlassMoeFCRunner 0, - "Grouped-GEMM LoRA dtype_bytes must be positive (grouped_gemm not fully populated?)."); - auto validateLoraModule = [](::tensorrt_llm::kernels::cutlass_kernels::MoeLoraGroupedGemmModule const& mod, - char const* name, int64_t expectedDimA, int64_t expectedDimB) - { - TLLM_CHECK_WITH_INFO(mod.ranks_src_dev != nullptr && mod.ptrs_src_dev != nullptr - && mod.permuted_ranks_dev != nullptr && mod.permuted_ptrs_dev != nullptr, - "Grouped-GEMM LoRA %s module is missing pointer-expand buffers.", name); - TLLM_CHECK_WITH_INFO(mod.dim_a == expectedDimA && mod.dim_b == expectedDimB, - "Grouped-GEMM LoRA %s module dimensions do not match the MoE runner dimensions.", name); - }; - validateLoraModule(grouped_gemm.fc1, "fc1", hidden_size, inter_size); - validateLoraModule(grouped_gemm.fc2, "fc2", inter_size, hidden_size); - if (is_gated_activation) - { - validateLoraModule(grouped_gemm.gated, "gated", hidden_size, inter_size); - } - - // Translate per-module grouped-GEMM metadata into the MoeLoraExpandModule + if (lora_params.device_path.enabled) + { + auto const& dp = lora_params.device_path; + // Translate per-module device-path metadata into the MoeLoraExpandModule // API that the pointer-expand kernel expects. ::tensorrt_llm::kernels::cutlass_kernels::MoeLoraExpandModule fc1_mod{}; - fc1_mod.ranks_src = grouped_gemm.fc1.ranks_src_dev; - fc1_mod.ptrs_src = grouped_gemm.fc1.ptrs_src_dev; - fc1_mod.dim_a = grouped_gemm.fc1.dim_a; - fc1_mod.dim_b = grouped_gemm.fc1.dim_b; - fc1_mod.ranks_out = grouped_gemm.fc1.permuted_ranks_dev; - fc1_mod.ptrs_out = grouped_gemm.fc1.permuted_ptrs_dev; + fc1_mod.ranks_src = dp.fc1.ranks_src_dev; + fc1_mod.ptrs_src = dp.fc1.ptrs_src_dev; + fc1_mod.dim_a = dp.fc1.dim_a; + fc1_mod.dim_b = dp.fc1.dim_b; + fc1_mod.ranks_out = dp.fc1.permuted_ranks_dev; + fc1_mod.ptrs_out = dp.fc1.permuted_ptrs_dev; ::tensorrt_llm::kernels::cutlass_kernels::MoeLoraExpandModule fc2_mod{}; - fc2_mod.ranks_src = grouped_gemm.fc2.ranks_src_dev; - fc2_mod.ptrs_src = grouped_gemm.fc2.ptrs_src_dev; - fc2_mod.dim_a = grouped_gemm.fc2.dim_a; - fc2_mod.dim_b = grouped_gemm.fc2.dim_b; - fc2_mod.ranks_out = grouped_gemm.fc2.permuted_ranks_dev; - fc2_mod.ptrs_out = grouped_gemm.fc2.permuted_ptrs_dev; + fc2_mod.ranks_src = dp.fc2.ranks_src_dev; + fc2_mod.ptrs_src = dp.fc2.ptrs_src_dev; + fc2_mod.dim_a = dp.fc2.dim_a; + fc2_mod.dim_b = dp.fc2.dim_b; + fc2_mod.ranks_out = dp.fc2.permuted_ranks_dev; + fc2_mod.ptrs_out = dp.fc2.permuted_ptrs_dev; ::tensorrt_llm::kernels::cutlass_kernels::MoeLoraExpandModule gated_mod{}; if (is_gated_activation) { - gated_mod.ranks_src = grouped_gemm.gated.ranks_src_dev; - gated_mod.ptrs_src = grouped_gemm.gated.ptrs_src_dev; - gated_mod.dim_a = grouped_gemm.gated.dim_a; - gated_mod.dim_b = grouped_gemm.gated.dim_b; - gated_mod.ranks_out = grouped_gemm.gated.permuted_ranks_dev; - gated_mod.ptrs_out = grouped_gemm.gated.permuted_ptrs_dev; + gated_mod.ranks_src = dp.gated.ranks_src_dev; + gated_mod.ptrs_src = dp.gated.ptrs_src_dev; + gated_mod.dim_a = dp.gated.dim_a; + gated_mod.dim_b = dp.gated.dim_b; + gated_mod.ranks_out = dp.gated.permuted_ranks_dev; + gated_mod.ptrs_out = dp.gated.permuted_ptrs_dev; } ::tensorrt_llm::kernels::cutlass_kernels::launchMoeLoraPointerExpand(permuted_row_to_unpermuted_row_, - expert_first_token_offset_, num_experts_per_node, start_expert, num_rows, expanded_num_rows, - grouped_gemm.dtype_bytes, fc1_mod, fc2_mod, is_gated_activation ? &gated_mod : nullptr, stream); + expert_first_token_offset_, num_experts_per_node, start_expert, num_rows, expanded_num_rows, dp.dtype_bytes, + fc1_mod, fc2_mod, is_gated_activation ? &gated_mod : nullptr, stream); sync_check_cuda_error(stream); return /*all_token_without_lora=*/false; } @@ -3919,15 +3880,15 @@ auto CutlassMoeFCRunner(permuted_data_); } - // Grouped-GEMM branch, running entirely on the stream. setupLoraWorkspace + // Device-path branch, running entirely on the stream. setupLoraWorkspace // has already populated the per-permuted-row ranks and pointers for fc1 and // gated via launchMoeLoraPointerExpand. - if (lora_params.grouped_gemm.enabled) + if (lora_params.device_path.enabled) { - auto const& grouped_gemm = lora_params.grouped_gemm; + auto const& dp = lora_params.device_path; nvinfer1::DataType const data_type = moeLoraNvInferType(); - // The grouped-GEMM GEMM skips rank-0 rows, but the bias/reorder paths + // The device-path GEMM skips rank-0 rows, but the bias/reorder paths // read lora_fc1_result_ for every valid row. Zero the buffer first so // skipped rows are a deterministic no-op. It is contiguous and holds // both the gated and fc1 halves when gated, so one memset covers both. @@ -3935,17 +3896,15 @@ auto CutlassMoeFCRunner(input), - /*output_base=*/static_cast(lora_fc1_result), grouped_gemm.run, data_type, stream); + runMoeLoraDeviceModule(dp.fc1, expanded_num_rows, /*in_hidden_size=*/hidden_size, dp.max_lora_rank, + dp.dtype_bytes, dp.splitk_slices, /*input_base=*/static_cast(input), + /*output_base=*/static_cast(lora_fc1_result), dp.run, data_type, stream); if (is_gated_activation) { - runMoeLoraGroupedGemmModule(grouped_gemm.gated, expanded_num_rows, /*in_hidden_size=*/hidden_size, - grouped_gemm.max_lora_rank, grouped_gemm.dtype_bytes, grouped_gemm.splitk_slices, - /*input_base=*/static_cast(input), - /*output_base=*/static_cast(lora_gated_out), grouped_gemm.run, data_type, stream); + runMoeLoraDeviceModule(dp.gated, expanded_num_rows, /*in_hidden_size=*/hidden_size, dp.max_lora_rank, + dp.dtype_bytes, dp.splitk_slices, /*input_base=*/static_cast(input), + /*output_base=*/static_cast(lora_gated_out), dp.run, data_type, stream); } } else @@ -4013,13 +3972,13 @@ void CutlassMoeFCRunner(fc1_result_); } - // Grouped-GEMM branch, mirroring loraFC1's branch. It consumes the + // Device-path branch, mirroring loraFC1's branch. It consumes the // per-permuted-row ranks and pointers that setupLoraWorkspace produced via // launchMoeLoraPointerExpand. num_tokens here is expanded_num_rows from // runMoe (top_k * num_rows). - if (lora_params.grouped_gemm.enabled) + if (lora_params.device_path.enabled) { - auto const& grouped_gemm = lora_params.grouped_gemm; + auto const& dp = lora_params.device_path; nvinfer1::DataType const data_type = moeLoraNvInferType(); // As in loraFC1, zero the output so rank-0 rows the GEMM skips do not @@ -4028,10 +3987,9 @@ void CutlassMoeFCRunner(num_tokens) * static_cast(hidden_size) * sizeof(ScaleBiasType); TLLM_CUDA_CHECK(cudaMemsetAsync(lora_fc2_result_, 0, fc2_result_bytes, stream)); - runMoeLoraGroupedGemmModule(grouped_gemm.fc2, num_tokens, /*in_hidden_size=*/inter_size, - grouped_gemm.max_lora_rank, grouped_gemm.dtype_bytes, grouped_gemm.splitk_slices, - /*input_base=*/static_cast(input), - /*output_base=*/static_cast(lora_fc2_result_), grouped_gemm.run, data_type, stream); + runMoeLoraDeviceModule(dp.fc2, num_tokens, /*in_hidden_size=*/inter_size, dp.max_lora_rank, dp.dtype_bytes, + dp.splitk_slices, /*input_base=*/static_cast(input), + /*output_base=*/static_cast(lora_fc2_result_), dp.run, data_type, stream); sync_check_cuda_error(stream); return; } @@ -4097,16 +4055,13 @@ void CutlassMoeFCRunner::value` - // expression adjusts automatically via the template parameter. TLLM_CHECK_WITH_INFO(hidden_size % (64 * 8 / sizeof_bits::value) == 0, - "Hidden size %d does not meet minimum alignment requirements for MXFPX MOE GEMM %d", (int) hidden_size, - (int) (64 * 8 / sizeof_bits::value)); + "Hidden size %d does not meet minimum alignment requirements for MXFP8_MXFP4 MOE GEMM %d", + (int) hidden_size, (int) (64 * 8 / sizeof_bits::value)); TLLM_CHECK_WITH_INFO(inter_size % (64 * 8 / sizeof_bits::value) == 0, - "Inter size %d does not meet minimum alignment requirements for MXFPX MOE GEMM %d", (int) inter_size, + "Inter size %d does not meet minimum alignment requirements for MXFP8_MXFP4 MOE GEMM %d", (int) inter_size, (int) (64 * 8 / sizeof_bits::value)); } else @@ -4164,12 +4119,8 @@ void CutlassMoeFCRunner& host_permuted_rows = host_lora_workspace_.host_permuted_rows; std::vector& host_expert_first_token_offset = host_lora_workspace_.host_expert_first_token_offset; @@ -4396,7 +4347,7 @@ CutlassMoeFCRunner:: TmaWarpSpecializedGroupedGemmInput::ElementSF const* fp4_act_flat2, QuantParams quant_params, ScaleBiasType const* bias1, ScaleBiasType const* bias2, UnfusedGemmOutputType* gemm1_output, UnfusedGemmOutputType* gemm2_output, float const* router_scales, int const* permuted_row_to_unpermuted_row, - TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType scaling_type, cudaStream_t stream) + cudaStream_t stream) { // Always nullptr layout_info1.ptr_c = nullptr; @@ -4434,8 +4385,8 @@ CutlassMoeFCRunner:: layout_info1.int4_groupwise_params.use_wfp4a16 = use_wfp4a16; layout_info2.int4_groupwise_params.use_wfp4a16 = use_wfp4a16; - layout_info1.fpX_block_scaling_type = scaling_type; - layout_info2.fpX_block_scaling_type = scaling_type; + layout_info1.fpX_block_scaling_type = getScalingType(); + layout_info2.fpX_block_scaling_type = getScalingType(); int const threads = std::min(1024, num_experts_per_node); int const blocks = (num_experts_per_node + threads - 1) / threads; @@ -4586,7 +4537,7 @@ CutlassMoeFCRunner:: fc1_fp4_act_scale_, fc2_fp4_act_scale_, quant_params, fc1_expert_biases, fc2_bias, reinterpret_cast(gemm1_output), reinterpret_cast(fc2_result_), permuted_token_final_scales_, - permuted_row_to_unpermuted_row_, getScalingType(use_mxfp8_weight_scaling_), stream); + permuted_row_to_unpermuted_row_, stream); } } @@ -4831,7 +4782,6 @@ std::map> GemmProfilerBackend::getProfile } // FP8 sizes - bool is_mxfp8_w8a8 = is_fp8_w_quant && is_fp8_act_quant && mUseMxfp8WeightScaling; quant_1_size = is_fp8_w_quant ? num_experts_per_node * sizeof(float) : quant_1_size; quant_2_size = is_fp8_w_quant ? sizeof(float) : quant_2_size; size_t quant_3_size = is_fp8_w_quant ? num_experts_per_node * sizeof(float) : 0; @@ -4842,18 +4792,6 @@ std::map> GemmProfilerBackend::getProfile quant_4_size = quant_2_size; } - // MXFP8 W8A8 weight SF sizes (need full per-expert SF buffers, not just per-tensor - // dequant scalars). quant_1 = fc1 weight SF (n=inter_size, k=hidden_size), - // quant_2 = fc2 weight SF (n=hidden_size, k=inter_size). - if (is_mxfp8_w8a8) - { - quant_1_size = getOffsetWeightSF(num_experts_per_node, fc1_out_size, hidden_size, mScalingType) - * sizeof(TmaWarpSpecializedGroupedGemmInput::ElementSF); - quant_2_size = getOffsetWeightSF(num_experts_per_node, hidden_size, inter_size, mScalingType) - * sizeof(TmaWarpSpecializedGroupedGemmInput::ElementSF); - quant_3_size = 0; - } - // FP4 sizes quant_1_size = is_fp4_w_quant ? sizeof(float) : quant_1_size; quant_2_size = is_fp4_w_quant ? getOffsetWeightSF(num_experts_per_node, inter_size, hidden_size, mScalingType) @@ -5072,22 +5010,9 @@ void GemmProfilerBackend::prepareQuantParams(int num_tokens, char* workspace_ptr } else if (mWType == nvinfer1::DataType::kFP8) { - if (mUseMxfp8WeightScaling) - { - // MXFP8 W8A8: only need per-expert weight SF buffers (no per-tensor alpha, - // no activation global scale — block scales carry that information). Wire - // them through the dedicated mxfp8_mxfp8 slot so per-expert SF setup runs. - TLLM_CHECK(quant_1 && quant_2); - mQuantParams = QuantParams::MXFP8MXFP8( - static_cast(quant_1), - static_cast(quant_2)); - } - else - { - TLLM_CHECK(quant_1 && quant_2 && quant_3); - mQuantParams = QuantParams::FP8(static_cast(quant_1), static_cast(quant_2), - static_cast(quant_3), static_cast(quant_4)); - } + TLLM_CHECK(quant_1 && quant_2 && quant_3); + mQuantParams = QuantParams::FP8(static_cast(quant_1), static_cast(quant_2), + static_cast(quant_3), static_cast(quant_4)); } else if (mDType == nvinfer1::DataType::kFP8 && (mWType == nvinfer1::DataType::kFP4 || mWType == nvinfer1::DataType::kINT64)) diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_lora_pointer_expand.cu b/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_lora_pointer_expand.cu index 8a23e51a3dd0..d130d917ad33 100644 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_lora_pointer_expand.cu +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_lora_pointer_expand.cu @@ -60,7 +60,7 @@ __device__ inline void expandOneModule( mod.ranks_out[i] = rank; } -// Reset one module's output slot to a rank-0 no-op. The grouped-GEMM scratch is +// Reset one module's output slot to a rank-0 no-op. The device-path scratch is // persistent and reused, so ghost rows must be explicitly zeroed; otherwise // stale ranks or pointers survive into the next grouped GEMM. __device__ inline void zeroOneModule(MoeLoraExpandModule const& mod, int64_t i) diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_lora_problem_builder.cu b/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_lora_problem_builder.cu index e6b3f9b98fc0..f1b6da4ca189 100644 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_lora_problem_builder.cu +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_lora_problem_builder.cu @@ -68,16 +68,8 @@ __global__ void moeLoraProblemBuilderKernel(int32_t const* __restrict__ ranks, i // Problem sizes: each permuted token gets its own (M=1) GEMM. This matches // worst-case scheduling with no run-length aggregation; a future // optimization can aggregate consecutive identical-adapter tokens. - // - // Rank-0 rows carry no active adapter (base/no-LoRA request, padding, or - // warmup) and have null A/B pointers, so their delta is zero and the caller - // pre-zeroes the output. The in-GEMM already collapses to N=0 (rank is its - // N) and is skipped, but the out-GEMM's N is out_hidden_size; forcing it to - // zero here lets the grouped GEMM skip these rows too instead of launching - // tiles that dereference the null B pointer. - int const out_n = (rank > 0) ? static_cast(out_hidden_size) : 0; problem_sizes_in[i] = cutlass::gemm::GemmCoord(1, rank, static_cast(in_hidden_size)); - problem_sizes_out[i] = cutlass::gemm::GemmCoord(1, out_n, rank); + problem_sizes_out[i] = cutlass::gemm::GemmCoord(1, static_cast(out_hidden_size), rank); // Pointer rows. dtype_bytes scales the per-row stride so the same // builder serves bf16/fp16/fp32 adapters without templating. diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/python/generate_kernels.py b/cpp/tensorrt_llm/kernels/cutlass_kernels/python/generate_kernels.py index cfbf2ef774b4..0e4c76c2c92d 100644 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/python/generate_kernels.py +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/python/generate_kernels.py @@ -3,13 +3,7 @@ import os from itertools import chain, product -# Explicit re-imports for names used in MXFP8 / W4A8_MXFP4_MXFP8 paths so they -# are not flagged as F405 "may be undefined from star import". This is purely -# for the legacy lint baseline gate; functionally identical to picking them up -# from the star import. from cutlass_library import * -from cutlass_library import EpilogueScheduleType # noqa: F401, F811 -from cutlass_library import DataType, GemmKind, KernelScheduleType ################################################################################ @@ -387,17 +381,6 @@ def is_gemm_op_valid_sm100(op): if op.arch == 100 and op.epi_schedule == EpilogueScheduleType.PtrArrayNoSmemWarpSpecialized1Sm: return False - # MXFP8xMXFP8 uses the Mxf8f6f4 block-scaled tensor-op. Mirror the FP4 - # block-scaled path's shape constraints (kept in sync with the IsMXFPX - # branch in are_tile_shapes_supported_sm100 in moe_gemm_template_dispatch_tma_ws.h). - if op.is_mx_fpx and op.act_type == DataType.e4m3 and op.weight_type == DataType.e4m3: - if tile_m != 128: - return False - if tile_n not in [64, 128, 192, 256]: - return False - if op.arch == 100 and op.epi_schedule == EpilogueScheduleType.PtrArrayNoSmemWarpSpecialized1Sm: - return False - # Shapes for fp8 small N shapes if (op.act_type == DataType.e4m3) and (tile_n == 16 or tile_n == 8) and (cga_m == 1 @@ -802,43 +785,30 @@ def generate_sm100_grouped_gemm_operations(is_arch_enabled, arch): if dtype in [DataType.e4m3, e2m1]: otypes = [DataType.f16, DataType.bf16] - # `` serves BOTH per-tensor FP8 MoE (is_mx_fpx=False) and - # MXFP8xMXFP8 block-scaled MoE (is_mx_fpx=True). Runtime dispatcher - # selects between them via TmaWarpSpecializedGroupedGemmInput::fpX_block_scaling_type. - is_wmxfp8amxfp8_eligible = (dtype == DataType.e4m3 - and weight_type == DataType.e4m3) - if is_wmxfp8amxfp8_eligible: - is_mx_fpx_variants = [False, True] - elif dtype == DataType.e4m3 and weight_type == e2m1: - is_mx_fpx_variants = [True] - else: - is_mx_fpx_variants = [False] - for otype in otypes: - for is_mx_fpx_variant in is_mx_fpx_variants: - moe_gemm_operation = TrtLlm_GemmLauncher( - GemmKind.Grouped, - arch, - dtype, - weight_type, - otype, - otype, - otype, - quant_op, - epi_tag, - cta_shape_mnk, - warp_shape, - stages, - cga_shape, - mainloop_schedule, - epi_schedule, - epi_fusion, - is_mx_fpx=is_mx_fpx_variant, - dynamic_cga=dynamic_cga, - swap_ab=swap_ab) - - if is_op_valid(moe_gemm_operation): - operations.append(moe_gemm_operation) + moe_gemm_operation = TrtLlm_GemmLauncher( + GemmKind.Grouped, + arch, + dtype, + weight_type, + otype, + otype, + otype, + quant_op, + epi_tag, + cta_shape_mnk, + warp_shape, + stages, + cga_shape, + mainloop_schedule, + epi_schedule, + epi_fusion, + is_mx_fpx=(dtype == DataType.e4m3 and weight_type == e2m1), + dynamic_cga=dynamic_cga, + swap_ab=swap_ab) + + if is_op_valid(moe_gemm_operation): + operations.append(moe_gemm_operation) return operations diff --git a/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention.h b/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention.h index e8ec573cdab5..9ef6593d162c 100644 --- a/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention.h +++ b/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention.h @@ -209,10 +209,6 @@ struct Multihead_attention_params_base // threadblock counter to identify the complete of partial attention computations int* block_counter = nullptr; - float* cascade_partial_out = nullptr; // [batch_beam x num_heads x head_size] - float* cascade_partial_max = nullptr; // [batch_beam x num_heads] - float* cascade_partial_sum = nullptr; // [batch_beam x num_heads] - int const* memory_length_per_sample = nullptr; int32_t const* mrope_position_deltas = nullptr; }; diff --git a/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/CMakeLists.txt b/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/CMakeLists.txt index 51776bc00a50..9ad95211f86f 100644 --- a/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/CMakeLists.txt +++ b/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/CMakeLists.txt @@ -40,11 +40,6 @@ foreach(HEAD ${basic_heads}) list(APPEND SRC_CU ${HEAD_SRCS}) endforeach() -# Cascade attention kernel. Compiled together with the basic heads so it lives -# in the same shared library as the MMHA dispatcher that calls into it. Uses Dh -# in {64, 128}. -list(APPEND SRC_CU ${CMAKE_CURRENT_SOURCE_DIR}/cascadeAttentionKernel.cu) - # skip mmha 48, 80, 96, 104, 112, 144, 160, 192 and 224 for fast build if(FAST_BUILD) set(extra_heads 256) diff --git a/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/cascadeAttentionKernel.cu b/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/cascadeAttentionKernel.cu deleted file mode 100644 index 8bdcc9109aeb..000000000000 --- a/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/cascadeAttentionKernel.cu +++ /dev/null @@ -1,1235 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2024-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. - */ - -#include "cascadeAttentionKernel.h" -#include "cascadeMma.cuh" - -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/cudaTypeUtils.cuh" -#include "tensorrt_llm/common/envUtils.h" -#include "tensorrt_llm/common/logger.h" -#include "tensorrt_llm/kernels/decoderMaskedMultiheadAttentionUtils.h" -#include "tensorrt_llm/kernels/gptKernels.h" -#include "tensorrt_llm/kernels/kvCacheUtils.h" - -#include -#include -#ifdef ENABLE_BF16 -#include -#endif - -#include -#include -#include -#include -#include - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels -{ -namespace mmha -{ -namespace cascade -{ - -//////////////////////////////////////////////////////////////////////////////////////////////////// -// -// Internal helpers -// -//////////////////////////////////////////////////////////////////////////////////////////////////// - -// Per-block configuration for the Tensor-Core Phase 1 kernel. -// -// - THDS_PER_BLOCK = 128 = 4 warps -// - BEAM_TILE = 16 = MMA M dimension (1 m16n8 tile of beams) -// - TOKEN_TILE = 16 = MMA K dimension for P·V and N dimension for Q·K^T -// -// Phase 2 (suffix decode) still uses THDS_PER_BLOCK threads-per-block with the -// per-channel block_sum reduction. Supported Dh value is 128. -template -struct CascadeConfig -{ - static constexpr int THDS_PER_BLOCK = 128; - static constexpr int BEAM_TILE = 16; - static constexpr int TOKEN_TILE = 16; - static_assert(Dh_ == 128, "cascade kernel only supports Dh = 128"); -}; - -// Tag selecting which side of the KV cache to read from. -enum class KVKind -{ - K, - V -}; - -// Read one element of K[token, head, channel] or V[token, head, channel] from the cache. -// Dh is the head dimension (== params.hidden_size_per_head, enforced by the -// caller's runtime guard) and is propagated as a compile-time constant. -template -__device__ inline T_cache load_kv(KVCacheBuffer const& kv, int seqIdx, int tokenIdx, int headIdx, int channel) -{ - auto const localTokenIdx = kv.getKVTokenIdx(tokenIdx); - void* rawPtr; - if constexpr (Kind == KVKind::K) - { - rawPtr = kv.getKBlockPtr(seqIdx, localTokenIdx); - } - else - { - rawPtr = kv.getVBlockPtr(seqIdx, localTokenIdx); - } - auto* blockPtr = reinterpret_cast(rawPtr); - auto const localOffset = kv.getKVLocalIdx(localTokenIdx, headIdx, Dh, channel); - return blockPtr[localOffset]; -} - -template -__device__ inline float block_sum(float v, float* scratch) -{ - static_assert(THDS % 32 == 0, "block_sum requires THDS to be a multiple of warpSize"); - constexpr int N_WARPS = THDS / 32; - int const tid = threadIdx.x; - int const warp_id = tid >> 5; - int const lane_id = tid & 31; - - // Step 1: intra-warp shuffle reduction (no __syncthreads needed). -#pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) - { - v += __shfl_xor_sync(0xFFFFFFFFu, v, offset); - } - - // Step 2: each warp's lane 0 publishes the warp partial. - if (lane_id == 0) - { - scratch[warp_id] = v; - } - __syncthreads(); - - // Step 3: warp 0 reduces the N_WARPS partials and broadcasts the result - // back through scratch[0]. - if (warp_id == 0) - { - float block_val = (lane_id < N_WARPS) ? scratch[lane_id] : 0.f; -#pragma unroll - for (int offset = N_WARPS / 2; offset > 0; offset >>= 1) - { - block_val += __shfl_xor_sync(0xFFFFFFFFu, block_val, offset); - } - if (lane_id == 0) - { - scratch[0] = block_val; - } - } - __syncthreads(); - - float const result = scratch[0]; - __syncthreads(); // keep scratch free for the next invocation - return result; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// -// -// RoPE helpers (GPT-NeoX style, full-head rotation). -// -// Cached K in TRT-LLM is stored *post-RoPE* whenever POS_SHIFT == false -// (which we already enforce in cascade_eligible), so cascade only needs to -// apply RoPE to the live Q tensor. The pair layout is (c, c + rotary_dim/2). -// -//////////////////////////////////////////////////////////////////////////////////////////////////// - -// Resolve the per-token row pointer in the global cos/sin cache. -// Returns nullptr when the global cache is not provided (callers fall back to -// computing cos/sin from base/scale). -__device__ __forceinline__ float2 const* cascade_rope_cache_row(float2 const* cos_sin_cache, int t_step, int rotary_dim) -{ - return cos_sin_cache != nullptr ? cos_sin_cache + static_cast(t_step) * (rotary_dim / 2) : nullptr; -} - -// NeoX RoPE coefficient for a single pair index `freq_idx \in [0, rotary_dim/2)`. -// `cache_row` is the per-token row pointer returned by `cascade_rope_cache_row` -// (nullable). Returns {cos, sin}. -__device__ __forceinline__ float2 cascade_rope_neox_cs( - int freq_idx, int rotary_dim, int t_step, float base, float scale, float2 const* cache_row) -{ - return rotary_embedding_coefficient( - /*inv_freq_cache=*/nullptr, - /*cos_sin_cache=*/cache_row, - /*zid=*/2 * freq_idx, - /*rot_embed_dim=*/rotary_dim, - /*base=*/base, - /*scale=*/scale, - /*mscale=*/1.0f, - /*t_step=*/static_cast(t_step)); -} - -template -__device__ constexpr int cascade_smem_row_stride(int dh) -{ - // Keep 16B padding (= 16 / sizeof(T) elements): 8 for bf16/half. - return dh + (16 / static_cast(sizeof(T))); -} - -template -__device__ __forceinline__ void cascade_async_load_kv_tile( - KVCacheBuffer const& kv, int owner_seq, int kv_head_idx, int t0, int tile_end, int tid, T* smem_buf) -{ - constexpr int BYTES_PER_CHUNK = 16; - constexpr int ELEMS_PER_CHUNK = BYTES_PER_CHUNK / sizeof(T); - constexpr int CHUNKS_PER_TOK = Dh / ELEMS_PER_CHUNK; - constexpr int TOTAL_CHUNKS = TOKEN_TILE * CHUNKS_PER_TOK; - constexpr int ROW_STRIDE = cascade_smem_row_stride(Dh); - static_assert(Dh % ELEMS_PER_CHUNK == 0, "Dh must be a multiple of 8 (16B / bf16)"); - -#pragma unroll - for (int idx = tid; idx < TOTAL_CHUNKS; idx += THDS) - { - int const tok_rel = idx / CHUNKS_PER_TOK; - int const chunk = idx - tok_rel * CHUNKS_PER_TOK; - int const dh_base = chunk * ELEMS_PER_CHUNK; - int const tok = t0 + tok_rel; - T* s_ptr = smem_buf + tok_rel * ROW_STRIDE + dh_base; - - bool const in_bounds = (tok < tile_end); - int const tok_safe = in_bounds ? tok : t0; - auto const localTokenIdx = kv.getKVTokenIdx(tok_safe); - void* rawPtr; - if constexpr (Kind == KVKind::K) - { - rawPtr = kv.getKBlockPtr(owner_seq, localTokenIdx); - } - else - { - rawPtr = kv.getVBlockPtr(owner_seq, localTokenIdx); - } - auto* base = reinterpret_cast(rawPtr); - auto const off = kv.getKVLocalIdx(localTokenIdx, kv_head_idx, Dh, dh_base); - unsigned const src_size = in_bounds ? 16u : 0u; - mma::cp_async_16B(mma::smem_addr_of(s_ptr), base + off, src_size); - } -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// -// -// Phase 1: Shared-prefix multi-query attention (Tensor-Core accelerated). -// -// Grid: ( num_heads, num_requests, ceil(beam_width / BEAM_TILE) ) -// Block: 128 threads = 4 warps -// -// Each block processes BEAM_TILE=16 beams x TOKEN_TILE=16 prefix tokens per -// inner iteration, using m16n8k16 Tensor-Core MMAs for both Q·K^T and P·V. -// Online-softmax statistics (m, l) are maintained per beam in SMEM, and an -// FP32 output accumulator is kept in registers. -// -// * Warp 0 owns the Q·K^T MMA + online softmax + produces P_smem (bf16/half) -// and alpha[beam] (=exp(prev_m - new_m)) in SMEM. -// * All four warps own slices of Dh for P·V: each warp covers DH/4 channels. -// * At iteration end the four warps rescale their O_acc by alpha[beam] and -// issue PV MMAs reading A=P_smem and B=V_smem. -// -// SMEM layout (bytes for Dh=128, bf16; adds 16B padding per Q/K/V row): -// Q_smem : [16][136] = 4352 (Dh + 8 bf16 padding) -// K_smem : [16][136] = 4352 -// V_smem : [16][136] = 4352 -// P_smem : [16][16] = 512 (not padded; only 32B per row) -// stats : [16][2] = 128 (fp32) -// alpha : [16] = 64 (fp32) -// Total ≈ 13696 B single-buffer / 21.9 KB with K/V double-buffer -// (well under the 46 KB threshold that would trigger -// cudaFuncAttributeMaxDynamicSharedMemorySize). -// -//////////////////////////////////////////////////////////////////////////////////////////////////// - -template -__global__ void cascade_prefix_mqa_kernel(Multihead_attention_params params, KVCacheBuffer kv_cache_buffer, - int const* __restrict__ d_input_lengths, float* __restrict__ partial_out, float* __restrict__ partial_m, - float* __restrict__ partial_l) -{ -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 - if (blockIdx.x == 0 && blockIdx.y == 0 && blockIdx.z == 0 && threadIdx.x == 0 && threadIdx.y == 0 - && threadIdx.z == 0) - { - printf("cascade_prefix_mqa_kernel: requires SM_80+ but launched on an unsupported device.\n"); - __trap(); - } - return; -#else - using Cfg = CascadeConfig; - constexpr int THDS = Cfg::THDS_PER_BLOCK; // 128 - constexpr int BEAM_TILE = Cfg::BEAM_TILE; // 16 - constexpr int TOKEN_TILE = Cfg::TOKEN_TILE; // 16 - constexpr int WARPS = THDS / 32; // 4 - constexpr int DH_PER_WARP = Dh / WARPS; // 32 (Dh=128) - constexpr int PV_N_BLOCKS = DH_PER_WARP / 8; // m16n8 output tiles per warp - static_assert(BEAM_TILE == 16 && TOKEN_TILE == 16, "MMA tiles are fixed at 16"); - static_assert(Dh % 16 == 0, "Dh must be divisible by MMA k-dim 16"); - static_assert(DH_PER_WARP >= 8 && DH_PER_WARP % 8 == 0, "PV needs at least one n8 per warp"); - - int const head_idx = blockIdx.x; - int const req_idx = blockIdx.y; - int const beam_chunk = blockIdx.z; - int const beam_base = beam_chunk * BEAM_TILE; - int const tid = threadIdx.x; - int const warp_id = tid >> 5; - int const lane_id = tid & 31; - - int const beam_width = params.beam_width; - int const num_heads = params.num_heads; - int const num_kv_heads = params.num_kv_heads; - float const inv_sqrt_dh = params.inv_sqrt_dh; - - int const prefix_len = __ldg(&d_input_lengths[req_idx * beam_width]); - - if (beam_base >= beam_width || params.hidden_size_per_head != Dh) - { - return; - } - - int const kv_head_idx = (num_kv_heads > 0) ? (head_idx * num_kv_heads / num_heads) : head_idx; - // All beams share prompt KV via beam 0 (TRT-LLM context-phase convention). - int const owner_seq = req_idx * beam_width; - - // ------------------------------ SMEM layout ------------------------- - // P0-1: each Q/K/V row carries 16B of padding (8 bf16 for Dh=128) to - // break the 8-way bank conflict seen on Q·K^T LDS. Total SMEM grows from - // 21.2 KB -> 21.9 KB for Dh=128, still well under the 46 KB threshold. - // Q_smem [16][Dh+8] T (4352B for Dh=128 bf16) - // K_smem[2][16][Dh+8] T (2 * 4352 = 8704B) - // V_smem[2][16][Dh+8] T (2 * 4352 = 8704B) - // P_smem [16][16] T (512B, not padded: small row) - // stats [16][2] float (128B) - // alpha_s [16] float (64B) - constexpr int SMEM_STRIDE = cascade_smem_row_stride(Dh); - extern __shared__ char smem_raw[]; - T* Q_smem = reinterpret_cast(smem_raw); - T* K_smem[2] = {Q_smem + BEAM_TILE * SMEM_STRIDE, Q_smem + BEAM_TILE * SMEM_STRIDE + TOKEN_TILE * SMEM_STRIDE}; - T* V_smem[2] - = {K_smem[1] + TOKEN_TILE * SMEM_STRIDE, K_smem[1] + TOKEN_TILE * SMEM_STRIDE + TOKEN_TILE * SMEM_STRIDE}; - T* P_smem = V_smem[1] + TOKEN_TILE * SMEM_STRIDE; - float* stats = reinterpret_cast(P_smem + BEAM_TILE * TOKEN_TILE); - float* alpha_s = stats + BEAM_TILE * 2; - - // ------------------------------ O accumulator ---------------------------- - // Per warp: one m16 row-tile times PV_N_BLOCKS n8 col-tiles. - // Each MMA C fragment holds 4 f32 per lane. - float O_acc[PV_N_BLOCKS][4]; -#pragma unroll - for (int n = 0; n < PV_N_BLOCKS; ++n) - { - O_acc[n][0] = 0.f; - O_acc[n][1] = 0.f; - O_acc[n][2] = 0.f; - O_acc[n][3] = 0.f; - } - - // ------------------------------ stats init ------------------------------- - if (tid < BEAM_TILE) - { - stats[tid * 2 + 0] = -FLT_MAX; - stats[tid * 2 + 1] = 0.f; - } - - // ------------------------------ Load Q to SMEM --------------------------- - // Layout Q_smem[beam][dh] row-major. Every thread contributes BEAM_TILE*Dh/THDS elements. - uint32_t const q_stride = params.stride ? static_cast(params.stride) : (num_heads * Dh); - { - constexpr int N_Q_ELEMS = BEAM_TILE * Dh; -#pragma unroll - for (int i = tid; i < N_Q_ELEMS; i += THDS) - { - int const b = i / Dh; - int const c = i - b * Dh; - int const beam_idx = beam_base + b; - T q_val; - if (beam_idx < beam_width) - { - int const q_offset = (req_idx * beam_width + beam_idx) * q_stride + head_idx * Dh + c; - q_val = params.q[q_offset]; - if (params.q_bias != nullptr) - { - q_val = common::cuda_cast( - common::cuda_cast(q_val) + common::cuda_cast(params.q_bias[head_idx * Dh + c])); - } - } - else - { - q_val = common::cuda_cast(0.f); - } - Q_smem[b * SMEM_STRIDE + c] = q_val; - } - } - __syncthreads(); - - // ------------------------------ RoPE on Q (NeoX full-head) --------------- - // cascade_eligible enforces rotary_embedding_dim == hidden_size_per_head for NeoX, - // so every beam covers all Dh channels. Each thread handles Dh/(2*THDS/BEAM_TILE) - // pairs per beam across the entire block. - bool const rope_on = (params.position_embedding_type == PositionEmbeddingType::kROPE_GPT_NEOX); - if (rope_on) - { - int const rotary_dim = params.rotary_embedding_dim; - int const half = rotary_dim >> 1; - float const rope_base = params.rotary_embedding_base; - float const rope_scale = params.rotary_embedding_scale; - int const q_pos = (params.length_per_sample != nullptr) ? (params.length_per_sample[req_idx * beam_width] - 1) - : params.timestep; - float2 const* cache_row = cascade_rope_cache_row(params.rotary_embedding_cos_sin_cache, q_pos, rotary_dim); - - int const n_pairs = BEAM_TILE * half; -#pragma unroll - for (int i = tid; i < n_pairs; i += THDS) - { - int const b = i / half; - int const c = i - b * half; // c \in [0, half) - int const beam_idx = beam_base + b; - if (beam_idx >= beam_width) - continue; - float2 const cs = cascade_rope_neox_cs(c, rotary_dim, q_pos, rope_base, rope_scale, cache_row); - float const v0 = common::cuda_cast(Q_smem[b * SMEM_STRIDE + c]); - float const v1 = common::cuda_cast(Q_smem[b * SMEM_STRIDE + c + half]); - Q_smem[b * SMEM_STRIDE + c] = common::cuda_cast(v0 * cs.x - v1 * cs.y); - Q_smem[b * SMEM_STRIDE + c + half] = common::cuda_cast(v1 * cs.x + v0 * cs.y); - } - __syncthreads(); - } - - // ===================================================================== - // Main loop over prefix tokens in tiles of TOKEN_TILE. - // - // while block computes on K/V_smem[tile_idx & 1], tile t+1's - // HBM->SMEM transfer is issued via cp.async into K/V_smem[(tile_idx+1)&1]. - // Per-iter timeline: - // 1. (if next exists) issue async load for tile t+1, commit -> 2 groups - // in flight; wait<1> to block until tile t's group is done. - // 2. __syncthreads(): ensure tile t's K/V and OOB zero-stores are visible. - // 3. compute Q.K / softmax / P.V on K_cur/V_cur. - // 4. __syncthreads(): tile t's reads complete before next iter overwrites - // that buffer as its `buf_nxt`. - // ===================================================================== - int const n_tiles = (prefix_len + TOKEN_TILE - 1) / TOKEN_TILE; - - // ---- Prologue: kick off tile 0 into buffer 0 ---- - if (n_tiles > 0) - { - int const t0_0 = 0; - int const tile_end_0 = min(TOKEN_TILE, prefix_len); - cascade_async_load_kv_tile( - kv_cache_buffer, owner_seq, kv_head_idx, t0_0, tile_end_0, tid, K_smem[0]); - cascade_async_load_kv_tile( - kv_cache_buffer, owner_seq, kv_head_idx, t0_0, tile_end_0, tid, V_smem[0]); - mma::cp_async_commit(); - } - - for (int tile_idx = 0; tile_idx < n_tiles; ++tile_idx) - { - int const t0 = tile_idx * TOKEN_TILE; - int const buf_cur = tile_idx & 1; - int const buf_nxt = buf_cur ^ 1; - - // 1. Prefetch tile t+1 (if exists), then wait for tile t. - if (tile_idx + 1 < n_tiles) - { - int const t0_n = (tile_idx + 1) * TOKEN_TILE; - int const tile_end_n = min(t0_n + TOKEN_TILE, prefix_len); - cascade_async_load_kv_tile( - kv_cache_buffer, owner_seq, kv_head_idx, t0_n, tile_end_n, tid, K_smem[buf_nxt]); - cascade_async_load_kv_tile( - kv_cache_buffer, owner_seq, kv_head_idx, t0_n, tile_end_n, tid, V_smem[buf_nxt]); - mma::cp_async_commit(); - mma::cp_async_wait<1>(); // keep the just-issued prefetch in flight - } - else - { - mma::cp_async_wait<0>(); // drain on the last tile - } - __syncthreads(); - - T* const K_cur = K_smem[buf_cur]; - T* const V_cur = V_smem[buf_cur]; - - // ============ Q · K^T (warp 0) ============ - // C[M=beam, N=tok] = A[M, K=dh] * B[K=dh, N=tok] col-major. - // B col-major source layout: B[k, n] == K_cur[n][k]. - // 2 n-blocks (N=0..7, 8..15), 8 k-slices for Dh=128. - if (warp_id == 0) - { - float qk_acc[2][4]; - int const tm = lane_id >> 2; // owned rows: tm, tm+8 - int const tk4 = (lane_id & 3) << 1; // owned col base (0,2,4,6) -#pragma unroll - for (int n = 0; n < 2; ++n) - { - qk_acc[n][0] = 0.f; - qk_acc[n][1] = 0.f; - qk_acc[n][2] = 0.f; - qk_acc[n][3] = 0.f; - } - // Fetch A fragments from Q_smem per k-iter. -#pragma unroll - for (int kk = 0; kk < Dh; kk += 16) - { - unsigned a0 = mma::load_pack2(&Q_smem[tm * SMEM_STRIDE + kk + tk4]); - unsigned a1 = mma::load_pack2(&Q_smem[(tm + 8) * SMEM_STRIDE + kk + tk4]); - unsigned a2 = mma::load_pack2(&Q_smem[tm * SMEM_STRIDE + kk + tk4 + 8]); - unsigned a3 = mma::load_pack2(&Q_smem[(tm + 8) * SMEM_STRIDE + kk + tk4 + 8]); - - int const tn = tm; // col-index within n8 block (same lane pattern) -#pragma unroll - for (int n = 0; n < 2; ++n) - { - int const actual_tok = n * 8 + tn; - // K_cur is row-major [tok][dh]; B[k, n] = K_cur[actual_tok][k + kk]. - // b[0] covers k = tk4..tk4+1 (2 consecutive dh), b[1] covers k+8..k+9. - unsigned b0 = mma::load_pack2(&K_cur[actual_tok * SMEM_STRIDE + kk + tk4]); - unsigned b1 = mma::load_pack2(&K_cur[actual_tok * SMEM_STRIDE + kk + tk4 + 8]); - mma::mma_m16n8k16( - qk_acc[n][0], qk_acc[n][1], qk_acc[n][2], qk_acc[n][3], a0, a1, a2, a3, b0, b1); - } - } - - // ---------------- Online softmax (warp 0) ---------------- - // Owned cells per lane: rows (tm, tm+8), cols within n-block: - // c = 2*(lane%4)..2*(lane%4)+1 and c+8..c+9 - int const col0 = (lane_id & 3) << 1; // 0,2,4,6 - // 1. Apply inv_sqrt_dh and mask invalid tokens to -FLT_MAX. -#pragma unroll - for (int n = 0; n < 2; ++n) - { - int const col_lo = n * 8 + col0; - bool const v0v = (t0 + col_lo) < prefix_len; - bool const v1v = (t0 + col_lo + 1) < prefix_len; - qk_acc[n][0] = v0v ? qk_acc[n][0] * inv_sqrt_dh : -FLT_MAX; - qk_acc[n][1] = v1v ? qk_acc[n][1] * inv_sqrt_dh : -FLT_MAX; - qk_acc[n][2] = v0v ? qk_acc[n][2] * inv_sqrt_dh : -FLT_MAX; - qk_acc[n][3] = v1v ? qk_acc[n][3] * inv_sqrt_dh : -FLT_MAX; - } - // 2. Row-max: for each of the two rows (tm, tm+8), reduce 16 cols. - float row_max_lo = fmaxf(fmaxf(qk_acc[0][0], qk_acc[0][1]), fmaxf(qk_acc[1][0], qk_acc[1][1])); - float row_max_hi = fmaxf(fmaxf(qk_acc[0][2], qk_acc[0][3]), fmaxf(qk_acc[1][2], qk_acc[1][3])); - // Reduce across the 4 lanes that share the same tm (lane_id/4). - row_max_lo = fmaxf(row_max_lo, __shfl_xor_sync(0xFFFFFFFFu, row_max_lo, 1)); - row_max_lo = fmaxf(row_max_lo, __shfl_xor_sync(0xFFFFFFFFu, row_max_lo, 2)); - row_max_hi = fmaxf(row_max_hi, __shfl_xor_sync(0xFFFFFFFFu, row_max_hi, 1)); - row_max_hi = fmaxf(row_max_hi, __shfl_xor_sync(0xFFFFFFFFu, row_max_hi, 2)); - - // 3. Combine with previous online stats. - float const prev_m_lo = stats[tm * 2 + 0]; - float const prev_l_lo = stats[tm * 2 + 1]; - float const prev_m_hi = stats[(tm + 8) * 2 + 0]; - float const prev_l_hi = stats[(tm + 8) * 2 + 1]; - float const new_m_lo = fmaxf(prev_m_lo, row_max_lo); - float const new_m_hi = fmaxf(prev_m_hi, row_max_hi); - // expf(-FLT_MAX - finite) underflows cleanly to 0, so no extra guard needed. - float const alpha_lo = (prev_m_lo <= -FLT_MAX * 0.5f) ? 0.f : __expf(prev_m_lo - new_m_lo); - float const alpha_hi = (prev_m_hi <= -FLT_MAX * 0.5f) ? 0.f : __expf(prev_m_hi - new_m_hi); - - // 4. Compute p = exp(qk - new_m) per cell, track row-sum. - float p[2][4]; -#pragma unroll - for (int n = 0; n < 2; ++n) - { - p[n][0] = (qk_acc[n][0] <= -FLT_MAX * 0.5f) ? 0.f : __expf(qk_acc[n][0] - new_m_lo); - p[n][1] = (qk_acc[n][1] <= -FLT_MAX * 0.5f) ? 0.f : __expf(qk_acc[n][1] - new_m_lo); - p[n][2] = (qk_acc[n][2] <= -FLT_MAX * 0.5f) ? 0.f : __expf(qk_acc[n][2] - new_m_hi); - p[n][3] = (qk_acc[n][3] <= -FLT_MAX * 0.5f) ? 0.f : __expf(qk_acc[n][3] - new_m_hi); - } - float sum_lo = p[0][0] + p[0][1] + p[1][0] + p[1][1]; - float sum_hi = p[0][2] + p[0][3] + p[1][2] + p[1][3]; - sum_lo += __shfl_xor_sync(0xFFFFFFFFu, sum_lo, 1); - sum_lo += __shfl_xor_sync(0xFFFFFFFFu, sum_lo, 2); - sum_hi += __shfl_xor_sync(0xFFFFFFFFu, sum_hi, 1); - sum_hi += __shfl_xor_sync(0xFFFFFFFFu, sum_hi, 2); - float const new_l_lo = prev_l_lo * alpha_lo + sum_lo; - float const new_l_hi = prev_l_hi * alpha_hi + sum_hi; - - // 5. Commit per-row outputs (one lane per row writes). - if ((lane_id & 3) == 0) - { - stats[tm * 2 + 0] = new_m_lo; - stats[tm * 2 + 1] = new_l_lo; - stats[(tm + 8) * 2 + 0] = new_m_hi; - stats[(tm + 8) * 2 + 1] = new_l_hi; - alpha_s[tm] = alpha_lo; - alpha_s[tm + 8] = alpha_hi; - } - // 6. Write P_smem [beam][tok]. All 32 lanes cover the 256 cells. -#pragma unroll - for (int n = 0; n < 2; ++n) - { - int const col_lo = n * 8 + col0; - P_smem[tm * TOKEN_TILE + col_lo] = common::cuda_cast(p[n][0]); - P_smem[tm * TOKEN_TILE + col_lo + 1] = common::cuda_cast(p[n][1]); - P_smem[(tm + 8) * TOKEN_TILE + col_lo] = common::cuda_cast(p[n][2]); - P_smem[(tm + 8) * TOKEN_TILE + col_lo + 1] = common::cuda_cast(p[n][3]); - } - } // warp 0 - __syncthreads(); - - // ============ Rescale O_acc + P · V (all warps) ============ - int const dh_base = warp_id * DH_PER_WARP; - { - int const tm = lane_id >> 2; - float const alpha_lo = alpha_s[tm]; - float const alpha_hi = alpha_s[tm + 8]; -#pragma unroll - for (int n = 0; n < PV_N_BLOCKS; ++n) - { - O_acc[n][0] *= alpha_lo; - O_acc[n][1] *= alpha_lo; - O_acc[n][2] *= alpha_hi; - O_acc[n][3] *= alpha_hi; - } - } - - // Fetch A fragment from P_smem (shared by all warps). - // P_smem shape [beam=16][tok=16] row-major => matches A row-major m16 k16. - unsigned pa0, pa1, pa2, pa3; - { - int const tm = lane_id >> 2; - int const tk4 = (lane_id & 3) << 1; - pa0 = mma::load_pack2(&P_smem[tm * TOKEN_TILE + tk4]); - pa1 = mma::load_pack2(&P_smem[(tm + 8) * TOKEN_TILE + tk4]); - pa2 = mma::load_pack2(&P_smem[tm * TOKEN_TILE + tk4 + 8]); - pa3 = mma::load_pack2(&P_smem[(tm + 8) * TOKEN_TILE + tk4 + 8]); - } - // PV MMA: C[M=beam, N=dh_slice] += A[M,K=tok] * B[K=tok, N=dh] col-major. - // V_cur[tok][dh] row-major => B[k, n] = V_cur[k][dh_base + n*8 + tn]. - // These span 2 different tok rows (stride Dh) and must be manually packed. - { - int const tn = lane_id >> 2; // 0..7 (N within n8 block) - int const tk4 = (lane_id & 3) << 1; // K-row base (0,2,4,6) -#pragma unroll - for (int n = 0; n < PV_N_BLOCKS; ++n) - { - int const actual_dh = dh_base + n * 8 + tn; - T v00 = V_cur[tk4 * SMEM_STRIDE + actual_dh]; - T v01 = V_cur[(tk4 + 1) * SMEM_STRIDE + actual_dh]; - T v10 = V_cur[(tk4 + 8) * SMEM_STRIDE + actual_dh]; - T v11 = V_cur[(tk4 + 9) * SMEM_STRIDE + actual_dh]; - unsigned vb0 = mma::pack2(v00, v01); - unsigned vb1 = mma::pack2(v10, v11); - mma::mma_m16n8k16(O_acc[n][0], O_acc[n][1], O_acc[n][2], O_acc[n][3], pa0, pa1, pa2, pa3, vb0, vb1); - } - } - __syncthreads(); - } // end prefix tile loop - - // ===================================================================== - // Write out: partial_out[(req*beam + beam_idx) * num_heads + head][dh] - // partial_m/l[(req*beam + beam_idx) * num_heads + head] - // - // Thread ownership (per warp): - // tm = lane/4 -> two beams: beam_base + tm, beam_base + tm + 8 - // col0 = 2*(lane%4) -> dh offsets within each n8 block: col0, col0+1 - // across PV_N_BLOCKS n8 blocks within the warp's DH_PER_WARP range. - // ===================================================================== - { - int const tm = lane_id >> 2; - int const col0 = (lane_id & 3) << 1; - int const beam_lo = beam_base + tm; - int const beam_hi = beam_base + tm + 8; - int const dh_base = warp_id * DH_PER_WARP; - -#pragma unroll - for (int n = 0; n < PV_N_BLOCKS; ++n) - { - int const dh_lo = dh_base + n * 8 + col0; - int const dh_hi = dh_lo + 1; - if (beam_lo < beam_width) - { - int const row = (req_idx * beam_width + beam_lo) * num_heads + head_idx; - partial_out[row * Dh + dh_lo] = O_acc[n][0]; - partial_out[row * Dh + dh_hi] = O_acc[n][1]; - } - if (beam_hi < beam_width) - { - int const row = (req_idx * beam_width + beam_hi) * num_heads + head_idx; - partial_out[row * Dh + dh_lo] = O_acc[n][2]; - partial_out[row * Dh + dh_hi] = O_acc[n][3]; - } - } - - // Stats: only warp 0's (lane%4 == 0) threads (one per beam row). - if (warp_id == 0 && (lane_id & 3) == 0) - { - if (beam_lo < beam_width) - { - int const row = (req_idx * beam_width + beam_lo) * num_heads + head_idx; - partial_m[row] = stats[tm * 2 + 0]; - partial_l[row] = stats[tm * 2 + 1]; - } - if (beam_hi < beam_width) - { - int const row = (req_idx * beam_width + beam_hi) * num_heads + head_idx; - partial_m[row] = stats[(tm + 8) * 2 + 0]; - partial_l[row] = stats[(tm + 8) * 2 + 1]; - } - } - } -#endif // __CUDA_ARCH__ < 800 -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// -// -// Phase 2: per-beam suffix decode. -// -// Grid: ( num_heads, batch_size * beam_width ) -// Block: ( Dh threads ) -// -// Walks the suffix tokens [prefix_len, T) following `cache_indir` to resolve -// the source beam for each cached token. -// -//////////////////////////////////////////////////////////////////////////////////////////////////// - -template -// fusion: suffix kernel now merges the prefix partial state in-register at -// the end and writes the final params.out directly. This eliminates Phase 3 -// (cascade_merge_kernel) and saves a full kernel launch + one DRAM write + -// three DRAM reads per (seq, head). The suffix partial values (m_s, l_s, -// v_acc) never leave registers. -__global__ void cascade_suffix_decode_kernel(Multihead_attention_params params, KVCacheBuffer kv_cache_buffer, - int const* __restrict__ d_input_lengths, float const* __restrict__ partial_out_p, - float const* __restrict__ partial_m_p, float const* __restrict__ partial_l_p) -{ -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 - if (blockIdx.x == 0 && blockIdx.y == 0 && blockIdx.z == 0 && threadIdx.x == 0 && threadIdx.y == 0 - && threadIdx.z == 0) - { - printf("cascade_suffix_decode_kernel: requires SM_80+ but launched on an unsupported device.\n"); - __trap(); - } - return; -#else - using Cfg = CascadeConfig; - constexpr int THDS = Cfg::THDS_PER_BLOCK; - - int const head_idx = blockIdx.x; - int const seq = blockIdx.y; // = req_idx * beam_width + beam_idx - int const tid = threadIdx.x; - - int const beam_width = params.beam_width; - int const req_idx = seq / beam_width; - int const beam_idx = seq % beam_width; - int const num_heads = params.num_heads; - int const num_kv_heads = params.num_kv_heads; - float const inv_sqrt_dh = params.inv_sqrt_dh; - - // Read prefix_len from device memory (Graph-safe: no host copy needed). - int const prefix_len = __ldg(&d_input_lengths[req_idx * beam_width]); - - if (params.hidden_size_per_head != Dh) - { - return; - } - - // Standard MMHA convention: length_per_sample INCLUDES the current timestep, - // but we only attend over PAST tokens. Subtract 1 to get the KV cache length. - int const tlength = (params.length_per_sample != nullptr) ? (params.length_per_sample[seq] - 1) : params.timestep; - int const kv_head_idx = (num_kv_heads > 0) ? (head_idx * num_kv_heads / num_heads) : head_idx; - - extern __shared__ float smem[]; - float* reduce = smem; // THDS floats - float* q_smem = reduce + THDS; // Dh floats - float* k_smem = q_smem + Dh; // Dh floats (for K RoPE partner exchange) - - // Load Q for this (head, beam) into shared mem so threads can read each - // other's channels during RoPE rotation. - // NOTE: Must use params.stride (packed QKV per-sample stride) not num_heads*Dh. - uint32_t const q_stride = params.stride ? static_cast(params.stride) : (num_heads * Dh); - int const q_offset = seq * q_stride + head_idx * Dh + tid; - q_smem[tid] = common::cuda_cast(params.q[q_offset]); - if (params.q_bias != nullptr) - { - q_smem[tid] += common::cuda_cast(params.q_bias[head_idx * Dh + tid]); - } - __syncthreads(); - - // Apply RoPE to Q if enabled. Cached K is already post-RoPE. - if (params.position_embedding_type == PositionEmbeddingType::kROPE_GPT_NEOX) - { - int const rotary_dim = params.rotary_embedding_dim; - int const half = rotary_dim / 2; - float const rope_base = params.rotary_embedding_base; - float const rope_scale = params.rotary_embedding_scale; - int const q_pos = (params.length_per_sample != nullptr) ? (params.length_per_sample[seq] - 1) : params.timestep; - float2 const* cache_row = cascade_rope_cache_row(params.rotary_embedding_cos_sin_cache, q_pos, rotary_dim); - - // Each thread owns channel `tid` of q_smem. Threads with tid >= rotary_dim - // are inactive but must still hit the __syncthreads() below. - bool const active = (tid < rotary_dim); - int const freq_idx = active ? ((tid < half) ? tid : (tid - half)) : 0; - float2 cs = make_float2(1.f, 0.f); - float val = 0.f; - float partner = 0.f; - if (active) - { - cs = cascade_rope_neox_cs(freq_idx, rotary_dim, q_pos, rope_base, rope_scale, cache_row); - val = q_smem[tid]; - partner = (tid < half) ? q_smem[tid + half] : q_smem[tid - half]; - } - __syncthreads(); - if (active) - { - q_smem[tid] = (tid < half) ? (val * cs.x - partner * cs.y) : (val * cs.x + partner * cs.y); - } - __syncthreads(); - } - - float const q_val = q_smem[tid]; - - // ===================================================================== - // Fused Phase 0: compute the current step's K/V (bias + RoPE) and have - // the leader Q-head in each GQA group persist them to the KV cache so - // the next decode step sees them. The per-channel values are also - // kept in registers (k_cur / v_cur) and reused for the tok==tlength - // round of the attention loop, avoiding an HBM write+read round-trip. - // ===================================================================== - float k_cur = 0.f; - float v_cur = 0.f; - { - int const num_kv_heads_eff = (num_kv_heads > 0) ? num_kv_heads : num_heads; - int const kv_group = num_heads / num_kv_heads_eff; - bool const is_leader = ((head_idx % kv_group) == 0); - - // K/V share the same packed QKV per-sample stride as Q. - uint32_t const kv_stride - = params.stride ? static_cast(params.stride) : static_cast(num_kv_heads_eff * Dh); - int const k_off = seq * kv_stride + kv_head_idx * Dh + tid; - int const v_off = seq * kv_stride + kv_head_idx * Dh + tid; - k_cur = common::cuda_cast(params.k[k_off]); - v_cur = common::cuda_cast(params.v[v_off]); - - if (params.k_bias != nullptr) - { - k_cur += common::cuda_cast(params.k_bias[kv_head_idx * Dh + tid]); - } - if (params.v_bias != nullptr) - { - v_cur += common::cuda_cast(params.v_bias[kv_head_idx * Dh + tid]); - } - - // Apply NeoX RoPE to K. Matches baseline numerics by going through the - // shared mmha::rotary_embedding_coefficient helper, which prefers the - // framework-provided cos_sin_cache when available and falls back to - // computing inv_freq from base/scale otherwise. - if (params.position_embedding_type == PositionEmbeddingType::kROPE_GPT_NEOX) - { - k_smem[tid] = k_cur; - __syncthreads(); - - int const rotary_dim = params.rotary_embedding_dim; - if (tid < rotary_dim) - { - int const half = rotary_dim / 2; - int const freq_idx = (tid < half) ? tid : (tid - half); - float2 const* cache_row - = cascade_rope_cache_row(params.rotary_embedding_cos_sin_cache, tlength, rotary_dim); - float2 const cs = cascade_rope_neox_cs(freq_idx, rotary_dim, tlength, params.rotary_embedding_base, - params.rotary_embedding_scale, cache_row); - float const val = k_smem[tid]; - float const partner = (tid < half) ? k_smem[tid + half] : k_smem[tid - half]; - k_cur = (tid < half) ? (val * cs.x - partner * cs.y) : (val * cs.x + partner * cs.y); - } - __syncthreads(); - } - - // Leader (one Q-head per GQA group) persists current K/V to cache. - if (is_leader) - { - auto const localTokenIdx = kv_cache_buffer.getKVTokenIdx(tlength); - auto* kPtr = reinterpret_cast(kv_cache_buffer.getKBlockPtr(seq, localTokenIdx)); - auto* vPtr = reinterpret_cast(kv_cache_buffer.getVBlockPtr(seq, localTokenIdx)); - auto const off = kv_cache_buffer.getKVLocalIdx(localTokenIdx, kv_head_idx, Dh, tid); - kPtr[off] = common::cuda_cast(k_cur); - vPtr[off] = common::cuda_cast(v_cur); - } - // No __threadfence needed: the next decode step is launched on the - // same stream and stream ordering guarantees cross-kernel visibility. - } - - float m = -FLT_MAX; - float l = 0.f; - float v_acc = 0.f; - - int const max_attention_window = params.max_attention_window_size; - int const* cache_indir = params.cache_indir; - - // tlength = position of the current decode token in the cache. - // Attention covers all suffix tokens INCLUDING the current one. Past - // tokens [prefix_len, tlength) are read from cache via cache_indir; - // the current token (tok == tlength) is handled separately below using - // k_cur / v_cur that we computed and (if leader) wrote to cache above. - for (int tok = prefix_len; tok < tlength; ++tok) - { - // Resolve the physical beam this cached token came from via - // cache_indir (beam search's indirection table). - int src_beam = beam_idx; - if (cache_indir != nullptr) - { - int const indir_offset - = req_idx * beam_width * max_attention_window + beam_idx * max_attention_window + tok; - src_beam = cache_indir[indir_offset]; - } - int const src_seq = req_idx * beam_width + src_beam; - - float k_val = common::cuda_cast( - load_kv(kv_cache_buffer, src_seq, tok, kv_head_idx, tid)); - float v_val = common::cuda_cast( - load_kv(kv_cache_buffer, src_seq, tok, kv_head_idx, tid)); - - float partial = q_val * k_val; - float qk = block_sum(partial, reduce) * inv_sqrt_dh; - - float new_m = fmaxf(m, qk); - float scale = expf(m - new_m); - float p = expf(qk - new_m); - l = l * scale + p; - v_acc = v_acc * scale + p * v_val; - m = new_m; - } - - // Current step (tok == tlength): use K/V from registers — cache_indir - // has not yet been updated for this position (beam search runs AFTER - // attention), so the current step is always self-pointing. - { - float partial = q_val * k_cur; - float qk = block_sum(partial, reduce) * inv_sqrt_dh; - - float new_m = fmaxf(m, qk); - float scale = expf(m - new_m); - float p = expf(qk - new_m); - l = l * scale + p; - v_acc = v_acc * scale + p * v_cur; - m = new_m; - } - - // ------------------------------------------------------------------- - // fusion: merge with prefix partial (formerly cascade_merge_kernel) - // ------------------------------------------------------------------- - // Suffix partial is in registers: m (suffix max), l (suffix denom), v_acc - // (suffix weighted V sum). Prefix partial comes from Phase 1's workspace. - int const out_row = seq * num_heads + head_idx; - float const m_s = m; - float const l_s = l; - float const v_s = v_acc; - - float const m_p = __ldg(&partial_m_p[out_row]); - float const l_p = __ldg(&partial_l_p[out_row]); - float const v_p = __ldg(&partial_out_p[out_row * Dh + tid]); - - // Empty-side guards preserve semantics when prefix or suffix contributes - // zero tokens (e.g. first decode step, or prefix_len > tlength). - bool const has_p = (l_p > 0.f) && (m_p > -FLT_MAX / 2.f); - bool const has_s = (l_s > 0.f) && (m_s > -FLT_MAX / 2.f); - - float out_val; - if (has_p && has_s) - { - float const new_m = fmaxf(m_p, m_s); - float const ep = expf(m_p - new_m); - float const es = expf(m_s - new_m); - float const new_l = l_p * ep + l_s * es; - out_val = (v_p * ep + v_s * es) / fmaxf(new_l, 1e-30f); - } - else if (has_p) - { - out_val = v_p / fmaxf(l_p, 1e-30f); - } - else if (has_s) - { - out_val = v_s / fmaxf(l_s, 1e-30f); - } - else - { - out_val = 0.f; - } - - int const out_offset = seq * num_heads * Dh + head_idx * Dh + tid; - reinterpret_cast(params.out)[out_offset] = common::cuda_cast(out_val); -#endif // __CUDA_ARCH__ < 800 -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// -// -// Eligibility & launcher. -// -// Note: The former Phase 3 `cascade_merge_kernel` has been removed. Its -// online-softmax merge logic is now fused into the tail of -// `cascade_suffix_decode_kernel` , so the merge -// happens in registers without an extra kernel launch / DRAM round-trip. -// -//////////////////////////////////////////////////////////////////////////////////////////////////// - -template -bool cascade_eligible(KernelParamsType const& params) -{ - if constexpr (KernelParamsType::DO_CROSS_ATTENTION) - { - return false; - } - if (!tensorrt_llm::common::getEnvEnableCascadeMmha()) - { - return false; - } - { - static std::atomic cached_sm{-1}; - int sm = cached_sm.load(std::memory_order_relaxed); - if (sm < 0) - { - sm = tensorrt_llm::common::getSMVersion(); - cached_sm.store(sm, std::memory_order_relaxed); - } - if (sm < 80) - { - return false; - } - } - if (params.position_shift_enabled || params.block_sparse_attention) - { - return false; - } - // v0.1 supports LEARNED_ABSOLUTE and full-head GPT-NeoX RoPE (Qwen / Llama / - // Mistral style). Other variants (GPTJ interleaved, YARN, LONG_ROPE, - // M-RoPE, ALIBI, ...) fall back to the baseline MMHA path. - if (params.position_embedding_type != PositionEmbeddingType::kLEARNED_ABSOLUTE - && params.position_embedding_type != PositionEmbeddingType::kROPE_GPT_NEOX) - { - return false; - } - if (params.position_embedding_type == PositionEmbeddingType::kROPE_GPT_NEOX) - { - // Require full-head rotation so the pair (c, c + Dh/2) is resolvable - // in-place within shared memory. - if (params.rotary_embedding_dim != params.hidden_size_per_head) - { - return false; - } - // Dynamic / long / yarn / m-scaling require cached cos-sin tables and - // per-request base updates that v0.1 does not implement. - if (params.rotary_embedding_scale_type != RotaryScalingType::kNONE) - { - return false; - } - // When scale_type == kNONE, the cos_sin_cache (when provided by the - // framework) is consumed via mmha::rotary_embedding_coefficient inside - // the kernel; if it is null, we fall back to computing cos/sin on the - // fly from base/scale, matching the baseline numerics. inv_freq_cache - // is currently ignored (we recompute inv_freq from base) but this is - // numerically equivalent under kNONE scaling. - if (params.mrope_position_deltas != nullptr) - { - return false; - } - } - if (params.attn_logit_softcapping_scale != 0.0f) - { - return false; - } - if (params.relative_attention_bias != nullptr || params.linear_bias_slopes != nullptr) - { - return false; - } - if (params.attention_mask != nullptr || params.attention_sinks != nullptr) - { - return false; - } - if (params.logn_scaling_ptr != nullptr) - { - return false; - } - if (params.ia3_tasks != nullptr) - { - return false; - } - if (params.timestep >= params.cyclic_attention_window_size) - { - return false; - } - if (params.timestep >= params.chunked_attention_size) - { - return false; - } - if (params.int8_kv_cache || params.fp8_kv_cache) - { - return false; - } - if (params.input_lengths == nullptr || params.length_per_sample == nullptr) - { - return false; - } - int const dh = params.hidden_size_per_head; - if (dh != 128) - { - return false; - } - - return true; -} - -namespace -{ - -// Cascade prefix-side per-token O/m/l accumulator footprints. Mirrors the -// public CascadeWorkspaceSizes reported via cascadeAttentionKernel.h: out is -// the per-token fp32 O accumulator, mMax / lSum are the running max and -// sum-of-exp. The values are independent of T because the partials are -// always materialized in fp32 inside the kernel. -constexpr size_t cascade_workspace_out_bytes(int batch_beam, int num_heads, int head_size) noexcept -{ - return static_cast(batch_beam) * num_heads * head_size * sizeof(float); -} - -constexpr size_t cascade_workspace_stat_bytes(int batch_beam, int num_heads) noexcept -{ - return static_cast(batch_beam) * num_heads * sizeof(float); -} - -} // namespace - -CascadeWorkspaceSizes getCascadeWorkspaceSizes(int batch_beam, int num_heads, int head_size) noexcept -{ - CascadeWorkspaceSizes s{}; - if (batch_beam <= 0 || num_heads <= 0 || head_size <= 0) - { - return s; - } - s.out = cascade_workspace_out_bytes(batch_beam, num_heads, head_size); - s.mMax = cascade_workspace_stat_bytes(batch_beam, num_heads); - s.lSum = cascade_workspace_stat_bytes(batch_beam, num_heads); - return s; -} - -template -bool launch_cascade_attention( - Multihead_attention_params const& params, KVCacheBuffer const& kv_cache_buffer, cudaStream_t stream) -{ - static_assert(std::is_same_v, "cascade kernel requires T_cache == T"); - - // IMPORTANT: In TRT-LLM MMHA, params.batch_size = total_sequences = - // num_requests × beam_width. Our cascade design separates the "shared - // prefix" (per-request) from "per-beam suffix", so we need num_requests. - int const total_seqs = params.batch_size; // = num_requests × beam_width - int const beam = params.beam_width; - int const num_requests = total_seqs / beam; - int const num_heads = params.num_heads; - int const dh = params.hidden_size_per_head; - - // Compute workspace layout: 3 buffers packed contiguously. - // Each buffer is indexed by [total_seqs × num_heads (× Dh for out)]. - size_t const out_elems = static_cast(total_seqs) * num_heads * dh; - size_t const stat_elems = static_cast(total_seqs) * num_heads; - size_t const out_bytes = out_elems * sizeof(float); - size_t const stat_bytes = stat_elems * sizeof(float); - - float* const ws_out_p = params.cascade_partial_out; - float* const ws_m_p = params.cascade_partial_max; - float* const ws_l_p = params.cascade_partial_sum; - if (ws_out_p == nullptr || ws_m_p == nullptr || ws_l_p == nullptr) - { - TLLM_LOG_WARNING( - "cascade_attention: cascade workspace not provisioned by AttentionOp (need %zu bytes), falling back", - out_bytes + 2 * stat_bytes); - return false; - } - - // The prefix_len is NOT fetched to host. Instead, the device pointer - // params.input_lengths is passed directly to the kernels, and each kernel - // reads the shared prefix length via __ldg (a single cached global load). - // This keeps the entire launch path free of D2H copies and stream-sync, - // making it fully compatible with CUDA Graph capture. - int const* d_input_lengths = params.input_lengths; - - using Cfg = CascadeConfig; - constexpr int THDS = Cfg::THDS_PER_BLOCK; - constexpr int BEAM_TILE = Cfg::BEAM_TILE; - - // Phase 1: shared-prefix attention (all beams share the same KV prefix). - // Grid Y = num_requests (NOT total_seqs!): each Y-block handles one request's - // shared prefix, processing BEAM_TILE beams per Z-block. - { - dim3 grid(num_heads, num_requests, (beam + BEAM_TILE - 1) / BEAM_TILE); - dim3 block(THDS); - // Tensor-Core SMEM layout: Q + K[2] + V[2] + P (bf16/half), then - // stats/alpha (fp32). K/V double-buffered for cp.async pipelining. - // P0-1: Q/K/V rows carry 16B of padding to break bank conflicts, so - // account for it via cascade_smem_row_stride(Dh) here as well. - constexpr int SMEM_STRIDE = cascade_smem_row_stride(Dh); - size_t const t_bytes = (BEAM_TILE * SMEM_STRIDE // Q_smem - + 2 * Cfg::TOKEN_TILE * SMEM_STRIDE // K_smem[2] - + 2 * Cfg::TOKEN_TILE * SMEM_STRIDE // V_smem[2] - + BEAM_TILE * Cfg::TOKEN_TILE) // P_smem (not padded) - * sizeof(T); - size_t const f_bytes = (BEAM_TILE * 2 // stats(m, l) - + BEAM_TILE) // alpha_s - * sizeof(float); - size_t const smem_bytes = t_bytes + f_bytes; - TLLM_CUDA_CHECK(cudaFuncSetAttribute(cascade_prefix_mqa_kernel, - cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(smem_bytes))); - cascade_prefix_mqa_kernel - <<>>(params, kv_cache_buffer, d_input_lengths, ws_out_p, ws_m_p, ws_l_p); - } - - // Phase 2 (FUSED): Phase-0 write_kv + per-beam suffix decode + in-register - // merge with prefix. Grid Y = total_seqs: one block per (request, beam). - // The leader Q-head in each GQA group writes the current step's K/V to - // cache; the other Q-heads keep K/V only in registers for their own use. - { - dim3 grid(num_heads, total_seqs); - dim3 block(THDS); - // reduce(THDS) + q_smem(Dh) + k_smem(Dh, for K RoPE partner exchange) - size_t const smem_bytes = (THDS + 2 * Dh) * sizeof(float); - cascade_suffix_decode_kernel - <<>>(params, kv_cache_buffer, d_input_lengths, ws_out_p, ws_m_p, ws_l_p); - } - - return true; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// -// -// Explicit instantiations. -// -//////////////////////////////////////////////////////////////////////////////////////////////////// - -#define INSTANTIATE_CASCADE(T, KVB, DH) \ - template bool launch_cascade_attention( \ - Multihead_attention_params const&, KVB const&, cudaStream_t); - -INSTANTIATE_CASCADE(half, KVLinearBuffer, 128) -INSTANTIATE_CASCADE(half, KVBlockArray, 128) - -#ifdef ENABLE_BF16 -INSTANTIATE_CASCADE(__nv_bfloat16, KVLinearBuffer, 128) -INSTANTIATE_CASCADE(__nv_bfloat16, KVBlockArray, 128) -#endif - -template bool cascade_eligible>(Masked_multihead_attention_params const&); -#ifdef ENABLE_BF16 -template bool cascade_eligible>( - Masked_multihead_attention_params<__nv_bfloat16> const&); -#endif -template bool cascade_eligible>( - Masked_multihead_attention_params const&); - -#undef INSTANTIATE_CASCADE - -} // namespace cascade -} // namespace mmha -} // namespace kernels - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/cascadeAttentionKernel.h b/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/cascadeAttentionKernel.h deleted file mode 100644 index 5fcca9de75ac..000000000000 --- a/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/cascadeAttentionKernel.h +++ /dev/null @@ -1,100 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2024-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. - */ -#pragma once - -#include "tensorrt_llm/common/config.h" -#include "tensorrt_llm/kernels/decoderMaskedMultiheadAttention.h" -#include "tensorrt_llm/kernels/kvCacheUtils.h" - -#include - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels -{ -namespace mmha -{ -namespace cascade -{ - -//////////////////////////////////////////////////////////////////////////////////////////////////// -// -// Cascade Attention Kernel -// -// Implements the divide-and-conquer attention from -// "Cascade Inference: Memory Bandwidth Efficient Shared Prefix Batch Decoding" -// (https://flashinfer.ai/2024/02/02/cascade-inference.html), specialized to -// beam-search decode and fused down to two kernels (the original Phase 0 -// KV-write and Phase 3 merge are absorbed into the suffix-decode kernel). -// -// Phase 1 (cascade_prefix_mqa_kernel): -// For each request, all `beam_width` queries attend to the *shared* prompt -// KV (token range [0, L_p)). KV is loaded once per (head, request, token tile) -// into shared memory and reused by every beam, eliminating the O(beam) HBM -// traffic that the baseline `masked_multihead_attention_kernel` pays. -// -// Phase 2 (cascade_suffix_decode_kernel): -// Each (head, beam) pair processes its own suffix KV [L_p, T) following -// `cache_indir`. This is a regular single-query decode with no sharing. -// The numerically stable online-softmax merge with the prefix partial -// state is fused at the end of this kernel (in-register), so no separate -// merge kernel launch is needed. -// -// This first version intentionally targets a *narrow* feature subset and -// falls back to MMHA otherwise: -// - DO_CROSS_ATTENTION = false -// - POS_SHIFT = false -// - BLOCK_SPARSE_ATTN = false -// - IMPLICIT_REL_ATTN = false -// - ATTN_LOGIT_SOFTCAP = off -// - PositionEmbedding in { LEARNED_ABSOLUTE, ROPE_GPT_NEOX } with full-head -// rotation and RotaryScalingType::kNONE (covers Qwen, -// Llama3, Mistral default configs). -// - T_cache = T (no INT8 / FP8 KV cache yet) -// - Dh = 128 -// - T in { half, __nv_bfloat16 } -// -// Activation is gated by the env var `TRTLLM_ENABLE_CASCADE_MMHA` (default off). -// -//////////////////////////////////////////////////////////////////////////////////////////////////// - -// Returns true if the cascade pipeline launched and produced the final output. -// Returns false if the caller should fall back to the standard MMHA path -// (either because gating is off or the params are unsupported). -template -bool launch_cascade_attention( - Multihead_attention_params const& params, KVCacheBuffer const& kv_cache_buffer, cudaStream_t stream); - -// Convenience predicate that does *not* allocate any workspace and only -// inspects host-side params. Used by the dispatcher to short-circuit cleanly. -template -bool cascade_eligible(KernelParamsType const& params); - -struct CascadeWorkspaceSizes -{ - size_t out{}; - size_t mMax{}; - size_t lSum{}; -}; - -CascadeWorkspaceSizes getCascadeWorkspaceSizes(int batch_beam, int num_heads, int head_size) noexcept; - -} // namespace cascade -} // namespace mmha -} // namespace kernels - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/cascadeMma.cuh b/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/cascadeMma.cuh deleted file mode 100644 index 476ae3ce139b..000000000000 --- a/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/cascadeMma.cuh +++ /dev/null @@ -1,220 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2024-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. - */ -#pragma once - -#include "tensorrt_llm/common/config.h" - -#include -#include -#include -#ifdef ENABLE_BF16 -#include -#endif - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels -{ -namespace mmha -{ -namespace cascade -{ -namespace mma -{ - -//////////////////////////////////////////////////////////////////////////////////////////////////// -// -// Low-level PTX wrappers for Tensor Core attention on SM80+. Adapted from -// `cpp/tensorrt_llm/kernels/selectiveScan/Common.h` and kept self-contained -// so the cascade module does not leak a dependency on the Mamba kernel dir. -// -// Supported ops (bf16 + f16 via mma): -// - mma.sync.aligned.m16n8k16.row.col.f32.{bf16|f16}... (16x8 output) -// - cp.async.ca.shared.global {.b4, .b8, .b16} (SM80+ pipelined) -// -//////////////////////////////////////////////////////////////////////////////////////////////////// - -// Convert a generic pointer into a 32-bit .shared address suitable for PTX. -__device__ __forceinline__ unsigned smem_addr_of(void const* smem_ptr) -{ - return static_cast(__cvta_generic_to_shared(smem_ptr)); -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// -// Element packing helpers: two 16-bit elements -> one uint32 (little-endian). -// Used to manually build MMA A/B fragments without ldmatrix. -//////////////////////////////////////////////////////////////////////////////////////////////////// - -template -__device__ __forceinline__ unsigned pack2(T a, T b); - -template <> -__device__ __forceinline__ unsigned pack2(half a, half b) -{ - unsigned short ua = __half_as_ushort(a); - unsigned short ub = __half_as_ushort(b); - return (static_cast(ub) << 16) | static_cast(ua); -} - -#ifdef ENABLE_BF16 -template <> -__device__ __forceinline__ unsigned pack2<__nv_bfloat16>(__nv_bfloat16 a, __nv_bfloat16 b) -{ - unsigned short ua = __bfloat16_as_ushort(a); - unsigned short ub = __bfloat16_as_ushort(b); - return (static_cast(ub) << 16) | static_cast(ua); -} -#endif - -// Read two 16-bit elements from SMEM at adjacent addresses (ptr[0], ptr[1]) as a -// packed uint32. Assumes 32-bit alignment of ptr (i.e. even index in T array). -template -__device__ __forceinline__ unsigned load_pack2(T const* ptr) -{ - return *reinterpret_cast(ptr); -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// -// mma.sync.aligned.m16n8k16 -// -// Thread-fragment layout (per PTX 8.x spec): -// A (M=16, K=16, row-major, bf16): -// Thread t owns 8 bf16 = 4 uint32: -// a[0] -> (row = t/4, col = 2*(t%4) + {0,1}) (within K=0..7) -// a[1] -> (row = t/4 + 8, col = 2*(t%4) + {0,1}) (within K=0..7) -// a[2] -> (row = t/4, col = 2*(t%4) + {8,9}) (within K=8..15) -// a[3] -> (row = t/4 + 8, col = 2*(t%4) + {8,9}) (within K=8..15) -// B (K=16, N=8, col-major, bf16): -// Thread t owns 4 bf16 = 2 uint32: -// b[0] -> (col = t/4, row = 2*(t%4) + {0,1}) covering K=0..7 -// b[1] -> (col = t/4, row = 2*(t%4) + {8,9}) covering K=8..15 -// C (M=16, N=8, f32): -// Thread t owns 4 f32: -// c[0] -> (row = t/4, col = 2*(t%4)) -// c[1] -> (row = t/4, col = 2*(t%4) + 1) -// c[2] -> (row = t/4 + 8, col = 2*(t%4)) -// c[3] -> (row = t/4 + 8, col = 2*(t%4) + 1) -//////////////////////////////////////////////////////////////////////////////////////////////////// - -__device__ __forceinline__ void mma_m16n8k16_bf16(float& c0, float& c1, float& c2, float& c3, unsigned a0, unsigned a1, - unsigned a2, unsigned a3, unsigned b0, unsigned b1) -{ -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 - asm volatile( - "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 \n" - "{%0, %1, %2, %3},\n" - "{%4, %5, %6, %7},\n" - "{%8, %9},\n" - "{%0, %1, %2, %3};\n" - : "+f"(c0), "+f"(c1), "+f"(c2), "+f"(c3) - : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1)); -#endif -} - -__device__ __forceinline__ void mma_m16n8k16_f16(float& c0, float& c1, float& c2, float& c3, unsigned a0, unsigned a1, - unsigned a2, unsigned a3, unsigned b0, unsigned b1) -{ -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 - asm volatile( - "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 \n" - "{%0, %1, %2, %3},\n" - "{%4, %5, %6, %7},\n" - "{%8, %9},\n" - "{%0, %1, %2, %3};\n" - : "+f"(c0), "+f"(c1), "+f"(c2), "+f"(c3) - : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1)); -#endif -} - -// Type-dispatched wrapper. T is the element type (__nv_bfloat16 or half). -template -__device__ __forceinline__ void mma_m16n8k16(float& c0, float& c1, float& c2, float& c3, unsigned a0, unsigned a1, - unsigned a2, unsigned a3, unsigned b0, unsigned b1); - -#ifdef ENABLE_BF16 -template <> -__device__ __forceinline__ void mma_m16n8k16<__nv_bfloat16>(float& c0, float& c1, float& c2, float& c3, unsigned a0, - unsigned a1, unsigned a2, unsigned a3, unsigned b0, unsigned b1) -{ - mma_m16n8k16_bf16(c0, c1, c2, c3, a0, a1, a2, a3, b0, b1); -} -#endif - -template <> -__device__ __forceinline__ void mma_m16n8k16(float& c0, float& c1, float& c2, float& c3, unsigned a0, unsigned a1, - unsigned a2, unsigned a3, unsigned b0, unsigned b1) -{ - mma_m16n8k16_f16(c0, c1, c2, c3, a0, a1, a2, a3, b0, b1); -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// -// cp.async (SM80+). Transfers from global to shared memory without staging -// through registers, enabling SW pipelining of HBM LDG with MMA compute. -//////////////////////////////////////////////////////////////////////////////////////////////////// - -// Issue a 16-byte async copy. smem_addr must be 16B aligned (shared-space), -// global_ptr must be 16B aligned (generic). -// -// Optional `src_size` (PTX `cp.async` 4th operand) controls hardware zero-fill: -// * src_size == 16 (default): full 16B copy, equivalent to omitting the operand. -// * src_size < 16 : bytes [src_size, 16) of the destination SMEM -// are zeroed by the LSU. -// * src_size == 0 : no global load is issued at all (per PTX ISA: -// "if src-size is zero, the access is a no-op"); -// the destination 16B are simply zeroed. -// Callers can therefore use a single uniform issue for both in-bounds and -// out-of-bounds chunks (passing 16 or 0 respectively) instead of branching -// on an explicit SMEM zero store. When `src_size` is the immediate `16`, -// ptxas emits the compact form without the `src-size` operand. -__device__ __forceinline__ void cp_async_16B(unsigned smem_addr, void const* global_ptr, unsigned src_size = 16u) -{ -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 - asm volatile("cp.async.cg.shared.global [%0], [%1], 16, %2;\n" ::"r"(smem_addr), "l"(global_ptr), "r"(src_size)); -#else - unsigned tmp[4] = {0u, 0u, 0u, 0u}; - if (src_size >= 16u) - { - asm volatile("ld.global.v4.b32 {%0, %1, %2, %3}, [%4];\n" - : "=r"(tmp[0]), "=r"(tmp[1]), "=r"(tmp[2]), "=r"(tmp[3]) - : "l"(global_ptr)); - } - asm volatile("st.shared.v4.b32 [%0], {%1, %2, %3, %4};\n" ::"r"(smem_addr), "r"(tmp[0]), "r"(tmp[1]), "r"(tmp[2]), - "r"(tmp[3])); -#endif -} - -__device__ __forceinline__ void cp_async_commit() -{ -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 - asm volatile("cp.async.commit_group;\n"); -#endif -} - -template -__device__ __forceinline__ void cp_async_wait() -{ -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 - asm volatile("cp.async.wait_group %0;\n" ::"n"(remain)); -#endif -} - -} // namespace mma -} // namespace cascade -} // namespace mmha -} // namespace kernels - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderMaskedMultiheadAttentionLaunch.h b/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderMaskedMultiheadAttentionLaunch.h index f36d3e61f23c..bf6b22385ecc 100644 --- a/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderMaskedMultiheadAttentionLaunch.h +++ b/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderMaskedMultiheadAttentionLaunch.h @@ -15,7 +15,6 @@ */ #pragma once -#include "cascadeAttentionKernel.h" #include "decoderMaskedMultiheadAttentionTemplate.h" #include "tensorrt_llm/common/assert.h" #include "tensorrt_llm/common/config.h" @@ -423,25 +422,6 @@ void mmha_launch_kernel(KernelParamsType const& params, KVCacheBuffer const& kv_ || params.position_embedding_type == PositionEmbeddingType::kROPE_GPTJ || params.position_embedding_type == PositionEmbeddingType::kLONG_ROPE || params.position_embedding_type == PositionEmbeddingType::kROPE_M)); - // Try cascade-attention path before falling through to the standard MMHA. - // Restricted to self-attention, beam search, narrow type/Dh subset and - // gated by env vars; otherwise short-circuits cheaply. - constexpr bool kCascadeTypeOk = std::is_same_v -#ifdef ENABLE_BF16 - || std::is_same_v -#endif - ; - if constexpr (!KernelParamsType::DO_CROSS_ATTENTION && !BLOCK_SPARSE_ATTN && !IMPLICIT_REL_ATTN_BIAS - && !ATTN_LOGIT_SOFTCAPPING && (Dh == 128) && kCascadeTypeOk) - { - if (params.beam_width > 1 && cascade::cascade_eligible(params)) - { - if (cascade::launch_cascade_attention(params, kv_cache_buffer, stream)) - { - return; - } - } - } if (params.beam_width == 1) { mmha_launch_kernel_dispatch= param.beamWidth) + if (parent < 0 || parent > param.beamWidth) { param.outputIds[levelBeamIx] = param.endTokens[batch]; parent = -1; @@ -702,11 +702,10 @@ __global__ void transposeLogProbs(float* outputLogProbs, float* outputLogProbsTi } auto const batchSlot = batchSlots[batchIdx]; - auto const batchBeamIdx = batchSlot * beamWidth + beamIdx; - if (pos < sequenceLengths[batchBeamIdx]) + if (pos < sequenceLengths[batchSlot]) { - auto const outputIndex = batchSlot * beamWidth * maxSeqLen + beamIdx * maxSeqLen + pos; - outputLogProbs[outputIndex] + auto const batchBeamIdx = batchSlot * beamWidth * maxSeqLen + beamIdx * maxSeqLen + pos; + outputLogProbs[batchBeamIdx] = outputLogProbsTiled[pos * maxBatchSize * beamWidth + batchSlot * beamWidth + beamIdx]; } } @@ -729,7 +728,7 @@ namespace tensorrt_llm::runtime::kernels { // Must be similar to [cpp/tensorrt_llm/thop/gatherTreeOp.cpp] gatherTree void gatherTree(DecodingOutput const& decodingOutput, DecodingInput const& decodingInput, - SamplingConfig const& samplingConfig, runtime::CudaStream const& cudaStream, runtime::SizeType32 batchSlot) + SamplingConfig const& samplingConfig, runtime::CudaStream const& cudaStream) { TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); @@ -782,24 +781,7 @@ void gatherTree(DecodingOutput const& decodingOutput, DecodingInput const& decod lengthPenaltyPtr = manager.copyFrom(lengthPenaltyVec, ITensor::makeShape({batchSize}), runtime::MemoryType::kGPU); tensorrt_llm::kernels::BeamHypotheses bh; - // logProbsTiled has shape [MSL, maxNumSequences, BM] and is passed unsliced. - // nMaxBatchSize must equal the allocation stride (dim-1), not the per-slot batchSize=1. - // The pointer is pre-offset by batchSlot*BM so that insertUnfinishedPathKernel, - // which uses bid=0 / nBatchSize=1, computes: - // (base + batchSlot*BM)[step * maxBS * BM + 0*BM + beamIdx] - // = base[step * maxBS * BM + batchSlot * BM + beamIdx] - // = logProbsTiled[step][batchSlot][beamIdx] ✓ - auto const logProbsTiledMaxBatchSize = static_cast(decodingOutput.logProbsTiled->getShape().d[1]); - auto const logProbsTiledBeamWidth = static_cast(decodingOutput.logProbsTiled->getShape().d[2]); - TLLM_CHECK_WITH_INFO(batchSlot < logProbsTiledMaxBatchSize, - "batchSlot (%d) must be < logProbsTiled maxBatchSize (%d); " - "logProbsTiled would be accessed out of bounds.", - batchSlot, logProbsTiledMaxBatchSize); - TLLM_CHECK_WITH_INFO(beamWidth == logProbsTiledBeamWidth, - "beamWidth (%d) must equal logProbsTiled BM dimension (%d); " - "pointer offset batchSlot*beamWidth would be misaligned.", - beamWidth, logProbsTiledBeamWidth); - bh.nMaxBatchSize = logProbsTiledMaxBatchSize; + bh.nMaxBatchSize = batchSize; bh.nBatchSize = batchSize; bh.nBeamWidth = beamWidth; bh.nMaxSeqLen = maxSeqLength; @@ -807,7 +789,7 @@ void gatherTree(DecodingOutput const& decodingOutput, DecodingInput const& decod bh.inputLengths = bufferCast(*decodingInput.lengths); bh.outputIds = bufferCast(finalOutputIds); bh.logProbs = bufferCastOrNull(decodingOutput.logProbs); - bh.logProbsTiled = bufferCast(*decodingOutput.logProbsTiled) + batchSlot * beamWidth; + bh.logProbsTiled = bufferCast(*decodingOutput.logProbsTiled); bh.sequenceLengths = bufferCast(*decodingOutput.lengths); bh.cumLogProbs = bufferCast(*decodingOutput.cumLogProbs); bh.outputIdsCBA = bufferCast(*decodingOutput.beamHypotheses.outputIdsCBA); diff --git a/cpp/tensorrt_llm/kernels/decodingKernels.h b/cpp/tensorrt_llm/kernels/decodingKernels.h index 25fca71ee267..0e4fded936bc 100644 --- a/cpp/tensorrt_llm/kernels/decodingKernels.h +++ b/cpp/tensorrt_llm/kernels/decodingKernels.h @@ -1,5 +1,5 @@ /* - * Copyright (c) 2019-2026, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2019-2024, NVIDIA CORPORATION. All rights reserved. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -133,5 +133,5 @@ namespace tensorrt_llm::runtime::kernels //! \param cudaStream the CUDA stream on which to perform the operation. void gatherTree(DecodingOutput const& decodingOutput, DecodingInput const& decodingInput, - SamplingConfig const& samplingConfig, runtime::CudaStream const& cudaStream, runtime::SizeType32 batchSlot = 0); + SamplingConfig const& samplingConfig, runtime::CudaStream const& cudaStream); } // namespace tensorrt_llm::runtime::kernels diff --git a/cpp/tensorrt_llm/kernels/deepseekV4BlockTable.cu b/cpp/tensorrt_llm/kernels/deepseekV4BlockTable.cu deleted file mode 100644 index 7b6188c74eca..000000000000 --- a/cpp/tensorrt_llm/kernels/deepseekV4BlockTable.cu +++ /dev/null @@ -1,372 +0,0 @@ -/* - * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. - * - * 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. - */ - -#include "tensorrt_llm/kernels/deepseekV4BlockTable.h" - -#include - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels -{ -namespace -{ - -constexpr int32_t kBadPageIndex = -1; -constexpr int32_t kThreadsPerBlock = 256; -constexpr int32_t kVecThreadsPerBlock = 128; -constexpr int32_t kRowKernelMinBlocks = 256; -constexpr int32_t kVecRowsPerBlock = 8; - -__device__ __forceinline__ int32_t computeBasePageIndex(int32_t const* __restrict__ blockOffsets, - int32_t const* __restrict__ copyIdx, int64_t const* __restrict__ poolIds, bool const* __restrict__ validPool, - int32_t const* __restrict__ scales, int32_t const* __restrict__ layerOffsets, int32_t numPools, - int32_t copyIdxCapacity, int32_t numAttnTypes, int32_t maxBlocksPerSeq, int32_t layerId, int32_t attnTypeId, - int32_t tableId, int32_t blockId) -{ - int32_t const layerAttnOffset = layerId * numAttnTypes + attnTypeId; - int64_t const poolId64 = poolIds[layerAttnOffset]; - bool const isValidPool = validPool[layerAttnOffset] && poolId64 >= 0 && poolId64 < numPools; - if (!isValidPool) - { - return kBadPageIndex; - } - - int32_t const mappedTableId = copyIdx[tableId]; - if (mappedTableId < 0 || mappedTableId >= copyIdxCapacity) - { - return kBadPageIndex; - } - - auto const poolId = static_cast(poolId64); - int64_t const blockOffsetsIndex - = (((static_cast(poolId) * copyIdxCapacity + mappedTableId) * 2) * maxBlocksPerSeq) + blockId; - int32_t const base = blockOffsets[blockOffsetsIndex]; - if (base == kBadPageIndex) - { - return kBadPageIndex; - } - - return base * scales[layerAttnOffset] + layerOffsets[layerAttnOffset]; -} - -__device__ __forceinline__ int32_t applyScaleAndOffset(int32_t base, int32_t scale, int32_t layerOffset) -{ - return base == kBadPageIndex ? kBadPageIndex : base * scale + layerOffset; -} - -__device__ __forceinline__ void fillBadSlidingBlockTableRow(int32_t* outputRow, int32_t maxBlocksPerSeq, bool useVec4) -{ - if (useVec4) - { - int4 const bad = {kBadPageIndex, kBadPageIndex, kBadPageIndex, kBadPageIndex}; - auto* outputVec = reinterpret_cast(outputRow); - int32_t const vecsPerRow = maxBlocksPerSeq / 4; - for (int32_t vecId = threadIdx.x; vecId < vecsPerRow; vecId += blockDim.x) - { - outputVec[vecId] = bad; - } - return; - } - - for (int32_t blockId = threadIdx.x; blockId < maxBlocksPerSeq; blockId += blockDim.x) - { - outputRow[blockId] = kBadPageIndex; - } -} - -__global__ void computeSlidingBlockTablesRowsTiledKernel(int32_t const* __restrict__ blockOffsets, - int32_t const* __restrict__ copyIdx, int64_t const* __restrict__ poolIds, bool const* __restrict__ validPool, - int32_t const* __restrict__ scales, int32_t const* __restrict__ layerOffsets, int32_t* __restrict__ output, - int32_t numPools, int32_t copyIdxCapacity, int32_t numLayerAttn, int32_t numTables, int32_t maxBlocksPerSeq) -{ - bool const useVec4 = maxBlocksPerSeq % 4 == 0; - int32_t const vecsPerRow = maxBlocksPerSeq / 4; - int32_t const firstTableId = static_cast(blockIdx.x) * kVecRowsPerBlock; - int32_t const layerAttnOffset = static_cast(blockIdx.y); - if (layerAttnOffset >= numLayerAttn) - { - return; - } - - int64_t const poolId64 = poolIds[layerAttnOffset]; - bool const isValidPool = validPool[layerAttnOffset] && poolId64 >= 0 && poolId64 < numPools; - if (!isValidPool) - { -#pragma unroll - for (int32_t localRow = 0; localRow < kVecRowsPerBlock; ++localRow) - { - int32_t const tableId = firstTableId + localRow; - if (tableId >= numTables) - { - continue; - } - - int64_t const outputOffset - = (static_cast(layerAttnOffset) * numTables + tableId) * maxBlocksPerSeq; - fillBadSlidingBlockTableRow(output + outputOffset, maxBlocksPerSeq, useVec4); - } - return; - } - - auto const poolId = static_cast(poolId64); - int32_t const scale = scales[layerAttnOffset]; - int32_t const layerOffset = layerOffsets[layerAttnOffset]; - -#pragma unroll - for (int32_t localRow = 0; localRow < kVecRowsPerBlock; ++localRow) - { - int32_t const tableId = firstTableId + localRow; - if (tableId >= numTables) - { - continue; - } - - int64_t const outputOffset = (static_cast(layerAttnOffset) * numTables + tableId) * maxBlocksPerSeq; - auto* outputRow = output + outputOffset; - int32_t const mappedTableId = copyIdx[tableId]; - bool const isValidTable = mappedTableId >= 0 && mappedTableId < copyIdxCapacity; - if (!isValidTable) - { - fillBadSlidingBlockTableRow(outputRow, maxBlocksPerSeq, useVec4); - continue; - } - - int64_t const blockOffsetsOffset - = ((static_cast(poolId) * copyIdxCapacity + mappedTableId) * 2) * maxBlocksPerSeq; - auto const* blockOffsetsRow = blockOffsets + blockOffsetsOffset; - if (useVec4) - { - auto const* blockOffsetsVec = reinterpret_cast(blockOffsetsRow); - auto* outputVec = reinterpret_cast(outputRow); - for (int32_t vecId = threadIdx.x; vecId < vecsPerRow; vecId += blockDim.x) - { - int4 const base = blockOffsetsVec[vecId]; - int4 const value = {applyScaleAndOffset(base.x, scale, layerOffset), - applyScaleAndOffset(base.y, scale, layerOffset), applyScaleAndOffset(base.z, scale, layerOffset), - applyScaleAndOffset(base.w, scale, layerOffset)}; - outputVec[vecId] = value; - } - continue; - } - - for (int32_t blockId = threadIdx.x; blockId < maxBlocksPerSeq; blockId += blockDim.x) - { - int32_t const base = blockOffsetsRow[blockId]; - outputRow[blockId] = applyScaleAndOffset(base, scale, layerOffset); - } - } -} - -__global__ void computeSlidingBlockTablesWithScratchKernel(int32_t const* __restrict__ blockOffsets, - int32_t const* __restrict__ copyIdx, int64_t const* __restrict__ poolIds, bool const* __restrict__ validPool, - int32_t const* __restrict__ scales, int32_t const* __restrict__ layerOffsets, - int32_t const* __restrict__ scratchPages, int32_t const* __restrict__ scratchBegs, - int32_t const* __restrict__ scratchEnds, int32_t const* __restrict__ scratchSlots, - int32_t const* __restrict__ numContexts, int32_t* __restrict__ output, int64_t totalElements, int32_t numPools, - int32_t copyIdxCapacity, int32_t numAttnTypes, int32_t numTables, int32_t maxBlocksPerSeq, int32_t scratchCapacity, - int32_t maxScratchSlots) -{ - int64_t const linearIdx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - if (linearIdx >= totalElements) - { - return; - } - - int64_t remaining = linearIdx; - int32_t const blockId = static_cast(remaining % maxBlocksPerSeq); - remaining /= maxBlocksPerSeq; - int32_t const tableId = static_cast(remaining % numTables); - remaining /= numTables; - int32_t const attnTypeId = static_cast(remaining % numAttnTypes); - int32_t const layerId = static_cast(remaining / numAttnTypes); - - int32_t const layerAttnOffset = layerId * numAttnTypes + attnTypeId; - int32_t const basePageIndex = computeBasePageIndex(blockOffsets, copyIdx, poolIds, validPool, scales, layerOffsets, - numPools, copyIdxCapacity, numAttnTypes, maxBlocksPerSeq, layerId, attnTypeId, tableId, blockId); - - int64_t const poolId64 = poolIds[layerAttnOffset]; - bool const isValidPool = validPool[layerAttnOffset] && poolId64 >= 0 && poolId64 < numPools; - int32_t const activeContexts = numContexts[0]; - bool const canUseScratch = isValidPool && tableId < scratchCapacity && tableId < activeContexts; - - if (!canUseScratch) - { - output[linearIdx] = basePageIndex; - return; - } - - auto const poolId = static_cast(poolId64); - int64_t const scratchRangeOffset = static_cast(poolId) * scratchCapacity + tableId; - int32_t const scratchBeg = scratchBegs[scratchRangeOffset]; - int32_t const scratchEnd = scratchEnds[scratchRangeOffset]; - bool const inScratchRange = blockId >= scratchBeg && blockId < scratchEnd; - if (!inScratchRange) - { - output[linearIdx] = basePageIndex; - return; - } - - int32_t const scale = scales[layerAttnOffset]; - int32_t const rangeIndex = blockId - scratchBeg; - int32_t const totalOffset = rangeIndex * scratchPages[layerAttnOffset]; - int32_t slotIdx = totalOffset / scale; - if (slotIdx >= maxScratchSlots) - { - slotIdx = maxScratchSlots - 1; - } - - int64_t const slotOffset = scratchRangeOffset * maxScratchSlots + slotIdx; - int32_t const slotId = scratchSlots[slotOffset]; - int32_t const offset = totalOffset % scale; - output[linearIdx] = slotId * scale + ((offset + layerOffsets[layerAttnOffset]) % scale); -} - -__global__ void computeSlidingBlockTablesWithScratchRowsKernel(int32_t const* __restrict__ blockOffsets, - int32_t const* __restrict__ copyIdx, int64_t const* __restrict__ poolIds, bool const* __restrict__ validPool, - int32_t const* __restrict__ scales, int32_t const* __restrict__ layerOffsets, - int32_t const* __restrict__ scratchPages, int32_t const* __restrict__ scratchBegs, - int32_t const* __restrict__ scratchEnds, int32_t const* __restrict__ scratchSlots, - int32_t const* __restrict__ numContexts, int32_t* __restrict__ output, int32_t numPools, int32_t copyIdxCapacity, - int32_t numAttnTypes, int32_t numTables, int32_t maxBlocksPerSeq, int32_t scratchCapacity, int32_t maxScratchSlots) -{ - int32_t const rowIdx = static_cast(blockIdx.x); - int32_t const tableId = rowIdx % numTables; - int32_t const layerAttnIdx = rowIdx / numTables; - int32_t const attnTypeId = layerAttnIdx % numAttnTypes; - int32_t const layerId = layerAttnIdx / numAttnTypes; - int32_t const layerAttnOffset = layerId * numAttnTypes + attnTypeId; - int32_t const outputOffset = rowIdx * maxBlocksPerSeq; - - int64_t const poolId64 = poolIds[layerAttnOffset]; - bool const isValidPool = validPool[layerAttnOffset] && poolId64 >= 0 && poolId64 < numPools; - if (!isValidPool) - { - for (int32_t blockId = threadIdx.x; blockId < maxBlocksPerSeq; blockId += blockDim.x) - { - output[outputOffset + blockId] = kBadPageIndex; - } - return; - } - - auto const poolId = static_cast(poolId64); - int32_t const scale = scales[layerAttnOffset]; - int32_t const layerOffset = layerOffsets[layerAttnOffset]; - int32_t const activeContexts = numContexts[0]; - bool const canUseScratch = tableId < scratchCapacity && tableId < activeContexts; - int64_t const scratchRangeOffset = static_cast(poolId) * scratchCapacity + tableId; - int32_t const scratchBeg = canUseScratch ? scratchBegs[scratchRangeOffset] : 0; - int32_t const scratchEnd = canUseScratch ? scratchEnds[scratchRangeOffset] : 0; - int32_t const scratchPageCount = scratchPages[layerAttnOffset]; - - int32_t const mappedTableId = copyIdx[tableId]; - bool const isValidTable = mappedTableId >= 0 && mappedTableId < copyIdxCapacity; - int64_t const blockOffsetsOffset - = ((static_cast(poolId) * copyIdxCapacity + mappedTableId) * 2) * maxBlocksPerSeq; - - for (int32_t blockId = threadIdx.x; blockId < maxBlocksPerSeq; blockId += blockDim.x) - { - bool const inScratchRange = canUseScratch && blockId >= scratchBeg && blockId < scratchEnd; - if (inScratchRange) - { - int32_t const rangeIndex = blockId - scratchBeg; - int32_t const totalOffset = rangeIndex * scratchPageCount; - int32_t slotIdx = totalOffset / scale; - if (slotIdx >= maxScratchSlots) - { - slotIdx = maxScratchSlots - 1; - } - - int64_t const slotOffset = scratchRangeOffset * maxScratchSlots + slotIdx; - int32_t const slotId = scratchSlots[slotOffset]; - int32_t const offset = totalOffset % scale; - output[outputOffset + blockId] = slotId * scale + ((offset + layerOffset) % scale); - continue; - } - - if (!isValidTable) - { - output[outputOffset + blockId] = kBadPageIndex; - continue; - } - - int32_t const base = blockOffsets[blockOffsetsOffset + blockId]; - output[outputOffset + blockId] = base == kBadPageIndex ? kBadPageIndex : base * scale + layerOffset; - } -} - -} // namespace - -void invokeDeepseekV4ComputeSlidingBlockTables(int32_t const* blockOffsets, int32_t const* copyIdx, - int64_t const* poolIds, bool const* validPool, int32_t const* scales, int32_t const* layerOffsets, int32_t* output, - int32_t numPools, int32_t copyIdxCapacity, int32_t numLayers, int32_t numAttnTypes, int32_t numTables, - int32_t maxBlocksPerSeq, cudaStream_t stream) -{ - int64_t const totalElements = static_cast(numLayers) * numAttnTypes * numTables * maxBlocksPerSeq; - if (totalElements == 0) - { - return; - } - - int32_t const numLayerAttn = numLayers * numAttnTypes; - int32_t const itemsPerRow = maxBlocksPerSeq % 4 == 0 ? maxBlocksPerSeq / 4 : maxBlocksPerSeq; - int32_t threadsPerBlock = itemsPerRow >= kVecThreadsPerBlock ? kVecThreadsPerBlock : itemsPerRow; - if (threadsPerBlock < 64) - { - threadsPerBlock = 64; - } - - dim3 const block(static_cast(threadsPerBlock)); - dim3 const grid(static_cast((numTables + kVecRowsPerBlock - 1) / kVecRowsPerBlock), - static_cast(numLayerAttn)); - computeSlidingBlockTablesRowsTiledKernel<<>>(blockOffsets, copyIdx, poolIds, validPool, - scales, layerOffsets, output, numPools, copyIdxCapacity, numLayerAttn, numTables, maxBlocksPerSeq); -} - -void invokeDeepseekV4ComputeSlidingBlockTablesWithScratch(int32_t const* blockOffsets, int32_t const* copyIdx, - int64_t const* poolIds, bool const* validPool, int32_t const* scales, int32_t const* layerOffsets, - int32_t const* scratchPages, int32_t const* scratchBegs, int32_t const* scratchEnds, int32_t const* scratchSlots, - int32_t const* numContexts, int32_t* output, int32_t numPools, int32_t copyIdxCapacity, int32_t numLayers, - int32_t numAttnTypes, int32_t numTables, int32_t maxBlocksPerSeq, int32_t scratchCapacity, int32_t maxScratchSlots, - cudaStream_t stream) -{ - int64_t const totalElements = static_cast(numLayers) * numAttnTypes * numTables * maxBlocksPerSeq; - if (totalElements == 0) - { - return; - } - - if (maxBlocksPerSeq >= kRowKernelMinBlocks) - { - int32_t const numRows = numLayers * numAttnTypes * numTables; - dim3 const block(kThreadsPerBlock); - dim3 const grid(static_cast(numRows)); - computeSlidingBlockTablesWithScratchRowsKernel<<>>(blockOffsets, copyIdx, poolIds, - validPool, scales, layerOffsets, scratchPages, scratchBegs, scratchEnds, scratchSlots, numContexts, output, - numPools, copyIdxCapacity, numAttnTypes, numTables, maxBlocksPerSeq, scratchCapacity, maxScratchSlots); - return; - } - - dim3 const block(kThreadsPerBlock); - dim3 const grid(static_cast((totalElements + kThreadsPerBlock - 1) / kThreadsPerBlock)); - computeSlidingBlockTablesWithScratchKernel<<>>(blockOffsets, copyIdx, poolIds, validPool, - scales, layerOffsets, scratchPages, scratchBegs, scratchEnds, scratchSlots, numContexts, output, totalElements, - numPools, copyIdxCapacity, numAttnTypes, numTables, maxBlocksPerSeq, scratchCapacity, maxScratchSlots); -} - -} // namespace kernels - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/deepseekV4BlockTable.h b/cpp/tensorrt_llm/kernels/deepseekV4BlockTable.h deleted file mode 100644 index f57b757faaf3..000000000000 --- a/cpp/tensorrt_llm/kernels/deepseekV4BlockTable.h +++ /dev/null @@ -1,43 +0,0 @@ -/* - * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. - * - * 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. - */ - -#pragma once - -#include "tensorrt_llm/common/config.h" - -#include -#include - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels -{ - -void invokeDeepseekV4ComputeSlidingBlockTables(int32_t const* blockOffsets, int32_t const* copyIdx, - int64_t const* poolIds, bool const* validPool, int32_t const* scales, int32_t const* layerOffsets, int32_t* output, - int32_t numPools, int32_t copyIdxCapacity, int32_t numLayers, int32_t numAttnTypes, int32_t numTables, - int32_t maxBlocksPerSeq, cudaStream_t stream); - -void invokeDeepseekV4ComputeSlidingBlockTablesWithScratch(int32_t const* blockOffsets, int32_t const* copyIdx, - int64_t const* poolIds, bool const* validPool, int32_t const* scales, int32_t const* layerOffsets, - int32_t const* scratchPages, int32_t const* scratchBegs, int32_t const* scratchEnds, int32_t const* scratchSlots, - int32_t const* numContexts, int32_t* output, int32_t numPools, int32_t copyIdxCapacity, int32_t numLayers, - int32_t numAttnTypes, int32_t numTables, int32_t maxBlocksPerSeq, int32_t scratchCapacity, int32_t maxScratchSlots, - cudaStream_t stream); - -} // namespace kernels - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/deepseekV4QNormKernel.cu b/cpp/tensorrt_llm/kernels/deepseekV4QNormKernel.cu deleted file mode 100644 index b9fb642bd69d..000000000000 --- a/cpp/tensorrt_llm/kernels/deepseekV4QNormKernel.cu +++ /dev/null @@ -1,274 +0,0 @@ -/* - * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. - * - * 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. - */ - -#include "tensorrt_llm/kernels/deepseekV4QNormKernel.h" - -#include "tensorrt_llm/common/assert.h" - -#include -#include -#include -#include - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels -{ -namespace -{ - -constexpr int kWarpSize = 32; -constexpr int kWarpsPerBlock = 4; -constexpr int kThreadsPerBlock = kWarpSize * kWarpsPerBlock; - -template -struct Vec2Traits; - -template <> -struct Vec2Traits -{ - using Type = half2; - - __device__ static float2 toFloat2(Type value) - { - return __half22float2(value); - } - - __device__ static Type fromFloat2(float2 value) - { - return __floats2half2_rn(value.x, value.y); - } -}; - -template <> -struct Vec2Traits<__nv_bfloat16> -{ - using Type = __nv_bfloat162; - - __device__ static float2 toFloat2(Type value) - { - return __bfloat1622float2(value); - } - - __device__ static Type fromFloat2(float2 value) - { - return __floats2bfloat162_rn(value.x, value.y); - } -}; - -__device__ __forceinline__ float warpReduceSum(float value) -{ - for (int mask = kWarpSize / 2; mask > 0; mask >>= 1) - { - value += __shfl_xor_sync(0xFFFFFFFF, value, mask); - } - return value; -} - -template -__global__ void deepseekV4QNormKernel(T const* input, T* output, int totalRows, float eps) -{ - static_assert(kHeadDim % (2 * kWarpSize) == 0); - constexpr int kPairsPerRow = kHeadDim / 2; - constexpr int kPairsPerLane = kPairsPerRow / kWarpSize; - - using Vec2 = typename Vec2Traits::Type; - - int const warpId = threadIdx.x / kWarpSize; - int const laneId = threadIdx.x % kWarpSize; - int const row = blockIdx.x * kWarpsPerBlock + warpId; - - if (row >= totalRows) - { - return; - } - - auto const* inputPair = reinterpret_cast(input + static_cast(row) * kHeadDim); - auto* outputPair = reinterpret_cast(output + static_cast(row) * kHeadDim); - - float2 values[kPairsPerLane]; - float sumSquares = 0.0F; - -#pragma unroll - for (int i = 0; i < kPairsPerLane; ++i) - { - int const pairIdx = i * kWarpSize + laneId; - values[i] = Vec2Traits::toFloat2(inputPair[pairIdx]); - sumSquares += values[i].x * values[i].x + values[i].y * values[i].y; - } - - sumSquares = warpReduceSum(sumSquares); - float const scale = rsqrtf(sumSquares / static_cast(kHeadDim) + eps); - -#pragma unroll - for (int i = 0; i < kPairsPerLane; ++i) - { - int const pairIdx = i * kWarpSize + laneId; - float2 value{values[i].x * scale, values[i].y * scale}; - outputPair[pairIdx] = Vec2Traits::fromFloat2(value); - } -} - -template -void dispatchDeepseekV4QNorm( - void const* input, void* output, int totalRows, int headDim, float eps, cudaStream_t stream) -{ - TLLM_CHECK_WITH_INFO(headDim == 512, "deepseek_v4_q_norm only supports head_dim=512, got %d", headDim); - - dim3 const block(kThreadsPerBlock); - dim3 const grid((totalRows + kWarpsPerBlock - 1) / kWarpsPerBlock); - deepseekV4QNormKernel - <<>>(static_cast(input), static_cast(output), totalRows, eps); -} - -// Fused q-norm + FP8 quant of nope segment. Row layout [nope|rope]; writes FP8 -// nope (scaled by inv_rms * quant_scale_qkv) to `quant_q_nope` with per-row -// stride `quantQNopeRowStrideBytes`, and bf16/fp16 rope to `q_pe_out`. -// Requires kHeadDim==512, kNopeDim==448, kRopeDim==64 so each lane's -// (kPairsPerLane-1) iterations cover the nope range and the last iteration -// covers the rope range exactly. - -template -__global__ void deepseekV4QNormFusedKernel(T const* __restrict__ input, __nv_fp8_e4m3* __restrict__ quant_q_nope, - T* __restrict__ q_pe_out, float const* __restrict__ quant_scale_qkv_ptr, int totalRows, - int quantQNopeRowStrideBytes, float eps) -{ - static_assert(kHeadDim % (2 * kWarpSize) == 0); - static_assert(kNopeDim > 0 && kNopeDim < kHeadDim); - constexpr int kRopeDim = kHeadDim - kNopeDim; - constexpr int kPairsPerRow = kHeadDim / 2; - constexpr int kPairsPerLane = kPairsPerRow / kWarpSize; - constexpr int kNopePairs = kNopeDim / 2; - constexpr int kRopePairs = kRopeDim / 2; - static_assert(kPairsPerLane >= 2); - static_assert(kNopePairs == (kPairsPerLane - 1) * kWarpSize, - "Fused kernel assumes the last per-lane iteration covers the rope segment."); - static_assert(kRopePairs == kWarpSize, "Each lane should own exactly one rope pair."); - - using Vec2 = typename Vec2Traits::Type; - - int const warpId = threadIdx.x / kWarpSize; - int const laneId = threadIdx.x % kWarpSize; - int const row = blockIdx.x * kWarpsPerBlock + warpId; - - if (row >= totalRows) - { - return; - } - - auto const* inputPair = reinterpret_cast(input + static_cast(row) * kHeadDim); - // Nope output: row stride is caller-controlled (kNopeDim for packed, kHeadDim - // when interleaved with the rope segment of a full Q-buffer that RoPE writes). - auto* nopeOutPair = reinterpret_cast<__nv_fp8x2_e4m3*>( - reinterpret_cast<__nv_fp8_e4m3*>(quant_q_nope) + static_cast(row) * quantQNopeRowStrideBytes); - auto* ropeOutPair = reinterpret_cast(q_pe_out + static_cast(row) * kRopeDim); - - float const quantScale = quant_scale_qkv_ptr ? quant_scale_qkv_ptr[0] : 1.0F; - - float2 values[kPairsPerLane]; - float sumSquares = 0.0F; - -#pragma unroll - for (int i = 0; i < kPairsPerLane; ++i) - { - int const pairIdx = i * kWarpSize + laneId; - values[i] = Vec2Traits::toFloat2(inputPair[pairIdx]); - sumSquares += values[i].x * values[i].x + values[i].y * values[i].y; - } - - sumSquares = warpReduceSum(sumSquares); - float const normScale = rsqrtf(sumSquares / static_cast(kHeadDim) + eps); - float const fp8Scale = normScale * quantScale; - - // First kPairsPerLane-1 iters land in the nope range -> FP8 STG. -#pragma unroll - for (int i = 0; i < kPairsPerLane - 1; ++i) - { - int const pairIdx = i * kWarpSize + laneId; - float2 scaled{values[i].x * fp8Scale, values[i].y * fp8Scale}; - nopeOutPair[pairIdx] = __nv_fp8x2_e4m3(scaled); - } - - // Last iter is the rope segment -> bf16/fp16 STG (no extra quant scale). - { - constexpr int i = kPairsPerLane - 1; - int const pairIdx = i * kWarpSize + laneId; // in [kNopePairs, kPairsPerRow) - int const ropePairIdx = pairIdx - kNopePairs; // in [0, kRopePairs) - float2 normalized{values[i].x * normScale, values[i].y * normScale}; - ropeOutPair[ropePairIdx] = Vec2Traits::fromFloat2(normalized); - } -} - -template -void dispatchDeepseekV4QNormFused(void const* input, void* quant_q_nope, void* q_pe_out, - void const* quant_scale_qkv_ptr, int totalRows, int headDim, int nopeDim, int quantQNopeRowStrideBytes, float eps, - cudaStream_t stream) -{ - assert(headDim == 512); - assert(nopeDim == 448); - assert(quantQNopeRowStrideBytes >= nopeDim); - - dim3 const block(kThreadsPerBlock); - dim3 const grid((totalRows + kWarpsPerBlock - 1) / kWarpsPerBlock); - deepseekV4QNormFusedKernel<<>>(static_cast(input), - static_cast<__nv_fp8_e4m3*>(quant_q_nope), static_cast(q_pe_out), - static_cast(quant_scale_qkv_ptr), totalRows, quantQNopeRowStrideBytes, eps); -} - -} // namespace - -void invokeDeepseekV4QNorm( - void const* input, void* output, int totalRows, int headDim, bool isBfloat16, float eps, cudaStream_t stream) -{ - if (totalRows == 0) - { - return; - } - - if (isBfloat16) - { - dispatchDeepseekV4QNorm<__nv_bfloat16>(input, output, totalRows, headDim, eps, stream); - } - else - { - dispatchDeepseekV4QNorm(input, output, totalRows, headDim, eps, stream); - } -} - -void invokeDeepseekV4QNormFusedFp8(void const* input, void* quant_q_nope, void* q_pe_out, - void const* quant_scale_qkv_ptr, int totalRows, int headDim, int nopeDim, int quantQNopeRowStrideBytes, - bool isBfloat16, float eps, cudaStream_t stream) -{ - if (totalRows == 0) - { - return; - } - - if (isBfloat16) - { - dispatchDeepseekV4QNormFused<__nv_bfloat16>(input, quant_q_nope, q_pe_out, quant_scale_qkv_ptr, totalRows, - headDim, nopeDim, quantQNopeRowStrideBytes, eps, stream); - } - else - { - dispatchDeepseekV4QNormFused(input, quant_q_nope, q_pe_out, quant_scale_qkv_ptr, totalRows, headDim, - nopeDim, quantQNopeRowStrideBytes, eps, stream); - } -} - -} // namespace kernels - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/deepseekV4QNormKernel.h b/cpp/tensorrt_llm/kernels/deepseekV4QNormKernel.h deleted file mode 100644 index 1c47e383a907..000000000000 --- a/cpp/tensorrt_llm/kernels/deepseekV4QNormKernel.h +++ /dev/null @@ -1,50 +0,0 @@ -/* - * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. - * - * 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. - */ - -#pragma once - -#include "tensorrt_llm/common/config.h" - -#include - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels -{ - -void invokeDeepseekV4QNorm( - void const* input, void* output, int totalRows, int headDim, bool isBfloat16, float eps, cudaStream_t stream); - -// Fused variant: in one pass, performs per-row RMSNorm and writes: -// * the first `nopeDim` columns of each normalized row as FP8E4M3 into `quant_q_nope` -// (scaled by `*quant_scale_qkv_ptr` if non-null, otherwise 1.0f). The per-row -// stride in bytes is `quantQNopeRowStrideBytes`; pass `nopeDim` for a packed -// [totalRows, nopeDim] output buffer, or `headDim` to interleave with the rope -// segment of a shared `[totalRows, headDim]` FP8 Q buffer (consumed by FMHA -// after applyMLARopeAndAssignQKVKernelOptContext fills the rope slot). -// * the remaining `headDim - nopeDim` columns as the input dtype into `q_pe_out` -// (always packed [totalRows, headDim - nopeDim]). -// Eliminates the trailing bf16→FP8 quant_copy pass that follows q_b_layernorm in the -// MLA absorption-mode prefill path. -// -// Only headDim==512, nopeDim==448 is currently supported (DeepSeek-V4 absorption shape). -void invokeDeepseekV4QNormFusedFp8(void const* input, void* quant_q_nope, void* q_pe_out, - void const* quant_scale_qkv_ptr, int totalRows, int headDim, int nopeDim, int quantQNopeRowStrideBytes, - bool isBfloat16, float eps, cudaStream_t stream); - -} // namespace kernels - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3RouterGemm.cu b/cpp/tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3RouterGemm.cu index 102866de363f..e33f46e1e152 100644 --- a/cpp/tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3RouterGemm.cu +++ b/cpp/tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3RouterGemm.cu @@ -291,55 +291,6 @@ template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__n template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__nv_bfloat16, 16, 256, 6144>( float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); -// hidden_dim=4096 instantiations (DeepSeek-V4). -template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__nv_bfloat16, 1, 256, 4096>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__nv_bfloat16, 2, 256, 4096>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__nv_bfloat16, 3, 256, 4096>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__nv_bfloat16, 4, 256, 4096>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__nv_bfloat16, 5, 256, 4096>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__nv_bfloat16, 6, 256, 4096>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__nv_bfloat16, 7, 256, 4096>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__nv_bfloat16, 8, 256, 4096>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__nv_bfloat16, 9, 256, 4096>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__nv_bfloat16, 10, 256, 4096>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__nv_bfloat16, 11, 256, 4096>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__nv_bfloat16, 12, 256, 4096>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__nv_bfloat16, 13, 256, 4096>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__nv_bfloat16, 14, 256, 4096>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__nv_bfloat16, 15, 256, 4096>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__nv_bfloat16, 16, 256, 4096>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - } // namespace kernels::dsv3MinLatencyKernels TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/fmhaDispatcher.cpp b/cpp/tensorrt_llm/kernels/fmhaDispatcher.cpp index eb8f258d0f1d..88ff741d6941 100644 --- a/cpp/tensorrt_llm/kernels/fmhaDispatcher.cpp +++ b/cpp/tensorrt_llm/kernels/fmhaDispatcher.cpp @@ -53,14 +53,13 @@ FmhaDispatcher::FmhaDispatcher(MHARunnerFixedParams fixedParams) // The exception will fall back to fmha v2. // Please update fmha_v2/setup.py if you want to add more supported head sizes. , mUseTllmGen(tensorrt_llm::common::isSM100Family() && fixedParams.headSize != 72) - , mMultiProcessorCount(tensorrt_llm::common::getMultiProcessorCount()) { if (mUseTllmGen) { auto [dataTypeK, dataTypeV] = unpack_kv_data_type(mFixedParams.dataTypeKv); - mTllmGenFMHARunner.reset(new TllmGenFmhaRunner(mFixedParams.dataType, dataTypeK, dataTypeV, - mFixedParams.dataTypeOut, mFixedParams.sageBlockSizeQ, mFixedParams.sageBlockSizeK, 0, - mFixedParams.sageBlockSizeV, mFixedParams.fusesDsv4InvRopeFp8Quant)); + mTllmGenFMHARunner.reset( + new TllmGenFmhaRunner(mFixedParams.dataType, dataTypeK, dataTypeV, mFixedParams.dataTypeOut, + mFixedParams.sageBlockSizeQ, mFixedParams.sageBlockSizeK, 0, mFixedParams.sageBlockSizeV)); if (!isSupported()) { TLLM_LOG_WARNING("TRTLLM-GEN does not support the requested kernels."); @@ -129,7 +128,7 @@ bool FmhaDispatcher::isSupported() tllmRunnerParams.mHeadDimV = mFixedParams.headSizeV; tllmRunnerParams.mNumTokensPerPage = (qkvLayout == QkvLayout::PagedKv) ? mFixedParams.numTokensPerBlock : 0; tllmRunnerParams.mNumHeadsQPerKv = mFixedParams.numQHeads / mFixedParams.numKvHeads; - tllmRunnerParams.mMultiProcessorCount = mMultiProcessorCount; + tllmRunnerParams.mMultiProcessorCount = tensorrt_llm::common::getMultiProcessorCount(); // Set the chunked attention size and sliding window size to INT_MAX to disable them when checking if // the kernel is supported. tllmRunnerParams.mChunkedAttentionSize = INT_MAX; @@ -141,7 +140,6 @@ bool FmhaDispatcher::isSupported() = (mFixedParams.useSparseMLA && mFixedParams.headSizeV == mFixedParams.headSize) ? SparseType::DynamicTokenSparse : SparseType::StaticTokenSparse; - tllmRunnerParams.mDsv4EpilogueFusion.enabled = mFixedParams.fusesDsv4InvRopeFp8Quant; tllmRunnerParams.mKernelType = FmhaKernelType::Generation; tllmRunnerParams.mMaskType = TrtllmGenAttentionMaskType::Causal; // Generation-style kernels on long KV can pick MultiCtasKv cubins @@ -222,12 +220,6 @@ void FmhaDispatcher::run(MHARunnerParams runnerParams) tllmRunnerParams.oSfScalePtr = runnerParams.oSfScalePtr; tllmRunnerParams.oPtr = runnerParams.outputPtr; tllmRunnerParams.oSfPtr = runnerParams.outputSfPtr; - if (runnerParams.dsv4EpilogueFusion.enabled) - { - tllmRunnerParams.mDsv4EpilogueFusion.enabled = true; - tllmRunnerParams.mDsv4EpilogueFusion.cosSinCache = runnerParams.dsv4EpilogueFusion.cosSinCache; - tllmRunnerParams.mDsv4EpilogueFusion.scaleBufM = runnerParams.dsv4EpilogueFusion.scaleBufM; - } // The sequence lengths for K/V. tllmRunnerParams.seqLensKvPtr = reinterpret_cast(runnerParams.kvSeqLenPtr); // Assume same headDim for Qk and V here. @@ -255,7 +247,7 @@ void FmhaDispatcher::run(MHARunnerParams runnerParams) tllmRunnerParams.mScaleQ = mFixedParams.qScaling; // Set it to INT_MAX as the kv cache pageOffsets will ensure that there is no out-of-bounds access. tllmRunnerParams.mNumPagesInMemPool = INT_MAX; - tllmRunnerParams.mMultiProcessorCount = mMultiProcessorCount; + tllmRunnerParams.mMultiProcessorCount = tensorrt_llm::common::getMultiProcessorCount(); tllmRunnerParams.mSfStartTokenIdx = 0; // SageAttention scaling factors. tllmRunnerParams.sageAttnSfsQPtr = runnerParams.qScalePtr; diff --git a/cpp/tensorrt_llm/kernels/fmhaDispatcher.h b/cpp/tensorrt_llm/kernels/fmhaDispatcher.h index f1291d6bc1cb..26a40411fdf3 100644 --- a/cpp/tensorrt_llm/kernels/fmhaDispatcher.h +++ b/cpp/tensorrt_llm/kernels/fmhaDispatcher.h @@ -61,9 +61,6 @@ class FmhaDispatcher UniqPtrWNullCopy mFMHARunner; // Runner for trtllm-gen fmha kernels (for SM == 100) UniqPtrWNullCopy mTllmGenFMHARunner; - // Cached SM count to avoid repeated cudaDeviceGetAttribute calls in the per-iter - // FMHA dispatch hot path (isSupported / run). - int mMultiProcessorCount{0}; }; //////////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/cpp/tensorrt_llm/kernels/fusedDiTGateResidNormShiftScaleKernel.cu b/cpp/tensorrt_llm/kernels/fusedDiTGateResidNormShiftScaleKernel.cu deleted file mode 100644 index c0437dee3c57..000000000000 --- a/cpp/tensorrt_llm/kernels/fusedDiTGateResidNormShiftScaleKernel.cu +++ /dev/null @@ -1,1116 +0,0 @@ -/* - * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. - * - * 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. - */ - -#include "fusedDiTGateResidNormShiftScaleKernel.h" -#include "tensorrt_llm/common/cudaUtils.h" -#include "tensorrt_llm/common/reduceKernelUtils.cuh" -#include "tensorrt_llm/kernels/quantization.cuh" -#include -#include -#include - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels -{ - -namespace -{ - -__device__ __forceinline__ uint32_t cvta_to_smem(void const* ptr) -{ - return static_cast(__cvta_generic_to_shared(const_cast(ptr))); -} - -__device__ __forceinline__ void mbar_init(uint64_t* bar, uint32_t count) -{ -#if (defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900) - asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" : : "r"(cvta_to_smem(bar)), "r"(count)); -#endif -} - -__device__ __forceinline__ void mbar_arrive(uint64_t* bar) -{ -#if (defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900) - asm volatile("mbarrier.arrive.shared::cta.b64 _, [%0];" : : "r"(cvta_to_smem(bar))); -#endif -} - -__device__ __forceinline__ void mbar_arrive_expect_tx(uint64_t* bar, uint32_t tx_bytes) -{ -#if (defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900) - asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;" : : "r"(cvta_to_smem(bar)), "r"(tx_bytes)); -#endif -} - -__device__ __forceinline__ void mbar_wait(uint64_t* bar, uint32_t phase) -{ -#if (defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900) - asm volatile( - "{ .reg .pred P; \n" - " WAIT: mbarrier.try_wait.parity.shared::cta.b64 P, [%0], %1; \n" - " @P bra DONE; \n" - " bra WAIT; \n" - " DONE: }" - : - : "r"(cvta_to_smem(bar)), "r"(phase)); -#endif -} - -__device__ __forceinline__ void cp_async_bulk(void* smem_dst, void const* global_src, uint32_t bytes, uint64_t* bar) -{ -#if (defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900) - asm volatile( - "cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes [%0], [%1], %2, [%3];" - : - : "r"(cvta_to_smem(smem_dst)), "l"(reinterpret_cast(global_src)), "r"(bytes), "r"(cvta_to_smem(bar)) - : "memory"); -#endif -} - -// Sync only the consumer warpgroup on named barrier 1; producer warp must not participate. -__device__ __forceinline__ void bar_sync_consumer(int count) -{ - asm volatile("bar.sync 1, %0;" : : "r"(count)); -} - -// combine_modulator_chunk: combine 8 bf16 ts + 8 fp32 table into 8 bf16 (one uint4). -// Matches PyTorch eager `_get_ada_values` semantics: narrow fp32 table to -// bf16 FIRST, then bf16 hw add. Used for gate / scale / shift modulators alike. -__device__ __forceinline__ uint4 combine_modulator_chunk(uint4 const& ts_v, float4 const& tbl_lo, float4 const& tbl_hi) -{ - __nv_bfloat162 const* ts_b2 = reinterpret_cast<__nv_bfloat162 const*>(&ts_v); - __nv_bfloat162 out_b2[4]; - out_b2[0] = __hadd2(__float22bfloat162_rn(make_float2(tbl_lo.x, tbl_lo.y)), ts_b2[0]); - out_b2[1] = __hadd2(__float22bfloat162_rn(make_float2(tbl_lo.z, tbl_lo.w)), ts_b2[1]); - out_b2[2] = __hadd2(__float22bfloat162_rn(make_float2(tbl_hi.x, tbl_hi.y)), ts_b2[2]); - out_b2[3] = __hadd2(__float22bfloat162_rn(make_float2(tbl_hi.z, tbl_hi.w)), ts_b2[3]); - uint4 out_v; - *reinterpret_cast<__nv_bfloat162*>(reinterpret_cast(&out_v) + 0) = out_b2[0]; - *reinterpret_cast<__nv_bfloat162*>(reinterpret_cast(&out_v) + 1) = out_b2[1]; - *reinterpret_cast<__nv_bfloat162*>(reinterpret_cast(&out_v) + 2) = out_b2[2]; - *reinterpret_cast<__nv_bfloat162*>(reinterpret_cast(&out_v) + 3) = out_b2[3]; - return out_v; -} - -} // anonymous namespace - -template -__global__ void fusedDiTGateResidNormShiftScaleKernel(AdaLNNormParams p) -{ - static_assert(NUM_OUT == 1 || NUM_OUT == 2, "NUM_OUT must be 1 or 2"); - static_assert(!HAS_GATE || HAS_RESIDUAL, "HAS_GATE requires HAS_RESIDUAL"); - static_assert(HAS_NORM || HAS_RESIDUAL, "HAS_NORM=false requires HAS_RESIDUAL (the only output is x_new)"); - static_assert(HAS_NORM || !HAS_SHIFT_SCALE, "HAS_NORM=false implies HAS_SHIFT_SCALE=false"); - static_assert(HAS_NORM || !HAS_QUANT, "HAS_NORM=false implies HAS_QUANT=false"); - - constexpr int THREADS_PER_ROW = BLOCK_SIZE / ROWS_PER_BLOCK; - constexpr int WARPS_PER_ROW = THREADS_PER_ROW / 32; - constexpr int CHUNK_ELEMS = 8; // uint4 = 8 bf16 - constexpr int CHUNKS_PER_ROW = (D + THREADS_PER_ROW * CHUNK_ELEMS - 1) / (THREADS_PER_ROW * CHUNK_ELEMS); - constexpr int SF_VEC_SIZE = 16; - constexpr int SF_PER_ROW = D / SF_VEC_SIZE; - - static_assert(D % CHUNK_ELEMS == 0, "D must be multiple of 8"); - static_assert(D % SF_VEC_SIZE == 0, "D must be multiple of 16 (NVFP4 group size)"); - -#if (defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900) - asm volatile("griddepcontrol.wait;"); -#endif - - int const tid = threadIdx.x; - int const row_in_block = tid / THREADS_PER_ROW; - int const lane_in_row = tid % THREADS_PER_ROW; - int const row_warp = lane_in_row >> 5; - int const row_lane = lane_in_row & 31; - - int const tokenIdx = blockIdx.x * ROWS_PER_BLOCK + row_in_block; - bool const valid = (tokenIdx < p.num_tokens); - int const safeTokenIdx = valid ? tokenIdx : 0; - - int const batchIdx = safeTokenIdx / p.tokens_per_batch; - int64_t const tokenBase = static_cast(safeTokenIdx) * D; - - // Hybrid TMA / cp.async gating, picked empirically from paired bench on B200: - // USE_TMA : single-instruction cp.async.bulk for the X load. Wins on D=4096 - // except when (HAS_QUANT && HAS_SHIFT_SCALE && NUM_OUT==1) where the - // kernel is HBM-light and the mbarrier setup cost exceeds the LSU - // saving. Disabled on D=2048 (audio path, 2-4us kernels). - // USE_TMA_ATTN : also TMA-bulk the residual `attn` tensor into smem. Only enabled - // on (USE_TMA && HAS_RESIDUAL && HAS_QUANT) -- bf16 paths regress - // because the per-thread LDG for attn was already overlapping - // the TMA-X load via different LSU pipes, while quant's heavy - // Phase 2 hides the extra mbarrier wait. - constexpr bool USE_TMA = (D >= 4096) && !(HAS_QUANT && HAS_SHIFT_SCALE && NUM_OUT == 1); - constexpr bool USE_TMA_ATTN = USE_TMA && HAS_RESIDUAL && HAS_QUANT; - constexpr int kAttnSmemBytes = USE_TMA_ATTN ? (ROWS_PER_BLOCK * D * static_cast(sizeof(__nv_bfloat16))) : 0; - - extern __shared__ __align__(16) unsigned char smem_raw[]; - __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw); - __nv_bfloat16* smem_attn = reinterpret_cast<__nv_bfloat16*>(smem_raw + ROWS_PER_BLOCK * D * sizeof(__nv_bfloat16)); - float* warp_sums = reinterpret_cast(smem_raw + ROWS_PER_BLOCK * D * sizeof(__nv_bfloat16) + kAttnSmemBytes); - - // mbarrier in static SMEM (8B), only used when USE_TMA. Compiler elides when unused. - __shared__ alignas(8) uint64_t mbar; - - // Phase 0a: load X -> SMEM. Either TMA bulk (1 instruction, mbarrier-synced) or - // cp.async (32 per-thread issues, __pipeline_commit/wait_prior synced). -#if (defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900) - if constexpr (USE_TMA) - { - constexpr uint32_t kXBytes = ROWS_PER_BLOCK * D * static_cast(sizeof(__nv_bfloat16)); - constexpr uint32_t kAttnBytes = USE_TMA_ATTN ? kXBytes : 0; - constexpr uint32_t kTotalBytes = kXBytes + kAttnBytes; - static_assert(kXBytes % 16 == 0, "cp.async.bulk requires nbBytes multiple of 16"); - if (tid == 0) - { - asm volatile( - "mbarrier.init.shared.b64 [%0], 1;\n" ::"r"(static_cast(__cvta_generic_to_shared(&mbar))) - : "memory"); - asm volatile("mbarrier.arrive.expect_tx.shared.b64 _, [%0], %1;\n" ::"r"( - static_cast(__cvta_generic_to_shared(&mbar))), - "r"(kTotalBytes) - : "memory"); - // Bulk load X. - asm volatile("cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes [%0], [%1], %2, [%3];\n" ::"l"( - __cvta_generic_to_shared(smem_x + row_in_block * D)), - "l"(reinterpret_cast(p.x + tokenBase)), "r"(kXBytes), - "l"(__cvta_generic_to_shared(&mbar)) - : "memory"); - // Bulk load attn (only when USE_TMA_ATTN), into smem_attn slot. - if constexpr (USE_TMA_ATTN) - { - asm volatile( - "cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes [%0], [%1], %2, [%3];\n" ::"l"( - __cvta_generic_to_shared(smem_attn + row_in_block * D)), - "l"(reinterpret_cast(p.attn + tokenBase)), "r"(kAttnBytes), - "l"(__cvta_generic_to_shared(&mbar)) - : "memory"); - } - } - __syncthreads(); - } - else -#endif - { -#pragma unroll - for (int chunk = 0; chunk < CHUNKS_PER_ROW; chunk++) - { - int const elemBase = chunk * THREADS_PER_ROW * CHUNK_ELEMS + lane_in_row * CHUNK_ELEMS; - if (elemBase >= D) - continue; - __pipeline_memcpy_async(smem_x + row_in_block * D + elemBase, p.x + tokenBase + elemBase, 16); - } - __pipeline_commit(); - } - - // Phase 0b: load gate / scale / shift modulators into register caches. - // Per-thread storage scaled by HAS_GATE / HAS_SHIFT_SCALE flags; compiler elides - // unused slots when the corresponding flag is false. - uint4 gate_cache[HAS_GATE ? CHUNKS_PER_ROW : 1]; - uint4 scale_cache[HAS_SHIFT_SCALE ? NUM_OUT * CHUNKS_PER_ROW : 1]; - uint4 shift_cache[HAS_SHIFT_SCALE ? NUM_OUT * CHUNKS_PER_ROW : 1]; - - if constexpr (HAS_GATE || HAS_SHIFT_SCALE) - { - int64_t const gateBase = HAS_GATE ? static_cast(batchIdx) * p.gate_ts_stride : 0; - int64_t scaleBase[NUM_OUT]; - int64_t shiftBase[NUM_OUT]; - if constexpr (HAS_SHIFT_SCALE) - { -#pragma unroll - for (int k = 0; k < NUM_OUT; k++) - { - scaleBase[k] = static_cast(batchIdx) * p.scale_ts_stride[k]; - shiftBase[k] = static_cast(batchIdx) * p.shift_ts_stride[k]; - } - } -#pragma unroll - for (int chunk = 0; chunk < CHUNKS_PER_ROW; chunk++) - { - int const elemBase = chunk * THREADS_PER_ROW * CHUNK_ELEMS + lane_in_row * CHUNK_ELEMS; - if (elemBase >= D) - continue; - if constexpr (HAS_GATE) - { - uint4 const ts_v = *reinterpret_cast(&p.gate_ts[gateBase + elemBase]); - float4 const tbl_lo = *reinterpret_cast(&p.gate_table[elemBase + 0]); - float4 const tbl_hi = *reinterpret_cast(&p.gate_table[elemBase + 4]); - gate_cache[chunk] = combine_modulator_chunk(ts_v, tbl_lo, tbl_hi); - } - if constexpr (HAS_SHIFT_SCALE) - { -#pragma unroll - for (int k = 0; k < NUM_OUT; k++) - { - uint4 const s_ts_v = *reinterpret_cast(&p.scale_ts[k][scaleBase[k] + elemBase]); - float4 const s_tbl_lo = *reinterpret_cast(&p.scale_table[k][elemBase + 0]); - float4 const s_tbl_hi = *reinterpret_cast(&p.scale_table[k][elemBase + 4]); - scale_cache[k * CHUNKS_PER_ROW + chunk] = combine_modulator_chunk(s_ts_v, s_tbl_lo, s_tbl_hi); - - uint4 const h_ts_v = *reinterpret_cast(&p.shift_ts[k][shiftBase[k] + elemBase]); - float4 const h_tbl_lo = *reinterpret_cast(&p.shift_table[k][elemBase + 0]); - float4 const h_tbl_hi = *reinterpret_cast(&p.shift_table[k][elemBase + 4]); - shift_cache[k * CHUNKS_PER_ROW + chunk] = combine_modulator_chunk(h_ts_v, h_tbl_lo, h_tbl_hi); - } - } - } - } - - // Phase 0c: load attn into regs (HAS_RESIDUAL only). Either from smem (USE_TMA_ATTN — - // attn was bulk-loaded into smem_attn in Phase 0a) or directly from GMEM (cp.async path). - // The smem read happens AFTER the mbarrier wait, so smem_attn is guaranteed populated. - uint4 attn_reg[HAS_RESIDUAL ? CHUNKS_PER_ROW : 1]; - if constexpr (HAS_RESIDUAL && !USE_TMA_ATTN) - { -#pragma unroll - for (int chunk = 0; chunk < CHUNKS_PER_ROW; chunk++) - { - int const elemBase = chunk * THREADS_PER_ROW * CHUNK_ELEMS + lane_in_row * CHUNK_ELEMS; - if (elemBase >= D) - continue; - attn_reg[chunk] = *reinterpret_cast(&p.attn[tokenBase + elemBase]); - } - } - - // Wait on Phase 0a load completion (TMA mbarrier OR cp.async pipeline, gated by USE_TMA). -#if (defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900) - if constexpr (USE_TMA) - { - uint32_t const mbar_smem = static_cast(__cvta_generic_to_shared(&mbar)); - asm volatile( - "{\n" - ".reg .pred P1;\n" - "WAIT_LOOP:\n" - "mbarrier.try_wait.parity.shared.b64 P1, [%0], 0;\n" - "@P1 bra DONE;\n" - "bra WAIT_LOOP;\n" - "DONE:\n" - "}\n" ::"r"(mbar_smem) - : "memory"); - __syncthreads(); - } - else -#endif - { - __pipeline_wait_prior(0); - __syncthreads(); - } - - // Phase 1: compute x_new and sum^2. - // HAS_RESIDUAL=false: x_new = x - // HAS_RESIDUAL, !HAS_GATE: x_new = x + attn - // HAS_RESIDUAL, HAS_GATE: x_new = x + attn * gate - // x_new is cached in regs (xnew_cache) for Phase 2 to avoid a re-read; if - // HAS_RESIDUAL also write x_new to the separate x_out buffer so the downstream - // residual chain sees it (x stays read-only -- the op is functional). - uint4 xnew_cache[CHUNKS_PER_ROW]; - float sum2 = 0.0f; -#pragma unroll - for (int chunk = 0; chunk < CHUNKS_PER_ROW; chunk++) - { - int const elemBase = chunk * THREADS_PER_ROW * CHUNK_ELEMS + lane_in_row * CHUNK_ELEMS; - if (elemBase >= D) - continue; - uint4 const xv = *reinterpret_cast(&smem_x[row_in_block * D + elemBase]); - // USE_TMA_ATTN path: attn was TMA-bulk-loaded into smem_attn in Phase 0a; - // populate attn_reg from smem here (post-mbarrier-wait, data is valid). - if constexpr (USE_TMA_ATTN) - { - attn_reg[chunk] = *reinterpret_cast(&smem_attn[row_in_block * D + elemBase]); - } - uint const* xu = reinterpret_cast(&xv); - uint4 new_vec; - uint* nu = reinterpret_cast(&new_vec); - -#pragma unroll - for (int i = 0; i < 4; i++) - { - float2 xf = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&xu[i])); - float2 nf; - if constexpr (HAS_RESIDUAL) - { - uint4 const av = attn_reg[chunk]; - uint const* au = reinterpret_cast(&av); - float2 af = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&au[i])); - if constexpr (HAS_GATE) - { - uint4 const gv = gate_cache[chunk]; - uint const* gu = reinterpret_cast(&gv); - float2 gf = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&gu[i])); - nf.x = xf.x + af.x * gf.x; - nf.y = xf.y + af.y * gf.y; - } - else - { - nf.x = xf.x + af.x; - nf.y = xf.y + af.y; - } - } - else - { - nf = xf; - } - sum2 += nf.x * nf.x + nf.y * nf.y; - __nv_bfloat162 bf = __float22bfloat162_rn(nf); - reinterpret_cast<__nv_bfloat162&>(nu[i]) = bf; - } - xnew_cache[chunk] = new_vec; - if constexpr (HAS_RESIDUAL) - { - if (valid) - *reinterpret_cast(&p.x_out[tokenBase + elemBase]) = new_vec; - } - } - - // HAS_NORM=false variant: x_new was written to p.x_out above (HAS_RESIDUAL is required when - // !HAS_NORM). Skip the rms reduce + Phase 2 shift_scale/store; the gate-residual fused result - // lives entirely in the x_out residual-stream buffer. - if constexpr (!HAS_NORM) - { - return; - } - - // Per-row warp reduce + cross-warp reduce within the row. - sum2 = tensorrt_llm::common::warpReduceSum(sum2); - if (row_lane == 0) - warp_sums[row_in_block * WARPS_PER_ROW + row_warp] = sum2; - __syncthreads(); - float total = 0.0f; -#pragma unroll - for (int w = 0; w < WARPS_PER_ROW; w++) - total += warp_sums[row_in_block * WARPS_PER_ROW + w]; - float const rms_rcp = rsqrtf(total / static_cast(D) + p.eps); - - // Pre-read sf_scale broadcast scalars (HAS_QUANT only). - float sf_scale_val[NUM_OUT]; - if constexpr (HAS_QUANT) - { -#pragma unroll - for (int k = 0; k < NUM_OUT; k++) - sf_scale_val[k] = (p.sf_scale[k] != nullptr) ? *p.sf_scale[k] : 1.0f; - } - - // Phase 2: optional shift_scale + write outputs (NUM_OUT entries). -#pragma unroll - for (int chunk = 0; chunk < CHUNKS_PER_ROW; chunk++) - { - int const elemBase = chunk * THREADS_PER_ROW * CHUNK_ELEMS + lane_in_row * CHUNK_ELEMS; - if (elemBase >= D) - continue; - - uint4 const in_vec = xnew_cache[chunk]; - uint const* x_uints = reinterpret_cast(&in_vec); - - if constexpr (HAS_QUANT) - { - // FP4 quant path: fp32 shift_scale + max-abs scan + e2m1 pack. Inlining the - // max scan with shift_scale shares the pair-lane reduction with the FP4 - // conversion (saves ~10us/call vs cvt_float_to_fp4_inline + separate scan). - float vals[NUM_OUT][8]; - float localMax[NUM_OUT]; -#pragma unroll - for (int k = 0; k < NUM_OUT; k++) - localMax[k] = 0.0f; - -#pragma unroll - for (int i = 0; i < 4; i++) - { - float2 xv = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&x_uints[i])); - float const nx = xv.x * rms_rcp; - float const ny = xv.y * rms_rcp; -#pragma unroll - for (int k = 0; k < NUM_OUT; k++) - { - float yx; - float yy; - if constexpr (HAS_SHIFT_SCALE) - { - uint4 const sv = scale_cache[k * CHUNKS_PER_ROW + chunk]; - uint4 const hv = shift_cache[k * CHUNKS_PER_ROW + chunk]; - uint const* s_uints = reinterpret_cast(&sv); - uint const* h_uints = reinterpret_cast(&hv); - float2 sf = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&s_uints[i])); - float2 hf = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&h_uints[i])); - yx = nx * (1.0f + sf.x) + hf.x; - yy = ny * (1.0f + sf.y) + hf.y; - } - else - { - yx = nx; - yy = ny; - } - vals[k][2 * i + 0] = yx; - vals[k][2 * i + 1] = yy; - localMax[k] = fmaxf(localMax[k], fmaxf(fabsf(yx), fabsf(yy))); - } - } - -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) - uint32_t fp4_packed[NUM_OUT]; - uint8_t sfBytes[NUM_OUT]; -#pragma unroll - for (int k = 0; k < NUM_OUT; k++) - { - // Pair-lane max-abs across 16 elements (one SF block). - float const blockMax = fmaxf(__shfl_xor_sync(0xffffffff, localMax[k], 1), localMax[k]); - constexpr float kE2M1MaxRcp = 1.0f / 6.0f; - float const sfValue = sf_scale_val[k] * (blockMax * kE2M1MaxRcp); - __nv_fp8_e4m3 const sfFp8 = __nv_fp8_e4m3(sfValue); - sfBytes[k] = sfFp8.__x; - float const sfQuant = static_cast(sfFp8); - float const outScale = (blockMax != 0.0f) ? (sf_scale_val[k] / sfQuant) : 0.0f; -#pragma unroll - for (int i = 0; i < 8; i++) - vals[k][i] *= outScale; - fp4_packed[k] = fp32_vec_to_e2m1(vals[k]); - } - - int const colVecIdx = (chunk * THREADS_PER_ROW) + lane_in_row; - uint8_t* first_sf_ptr = cvt_quant_get_sf_out_offset(std::nullopt, - safeTokenIdx, colVecIdx, std::optional(p.num_tokens), SF_PER_ROW, p.out_sf[0], - QuantizationSFLayout::SWIZZLED); - if (valid && first_sf_ptr != nullptr) - { - *first_sf_ptr = sfBytes[0]; - if constexpr (NUM_OUT == 2) - { - uint8_t* second_sf_ptr = cvt_quant_get_sf_out_offset( - std::nullopt, safeTokenIdx, colVecIdx, std::optional(p.num_tokens), SF_PER_ROW, - p.out_sf[1], QuantizationSFLayout::SWIZZLED); - *second_sf_ptr = sfBytes[1]; - } - } - if (valid) - { - int64_t const fp4_row_off = static_cast(safeTokenIdx) * (D / 8); -#pragma unroll - for (int k = 0; k < NUM_OUT; k++) - p.out_fp4[k][fp4_row_off + colVecIdx] = fp4_packed[k]; - } -#endif - } - else - { - // bf16 shift_scale path: byte-matches eager apply_fused_*_shift_scale semantics. - // normed_bf16 = bf16(x_new_fp32 * rms_rcp_fp32) - // y_bf16 = bf16(normed_bf16 * bf16(1 + scale_bf16) + shift_bf16) if HAS_SHIFT_SCALE - // y_bf16 = normed_bf16 else - __nv_bfloat162 const one_b2 = __float2bfloat162_rn(1.0f); - uint4 out_vecs[NUM_OUT]; -#pragma unroll - for (int i = 0; i < 4; i++) - { - __nv_bfloat162 x_b2 = *reinterpret_cast<__nv_bfloat162 const*>(&x_uints[i]); - float2 xv = __bfloat1622float2(x_b2); - __nv_bfloat162 normed_b2 = __float22bfloat162_rn(make_float2(xv.x * rms_rcp, xv.y * rms_rcp)); -#pragma unroll - for (int k = 0; k < NUM_OUT; k++) - { - uint* o_uints = reinterpret_cast(&out_vecs[k]); - if constexpr (HAS_SHIFT_SCALE) - { - uint4 const sv = scale_cache[k * CHUNKS_PER_ROW + chunk]; - uint4 const hv = shift_cache[k * CHUNKS_PER_ROW + chunk]; - uint const* s_uints = reinterpret_cast(&sv); - uint const* h_uints = reinterpret_cast(&hv); - __nv_bfloat162 s_b2 = *reinterpret_cast<__nv_bfloat162 const*>(&s_uints[i]); - __nv_bfloat162 h_b2 = *reinterpret_cast<__nv_bfloat162 const*>(&h_uints[i]); - __nv_bfloat162 one_p_s = __hadd2(one_b2, s_b2); - __nv_bfloat162 y_b2 = __hadd2(__hmul2(normed_b2, one_p_s), h_b2); - reinterpret_cast<__nv_bfloat162&>(o_uints[i]) = y_b2; - } - else - { - reinterpret_cast<__nv_bfloat162&>(o_uints[i]) = normed_b2; - } - } - } - if (valid) - { -#pragma unroll - for (int k = 0; k < NUM_OUT; k++) - *reinterpret_cast(&p.out_bf16[k][tokenBase + elemBase]) = out_vecs[k]; - } - } - } - -#if (defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900) - asm volatile("griddepcontrol.launch_dependents;"); -#endif -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// -// Pipelined variant: multi-row CTA + circular SMEM stages + 1 producer warp / 8 consumer warps. -// -// Pipeline: -// - Per-stage SMEM: [X (D bf16) | attn (D bf16) if HAS_RESIDUAL]. -// - full_bar[NUM_STAGES]: producer signals "TMA done". -// - empty_bar[NUM_STAGES]: consumers signal "slot reusable" (init=CONSUMER_WARPS arrives). -// - Producer warp (lane 0 active): wait empty_bar, issue cp.async.bulk for X (and attn when -// HAS_RESIDUAL), arrive full_bar with the expected tx-bytes. -// - Consumer warps: wait full_bar, compute x_new + sum^2, normalize + shift_scale + write, -// arrive empty_bar. -// -// Caller contract: tokens_per_batch >= R_CTA && tokens_per_batch % R_CTA == 0 so all R_CTA -// rows in a CTA share batchIdx and modulator load amortizes once per CTA. -template -__global__ __launch_bounds__(288, 2) void fusedDiTGateResidNormShiftScaleKernelPipelined(AdaLNNormParams p) -{ - static_assert(D == 4096, "Pipelined variant requires D=4096"); - static_assert(R_CTA == 4, "Pipelined variant requires R_CTA=4"); - static_assert(NUM_STAGES == 2 || NUM_STAGES == 3, "NUM_STAGES must be 2 or 3"); - static_assert(NUM_OUT == 1 || NUM_OUT == 2, "NUM_OUT must be 1 or 2"); - static_assert(!HAS_GATE || HAS_RESIDUAL, "HAS_GATE requires HAS_RESIDUAL"); - - constexpr int CONSUMER_WARPS = 8; - constexpr int CONSUMER_THREADS = CONSUMER_WARPS * 32; // 256 - constexpr int VEC_ELEMS = 8; // bf16 per uint4 - constexpr int VEC_PER_THREAD = D / (CONSUMER_THREADS * VEC_ELEMS); // 2 for D=4096 - constexpr int SF_VEC_SIZE = 16; - constexpr int SF_PER_ROW = D / SF_VEC_SIZE; - static_assert(D % (CONSUMER_THREADS * VEC_ELEMS) == 0, "D must split evenly across 256 threads x 8 bf16"); - -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 900) - __trap(); // pipelined variant is sm_90+-only (TMA/mbarrier); never dispatched below Blackwell -#endif - -#if (defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900) - asm volatile("griddepcontrol.wait;"); -#endif - - int const tid = threadIdx.x; - int const warp_id = tid / 32; - int const lane_id = tid % 32; - bool const is_producer = (warp_id == 0); - int const consumer_id = tid - 32; // valid for warp_id >= 1 - int const consumer_warp = warp_id - 1; // 0..7 - - int const base_token = blockIdx.x * R_CTA; - int const batchIdx = base_token / p.tokens_per_batch; - - constexpr int kStageElems = HAS_RESIDUAL ? (2 * D) : D; - constexpr size_t kStageBytes = static_cast(kStageElems) * sizeof(__nv_bfloat16); - extern __shared__ __align__(16) unsigned char smem_raw[]; - __nv_bfloat16* smem_x_stage[NUM_STAGES]; - __nv_bfloat16* smem_attn_stage[NUM_STAGES]; -#pragma unroll - for (int s = 0; s < NUM_STAGES; s++) - { - smem_x_stage[s] = reinterpret_cast<__nv_bfloat16*>(smem_raw + s * kStageBytes); - smem_attn_stage[s] = HAS_RESIDUAL ? (smem_x_stage[s] + D) : nullptr; - } - uint64_t* full_bar = reinterpret_cast(smem_raw + NUM_STAGES * kStageBytes); - uint64_t* empty_bar = full_bar + NUM_STAGES; - float* warp_sums = reinterpret_cast(empty_bar + NUM_STAGES); - - if (tid == 0) - { -#pragma unroll - for (int s = 0; s < NUM_STAGES; s++) - { - mbar_init(&full_bar[s], 1); - mbar_init(&empty_bar[s], CONSUMER_WARPS); - } - } - - // Pre-load modulator caches into REGs once per CTA. All R_CTA rows share batchIdx. - constexpr int kScaleCacheSize = HAS_SHIFT_SCALE ? (NUM_OUT * VEC_PER_THREAD) : 1; - constexpr int kGateCacheSize = HAS_GATE ? VEC_PER_THREAD : 1; - uint4 scale_cache[kScaleCacheSize]; - uint4 shift_cache[kScaleCacheSize]; - uint4 gate_cache[kGateCacheSize]; - - if (!is_producer) - { - if constexpr (HAS_SHIFT_SCALE) - { -#pragma unroll - for (int k = 0; k < NUM_OUT; k++) - { - int64_t const scaleBase = static_cast(batchIdx) * p.scale_ts_stride[k]; - int64_t const shiftBase = static_cast(batchIdx) * p.shift_ts_stride[k]; -#pragma unroll - for (int v = 0; v < VEC_PER_THREAD; v++) - { - int const elem = (v * CONSUMER_THREADS + consumer_id) * VEC_ELEMS; - uint4 const s_ts = *reinterpret_cast(&p.scale_ts[k][scaleBase + elem]); - float4 const s_tbl_lo = *reinterpret_cast(&p.scale_table[k][elem + 0]); - float4 const s_tbl_hi = *reinterpret_cast(&p.scale_table[k][elem + 4]); - scale_cache[k * VEC_PER_THREAD + v] = combine_modulator_chunk(s_ts, s_tbl_lo, s_tbl_hi); - - uint4 const h_ts = *reinterpret_cast(&p.shift_ts[k][shiftBase + elem]); - float4 const h_tbl_lo = *reinterpret_cast(&p.shift_table[k][elem + 0]); - float4 const h_tbl_hi = *reinterpret_cast(&p.shift_table[k][elem + 4]); - shift_cache[k * VEC_PER_THREAD + v] = combine_modulator_chunk(h_ts, h_tbl_lo, h_tbl_hi); - } - } - } - if constexpr (HAS_GATE) - { - int64_t const gateBase = static_cast(batchIdx) * p.gate_ts_stride; -#pragma unroll - for (int v = 0; v < VEC_PER_THREAD; v++) - { - int const elem = (v * CONSUMER_THREADS + consumer_id) * VEC_ELEMS; - uint4 const g_ts = *reinterpret_cast(&p.gate_ts[gateBase + elem]); - float4 const g_tbl_lo = *reinterpret_cast(&p.gate_table[elem + 0]); - float4 const g_tbl_hi = *reinterpret_cast(&p.gate_table[elem + 4]); - gate_cache[v] = combine_modulator_chunk(g_ts, g_tbl_lo, g_tbl_hi); - } - } - } - - __syncthreads(); - - // Pre-fill empty_bar so producer's first iteration doesn't block. - if (!is_producer && lane_id == 0) - { -#pragma unroll - for (int s = 0; s < NUM_STAGES; s++) - mbar_arrive(&empty_bar[s]); - } - __syncthreads(); - - uint32_t stage = 0; - uint32_t phase_full = 0; - uint32_t phase_empty = 0; - constexpr uint32_t kXBytes = D * sizeof(__nv_bfloat16); - constexpr uint32_t kAttnBytes = HAS_RESIDUAL ? kXBytes : 0; - constexpr uint32_t kTotalBytes = kXBytes + kAttnBytes; - - if (is_producer) - { - if (lane_id == 0) - { -#pragma unroll - for (int sub = 0; sub < R_CTA; sub++) - { - int const token = base_token + sub; - if (token >= p.num_tokens) - break; - int64_t const tokenBase = static_cast(token) * D; - mbar_wait(&empty_bar[stage], phase_empty); - mbar_arrive_expect_tx(&full_bar[stage], kTotalBytes); - cp_async_bulk(smem_x_stage[stage], p.x + tokenBase, kXBytes, &full_bar[stage]); - if constexpr (HAS_RESIDUAL) - { - cp_async_bulk(smem_attn_stage[stage], p.attn + tokenBase, kAttnBytes, &full_bar[stage]); - } - stage = stage + 1; - if (stage == NUM_STAGES) - { - stage = 0; - phase_empty ^= 1u; - } - } - } - } - else - { -#pragma unroll - for (int sub = 0; sub < R_CTA; sub++) - { - int const token = base_token + sub; - bool const valid = (token < p.num_tokens); - int const safeToken = valid ? token : 0; - int64_t const tokenBase = static_cast(safeToken) * D; - - mbar_wait(&full_bar[stage], phase_full); - - // Phase 1: x_new = x [+ attn [* gate]]; cache x_new in regs, accumulate sum^2. - uint4 xnew_cache[VEC_PER_THREAD]; - float sum2 = 0.0f; -#pragma unroll - for (int v = 0; v < VEC_PER_THREAD; v++) - { - int const elem = (v * CONSUMER_THREADS + consumer_id) * VEC_ELEMS; - uint4 const xv = *reinterpret_cast(&smem_x_stage[stage][elem]); - uint const* xu = reinterpret_cast(&xv); - uint4 new_vec; - uint* nu = reinterpret_cast(&new_vec); - - if constexpr (HAS_RESIDUAL) - { - uint4 const av = *reinterpret_cast(&smem_attn_stage[stage][elem]); - uint const* au = reinterpret_cast(&av); - if constexpr (HAS_GATE) - { - uint4 const gv = gate_cache[v]; - uint const* gu = reinterpret_cast(&gv); -#pragma unroll - for (int i = 0; i < 4; i++) - { - float2 xf = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&xu[i])); - float2 af = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&au[i])); - float2 gf = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&gu[i])); - float2 nf; - nf.x = xf.x + af.x * gf.x; - nf.y = xf.y + af.y * gf.y; - sum2 += nf.x * nf.x + nf.y * nf.y; - reinterpret_cast<__nv_bfloat162&>(nu[i]) = __float22bfloat162_rn(nf); - } - } - else - { -#pragma unroll - for (int i = 0; i < 4; i++) - { - float2 xf = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&xu[i])); - float2 af = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&au[i])); - float2 nf; - nf.x = xf.x + af.x; - nf.y = xf.y + af.y; - sum2 += nf.x * nf.x + nf.y * nf.y; - reinterpret_cast<__nv_bfloat162&>(nu[i]) = __float22bfloat162_rn(nf); - } - } - } - else - { -#pragma unroll - for (int i = 0; i < 4; i++) - { - float2 xf = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&xu[i])); - sum2 += xf.x * xf.x + xf.y * xf.y; - reinterpret_cast(nu[i]) = xu[i]; - } - } - xnew_cache[v] = new_vec; - // Write x_new to the separate p.x_out buffer when HAS_RESIDUAL (op stays functional). - if constexpr (HAS_RESIDUAL) - { - if (valid) - *reinterpret_cast(&p.x_out[tokenBase + elem]) = new_vec; - } - } - - sum2 = tensorrt_llm::common::warpReduceSum(sum2); - if (lane_id == 0) - warp_sums[consumer_warp] = sum2; - bar_sync_consumer(CONSUMER_THREADS); - - float total = 0.0f; -#pragma unroll - for (int w = 0; w < CONSUMER_WARPS; w++) - total += warp_sums[w]; - float const rms_rcp = rsqrtf(total / static_cast(D) + p.eps); - - // ---- Phase 2: normalize + (optional shift_scale) + write NUM_OUT outputs ---- - float sf_scale_val[NUM_OUT]; - if constexpr (HAS_QUANT) - { -#pragma unroll - for (int k = 0; k < NUM_OUT; k++) - sf_scale_val[k] = (p.sf_scale[k] != nullptr) ? *p.sf_scale[k] : 1.0f; - } - - __nv_bfloat162 const one_b2 = __float2bfloat162_rn(1.0f); -#pragma unroll - for (int v = 0; v < VEC_PER_THREAD; v++) - { - int const elem = (v * CONSUMER_THREADS + consumer_id) * VEC_ELEMS; - uint4 const in_vec = xnew_cache[v]; - uint const* xu = reinterpret_cast(&in_vec); - - if constexpr (HAS_QUANT) - { - // Per-v NVFP4 path: fp32 shift_scale + max-abs scan + e2m1 pack. - // Each thread covers 8 bf16 (one uint4) per v, half SF block. - // Pair of adjacent lanes (lane ^ 1) shares one 16-element SF block. - float vals[NUM_OUT][8]; - float localMax[NUM_OUT]; -#pragma unroll - for (int k = 0; k < NUM_OUT; k++) - localMax[k] = 0.0f; - -#pragma unroll - for (int i = 0; i < 4; i++) - { - float2 xf = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&xu[i])); - float const nx = xf.x * rms_rcp; - float const ny = xf.y * rms_rcp; -#pragma unroll - for (int k = 0; k < NUM_OUT; k++) - { - float yx, yy; - if constexpr (HAS_SHIFT_SCALE) - { - uint4 const sv = scale_cache[k * VEC_PER_THREAD + v]; - uint4 const hv = shift_cache[k * VEC_PER_THREAD + v]; - uint const* su = reinterpret_cast(&sv); - uint const* hu = reinterpret_cast(&hv); - float2 sf = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&su[i])); - float2 hf = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&hu[i])); - yx = nx * (1.0f + sf.x) + hf.x; - yy = ny * (1.0f + sf.y) + hf.y; - } - else - { - yx = nx; - yy = ny; - } - vals[k][2 * i + 0] = yx; - vals[k][2 * i + 1] = yy; - localMax[k] = fmaxf(localMax[k], fmaxf(fabsf(yx), fabsf(yy))); - } - } - -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) - uint32_t fp4_packed[NUM_OUT]; - uint8_t sfBytes[NUM_OUT]; -#pragma unroll - for (int k = 0; k < NUM_OUT; k++) - { - // Pair-lane max-abs across 16 elements (one SF block). - float const blockMax = fmaxf(__shfl_xor_sync(0xffffffff, localMax[k], 1), localMax[k]); - constexpr float kE2M1MaxRcp = 1.0f / 6.0f; - float const sfValue = sf_scale_val[k] * (blockMax * kE2M1MaxRcp); - __nv_fp8_e4m3 const sfFp8 = __nv_fp8_e4m3(sfValue); - sfBytes[k] = sfFp8.__x; - float const sfQuant = static_cast(sfFp8); - float const outScale = (blockMax != 0.0f) ? (sf_scale_val[k] / sfQuant) : 0.0f; -#pragma unroll - for (int i = 0; i < 8; i++) - vals[k][i] *= outScale; - fp4_packed[k] = fp32_vec_to_e2m1(vals[k]); - } - - // colVecIdx: one uint32 of packed FP4 per (v, thread). Pair of adjacent - // lanes (lane ^ 1) shares one 16-element SF block (CVT_NUM_THREADS_PER_SF=2). - int const colVecIdx = v * CONSUMER_THREADS + consumer_id; - uint8_t* first_sf_ptr = cvt_quant_get_sf_out_offset( - std::nullopt, safeToken, colVecIdx, std::optional(p.num_tokens), SF_PER_ROW, p.out_sf[0], - QuantizationSFLayout::SWIZZLED); - if (valid && first_sf_ptr != nullptr) - { - *first_sf_ptr = sfBytes[0]; - if constexpr (NUM_OUT == 2) - { - uint8_t* second_sf_ptr - = cvt_quant_get_sf_out_offset(std::nullopt, - safeToken, colVecIdx, std::optional(p.num_tokens), SF_PER_ROW, p.out_sf[1], - QuantizationSFLayout::SWIZZLED); - *second_sf_ptr = sfBytes[1]; - } - } - if (valid) - { - int64_t const fp4_row_off = static_cast(safeToken) * (D / 8); -#pragma unroll - for (int k = 0; k < NUM_OUT; k++) - p.out_fp4[k][fp4_row_off + colVecIdx] = fp4_packed[k]; - } -#endif - } - else - { - uint4 out_vecs[NUM_OUT]; -#pragma unroll - for (int i = 0; i < 4; i++) - { - __nv_bfloat162 x_b2 = *reinterpret_cast<__nv_bfloat162 const*>(&xu[i]); - float2 xf = __bfloat1622float2(x_b2); - __nv_bfloat162 normed = __float22bfloat162_rn(make_float2(xf.x * rms_rcp, xf.y * rms_rcp)); -#pragma unroll - for (int k = 0; k < NUM_OUT; k++) - { - uint* o_u = reinterpret_cast(&out_vecs[k]); - if constexpr (HAS_SHIFT_SCALE) - { - uint4 const sv = scale_cache[k * VEC_PER_THREAD + v]; - uint4 const hv = shift_cache[k * VEC_PER_THREAD + v]; - uint const* su = reinterpret_cast(&sv); - uint const* hu = reinterpret_cast(&hv); - __nv_bfloat162 s_b2 = *reinterpret_cast<__nv_bfloat162 const*>(&su[i]); - __nv_bfloat162 h_b2 = *reinterpret_cast<__nv_bfloat162 const*>(&hu[i]); - __nv_bfloat162 one_p_s = __hadd2(one_b2, s_b2); - __nv_bfloat162 y_b2 = __hadd2(__hmul2(normed, one_p_s), h_b2); - reinterpret_cast<__nv_bfloat162&>(o_u[i]) = y_b2; - } - else - { - reinterpret_cast<__nv_bfloat162&>(o_u[i]) = normed; - } - } - } - if (valid) - { -#pragma unroll - for (int k = 0; k < NUM_OUT; k++) - *reinterpret_cast(&p.out_bf16[k][tokenBase + elem]) = out_vecs[k]; - } - } - } - - bar_sync_consumer(CONSUMER_THREADS); - if (lane_id == 0) - mbar_arrive(&empty_bar[stage]); - - stage = stage + 1; - if (stage == NUM_STAGES) - { - stage = 0; - phase_full ^= 1u; - } - } - } - -#if (defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900) - asm volatile("griddepcontrol.launch_dependents;"); -#endif -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// -// Templated host launcher (compile-time variant flags) + a runtime-flag dispatch entry point that -// maps the op's inferred flags onto the right specialization. Only the variants the dispatcher -// references are instantiated. Consumed by LTX-2's transformer block. - -namespace -{ -template -void launchFusedDiTGateResidNormShiftScaleKernelImpl(AdaLNNormParams const& params, int hidden_dim, cudaStream_t stream) -{ - TLLM_CHECK_WITH_INFO(params.num_tokens >= 1, "num_tokens must be >= 1, got %d", params.num_tokens); - TLLM_CHECK_WITH_INFO( - params.tokens_per_batch >= 1, "tokens_per_batch must be >= 1, got %d", params.tokens_per_batch); - TLLM_CHECK_WITH_INFO(params.num_tokens % params.tokens_per_batch == 0, - "num_tokens (%d) must be divisible by tokens_per_batch (%d)", params.num_tokens, params.tokens_per_batch); - - // Production tile selected via NCU sweep at D=4096, B=1, T=15360: 1-row/CTA, 256 threads/CTA - // gives 2 CTAs/SM (regs <= 80) and the highest per-element bench rate among {2r256, 1r256, 1r512}. - constexpr int ROWS_PER_BLOCK = 1; - constexpr int BLOCK_SIZE = 256; - constexpr int WARPS_PER_ROW = (BLOCK_SIZE / ROWS_PER_BLOCK) / 32; - - cudaLaunchConfig_t cfg = {}; - cfg.stream = stream; - cudaLaunchAttribute attrs[1] = {}; - attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; - attrs[0].val.programmaticStreamSerializationAllowed = 1; - cfg.attrs = attrs; - cfg.numAttrs = 1; - - // Pipelined dispatch (multi-row CTA + circular TMA stages + warp specialization). - // Auto-selected on shapes where it wins; otherwise falls through to the default path below. - // - hidden_dim == 4096 (video path; audio D=2048 has too small a grid to amortize the - // warp-specialization overhead and bench-regresses) - // - HAS_SHIFT_SCALE || !HAS_GATE (the HAS_GATE && !HAS_SHIFT_SCALE variant is HBM-bound and bench-regresses) - // - tokens_per_batch >= 4 && % 4 == 0 (R_CTA=4 rows share batchIdx for per-CTA mod cache) - // - num_tokens % 4 == 0 (whole grid must be R_CTA-aligned) - constexpr bool kPipelinedVariantOK = HAS_SHIFT_SCALE || !HAS_GATE; - if constexpr (kPipelinedVariantOK) - { - // NUM_STAGES chosen per HAS_QUANT: bf16 path is HBM-latency-bound so deeper pipeline - // (3 stages) helps; quant Phase 2 is compute-heavy and extra stages just add SMEM - // pressure without commensurate latency hiding. - constexpr int PIPE_NUM_STAGES = HAS_QUANT ? 2 : 3; - constexpr int PIPE_BLOCK_SIZE = 288; - constexpr int PIPE_D = 4096; - constexpr int PIPE_R_CTA = 4; - // Pipelined variant uses sm_90+ TMA/mbarrier; on older archs (sm_80/86/89) fall through to the - // V1 cp.async path below (functionally correct, just no TMA/pipelining). Cache the SM query -- - // getSMVersion() issues cudaGetDevice + cudaDeviceGetAttribute, too costly on the per-launch path. - static int const kSmVersion = tensorrt_llm::common::getSMVersion(); - if (kSmVersion >= 90 && hidden_dim == PIPE_D && (params.num_tokens % PIPE_R_CTA == 0) - && (params.tokens_per_batch >= PIPE_R_CTA) && (params.tokens_per_batch % PIPE_R_CTA == 0)) - { - constexpr int pipe_stage_elems = HAS_RESIDUAL ? (2 * PIPE_D) : PIPE_D; - size_t const pipe_smem_bytes - = static_cast(PIPE_NUM_STAGES) * pipe_stage_elems * sizeof(__nv_bfloat16) - + 2 * PIPE_NUM_STAGES * sizeof(uint64_t) + 8 * sizeof(float) + 16; - cudaLaunchConfig_t pipe_cfg = cfg; - pipe_cfg.gridDim = dim3((params.num_tokens + PIPE_R_CTA - 1) / PIPE_R_CTA); - pipe_cfg.blockDim = dim3(PIPE_BLOCK_SIZE); - pipe_cfg.dynamicSmemBytes = static_cast(pipe_smem_bytes); - // Pipelined kernel is only built for HAS_NORM=true variants. The dispatch predicate - // above requires HAS_SHIFT_SCALE || !HAS_GATE which is satisfied only by HAS_NORM=true - // paths; the gate-residual-only variant (HAS_SHIFT_SCALE=false, HAS_GATE=true) fails the - // predicate and falls through to the default tile. - static_assert(HAS_NORM, "pipelined dispatch requires HAS_NORM"); - auto* pipe_kp = fusedDiTGateResidNormShiftScaleKernelPipelined; - cudaFuncSetAttribute( - pipe_kp, cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(pipe_smem_bytes)); - cudaLaunchKernelEx(&pipe_cfg, pipe_kp, params); - return; - } - } - -#define LAUNCH(D_VAL) \ - do \ - { \ - constexpr bool USE_TMA_HOST = (D_VAL >= 4096) && !(HAS_QUANT && HAS_SHIFT_SCALE && NUM_OUT == 1); \ - constexpr bool USE_TMA_ATTN_HOST = USE_TMA_HOST && HAS_RESIDUAL && HAS_QUANT; \ - int const attn_extra_bytes \ - = USE_TMA_ATTN_HOST ? (ROWS_PER_BLOCK * D_VAL * static_cast(sizeof(__nv_bfloat16))) : 0; \ - cfg.gridDim = dim3((params.num_tokens + ROWS_PER_BLOCK - 1) / ROWS_PER_BLOCK); \ - cfg.blockDim = dim3(BLOCK_SIZE); \ - cfg.dynamicSmemBytes = ROWS_PER_BLOCK * D_VAL * static_cast(sizeof(__nv_bfloat16)) + attn_extra_bytes \ - + ROWS_PER_BLOCK * WARPS_PER_ROW * static_cast(sizeof(float)); \ - cudaLaunchKernelEx(&cfg, \ - fusedDiTGateResidNormShiftScaleKernel, \ - params); \ - } while (0) - - switch (hidden_dim) - { - case 2048: LAUNCH(2048); break; - case 4096: LAUNCH(4096); break; - default: - TLLM_THROW( - "Unsupported hidden_dim for fusedDiTGateResidNormShiftScaleKernel: %d (only 2048, 4096)", hidden_dim); - } -#undef LAUNCH -} - -} // anonymous namespace - -// Runtime-flag dispatch entry. The op has already validated (residual, gate, norm, shift_scale, quant) -// against kSupportedVariants, so every line below is a direct "args == this tuple -> launch this tuple" -// map: the compared flags and the template arguments are the same literals (no negation / elimination -// ordering). NUM_OUT is fixed by the variant those flags select (2 only for the residual-no-gate dual); -// gate_resid (no norm) is bf16-only, so it has no NVFP4 line. Only the tuples listed here are instantiated. -void launchFusedDiTGateResidNormShiftScaleKernel(AdaLNNormParams const& params, bool residual, bool gate, bool norm, - bool shift_scale, bool quant, int hidden_dim, cudaStream_t stream) -{ -#define DISPATCH(R, G, N, SS, NO, QUANT) \ - if (residual == (R) && gate == (G) && norm == (N) && shift_scale == (SS) && quant == (QUANT)) \ - return launchFusedDiTGateResidNormShiftScaleKernelImpl(params, hidden_dim, stream) - - DISPATCH(false, false, true, true, 1, false); // rmsnorm_shift_scale - DISPATCH(false, false, true, true, 1, true); // rmsnorm_shift_scale (NVFP4) - DISPATCH(true, false, true, true, 2, false); // resid_rmsnorm_shift_scale_dual - DISPATCH(true, false, true, true, 2, true); // resid_rmsnorm_shift_scale_dual (NVFP4) - DISPATCH(true, true, true, true, 1, false); // gate_resid_rmsnorm_shift_scale - DISPATCH(true, true, true, true, 1, true); // gate_resid_rmsnorm_shift_scale (NVFP4) - DISPATCH(true, true, true, false, 1, false); // gate_resid_rmsnorm - DISPATCH(true, true, true, false, 1, true); // gate_resid_rmsnorm (NVFP4) - DISPATCH(true, true, false, false, 1, false); // gate_resid (bf16-only) -#undef DISPATCH - - TLLM_THROW( - "Unsupported fusedDiTGateResidNormShiftScale variant: residual=%d gate=%d norm=%d shift_scale=%d " - "quant=%d (op-side validation should have rejected this)", - residual, gate, norm, shift_scale, quant); -} - -} // namespace kernels - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/fusedDiTGateResidNormShiftScaleKernel.h b/cpp/tensorrt_llm/kernels/fusedDiTGateResidNormShiftScaleKernel.h deleted file mode 100644 index 9e4500df2f78..000000000000 --- a/cpp/tensorrt_llm/kernels/fusedDiTGateResidNormShiftScaleKernel.h +++ /dev/null @@ -1,110 +0,0 @@ -/* - * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. - * - * 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. - */ - -#ifndef TRTLLM_FUSEDDITGATERESIDNORMSHIFTSCALEKERNEL_H -#define TRTLLM_FUSEDDITGATERESIDNORMSHIFTSCALEKERNEL_H - -#include "tensorrt_llm/common/config.h" -#include -#include -#include - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels -{ - -// Single fused DiT pre-block kernel covering all AdaLN variants via template flags. -// -// Pipeline (compile-time selected, all "pluggable" except RmsNorm which is always on): -// 1. Phase 0a -- cp.async x -> SMEM (always) -// 2. Phase 0b -- combine modulator (table, ts) pairs into bf16 caches: -// gate (if HAS_GATE) -// NUM_OUT scale + shift (if HAS_SHIFT_SCALE) -// 3. Phase 0c -- load attn into reg cache (if HAS_RESIDUAL) -// 4. Phase 1 -- x_new = x [+ attn [* gate]]; sum^2; write x_new to x_out if HAS_RESIDUAL -// 5. Reduce -- per-row warp + cross-warp reduce -> rms_rcp (always) -// 6. Phase 2 -- write NUM_OUT outputs: -// if HAS_SHIFT_SCALE: y[k] = (1 + scale[k]) * normed + shift[k] -// else : y = normed -// HAS_QUANT=false: bf16 store to out_bf16[k] -// HAS_QUANT=true : NVFP4 + 128x4 swizzled SF to (out_fp4[k], out_sf[k]) -// -// Specializations used by LTX-2: -// rmsnorm_shift_scale: HAS_RESIDUAL=false, HAS_GATE=false, HAS_SHIFT_SCALE=true, NUM_OUT=1 -// resid_rmsnorm_shift_scale_dual: HAS_RESIDUAL=true, HAS_GATE=false, HAS_SHIFT_SCALE=true, NUM_OUT=2 -// gate_resid_rmsnorm_shift_scale: HAS_RESIDUAL=true, HAS_GATE=true, HAS_SHIFT_SCALE=true, NUM_OUT=1 -// gate_resid_rmsnorm: HAS_RESIDUAL=true, HAS_GATE=true, HAS_SHIFT_SCALE=false, NUM_OUT=1 -// -// HAS_GATE implies HAS_RESIDUAL (asserted at compile time). -// -// Modulator combine matches PyTorch eager `_get_all_ada_values` semantics: -// narrow fp32 table to bf16 first, then bf16 hw add (`__hadd2`). -// -// Tile: production hardcoded to (ROWS_PER_BLOCK=1, BLOCK_SIZE=256). NCU sweep on B200 at -// the production shape found 1r256 the best balance: 2r256 has too much per-thread register -// pressure (12.5% theoretical occupancy on the dual-output bf16 variant); 1r512 reaches higher occupancy but its -// 16-warp CTAs starve the SM warp schedulers and pay a heavier __syncthreads() barrier. -// -// Supported hidden_dim: 2048 (LTX-2 audio) and 4096 (LTX-2 video). - -struct AdaLNNormParams -{ - // === Input === - __nv_bfloat16* x = nullptr; // [num_tokens, D] bf16. Read-only. - __nv_bfloat16* x_out = nullptr; // HAS_RESIDUAL: [num_tokens, D] bf16. x_new = x [+ attn [* gate]] - // written here (separate buffer -> op stays functional, no clone of x). - __nv_bfloat16 const* attn = nullptr; // HAS_RESIDUAL: [num_tokens, D] bf16 - - // === Gate modulator (HAS_GATE => HAS_RESIDUAL) === - float const* gate_table = nullptr; // [D] fp32, broadcast over batch - __nv_bfloat16 const* gate_ts = nullptr; // [batch, D] bf16 - int gate_ts_stride = 0; // inner stride between batches - - // === Affine modulators (HAS_SHIFT_SCALE) -- up to NUM_OUT={1,2} entries === - float const* scale_table[2] = {nullptr, nullptr}; - __nv_bfloat16 const* scale_ts[2] = {nullptr, nullptr}; - int scale_ts_stride[2] = {0, 0}; - float const* shift_table[2] = {nullptr, nullptr}; - __nv_bfloat16 const* shift_ts[2] = {nullptr, nullptr}; - int shift_ts_stride[2] = {0, 0}; - - // === Outputs (NUM_OUT entries) === - __nv_bfloat16* out_bf16[2] = {nullptr, nullptr}; // HAS_QUANT=false - uint32_t* out_fp4[2] = {nullptr, nullptr}; // HAS_QUANT=true - uint32_t* out_sf[2] = {nullptr, nullptr}; // HAS_QUANT=true: SWIZZLED 128x4 SF - float const* sf_scale[2] = {nullptr, nullptr}; // HAS_QUANT=true: scalar broadcast - - // === Shape === - int num_tokens = 0; - int tokens_per_batch = 0; - float eps = 1e-6f; -}; - -// Launch the unified kernel. Pass the variant as runtime flags; the launcher maps them onto the -// matching compile-time specialization (selects production tile 1r256, dispatches on hidden_dim). -// The caller populates only the params relevant to the chosen variant; unused fields stay default. -// -// Supported hidden_dim: 2048, 4096. NUM_OUT is implied by the variant (2 for residual-no-gate dual, -// else 1). gate_resid (norm=false) has no NVFP4-quant variant -- pass quant=false for it. -void launchFusedDiTGateResidNormShiftScaleKernel(AdaLNNormParams const& params, bool residual, bool gate, bool norm, - bool shift_scale, bool quant, int hidden_dim, cudaStream_t stream); - -} // namespace kernels - -TRTLLM_NAMESPACE_END - -#endif // TRTLLM_FUSEDDITGATERESIDNORMSHIFTSCALEKERNEL_H diff --git a/cpp/tensorrt_llm/kernels/gptKernels.cu b/cpp/tensorrt_llm/kernels/gptKernels.cu index c30239bf827f..082709e7af35 100644 --- a/cpp/tensorrt_llm/kernels/gptKernels.cu +++ b/cpp/tensorrt_llm/kernels/gptKernels.cu @@ -79,13 +79,8 @@ __global__ __launch_bounds__(THREADS_PER_BLOCK) void computeSeqAndPaddingOffsets // Fixed Q sequence lengths. bool const fixed_q_seqlen = params.seqQLengths == nullptr; - // Whether to reuse externally computed cumulative Q/KV sequence lengths. - bool const usePrecomputedQOffsets = params.precomputedSeqQOffsets != nullptr; - bool const usePrecomputedKVOffsets = params.precomputedSeqKVOffsets != nullptr; - // Whether to calculate cumulative KV sequence lengths. - bool const calculate_kv_offsets = params.seqKVOffsets != nullptr && !usePrecomputedKVOffsets; - bool const needKVOffsets = calculate_kv_offsets || usePrecomputedKVOffsets; + bool const calculate_kv_offsets = params.seqKVOffsets != nullptr; // Whether to calculate cumulative packed mask rows. bool const calculate_packed_mask_row_offsets = params.packedMaskRowOffsets != nullptr; @@ -95,7 +90,7 @@ __global__ __launch_bounds__(THREADS_PER_BLOCK) void computeSeqAndPaddingOffsets bool const calculate_cp_offsets = cpSize > 1 && params.seqCpPartialOffsets != nullptr; // Compute the padding offsets for Encoder Inputs. - bool const need_encoder_padding_offsets = (params.encoderPaddingOffsets != nullptr) && needKVOffsets; + bool const need_encoder_padding_offsets = (params.encoderPaddingOffsets != nullptr) && calculate_kv_offsets; [[maybe_unused]] int* smemEncoderSeqQOffsets; // The implementation of the parallel scan in the thread block (see CUB for details). @@ -147,21 +142,11 @@ __global__ __launch_bounds__(THREADS_PER_BLOCK) void computeSeqAndPaddingOffsets } // Do the prefix-scan (it calls syncthreads internally). - int seqQOffset = 0; + int seqQOffset; [[maybe_unused]] int packedMaskRowOffset; - [[maybe_unused]] int seqKVOffset = 0; + [[maybe_unused]] int seqKVOffset; [[maybe_unused]] int seqCpPartialOffset; - if (usePrecomputedQOffsets) - { - if (batchIdx <= batchSizeBound) - { - seqQOffset = params.precomputedSeqQOffsets[batchIdx]; - } - } - else - { - BlockScan(tempQStorage).ExclusiveSum(seqQLength, seqQOffset, prefixQOp); - } + BlockScan(tempQStorage).ExclusiveSum(seqQLength, seqQOffset, prefixQOp); if (calculate_packed_mask_row_offsets) { BlockScan(tempMaskStorage).ExclusiveSum(packedMaskRows, packedMaskRowOffset, prefixMaskOp); @@ -170,10 +155,6 @@ __global__ __launch_bounds__(THREADS_PER_BLOCK) void computeSeqAndPaddingOffsets { BlockScan(tempKVStorage).ExclusiveSum(seqKVLength, seqKVOffset, prefixKVOp); } - else if (usePrecomputedKVOffsets && batchIdx <= batchSizeBound) - { - seqKVOffset = params.precomputedSeqKVOffsets[batchIdx]; - } if (calculate_cp_offsets) { BlockScan(tempKVStorage).ExclusiveSum(seqCpPartialLength, seqCpPartialOffset, prefixCpPartialOp); @@ -192,10 +173,7 @@ __global__ __launch_bounds__(THREADS_PER_BLOCK) void computeSeqAndPaddingOffsets // Store the result. if (batchIdx <= batchSizeBound && storeSeqOffsets) { - if (!usePrecomputedQOffsets) - { - params.seqQOffsets[batchIdx] = seqQOffset; - } + params.seqQOffsets[batchIdx] = seqQOffset; if (calculate_packed_mask_row_offsets) { params.packedMaskRowOffsets[batchIdx] = packedMaskRowOffset; @@ -332,8 +310,8 @@ void invokeBuildDecoderInfo(BuildDecoderInfoParams const& params, cudaStream_ "Rotary embedding dim is assumed to be smaller than 512 and multiple of 2."); TLLM_CHECK_WITH_INFO( !(params.seqKVLengths == nullptr && params.rotaryEmbeddingDim > 0), "KV sequence lengths buffer is invalid."); - bool const needKVOffsets = params.seqKVOffsets != nullptr || params.precomputedSeqKVOffsets != nullptr; - bool const need_encoder_padding_offsets = (params.encoderPaddingOffsets != nullptr) && needKVOffsets; + bool const need_encoder_padding_offsets + = (params.encoderPaddingOffsets != nullptr) && (params.seqKVOffsets != nullptr); const size_t smem_size = (need_encoder_padding_offsets ? (params.batchSize + 1) * 2 : (params.batchSize + 1)) * sizeof(int); computeSeqAndPaddingOffsets diff --git a/cpp/tensorrt_llm/kernels/gptKernels.h b/cpp/tensorrt_llm/kernels/gptKernels.h index e13e9bca4d6a..c8d9e7e202ee 100644 --- a/cpp/tensorrt_llm/kernels/gptKernels.h +++ b/cpp/tensorrt_llm/kernels/gptKernels.h @@ -115,10 +115,6 @@ struct BuildDecoderInfoParams int* seqQOffsets; // The offsets to the 1st token in each sequence of KV buffer. Shape: [batchSize+1]. int* seqKVOffsets; - // Precomputed offsets to the 1st token in each sequence of Q buffer. Shape: [batchSize+1]. - int const* precomputedSeqQOffsets; - // Precomputed offsets to the 1st token in each sequence of KV buffer. Shape: [batchSize+1]. - int const* precomputedSeqKVOffsets; // The number of padded tokens in the corresponding padded tensor before the current token, for Decoder. Shape: // [numTokens]. int* paddingOffsets; diff --git a/cpp/tensorrt_llm/kernels/heuristicTopKDecode.cu b/cpp/tensorrt_llm/kernels/heuristicTopKDecode.cu index 839ae5483c05..25b0de3ef0a5 100644 --- a/cpp/tensorrt_llm/kernels/heuristicTopKDecode.cu +++ b/cpp/tensorrt_llm/kernels/heuristicTopKDecode.cu @@ -1,5 +1,5 @@ /* - * Copyright (c) 2019-2026, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2019-2025, NVIDIA CORPORATION. All rights reserved. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -47,19 +47,15 @@ using heuristic_topk::KernelSmemTplK; // same kernel template. Smem layout is derived from GvrParams // at compile time. template -__global__ void __launch_bounds__(BLOCK_SIZE) - heuristicTopKMultiRowKernel(float const* __restrict__ logits, int const* __restrict__ seqLens, - int const* __restrict__ preIdx, float* __restrict__ scratchValues, int* __restrict__ outIndices, int stride0, - int next_n, int topK, int preIdxStride, int preIdxCount, int compressRatio) +__global__ void __launch_bounds__(BLOCK_SIZE) heuristicTopKMultiRowKernel(float const* __restrict__ logits, + int const* __restrict__ seqLens, int const* __restrict__ preIdx, float* __restrict__ scratchValues, + int* __restrict__ outIndices, int stride0, int next_n, int topK, int preIdxStride, int preIdxCount) { using SmemT = KernelSmemTplK::kC, GvrParams::kNumBins>; int const rowIdx = blockIdx.x; int const seq_len = seqLens[rowIdx / next_n]; - // seqLens is in uncompressed token space; the logits/preIdx live in - // compressed-index space when compressRatio > 1 (DSv4 indexer). - int const actual_kv_len = seq_len - next_n + (rowIdx % next_n) + 1; - int const N = actual_kv_len / compressRatio; + int const N = seq_len - next_n + (rowIdx % next_n) + 1; float const* __restrict__ input = logits + static_cast(rowIdx) * stride0; int const* __restrict__ rowPreIdx = preIdx + static_cast(rowIdx / next_n) * preIdxStride; @@ -88,19 +84,9 @@ __global__ void __launch_bounds__(BLOCK_SIZE) return; } - // Temporal-shift offset to map prev-step's top-K indices into this step's - // KV index space. - // compressRatio == 1 (DSv3.2): +1 — KV grew by exactly 1 token per - // decode step; prev indices were at seq_len-1 so a uniform +1 maps - // them to the equivalent positions under the indexer's "newest-first" - // layout. The (rowIdx % next_n) addend extends this to MTP windows. - // compressRatio == 4 (DSv4): 0 — in compressed-index space new - // compressed entries are appended at the end; prev indices in - // [0, c_prev-1] remain valid as-is. Per-row Δc varies (0 or 1) with - // prev kv_len mod 4 alignment, but a uniform offset of 0 stays - // within-bounds for all rows and preserves the temporal-correlation - // hint (vertical top-K consistency validated offline). - int const preIdxOffset = (compressRatio == 1) ? ((rowIdx % next_n) + 1) : 0; + // +1 accounts for the temporal shift: prev_topk indices were computed at + // seq_len-1, but the current step has one additional KV token appended. + int const preIdxOffset = (rowIdx % next_n) + 1; gvrTopKJob(input, N, rowPreIdx, preIdxCount, topK, outputValues, outputIndices, smem, preIdxOffset); #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) cudaTriggerProgrammaticLaunchCompletion(); @@ -117,18 +103,16 @@ __global__ void __launch_bounds__(BLOCK_SIZE) // Templated on (InputT, TopK). Smem layout is derived from // GvrParams. template -__global__ void __launch_bounds__(BLOCK_SIZE) - heuristicTopKMultiRowKernelDtype(InputT const* __restrict__ logits, int const* __restrict__ seqLens, - int const* __restrict__ preIdx, InputT* __restrict__ scratchValues, int* __restrict__ outIndices, int stride0, - int next_n, int topK, int preIdxStride, int preIdxCount, int compressRatio) +__global__ void __launch_bounds__(BLOCK_SIZE) heuristicTopKMultiRowKernelDtype(InputT const* __restrict__ logits, + int const* __restrict__ seqLens, int const* __restrict__ preIdx, InputT* __restrict__ scratchValues, + int* __restrict__ outIndices, int stride0, int next_n, int topK, int preIdxStride, int preIdxCount) { // dtype path uses fp32 keys[] in smem (down-conversion deferred to writeback). using SmemT = KernelSmemTplK::kC, GvrParams::kNumBins>; int const rowIdx = blockIdx.x; int const seq_len = seqLens[rowIdx / next_n]; - int const actual_kv_len = seq_len - next_n + (rowIdx % next_n) + 1; - int const N = actual_kv_len / compressRatio; + int const N = seq_len - next_n + (rowIdx % next_n) + 1; InputT const* __restrict__ input = logits + static_cast(rowIdx) * stride0; int const* __restrict__ rowPreIdx = preIdx + static_cast(rowIdx / next_n) * preIdxStride; @@ -158,8 +142,7 @@ __global__ void __launch_bounds__(BLOCK_SIZE) return; } - // See fp32 path: cr==1 → (rowIdx % next_n)+1; cr!=1 (DSv4) → 0. - int const preIdxOffset = (compressRatio == 1) ? ((rowIdx % next_n) + 1) : 0; + int const preIdxOffset = (rowIdx % next_n) + 1; gvrTopKJobDtype( input, N, rowPreIdx, preIdxCount, topK, outputValues, outputIndices, smem, preIdxOffset); #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) @@ -169,25 +152,24 @@ __global__ void __launch_bounds__(BLOCK_SIZE) // Explicit instantiations — 6 (dtype × K) combos. Launchers dispatch on // runtime topK via switch, so all 6 must be available at link time. -// Trailing `int` is the compressRatio parameter (1 = V3.2, 4 = V4 indexer). template __global__ void heuristicTopKMultiRowKernelDtype<__nv_bfloat16, 512>( - __nv_bfloat16 const*, int const*, int const*, __nv_bfloat16*, int*, int, int, int, int, int, int); + __nv_bfloat16 const*, int const*, int const*, __nv_bfloat16*, int*, int, int, int, int, int); template __global__ void heuristicTopKMultiRowKernelDtype<__nv_bfloat16, 1024>( - __nv_bfloat16 const*, int const*, int const*, __nv_bfloat16*, int*, int, int, int, int, int, int); + __nv_bfloat16 const*, int const*, int const*, __nv_bfloat16*, int*, int, int, int, int, int); template __global__ void heuristicTopKMultiRowKernelDtype<__nv_bfloat16, 2048>( - __nv_bfloat16 const*, int const*, int const*, __nv_bfloat16*, int*, int, int, int, int, int, int); + __nv_bfloat16 const*, int const*, int const*, __nv_bfloat16*, int*, int, int, int, int, int); template __global__ void heuristicTopKMultiRowKernelDtype<__half, 512>( - __half const*, int const*, int const*, __half*, int*, int, int, int, int, int, int); + __half const*, int const*, int const*, __half*, int*, int, int, int, int, int); template __global__ void heuristicTopKMultiRowKernelDtype<__half, 1024>( - __half const*, int const*, int const*, __half*, int*, int, int, int, int, int, int); + __half const*, int const*, int const*, __half*, int*, int, int, int, int, int); template __global__ void heuristicTopKMultiRowKernelDtype<__half, 2048>( - __half const*, int const*, int const*, __half*, int*, int, int, int, int, int, int); + __half const*, int const*, int const*, __half*, int*, int, int, int, int, int); template __global__ void heuristicTopKMultiRowKernel<512>( - float const*, int const*, int const*, float*, int*, int, int, int, int, int, int); + float const*, int const*, int const*, float*, int*, int, int, int, int, int); template __global__ void heuristicTopKMultiRowKernel<1024>( - float const*, int const*, int const*, float*, int*, int, int, int, int, int, int); + float const*, int const*, int const*, float*, int*, int, int, int, int, int); template __global__ void heuristicTopKMultiRowKernel<2048>( - float const*, int const*, int const*, float*, int*, int, int, int, int, int, int); + float const*, int const*, int const*, float*, int*, int, int, int, int, int); // Dispatch on topK at runtime — each TopK-instantiation gets its own smem // size (driven by GvrParams::kC/kNumBins) and own kfn pointer @@ -202,7 +184,7 @@ template __global__ void heuristicTopKMultiRowKernel<2048>( template void launchHeuristicTopKDecodeImpl(InputT const* logits, int const* seqLens, int const* preIdx, int* outIndices, InputT* scratchValues, int stride0, int next_n, int topK, int preIdxStride, int preIdxCount, int numRows, - int compressRatio, cudaStream_t stream) + cudaStream_t stream) { TLLM_CHECK_WITH_INFO( topK == 512 || topK == 1024 || topK == 2048, "heuristicTopKDecode requires topK ∈ {512, 1024, 2048}"); @@ -242,7 +224,7 @@ void launchHeuristicTopKDecodeImpl(InputT const* logits, int const* seqLens, int config.attrs = attrs; cudaLaunchKernelEx(&config, kfn, logits, seqLens, preIdx, scratchValues, outIndices, stride0, next_n, topK, - preIdxStride, preIdxCount, compressRatio); + preIdxStride, preIdxCount); }; switch (topK) @@ -258,26 +240,26 @@ void launchHeuristicTopKDecodeImpl(InputT const* logits, int const* seqLens, int void launchHeuristicTopKDecode(float const* logits, int const* seqLens, int const* preIdx, int* outIndices, float* scratchValues, int stride0, int next_n, int topK, int preIdxStride, int preIdxCount, int numRows, - int compressRatio, cudaStream_t stream) + cudaStream_t stream) { launchHeuristicTopKDecodeImpl(logits, seqLens, preIdx, outIndices, scratchValues, stride0, next_n, topK, - preIdxStride, preIdxCount, numRows, compressRatio, stream); + preIdxStride, preIdxCount, numRows, stream); } void launchHeuristicTopKDecode(__nv_bfloat16 const* logits, int const* seqLens, int const* preIdx, int* outIndices, __nv_bfloat16* scratchValues, int stride0, int next_n, int topK, int preIdxStride, int preIdxCount, int numRows, - int compressRatio, cudaStream_t stream) + cudaStream_t stream) { launchHeuristicTopKDecodeImpl<__nv_bfloat16>(logits, seqLens, preIdx, outIndices, scratchValues, stride0, next_n, - topK, preIdxStride, preIdxCount, numRows, compressRatio, stream); + topK, preIdxStride, preIdxCount, numRows, stream); } void launchHeuristicTopKDecode(__half const* logits, int const* seqLens, int const* preIdx, int* outIndices, __half* scratchValues, int stride0, int next_n, int topK, int preIdxStride, int preIdxCount, int numRows, - int compressRatio, cudaStream_t stream) + cudaStream_t stream) { launchHeuristicTopKDecodeImpl<__half>(logits, seqLens, preIdx, outIndices, scratchValues, stride0, next_n, topK, - preIdxStride, preIdxCount, numRows, compressRatio, stream); + preIdxStride, preIdxCount, numRows, stream); } } // namespace kernels diff --git a/cpp/tensorrt_llm/kernels/heuristicTopKDecode.h b/cpp/tensorrt_llm/kernels/heuristicTopKDecode.h index 0d2330f76545..5e38c64cba06 100644 --- a/cpp/tensorrt_llm/kernels/heuristicTopKDecode.h +++ b/cpp/tensorrt_llm/kernels/heuristicTopKDecode.h @@ -1,5 +1,5 @@ /* - * Copyright (c) 2019-2026, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2019-2025, NVIDIA CORPORATION. All rights reserved. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -33,28 +33,21 @@ inline constexpr int kHeuristicSize = 2048; /// Launch heuristic TopK decode kernel — fp32 input. /// @param scratchValues Caller-owned buffer of size [numRows * topK] floats. /// Required for CUDA Graph compatibility — must have a stable device address. -/// @param compressRatio KV compression ratio (1 = V3.2 indexer; 4 = V4 indexer -/// whose logits/preIdx live in compressed-token-index space). For -/// compressRatio != 1, preIdxOffset is forced to 0 (append-at-end in -/// compressed space → prev-step indices remain valid as-is); the -/// existing (rowIdx % next_n)+1 shift is used only when compressRatio==1. void launchHeuristicTopKDecode(float const* logits, int const* seqLens, int const* preIdx, int* outIndices, float* scratchValues, int stride0, int next_n, int topK, int preIdxStride, int preIdxCount, int numRows, - int compressRatio, cudaStream_t stream); + cudaStream_t stream); /// Launch heuristic TopK decode kernel — bf16 input. /// scratchValues is [numRows * topK] of bf16 (matches input dtype). -/// @param compressRatio See fp32 overload. void launchHeuristicTopKDecode(__nv_bfloat16 const* logits, int const* seqLens, int const* preIdx, int* outIndices, __nv_bfloat16* scratchValues, int stride0, int next_n, int topK, int preIdxStride, int preIdxCount, int numRows, - int compressRatio, cudaStream_t stream); + cudaStream_t stream); /// Launch heuristic TopK decode kernel — fp16 input. /// scratchValues is [numRows * topK] of fp16 (matches input dtype). -/// @param compressRatio See fp32 overload. void launchHeuristicTopKDecode(__half const* logits, int const* seqLens, int const* preIdx, int* outIndices, __half* scratchValues, int stride0, int next_n, int topK, int preIdxStride, int preIdxCount, int numRows, - int compressRatio, cudaStream_t stream); + cudaStream_t stream); } // namespace kernels diff --git a/cpp/tensorrt_llm/kernels/heuristic_topk.cuh b/cpp/tensorrt_llm/kernels/heuristic_topk.cuh index f21a4d7cf2c6..ed1364096e04 100644 --- a/cpp/tensorrt_llm/kernels/heuristic_topk.cuh +++ b/cpp/tensorrt_llm/kernels/heuristic_topk.cuh @@ -224,12 +224,7 @@ struct GvrParams; // primary undefined → compile-time error for bad combos template <> struct GvrParams { - // kFTarget=kK aligns the secant's soft steering target with the band's - // lower edge; eliminates the upper-clamp saturation on tight-σ + high-A2 - // layers (L36/L42/L28). Cross-prompt simulator validation on swe-bench - // 32k/64k/100k showed 2.19× / 1.77× / 1.51× total P2-iter reduction with - // zero cap-hits, zero per-layer regression vs the prior kFTarget=384. - static constexpr int kFTarget = 512; + static constexpr int kFTarget = 384; static constexpr int kC = 5120; static constexpr int kNumBins = 1024; }; @@ -237,14 +232,7 @@ struct GvrParams template <> struct GvrParams { - // kFTarget = kK (see GvrParams rationale). Q9k Pro 32k - // K=1024 native sweep (M=K=1024) finds kFT=1024 reduces sum_mean - // P2 iters from 35.33 (kFT=2560) → 30.21 (1.17× speedup) with zero - // per-layer regression and zero cap-hits. The prior kFT=2560 setting - // was tuned with M=512 K=1024 (sparse_attention_config default - // index_topk=512 inherited Flash's K), which does not represent - // production Pro behavior (production: M = K). - static constexpr int kFTarget = 1024; + static constexpr int kFTarget = 2560; static constexpr int kC = 5120; static constexpr int kNumBins = 1024; }; @@ -263,8 +251,7 @@ struct GvrParams template <> struct GvrParams<__nv_bfloat16, 512> { - // kFTarget aligned to kK — see GvrParams rationale. - static constexpr int kFTarget = 512; + static constexpr int kFTarget = 384; static constexpr int kC = 5120; static constexpr int kNumBins = 512; }; @@ -272,8 +259,7 @@ struct GvrParams<__nv_bfloat16, 512> template <> struct GvrParams<__nv_bfloat16, 1024> { - // kFTarget = kK — see GvrParams rationale. - static constexpr int kFTarget = 1024; + static constexpr int kFTarget = 2560; static constexpr int kC = 5120; static constexpr int kNumBins = 512; }; @@ -289,8 +275,7 @@ struct GvrParams<__nv_bfloat16, 2048> template <> struct GvrParams<__half, 512> { - // kFTarget aligned to kK — see GvrParams rationale. - static constexpr int kFTarget = 512; + static constexpr int kFTarget = 384; static constexpr int kC = 5120; static constexpr int kNumBins = 512; }; @@ -298,8 +283,7 @@ struct GvrParams<__half, 512> template <> struct GvrParams<__half, 1024> { - // kFTarget = kK — see GvrParams rationale. - static constexpr int kFTarget = 1024; + static constexpr int kFTarget = 2560; static constexpr int kC = 5120; static constexpr int kNumBins = 1024; }; diff --git a/cpp/tensorrt_llm/kernels/indexerTopK.cu b/cpp/tensorrt_llm/kernels/indexerTopK.cu index d63f203a048c..ad1e728b3298 100644 --- a/cpp/tensorrt_llm/kernels/indexerTopK.cu +++ b/cpp/tensorrt_llm/kernels/indexerTopK.cu @@ -39,13 +39,6 @@ TRTLLM_NAMESPACE_BEGIN namespace kernels { -// Radix histogram bin count, used by topKPerRowJob's per-step distribution -// pass. 2048 = 2^11 maps directly to the fp16 fast path's 11-bit bin index -// (sign + exponent + top mantissa bits, mask 0x7FF). All file-scope dispatch -// thresholds below derive from this so that a future change to the histogram -// width does not require re-tuning the heuristics. -constexpr int kNumBins = 2048; - namespace { @@ -176,20 +169,6 @@ __device__ bool processHistogramStep(int const* indices, InputT const* logits, i int& thresholdBinIdx, SmemOutputType& smemOutput, int* smemThresholdBinIdx, int* smemFinalDstIdx, int* smemFinalBinSize, int* smemFoundTopKValues, SmemFinalType& smemFinal, int stride1, int rowStart, int topK) { - // Step 0 is the fp16 fast path; if it could not resolve top-K (threshold bin - // exceeded kNumFinalItems) we restart the fp32 radix from scratch in steps 1-3. - // Discard any candidates step 0 wrote into smemOutput so step 1 doesn't - // double-count valid entries that fall under both fp16 and fp32 thresholds - // (this is what produced 2x duplicated indices when many -FLT_MAX padding - // entries dominated the threshold bin in the multi-block merge path). - if constexpr (step == 1) - { - if (threadIdx.x == 0) - { - smemFoundTopKValues[0] = 0; - smemFinalDstIdx[0] = 0; - } - } // Clear the histogram. #pragma unroll for (int idx = threadIdx.x; idx < kNumBins; idx += kNumThreadsPerBlock) @@ -292,13 +271,18 @@ __device__ bool processHistogramStep(int const* indices, InputT const* logits, i // The threshold bin. thresholdBinIdx = smemThresholdBinIdx[0]; + // Skip auto-promote at step 0 when we'll continue: half-precision bins + // don't align with step 2's full-precision bit-pattern filter, so a + // step-0 promote would be double-counted at step 2. + bool const step0WillContinue = (step == 0) && (smemFinalBinSize[0] > kNumFinalItems); + auto processBins = [&](InputT logitIn, int idx) { float const logit = static_cast(logitIn); if (isPartialMatch(logit, logitPattern)) { uint32_t binIdx = extractBinIdx(logit); - if (binIdx < thresholdBinIdx) + if (binIdx < thresholdBinIdx && !step0WillContinue) { // The element is part of the top-k selection int dstIdx = atomicAdd(&smemFoundTopKValues[0], 1); @@ -384,28 +368,21 @@ __device__ bool processHistogramStep(int const* indices, InputT const* logits, i return smemFinalBinSize[0] > kNumFinalItems; } -// Follows half - 11 - 11 - 10 bit iterations -template -static __device__ void topKPerRowJob(int const* indices, InputT const* logits, int rowStart, int rowEnd, - int* outIndices, float* outLogits, int stride1, int topK) +// Smem holder for topKPerRowJob's final-sort. Always reserves +// BlockRadixSort::TempStorage so the sort algorithm can be picked at runtime +// (the union's size is dominated by FinalItems = 16 KB at our shapes). +template +struct TopKSmem { - // The number of slots for the final pass. static constexpr int kNumFinalItems = 2048; - // The number of elements per thread for the final sort. static constexpr int kNumFinalItemsPerThread = kNumFinalItems / kNumThreadsPerBlock; - // The class to sort the elements during the final pass. using FinalSort = cub::BlockRadixSort; - using FinalSortTempStorage = std::conditional_t; - // The class to compute the inclusive prefix-sum over the histogram. + using FinalSortTempStorage = typename FinalSort::TempStorage; using Scan = cub::BlockScan; - // The structure to store the final items (for the final pass). struct FinalItems { - // Shared memory to store the indices for the final pass. int indices[kNumFinalItems]; - // Shared memory to store the logits for the final pass. float logits[kNumFinalItems]; }; @@ -415,34 +392,47 @@ static __device__ void topKPerRowJob(int const* indices, InputT const* logits, i int data[kNumBins]; }; - // Shared memory to compute the block sort. - __shared__ union + union Final { FinalItems items; FinalSortTempStorage finalSort; Histogram histo; - } smemFinal; + }; + + Final smemFinal; + int smemThresholdBinIdx[1]; + int smemFinalDstIdx[1]; + int smemFinalBinSize[1]; + int smemFoundTopKValues[1]; +}; + +// Follows half - 11 - 11 - 10 bit iterations +template +static __device__ void topKPerRowJob(int const* indices, InputT const* logits, int rowStart, int rowEnd, + int* outIndices, float* outLogits, int stride1, int topK, TopKSmem& smem) +{ + static constexpr int kNumFinalItems = TopKSmem::kNumFinalItems; + static constexpr int kNumFinalItemsPerThread = TopKSmem::kNumFinalItemsPerThread; + using FinalSort = typename TopKSmem::FinalSort; + + auto& smemFinal = smem.smemFinal; + int* smemThresholdBinIdx = smem.smemThresholdBinIdx; + int* smemFinalDstIdx = smem.smemFinalDstIdx; + int* smemFinalBinSize = smem.smemFinalBinSize; + int* smemFoundTopKValues = smem.smemFoundTopKValues; // Shared memory to store the selected indices. // If we are processing using multiple blocks, we need to store the logits and // indices. extern __shared__ int32_t smemOutput[]; - // Shared memory to store the threshold bin. - __shared__ int smemThresholdBinIdx[1]; - // Shared memory counter to register the candidates for the final phase. - __shared__ int smemFinalDstIdx[1]; - // Shared memory to determine if the threshold bin fits in the final items. - __shared__ int smemFinalBinSize[1]; - // Shared memory to keep track of the top-k values found so far by the - // previous iterations - __shared__ int smemFoundTopKValues[1]; - // The length of the row. int rowLen = rowEnd - rowStart; // Shortcut if the length of the row is smaller than Top-K. Indices are not - // sorted by their corresponding logit. + // sorted by their corresponding logit. Unreachable when mergeBlocks=true: + // both merge callers pass rowLen = numBlocksPerRow * topK > topK. if (rowLen <= topK) { for (int rowIt = threadIdx.x; rowIt < rowLen; rowIt += kNumThreadsPerBlock) @@ -512,10 +502,12 @@ static __device__ void topKPerRowJob(int const* indices, InputT const* logits, i if (!continueToNextStep) { - // The histogram did not proceed to the final 10 bits, therefore we need to - // sort the final items The logits of the elements to be sorted in the final - // pass. - if constexpr (useRadixSort) + // Sort the threshold-bin candidates. Insertion sort wins below ~512 + // items (O(n^2/T) with no fixed cost); BlockRadixSort wins above + // (constant cost padded to kNumFinalItems = 2048). + constexpr int kInsertionSortBranchThreshold = 512; + int const finalCount = smemFinalDstIdx[0]; + if (finalCount > kInsertionSortBranchThreshold) { // Sorting with radix sort float finalLogits[kNumFinalItemsPerThread]; @@ -526,12 +518,6 @@ static __device__ void topKPerRowJob(int const* indices, InputT const* logits, i for (int ii = 0; ii < kNumFinalItemsPerThread; ++ii) { finalLogits[ii] = -FLT_MAX; - // Indices must be paired with the -FLT_MAX sentinel logit so unused - // slots survive SortDescendingBlockedToStriped and the post-sort copy - // as -1 rather than garbage. Without this, when the threshold bin is - // dominated by -FLT_MAX padding (multi-block merge path) the trailing - // top-K slots receive arbitrary register contents. - finalIndices[ii] = -1; } // Read the elements from SMEM. @@ -539,7 +525,7 @@ static __device__ void topKPerRowJob(int const* indices, InputT const* logits, i for (int ii = 0; ii < kNumFinalItemsPerThread; ++ii) { int srcIdx = ii * kNumThreadsPerBlock + threadIdx.x; - if (srcIdx < smemFinalDstIdx[0]) + if (srcIdx < finalCount) { finalLogits[ii] = smemFinal.items.logits[srcIdx]; finalIndices[ii] = smemFinal.items.indices[srcIdx]; @@ -572,13 +558,13 @@ static __device__ void topKPerRowJob(int const* indices, InputT const* logits, i } else { - // Sorting with insertion sort + // Sorting with insertion sort. auto baseIdx = smemFoundTopKValues[0]; - for (int i = threadIdx.x; i < smemFinalDstIdx[0]; i += kNumThreadsPerBlock) + for (int i = threadIdx.x; i < finalCount; i += kNumThreadsPerBlock) { int outIndex = 0; auto logit = smemFinal.items.logits[i]; - for (int j = 0; j < smemFinalDstIdx[0]; j++) + for (int j = 0; j < finalCount; j++) { auto otherLogit = smemFinal.items.logits[j]; if (logit < otherLogit || (logit == otherLogit && i < j)) @@ -586,7 +572,6 @@ static __device__ void topKPerRowJob(int const* indices, InputT const* logits, i outIndex++; } } - // Store if outIndex is in bounds if (outIndex + baseIdx < topK) { smemOutput[outIndex + baseIdx] = smemFinal.items.indices[i]; @@ -608,6 +593,10 @@ static __device__ void topKPerRowJob(int const* indices, InputT const* logits, i outIndices[i] = smemOutput[i]; outLogits[i] = reinterpret_cast(smemOutput + topK)[i]; } + else if constexpr (mergeBlocks) + { + outIndices[i] = smemOutput[i]; + } else { if (stride1 == 1) @@ -624,7 +613,7 @@ static __device__ void topKPerRowJob(int const* indices, InputT const* logits, i } } // namespace -template +template static __global__ __launch_bounds__(kNumThreadsPerBlock) void topKPerRowPrefill(InputT const* logits, int const* rowStarts, int const* rowEnds, int* outIndices, int stride0, int stride1, int const topK, int const offsetIndex) @@ -632,6 +621,9 @@ static __global__ __launch_bounds__(kNumThreadsPerBlock) void topKPerRowPrefill( #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) cudaGridDependencySynchronize(); #endif + // The number of bins in the histogram. + static constexpr int kNumBins = 2048; + // The row computed by this block. int rowIdx = blockIdx.x + offsetIndex; @@ -643,30 +635,32 @@ static __global__ __launch_bounds__(kNumThreadsPerBlock) void topKPerRowPrefill( outIndices += static_cast(rowIdx) * topK; logits += static_cast(rowIdx) * stride0; - topKPerRowJob( - nullptr, logits, rowStart, rowEnd, outIndices, nullptr, stride1, topK); + __shared__ TopKSmem smem; + topKPerRowJob( + nullptr, logits, rowStart, rowEnd, outIndices, nullptr, stride1, topK, smem); #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) cudaTriggerProgrammaticLaunchCompletion(); #endif } -template +template static __global__ __launch_bounds__(kNumThreadsPerBlock) void topKPerRowDecode(InputT const* logits, int const* seqLens, - int* outIndices, int stride0, int stride1, int const topK, int next_n, int compressRatio, - float* outLogits = nullptr, int const numBlocksToMerge = 0, int const* indices = nullptr) + int* outIndices, int stride0, int stride1, int const topK, int next_n, float* outLogits = nullptr, + int const numBlocksToMerge = 0, int const* indices = nullptr) { #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) cudaGridDependencySynchronize(); #endif + // The number of bins in the histogram. + static constexpr int kNumBins = 2048; + // The row computed by this block. int rowIdx = blockIdx.x; // The range of logits within the row. int rowStart = 0; int seq_len = seqLens[rowIdx / next_n]; - int actual_kv_len = seq_len - next_n + (rowIdx % next_n) + 1; - int rowEnd = actual_kv_len / compressRatio; + int rowEnd = seq_len - next_n + (rowIdx % next_n) + 1; // Local pointers to this block if constexpr (!multipleBlocksPerRow && !mergeBlocks) @@ -689,8 +683,9 @@ static __global__ __launch_bounds__(kNumThreadsPerBlock) void topKPerRowDecode(I } logits += static_cast(rowIdx) * stride0; - topKPerRowJob( - indices, logits, rowStart, rowEnd, outIndices, outLogits, stride1, topK); + __shared__ TopKSmem smem; + topKPerRowJob( + indices, logits, rowStart, rowEnd, outIndices, outLogits, stride1, topK, smem); #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) cudaTriggerProgrammaticLaunchCompletion(); #endif @@ -699,29 +694,415 @@ static __global__ __launch_bounds__(kNumThreadsPerBlock) void topKPerRowDecode(I namespace { -// Insertion vs radix crossover for the single-block path. Radix always pays a -// histogram-clear + per-step scan over kNumBins bins (~4 refinement passes); -// insertion only maintains a topK SMEM array. The crossover is where insertion's -// O(numColumns * lg topK) catches up to radix's O(numColumns + kNumBins) per -// pass — measured empirically around 6 histograms of work. -constexpr int kSortingAlgorithmThreshold = 6 * kNumBins; -// Force the multi-block split-and-merge path above this column count: per-block -// work would otherwise dwarf the merge-pass cost. Callers may override via the -// splitWorkThreshold parameter. Tuned empirically; below this width the merge -// launch overhead isn't worth saving the per-block radix cost. -constexpr int kDefaultSplitWorkThreshold = 200 * 1000; -// Cap blocks-per-row in the multi-block path. Bounds the aux-buffer size -// (numRows * blocksPerRow * topK * sizeof(int32)) and the merge-step input width -// so the second-pass merge kernel stays SMEM-resident. -constexpr int kMaxBlocksPerRowDecode = 10; -// Each sub-block must amortize the radix histogram overhead, i.e. cover at least -// one full histogram pass worth of columns. -constexpr int kDecodeMinColsPerSubBlock = kNumBins; - -// Scheme X bound calculator — shared between fp32 and bf16/fp16 dispatchers. -// Caches hardware attrs (SM count, L2 capacity) and the small-N threshold -// once per process via std::call_once. Per-call cost is just two reads -// from cached static variables plus a small arithmetic block, no syscalls. +// Multi-pass radix: 4 launches (half-precision top-11 bits, then float bits +// [21..32), [10..21), [0..10)). Per-row state lives in DRAM scratch; candidate +// data uses two ping-pong buffers. Pass 3's last block emits the final top-K. +static constexpr int kRadixBins = 2048; + +// Per-row scratch state. The last block of each pass (picked via +// `finishedBlocks`) scans the global histogram and writes the next pass's +// threshold. +struct alignas(64) RadixState +{ + int candCount; // candidates entering current pass + int outIdx; // running outIndices write position + int kRemaining; // topK - outIdx + int filterCnt; // running candidate-buf write position + int thresholdBin; // prior pass's threshold; overwritten this pass's last block + int finishedBlocks; // last-block atomic, reset between passes + int thresholdLess; // count of bins below thresholdBin; pass-3 emit routes + // ties (bin == threshold) into a disjoint slot range. + int padding[1]; +}; + +// Common last-block trailer: prefix-scan the merged global histogram, locate +// the bin where the running count crosses kRemaining, stash to state for the +// next pass. step < 3 also resets the per-row global histogram. In step 1 +// the trailer writes st.candCount = full row length for pass 2 to consume +// (pass 1 reads the length inline from seqLens, so no init kernel is needed). +template +__device__ __forceinline__ void radixLastBlockTrailer(int* gHist, RadixState& st, int topK, int rowFullLen) +{ + using Scan = cub::BlockScan; + __shared__ typename Scan::TempStorage scanStorage; + __shared__ int s_thresholdBin; + __shared__ int s_runningBefore; + __shared__ int s_thresholdCount; + + if (threadIdx.x == 0) + { + s_thresholdBin = -1; + s_runningBefore = 0; + s_thresholdCount = 0; + } + __syncthreads(); + + // kRemaining for THIS pass: + // step 1: topK (no auto-promotes have happened yet). + // step 2/3: topK - outIdx (where outIdx was atomic-incremented during + // the filter loop). The histogram for the next pass's + // threshold pick must target this fresh value, not the + // stale state.kRemaining left over from the previous pass. + int const kRem = (step == 1) ? topK : (topK - st.outIdx); + if (kRem <= 0) + { + // All top-k slots already filled by auto-promotes in this pass — + // no need for a next pass to emit anything. + if (threadIdx.x == 0) + { + st.thresholdBin = -1; + st.candCount = 0; + st.kRemaining = 0; + st.finishedBlocks = 0; + st.filterCnt = 0; + (void) rowFullLen; + } + if constexpr (step < 3) + { + __syncthreads(); + for (int i = threadIdx.x; i < kRadixBins; i += kThreads) + gHist[i] = 0; + } + return; + } + constexpr int kRoundsPerScan = kRadixBins / kThreads; + int running = 0; + for (int r = 0; r < kRoundsPerScan; ++r) + { + int bin = r * kThreads + threadIdx.x; + int c = gHist[bin]; + int prefix, total; + Scan(scanStorage).ExclusiveSum(c, prefix, total); + prefix += running; + int next = prefix + c; + if (prefix < kRem && next >= kRem && s_thresholdBin == -1) + { + atomicCAS(&s_thresholdBin, -1, bin); + if (s_thresholdBin == bin) + { + s_runningBefore = prefix; + s_thresholdCount = c; + } + } + running += total; + __syncthreads(); + if (s_thresholdBin != -1) + break; + } + __syncthreads(); + + if (threadIdx.x == 0) + { + // If the cumsum over the whole histogram never reached kRem the row + // has fewer items than we still need (e.g. decode rows shorter than + // topK). Set thresholdBin to a sentinel above any valid bin so the + // next pass's `bin < thresholdBin` test accepts every surviving + // candidate as auto-promote. + st.thresholdBin = (s_thresholdBin == -1) ? kRadixBins : s_thresholdBin; + // s_runningBefore is the count of histogram items in bins < threshold. + // In the sentinel case it stays at its init 0; that's fine because + // pass-3's inline final-emit only uses thresholdLess to position the + // ties (bin == threshold) write base, and the sentinel branch has + // no items in the threshold bin (the sentinel is above all valid bins). + st.thresholdLess = s_runningBefore; + st.finishedBlocks = 0; + if constexpr (step == 1) + { + // Pass 2 also scans the full row from `logits`, so it reads + // st.candCount = full row length. Pass 1 did not write to it + // and the state struct started at zero (cudaMemsetAsync). + st.candCount = rowFullLen; + (void) s_thresholdCount; + (void) kRem; + } + else + { + st.candCount = st.filterCnt; + int newKRem = topK - st.outIdx; + if (newKRem < 0) + newKRem = 0; + st.kRemaining = newKRem; + st.filterCnt = 0; + } + } + if constexpr (step < 3) + { + __syncthreads(); + for (int i = threadIdx.x; i < kRadixBins; i += kThreads) + gHist[i] = 0; + } +} + +// Fused histogram + filter pass kernel. +// +// step == 1: pass 1, no filter, just histogram top-11 bits of the row. +// step == 2: pass 2 reads `logits` (pass 1 didn't write a candidate buffer), +// for each item: +// bin1 < thresholdBin1 → write to outIndices (auto-promote) +// bin1 == thresholdBin1 → append to candBufOut, count its bin2 +// in the histogram for pass 2 +// else → drop +// step == 3: same as step 2 but reads `candBufIn` (pass 2's output) and +// uses extractBinIdx<2> for the prior-bits check, extractBinIdx<3> +// for the histogram. +// +// Last block of every pass runs `radixLastBlockTrailer` to compute the +// next pass's threshold and reset cross-pass state. +template +static __global__ __launch_bounds__(kThreads) void radixPassKernel(InputT const* logits, int const* seqLens, + int* outIndices, int const* candBufIn, int* candBufOut, int* histograms, RadixState* state, int stride0, int next_n, + int topK) +{ + int rowIdx = blockIdx.y; + int blockInRow = blockIdx.x; + int blocksPerRow = gridDim.x; + + RadixState& st = state[rowIdx]; + int* gHist = histograms + static_cast(rowIdx) * kRadixBins; + + __shared__ int sHist[kRadixBins]; + for (int i = threadIdx.x; i < kRadixBins; i += kThreads) + sHist[i] = 0; + __syncthreads(); + + if constexpr (step == 1) + { + // Read seqLen inline so pass 1 does not depend on st.candCount being + // pre-initialised by a separate init kernel; the cudaMemsetAsync that + // zeroes state+histograms together is enough. Pass-1 trailer below + // writes st.candCount = seqLens[rowIdx] for pass 2 to consume. + int const rowEnd = seqLens[rowIdx / next_n] - next_n + (rowIdx % next_n) + 1; + InputT const* in = logits + static_cast(rowIdx) * stride0; + size_t threadRank = static_cast(blockInRow) * kThreads + threadIdx.x; + size_t numThreads = static_cast(blocksPerRow) * kThreads; + auto f = [&](InputT vIn, size_t /*idx*/) + { + float const v = static_cast(vIn); + uint32_t bin = extractBinIdx(v); + atomicAdd(&sHist[bin], 1); + }; + vectorized_process(threadRank, numThreads, in, static_cast(rowEnd), f); + } + else if constexpr (step == 2) + { + int const rowEnd = st.candCount; + InputT const* in = logits + static_cast(rowIdx) * stride0; + int* outIdxArr = outIndices + static_cast(rowIdx) * topK; + int* candArr = candBufOut + static_cast(rowIdx) * stride0; + int const prevThresh = st.thresholdBin; + size_t threadRank = static_cast(blockInRow) * kThreads + threadIdx.x; + size_t numThreads = static_cast(blocksPerRow) * kThreads; + auto f = [&](InputT vIn, size_t i) + { + float const v = static_cast(vIn); + int bin1 = static_cast(extractBinIdx<1>(v)); + if (bin1 < prevThresh) + { + int pos = atomicAdd(&st.outIdx, 1); + if (pos < topK) + outIdxArr[pos] = static_cast(i); + } + else if (bin1 == prevThresh) + { + int pos = atomicAdd(&st.filterCnt, 1); + candArr[pos] = static_cast(i); + uint32_t bin2 = extractBinIdx(v); + atomicAdd(&sHist[bin2], 1); + } + }; + vectorized_process(threadRank, numThreads, in, static_cast(rowEnd), f); + } + else // step == 3 + { + int const candCnt = st.candCount; + int const* candArrIn = candBufIn + static_cast(rowIdx) * stride0; + InputT const* in = logits + static_cast(rowIdx) * stride0; + int* outIdxArr = outIndices + static_cast(rowIdx) * topK; + int* candArrOut = candBufOut + static_cast(rowIdx) * stride0; + int const prevThresh = st.thresholdBin; + for (int i = blockInRow * kThreads + threadIdx.x; i < candCnt; i += blocksPerRow * kThreads) + { + int srcIdx = candArrIn[i]; + float v = static_cast(in[srcIdx]); + int bin2 = static_cast(extractBinIdx<2>(v)); + if (bin2 < prevThresh) + { + int pos = atomicAdd(&st.outIdx, 1); + if (pos < topK) + outIdxArr[pos] = srcIdx; + } + else if (bin2 == prevThresh) + { + int pos = atomicAdd(&st.filterCnt, 1); + candArrOut[pos] = srcIdx; + uint32_t bin3 = extractBinIdx(v); + atomicAdd(&sHist[bin3], 1); + } + } + } + __syncthreads(); + + for (int i = threadIdx.x; i < kRadixBins; i += kThreads) + { + int c = sHist[i]; + if (c) + atomicAdd(&gHist[i], c); + } + + __threadfence(); + __shared__ int isLast; + if (threadIdx.x == 0) + { + int prev = atomicAdd(&st.finishedBlocks, 1); + isLast = (prev == blocksPerRow - 1) ? 1 : 0; + } + __syncthreads(); + if (!isLast) + return; + + int const rowFullLen = (step == 1) ? (seqLens[rowIdx / next_n] - next_n + (rowIdx % next_n) + 1) : 0; + radixLastBlockTrailer(gHist, st, topK, rowFullLen); + + if constexpr (step == 3) + { + // Final emit, folded into the last block of pass 3. Scan candBufOut + // (top 22 bits == thresholdBin2) and route items into outIndices by + // bin3 vs thresholdBin3. Two-counter scheme so threshold-bin ties + // don't race definite top-k items on the same atomic: + // bin3 < thresh3 → slots [ltBase, ltBase + prefix3) + // bin3 == thresh3 → slots [ltBase + prefix3, topK) + __syncthreads(); + int const filterCnt = st.candCount; + int const thresh3 = st.thresholdBin; + int const prefix3 = st.thresholdLess; + int const* candArr = candBufOut + static_cast(rowIdx) * stride0; + InputT const* in = logits + static_cast(rowIdx) * stride0; + int* outIdxArr = outIndices + static_cast(rowIdx) * topK; + int const ltBase = st.outIdx; // already at outBase here + int const eqBase = ltBase + prefix3; + int const eqCap = topK - eqBase; // ≥ 0 by trailer invariant + __shared__ int sEqEmitted; + if (threadIdx.x == 0) + sEqEmitted = 0; + __syncthreads(); + for (int i = threadIdx.x; i < filterCnt; i += kThreads) + { + int srcIdx = candArr[i]; + float v = static_cast(in[srcIdx]); + int bin3 = static_cast(extractBinIdx<3>(v)); + if (bin3 < thresh3) + { + // atomicAdd on st.outIdx is safe: by construction exactly + // prefix3 items fall in this branch, so pos stays in + // [ltBase, eqBase) which is strictly inside [0, topK). + int pos = atomicAdd(&st.outIdx, 1); + outIdxArr[pos] = srcIdx; + } + else if (bin3 == thresh3) + { + int pos = atomicAdd(&sEqEmitted, 1); + if (pos < eqCap) + outIdxArr[eqBase + pos] = srcIdx; + } + } + __syncthreads(); + if (threadIdx.x == 0) + { + int eq = sEqEmitted < eqCap ? sEqEmitted : eqCap; + int filled = eqBase + eq; + if (filled > topK) + filled = topK; + for (int i = filled; i < topK; ++i) + outIdxArr[i] = -1; + } + // Reset the per-row global histogram and st.outIdx so the next call + // sees a clean state without a per-call cudaMemsetAsync. Caller must + // zero-initialize the scratch buffer before the first call. + __syncthreads(); + for (int i = threadIdx.x; i < kRadixBins; i += kThreads) + gHist[i] = 0; + if (threadIdx.x == 0) + st.outIdx = 0; + } +} + +// Scratch layout (uint8 buffer, 64-byte aligned regions): +// RadixState[numRows] +// int histograms[numRows * kRadixBins] (zeroed on first call by +// torch::zeros allocator; pass-3 +// trailer zeroes for subsequent +// calls) +// int candBuf1[numRows * stride0] (pass 2 → pass 3 input) +// int candBuf2[numRows * stride0] (pass 3 → fused final filter) +static size_t radixScratchBytes(int numRows, int numColumns) +{ + auto roundUp = [](size_t x) { return (x + 63) & ~size_t(63); }; + size_t s = 0; + s += roundUp(sizeof(RadixState) * numRows); + s += roundUp(sizeof(int) * static_cast(numRows) * kRadixBins); + s += roundUp(sizeof(int) * static_cast(numRows) * numColumns); + s += roundUp(sizeof(int) * static_cast(numRows) * numColumns); + return s; +} + +template +static void launchMultiPassRadix(void* scratch, InputT const* logits, int const* seqLens, int* outIndices, int numRows, + int numColumns, int topK, int stride0, int next_n, cudaLaunchAttribute const* attrs, cudaStream_t stream) +{ + auto roundUp = [](size_t x) { return (x + 63) & ~size_t(63); }; + char* base = static_cast(scratch); + RadixState* state = reinterpret_cast(base); + base += roundUp(sizeof(RadixState) * numRows); + int* histograms = reinterpret_cast(base); + base += roundUp(sizeof(int) * static_cast(numRows) * kRadixBins); + int* candBuf1 = reinterpret_cast(base); + base += roundUp(sizeof(int) * static_cast(numRows) * numColumns); + int* candBuf2 = reinterpret_cast(base); + + int sm_cnt = 132; + { + int dev = 0; + cudaGetDevice(&dev); + cudaDeviceGetAttribute(&sm_cnt, cudaDevAttrMultiProcessorCount, dev); + } + // Block fan-out heuristic: target ~4 active blocks/SM (one wave at the + // achievable occupancy of radixPassKernel<512, 1>), with a per-block work + // floor of 2048 items (4 items/thread at 512-wide). + int targetTotalBlocks = sm_cnt * 4; + int numBlocksPerRow = (targetTotalBlocks + numRows - 1) / numRows; + int maxByCols = numColumns / 2048; + if (numBlocksPerRow > maxByCols) + numBlocksPerRow = maxByCols; + if (numBlocksPerRow < 1) + numBlocksPerRow = 1; + + constexpr int kPassThreads = 512; + + auto launchPass = [&](void const* kernel, int const* candIn, int* candOut) + { + cudaLaunchConfig_t cfg{}; + cfg.gridDim = dim3(numBlocksPerRow, numRows); + cfg.blockDim = kPassThreads; + cfg.dynamicSmemBytes = 0; + cfg.stream = stream; + cfg.numAttrs = 1; + cfg.attrs = const_cast(attrs); + void* args[] = {(void*) &logits, (void*) &seqLens, (void*) &outIndices, (void*) &candIn, (void*) &candOut, + (void*) &histograms, (void*) &state, (void*) &stride0, (void*) &next_n, (void*) &topK}; + cudaLaunchKernelExC(&cfg, kernel, args); + }; + + launchPass( + reinterpret_cast(&radixPassKernel), (int const*) nullptr, (int*) nullptr); + launchPass( + reinterpret_cast(&radixPassKernel), (int const*) nullptr, candBuf1); + // Pass 3 emits the final top-K inline in its last-block trailer (see + // radixPassKernel) instead of requiring a separate filter launch. + launchPass( + reinterpret_cast(&radixPassKernel), (int const*) candBuf1, candBuf2); +} + +// Architecture-derived GVR eligibility bounds (cached per-process). struct SchemeXBounds { int smCount; @@ -732,53 +1113,12 @@ struct SchemeXBounds int kSeqSmall; }; -// Uniform small-N lower bound for the Heuristic GVR path across all K. -// Aligns the GVR routing boundary with the Radix multi-CTA split-work -// threshold (maxByCols = N / kDecodeMinColsPerSubBlock(=2048) ≥ 2 at -// N ≥ 4096), so the dispatcher's algorithmic-handoff point is consistent: -// below 4096 the Radix path resolves to single-CTA insertion-sort and GVR -// is not attempted; at or above 4096 GVR may be considered. -// DSv4 swe-bench synth sweeps on B200/B300 (V3.2-Q19c protocol, May 2026): -// N=4K cells across K ∈ {512, 1024, 2048} all win — GVR R/H bf16 = 3.07× -// (K=512) / 2.57× (K=1024) / 1.34× (K=2048). -// N=2K cells across the same 9 (K × dtype) combos all show GVR R/H < 1 -// (0.55× – 0.84×), justifying 4K as the floor. -inline int kSeqSmallDefaultForK(int /*topK*/) -{ - return 4096; -} - -inline SchemeXBounds getSchemeXBounds(int numColumns, int bytesPerElem, int topK) +inline SchemeXBounds getSchemeXBounds(int numColumns, int bytesPerElem) { static std::once_flag sOnce; static int sSm = 0; static int sL2 = 0; - // ----------------------------------------------------------------------- - // Diagnostic / tuning escape-hatch env overrides. Both are OFF by default - // and the K-aware / hardware-derived defaults below are expected to be - // optimal for production. Use only for microbenchmarks, regression - // bisection, or workload-specific tuning where the defaults are clearly - // suboptimal. - // - // TRTLLM_HEURISTIC_NMIN (valid range [1024, 200000]) - // Overrides `kSeqSmall` (Heuristic small-N threshold) for ALL K. - // Lower risk: only shifts a perf threshold; the kernel still - // produces an exact top-K either way. Setting it too low routes - // more N → Heuristic and may be slower than the fallback for - // small N; correctness is preserved. - // - // TRTLLM_HEURISTIC_BSMAX (valid range [1, 65536]) - // Overrides `kBsLarge` (BS upper bound for Heuristic) past the - // hardware-derived min(kBsWave, kBsL2). Higher risk: bypasses - // L2/occupancy safety bounds, so heuristic may run in working-set - // ranges where it has not been tuned (L2 thrash, suboptimal grid - // configs). Primary use is indexer microbenchmarks that need a - // BS-scaling comparison against the Radix path on identical inputs. - // ----------------------------------------------------------------------- - // sNMinEnv > 0 iff TRTLLM_HEURISTIC_NMIN is set to a valid value. When set, - // it overrides the per-K default for ALL K. - static int sNMinEnv = 0; - static int sBsMax = 0; + static int sNMin = 0; std::call_once(sOnce, []() { @@ -786,17 +1126,16 @@ inline SchemeXBounds getSchemeXBounds(int numColumns, int bytesPerElem, int topK cudaGetDevice(&dev); cudaDeviceGetAttribute(&sSm, cudaDevAttrMultiProcessorCount, dev); cudaDeviceGetAttribute(&sL2, cudaDevAttrL2CacheSize, dev); + constexpr int kSeqSmallDefault = 12288; char const* env = std::getenv("TRTLLM_HEURISTIC_NMIN"); if (env != nullptr) { int const v = std::atoi(env); - sNMinEnv = (v >= 1024 && v <= 200000) ? v : 0; + sNMin = (v >= 1024 && v <= 200000) ? v : kSeqSmallDefault; } - char const* env_bsmax = std::getenv("TRTLLM_HEURISTIC_BSMAX"); - if (env_bsmax != nullptr) + else { - int const v = std::atoi(env_bsmax); - sBsMax = (v >= 1 && v <= 65536) ? v : 0; + sNMin = kSeqSmallDefault; } }); @@ -808,163 +1147,37 @@ inline SchemeXBounds getSchemeXBounds(int numColumns, int bytesPerElem, int topK ? static_cast(static_cast(sL2) * 9 / 10 / (static_cast(numColumns) * bytesPerElem)) : b.kBsWave; b.kBsLarge = std::min(b.kBsWave, b.kBsL2 > 0 ? b.kBsL2 : b.kBsWave); - if (sBsMax > 0) - { - // BSMAX env override bypasses the hardware-derived L2/occupancy bound - // (see the BSMAX section in the call_once block above for risk notes). - b.kBsLarge = sBsMax; - } - // NMIN env override (if set) wins over the per-K default for ALL K. - b.kSeqSmall = (sNMinEnv > 0) ? sNMinEnv : kSeqSmallDefaultForK(topK); + b.kSeqSmall = sNMin; return b; } -} // namespace - -int computeIndexerTopKDecodeBlocksPerRow(int numRows, int numColumns, int splitWorkThreshold) -{ - if (numRows <= 0) - { - return 1; - } - - int const forceSplitThreshold = splitWorkThreshold > 0 ? splitWorkThreshold : kDefaultSplitWorkThreshold; - - // Preserve original behavior for very long sequences. - if (numColumns >= forceSplitThreshold) - { - return kMaxBlocksPerRowDecode; - } - - // Query the actual SM count from the driver so the dispatch tracks the - // hardware rather than a baked-in target (H100=132, B200=148, …). - // topK=0: blocks-per-row computation is K-agnostic; kSeqSmall is uniform - // 4096 across K, so the topK arg is unused for the kSeqSmall lookup as - // well, and only smCount/kBsWave/kBsL2 are consumed here. - auto const bounds = getSchemeXBounds(numColumns, /*bytesPerElem=*/4, /*topK=*/0); - TLLM_CHECK_WITH_INFO(bounds.smCount > 0, "indexerTopK: failed to query device SM count"); - int const smCount = bounds.smCount; - int const maxByCols = std::max(1, numColumns / kDecodeMinColsPerSubBlock); - int const maxBp = std::min(maxByCols, kMaxBlocksPerRowDecode); - - int blocksPerRow; - if (numRows < smCount / 2) - { - // Sub-half-wave band: bp=2 by itself leaves SMs idle (numRows × 2 < smCount), - // so sweep bp ∈ [2, maxBp] for the choice that minimizes waves(bp) / bp, - // where waves(bp) = ceil(numRows * bp / smCount). The wave-quantization-aware - // sweep avoids spilling one extra block per row across a wave boundary, which - // is what a naive ceil(smCount / numRows) target would do. - int bestBp = 1; - int bestWaves = 1; // numRows < smCount/2 → bp=1 always fits in a single wave - for (int bp = 2; bp <= maxBp; ++bp) - { - int const totalBlocks = numRows * bp; - int const waves = (totalBlocks + smCount - 1) / smCount; - // waves / bp < bestWaves / bestBp ⇔ waves * bestBp < bestWaves * bp - if (waves * bestBp < bestWaves * bp) - { - bestWaves = waves; - bestBp = bp; - } - } - blocksPerRow = bestBp; - } - else - { - // numRows >= smCount/2: bp=2 saturates SMs with one wave (numRows*2 >= smCount) - // and stays on the multi-block split+merge path. Crucially this path uses - // a different kernel instantiation than bp=1, and only the bp=1 single-block - // radix kernel pays the wave-scheduling cliff when gridDim.x approaches - // smCount (measured on B200, cols=196608, topK=2048: BS=131=125us, - // BS=132=312us, BS=148=390us — the cliff disappears entirely with bp=2). - // bp=2 is also the cheapest split (smallest merge input); larger bp piles - // on merge-pass overhead without proportional gain on the shapes measured. - // Falls back to 1 only when maxByCols caps it at 1 for very narrow rows - // (numColumns < kDecodeMinColsPerSubBlock). - blocksPerRow = std::min(2, maxByCols); - } - return std::max(1, blocksPerRow); -} - -void invokeIndexerTopKDecode(float const* logits, int const* seqLens, int* indices, float* outLogitsAux, - int* outIndicesAux, int const splitWorkThreshold, int const numRows, int const numColumns, int const stride0, - int const stride1, int const next_n, int const topK, int const* preIdx, int const preIdxStride, - int const preIdxCount, float* heuristicScratch, int const compressRatio, cudaStream_t const stream) +// Unified dispatcher (fp32 / bf16 / fp16). Each tier's kernel is templated on +// InputT and casts to float at HBM-read sites; the scratch buffer is uint8. +template +void invokeIndexerTopKDecodeImpl(InputT const* logits, int const* seqLens, int* indices, int const splitWorkThreshold, + int const numRows, int const numColumns, int const stride0, int const stride1, int const next_n, int const topK, + int const* preIdx, int const preIdxStride, int const preIdxCount, InputT* heuristicScratch, + cudaStream_t const stream, void* scratch, size_t scratchBytes, bool is_prefill) { + // Split-work cutoff: matches main's 200k default. is_prefill forces + // single-block via a 1<<30 threshold no shape can reach: prefill chunks are + // bounded by max_num_tokens, well below the multi-pass radix crossover at + // any practical setting. + int const adaptiveSplitWorkThreshold = is_prefill ? (1 << 30) : 200 * 1000; + int const effectiveSplitWorkThreshold = splitWorkThreshold > 0 ? splitWorkThreshold : adaptiveSplitWorkThreshold; constexpr int kNumThreadsPerBlock = 512; - int const effectiveSplitWorkThreshold = splitWorkThreshold > 0 ? splitWorkThreshold : kDefaultSplitWorkThreshold; - - // ======================================================================== - // Small-N dispatch axis. - // - // GVR Heuristic Top-K has a *fixed* per-launch overhead from Phase-1 - // (preIdx stats reduction over M=2048) and Phase-4 (2048-bin histogram - // snap), totaling ~11 µs regardless of N. For small N (≤16K), this - // fixed cost dominates and the kernel loses to the existing - // insertion-sort/radix path. Empirically (random data, B200 BS=1): - // N=8192 : Heuristic 16.5 µs vs Radix 11.2 µs (radix 1.47× faster) - // N=16384 : Heuristic 21.9 µs vs Radix 22.0 µs (parity) - // N=32768 : Heuristic 26.1 µs vs Radix 32.9 µs (heuristic 1.26× faster) - // N=131072 : Heuristic 43.4 µs vs Radix 76.1 µs (heuristic 1.75× faster) - // - // Route N < kSeqSmall to the existing Radix/Insertion path (which itself - // splits at kSortingAlgorithmThreshold=12288). kSeqSmall is set at the - // empirical crossover point. - // - // ======================================================================== - // Architecture-derived BS-threshold dispatch — jointly bounded by - // occupancy AND L2 cache capacity. - // - // Two physical constraints bound when the per-row heuristic kernel - // remains faster than a radix streaming kernel: - // - // (A) Occupancy bound — 3·SM − SM/8 (wave geometry + setup margin) - // Each CTA uses ~58 KB SMEM (fixed, independent of N), so B200's - // 228 KB dynamic SMEM allows max 3 CTA/SM. Above 3·SM rows per - // launch, tail-wave imbalance causes stragglers. The -SM/8 margin - // (~1/8 wave) covers CTA setup + L2 ingestion overhead. - // On B200(148 SM): 3×148 − 18 = 426. - // - // (B) L2 cache bound — 0.9·L2 / (4·N) per-CTA logits fit - // Each CTA streams its row (N×4B) through L2 per Phase-2 iter. - // With num_concurrent_CTAs × N × 4B > L2, eviction dominates. - // On B200(126 MB L2) with N=70K: 0.9·126MB/(4·70690) ≈ 440, - // which is ~ equal to (A)=426 — the two constraints cross over - // near the SWE-Bench data point. - // For N > 73K the L2 bound tightens below (A) and must take - // over; e.g. N=128K → kBsL2=238, N=196K → kBsL2=155. - // - // Dispatch threshold = min(kBsWave, kBsL2), still data-agnostic (only - // queries hardware attrs). At N≈70K both bounds produce ~426, so the - // L2 axis is a no-op there; for larger N it auto-tightens the threshold. - // - // Small-N lower bound `kSeqSmall` is uniform 4096 across all K (see - // kSeqSmallDefaultForK). 4K is the dispatcher's algorithmic-handoff - // point: below 4096 the Radix path resolves to single-CTA insertion-sort - // (maxByCols = N/2048 = 1 → bp=1; useRadixSort = N≥12288 = false), and - // GVR is empirically slower than insertion-sort below 4K across all - // K ∈ {512, 1024, 2048} × dtype ∈ {fp32, bf16, fp16} (R/H ∈ [0.55, 0.84] - // at N=2K; DSv4 V3.2-Q19c synth sweeps May 2026). Configurable via - // TRTLLM_HEURISTIC_NMIN env (>=1024), which overrides the default. - // ======================================================================== - auto const bounds = getSchemeXBounds(numColumns, /*bytesPerElem=*/4, topK); - int const kBsWave = bounds.kBsWave; - int const kBsL2 = bounds.kBsL2; - int const kBsLarge = bounds.kBsLarge; - int const kSeqSmall = bounds.kSeqSmall; + // GVR eligibility (matches main's rule): supported K, stride1 contiguous, + // preIdx + scratch provided, numColumns in [kSeqSmall, splitWorkThreshold), + // and numRows below the architecture-derived wave/L2 bound. is_prefill + // suppresses GVR through effectiveSplitWorkThreshold being huge. + auto const bounds = getSchemeXBounds(numColumns, /*bytesPerElem=*/static_cast(sizeof(InputT))); bool const isSupportedTopK = (topK == 512 || topK == 1024 || topK == 2048); - // compressRatio == 1: DSv3.2 indexer (no compressor). - // compressRatio == 4: DSv4 indexer (overlap compressor); logits/preIdx in - // compressed-token-index space. Kernel handles N = actual_kv_len/cr and - // forces preIdxOffset=0 internally for cr != 1. - bool const compressRatioOk = (compressRatio == 1 || compressRatio == 4); - bool const canUseHeuristic = compressRatioOk && preIdx != nullptr && stride1 == 1 && isSupportedTopK - && preIdxCount == topK && preIdxStride >= preIdxCount && numColumns < effectiveSplitWorkThreshold - && numColumns >= kSeqSmall && heuristicScratch != nullptr && numRows < kBsLarge; - - // Optional env-gated dispatch trace (set TRTLLM_SCHEMEX_DEBUG=1 to enable) + bool const canUseHeuristic = preIdx != nullptr && stride1 == 1 && isSupportedTopK && preIdxCount == topK + && preIdxStride >= preIdxCount && heuristicScratch != nullptr && numColumns >= bounds.kSeqSmall + && numColumns < effectiveSplitWorkThreshold && numRows < bounds.kBsLarge; + + // Env-gated dispatch trace (TRTLLM_SCHEMEX_DEBUG=1). { static std::once_flag sDebugOnceFlag; static bool sDebug = false; @@ -976,161 +1189,21 @@ void invokeIndexerTopKDecode(float const* logits, int const* seqLens, int* indic }); if (sDebug) { - fprintf(stderr, - "[Scheme X] numRows=%d numColumns=%d kBsWave=%d kBsL2=%d kBsLarge=%d kSeqSmall=%d smCount=%d " - "L2=%dMB -> %s path%s\n", - numRows, numColumns, kBsWave, kBsL2, kBsLarge, kSeqSmall, bounds.smCount, - bounds.l2Bytes / (1024 * 1024), canUseHeuristic ? "Heuristic" : "Radix", - (numColumns < kSeqSmall) ? " (small-N route)" : ""); + fprintf(stderr, "[Scheme X] numRows=%d numColumns=%d kBsLarge=%d kSeqSmall=%d -> %s path\n", numRows, + numColumns, bounds.kBsLarge, bounds.kSeqSmall, canUseHeuristic ? "Heuristic" : "Radix"); } } if (canUseHeuristic) { launchHeuristicTopKDecode(logits, seqLens, preIdx, indices, heuristicScratch, stride0, next_n, topK, - preIdxStride, preIdxCount, numRows, compressRatio, stream); - sync_check_cuda_error(stream); - return; - } - - int const blocksPerRow = computeIndexerTopKDecodeBlocksPerRow(numRows, numColumns, splitWorkThreshold); - - cudaLaunchAttribute attrs[1]; - attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; - attrs[0].val.programmaticStreamSerializationAllowed = tensorrt_llm::common::getEnvEnablePDL(); - - if (blocksPerRow == 1) - { - // Single block per row. Below kSortingAlgorithmThreshold use insertion sort, - // above use the histogram-radix path. - bool const useRadixSort = numColumns >= kSortingAlgorithmThreshold; - auto* kernel_instance = useRadixSort ? &topKPerRowDecode - : &topKPerRowDecode; - - cudaLaunchConfig_t config; - config.gridDim = numRows; - config.blockDim = kNumThreadsPerBlock; - config.dynamicSmemBytes = topK * sizeof(int32_t); - config.stream = stream; - config.numAttrs = 1; - config.attrs = attrs; - - cudaLaunchKernelEx(&config, kernel_instance, logits, seqLens, indices, stride0, stride1, topK, next_n, - compressRatio, nullptr, 0, nullptr); - } - else - { - // Split each row across `blocksPerRow` blocks, then merge with a second pass. - auto* kernel_instance_part1 = &topKPerRowDecode; - cudaLaunchConfig_t config_part1; - config_part1.gridDim = dim3(numRows, blocksPerRow); - config_part1.blockDim = kNumThreadsPerBlock; - config_part1.dynamicSmemBytes = 2 * topK * sizeof(int32_t); - config_part1.stream = stream; - config_part1.numAttrs = 1; - config_part1.attrs = attrs; - - cudaLaunchKernelEx(&config_part1, kernel_instance_part1, logits, seqLens, outIndicesAux, stride0, stride1, topK, - next_n, compressRatio, outLogitsAux, 0, nullptr); - - constexpr int kNumThreadsPerBlockMerge = 1024; - auto* kernel_instance_part2 = &topKPerRowDecode; - cudaLaunchConfig_t config_part2; - config_part2.gridDim = numRows; - config_part2.blockDim = kNumThreadsPerBlockMerge; - config_part2.dynamicSmemBytes = topK * sizeof(int32_t); - config_part2.stream = stream; - config_part2.numAttrs = 1; - config_part2.attrs = attrs; - - cudaLaunchKernelEx(&config_part2, kernel_instance_part2, outLogitsAux, seqLens, indices, blocksPerRow * topK, 1, - topK, next_n, 1, nullptr, blocksPerRow, outIndicesAux); - } - sync_check_cuda_error(stream); -} - -// ============================================================================ -// bf16 / fp16 dispatcher overloads -// ============================================================================ -// Reuses the BS-threshold + small-N dispatch axes (kBsLarge, kSeqSmall) from -// the fp32 dispatcher, except kBsL2 uses sizeof(InputT) bytes/element instead -// of 4 — L2 footprint is half, so bf16/fp16 path remains valid for larger BS -// than fp32 at the same N. -// -// Fallback chain when GVR-Heuristic preconditions are not met (preIdx -// missing, BS too large, or numColumns < kSeqSmall): -// numColumns < kSortingAlgorithmThreshold (12288) → insertion sort -// kSortingAlgorithmThreshold ≤ numColumns < splitWorkThreshold → radix sort -// numColumns ≥ splitWorkThreshold (200K default) → unsupported -// -// Insertion + radix tiers use the same topKPerRowDecode kernel as fp32 with -// InputT propagated through; the histogram and sort steps operate on float -// keys after a static_cast(InputT) at HBM-read sites, so accuracy is -// identical to casting input to fp32 before the kernel. -// -// The split-work tier requires float aux buffers (outLogitsAux / -// outIndicesAux) that the bf16/fp16 entry does not expose; callers in that -// regime must use the fp32 entry. - -namespace -{ - -template -void invokeIndexerTopKDecodeDtype(InputT const* logits, int const* seqLens, int* indices, int const splitWorkThreshold, - int const numRows, int const numColumns, int const stride0, int const stride1, int const next_n, int const topK, - int const* preIdx, int const preIdxStride, int const preIdxCount, InputT* heuristicScratch, int const compressRatio, - cudaStream_t const stream) -{ - static_assert(std::is_same_v || std::is_same_v, - "invokeIndexerTopKDecodeDtype is for bf16/fp16 only"); - - constexpr int kNumThreadsPerBlock = 512; - int const effectiveSplitWorkThreshold = splitWorkThreshold > 0 ? splitWorkThreshold : kDefaultSplitWorkThreshold; - - // bf16/fp16: bytes_per_element = sizeof(InputT) = 2 → kBsL2 doubles vs fp32. - // K-aware kSeqSmall — see fp32 dispatcher for rationale. - auto const bounds = getSchemeXBounds(numColumns, /*bytesPerElem=*/static_cast(sizeof(InputT)), topK); - int const kBsLarge = bounds.kBsLarge; - int const kSeqSmall = bounds.kSeqSmall; - - bool const isSupportedTopK = (topK == 512 || topK == 1024 || topK == 2048); - // See fp32 path: cr==1 (V3.2) and cr==4 (V4 indexer) are both supported. - bool const compressRatioOk = (compressRatio == 1 || compressRatio == 4); - bool const canUseHeuristic = compressRatioOk && preIdx != nullptr && stride1 == 1 && isSupportedTopK - && preIdxCount == topK && preIdxStride >= preIdxCount && numColumns < effectiveSplitWorkThreshold - && numColumns >= kSeqSmall && heuristicScratch != nullptr && numRows < kBsLarge; - - if (canUseHeuristic) - { - launchHeuristicTopKDecode(logits, seqLens, preIdx, indices, heuristicScratch, stride0, next_n, topK, - preIdxStride, preIdxCount, numRows, compressRatio, stream); - } - else if (numColumns < kSortingAlgorithmThreshold) - { - // Insertion sort path — InputT propagated; histogram/sort run on float keys. - auto* kernel_instance = &topKPerRowDecode; - - cudaLaunchConfig_t config; - config.gridDim = numRows; - config.blockDim = kNumThreadsPerBlock; - config.dynamicSmemBytes = topK * sizeof(int32_t); - config.stream = stream; - cudaLaunchAttribute attrs[1]; - attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; - attrs[0].val.programmaticStreamSerializationAllowed = tensorrt_llm::common::getEnvEnablePDL(); - config.numAttrs = 1; - config.attrs = attrs; - - cudaLaunchKernelEx(&config, kernel_instance, logits, seqLens, indices, stride0, stride1, topK, next_n, - compressRatio, nullptr, 0, nullptr); + preIdxStride, preIdxCount, numRows, stream); } else if (numColumns < effectiveSplitWorkThreshold) { - // Radix sort path — InputT propagated; histogram/sort run on float keys. - auto* kernel_instance = &topKPerRowDecode; - + // Single-block tier: one CTA per row. + auto* kernel_instance = &topKPerRowDecode; cudaLaunchConfig_t config; config.gridDim = numRows; config.blockDim = kNumThreadsPerBlock; @@ -1141,40 +1214,62 @@ void invokeIndexerTopKDecodeDtype(InputT const* logits, int const* seqLens, int* attrs[0].val.programmaticStreamSerializationAllowed = tensorrt_llm::common::getEnvEnablePDL(); config.numAttrs = 1; config.attrs = attrs; - cudaLaunchKernelEx(&config, kernel_instance, logits, seqLens, indices, stride0, stride1, topK, next_n, - compressRatio, nullptr, 0, nullptr); + /*outLogits=*/nullptr, /*numBlocksToMerge=*/0, /*indices=*/nullptr); } else { - TLLM_CHECK_WITH_INFO(false, - "indexer_topk_decode bf16/fp16 path does not support numColumns >= splitWorkThreshold " - "(split-work path requires float aux buffers not exposed in the bf16/fp16 entry). " - "Got numColumns=%d splitWorkThreshold=%d. Use the fp32 entry for this regime.", - numColumns, effectiveSplitWorkThreshold); + // Multi-pass radix. radixPassKernel reads logits contiguously, so + // strided inputs would rank the wrong values — gate on stride1 == 1. + // (The single-block tier handles stride1 != 1 via topKPerRowJob's + // strided fallback.) + TLLM_CHECK_WITH_INFO(stride1 == 1, "indexer top-k split-work tier (multi-pass radix) requires stride1 == 1."); + TLLM_CHECK_WITH_INFO(scratch != nullptr && scratchBytes >= radixScratchBytes(numRows, numColumns), + "indexer top-k split-work tier: scratch buffer missing or too small."); + cudaLaunchAttribute radixAttrs[1]; + radixAttrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + radixAttrs[0].val.programmaticStreamSerializationAllowed = tensorrt_llm::common::getEnvEnablePDL(); + launchMultiPassRadix( + scratch, logits, seqLens, indices, numRows, numColumns, topK, stride0, next_n, radixAttrs, stream); } - sync_check_cuda_error(stream); } } // anonymous namespace +void invokeIndexerTopKDecode(float const* logits, int const* seqLens, int* indices, int const splitWorkThreshold, + int const numRows, int const numColumns, int const stride0, int const stride1, int const next_n, int const topK, + int const* preIdx, int const preIdxStride, int const preIdxCount, float* heuristicScratch, + cudaStream_t const stream, void* scratch, size_t scratchBytes, bool is_prefill) +{ + invokeIndexerTopKDecodeImpl(logits, seqLens, indices, splitWorkThreshold, numRows, numColumns, stride0, + stride1, next_n, topK, preIdx, preIdxStride, preIdxCount, heuristicScratch, stream, scratch, scratchBytes, + is_prefill); +} + +size_t indexerTopKDecodeScratchBytes(int numRows, int numColumns, int /*topK*/) +{ + return radixScratchBytes(numRows, numColumns); +} + void invokeIndexerTopKDecode(__nv_bfloat16 const* logits, int const* seqLens, int* indices, int const splitWorkThreshold, int const numRows, int const numColumns, int const stride0, int const stride1, int const next_n, int const topK, int const* preIdx, int const preIdxStride, int const preIdxCount, - __nv_bfloat16* heuristicScratch, int const compressRatio, cudaStream_t const stream) + __nv_bfloat16* heuristicScratch, cudaStream_t const stream, void* scratch, size_t scratchBytes, bool is_prefill) { - invokeIndexerTopKDecodeDtype<__nv_bfloat16>(logits, seqLens, indices, splitWorkThreshold, numRows, numColumns, - stride0, stride1, next_n, topK, preIdx, preIdxStride, preIdxCount, heuristicScratch, compressRatio, stream); + invokeIndexerTopKDecodeImpl<__nv_bfloat16>(logits, seqLens, indices, splitWorkThreshold, numRows, numColumns, + stride0, stride1, next_n, topK, preIdx, preIdxStride, preIdxCount, heuristicScratch, stream, scratch, + scratchBytes, is_prefill); } void invokeIndexerTopKDecode(__half const* logits, int const* seqLens, int* indices, int const splitWorkThreshold, int const numRows, int const numColumns, int const stride0, int const stride1, int const next_n, int const topK, - int const* preIdx, int const preIdxStride, int const preIdxCount, __half* heuristicScratch, int const compressRatio, - cudaStream_t const stream) + int const* preIdx, int const preIdxStride, int const preIdxCount, __half* heuristicScratch, + cudaStream_t const stream, void* scratch, size_t scratchBytes, bool is_prefill) { - invokeIndexerTopKDecodeDtype<__half>(logits, seqLens, indices, splitWorkThreshold, numRows, numColumns, stride0, - stride1, next_n, topK, preIdx, preIdxStride, preIdxCount, heuristicScratch, compressRatio, stream); + invokeIndexerTopKDecodeImpl<__half>(logits, seqLens, indices, splitWorkThreshold, numRows, numColumns, stride0, + stride1, next_n, topK, preIdx, preIdxStride, preIdxCount, heuristicScratch, stream, scratch, scratchBytes, + is_prefill); } void invokeIndexerTopKPrefill(float const* logits, int const* rowStarts, int const* rowEnds, int* indices, @@ -1183,18 +1278,10 @@ void invokeIndexerTopKPrefill(float const* logits, int const* rowStarts, int con { constexpr int kNumThreadsPerBlock = 512; - int numInsertionBlocks = std::min(numRows, kSortingAlgorithmThreshold); - topKPerRowPrefill - <<>>( - logits, rowStarts, rowEnds, indices, stride0, stride1, topK, 0); - - if (numRows > kSortingAlgorithmThreshold) - { - int numRadixBlocks = numRows - kSortingAlgorithmThreshold; - topKPerRowPrefill - <<>>( - logits, rowStarts, rowEnds, indices, stride0, stride1, topK, kSortingAlgorithmThreshold); - } + // One launch over all rows; the per-row sort algorithm is picked at + // runtime inside topKPerRowJob. + topKPerRowPrefill<<>>( + logits, rowStarts, rowEnds, indices, stride0, stride1, topK, 0); sync_check_cuda_error(stream); } @@ -1206,7 +1293,9 @@ bool canIndexerTopKDecodeUseGvr(int numRows, int numColumns, int topK, int bytes { return false; } - auto const bounds = getSchemeXBounds(numColumns, bytesPerElem, topK); + // Mirrors the dispatcher's effectiveSplitWorkThreshold default. + constexpr int kDefaultSplitWorkThreshold = 200 * 1000; + auto const bounds = getSchemeXBounds(numColumns, bytesPerElem); return numColumns >= bounds.kSeqSmall && numColumns < kDefaultSplitWorkThreshold && numRows < bounds.kBsLarge; } diff --git a/cpp/tensorrt_llm/kernels/inverseRopeFp8QuantKernel.cu b/cpp/tensorrt_llm/kernels/inverseRopeFp8QuantKernel.cu deleted file mode 100644 index 245e573efeb7..000000000000 --- a/cpp/tensorrt_llm/kernels/inverseRopeFp8QuantKernel.cu +++ /dev/null @@ -1,426 +0,0 @@ -/* - * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. - * - * 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. - */ -#include "tensorrt_llm/kernels/inverseRopeFp8QuantKernel.h" - -#include -#include -#include -#include - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels -{ -namespace -{ - -// Fixed by the DSv4 inverse-RoPE op contract (matches the Triton reference): -// * QUANT_GROUP_SIZE = 128 -// * ROPE_DIM = 64 (with HALF_ROPE = 32, NEOX layout) -// * Rope occupies the LAST ROPE_DIM elements of the head (i.e. the second -// half of the final quant chunk). Earlier chunks are pure nope. -// * Per-lane-per-chunk: 4 fp32 elements, full warp covers all 128 elts. -constexpr int kQuantGroupSize = 128; -constexpr int kRopeDim = 64; -constexpr int kHalfRope = kRopeDim / 2; // 32 -constexpr int kCsStride = 2 * kHalfRope; // 64 fp32 per position (cos block + sin block) -constexpr int kRopeStartInChunk = kQuantGroupSize - kRopeDim; // 64; rope is [64,128) of last chunk -constexpr int kWarpSize = 32; -constexpr int kEltsPerThreadPerChunk = kQuantGroupSize / kWarpSize; // 4 -constexpr float kFp8Max = 448.0f; -constexpr float kFp8InvMax = 1.0f / 448.0f; -constexpr float kEps = 1e-12f; - -// Per-warp per-chunk processing. Each chunk = 128 elements = 32 lanes × 4 elts. -// Caller passes 4 bf16 inputs (already loaded into registers), and the -// chunk-relative `is_rope_chunk` flag controls whether the second half of -// the chunk gets the inverse-NEOX rotation. Within the rope chunk the -// elements [kRopeStartInChunk, 128) carry the rope, partner across the -// rope halves is lane ^ 8. -// -// On exit: x_quant_lo / x_quant_hi hold the packed fp8x2 outputs to store -// for this chunk, and `scale_out` is the per-chunk fp32 scale (written -// only by lane 0). -template -__device__ __forceinline__ void process_chunk(__nv_bfloat162 in_lo, __nv_bfloat162 in_hi, - float const* __restrict__ cs_row, bool is_rope_chunk, __nv_fp8x2_e4m3& out_lo, __nv_fp8x2_e4m3& out_hi, - float& scale_out) -{ - int const lane = threadIdx.x & 31; - int const e0_in_chunk = lane * kEltsPerThreadPerChunk; // [0..128) in steps of 4 - - float x[4]; - x[0] = __bfloat162float(in_lo.x); - x[1] = __bfloat162float(in_lo.y); - x[2] = __bfloat162float(in_hi.x); - x[3] = __bfloat162float(in_hi.y); - - // NEOX layout: partner across the two rope halves = `lane XOR 8`. All - // 32 lanes must issue the shuffle (mask=0xffffffff). Skipped under - // GPT-J / interleaved layout where the partner is intra-lane. - unsigned part_lo_u = 0u, part_hi_u = 0u; - if (IS_NEOX) - { - int const partner_lane = lane ^ 8; - unsigned in_lo_u = *reinterpret_cast(&in_lo); - unsigned in_hi_u = *reinterpret_cast(&in_hi); - part_lo_u = __shfl_sync(0xFFFFFFFFu, in_lo_u, partner_lane); - part_hi_u = __shfl_sync(0xFFFFFFFFu, in_hi_u, partner_lane); - } - - bool const lane_in_rope = is_rope_chunk && (lane >= 16); - - if (IS_NEOX && lane_in_rope) - { - // Inverse-NEOX rope: lanes 16..23 own the first rope half - // (rope_local 0..31, chunk-local 64..95); lanes 24..31 own the - // second half (rope_local 32..63, chunk-local 96..127). - bool const lane_in_first_half = (lane < 24); - __nv_bfloat162 p2_lo = *reinterpret_cast<__nv_bfloat162*>(&part_lo_u); - __nv_bfloat162 p2_hi = *reinterpret_cast<__nv_bfloat162*>(&part_hi_u); - - float xp[4]; - xp[0] = __bfloat162float(p2_lo.x); - xp[1] = __bfloat162float(p2_lo.y); - xp[2] = __bfloat162float(p2_hi.x); - xp[3] = __bfloat162float(p2_hi.y); - - int const cs_base = e0_in_chunk - kRopeStartInChunk - (lane_in_first_half ? 0 : kHalfRope); - float const sign = lane_in_first_half ? 1.0f : -1.0f; - // 16-byte (float4) coalesced load -- cs_base is always a multiple of 4. - float4 const cos4 = *reinterpret_cast(cs_row + cs_base); - float4 const sin4 = *reinterpret_cast(cs_row + kHalfRope + cs_base); - x[0] = x[0] * cos4.x + sign * sin4.x * xp[0]; - x[1] = x[1] * cos4.y + sign * sin4.y * xp[1]; - x[2] = x[2] * cos4.z + sign * sin4.z * xp[2]; - x[3] = x[3] * cos4.w + sign * sin4.w * xp[3]; - } - else if (!IS_NEOX && lane_in_rope) - { - // Interleaved (GPT-J) inverse RoPE -- partners are adjacent pairs - // (x[2i], x[2i+1]) within the rope segment. Each lane's 4 elements - // form two intra-lane pairs (x[0],x[1]) and (x[2],x[3]); no - // cross-lane shuffle is needed. - // for rope_local r: cs_idx = r >> 1 - // r even: new = x[r]*cos[cs_idx] + x[r+1]*sin[cs_idx] - // r odd : new = x[r]*cos[cs_idx] - x[r-1]*sin[cs_idx] - // Lane t in [16,31] owns rope_local in [(t-16)*4, (t-16)*4+4), so - // cs_idx ∈ {(t-16)*2, (t-16)*2 + 1} for the two intra-lane pairs. - int const cs_base = (lane - 16) * 2; - float2 const cos2 = *reinterpret_cast(cs_row + cs_base); - float2 const sin2 = *reinterpret_cast(cs_row + kHalfRope + cs_base); - float const x0_new = x[0] * cos2.x + x[1] * sin2.x; - float const x1_new = x[1] * cos2.x - x[0] * sin2.x; - float const x2_new = x[2] * cos2.y + x[3] * sin2.y; - float const x3_new = x[3] * cos2.y - x[2] * sin2.y; - x[0] = x0_new; - x[1] = x1_new; - x[2] = x2_new; - x[3] = x3_new; - } - - // Per-chunk warp absmax (128 elements). - float local_max = fmaxf(fabsf(x[0]), fmaxf(fabsf(x[1]), fmaxf(fabsf(x[2]), fabsf(x[3])))); -#pragma unroll - for (int mask = 16; mask > 0; mask >>= 1) - { - local_max = fmaxf(local_max, __shfl_xor_sync(0xFFFFFFFFu, local_max, mask)); - } - float const block_max = fmaxf(local_max, kEps); - scale_out = block_max * kFp8InvMax; - float const inv_scale = __fdividef(kFp8Max, block_max); - - out_lo = __nv_fp8x2_e4m3(float2{x[0] * inv_scale, x[1] * inv_scale}); - out_hi = __nv_fp8x2_e4m3(float2{x[2] * inv_scale, x[3] * inv_scale}); -} - -// Baseline 1-warp-per-token variant. Each warp handles one (head, token) -// pair and emits CHUNKS_PER_HEAD fp8 chunks + scales. The kernel hardcodes -// the DSv4 layout: rope_dim=64 lives in the second half of the last chunk. -template -__global__ __launch_bounds__(BLOCK_TOKENS* kWarpSize) void inverseRopeFp8QuantKernel( - __nv_bfloat16 const* __restrict__ o_ptr, // - int64_t const* __restrict__ positions_ptr, // - float const* __restrict__ cs_cache_ptr, // - __nv_fp8_e4m3* __restrict__ fp8_ptr, // - float* __restrict__ scale_ptr, // - int num_tokens, int scale_buf_m, int heads_per_group, // - int o_stride_token, int o_stride_head, // - int fp8_stride_group, int fp8_stride_token, // - int scale_stride_group, int scale_stride_k) -{ - constexpr int HEAD_DIM = CHUNKS_PER_HEAD * kQuantGroupSize; - int const warp_id = threadIdx.x >> 5; - int const lane = threadIdx.x & 31; - int const e0_in_chunk = lane * kEltsPerThreadPerChunk; - int const pid_token = blockIdx.x * BLOCK_TOKENS + warp_id; - int const head_idx = blockIdx.y; // global head: [0, num_heads) - int const g_idx = head_idx / heads_per_group; - int const h_in_g = head_idx - g_idx * heads_per_group; - int const qb_base = h_in_g * CHUNKS_PER_HEAD; - - if (pid_token >= scale_buf_m) - return; - - if (pid_token >= num_tokens) - { - // Pad row in [num_tokens, pad_up(num_tokens, 4)): zero out all - // CHUNKS_PER_HEAD scale slots for this (group, head_in_group). - if (lane == 0) - { - float* sc_grp = scale_ptr + static_cast(g_idx) * scale_stride_group; -#pragma unroll - for (int c = 0; c < CHUNKS_PER_HEAD; ++c) - { - sc_grp[(qb_base + c) * scale_stride_k + pid_token] = 0.0f; - } - } - return; - } - - long const pos = positions_ptr[pid_token]; - // Input layout: [num_tokens, num_heads, head_dim] flat — global head_idx works. - auto const* in_row - = o_ptr + static_cast(pid_token) * o_stride_token + static_cast(head_idx) * o_stride_head; - // Output layout: fp8_buf [n_groups, num_tokens, heads_per_group * head_dim] - // and scale_buf [n_groups, heads_per_group*chunks, pad_up(T,4)] — both with - // n_groups outermost (the BMM consumer's expected (G, T, K) view). - auto* fp8_row = fp8_ptr + static_cast(g_idx) * fp8_stride_group - + static_cast(pid_token) * fp8_stride_token + static_cast(h_in_g) * HEAD_DIM; - float* sc_grp = scale_ptr + static_cast(g_idx) * scale_stride_group; - auto const* cs_row = cs_cache_ptr + pos * kCsStride; - -#pragma unroll - for (int c = 0; c < CHUNKS_PER_HEAD; ++c) - { - // Load this chunk's 4 bf16 per lane. - auto const* in_pair = reinterpret_cast<__nv_bfloat162 const*>(in_row + c * kQuantGroupSize + e0_in_chunk); - __nv_bfloat162 in_lo = in_pair[0]; - __nv_bfloat162 in_hi = in_pair[1]; - - bool const is_rope = (c == CHUNKS_PER_HEAD - 1); - __nv_fp8x2_e4m3 out_lo, out_hi; - float scale; - process_chunk(in_lo, in_hi, cs_row, is_rope, out_lo, out_hi, scale); - - auto* out_pair = reinterpret_cast<__nv_fp8x2_e4m3*>(fp8_row + c * kQuantGroupSize + e0_in_chunk); - out_pair[0] = out_lo; - out_pair[1] = out_hi; - - if (lane == 0) - { - sc_grp[(qb_base + c) * scale_stride_k + pid_token] = scale; - } - } -} - -// Software-pipelined variant: each warp processes TOKENS_PER_WARP tokens -// with explicit double-buffered load/compute interleaving. Hides L1 -// scoreboard stalls by overlapping the next iter's input/cs LDGs with the -// current iter's compute+store. TPW=2 is the sweet spot — TPW>=4 spills -// registers and drops occupancy. -template -__global__ __launch_bounds__(BLOCK_TOKENS* kWarpSize) void inverseRopeFp8QuantKernelPipelined( - __nv_bfloat16 const* __restrict__ o_ptr, // - int64_t const* __restrict__ positions_ptr, // - float const* __restrict__ cs_cache_ptr, // - __nv_fp8_e4m3* __restrict__ fp8_ptr, // - float* __restrict__ scale_ptr, // - int num_tokens, int scale_buf_m, int heads_per_group, // - int o_stride_token, int o_stride_head, // - int fp8_stride_group, int fp8_stride_token, // - int scale_stride_group, int scale_stride_k) -{ - constexpr int HEAD_DIM = CHUNKS_PER_HEAD * kQuantGroupSize; - int const warp_id = threadIdx.x >> 5; - int const lane = threadIdx.x & 31; - int const e0_in_chunk = lane * kEltsPerThreadPerChunk; - int const head_idx = blockIdx.y; - int const g_idx = head_idx / heads_per_group; - int const h_in_g = head_idx - g_idx * heads_per_group; - int const qb_base = h_in_g * CHUNKS_PER_HEAD; - int const pid_token_base = blockIdx.x * (BLOCK_TOKENS * TOKENS_PER_WARP) + warp_id * TOKENS_PER_WARP; - - // ---- Stage 1: issue all input + position LDGs up front -------------- - __nv_bfloat162 in_lo_arr[TOKENS_PER_WARP][CHUNKS_PER_HEAD]; - __nv_bfloat162 in_hi_arr[TOKENS_PER_WARP][CHUNKS_PER_HEAD]; - long pos_arr[TOKENS_PER_WARP]; - bool valid[TOKENS_PER_WARP]; - bool in_range[TOKENS_PER_WARP]; - -#pragma unroll - for (int t = 0; t < TOKENS_PER_WARP; ++t) - { - int const pid = pid_token_base + t; - in_range[t] = pid < scale_buf_m; - valid[t] = pid < num_tokens; - if (valid[t]) - { - pos_arr[t] = positions_ptr[pid]; - auto const* in_row - = o_ptr + static_cast(pid) * o_stride_token + static_cast(head_idx) * o_stride_head; -#pragma unroll - for (int c = 0; c < CHUNKS_PER_HEAD; ++c) - { - auto const* in_pair - = reinterpret_cast<__nv_bfloat162 const*>(in_row + c * kQuantGroupSize + e0_in_chunk); - in_lo_arr[t][c] = in_pair[0]; - in_hi_arr[t][c] = in_pair[1]; - } - } - } - - // ---- Stage 2: compute + store, one token at a time ----------------- -#pragma unroll - for (int t = 0; t < TOKENS_PER_WARP; ++t) - { - int const pid = pid_token_base + t; - if (!in_range[t]) - continue; - - float* sc_grp = scale_ptr + static_cast(g_idx) * scale_stride_group; - if (!valid[t]) - { - if (lane == 0) - { -#pragma unroll - for (int c = 0; c < CHUNKS_PER_HEAD; ++c) - { - sc_grp[(qb_base + c) * scale_stride_k + pid] = 0.0f; - } - } - continue; - } - - auto const* cs_row = cs_cache_ptr + pos_arr[t] * kCsStride; - auto* fp8_row = fp8_ptr + static_cast(g_idx) * fp8_stride_group - + static_cast(pid) * fp8_stride_token + static_cast(h_in_g) * HEAD_DIM; - -#pragma unroll - for (int c = 0; c < CHUNKS_PER_HEAD; ++c) - { - bool const is_rope = (c == CHUNKS_PER_HEAD - 1); - __nv_fp8x2_e4m3 out_lo, out_hi; - float scale; - process_chunk(in_lo_arr[t][c], in_hi_arr[t][c], cs_row, is_rope, out_lo, out_hi, scale); - - auto* out_pair = reinterpret_cast<__nv_fp8x2_e4m3*>(fp8_row + c * kQuantGroupSize + e0_in_chunk); - out_pair[0] = out_lo; - out_pair[1] = out_hi; - - if (lane == 0) - { - sc_grp[(qb_base + c) * scale_stride_k + pid] = scale; - } - } - } -} - -template -inline void dispatchByM(int num_tokens, int num_heads, int heads_per_group, int scale_buf_m, // - __nv_bfloat16 const* o_p, int64_t const* pos_p, float const* cs_p, __nv_fp8_e4m3* fp8_p, float* sc_p, // - int o_stride_token, int o_stride_head, int fp8_stride_group, int fp8_stride_token, // - int scale_stride_group, int scale_stride_k, cudaStream_t stream) -{ - // BTM=4 with TPW=2 above M=4096, otherwise baseline. Tuned on B200. - if (num_tokens >= 4096) - { - constexpr int BTM = 4; - constexpr int TPW = 2; - int const tokens_per_cta = BTM * TPW; - int const grid_x = (scale_buf_m + tokens_per_cta - 1) / tokens_per_cta; - dim3 grid(grid_x, num_heads); - dim3 block(BTM * kWarpSize); - inverseRopeFp8QuantKernelPipelined<<>>( // - o_p, pos_p, cs_p, fp8_p, sc_p, // - num_tokens, scale_buf_m, heads_per_group, // - o_stride_token, o_stride_head, fp8_stride_group, fp8_stride_token, // - scale_stride_group, scale_stride_k); - } - else - { - constexpr int BTM = 4; - int const grid_x = (scale_buf_m + BTM - 1) / BTM; - dim3 grid(grid_x, num_heads); - dim3 block(BTM * kWarpSize); - inverseRopeFp8QuantKernel<<>>( // - o_p, pos_p, cs_p, fp8_p, sc_p, // - num_tokens, scale_buf_m, heads_per_group, // - o_stride_token, o_stride_head, fp8_stride_group, fp8_stride_token, // - scale_stride_group, scale_stride_k); - } -} - -} // namespace - -void invokeInverseRopeFp8Quant(void const* o, // - void const* positions, // - void const* cos_sin_cache, // - void* fp8_out, // - void* scale_out, // - int num_tokens, // - int num_heads, // - int heads_per_group, // - int chunks_per_head, // - bool is_neox, // - int scale_buf_m, // - int o_stride_token, // - int o_stride_head, // - int fp8_stride_group, // - int fp8_stride_token, // - int scale_stride_group, // - int scale_stride_k, // - cudaStream_t stream) -{ - auto const* o_p = reinterpret_cast<__nv_bfloat16 const*>(o); - auto const* pos_p = reinterpret_cast(positions); - auto const* cs_p = reinterpret_cast(cos_sin_cache); - auto* fp8_p = reinterpret_cast<__nv_fp8_e4m3*>(fp8_out); - auto* sc_p = reinterpret_cast(scale_out); - -#define TRTLLM_INV_ROPE_DISPATCH_CHUNK(CHUNKS, NEOX) \ - dispatchByM<(CHUNKS), (NEOX)>(num_tokens, num_heads, heads_per_group, scale_buf_m, o_p, pos_p, cs_p, fp8_p, sc_p, \ - o_stride_token, o_stride_head, fp8_stride_group, fp8_stride_token, scale_stride_group, scale_stride_k, stream) - - if (is_neox) - { - switch (chunks_per_head) - { - case 1: TRTLLM_INV_ROPE_DISPATCH_CHUNK(1, true); break; - case 2: TRTLLM_INV_ROPE_DISPATCH_CHUNK(2, true); break; - case 3: TRTLLM_INV_ROPE_DISPATCH_CHUNK(3, true); break; - case 4: TRTLLM_INV_ROPE_DISPATCH_CHUNK(4, true); break; - default: break; - } - } - else - { - switch (chunks_per_head) - { - case 1: TRTLLM_INV_ROPE_DISPATCH_CHUNK(1, false); break; - case 2: TRTLLM_INV_ROPE_DISPATCH_CHUNK(2, false); break; - case 3: TRTLLM_INV_ROPE_DISPATCH_CHUNK(3, false); break; - case 4: TRTLLM_INV_ROPE_DISPATCH_CHUNK(4, false); break; - default: break; - } - } - -#undef TRTLLM_INV_ROPE_DISPATCH_CHUNK -} - -} // namespace kernels - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/inverseRopeFp8QuantKernel.h b/cpp/tensorrt_llm/kernels/inverseRopeFp8QuantKernel.h deleted file mode 100644 index 91759d5472d0..000000000000 --- a/cpp/tensorrt_llm/kernels/inverseRopeFp8QuantKernel.h +++ /dev/null @@ -1,79 +0,0 @@ -/* - * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. - * - * 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. - */ -#pragma once - -#include "tensorrt_llm/common/config.h" - -#include -#include - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels -{ - -// Fused inverse-RoPE + 1x128 block-scaled FP8 quant for DeepSeek-V4 -// absorption-mode attention output. Reads bf16 attention output -// [num_tokens, num_heads, head_dim], applies NEOX inverse rotary embedding -// to the rope segment (the LAST 64 elements of each head, i.e. the second -// half of the final quant chunk), then per-128-element-chunk -// absmax-quantizes the result to FP8 e4m3. -// -// Layout constraints (asserted by caller; the kernel hardcodes the matching -// constants): -// * quant_group_size == 128 -// * rope_dim == 64 (half_rope == 32, NEOX) -// * head_dim == chunks_per_head * 128, with chunks_per_head in -// {1, 2, 3, 4}. Production DSv4-{Flash,Pro} use -// head_dim = 512 (kv_lora_rank=448 nope + 64 rope). -// * Per head: scales emitted for every quant chunk (so total scale slots -// per token = num_heads * chunks_per_head). -// -// Hardware: targets SM89+ (uses cvt.rn.satfinite.e4m3x2.f32, __shfl_sync, -// bf16 intrinsics). No TMA / cluster / tcgen05 dependencies. -// -// Layouts: -// o : bf16 [num_tokens, num_heads, head_dim] -// positions : int64 [num_tokens] -// cos_sin_cache: fp32 [max_positions, 2, 32] (cos block then sin block, NEOX) -// fp8_out : fp8 [..., num_tokens, num_heads * head_dim] -// scale_out : fp32 with stride scale_stride_k between adjacent quant -// blocks (qb index = head_idx * chunks_per_head + chunk), -// layout [..., num_heads * chunks_per_head, scale_buf_m] -// where scale_buf_m = pad_up(num_tokens, 4) per the BMM -// dequant consumer's hard-coded m-dim stride. -void invokeInverseRopeFp8Quant(void const* o, // - void const* positions, // - void const* cos_sin_cache, // - void* fp8_out, // - void* scale_out, // - int num_tokens, // - int num_heads, // - int heads_per_group, // - int chunks_per_head, // - bool is_neox, // - int scale_buf_m, // - int o_stride_token, // - int o_stride_head, // - int fp8_stride_group, // - int fp8_stride_token, // - int scale_stride_group, // - int scale_stride_k, // - cudaStream_t stream); - -} // namespace kernels - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/marlin/marlin.cuh b/cpp/tensorrt_llm/kernels/marlin/marlin.cuh deleted file mode 100644 index 63864080b0f2..000000000000 --- a/cpp/tensorrt_llm/kernels/marlin/marlin.cuh +++ /dev/null @@ -1,397 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2022-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. - */ - -// Internal device-side header for the Marlin NVFP4 kernels. - -#pragma once - -#ifndef _marlin_cuh -#define _marlin_cuh - -#include -#include -#include -#include -#include - -#include - -#ifndef MARLIN_NAMESPACE_NAME -#define MARLIN_NAMESPACE_NAME marlin -#endif - -namespace MARLIN_NAMESPACE_NAME -{ - -static constexpr int default_threads = 256; -static constexpr int pipe_stages = 4; - -static constexpr int min_thread_n = 64; -static constexpr int min_thread_k = 64; -static constexpr int max_thread_n = 256; - -static constexpr int tile_size = 16; -static constexpr int max_par = 16; - -static constexpr int repack_stages = 8; -static constexpr int repack_threads = 256; - -static constexpr int tile_k_size = tile_size; -static constexpr int tile_n_size = tile_k_size * 4; - -template -struct Vec -{ - T elems[n]; - - __device__ T& operator[](int i) - { - return elems[i]; - } -}; - -using I4 = Vec; - -constexpr int div_ceil(int a, int b) -{ - return (a + b - 1) / b; -} - -// cp.async wrappers (SM 7.x fallback / SM 8.x+ inline asm). - -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 - -__device__ inline void cp_async1_ca_pred(void* smem_ptr, void const* glob_ptr, bool pred = true) -{ - if (pred) - { - reinterpret_cast(smem_ptr)[0] = reinterpret_cast(glob_ptr)[0]; - } -} - -__device__ inline void cp_async2_ca_pred(void* smem_ptr, void const* glob_ptr, bool pred = true) -{ - if (pred) - { - reinterpret_cast(smem_ptr)[0] = reinterpret_cast(glob_ptr)[0]; - } -} - -__device__ inline void cp_async4_ca_pred(void* smem_ptr, void const* glob_ptr, bool pred = true) -{ - if (pred) - { - reinterpret_cast(smem_ptr)[0] = reinterpret_cast(glob_ptr)[0]; - } -} - -__device__ inline void cp_async4_pred(void* smem_ptr, void const* glob_ptr, bool pred = true) -{ - if (pred) - { - reinterpret_cast(smem_ptr)[0] = reinterpret_cast(glob_ptr)[0]; - } -} - -__device__ inline void cp_async4(void* smem_ptr, void const* glob_ptr) -{ - reinterpret_cast(smem_ptr)[0] = reinterpret_cast(glob_ptr)[0]; -} - -__device__ inline void cp_async_fence() {} - -template -__device__ inline void cp_async_wait() -{ -} - -#else - -__device__ inline void cp_async1_ca_pred(void* smem_ptr, void const* glob_ptr, bool pred = true) -{ - int const BYTES = 4; - uint32_t smem = static_cast(__cvta_generic_to_shared(smem_ptr)); - asm volatile( - "{\n" - " .reg .pred p;\n" - " setp.ne.b32 p, %0, 0;\n" - " @p cp.async.ca.shared.global [%1], [%2], %3;\n" - "}\n" ::"r"((int) pred), - "r"(smem), "l"(glob_ptr), "n"(BYTES)); -} - -__device__ inline void cp_async2_ca_pred(void* smem_ptr, void const* glob_ptr, bool pred = true) -{ - int const BYTES = 8; - uint32_t smem = static_cast(__cvta_generic_to_shared(smem_ptr)); - asm volatile( - "{\n" - " .reg .pred p;\n" - " setp.ne.b32 p, %0, 0;\n" - " @p cp.async.ca.shared.global [%1], [%2], %3;\n" - "}\n" ::"r"((int) pred), - "r"(smem), "l"(glob_ptr), "n"(BYTES)); -} - -__device__ inline void cp_async4_ca_pred(void* smem_ptr, void const* glob_ptr, bool pred = true) -{ - int const BYTES = 16; - uint32_t smem = static_cast(__cvta_generic_to_shared(smem_ptr)); - asm volatile( - "{\n" - " .reg .pred p;\n" - " setp.ne.b32 p, %0, 0;\n" - " @p cp.async.ca.shared.global [%1], [%2], %3;\n" - "}\n" ::"r"((int) pred), - "r"(smem), "l"(glob_ptr), "n"(BYTES)); -} - -__device__ inline void cp_async4_pred(void* smem_ptr, void const* glob_ptr, bool pred = true) -{ - int const BYTES = 16; - uint32_t smem = static_cast(__cvta_generic_to_shared(smem_ptr)); - asm volatile( - "{\n" - " .reg .pred p;\n" - " setp.ne.b32 p, %0, 0;\n" - " @p cp.async.cg.shared.global [%1], [%2], %3;\n" - "}\n" ::"r"((int) pred), - "r"(smem), "l"(glob_ptr), "n"(BYTES)); -} - -__device__ inline void cp_async4(void* smem_ptr, void const* glob_ptr) -{ - int const BYTES = 16; - uint32_t smem = static_cast(__cvta_generic_to_shared(smem_ptr)); - asm volatile( - "{\n" - " cp.async.cg.shared.global [%0], [%1], %2;\n" - "}\n" ::"r"(smem), - "l"(glob_ptr), "n"(BYTES)); -} - -__device__ inline void cp_async_fence() -{ - asm volatile("cp.async.commit_group;\n" ::); -} - -template -__device__ inline void cp_async_wait() -{ - asm volatile("cp.async.wait_group %0;\n" ::"n"(n)); -} - -#endif - -// MarlinType traits + fragment aliases. -// MMA m16n8k16 fragment layouts: -// https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#matrix-fragments-for-mma-m16n8k16-with-floating-point-type -template -struct MarlinType -{ -}; - -template <> -struct MarlinType -{ - using scalar_t = nv_bfloat16; - using scalar_t2 = nv_bfloat162; - using scalar_t4 = nv_bfloat162; - using scalar_32bit_t = nv_bfloat162; - - using FragA = Vec; - using FragB = Vec; - using FragC = Vec; - using FragS = Vec; - using FragS0 = Vec<__nv_fp8x2_e4m3, 1>; - using FragZP = Vec; - -#if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 800 - static __device__ float inline num2float(const nv_bfloat16 x) - { - return __bfloat162float(x); - } - - static __device__ nv_bfloat162 inline num2num2(const nv_bfloat16 x) - { - return __bfloat162bfloat162(x); - } - - static __device__ nv_bfloat162 inline nums2num2(const nv_bfloat16 x1, const nv_bfloat16 x2) - { - return __halves2bfloat162(x1, x2); - } - - static __host__ __device__ nv_bfloat16 inline float2num(float const x) - { - return __float2bfloat16(x); - } - - static __host__ __device__ float2 inline num22float2(const nv_bfloat162 x) - { - return __bfloat1622float2(x); - } -#endif -}; - -// Fast FP4 E2M1 -> BF16 and FP8 E4M3 -> BF16 dequantization. -// FP4->BF16 places the 3 FP4 bits into BF16's exponent/mantissa via bitwise -// ops; a subsequent multiply applies the exponent-bias correction (or -// ``skip_flop=true`` defers it to fuse with a scale multiply downstream). - -#if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 750 - -// Lookup-table based 3-input logical operation; the compiler does not always -// recognize the pattern automatically. -template -__device__ inline int lop3(int a, int b, int c) -{ - int res; - asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" : "=r"(res) : "r"(a), "r"(b), "r"(c), "n"(lut)); - return res; -} - -// FP4 E2M1 -> BF16 -// -// skip_flop=true: just place bits, caller multiplies exponent bias later. -template -__device__ inline void dequant_fp4(int q, nv_bfloat162* frag_b) -{ - // Constants for FP4 (E2M1) -> BF16 (E8M7) - constexpr int FP4_EXPONENT = 2, BF16_EXPONENT = 8; - constexpr int RIGHT_SHIFT = BF16_EXPONENT - FP4_EXPONENT; - constexpr int MASK = 0x70007000; - - // Extract and shift FP4 values to BF16 format - int Out1 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT); - q <<= 4; - int Out2 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT); - - // Note: reverse indexing is intentional because weights are permuted - frag_b[1] = *reinterpret_cast(&Out1); - frag_b[0] = *reinterpret_cast(&Out2); - - if constexpr (!skip_flop) - { - // Apply exponent bias correction - constexpr int BIAS_OFFSET = (1 << (BF16_EXPONENT - 1)) - (1 << (FP4_EXPONENT - 1)); - constexpr uint32_t BIAS = (BIAS_OFFSET + 127) << 23; - const nv_bfloat162 bias_reg = __float2bfloat162_rn(*reinterpret_cast(&BIAS)); - - frag_b[1] = __hmul2(frag_b[1], bias_reg); - frag_b[0] = __hmul2(frag_b[0], bias_reg); - } -} - -// FP8 E4M3 scale -> BF16 -__device__ inline void dequant_fp8_scales(int q, nv_bfloat162* frag_b) -{ - constexpr int FP8_EXPONENT = 4, BF16_EXPONENT = 8; - constexpr int RIGHT_SHIFT = BF16_EXPONENT - FP8_EXPONENT; - constexpr int MASK = 0x7F007F00; - - // Extract and shift FP8 values to BF16 format - int Out1 = ((q & 0x80008000) >> 1) | ((q & MASK) >> RIGHT_SHIFT); - q <<= 8; - int Out2 = ((q & 0x80008000) >> 1) | ((q & MASK) >> RIGHT_SHIFT); - - // Note: reverse indexing is intentional because weights are permuted - frag_b[1] = *reinterpret_cast(&Out1); - frag_b[0] = *reinterpret_cast(&Out2); -} - -#endif // __CUDA_ARCH__ >= 750 - -// m16n8k16 tensor-core MMA: BF16 inputs, FP32 accumulation. -template -__device__ inline void mma(const typename MarlinType::FragA& a_frag, - const typename MarlinType::FragB& frag_b, typename MarlinType::FragC& frag_c) -{ - uint32_t const* a = reinterpret_cast(&a_frag); - uint32_t const* b = reinterpret_cast(&frag_b); - - static_assert(std::is_same::value, "Only BF16 is supported for Marlin NVFP4 MMA"); - - float* c = reinterpret_cast(&frag_c); - asm volatile( - "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 " - "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n" - : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3]) - : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]), "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3])); -} - -// Transposed variant: column-major B weight loading. -template -__device__ inline void mma_trans(const typename MarlinType::FragA& a_frag, - const typename MarlinType::FragB& frag_b, const typename MarlinType::FragB& frag_b2, - typename MarlinType::FragC& frag_c) -{ - uint32_t const* a = reinterpret_cast(&a_frag); - uint32_t const* b = reinterpret_cast(&frag_b); - uint32_t const* b2 = reinterpret_cast(&frag_b2); - - static_assert(std::is_same::value, "Only BF16 is supported for Marlin NVFP4 MMA"); - - float* c = reinterpret_cast(&frag_c); - asm volatile( - "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 " - "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n" - : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3]) - : "r"(b[0]), "r"(b2[0]), "r"(b[1]), "r"(b2[1]), "r"(a[0]), "r"(a[1]), "f"(c[0]), "f"(c[1]), "f"(c[2]), - "f"(c[3])); -} - -} // namespace MARLIN_NAMESPACE_NAME - -// Single-expert kernel forward decl. Opt in with -// ``#define MARLIN_DECLARE_SINGLE_EXPERT_KERNEL`` before including. The MoE -// TU does NOT define it (the MoE kernel has a different parameter list, in -// marlin_nvfp4_moe_template.h). -#ifdef MARLIN_DECLARE_SINGLE_EXPERT_KERNEL - -#define MARLIN_KERNEL_PARAMS \ - const int4 *__restrict__ A, const int4 *__restrict__ B, int4 *__restrict__ C, int4 *__restrict__ C_tmp, \ - const int4 *__restrict__ b_bias_ptr, const float *__restrict__ a_scales_ptr, \ - const int4 *__restrict__ scales_ptr, const uint16_t *__restrict__ global_scale_ptr, \ - const int4 *__restrict__ zp_ptr, const int *__restrict__ g_idx, int num_groups, int prob_m, int prob_n, \ - int prob_k, int lda, int *locks, bool has_bias, bool use_atomic_add, bool use_fp32_reduce, int max_shared_mem - -namespace MARLIN_NAMESPACE_NAME -{ - -// clang-format off -// NOTE: keep this template parameter list out of clang-format. East-const -// (QualifierAlignment: Right) miscompiles the *last* NTTP as `int X const`, -// which is invalid syntax. Non-type template parameters are implicitly -// const anyway, so we omit it on the last entry. -template -__global__ void Marlin(MARLIN_KERNEL_PARAMS); -// clang-format on - -} // namespace MARLIN_NAMESPACE_NAME - -#endif // MARLIN_DECLARE_SINGLE_EXPERT_KERNEL - -#endif // _marlin_cuh diff --git a/cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4.h b/cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4.h deleted file mode 100644 index ad7d55471514..000000000000 --- a/cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4.h +++ /dev/null @@ -1,73 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2022-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. - */ - -// Host-side header for the Marlin NVFP4 W4A16 kernels. - -#pragma once - -#include -#include - -namespace marlin_nvfp4 -{ - -void dequantFp4Activations( - void const* act_fp4, void const* act_sf, float const* alpha, void* act_bf16, int m, int k, cudaStream_t stream); - -void marlinNvfp4Gemm(void const* act_bf16, void const* weight, void* output, void* C_tmp, void const* weight_sf, - void const* global_scale_bf16, int m, int n, int k, int* workspace, int num_groups, int group_size, - bool use_fp32_reduce, cudaStream_t stream); - -void marlinNvfp4MoeGemmDispatcher(void const* A, void const* B, void* C, void* C_tmp, void const* b_scales, - void const* global_scale, void const* sorted_token_ids, void const* expert_ids, void const* num_tokens_past_padded, - void const* topk_weights, int moe_block_size, int top_k, bool mul_topk_weights, int prob_m, int prob_n, int prob_k, - void* workspace, int num_groups, int group_size, bool use_fp32_reduce, bool use_atomic_add, cudaDataType_t outType, - cudaStream_t stream); - -void gptq_marlin_repack_dispatch(uint32_t const* b_q_weight_ptr, uint32_t const* perm_ptr, uint32_t* out_ptr, - int size_k, int size_n, int num_bits, bool has_perm, bool is_a_8bit, cudaStream_t stream); - -} // namespace marlin_nvfp4 - -namespace marlin_nvfp4_dispatch -{ - -struct thread_config_t -{ - int thread_k; - int thread_n; - int num_threads; -}; - -struct exec_config_t -{ - int blocks_per_sm; - thread_config_t tb_cfg; -}; - -extern thread_config_t const kSmallBatchConfigs[]; -extern thread_config_t const kLargeBatchConfigs[]; - -extern int const kSmallBatchConfigCount; -extern int const kLargeBatchConfigCount; - -int get_scales_cache_size( - thread_config_t const& th_config, int prob_n, int prob_k, int num_bits, int group_size, int stages); - -bool is_config_feasible(thread_config_t const& cfg, int prob_n, int prob_k); - -} // namespace marlin_nvfp4_dispatch diff --git a/cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_dispatch_utils.cpp b/cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_dispatch_utils.cpp deleted file mode 100644 index 9b3dac04ed25..000000000000 --- a/cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_dispatch_utils.cpp +++ /dev/null @@ -1,58 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2022-2025 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. - */ - -#include "marlin_nvfp4.h" - -namespace marlin_nvfp4_dispatch -{ - -thread_config_t const kSmallBatchConfigs[] = {{128, 128, 256}, {64, 128, 128}, {128, 64, 128}}; -thread_config_t const kLargeBatchConfigs[] = {{64, 256, 256}, {64, 128, 128}, {128, 64, 128}}; - -int const kSmallBatchConfigCount = sizeof(kSmallBatchConfigs) / sizeof(thread_config_t); -int const kLargeBatchConfigCount = sizeof(kLargeBatchConfigs) / sizeof(thread_config_t); - -int get_scales_cache_size( - thread_config_t const& th_config, int prob_n, int prob_k, int num_bits, int group_size, int stages) -{ - int tb_n = th_config.thread_n; - int tb_k = th_config.thread_k; - int tb_groups; - if (group_size == -1) - tb_groups = 1; - else if (group_size == 0) - tb_groups = (tb_k + 31) / 32; // div_ceil(tb_k, 32) - else - tb_groups = (tb_k + group_size - 1) / group_size; // div_ceil(tb_k, group_size) - return tb_groups * tb_n * 2 * stages; -} - -bool is_config_feasible(thread_config_t const& cfg, int prob_n, int prob_k) -{ - if (cfg.thread_k == -1 || cfg.thread_n == -1 || cfg.num_threads == -1) - return false; - if (prob_k % cfg.thread_k != 0 || prob_n % cfg.thread_n != 0) - return false; - // min_thread_n = 64, min_thread_k = 64 (from marlin.cuh constants) - if (cfg.thread_n < 64 || cfg.thread_k < 64) - return false; - if (cfg.num_threads < 128) - return false; - return true; -} - -} // namespace marlin_nvfp4_dispatch diff --git a/cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_gemm.cu b/cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_gemm.cu deleted file mode 100644 index f9cbbe3d64ab..000000000000 --- a/cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_gemm.cu +++ /dev/null @@ -1,340 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2022-2025 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. - */ - -#ifndef MARLIN_NAMESPACE_NAME -#define MARLIN_NAMESPACE_NAME marlin -#endif - -#define MARLIN_DECLARE_SINGLE_EXPERT_KERNEL -#include "marlin_nvfp4.h" -#include "marlin_nvfp4_template.h" -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/cudaUtils.h" - -#include -#include -#include -#include - -namespace marlin -{ - -using namespace marlin_nvfp4_dispatch; - -__global__ void MarlinDefault(MARLIN_KERNEL_PARAMS){}; - -using MarlinFuncPtr = void (*)(MARLIN_KERNEL_PARAMS); - -// Single-expert shared-memory size (no block-meta overhead). -int get_kernel_cache_size(thread_config_t const& th_config, int thread_m_blocks, int prob_n, int prob_k, int num_bits, - int group_size, int stages) -{ - int pack_factor = 32 / num_bits; - int tb_k = th_config.thread_k; - int tb_n = th_config.thread_n; - int tb_m = thread_m_blocks * 16; - int sh_a_size = stages * (tb_m * tb_k) * 2; - int sh_b_size = stages * (tb_k * tb_n / pack_factor) * 4; - int sh_red_size = tb_m * (tb_n + 8) * 2; - int sh_bias_size = tb_n * 2; - int tmp_size = (sh_b_size > sh_red_size ? sh_red_size : sh_b_size) + sh_bias_size; - tmp_size = std::max(std::max(sh_b_size, sh_red_size), tmp_size); - int sh_s_size = get_scales_cache_size(th_config, prob_n, prob_k, num_bits, group_size, stages); - return tmp_size + sh_a_size + sh_s_size; -} - -bool is_valid_config(thread_config_t const& th_config, int thread_m_blocks, int prob_n, int prob_k, int num_bits, - int group_size, int stages, int max_shared_mem) -{ - if (!is_config_feasible(th_config, prob_n, prob_k)) - return false; - return get_kernel_cache_size(th_config, thread_m_blocks, prob_n, prob_k, num_bits, group_size, stages) - <= max_shared_mem; -} - -MarlinFuncPtr get_marlin_kernel(int thread_m_blocks, int thread_n_blocks, int thread_k_blocks, bool m_block_size_8, - int group_blocks, int threads, int stages) -{ -#define MARLIN_KERNEL_MATCH(T, M, N, K, M8) \ - (threads == (T) && thread_m_blocks == (M) && thread_n_blocks == (N) && thread_k_blocks == (K) \ - && m_block_size_8 == (M8) && stages == 4 && group_blocks == 1) -#define MARLIN_KERNEL_IF(T, M, N, K, M8) \ - if (MARLIN_KERNEL_MATCH(T, M, N, K, M8)) \ - return Marlin; - - MARLIN_KERNEL_IF(256, 1, 8, 8, true) - MARLIN_KERNEL_IF(128, 1, 8, 4, true) - MARLIN_KERNEL_IF(128, 1, 4, 8, true) - MARLIN_KERNEL_IF(256, 1, 8, 8, false) - MARLIN_KERNEL_IF(128, 1, 8, 4, false) - MARLIN_KERNEL_IF(128, 1, 4, 8, false) - MARLIN_KERNEL_IF(256, 2, 16, 4, false) - MARLIN_KERNEL_IF(128, 2, 8, 4, false) - MARLIN_KERNEL_IF(128, 2, 4, 8, false) - MARLIN_KERNEL_IF(256, 3, 16, 4, false) - MARLIN_KERNEL_IF(128, 3, 8, 4, false) - MARLIN_KERNEL_IF(128, 3, 4, 8, false) - MARLIN_KERNEL_IF(256, 4, 16, 4, false) - MARLIN_KERNEL_IF(128, 4, 8, 4, false) - MARLIN_KERNEL_IF(128, 4, 4, 8, false) - -#undef MARLIN_KERNEL_MATCH -#undef MARLIN_KERNEL_IF - return MarlinDefault; -} - -exec_config_t determine_exec_config(int prob_m, int prob_n, int prob_k, int thread_m_blocks, bool m_block_size_8, - int num_bits, int group_size, int stages, int max_shared_mem, int sms) -{ - exec_config_t exec_cfg{1, {-1, -1, -1}}; - thread_config_t const* cfgs = thread_m_blocks > 1 ? kLargeBatchConfigs : kSmallBatchConfigs; - int cfg_count = thread_m_blocks > 1 ? kLargeBatchConfigCount : kSmallBatchConfigCount; - - for (int i = 0; i < cfg_count; i++) - { - thread_config_t th = cfgs[i]; - if (!is_valid_config(th, thread_m_blocks, prob_n, prob_k, num_bits, group_size, stages, max_shared_mem - 512)) - continue; - int group_blocks = group_size == -1 ? -1 : group_size / 16; - auto kernel = get_marlin_kernel( - thread_m_blocks, th.thread_n / 16, th.thread_k / 16, m_block_size_8, group_blocks, th.num_threads, stages); - if (kernel == MarlinDefault) - continue; - return {1, th}; - } - return exec_cfg; -} - -void marlin_mm_nvfp4(void const* A, void const* B, void* C, void* C_tmp, void const* b_s, void const* g_s, int prob_m, - int prob_n, int prob_k, int* locks, int num_groups, int group_size, bool use_fp32_reduce, int dev, - cudaStream_t stream) -{ - constexpr int num_bits = 4; - - int group_blocks = group_size == -1 ? -1 : group_size / 16; - - int4 const* A_ptr = (int4 const*) A; - int4 const* B_ptr = (int4 const*) B; - int4* C_ptr = (int4*) C; - int4* C_tmp_ptr = (int4*) C_tmp; - int4 const* b_s_ptr = (int4 const*) b_s; - uint16_t const* g_s_ptr = (uint16_t const*) g_s; - - int max_shared_mem = 0; - cudaDeviceGetAttribute(&max_shared_mem, cudaDevAttrMaxSharedMemoryPerBlockOptin, dev); - - int stages = 4; - int sms = -1; - cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, dev); - - int max_par_val = 16; - if (prob_n <= 4096) - max_par_val = 16 * 8; - int max_shared_mem_new = max_shared_mem; - int rest_m = prob_m; - int max_thread_m_blocks = 4; - int lda = prob_k; - - while (rest_m) - { - int par_count = std::min(rest_m / (max_thread_m_blocks * 16), max_par_val); - int prob_m_split = par_count > 0 ? (par_count * (max_thread_m_blocks * 16)) : rest_m; - - int thread_m_blocks = std::min(div_ceil(prob_m_split, 16), max_thread_m_blocks); - bool m_block_size_8 = prob_m_split <= 8; - - exec_config_t exec_cfg = determine_exec_config(prob_m_split, prob_n, prob_k, thread_m_blocks, m_block_size_8, - num_bits, group_size, stages, max_shared_mem, sms); - thread_config_t thread_tfg = exec_cfg.tb_cfg; - - if (thread_tfg.thread_k == -1 && max_thread_m_blocks > 1) - { - max_thread_m_blocks--; - continue; - } - - if (thread_tfg.thread_k == -1) - { - break; - } - - // Small wave optimization - if (thread_tfg.thread_n != -1) - { - if (prob_n / thread_tfg.thread_n * div_ceil(prob_m_split, thread_m_blocks * 16) * 4 <= sms) - { - if (is_valid_config({128, 64, 128}, thread_m_blocks, prob_n, prob_k, num_bits, group_size, stages, - max_shared_mem_new)) - { - thread_tfg = {128, 64, 128}; - exec_cfg = {1, thread_tfg}; - } - } - } - - int num_threads = thread_tfg.num_threads; - int thread_k = thread_tfg.thread_k; - int thread_n = thread_tfg.thread_n; - int blocks = sms * exec_cfg.blocks_per_sm; - if (exec_cfg.blocks_per_sm > 1) - max_shared_mem_new = max_shared_mem / exec_cfg.blocks_per_sm - 1024; - - int thread_k_blocks = thread_k / 16; - int thread_n_blocks = thread_n / 16; - - auto kernel = get_marlin_kernel( - thread_m_blocks, thread_n_blocks, thread_k_blocks, m_block_size_8, group_blocks, num_threads, stages); - - if (kernel == MarlinDefault) - { - break; - } - - cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, max_shared_mem_new); - - // clang-format off - kernel<<>>( - A_ptr, B_ptr, C_ptr, C_tmp_ptr, - nullptr, // b_bias - nullptr, // a_scales - b_s_ptr, g_s_ptr, - nullptr, // zp - nullptr, // g_idx - num_groups, prob_m_split, prob_n, prob_k, lda, locks, - false, // has_bias - false, // use_atomic_add - use_fp32_reduce, max_shared_mem_new); - // clang-format on - - A_ptr += prob_m_split * (lda / 8); - C_ptr += prob_m_split * (prob_n / 8); - rest_m -= prob_m_split; - } -} - -// Explicit template instantiations for BF16 + NVFP4 Marlin kernels. -// clang-format off -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -// clang-format on - -// FP4 E2M1 -> BF16 activation dequant: out[i] = fp4_to_bf16(act[i]) * -// block_scale[i/16] * global_scale. Block scales are FP8 E4M3 (swizzled). -__global__ void dequant_fp4_act_kernel(uint8_t const* __restrict__ act_fp4, // [M, K/2] packed FP4 - uint8_t const* __restrict__ act_sf, // FP8 E4M3 block scales (swizzled) - float const* __restrict__ alpha, // global scale - nv_bfloat16* __restrict__ out, // [M, K] BF16 output - int M, int K) -{ - - int idx = blockIdx.x * blockDim.x + threadIdx.x; - int total_pairs = M * (K / 2); - if (idx >= total_pairs) - return; - - float global_s = *alpha; - int row = idx / (K / 2); - int col_pair = idx % (K / 2); - - uint8_t packed = act_fp4[idx]; - - // Unpack two FP4 E2M1 values (low nibble first) - auto fp4_to_float = [](uint8_t nibble) -> float - { - // FP4 E2M1: 1 sign + 2 exponent + 1 mantissa - uint8_t sign = (nibble >> 3) & 1; - uint8_t exp = (nibble >> 1) & 0x3; - uint8_t mant = nibble & 1; - float val; - if (exp == 0) - { - // subnormal: (-1)^s * 0.mantissa * 2^(1-bias) = (-1)^s * mant * 0.5 - val = mant * 0.5f; - } - else - { - // normal: (-1)^s * 1.mantissa * 2^(exp-bias), bias=1 - val = (1.0f + mant * 0.5f) * (float) (1 << (exp - 1)); - } - return sign ? -val : val; - }; - - float v0 = fp4_to_float(packed & 0x0F); - float v1 = fp4_to_float((packed >> 4) & 0x0F); - - // Block scale: one FP8 E4M3 per 16 FP4 elements = per 8 bytes - // The scale layout is swizzled 128x4 — for now use linear indexing - // as a reasonable approximation. TODO: handle swizzled layout properly. - int elem0 = col_pair * 2; - int scale_idx = row * (K / 16) + elem0 / 16; - uint8_t sf_byte = act_sf[scale_idx]; - // FP8 E4M3 -> float: reinterpret as __nv_fp8_e4m3 - __nv_fp8_e4m3 sf_fp8 = *reinterpret_cast<__nv_fp8_e4m3 const*>(&sf_byte); - float sf = float(sf_fp8); - - float scale = sf * global_s; - int out_idx = row * K + col_pair * 2; - out[out_idx] = __float2bfloat16(v0 * scale); - out[out_idx + 1] = __float2bfloat16(v1 * scale); -} - -} // namespace marlin - -namespace marlin_nvfp4 -{ - -void dequantFp4Activations( - void const* act_fp4, void const* act_sf, float const* alpha, void* act_bf16, int m, int k, cudaStream_t stream) -{ - - int total_pairs = m * (k / 2); - int threads = 256; - int blocks = (total_pairs + threads - 1) / threads; - ::marlin::dequant_fp4_act_kernel<<>>( - (uint8_t const*) act_fp4, (uint8_t const*) act_sf, alpha, (nv_bfloat16*) act_bf16, m, k); -} - -void marlinNvfp4Gemm(void const* act_bf16, void const* weight, void* output, void* C_tmp, void const* weight_sf, - void const* global_scale_bf16, int m, int n, int k, int* workspace, int num_groups, int group_size, - bool use_fp32_reduce, cudaStream_t stream) -{ - int const sm = tensorrt_llm::common::getSMVersion(); - TLLM_CHECK_WITH_INFO( - sm >= 90 && sm < 100, "Marlin NVFP4 GEMM is only supported on Hopper (SM 9.x); current SM = %d", sm); - - int dev; - cudaGetDevice(&dev); - - ::marlin::marlin_mm_nvfp4(act_bf16, weight, output, C_tmp, weight_sf, global_scale_bf16, m, n, k, workspace, - num_groups, group_size, use_fp32_reduce, dev, stream); -} - -} // namespace marlin_nvfp4 diff --git a/cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_moe_gemm.cu b/cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_moe_gemm.cu deleted file mode 100644 index 35f28b309345..000000000000 --- a/cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_moe_gemm.cu +++ /dev/null @@ -1,269 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2022-2025 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. - */ - -#ifndef MARLIN_NAMESPACE_NAME -#define MARLIN_NAMESPACE_NAME marlin_moe_wna16 -#endif - -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/cudaUtils.h" -#include "tensorrt_llm/common/logger.h" - -#include "marlin_nvfp4.h" -#include "marlin_nvfp4_moe_template.h" - -#include -#include - -namespace marlin_moe_wna16 -{ - -using namespace marlin_nvfp4_dispatch; - -__global__ void MarlinDefault(MARLIN_KERNEL_PARAMS){}; - -using MarlinFuncPtr = void (*)(MARLIN_KERNEL_PARAMS); - -// MoE shared-memory size includes block-meta overhead for sorted_token_ids. -int get_kernel_cache_size(thread_config_t const& th_config, int thread_m_blocks, int prob_n, int prob_k, int num_bits, - int group_size, int stages) -{ - int pack_factor = 32 / num_bits; - int tb_k = th_config.thread_k; - int tb_n = th_config.thread_n; - int tb_m = thread_m_blocks * 16; - int sh_block_meta_size = tb_m * 16; - int sh_a_size = stages * (tb_m * tb_k) * 2; - int sh_b_size = stages * (tb_k * tb_n / pack_factor) * 4; - int sh_red_size = tb_m * (tb_n + 8) * 2; - int sh_bias_size = tb_n * 2; - int tmp_size = (sh_b_size > sh_red_size ? sh_red_size : sh_b_size) + sh_bias_size; - tmp_size = std::max(std::max(sh_b_size, sh_red_size), tmp_size); - int sh_s_size = get_scales_cache_size(th_config, prob_n, prob_k, num_bits, group_size, stages); - return tmp_size + sh_a_size + sh_s_size + sh_block_meta_size; -} - -bool is_valid_config(thread_config_t const& th_config, int thread_m_blocks, int prob_n, int prob_k, int num_bits, - int group_size, int stages, int max_shared_mem) -{ - if (!is_config_feasible(th_config, prob_n, prob_k)) - return false; - return get_kernel_cache_size(th_config, thread_m_blocks, prob_n, prob_k, num_bits, group_size, stages) - <= max_shared_mem; -} - -MarlinFuncPtr get_marlin_kernel(int thread_m_blocks, int thread_n_blocks, int thread_k_blocks, bool m_block_size_8, - int group_blocks, int threads, int stages) -{ - -#define MARLIN_KERNEL_MATCH(T, M, N, K, M8) \ - (threads == (T) && thread_m_blocks == (M) && thread_n_blocks == (N) && thread_k_blocks == (K) \ - && m_block_size_8 == (M8) && stages == 4 && group_blocks == 1) -#define MARLIN_KERNEL_IF(T, M, N, K, M8) \ - if (MARLIN_KERNEL_MATCH(T, M, N, K, M8)) \ - return Marlin; - - MARLIN_KERNEL_IF(256, 1, 8, 8, true) - MARLIN_KERNEL_IF(128, 1, 8, 4, true) - MARLIN_KERNEL_IF(128, 1, 4, 8, true) - MARLIN_KERNEL_IF(256, 1, 8, 8, false) - MARLIN_KERNEL_IF(128, 1, 8, 4, false) - MARLIN_KERNEL_IF(128, 1, 4, 8, false) - MARLIN_KERNEL_IF(256, 2, 16, 4, false) - MARLIN_KERNEL_IF(128, 2, 8, 4, false) - MARLIN_KERNEL_IF(128, 2, 4, 8, false) - MARLIN_KERNEL_IF(256, 3, 16, 4, false) - MARLIN_KERNEL_IF(128, 3, 8, 4, false) - MARLIN_KERNEL_IF(128, 3, 4, 8, false) - MARLIN_KERNEL_IF(256, 4, 16, 4, false) - MARLIN_KERNEL_IF(128, 4, 8, 4, false) - MARLIN_KERNEL_IF(128, 4, 4, 8, false) - -#undef MARLIN_KERNEL_MATCH -#undef MARLIN_KERNEL_IF - return MarlinDefault; -} - -// MoE config selection with occupancy-based multi-block logic. -exec_config_t determine_exec_config(int prob_m, int prob_n, int prob_k, int num_experts, int top_k, int thread_m_blocks, - bool m_block_size_8, int num_bits, int group_size, int stages, int max_shared_mem, int sms) -{ - exec_config_t exec_cfg{1, {-1, -1, -1}}; - thread_config_t const* cfgs = thread_m_blocks > 1 ? kLargeBatchConfigs : kSmallBatchConfigs; - int cfg_count = thread_m_blocks > 1 ? kLargeBatchConfigCount : kSmallBatchConfigCount; - - int count = 0; - constexpr int device_max_reg_size = 255 * 1024; - int group_blocks = group_size == -1 ? -1 : (group_size / 16); - - for (int i = 0; i < cfg_count; i++) - { - thread_config_t th = cfgs[i]; - if (!is_valid_config(th, thread_m_blocks, prob_n, prob_k, num_bits, group_size, stages, max_shared_mem - 512)) - continue; - - int cache_size = get_kernel_cache_size(th, thread_m_blocks, prob_n, prob_k, num_bits, group_size, stages); - - auto kernel = get_marlin_kernel( - thread_m_blocks, th.thread_n / 16, th.thread_k / 16, m_block_size_8, group_blocks, th.num_threads, stages); - if (kernel == MarlinDefault) - continue; - - cudaFuncAttributes attr; - cudaFuncGetAttributes(&attr, kernel); - int reg_size = std::max(attr.numRegs, 1) * th.num_threads * 4; - int allow_count = std::min(device_max_reg_size / reg_size, max_shared_mem / (cache_size + 1536)); - if (thread_m_blocks == 1) - allow_count = std::max(std::min(allow_count, 4), 1); - else - allow_count = std::max(std::min(allow_count, 2), 1); - - if (prob_n / th.thread_n * prob_m * top_k * 4 < sms * allow_count) - allow_count = std::max(prob_n / th.thread_n * prob_m * top_k * 4 / sms, 1); - - if (allow_count > count) - { - count = allow_count; - exec_cfg = {count, th}; - }; - } - return exec_cfg; -} - -void marlin_mm_moe_nvfp4(void const* A, void const* B, void* C, void* C_tmp, void const* b_s, void const* g_s, - void const* sorted_token_ids, void const* expert_ids, void const* num_tokens_past_padded, void const* topk_weights, - int moe_block_size, int num_experts, int top_k, bool mul_topk_weights, int prob_m, int prob_n, int prob_k, - int* locks, int num_groups, int group_size, bool use_fp32_reduce, bool use_atomic_add, int dev, cudaStream_t stream) -{ - - constexpr int num_bits = 4; - - int thread_m_blocks = div_ceil(moe_block_size, 16); - bool m_block_size_8 = moe_block_size == 8; - int group_blocks = group_size == -1 ? -1 : group_size / 16; - - int max_shared_mem = 0; - cudaDeviceGetAttribute(&max_shared_mem, cudaDevAttrMaxSharedMemoryPerBlockOptin, dev); - int stages = 4; - int sms = -1; - cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, dev); - - exec_config_t exec_cfg = determine_exec_config(prob_m, prob_n, prob_k, num_experts, top_k, thread_m_blocks, - m_block_size_8, num_bits, group_size, stages, max_shared_mem, sms); - thread_config_t thread_tfg = exec_cfg.tb_cfg; - - if (thread_tfg.thread_k == -1) - return; - - int num_threads = thread_tfg.num_threads; - int thread_k = thread_tfg.thread_k; - int thread_n = thread_tfg.thread_n; - int blocks = sms * exec_cfg.blocks_per_sm; - if (exec_cfg.blocks_per_sm > 1) - max_shared_mem = max_shared_mem / exec_cfg.blocks_per_sm - 1024; - - int thread_k_blocks = thread_k / 16; - int thread_n_blocks = thread_n / 16; - - auto kernel = get_marlin_kernel( - thread_m_blocks, thread_n_blocks, thread_k_blocks, m_block_size_8, group_blocks, num_threads, stages); - - if (kernel == MarlinDefault) - { - TLLM_LOG_ERROR( - "xuantengh debug error: kernel is MarlinDefault, cannot find corresponding instantiated kernel for threads " - "= %d, " - "thread_n_blocks = %d, thread_k_blocks = %d, m_block_size_8 = %d, group_blocks = %d, num_threads = %d, " - "stages = %d", - num_threads, thread_n_blocks, thread_k_blocks, m_block_size_8, group_blocks, num_threads, stages); - return; - } - - cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, max_shared_mem); - - int4 const* A_ptr = (int4 const*) A; - int4 const* B_ptr = (int4 const*) B; - int4* C_ptr = (int4*) C; - int4* C_tmp_ptr = (int4*) C_tmp; - int4 const* b_s_ptr = (int4 const*) b_s; - uint16_t const* g_s_ptr = (uint16_t const*) g_s; - int32_t const* sorted_token_ids_ptr = (int32_t const*) sorted_token_ids; - int32_t const* expert_ids_ptr = (int32_t const*) expert_ids; - int32_t const* num_tokens_past_padded_ptr = (int32_t const*) num_tokens_past_padded; - float const* topk_weights_ptr = (float const*) topk_weights; - - // clang-format off - kernel<<>>( - A_ptr, B_ptr, C_ptr, C_tmp_ptr, - nullptr, // b_bias - nullptr, // a_scales - b_s_ptr, g_s_ptr, - nullptr, // zp - nullptr, // g_idx - sorted_token_ids_ptr, expert_ids_ptr, num_tokens_past_padded_ptr, - topk_weights_ptr, top_k, mul_topk_weights, num_groups, - prob_m, prob_n, prob_k, locks, - false, // has_bias - use_atomic_add, use_fp32_reduce); - // clang-format on -} - -// Explicit template instantiations for BF16 + NVFP4 MoE Marlin kernels. -// clang-format off -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -// clang-format on - -} // namespace marlin_moe_wna16 - -namespace marlin_nvfp4 -{ - -void marlinNvfp4MoeGemmDispatcher(void const* A, void const* B, void* C, void* C_tmp, void const* b_scales, - void const* global_scale, void const* sorted_token_ids, void const* expert_ids, void const* num_tokens_past_padded, - void const* topk_weights, int moe_block_size, int top_k, bool mul_topk_weights, int prob_m, int prob_n, int prob_k, - void* workspace, int num_groups, int group_size, bool use_fp32_reduce, bool use_atomic_add, cudaDataType_t outType, - cudaStream_t stream) -{ - int const sm = tensorrt_llm::common::getSMVersion(); - TLLM_CHECK_WITH_INFO( - sm >= 90 && sm < 100, "Marlin NVFP4 MoE GEMM is only supported on Hopper (SM 9.x); current SM = %d", sm); - - int dev; - cudaGetDevice(&dev); - - int num_experts = 1; // Not used in kernel dispatch, only in config selection - - ::marlin_moe_wna16::marlin_mm_moe_nvfp4(A, B, C, C_tmp, b_scales, global_scale, sorted_token_ids, expert_ids, - num_tokens_past_padded, topk_weights, moe_block_size, num_experts, top_k, mul_topk_weights, prob_m, prob_n, - prob_k, (int*) workspace, num_groups, group_size, use_fp32_reduce, use_atomic_add, dev, stream); -} - -} // namespace marlin_nvfp4 diff --git a/cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_moe_template.h b/cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_moe_template.h deleted file mode 100644 index 9159e0feea61..000000000000 --- a/cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_moe_template.h +++ /dev/null @@ -1,2175 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2022-2025 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. - */ - -/* - * Modified by Neural Magic - * Copyright (C) Marlin.2024 Elias Frantar - * - * 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. - */ - -/* - * Adapted from https://github.com/IST-DASLab/marlin - */ - -#ifndef MARLIN_NAMESPACE_NAME -#define MARLIN_NAMESPACE_NAME marlin_moe_wna16 -#endif - -#include "marlin.cuh" - -#define MARLIN_KERNEL_PARAMS \ - const int4 *__restrict__ A, const int4 *__restrict__ B, int4 *__restrict__ C, int4 *__restrict__ C_tmp, \ - const int4 *__restrict__ b_bias_ptr, const float *__restrict__ a_scales_ptr, \ - const int4 *__restrict__ scales_ptr, const uint16_t *__restrict__ global_scale_ptr, \ - const int4 *__restrict__ zp_ptr, const int *__restrict__ g_idx, \ - const int32_t *__restrict__ sorted_token_ids_ptr, const int32_t *__restrict__ expert_ids_ptr, \ - const int32_t *__restrict__ num_tokens_past_padded_ptr, const float *__restrict__ topk_weights_ptr, int top_k, \ - bool mul_topk_weights, int num_groups, int prob_m, int prob_n, int prob_k, int *locks, bool has_bias, \ - bool use_atomic_add, bool use_fp32_reduce - -namespace MARLIN_NAMESPACE_NAME -{ - -template -__global__ void Marlin(MARLIN_KERNEL_PARAMS); - -#define STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t) \ - static_assert(std::is_same::value || std::is_same::value, \ - "only float16 and bfloat16 is supported"); - -// Empty kernel stub for non-Hopper device passes; see marlin.cuh. -#if defined(__CUDA_ARCH__) && !(__CUDA_ARCH__ >= 900 && __CUDA_ARCH__ < 1000) - -template -__global__ void Marlin(MARLIN_KERNEL_PARAMS) -{ -} - -#else - -// Instruction for loading a full 16x16 matrix fragment of operand A from shared -// memory, directly in tensor core layout. -template -__device__ inline void ldsm(typename MarlinType::FragA& frag_a, void const* smem_ptr) -{ - uint32_t* a = reinterpret_cast(&frag_a); - uint32_t smem = static_cast(__cvta_generic_to_shared(smem_ptr)); - if constexpr (count == 4) - { - asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];\n" - : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) - : "r"(smem)); - } - else if constexpr (count == 2) - { - asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0,%1}, [%2];\n" : "=r"(a[0]), "=r"(a[1]) : "r"(smem)); - } - else if constexpr (count == 1) - { - asm volatile("ldmatrix.sync.aligned.m8n8.x1.shared.b16 {%0}, [%1];\n" : "=r"(a[0]) : "r"(smem)); - } - else - { - static_assert(count == 1 || count == 2 || count == 4, "invalid count"); - } -} - -// Multiply dequantized values by the corresponding quantization scale; used -// only for grouped quantization. -template -__device__ inline void scale( - typename MarlinType::FragB& frag_b, typename MarlinType::FragS& frag_s, int i) -{ - using scalar_t2 = typename MarlinType::scalar_t2; - scalar_t2 s = MarlinType::num2num2(reinterpret_cast(&frag_s)[i]); - frag_b[0] = __hmul2(frag_b[0], s); - frag_b[1] = __hmul2(frag_b[1], s); -} - -template -__device__ inline void scale_and_sub(typename MarlinType::FragB& frag_b, scalar_t s, scalar_t zp) -{ - using scalar_t2 = typename MarlinType::scalar_t2; - scalar_t2 s2 = MarlinType::num2num2(s); - scalar_t2 zp2 = MarlinType::num2num2(zp); - frag_b[0] = __hfma2(frag_b[0], s2, __hneg2(zp2)); - frag_b[1] = __hfma2(frag_b[1], s2, __hneg2(zp2)); -} - -template -__device__ inline void sub_zp( - typename MarlinType::FragB& frag_b, typename MarlinType::scalar_t2& frag_zp, int i) -{ - using scalar_t2 = typename MarlinType::scalar_t2; - scalar_t2 zp = MarlinType::num2num2(reinterpret_cast(&frag_zp)[i]); - frag_b[0] = __hsub2(frag_b[0], zp); - frag_b[1] = __hsub2(frag_b[1], zp); -} - -// Same as above, but for act_order (each K is multiplied individually) -template -__device__ inline void scale4(typename MarlinType::FragB& frag_b, - typename MarlinType::FragS& frag_s_1, typename MarlinType::FragS& frag_s_2, - typename MarlinType::FragS& frag_s_3, typename MarlinType::FragS& frag_s_4, int i) -{ - using scalar_t2 = typename MarlinType::scalar_t2; - - scalar_t2 s_val_1_2; - s_val_1_2.x = reinterpret_cast(&frag_s_1)[i]; - s_val_1_2.y = reinterpret_cast(&frag_s_2)[i]; - - scalar_t2 s_val_3_4; - s_val_3_4.x = reinterpret_cast(&frag_s_3)[i]; - s_val_3_4.y = reinterpret_cast(&frag_s_4)[i]; - - frag_b[0] = __hmul2(frag_b[0], s_val_1_2); - frag_b[1] = __hmul2(frag_b[1], s_val_3_4); -} - -// Given 2 floats multiply by 2 scales (halves) -template -__device__ inline void scale_float(float* c, typename MarlinType::FragS& s) -{ - scalar_t* s_ptr = reinterpret_cast(&s); - c[0] = __fmul_rn(c[0], MarlinType::num2float(s_ptr[0])); - c[1] = __fmul_rn(c[1], MarlinType::num2float(s_ptr[1])); -} - -// Wait until barrier reaches `count`, then lock for current threadblock. -__device__ inline void barrier_acquire(int* lock, int count) -{ - if (threadIdx.x == 0) - { - int state = -1; - do - // Guarantee that subsequent writes by this threadblock will be visible - // globally. - asm volatile("ld.global.acquire.gpu.b32 %0, [%1];\n" : "=r"(state) : "l"(lock)); - while (state != count); - } - __syncthreads(); -} - -// Release barrier and increment visitation count. -__device__ inline void barrier_release(int* lock, bool reset = false) -{ - __syncthreads(); - if (threadIdx.x == 0) - { - if (reset) - { - lock[0] = 0; - return; - } - int val = 1; - // Make sure that all writes since acquiring this barrier are visible - // globally, while releasing the barrier. - asm volatile("fence.acq_rel.gpu;\n"); - asm volatile("red.relaxed.gpu.global.add.s32 [%0], %1;\n" : : "l"(lock), "r"(val)); - } -} - -// Wait until value of lock to be negative, and then add 1 -__device__ inline void wait_negative_and_add(int* lock) -{ - if (threadIdx.x == 0) - { - int state = 0; - do - // Guarantee that subsequent writes by this threadblock will be visible - // globally. - asm volatile("ld.global.acquire.gpu.b32 %0, [%1];\n" : "=r"(state) : "l"(lock)); - while (state >= 0); - atomicAdd(lock, 1); - } - __syncthreads(); -} - -template -__global__ void Marlin(MARLIN_KERNEL_PARAMS) -{ - // Each threadblock processes one "stripe" of the B matrix with (roughly) the - // same size, which might involve multiple column "slices" (of width 16 * - // `thread_n_blocks`). Stripes are defined as shown in the 3x3 matrix 5 SM - // example: - // 0 1 3 - // 0 2 3 - // 1 2 4 - // While this kind of partitioning makes things somewhat more complicated, it - // ensures good utilization of all SMs for many kinds of shape and GPU - // configurations, while requiring as few slow global cross-threadblock - // reductions as possible. - - // NVFP4 kernel: BF16 activations only, no FP8/Turing arch guards needed. - static_assert(std::is_same::value, "NVFP4 kernel only supports BF16 compute type"); - - int num_tokens_past_padded = num_tokens_past_padded_ptr[0]; - constexpr int moe_block_size = m_block_size_8 ? 8 : (16 * thread_m_blocks); - - constexpr bool use_fp16_accum = false; - using Dtype = MarlinType; - - using scalar_t2 = typename MarlinType::scalar_t2; - using scalar_32bit_t = typename MarlinType::scalar_32bit_t; - - using FragA = typename MarlinType::FragA; - using FragB = typename MarlinType::FragB; - using FragC = typename MarlinType::FragC; - using FragS = typename MarlinType::FragS; - using FragZP = typename MarlinType::FragZP; - - extern __shared__ int4 sh[]; - // NVFP4: b_type=FP4_E2M1, s_type=FP8_E4M3, a_type=c_type=BF16 - constexpr bool is_a_8bit = false; // BF16 activations are 16-bit - constexpr bool has_zp = false; // FP4 E2M1 has no zero-points - constexpr bool is_int_type = false; // FP4 E2M1 is not int type - constexpr bool dequant_skip_flop = true; // FP4 E2M1 + FP8 E4M3 scales - - scalar_t2 global_scale; - - constexpr bool has_act_order = group_blocks == 0; - - constexpr int pack_factor = 8; // 32 / 4 bits for FP4 E2M1 - static_assert(thread_m_blocks == 1 || !m_block_size_8); - int const group_size = (!has_act_order && group_blocks == -1) ? prob_k : prob_k / num_groups; - int const scales_expert_stride = prob_n * prob_k / group_size / 16; - int const zp_expert_stride = 0; // No zero-points for NVFP4 - int const b_bias_expert_stride = prob_n / 8; - - // parallel: num valid moe blocks - int parallel = num_tokens_past_padded / moe_block_size; - - int k_tiles = prob_k / 16 / thread_k_blocks; - int n_tiles = prob_n / 16 / thread_n_blocks; - - int global_mn_tiles = parallel * n_tiles; - int part2_mn_tiles = global_mn_tiles; - int part1_mn_iters = 0; - bool in_part2 = false; - - // we use DP + two-tile SK here - // part1: DP - // part2: two-tile SK - // see https://github.com/vllm-project/vllm/pull/24722 for more details - if (global_mn_tiles > gridDim.x) - { - part2_mn_tiles = global_mn_tiles % gridDim.x; - if (part2_mn_tiles * 3 <= gridDim.x) - part2_mn_tiles += gridDim.x; - part1_mn_iters = (global_mn_tiles - part2_mn_tiles) / gridDim.x; - } - - int iters = div_ceil(k_tiles * part2_mn_tiles, gridDim.x); - - if constexpr (!has_act_order && group_blocks != -1) - { - if (group_blocks >= thread_k_blocks) - { - // Ensure that the number of tiles in each stripe is a multiple of the - // groupsize; this avoids an annoying special case where a stripe starts - // in the middle of group. - iters = (group_blocks / thread_k_blocks) * div_ceil(iters, (group_blocks / thread_k_blocks)); - } - } - - int slice_row = 0; - int slice_col_par = blockIdx.x; - int slice_col; - int slice_iters = k_tiles; // number of threadblock tiles in the current slice - // total number of active threadblocks in the current slice - int slice_count = 1; - // index of threadblock in current slice; numbered bottom to top - int slice_idx = 0; - - int par_id = 0; - int block_id = -1; - int64_t expert_id = 0; // use int64 to avoid computation result overflow - int old_expert_id = 0; - int64_t B_expert_off = 0; - - float* sh_a_s = reinterpret_cast(sh); - int4* sh_block_sorted_ids_int4 = sh + (is_a_8bit ? (4 * thread_m_blocks) : 0); - int4* sh_rd_block_sorted_ids_int4 = sh_block_sorted_ids_int4 + moe_block_size / 4; - int4* sh_block_topk_weights_int4 = sh_rd_block_sorted_ids_int4 + moe_block_size / 4; - // sh_block_topk_weights_int4 only need (moe_block_size / 4); - // but we pad to align to 256 bytes - int4* sh_new = sh_block_topk_weights_int4 + moe_block_size / 2; - int32_t* sh_block_sorted_ids = reinterpret_cast(sh_block_sorted_ids_int4); - int32_t* sh_rd_block_sorted_ids = reinterpret_cast(sh_rd_block_sorted_ids_int4); - scalar_t2* sh_block_topk_weights = reinterpret_cast(sh_block_topk_weights_int4); - - int32_t block_num_valid_tokens = 0; - int32_t locks_off = 0; - - // We can easily implement parallel problem execution by just remapping - // indices and advancing global pointers - if (part2_mn_tiles >= gridDim.x) - { - // when part2_mn_tiles >= sms - // then there are at most $sms$ conflict tile blocks - locks_off = blockIdx.x; - } - else - { - locks_off = (iters * blockIdx.x) / k_tiles - 1; - } - - int prob_m_top_k = prob_m * top_k; - // read moe block data given block_id - // block_sorted_ids / block_num_valid_tokens / block_topk_weights - auto read_moe_block_data = [&](int block_id) - { - block_num_valid_tokens = moe_block_size; - - cp_async4_pred(sh_block_sorted_ids_int4 + threadIdx.x, - reinterpret_cast(sorted_token_ids_ptr) + (block_id * moe_block_size / 4 + threadIdx.x), - threadIdx.x < moe_block_size / 4); - - cp_async_fence(); - cp_async_wait<0>(); - - __syncthreads(); - - if (threadIdx.x >= threads - 32) - { - constexpr int size_per_thread = div_ceil(moe_block_size, 32); - int lane_id = threadIdx.x - (threads - 32); - - int local_count = 0; -#pragma unroll - for (int i = 0; i < size_per_thread; i++) - { - int j = lane_id * size_per_thread + i; - if (j < moe_block_size) - { - int idx = sh_block_sorted_ids[j]; - if (idx < prob_m_top_k) - local_count++; - } - } - - block_num_valid_tokens = __reduce_add_sync(0xffffffff, local_count); - - if (lane_id == 0) - reinterpret_cast(sh_new)[0] = block_num_valid_tokens; - } - - if (threadIdx.x < moe_block_size) - { - int idx = sh_block_sorted_ids[threadIdx.x]; - sh_rd_block_sorted_ids[threadIdx.x] = idx / top_k; - - if (mul_topk_weights) - { - idx = idx < prob_m_top_k ? idx : 0; - scalar_t2 topk_weight_val = Dtype::num2num2(Dtype::float2num(topk_weights_ptr[idx])); - topk_weight_val = __hmul2(topk_weight_val, global_scale); - sh_block_topk_weights[threadIdx.x] = topk_weight_val; - } - } - - __syncthreads(); - - block_num_valid_tokens = reinterpret_cast(sh_new)[0]; - __syncthreads(); - }; - - // when move to next moe block, find the next block_id and expert_id - // and then read moe block data - auto update_next_moe_block_data = [&]() - { - if (par_id >= parallel) - return; - - old_expert_id = expert_id; - block_id = par_id; - expert_id = expert_ids_ptr[block_id]; - - { - uint16_t val = global_scale_ptr[expert_id]; - global_scale = Dtype::num2num2(*reinterpret_cast(&val)); - } - - B_expert_off = expert_id * prob_n * prob_k / (pack_factor * 4); - scales_ptr += (expert_id - old_expert_id) * scales_expert_stride; - if constexpr (has_zp) - { - zp_ptr += (expert_id - old_expert_id) * zp_expert_stride; - } - if constexpr (has_act_order) - { - g_idx += (expert_id - old_expert_id) * prob_k; - } - if (has_bias) - { - b_bias_ptr += (expert_id - old_expert_id) * b_bias_expert_stride; - } - - read_moe_block_data(block_id); - }; - - // Compute all information about the current slice which is required for - // synchronization. - bool first_init = true; - auto init_part2_slice = [&]() - { - slice_iters = iters * (blockIdx.x + 1) - (k_tiles * slice_col_par + slice_row); - if (slice_iters < 0 || slice_col_par >= part2_mn_tiles) - slice_iters = 0; - if (slice_iters == 0) - return; - if (slice_row + slice_iters > k_tiles) - slice_iters = k_tiles - slice_row; - slice_count = 1; - slice_idx = 0; - int col_first = iters * div_ceil(k_tiles * slice_col_par, iters); - if (col_first <= k_tiles * (slice_col_par + 1)) - { - int col_off = col_first - k_tiles * slice_col_par; - slice_count = div_ceil(k_tiles - col_off, iters); - if (col_off > 0) - slice_count++; - int delta_first = iters * blockIdx.x - col_first; - if (delta_first < 0 || (col_off == 0 && delta_first == 0)) - slice_idx = slice_count - 1; - else - { - slice_idx = slice_count - 1 - delta_first / iters; - if (col_off > 0) - slice_idx--; - } - } - if (part2_mn_tiles >= gridDim.x) - { - if (slice_count > 1 && slice_idx == slice_count - 1) - { - locks_off++; - } - } - else - { - locks_off++; - } - - if (first_init && use_atomic_add && slice_count > 1 && slice_idx == 0) - { - constexpr int threads_per_m = 16 * thread_n_blocks / 8; - int m_per_thread = div_ceil(block_num_valid_tokens, threads / threads_per_m); - for (int i = 0; i < m_per_thread; i++) - { - int row = threads / threads_per_m * i + threadIdx.x / threads_per_m; - if (row < block_num_valid_tokens) - { - int64_t sorted_row = sh_block_sorted_ids[row]; - int col = slice_col * 16 * thread_n_blocks / 8 + threadIdx.x % threads_per_m; - C[sorted_row * prob_n / 8 + col] = {0, 0, 0, 0}; - } - } - // After write zero to output, write a negative value to lock. - // Every SM that processes the same slice would wait for - // the negative value, and then atomicAdd 1 to it. - // After all SMs are processed, the lock value would back to 0 again. - __syncthreads(); - if (threadIdx.x == 0) - locks[locks_off] = 1 - slice_count; - } - - if (slice_col == n_tiles) - { - slice_col = 0; - par_id++; - update_next_moe_block_data(); - } - if (is_a_8bit && (first_init || slice_col == 0)) - { - __syncthreads(); - cp_async1_ca_pred(&sh_a_s[threadIdx.x], &a_scales_ptr[sh_rd_block_sorted_ids[threadIdx.x]], - threadIdx.x < block_num_valid_tokens); - } - }; - - auto init_part1_slice = [&]() - { - if (part1_mn_iters) - { - part1_mn_iters--; - par_id = slice_col_par / n_tiles; - slice_col = slice_col_par % n_tiles; - slice_iters = k_tiles; - update_next_moe_block_data(); - if (is_a_8bit) - { - __syncthreads(); - cp_async1_ca_pred(&sh_a_s[threadIdx.x], &a_scales_ptr[sh_rd_block_sorted_ids[threadIdx.x]], - threadIdx.x < block_num_valid_tokens); - } - } - }; - - auto init_slice = [&]() - { - if (!in_part2 && !part1_mn_iters) - { - in_part2 = true; - slice_col_par = (iters * blockIdx.x) / k_tiles; - slice_row = (iters * blockIdx.x) % k_tiles; - slice_col = (slice_col_par + global_mn_tiles - part2_mn_tiles) % n_tiles; - par_id = (slice_col_par + global_mn_tiles - part2_mn_tiles) / n_tiles; - update_next_moe_block_data(); - } - if (!in_part2) - { - init_part1_slice(); - } - else - { - init_part2_slice(); - first_init = false; - } - }; - - init_slice(); - - // A sizes/strides - - // stride of the A matrix in global memory - int a_gl_stride = prob_k / (is_a_8bit ? 16 : 8); - // stride of an A matrix tile in shared memory - constexpr int a_sh_stride = 16 * thread_k_blocks / (is_a_8bit ? 16 : 8); - // delta between subsequent A tiles in global memory - constexpr int a_gl_rd_delta_o = 16 * thread_k_blocks / (is_a_8bit ? 16 : 8); - // between subsequent accesses within a tile - int a_gl_rd_delta_i = a_gl_stride * (threads / a_gl_rd_delta_o); - // between shared memory writes - constexpr int a_sh_wr_delta = a_sh_stride * (threads / a_gl_rd_delta_o); - // within a shared memory tile - constexpr int a_sh_rd_delta_i = a_sh_stride * 16; - // overall size of a tile - constexpr int a_sh_stage = a_sh_stride * (16 * thread_m_blocks); - // number of shared write iterations for a tile - constexpr int a_sh_wr_iters = div_ceil(a_sh_stage, a_sh_wr_delta); - - // B sizes/strides - int b_gl_stride = 16 * prob_n / (pack_factor * (is_a_8bit ? 2 : 4)); - constexpr int b_sh_stride = ((thread_n_blocks * 16) * 16 / pack_factor) / (is_a_8bit ? 2 : 4); - constexpr int b_thread_vecs = 1; - constexpr int b_sh_stride_threads = b_sh_stride / b_thread_vecs; - - int b_gl_rd_delta_o = b_gl_stride * thread_k_blocks / (is_a_8bit ? 2 : 1); - constexpr int b_sh_wr_delta = threads * b_thread_vecs; - constexpr int b_sh_stage = b_sh_stride * thread_k_blocks / (is_a_8bit ? 2 : 1); - constexpr int b_sh_wr_iters = b_sh_stage / b_sh_wr_delta; - - // Scale sizes/strides without act_order - int s_gl_stride = prob_n / 16; - constexpr int s_sh_stride = 16 * thread_n_blocks / 16; - constexpr int s_tb_groups - = !has_act_order && group_blocks != -1 && group_blocks < thread_k_blocks ? thread_k_blocks / group_blocks : 1; - constexpr int s_sh_stage = s_tb_groups * s_sh_stride; - int s_gl_rd_delta = s_gl_stride; - - // Scale size/strides with act_order - constexpr int tb_k = 16 * thread_k_blocks; - constexpr int g_idx_stage = has_act_order ? (tb_k * sizeof(int)) / 16 : 0; - // constexpr int act_s_row_stride = 1; - // int act_s_col_stride = act_s_row_stride * num_groups; - constexpr int act_s_max_num_groups = 32; - int act_s_col_stride = 1; - int act_s_col_warp_stride = act_s_col_stride * 8; - - constexpr int tb_n_warps = thread_n_blocks / (is_a_8bit ? 2 : 4); - int act_s_col_tb_stride = act_s_col_warp_stride * tb_n_warps; - - // Zero-points sizes/strides - int zp_gl_stride = 0; // No zero-points for NVFP4 - constexpr int zp_sh_stride = 0; // No zero-points for NVFP4 - constexpr int zp_tb_groups = s_tb_groups; - constexpr int zp_sh_stage = has_zp ? zp_tb_groups * zp_sh_stride : 0; - int zp_gl_rd_delta = zp_gl_stride; - - // Global A read index of current thread. - int a_gl_rd_row = threadIdx.x / a_gl_rd_delta_o; - int a_gl_rd_col = a_gl_rd_delta_o * slice_row + threadIdx.x % a_gl_rd_delta_o; - // Shared write index of current thread. - int a_sh_wr = a_sh_stride * (threadIdx.x / a_gl_rd_delta_o) + (threadIdx.x % a_gl_rd_delta_o); - // Shared read index. - int a_sh_rd = a_sh_stride * ((threadIdx.x % 32) % (16 / (m_block_size_8 ? 2 : 1))) - + (threadIdx.x % 32) / (16 / (m_block_size_8 ? 2 : 1)); - a_sh_rd += 2 * ((threadIdx.x / 32) / tb_n_warps) * b_sh_wr_iters; - - int b_gl_rd; - if (threads <= b_sh_stride) - { - b_gl_rd = threadIdx.x; - } - else - { - b_gl_rd = b_gl_stride * (threadIdx.x / b_sh_stride) + (threadIdx.x % b_sh_stride); - } - - b_gl_rd += B_expert_off + b_sh_stride * slice_col; - b_gl_rd += b_gl_rd_delta_o * slice_row; - auto b_sh_rd = threadIdx.x * b_thread_vecs; - b_sh_rd += b_sh_rd / b_sh_stride * (b_sh_stride * (b_sh_wr_iters - 1)); - - // For act_order - int slice_k_start = tb_k * slice_row; - int slice_k_finish = slice_k_start + tb_k * slice_iters; - int slice_k_start_shared_fetch = slice_k_start; - int slice_n_offset = act_s_col_tb_stride * slice_col; - - // No act_order - int s_gl_rd; - if constexpr (!has_act_order) - { - if constexpr (group_blocks == -1) - { - s_gl_rd = s_sh_stride * slice_col + threadIdx.x; - } - else if constexpr (group_blocks >= thread_k_blocks) - { - s_gl_rd - = s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) + s_sh_stride * slice_col + threadIdx.x; - } - else - { - s_gl_rd = s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks + threadIdx.x / s_sh_stride) - + s_sh_stride * slice_col + threadIdx.x % s_sh_stride; - } - } - auto s_sh_wr = threadIdx.x; - bool s_sh_wr_pred = threadIdx.x < s_sh_stage; - - // Zero-points - int zp_gl_rd; - if constexpr (has_zp) - { - if constexpr (group_blocks == -1) - { - zp_gl_rd = zp_sh_stride * slice_col + threadIdx.x; - } - else if constexpr (group_blocks >= thread_k_blocks) - { - zp_gl_rd = zp_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) + zp_sh_stride * slice_col - + threadIdx.x; - } - else - { - zp_gl_rd = zp_gl_stride * ((thread_k_blocks * slice_row) / group_blocks + threadIdx.x / zp_sh_stride) - + zp_sh_stride * slice_col + threadIdx.x % zp_sh_stride; - } - } - auto zp_sh_wr = threadIdx.x; - bool zp_sh_wr_pred = zp_sh_stage > 0 && threadIdx.x < zp_sh_stage; - - // We use a different scale layout for grouped and column-wise quantization as - // we scale a `half2` tile in column-major layout in the former and in - // row-major in the latter case. - int s_sh_rd; - if constexpr (is_a_8bit) - { - s_sh_rd = 4 * ((threadIdx.x / 32) % tb_n_warps) + (threadIdx.x % 4); - } - else if constexpr (group_blocks != -1) - s_sh_rd = 8 * ((threadIdx.x / 32) % tb_n_warps) + (threadIdx.x % 32) / 4; - else if constexpr (group_blocks == -1 && (m_block_size_8 || (has_zp && !dequant_skip_flop))) - s_sh_rd = 8 * ((threadIdx.x / 32) % tb_n_warps) + (threadIdx.x % 32) / 8; - else - s_sh_rd = 8 * ((threadIdx.x / 32) % tb_n_warps) + (threadIdx.x % 32) % 4; - - int bias_sh_rd; - if constexpr (m_block_size_8) - { - bias_sh_rd = 8 * ((threadIdx.x / 32) % tb_n_warps) + (threadIdx.x % 32) / 8; - } - else - { - bias_sh_rd = (is_a_8bit ? 4 : 8) * ((threadIdx.x / 32) % tb_n_warps) + (threadIdx.x % 32) % 4; - } - - int bias_sh_wr = threadIdx.x; - int bias_gl_rd = (thread_n_blocks * 16 / 8) * slice_col + threadIdx.x; - - // Zero-points have the same read layout as the scales - // (without column-wise case) - constexpr int num_col_threads = 8; - constexpr int num_row_threads = 4; - constexpr int num_ints_per_thread = 8 / pack_factor; - int zp_sh_rd; - if constexpr (has_zp) - { - if (is_a_8bit) - { - zp_sh_rd = num_ints_per_thread * num_col_threads * ((threadIdx.x / 32) % tb_n_warps / 2) - + num_ints_per_thread * ((threadIdx.x % 32) / num_row_threads); - } - else - { - zp_sh_rd = num_ints_per_thread * num_col_threads * ((threadIdx.x / 32) % tb_n_warps) - + num_ints_per_thread * ((threadIdx.x % 32) / num_row_threads); - } - } - - // To ensure that writing and reading A tiles to/from shared memory, the - // latter in fragment format, is fully bank conflict free, we need to use a - // rather fancy XOR-based layout. The key here is that neither reads nor - // writes of the 16-byte `int4` blocks of 8 consecutive threads involve the - // same shared memory banks. Further, it seems (based on NSight-Compute) that - // each warp must also write a consecutive memory segment? - auto transform_a = [&](int i) - { - int row = i / a_gl_rd_delta_o; - return a_gl_rd_delta_o * row + (i % a_gl_rd_delta_o) ^ (row % 8); - }; - // Since the computation of this remapping is non-trivial and, due to our main - // loop unrolls, all shared memory accesses are static, we simply precompute - // both transformed reads and writes. - int a_sh_wr_trans[a_sh_wr_iters]; -#pragma unroll - for (int i = 0; i < a_sh_wr_iters; i++) - a_sh_wr_trans[i] = transform_a(a_sh_wr_delta * i + a_sh_wr); - int a_sh_rd_trans[b_sh_wr_iters][thread_m_blocks]; -#pragma unroll - for (int i = 0; i < b_sh_wr_iters; i++) - { -#pragma unroll - for (int j = 0; j < thread_m_blocks; j++) - a_sh_rd_trans[i][j] = transform_a(2 * i + a_sh_rd_delta_i * j + a_sh_rd); - } - - // Since B-accesses have non-constant stride they have to be computed at - // runtime; we break dependencies between subsequent accesses with a tile by - // maintining multiple pointers (we have enough registers), a tiny - // optimization. - - // Shared memory storage for global fetch pipelines. - constexpr int sh_red_size = (2 * thread_n_blocks + 1) * 16 * thread_m_blocks; - constexpr int sh_b_size = stages * b_sh_stage; - int4* sh_b = sh_new; - int4* sh_red = sh_new; - - constexpr int sh_size_b_red_min = (sh_red_size < sh_b_size ? sh_red_size : sh_b_size); - constexpr int sh_size_b_red_max = (sh_red_size > sh_b_size ? sh_red_size : sh_b_size); - constexpr int sh_bias_size = (thread_n_blocks * 16 / 8); - constexpr int sh_b_red_bias_size = sh_size_b_red_max > (sh_size_b_red_min + sh_bias_size) - ? sh_size_b_red_max - : (sh_size_b_red_min + sh_bias_size); - - int4* sh_bias = sh_new + sh_size_b_red_min; - int4* sh_g_idx = sh_new + sh_b_red_bias_size; - int4* sh_zp = sh_g_idx + (stages * g_idx_stage); - constexpr int sh_s_size = has_act_order ? (act_s_max_num_groups * s_sh_stride) : (stages * s_sh_stage); - int4* sh_s = sh_zp + (stages * zp_sh_stage); - int4* sh_a = sh_s + sh_s_size; - - // Register storage for double buffer of shared memory reads. - FragA frag_a[2][thread_m_blocks]; - I4 frag_b_quant[2][b_thread_vecs]; - FragC frag_c[thread_m_blocks][is_a_8bit ? 2 : 4][2]; - FragC frag_c_tmp[thread_m_blocks][is_a_8bit ? 2 : 4][2]; - FragS frag_s[2][4]; // No act-order - FragS frag_bias[2][4]; - FragS act_frag_s[2][4][4]; // For act-order - int frag_qzp[2][num_ints_per_thread]; // Zero-points - FragZP frag_zp; // Zero-points in fp16 - FragZP frag_zpf[2]; // Zero-points in fp16 in HQQ - - if constexpr (is_a_8bit && group_blocks != -1) - { -#pragma unroll - for (int j = 0; j < 2; j++) - { -#pragma unroll - for (int i = 0; i < thread_m_blocks; i++) - { -#pragma unroll - for (int g = 0; g < 4; g++) - { - frag_c_tmp[i][j][0][g] = 0.0f; - } - -#pragma unroll - for (int g = 0; g < 4; g++) - { - frag_c_tmp[i][j][1][g] = 0.0f; - } - } - } - } - - // Zero accumulators. - auto zero_accums = [&]() - { -#pragma unroll - for (int i = 0; i < thread_m_blocks * 4 * 2 * 4; i++) - reinterpret_cast(frag_c)[i] = 0; - }; - - int sh_first_group_id = -1; - int sh_num_groups = -1; - - auto fetch_act_order_scales_to_shared = [&](bool is_async, int first_group_id, int last_group_id) - { - sh_first_group_id = first_group_id; - sh_num_groups = last_group_id - first_group_id + 1; - - if (sh_num_groups > act_s_max_num_groups) - { - sh_num_groups = act_s_max_num_groups; - } - - if (sh_first_group_id + sh_num_groups > num_groups) - { - sh_num_groups = num_groups - sh_first_group_id; - } - - int row_offset = first_group_id * s_gl_stride; - - if (is_async) - { - for (int i = 0; i < sh_num_groups; i++) - { - if (threadIdx.x < s_sh_stride) - { - cp_async4_pred(&sh_s[(i * s_sh_stride) + threadIdx.x], - &scales_ptr[row_offset + (i * s_gl_stride) + slice_n_offset + threadIdx.x]); - } - } - } - else - { - for (int i = 0; i < sh_num_groups; i++) - { - if (threadIdx.x < s_sh_stride) - { - sh_s[(i * s_sh_stride) + threadIdx.x] - = scales_ptr[row_offset + (i * s_gl_stride) + slice_n_offset + threadIdx.x]; - } - } - } - }; - // Asynchronously fetch the next A, B and s tile from global to the next - // shared memory pipeline location. - auto fetch_to_shared = [&](int pipe, int a_off, bool pred = true) - { - if (pred) - { - int4* sh_a_stage = sh_a + moe_block_size * a_sh_stride * pipe; -#pragma unroll - for (int i = 0; i < a_sh_wr_iters; i++) - { - int row = a_gl_rd_delta_i / a_gl_stride * i + a_gl_rd_row; - int64_t sorted_row = 0; - if (!m_block_size_8 || row < 8) - sorted_row = sh_rd_block_sorted_ids[row]; - int64_t true_idx = sorted_row * a_gl_stride + a_gl_rd_col + a_gl_rd_delta_o * a_off; - cp_async4_pred(&sh_a_stage[a_sh_wr_trans[i]], &A[true_idx], row < block_num_valid_tokens); - } - - int4* sh_b_stage = sh_b + b_sh_stage * pipe; -#pragma unroll - for (int i = 0; i < (b_sh_wr_iters * b_thread_vecs); i++) - { - constexpr int count = div_ceil(b_sh_stride, threads); - int b_gl_idx - = b_gl_rd + (i % count) * threads + b_gl_stride * (i / count) * div_ceil(threads, b_sh_stride); - - cp_async4(&sh_b_stage[threads * i + threadIdx.x], &B[b_gl_idx]); - } - - b_gl_rd += b_gl_rd_delta_o; - - if constexpr (has_act_order) - { - // Fetch g_idx thread-block portion - int full_pipe = a_off; - int cur_k = slice_k_start_shared_fetch + tb_k * full_pipe; - if (cur_k < prob_k && cur_k < slice_k_finish) - { - int4* sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe; - - int4 const* cur_g_idx_stage_ptr = reinterpret_cast(&g_idx[cur_k]); - - if (threadIdx.x < g_idx_stage) - { - cp_async4_pred(&sh_g_idx_stage[threadIdx.x], &cur_g_idx_stage_ptr[threadIdx.x]); - } - } - } - else - { - if constexpr (group_blocks != -1) - { - int4* sh_s_stage = sh_s + s_sh_stage * pipe; - - // Only fetch scales if this tile starts a new group - if (pipe % div_ceil(group_blocks, thread_k_blocks) == 0) - { - if (s_sh_wr_pred) - { - cp_async4(&sh_s_stage[s_sh_wr], &scales_ptr[s_gl_rd]); - } - s_gl_rd += s_gl_rd_delta * s_tb_groups; - } - } - - if constexpr (has_zp && group_blocks != -1) - { - int4* sh_zp_stage = sh_zp + zp_sh_stage * pipe; - - // Only fetch zero points if this tile starts a new group - if (pipe % div_ceil(group_blocks, thread_k_blocks) == 0) - { - if (zp_sh_wr_pred) - { - cp_async4(&sh_zp_stage[zp_sh_wr], &zp_ptr[zp_gl_rd]); - } - zp_gl_rd += zp_gl_rd_delta * zp_tb_groups; - } - } - } - } - // Insert a fence even when we are winding down the pipeline to ensure that - // waiting is also correct at this point. - cp_async_fence(); - }; - - auto fetch_col_zp_to_shared = [&]() - { - if (zp_sh_wr_pred) - { - cp_async4(&sh_zp[zp_sh_wr], &zp_ptr[zp_gl_rd]); - } - }; - - auto fetch_col_scale_to_shared = [&]() - { - if (s_sh_wr_pred) - { - cp_async4(&sh_s[s_sh_wr], &scales_ptr[s_gl_rd]); - } - }; - - // Wait until the next thread tile has been loaded to shared memory. - auto wait_for_stage = [&]() - { - // We only have `stages - 2` active fetches since we are double buffering - // and can only issue the next fetch when it is guaranteed that the previous - // shared memory load is fully complete (as it may otherwise be - // overwritten). - cp_async_wait(); - __syncthreads(); - }; - - // Load the next sub-tile from the current location in the shared memory pipe - // into the current register buffer. - auto fetch_to_registers = [&](int k, int pipe) - { - int4* sh_a_stage = sh_a + moe_block_size * a_sh_stride * pipe; -#pragma unroll - for (int i = 0; i < thread_m_blocks; i++) - ldsm(frag_a[k % 2][i], &sh_a_stage[a_sh_rd_trans[k % b_sh_wr_iters][i]]); - int4* sh_b_stage = sh_b + b_sh_stage * pipe; - -#pragma unroll - for (int i = 0; i < b_thread_vecs; i++) - { - frag_b_quant[k % 2][i] - = *reinterpret_cast(&sh_b_stage[b_sh_stride * (k % b_sh_wr_iters) + b_sh_rd + i]); - } - }; - - bool is_same_group[stages]; - int same_group_id[stages]; - - auto init_same_group = [&](int pipe) - { - if constexpr (!has_act_order) - { - return; - } - - int4* sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe; - int* sh_g_idx_int_ptr = reinterpret_cast(sh_g_idx_stage); - - int group_id_1 = sh_g_idx_int_ptr[0]; - int group_id_2 = sh_g_idx_int_ptr[tb_k - 1]; - - is_same_group[pipe] = group_id_1 == group_id_2; - same_group_id[pipe] = group_id_1; - }; - - auto fetch_scales_to_registers = [&](int k, int full_pipe) - { - int pipe = full_pipe % stages; - using IT1 = typename std::conditional_t; - using IT0 = typename std::conditional_t; - constexpr int group_blocks2 = div_ceil(group_blocks, is_a_8bit ? 2 : 1); - - if constexpr (!has_act_order) - { - // No act-order case - if constexpr (group_blocks == -1) - { - // load only when starting a new slice - if (k == 0 && full_pipe == 0 && dequant_skip_flop) - { - reinterpret_cast(&frag_s)[0] = sh_s[s_sh_rd]; - reinterpret_cast(&frag_s)[1] = sh_s[s_sh_rd + 4]; - } - } - else if constexpr (group_blocks != -1) - { - if constexpr (group_blocks >= thread_k_blocks) - { - constexpr int g = group_blocks / thread_k_blocks; - if (pipe % g == 0) - { - if (k % b_sh_wr_iters == 0) - { - int4* sh_s_stage = sh_s + s_sh_stage * (g * (pipe / g)); - reinterpret_cast(&frag_s[k % 2])[0] = sh_s_stage[s_sh_rd]; - } - else - { - reinterpret_cast(&frag_s[1])[0] = reinterpret_cast(&frag_s[0])[0]; - } - } - } - else if (group_blocks2 < b_sh_wr_iters || k % b_sh_wr_iters == 0) - { - auto warp_id = threadIdx.x / 32; - int warp_row = warp_id / tb_n_warps; - - int k_blocks = b_sh_wr_iters * warp_row + k % b_sh_wr_iters; - int cur_group_id = k_blocks / group_blocks2; - - int4* sh_s_stage = sh_s + s_sh_stage * pipe; - - reinterpret_cast(&frag_s[k % 2])[0] - = reinterpret_cast(sh_s_stage)[s_sh_rd + cur_group_id * (2 * s_sh_stride)]; - } - else if (group_blocks >= b_sh_wr_iters) - { - reinterpret_cast(&frag_s[1])[0] = reinterpret_cast(&frag_s[0])[0]; - } - } - - return; - } - - // Act-order case - - // Determine K of the "current" thread-block - int cur_k = slice_k_start + tb_k * full_pipe; - if (cur_k >= prob_k || cur_k >= slice_k_finish) - { - return; - } - - // Reset (to current thread-block) since we read g_idx portion from the - // shared memory - cur_k = 0; - - // Progress to current iteration - cur_k += k % b_sh_wr_iters; - - // Determine "position" inside the thread-block (based on warp and - // thread-id) - auto warp_id = threadIdx.x / 32; - int warp_row = warp_id / tb_n_warps; - int warp_col = warp_id % tb_n_warps; - - cur_k += warp_row * 16 * b_sh_wr_iters; - - auto th_id = threadIdx.x % 32; - cur_k += (th_id % 4) * 2; // Due to tensor-core layout for fp16 B matrix - - int s_col_shift = - /*slice_n_offset +*/ (act_s_col_warp_stride * warp_col) + (th_id / 4) * act_s_col_stride; - - if (is_same_group[pipe]) - { - if (k % 2 == 0) - { - *(reinterpret_cast(&(act_frag_s[k % 2][0][0]))) - = sh_s[(same_group_id[pipe] - sh_first_group_id) * s_sh_stride + s_col_shift]; - } - else - { - *(reinterpret_cast(&(act_frag_s[k % 2][0][0]))) - = *(reinterpret_cast(&(act_frag_s[(k - 1) % 2][0][0]))); - } - - for (int i = 1; i < 4; i++) - { - *(reinterpret_cast(&(act_frag_s[k % 2][i][0]))) - = *(reinterpret_cast(&(act_frag_s[k % 2][0][0]))); - } - return; - } - - int4* sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe; - int* sh_g_idx_int_ptr = reinterpret_cast(sh_g_idx_stage); - - constexpr int k_frag_offsets[4] = {0, 1, 8, 9}; // Tensor core offsets per thread - -#pragma unroll - for (int i = 0; i < 4; i++) - { - int actual_k = cur_k + k_frag_offsets[i]; - - int group_id = sh_g_idx_int_ptr[actual_k]; - int rel_group_id = group_id - sh_first_group_id; - - *(reinterpret_cast(&(act_frag_s[k % 2][i][0]))) = sh_s[rel_group_id * s_sh_stride + s_col_shift]; - } - }; - - auto fetch_zp_to_registers = [&](int k, int full_pipe) - { - // This code does not handle group_blocks == 0, - // which signifies act_order. - // has_zp implies AWQ, which doesn't have act_order, - static_assert(!has_zp || group_blocks != 0); - - if constexpr (has_zp) - { - int pipe = full_pipe % stages; - - if constexpr (group_blocks == -1) - { - // load only when starting a new slice - if (k == 0 && full_pipe == 0 || is_a_8bit) - { -#pragma unroll - for (int i = 0; i < num_ints_per_thread; i++) - { - frag_qzp[k % 2][i] = (reinterpret_cast(sh_zp))[zp_sh_rd + i]; - } - } - } - else if constexpr (group_blocks >= thread_k_blocks) - { - constexpr int g = group_blocks / thread_k_blocks; - if (pipe % g == 0 && k % b_sh_wr_iters == 0 || is_a_8bit) - { - int4* sh_zp_stage = sh_zp + zp_sh_stage * (g * (pipe / g)); -#pragma unroll - for (int i = 0; i < num_ints_per_thread; i++) - { - frag_qzp[k % 2][i] = (reinterpret_cast(sh_zp_stage))[zp_sh_rd + i]; - } - } - } - else - { - auto warp_id = threadIdx.x / 32; - - int warp_row = warp_id / tb_n_warps; - - int k_blocks = b_sh_wr_iters * warp_row + k % b_sh_wr_iters; - int cur_group_id = k_blocks / div_ceil(group_blocks, is_a_8bit ? 2 : 1); - - int4* sh_zp_stage = sh_zp + zp_sh_stage * pipe; - - sh_zp_stage += cur_group_id * zp_sh_stride; - -#pragma unroll - for (int i = 0; i < num_ints_per_thread; i++) - { - frag_qzp[k % 2][i] = (reinterpret_cast(sh_zp_stage))[zp_sh_rd + i]; - } - } - } - }; - - auto dequant_data = [&](int q, scalar_32bit_t* frag_b_ptr, int zp = 0) - { - if constexpr (is_a_8bit && has_zp) - { - dequant_fp4(q, frag_b_ptr, zp); - } - else - { - dequant_fp4(q, frag_b_ptr); - } - }; - - // Execute the actual tensor core matmul of a sub-tile. - bool is_first_matmul_in_slice = true; - auto matmul = [&](int k, int pipe) - { - if (is_a_8bit) - return; - int k2 = k % 2; - constexpr int g = group_blocks > 0 ? div_ceil(group_blocks, thread_k_blocks) : 1; - bool const is_new_zp = (group_blocks == 0) - || ((group_blocks > 0) && (group_blocks < b_sh_wr_iters || k == 0)) && (pipe % g == 0) - || (group_blocks == -1 && is_first_matmul_in_slice); - if constexpr (has_zp) - { - if (is_new_zp) - { - if constexpr (group_blocks == -1) - is_first_matmul_in_slice = false; - int zp_quant_0 = frag_qzp[k2][0]; - int zp_quant_1 = zp_quant_0 >> 8; - - dequant_data(zp_quant_0, reinterpret_cast(&frag_zp)); - dequant_data(zp_quant_1, reinterpret_cast(&frag_zp) + 2); - } - } - { - int s_quant_0 = reinterpret_cast(frag_s[k2])[0]; - int s_quant_1 = reinterpret_cast(frag_s[k2])[1]; - - dequant_fp8_scales(s_quant_0, reinterpret_cast(&frag_s[k2])); - dequant_fp8_scales(s_quant_1, reinterpret_cast(&frag_s[k2]) + 2); - } - -// We have the m dimension as the inner loop in order to encourage overlapping -// dequantization and matmul operations. -#pragma unroll - for (int j = 0; j < 4; j++) - { - FragB frag_b0; - FragB frag_b1; - int b_quant_1 = frag_b_quant[k2][0][j]; - int b_quant_0 = b_quant_1 << 8; - - dequant_data(b_quant_0, reinterpret_cast(&frag_b0)); - dequant_data(b_quant_1, reinterpret_cast(&frag_b1)); - - if constexpr (dequant_skip_flop && has_zp && !is_a_8bit) - { - sub_zp(frag_b0, frag_zp[j], 0); - sub_zp(frag_b1, frag_zp[j], 1); - } - - // Apply scale to frag_b0 - if constexpr (has_act_order && !is_a_8bit) - { - static_assert(group_blocks != -1); - scale4( - frag_b0, act_frag_s[k2][0][j], act_frag_s[k2][1][j], act_frag_s[k2][2][j], act_frag_s[k2][3][j], 0); - scale4( - frag_b1, act_frag_s[k2][0][j], act_frag_s[k2][1][j], act_frag_s[k2][2][j], act_frag_s[k2][3][j], 1); - } - else if constexpr (!dequant_skip_flop && has_zp && group_blocks == -1 && !is_a_8bit) - { - int idx = (threadIdx.x / 4) % 2; - scalar_t2 s2 = Dtype::nums2num2(reinterpret_cast(&frag_s[j / 2][j % 2 * 2 + 0])[idx], - reinterpret_cast(&frag_s[j / 2][j % 2 * 2 + 1])[idx]); - if (is_new_zp) - frag_zp[j] = __hmul2(frag_zp[j], s2); - scale_and_sub(frag_b0, s2.x, frag_zp[j].x); - scale_and_sub(frag_b1, s2.y, frag_zp[j].y); - } - else if constexpr (!dequant_skip_flop && has_zp && group_blocks != -1 && !is_a_8bit) - { - if (is_new_zp) - frag_zp[j] = __hmul2(frag_zp[j], *reinterpret_cast(&frag_s[k2][j])); - scale_and_sub(frag_b0, frag_s[k2][j][0].x, frag_zp[j].x); - scale_and_sub(frag_b1, frag_s[k2][j][0].y, frag_zp[j].y); - } - else if constexpr (group_blocks != -1 && !is_a_8bit) - { - scale(frag_b0, frag_s[k2][j], 0); - scale(frag_b1, frag_s[k2][j], 1); - } - -#pragma unroll - for (int i = 0; i < thread_m_blocks; i++) - { - if constexpr (m_block_size_8) - { - mma_trans(frag_a[k2][i], frag_b0, frag_b1, frag_c[i][j][0]); - } - else - { - mma(frag_a[k2][i], frag_b0, frag_c[i][j][0]); - mma(frag_a[k2][i], frag_b1, frag_c[i][j][1]); - } - } - } - }; - - auto matmul_a8 = [&](int k) - { - int k2 = k % 2; -#pragma unroll - for (int j = 0; j < 2; j++) - { - FragB frag_b[2]; - - if (is_a_8bit && !has_zp) - { - dequant_data(frag_b_quant[k2][0][j * 2], reinterpret_cast(&frag_b)); - dequant_data(frag_b_quant[k2][0][j * 2 + 1], reinterpret_cast(&frag_b) + 2); - } - else if (is_a_8bit && has_zp) - { - int off = (threadIdx.x / 32) % 2 * 2 + j; - int zp = (frag_qzp[k2][0] >> (off * 8)) & 0xF; - dequant_data(frag_b_quant[k2][0][j * 2], reinterpret_cast(&frag_b), zp); - zp = (frag_qzp[k2][0] >> (off * 8 + 4)) & 0xF; - dequant_data(frag_b_quant[k2][0][j * 2 + 1], reinterpret_cast(&frag_b) + 2, zp); - } - else - { - reinterpret_cast(&frag_b)[0] = reinterpret_cast(&frag_b_quant[k2][j])[0]; - reinterpret_cast(&frag_b)[1] = reinterpret_cast(&frag_b_quant[k2][j])[1]; - } - -#pragma unroll - for (int i = 0; i < thread_m_blocks; i++) - { - mma(frag_a[k2][i], frag_b[0], (group_blocks == -1 ? frag_c : frag_c_tmp)[i][j][0]); - mma(frag_a[k2][i], frag_b[1], (group_blocks == -1 ? frag_c : frag_c_tmp)[i][j][1]); - } - - if constexpr (group_blocks != -1) - { - if (group_blocks == 2 || k == 1) - { - { - float2 s_vals[2]; - s_vals[0] = Dtype::num22float2(frag_s[k2][j * 2][0]); - s_vals[1] = Dtype::num22float2(frag_s[k2][j * 2 + 1][0]); - -#pragma unroll - for (int i = 0; i < thread_m_blocks; i++) - { -#pragma unroll - for (int g = 0; g < 4; g++) - { - float scale = reinterpret_cast(&s_vals[0])[g % 2]; - frag_c[i][j][0][g] += frag_c_tmp[i][j][0][g] * scale; - frag_c_tmp[i][j][0][g] = 0.0f; - } - -#pragma unroll - for (int g = 0; g < 4; g++) - { - float scale = reinterpret_cast(&s_vals[1])[g % 2]; - frag_c[i][j][1][g] += frag_c_tmp[i][j][1][g] * scale; - frag_c_tmp[i][j][1][g] = 0.0f; - } - } - } - } - } - } - }; - - // Since we slice across the k dimension of a tile in order to increase the - // number of warps while keeping the n dimension of a tile reasonable, we have - // multiple warps that accumulate their partial sums of the same output - // location; which we have to reduce over in the end. We do in shared memory. - auto thread_block_reduce = [&]() - { - constexpr int red_off = threads / b_sh_stride_threads / 2; - if (red_off >= 1) - { - auto red_idx = threadIdx.x / b_sh_stride_threads; - constexpr int red_sh_stride = b_sh_stride_threads * (is_a_8bit ? 2 : 4) * 2; - constexpr int red_sh_delta = b_sh_stride_threads; - int red_sh_rd = red_sh_stride * (threadIdx.x / b_sh_stride_threads) + (threadIdx.x % b_sh_stride_threads); - - // Parallel logarithmic shared memory reduction. We make sure to avoid any - // unnecessary read or write iterations, e.g., for two warps we write only - // once by warp 1 and read only once by warp 0. - -#pragma unroll - for (int m_block = 0; m_block < thread_m_blocks; m_block++) - { -#pragma unroll - for (int i = red_off; i > 0; i /= 2) - { - if (i <= red_idx && red_idx < 2 * i) - { -#pragma unroll - for (int j = 0; j < (is_a_8bit ? 2 : 4) * 2; j += (m_block_size_8 ? 2 : 1)) - { - int red_sh_wr = red_sh_delta * j + (red_sh_rd - red_sh_stride * i); - if (i < red_off) - { - float* c_rd = reinterpret_cast(&sh_red[red_sh_delta * j + red_sh_rd]); - float* c_wr = reinterpret_cast(&sh_red[red_sh_wr]); -#pragma unroll - for (int k = 0; k < 4; k++) - reinterpret_cast(frag_c)[(is_a_8bit ? 2 : 4) * 2 * m_block + j][k] - += c_rd[k] + c_wr[k]; - } - sh_red[red_sh_wr] = reinterpret_cast(&frag_c)[(is_a_8bit ? 2 : 4) * 2 * m_block + j]; - } - } - __syncthreads(); - } - if (red_idx == 0) - { -#pragma unroll - for (int i = 0; i < (is_a_8bit ? 2 : 4) * 2; i += (m_block_size_8 ? 2 : 1)) - { - float* c_rd = reinterpret_cast(&sh_red[red_sh_delta * i + red_sh_rd]); -#pragma unroll - for (int j = 0; j < 4; j++) - reinterpret_cast(frag_c)[(is_a_8bit ? 2 : 4) * 2 * m_block + i][j] += c_rd[j]; - } - } - __syncthreads(); - } - } - }; - - // Since multiple threadblocks may process parts of the same column slice, we - // finally have to globally reduce over the results. As the striped - // partitioning minimizes the number of such reductions and our outputs are - // usually rather small, we perform this reduction serially in L2 cache. - auto global_reduce_fp16 = [&](bool first = false, bool last = false) - { - // We are very careful here to reduce directly in the output buffer to - // maximize L2 cache utilization in this step. To do this, we write out - // results in FP16 (but still reduce with FP32 compute). - constexpr int active_threads = 32 * tb_n_warps; - bool is_th_active = threadIdx.x < active_threads; - if (!is_th_active) - { - return; - } - - int c_gl_stride = prob_n / 8 * (is_a_8bit ? 2 : 1); - int c_gl_wr_delta_o = 8 * c_gl_stride; - int c_gl_wr_delta_i = 4 * (active_threads / 32); - int c_gl_wr; - if constexpr (m_block_size_8) - { - c_gl_wr = c_gl_stride * ((threadIdx.x % 4) * 2) + 4 * (threadIdx.x / 32) + (threadIdx.x % 32) / 8; - c_gl_wr += (2 * thread_n_blocks) * slice_col; - } - else - { - c_gl_wr = c_gl_stride * ((threadIdx.x % 32) / 4) + 4 * (threadIdx.x / 32) + threadIdx.x % 4; - c_gl_wr += (2 * thread_n_blocks) * slice_col * (is_a_8bit ? 2 : 1); - } - constexpr int c_sh_wr_delta = active_threads; - int c_sh_wr = threadIdx.x; - - if (!first) - { - -#pragma unroll - for (int i = 0; i < (m_block_size_8 ? 2 : thread_m_blocks * 4); i++) - { - int c_idx; - if constexpr (m_block_size_8) - c_idx = c_gl_wr + i * c_gl_stride + (threadIdx.x % 8) / 4 * c_gl_wr_delta_i; - else - c_idx = c_gl_wr + c_gl_wr_delta_o * (i / 2) + c_gl_wr_delta_i * (i % 2); - if (c_idx / c_gl_stride < block_num_valid_tokens) - { - int64_t sorted_row = sh_block_sorted_ids[c_idx / c_gl_stride]; - int64_t true_idx = sorted_row * c_gl_stride + c_idx % c_gl_stride; - if constexpr (is_a_8bit) - { - int2* sh_red_int2 = reinterpret_cast(sh_red); - int2* c_int2 = reinterpret_cast(C); - sh_red_int2[c_sh_wr + c_sh_wr_delta * i] = c_int2[true_idx]; - } - else - { - sh_red[c_sh_wr + c_sh_wr_delta * i] = C[true_idx]; - } - } - } - } - -#pragma unroll - for (int i = 0; i < (m_block_size_8 ? 2 : thread_m_blocks * 4); i++) - { - if (!first) - { - scalar_t* c_red_f16; - if constexpr (is_a_8bit) - { - int2 tmp = reinterpret_cast(sh_red)[c_sh_wr + i * c_sh_wr_delta]; - c_red_f16 = reinterpret_cast(&tmp); - } - else - { - int4 tmp = sh_red[c_sh_wr + i * c_sh_wr_delta]; - c_red_f16 = reinterpret_cast(&tmp); - } -#pragma unroll - for (int j = 0; j < 2 * (is_a_8bit ? 2 : 4); j++) - { - int delta = 0; - if constexpr (m_block_size_8) - { - delta = j % 2 == 1 ? -2 : 0; - } - reinterpret_cast(&frag_c)[(is_a_8bit ? 2 : 4) * 2 * 4 * (i / 4) + 4 * j + (i % 4) + delta] - += Dtype::num2float(c_red_f16[j]); - } - } - if (!last) - { - scalar_t c_f16[is_a_8bit ? 4 : 8]; -#pragma unroll - for (int j = 0; j < 2 * (is_a_8bit ? 2 : 4); j++) - { - int delta = 0; - if constexpr (m_block_size_8) - { - delta = j % 2 == 1 ? -2 : 0; - } - c_f16[j] = Dtype::float2num(reinterpret_cast( - &frag_c)[(is_a_8bit ? 2 : 4) * 2 * 4 * (i / 4) + 4 * j + (i % 4) + delta]); - } - - int c_idx; - if constexpr (m_block_size_8) - c_idx = c_gl_wr + i * c_gl_stride + (threadIdx.x % 8) / 4 * c_gl_wr_delta_i; - else - c_idx = c_gl_wr + c_gl_wr_delta_o * (i / 2) + c_gl_wr_delta_i * (i % 2); - if (c_idx / c_gl_stride < block_num_valid_tokens) - { - int64_t sorted_row = sh_block_sorted_ids[c_idx / c_gl_stride]; - int64_t true_idx = sorted_row * c_gl_stride + c_idx % c_gl_stride; - if constexpr (is_a_8bit) - { - int2* c_int2 = reinterpret_cast(C); - c_int2[true_idx] = *reinterpret_cast(c_f16); - } - else - { - C[true_idx] = *reinterpret_cast(c_f16); - } - } - } - } - }; - - // Globally reduce over threadblocks that compute the same column block. - // We use a tmp C buffer to reduce in full fp32 precision. - auto global_reduce_fp32 = [&](bool first = false, bool last = false) - { - constexpr int tb_m = thread_m_blocks * 16; - constexpr int tb_n = thread_n_blocks * 16; - - constexpr int c_size = tb_m * tb_n * sizeof(float) / 16; - - constexpr int active_threads = 32 * tb_n_warps; - bool is_th_active = threadIdx.x < active_threads; - - constexpr int num_floats = thread_m_blocks * (is_a_8bit ? 2 : 4) * 2 * 4; - constexpr int th_size = num_floats * sizeof(float) / 16; - - int c_cur_offset = locks_off * c_size; - - if (!is_th_active) - { - return; - } - - if (!first) - { - float* frag_c_ptr = reinterpret_cast(&frag_c); -#pragma unroll - for (int k = 0; k < th_size; k++) - { - if constexpr (m_block_size_8) - { - if (k % 2) - continue; - } - else - { - if (k / 8 * 16 + (threadIdx.x % 32) / 4 >= block_num_valid_tokens) - continue; - } - - sh_red[threadIdx.x] = C_tmp[c_cur_offset + active_threads * k + threadIdx.x]; - - float* sh_c_ptr = reinterpret_cast(&sh_red[threadIdx.x]); -#pragma unroll - for (int f = 0; f < 4; f++) - { - frag_c_ptr[k * 4 + f] += sh_c_ptr[f]; - } - } - } - - if (!last) - { - int4* frag_c_ptr = reinterpret_cast(&frag_c); -#pragma unroll - for (int k = 0; k < th_size; k++) - { - if constexpr (m_block_size_8) - { - if (k % 2) - continue; - } - else - { - if (k / 8 * 16 + (threadIdx.x % 32) / 4 >= block_num_valid_tokens) - continue; - } - - C_tmp[c_cur_offset + active_threads * k + threadIdx.x] = frag_c_ptr[k]; - } - } - }; - - // Write out the reduce final result in the correct layout. We only actually - // reshuffle matrix fragments in this step, the reduction above is performed - // in fragment layout. - auto write_result = [&](bool last) - { - int c_gl_stride = prob_n / 8; - constexpr int c_sh_stride = 2 * thread_n_blocks + 1; - int c_gl_wr_delta = c_gl_stride * (threads / (2 * thread_n_blocks)); - constexpr int c_sh_rd_delta = c_sh_stride * (threads / (2 * thread_n_blocks)); - - int c_gl_wr = c_gl_stride * (threadIdx.x / (2 * thread_n_blocks)) + (threadIdx.x % (2 * thread_n_blocks)); - c_gl_wr += (2 * thread_n_blocks) * slice_col; - int c_sh_wr; - if constexpr (m_block_size_8) - { - c_sh_wr = (8 * c_sh_stride) * ((threadIdx.x % 32) % 4 * 2) + (threadIdx.x % 32) / 4; - c_sh_wr += 64 * (threadIdx.x / 32); - } - else - { - c_sh_wr = (4 * c_sh_stride) * ((threadIdx.x % 32) / 4) + (threadIdx.x % 32) % 4; - c_sh_wr += (is_a_8bit ? 16 : 32) * (threadIdx.x / 32); - } - - int c_sh_rd = c_sh_stride * (threadIdx.x / (2 * thread_n_blocks)) + (threadIdx.x % (2 * thread_n_blocks)); - - // We first reorder in shared memory to guarantee the most efficient final - // global write patterns - auto write = [&](int idx, float c0, float c1, FragS& s, FragS& b_bias) - { - scalar_t2 res = Dtype::nums2num2(Dtype::float2num(c0), Dtype::float2num(c1)); - - // For per-column quantization we finally apply the scale here (only for - // 4-bit) - if constexpr (!has_act_order && group_blocks == -1 && !is_a_8bit - && (has_zp && dequant_skip_flop || !has_zp)) - { - scalar_t2 tmp_scale = s[0]; - if constexpr (m_block_size_8) - { - tmp_scale = Dtype::num2num2(reinterpret_cast(&s[0])[(threadIdx.x % 8) / 4]); - } - res = __hmul2(res, tmp_scale); - } - - if (!mul_topk_weights) - { - res = __hmul2(res, global_scale); - } - if (has_bias && last) - { - scalar_t2 tmp_bias = b_bias[0]; - if constexpr (m_block_size_8) - { - tmp_bias = Dtype::num2num2(reinterpret_cast(&b_bias[0])[(threadIdx.x % 8) / 4]); - } - res = __hadd2(res, tmp_bias); - } - - if constexpr (m_block_size_8) - { - ((scalar_t*) sh_red)[idx] = res.x; - ((scalar_t*) sh_red)[idx + 8 * c_sh_stride] = res.y; - } - else - { - ((scalar_t2*) sh_red)[idx] = res; - } - }; - - if (threadIdx.x / 32 < tb_n_warps) - { -#pragma unroll - for (int i = 0; i < thread_m_blocks; i++) - { -#pragma unroll - for (int j = 0; j < (is_a_8bit ? 2 : 4); j++) - { - if constexpr (m_block_size_8) - { - int wr = c_sh_wr + 16 * j; - write(wr, frag_c[i][j][0][0], frag_c[i][j][0][1], frag_s[j / 2][2 * (j % 2) + 0], - frag_bias[j / 2][2 * (j % 2) + 0]); - write(wr + 8, frag_c[i][j][0][2], frag_c[i][j][0][3], frag_s[j / 2][2 * (j % 2) + 1], - frag_bias[j / 2][2 * (j % 2) + 1]); - } - else - { - int wr = c_sh_wr + 8 * j; - write(wr + (4 * c_sh_stride) * 0 + 0, frag_c[i][j][0][0], frag_c[i][j][0][1], - frag_s[j / 2][2 * (j % 2) + 0], frag_bias[j / 2][2 * (j % 2) + 0]); - write(wr + (4 * c_sh_stride) * 8 + 0, frag_c[i][j][0][2], frag_c[i][j][0][3], - frag_s[j / 2][2 * (j % 2) + 0], frag_bias[j / 2][2 * (j % 2) + 0]); - write(wr + (4 * c_sh_stride) * 0 + 4, frag_c[i][j][1][0], frag_c[i][j][1][1], - frag_s[j / 2][2 * (j % 2) + 1], frag_bias[j / 2][2 * (j % 2) + 1]); - write(wr + (4 * c_sh_stride) * 8 + 4, frag_c[i][j][1][2], frag_c[i][j][1][3], - frag_s[j / 2][2 * (j % 2) + 1], frag_bias[j / 2][2 * (j % 2) + 1]); - } - } - c_sh_wr += 16 * (4 * c_sh_stride); - } - } - __syncthreads(); - -#pragma unroll - for (int i = 0; i < div_ceil(16 * thread_m_blocks, threads / (2 * thread_n_blocks)); i++) - { - int row = c_gl_wr / c_gl_stride; - if (row < block_num_valid_tokens) - { - int64_t sorted_row = sh_block_sorted_ids[row]; - int64_t true_idx = sorted_row * c_gl_stride + c_gl_wr % c_gl_stride; - scalar_t2 topk_weight_score; - if (mul_topk_weights) - topk_weight_score = sh_block_topk_weights[row]; - if (use_atomic_add && slice_count > 1 || mul_topk_weights) - { - scalar_t2* C_half2 = reinterpret_cast(&C[true_idx]); - scalar_t2* sh_red_half2 = reinterpret_cast(&sh_red[c_sh_rd]); - if (mul_topk_weights) - { -#pragma unroll - for (int a = 0; a < 4; a++) - { - sh_red_half2[a] = __hmul2(sh_red_half2[a], topk_weight_score); - } - } - - if (use_atomic_add && slice_count > 1) - { -#pragma unroll - for (int a = 0; a < 4; a++) - { - atomicAdd(&C_half2[a], sh_red_half2[a]); - } - } - else - { - C[true_idx] = *reinterpret_cast(sh_red_half2); - } - } - else - { - C[true_idx] = sh_red[c_sh_rd]; - } - c_gl_wr += c_gl_wr_delta; - c_sh_rd += c_sh_rd_delta; - } - } - __syncthreads(); - }; - - // Start global fetch and register load pipelines. - auto start_pipes = [&]() - { - -#pragma unroll - for (int i = 0; i < stages - 1; i++) - { - if (has_act_order && i == 0) - { - int last_g_idx = slice_k_start + stages * tb_k * 2; - if (last_g_idx >= prob_k) - { - last_g_idx = prob_k - 1; - } - fetch_act_order_scales_to_shared(true, g_idx[slice_k_start], g_idx[last_g_idx]); - } - - if constexpr (has_zp && group_blocks == -1) - { - if (i == 0) - { - fetch_col_zp_to_shared(); - if constexpr (!dequant_skip_flop) - { - fetch_col_scale_to_shared(); - } - } - } - fetch_to_shared(i, i, i < slice_iters); - } - - zero_accums(); - wait_for_stage(); - init_same_group(0); - fetch_to_registers(0, 0); - fetch_scales_to_registers(0, 0); - fetch_zp_to_registers(0, 0); - a_gl_rd_col += a_gl_rd_delta_o * (stages - 1); - if constexpr (has_act_order) - { - slice_k_start_shared_fetch += tb_k * (stages - 1); - } - }; - if (slice_iters) - { - start_pipes(); - } - - // Main loop. - while (slice_iters) - { - // We unroll over both the global fetch and the register load pipeline to - // ensure all shared memory accesses are static. Note that both pipelines - // have even length meaning that the next iteration will always start at - // index 0. - -#pragma unroll - for (int pipe = 0; pipe < stages;) - { -#pragma unroll - for (int k = 0; k < b_sh_wr_iters; k++) - { - fetch_to_registers(k + 1, pipe % stages); - fetch_scales_to_registers(k + 1, pipe); - fetch_zp_to_registers(k + 1, pipe); - if (k == b_sh_wr_iters - 2) - { - fetch_to_shared((pipe + stages - 1) % stages, pipe, slice_iters >= stages); - pipe++; - wait_for_stage(); - init_same_group(pipe % stages); - } - - if constexpr (!is_a_8bit) - { - matmul(k, pipe - (k >= b_sh_wr_iters - 2 ? 1 : 0)); - } - else - { - static_assert(group_blocks != 0 && group_blocks != 1); - matmul_a8(k); - } - } - slice_iters--; - if (slice_iters == 0) - { - break; - } - } - - a_gl_rd_col += a_gl_rd_delta_o * stages; - - if constexpr (has_act_order) - { - slice_k_start += tb_k * stages; - - if (slice_k_start < prob_k) - { - slice_k_start_shared_fetch += tb_k * stages; - int first_group_id = g_idx[slice_k_start]; - int last_g_idx = slice_k_start + stages * tb_k * 2; - if (last_g_idx >= prob_k) - { - last_g_idx = prob_k - 1; - } - int last_group_id = g_idx[last_g_idx]; - if (last_group_id >= sh_first_group_id + sh_num_groups) - { - fetch_act_order_scales_to_shared(false, first_group_id, last_group_id); - __syncthreads(); - } - } - } - - // Process results and, if necessary, proceed to the next column slice. - // While this pattern may not be the most readable, other ways of writing - // the loop seemed to noticeably worse performance after compilation. - if (slice_iters == 0) - { - // convert fp16 accum to fp32 for reduction - if constexpr (use_fp16_accum) - { -#pragma unroll - for (int i = 0; i < (thread_m_blocks * (is_a_8bit ? 2 : 4) * 2); i++) - { - float* frag_c_part_float = reinterpret_cast(frag_c) + i * 4; - scalar_t* frag_c_part_half = reinterpret_cast(frag_c_part_float); - -#pragma unroll - for (int i = 3; i >= 0; i--) - { - frag_c_part_float[i] = Dtype::num2float(frag_c_part_half[i]); - } - } - } - - if constexpr (is_a_8bit) - { - float frag_a_s[2 * thread_m_blocks]; - - for (int i = 0; i < 2 * thread_m_blocks; i++) - frag_a_s[i] = sh_a_s[i * 8 + (threadIdx.x % 32) / 4]; - -#pragma unroll - for (int j = 0; j < 2; j++) - { -#pragma unroll - for (int i = 0; i < thread_m_blocks; i++) - { -#pragma unroll - for (int g = 0; g < 4; g++) - { - float c_val = frag_c[i][j][0][g]; - float s_val = frag_a_s[i * 2 + g / 2]; - frag_c[i][j][0][g] = c_val * s_val; - } -#pragma unroll - for (int g = 0; g < 4; g++) - { - float c_val = frag_c[i][j][1][g]; - float s_val = frag_a_s[i * 2 + g / 2]; - frag_c[i][j][1][g] = c_val * s_val; - } - } - } - } - - cp_async_wait<0>(); - bool last = slice_idx == slice_count - 1; - // For per-column scales, we only fetch them here in the final step before - // write-out - if constexpr (!has_act_order && group_blocks == -1 && (has_zp && dequant_skip_flop || !has_zp)) - { - if ((last || use_atomic_add) || is_a_8bit) - { - if (s_sh_wr_pred) - { - cp_async4(&sh_s[s_sh_wr], &scales_ptr[s_gl_rd]); - } - cp_async_fence(); - } - } - - thread_block_reduce(); - - if (has_bias && last) - { - __syncthreads(); - cp_async4_pred(&sh_bias[bias_sh_wr], &b_bias_ptr[bias_gl_rd], threadIdx.x < 16 * thread_n_blocks / 8); - cp_async_fence(); - } - - if constexpr (!has_act_order && group_blocks == -1 && (has_zp && dequant_skip_flop || !has_zp || is_a_8bit)) - { - if constexpr (is_a_8bit) - { - cp_async_wait<0>(); - __syncthreads(); - if (threadIdx.x / 32 < tb_n_warps) - { - reinterpret_cast(&frag_s)[0] = sh_s[s_sh_rd + 0]; - } - } - else if (last || use_atomic_add) - { - cp_async_wait<0>(); - __syncthreads(); - if (threadIdx.x / 32 < tb_n_warps) - { - reinterpret_cast(&frag_s)[0] = sh_s[s_sh_rd + 0]; - reinterpret_cast(&frag_s)[1] = sh_s[s_sh_rd + 4]; - if constexpr (m_block_size_8) - { - int idx = (threadIdx.x / 4) % 2; - scalar_t2* frag_s_half2 = reinterpret_cast(frag_s); -#pragma unroll - for (int i = 0; i < 8; i++) - { - frag_s_half2[i] = Dtype::num2num2(reinterpret_cast(&frag_s_half2[i])[idx]); - } - } - } - } - } - - // For 8-bit channelwise, we apply the scale before the global reduction - // that converts the fp32 results to fp16 (so that we avoid possible - // overflow in fp16) - if constexpr (!has_act_order && group_blocks == -1 && is_a_8bit) - { -#pragma unroll - for (int j = 0; j < 2; j++) - { - float2 aa[2]; - aa[0] = Dtype::num22float2(frag_s[0][j * 2][0]); - aa[1] = Dtype::num22float2(frag_s[0][j * 2 + 1][0]); - -#pragma unroll - for (int i = 0; i < thread_m_blocks; i++) - { -#pragma unroll - for (int g = 0; g < 4; g++) - { - float scale = reinterpret_cast(&aa[0])[g % 2]; - frag_c[i][j][0][g] *= scale; - } - -#pragma unroll - for (int g = 0; g < 4; g++) - { - float scale = reinterpret_cast(&aa[1])[g % 2]; - frag_c[i][j][1][g] *= scale; - } - } - } - } - - if (slice_count > 1 && !use_atomic_add) - { - // only globally reduce if there is more than one block in a slice - barrier_acquire(&locks[locks_off], slice_idx); - if (use_fp32_reduce) - { - global_reduce_fp32(slice_idx == 0, last); - } - else - { - global_reduce_fp16(slice_idx == 0, last); - } - barrier_release(&locks[locks_off], last); - } - - if (has_bias && last) - { - cp_async_wait<0>(); - __syncthreads(); - reinterpret_cast(&frag_bias)[0] = sh_bias[bias_sh_rd]; - if constexpr (!is_a_8bit) - reinterpret_cast(&frag_bias)[1] = sh_bias[bias_sh_rd + 4]; - __syncthreads(); - } - - if (use_atomic_add && slice_count > 1 && slice_idx != 0) - wait_negative_and_add(&locks[locks_off]); - if (last || use_atomic_add) - // only the last block in a slice actually writes the result - write_result(last); - slice_row = 0; - if (!in_part2) - { - slice_col_par += gridDim.x; - } - else - { - slice_col_par++; - slice_col++; - } - is_first_matmul_in_slice = true; - init_slice(); - - if (slice_iters) - { - a_gl_rd_col = a_gl_rd_delta_o * slice_row + threadIdx.x % a_gl_rd_delta_o; - b_gl_rd = B_expert_off + b_gl_stride * (threadIdx.x / b_sh_stride) + (threadIdx.x % b_sh_stride); - b_gl_rd += b_sh_stride * slice_col + b_gl_rd_delta_o * slice_row; - - bias_gl_rd = (thread_n_blocks * 16 / 8) * slice_col + threadIdx.x; - // Update slice k/n for scales loading - if constexpr (has_act_order) - { - slice_k_start = tb_k * slice_row; - slice_k_finish = slice_k_start + tb_k * slice_iters; - slice_k_start_shared_fetch = slice_k_start; - slice_n_offset = act_s_col_tb_stride * slice_col; - } - else - { - if constexpr (group_blocks == -1) - { - s_gl_rd = s_sh_stride * slice_col + threadIdx.x; - zp_gl_rd = zp_sh_stride * slice_col + threadIdx.x; - } - else if constexpr (group_blocks >= thread_k_blocks) - { - s_gl_rd = s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) + s_sh_stride * slice_col - + threadIdx.x; - zp_gl_rd = zp_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) - + zp_sh_stride * slice_col + threadIdx.x; - } - else - { - s_gl_rd - = s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks + threadIdx.x / s_sh_stride) - + s_sh_stride * slice_col + threadIdx.x % s_sh_stride; - zp_gl_rd - = zp_gl_stride * ((thread_k_blocks * slice_row) / group_blocks + threadIdx.x / zp_sh_stride) - + zp_sh_stride * slice_col + threadIdx.x % zp_sh_stride; - } - } - start_pipes(); - } - } - } -} - -#endif - -} // namespace MARLIN_NAMESPACE_NAME diff --git a/cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_template.h b/cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_template.h deleted file mode 100644 index 9fdc9d2a6d38..000000000000 --- a/cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_template.h +++ /dev/null @@ -1,2059 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2022-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. - */ - -/* - * Modified by Neural Magic - * Copyright (C) Marlin.2024 Elias Frantar - * - * 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. - */ - -/* - * Adapted from https://github.com/IST-DASLab/marlin - */ - -#ifndef MARLIN_NAMESPACE_NAME -#define MARLIN_NAMESPACE_NAME marlin -#endif - -#include "marlin.cuh" - -#define STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t) \ - static_assert(std::is_same::value || std::is_same::value, \ - "only float16 and bfloat16 is supported"); - -namespace MARLIN_NAMESPACE_NAME -{ - -// Empty kernel stub for non-Hopper device passes; see marlin.cuh. -#if defined(__CUDA_ARCH__) && !(__CUDA_ARCH__ >= 900 && __CUDA_ARCH__ < 1000) - -template shared - // fetch pipeline - int group_blocks // number of consecutive 16x16 blocks - // with a separate quantization scale - // (implicit const: trailing NTTP) - > -__global__ void Marlin(MARLIN_KERNEL_PARAMS) -{ -} - -} // namespace MARLIN_NAMESPACE_NAME - -#else - -// Instruction for loading a full 16x16 matrix fragment of operand A from shared -// memory, directly in tensor core layout. -template -__device__ inline void ldsm(typename MarlinType::FragA& frag_a, void const* smem_ptr) -{ - uint32_t* a = reinterpret_cast(&frag_a); - uint32_t smem = static_cast(__cvta_generic_to_shared(smem_ptr)); - if constexpr (count == 4) - { - asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];\n" - : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) - : "r"(smem)); - } - else if constexpr (count == 2) - { - asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0,%1}, [%2];\n" : "=r"(a[0]), "=r"(a[1]) : "r"(smem)); - } - else if constexpr (count == 1) - { - asm volatile("ldmatrix.sync.aligned.m8n8.x1.shared.b16 {%0}, [%1];\n" : "=r"(a[0]) : "r"(smem)); - } - else - { - static_assert(count == 1 || count == 2 || count == 4, "invalid count"); - } -} - -// Multiply dequantized values by the corresponding quantization scale; used -// only for grouped quantization. -template -__device__ inline void scale( - typename MarlinType::FragB& frag_b, typename MarlinType::FragS& frag_s, int i) -{ - using scalar_t2 = typename MarlinType::scalar_t2; - scalar_t2 s = MarlinType::num2num2(reinterpret_cast(&frag_s)[i]); - frag_b[0] = __hmul2(frag_b[0], s); - frag_b[1] = __hmul2(frag_b[1], s); -} - -template -__device__ inline void scale_and_sub(typename MarlinType::FragB& frag_b, scalar_t s, scalar_t zp) -{ - using scalar_t2 = typename MarlinType::scalar_t2; - scalar_t2 s2 = MarlinType::num2num2(s); - scalar_t2 zp2 = MarlinType::num2num2(zp); - frag_b[0] = __hfma2(frag_b[0], s2, __hneg2(zp2)); - frag_b[1] = __hfma2(frag_b[1], s2, __hneg2(zp2)); -} - -template -__device__ inline void sub_zp( - typename MarlinType::FragB& frag_b, typename MarlinType::scalar_t2& frag_zp, int i) -{ - using scalar_t2 = typename MarlinType::scalar_t2; - scalar_t2 zp = MarlinType::num2num2(reinterpret_cast(&frag_zp)[i]); - frag_b[0] = __hsub2(frag_b[0], zp); - frag_b[1] = __hsub2(frag_b[1], zp); -} - -// Same as above, but for act_order (each K is multiplied individually) -template -__device__ inline void scale4(typename MarlinType::FragB& frag_b, - typename MarlinType::FragS& frag_s_1, typename MarlinType::FragS& frag_s_2, - typename MarlinType::FragS& frag_s_3, typename MarlinType::FragS& frag_s_4, int i) -{ - using scalar_t2 = typename MarlinType::scalar_t2; - - scalar_t2 s_val_1_2; - s_val_1_2.x = reinterpret_cast(&frag_s_1)[i]; - s_val_1_2.y = reinterpret_cast(&frag_s_2)[i]; - - scalar_t2 s_val_3_4; - s_val_3_4.x = reinterpret_cast(&frag_s_3)[i]; - s_val_3_4.y = reinterpret_cast(&frag_s_4)[i]; - - frag_b[0] = __hmul2(frag_b[0], s_val_1_2); - frag_b[1] = __hmul2(frag_b[1], s_val_3_4); -} - -// Given 2 floats multiply by 2 scales (halves) -template -__device__ inline void scale_float(float* c, typename MarlinType::FragS& s) -{ - scalar_t* s_ptr = reinterpret_cast(&s); - c[0] = __fmul_rn(c[0], MarlinType::num2float(s_ptr[0])); - c[1] = __fmul_rn(c[1], MarlinType::num2float(s_ptr[1])); -} - -// Wait until barrier reaches `count`, then lock for current threadblock. -__device__ inline void barrier_acquire(int* lock, int count) -{ - if (threadIdx.x == 0) - { - int state = -1; - do - // Guarantee that subsequent writes by this threadblock will be visible - // globally. - asm volatile("ld.global.acquire.gpu.b32 %0, [%1];\n" : "=r"(state) : "l"(lock)); - while (state != count); - } - __syncthreads(); -} - -// Release barrier and increment visitation count. -__device__ inline void barrier_release(int* lock, bool reset = false) -{ - __syncthreads(); - if (threadIdx.x == 0) - { - if (reset) - { - lock[0] = 0; - return; - } - int val = 1; - // Make sure that all writes since acquiring this barrier are visible - // globally, while releasing the barrier. - asm volatile("fence.acq_rel.gpu;\n"); - asm volatile("red.relaxed.gpu.global.add.s32 [%0], %1;\n" : : "l"(lock), "r"(val)); - } -} - -// Wait until value of lock to be negative, and then add 1 -__device__ inline void wait_negative_and_add(int* lock) -{ - if (threadIdx.x == 0) - { - int state = 0; - do - // Guarantee that subsequent writes by this threadblock will be visible - // globally. - asm volatile("ld.global.acquire.gpu.b32 %0, [%1];\n" : "=r"(state) : "l"(lock)); - while (state >= 0); - atomicAdd(lock, 1); - } - __syncthreads(); -} - -template shared - // fetch pipeline - int group_blocks // number of consecutive 16x16 blocks - // with a separate quantization scale - // (implicit const: trailing NTTP) - > -__global__ void Marlin(int4 const* __restrict__ A0, // fp16 input matrix of shape mxk - int4 const* __restrict__ B, // 4bit quantized weight matrix of shape kxn - int4* __restrict__ C0, // fp16 output buffer of shape mxn - int4* __restrict__ C_tmp, // fp32 tmp output buffer (for reduce) - int4 const* __restrict__ b_bias_ptr, - // float scales of input matrix, only used when is_a_8bit == true. - // shape (m,) - float const* __restrict__ a_scales_ptr, - // fp16 quantization scales. shape (k/groupsize, n) - int4 const* __restrict__ scales_ptr, - // fp16 global scale (for nvfp4// only) - uint16_t const* __restrict__ global_scale_ptr, - // 4bit packed zero-points of shape - // (k/groupsize, n/pack_factor) - int4 const* __restrict__ zp_ptr, - // int32 group indices of shape k - int const* __restrict__ g_idx, - int num_groups, // number of scale groups per output channel - int prob_m, // batch dimension m - int prob_n, // output dimension n - int prob_k, // reduction dimension k - int lda, // A.stride(0), equal to prob_k is A is contiguous - int* locks, // extra global storage for barrier synchronization - bool has_bias, - bool use_atomic_add, // whether to use atomic add to reduce - bool use_fp32_reduce, // whether to use fp32 global reduce - int max_shared_mem) -{ - // Each threadblock processes one "stripe" of the B matrix with (roughly) the - // same size, which might involve multiple column "slices" (of width 16 * - // `thread_n_blocks`). Stripes are defined as shown in the 3x3 matrix 5 SM - // example: - // 0 1 3 - // 0 2 3 - // 1 2 4 - // While this kind of partitioning makes things somewhat more complicated, it - // ensures good utilization of all SMs for many kinds of shape and GPU - // configurations, while requiring as few slow global cross-threadblock - // reductions as possible. - - constexpr bool use_fp16_accum = false; - - using scalar_t2 = typename MarlinType::scalar_t2; - using scalar_32bit_t = typename MarlinType::scalar_32bit_t; - - using c_scalar_t = scalar_t; - using c_scalar_t2 = scalar_t2; - - using FragA = typename MarlinType::FragA; - using FragB = typename MarlinType::FragB; - using FragC = typename MarlinType::FragC; - using FragS = typename MarlinType::FragS; - using FragZP = typename MarlinType::FragZP; - - int4 const* A = A0; - int4* C = C0; - - constexpr bool is_a_8bit = false; // BF16 is 16-bit - constexpr bool has_zp = false; // NVFP4 has no zero-points - // For NVFP4: dequant places bits, scale applied separately - constexpr bool dequant_skip_flop = true; - - c_scalar_t2 global_scale; - - { - uint16_t val = global_scale_ptr[0]; - global_scale = MarlinType::num2num2(*reinterpret_cast(&val)); - } - - constexpr bool has_act_order = group_blocks == 0; - constexpr int m_block_size = m_block_size_8 ? 8 : (16 * thread_m_blocks); - - extern __shared__ int4 sh[]; - float* sh_a_s = reinterpret_cast(sh); - int4* sh_new = sh + (is_a_8bit ? (4 * thread_m_blocks) : 0); - constexpr int pack_factor = 8; // 32 / 4 (FP4) - static_assert(thread_m_blocks == 1 || !m_block_size_8); - - // For larger GEMMs we run multiple batchsize 64 versions in parallel for a - // better partitioning with less reductions - int parallel = 1; - if (prob_m > m_block_size) - { - parallel = prob_m / m_block_size; - prob_m = m_block_size; - } - - int k_tiles = prob_k / 16 / thread_k_blocks; - int n_tiles = prob_n / 16 / thread_n_blocks; - - int global_mn_tiles = parallel * n_tiles; - int part2_mn_tiles = global_mn_tiles; - int part1_mn_iters = 0; - bool in_part2 = false; - - if (global_mn_tiles > gridDim.x) - { - part2_mn_tiles = global_mn_tiles % gridDim.x; - if (part2_mn_tiles * 3 <= gridDim.x) - part2_mn_tiles += gridDim.x; - part1_mn_iters = (global_mn_tiles - part2_mn_tiles) / gridDim.x; - } - - int iters = div_ceil(k_tiles * part2_mn_tiles, gridDim.x); - - if constexpr (!has_act_order && group_blocks != -1) - { - if (group_blocks >= thread_k_blocks) - { - // Ensure that the number of tiles in each stripe is a multiple of the - // groupsize; this avoids an annoying special case where a stripe starts - // in the middle of group. - iters = (group_blocks / thread_k_blocks) * div_ceil(iters, (group_blocks / thread_k_blocks)); - } - } - - int slice_row = 0; - int slice_col_par = blockIdx.x; - int slice_col; - int slice_iters = k_tiles; // number of threadblock tiles in the current slice - // total number of active threadblocks in the current slice - int slice_count = 1; - // index of threadblock in current slice; numbered bottom to top - int slice_idx = 0; - - int par_id = 0; - int locks_off = 0; - - if (part2_mn_tiles >= gridDim.x) - { - // when part2_mn_tiles >= sms - // then there are at most $sms$ conflict tile blocks - locks_off = blockIdx.x; - } - else - { - locks_off = (iters * blockIdx.x) / k_tiles - 1; - } - - // Compute all information about the current slice which is required for - // synchronization. - bool first_init = true; - auto init_part2_slice = [&]() - { - slice_iters = iters * (blockIdx.x + 1) - (k_tiles * slice_col_par + slice_row); - if (slice_iters < 0 || slice_col_par >= part2_mn_tiles) - slice_iters = 0; - if (slice_iters == 0) - return; - if (slice_row + slice_iters > k_tiles) - slice_iters = k_tiles - slice_row; - slice_count = 1; - slice_idx = 0; - int col_first = iters * div_ceil(k_tiles * slice_col_par, iters); - if (col_first <= k_tiles * (slice_col_par + 1)) - { - int col_off = col_first - k_tiles * slice_col_par; - slice_count = div_ceil(k_tiles - col_off, iters); - if (col_off > 0) - slice_count++; - int delta_first = iters * blockIdx.x - col_first; - if (delta_first < 0 || (col_off == 0 && delta_first == 0)) - slice_idx = slice_count - 1; - else - { - slice_idx = slice_count - 1 - delta_first / iters; - if (col_off > 0) - slice_idx--; - } - } - if (part2_mn_tiles >= gridDim.x) - { - if (slice_count > 1 && slice_idx == slice_count - 1) - { - locks_off++; - } - } - else - { - locks_off++; - } - - if (first_init && use_atomic_add && slice_count > 1 && slice_idx == 0) - { - constexpr int threads_per_m = 16 * thread_n_blocks / 8; - int m_per_thread = div_ceil(thread_m_blocks * 16, threads / threads_per_m); - if (m_block_size_8) - m_per_thread = div_ceil(8, threads / threads_per_m); - for (int i = 0; i < m_per_thread; i++) - { - int row = threads / threads_per_m * i + threadIdx.x / threads_per_m; - if (row < prob_m) - { - int col = slice_col * 16 * thread_n_blocks / 8 + threadIdx.x % threads_per_m; - C[row * prob_n / 8 + col] = {0, 0, 0, 0}; - } - } - // After write zero to output, write a negative value to lock. - // Every SM that processes the same slice would wait for - // the negative value, and then atomicAdd 1 to it. - // After all SMs are processed, the lock value would back to 0 again. - __syncthreads(); - if (threadIdx.x == 0) - locks[locks_off] = 1 - slice_count; - } - - if (slice_col == n_tiles) - { - A += 16 * thread_m_blocks * lda / (is_a_8bit ? 16 : 8); - C += 16 * thread_m_blocks * prob_n / 8; - slice_col = 0; - par_id++; - } - if (is_a_8bit && (first_init || slice_col == 0)) - { - __syncthreads(); - int a_s_gl_rd = par_id * 16 * thread_m_blocks + threadIdx.x; - cp_async1_ca_pred(&sh_a_s[threadIdx.x], &a_scales_ptr[a_s_gl_rd], threadIdx.x < prob_m); - } - }; - - auto init_part1_slice = [&]() - { - if (part1_mn_iters) - { - part1_mn_iters--; - par_id = slice_col_par / n_tiles; - slice_col = slice_col_par % n_tiles; - slice_iters = k_tiles; - A = A0 + 16 * thread_m_blocks / (is_a_8bit ? 16 : 8) * par_id * lda; - C = C0 + 16 * thread_m_blocks / 8 * par_id * prob_n; - if (is_a_8bit) - { - __syncthreads(); - int a_s_gl_rd = par_id * 16 * thread_m_blocks + threadIdx.x; - cp_async1_ca_pred(&sh_a_s[threadIdx.x], &a_scales_ptr[a_s_gl_rd], threadIdx.x < prob_m); - } - } - }; - - auto init_slice = [&]() - { - if (!in_part2 && !part1_mn_iters) - { - in_part2 = true; - slice_col_par = (iters * blockIdx.x) / k_tiles; - slice_row = (iters * blockIdx.x) % k_tiles; - slice_col = (slice_col_par + global_mn_tiles - part2_mn_tiles) % n_tiles; - par_id = (slice_col_par + global_mn_tiles - part2_mn_tiles) / n_tiles; - A = A0 + 16 * thread_m_blocks / (is_a_8bit ? 16 : 8) * par_id * lda; - C = C0 + 16 * thread_m_blocks / 8 * par_id * prob_n; - } - if (!in_part2) - { - init_part1_slice(); - } - else - { - init_part2_slice(); - first_init = false; - } - }; - - init_slice(); - - // A sizes/strides - - // stride of the A matrix in global memory - int a_gl_stride = lda / (is_a_8bit ? 16 : 8); - // stride of an A matrix tile in shared memory - constexpr int a_sh_stride = 16 * thread_k_blocks / (is_a_8bit ? 16 : 8); - // delta between subsequent A tiles in global memory - constexpr int a_gl_rd_delta_o = 16 * thread_k_blocks / (is_a_8bit ? 16 : 8); - // between subsequent accesses within a tile - int a_gl_rd_delta_i = a_gl_stride * (threads / a_gl_rd_delta_o); - // between shared memory writes - constexpr int a_sh_wr_delta = a_sh_stride * (threads / a_gl_rd_delta_o); - // within a shared memory tile - constexpr int a_sh_rd_delta_i = a_sh_stride * 16; - // overall size of a tile - constexpr int a_sh_stage = a_sh_stride * m_block_size; - // number of shared write iterations for a tile - constexpr int a_sh_wr_iters = div_ceil(a_sh_stage, a_sh_wr_delta); - - // B sizes/strides - int b_gl_stride = 16 * prob_n / (pack_factor * (is_a_8bit ? 2 : 4)); - constexpr int b_sh_stride = ((thread_n_blocks * 16) * 16 / pack_factor) / (is_a_8bit ? 2 : 4); - constexpr int b_thread_vecs = 1; // FP4: 1 vec per thread - constexpr int b_sh_stride_threads = b_sh_stride / b_thread_vecs; - - int b_gl_rd_delta_o = b_gl_stride * thread_k_blocks / (is_a_8bit ? 2 : 1); - constexpr int b_sh_wr_delta = threads * b_thread_vecs; - constexpr int b_sh_stage = b_sh_stride * thread_k_blocks / (is_a_8bit ? 2 : 1); - constexpr int b_sh_wr_iters = b_sh_stage / b_sh_wr_delta; - - // Scale sizes/strides without act_order - int s_gl_stride = prob_n / (16 /* NVFP4 scale stride */); - constexpr int s_sh_stride = 16 * thread_n_blocks / (16 /* NVFP4 scale stride */); - constexpr int s_tb_groups - = !has_act_order && group_blocks != -1 && group_blocks < thread_k_blocks ? thread_k_blocks / group_blocks : 1; - constexpr int s_sh_stage = s_tb_groups * s_sh_stride; - int s_gl_rd_delta = s_gl_stride; - - // Scale size/strides with act_order - constexpr int tb_k = 16 * thread_k_blocks; - constexpr int g_idx_stage = has_act_order ? (tb_k * sizeof(int)) / 16 : 0; - // constexpr int act_s_row_stride = 1; - // int act_s_col_stride = act_s_row_stride * num_groups; - constexpr int act_s_max_num_groups = 32; - int act_s_col_stride = 1; - int act_s_col_warp_stride = act_s_col_stride * 8; - - constexpr int tb_n_warps = thread_n_blocks / (is_a_8bit ? 2 : 4); - int act_s_col_tb_stride = act_s_col_warp_stride * tb_n_warps; - - // Zero-points sizes/strides - int zp_gl_stride = (prob_n / pack_factor) / 4; - constexpr int zp_sh_stride = ((16 * thread_n_blocks) / pack_factor) / 4; - constexpr int zp_tb_groups = s_tb_groups; - constexpr int zp_sh_stage = has_zp ? zp_tb_groups * zp_sh_stride : 0; - int zp_gl_rd_delta = zp_gl_stride; - - // Global A read index of current thread. - int a_gl_rd = a_gl_stride * (threadIdx.x / a_gl_rd_delta_o) + (threadIdx.x % a_gl_rd_delta_o); - a_gl_rd += a_gl_rd_delta_o * slice_row; - // Shared write index of current thread. - int a_sh_wr = a_sh_stride * (threadIdx.x / a_gl_rd_delta_o) + (threadIdx.x % a_gl_rd_delta_o); - // Shared read index. - int a_sh_rd = a_sh_stride * ((threadIdx.x % 32) % (16 / (m_block_size_8 ? 2 : 1))) - + (threadIdx.x % 32) / (16 / (m_block_size_8 ? 2 : 1)); - a_sh_rd += 2 * ((threadIdx.x / 32) / tb_n_warps) * b_sh_wr_iters; - - int b_gl_rd; - if (threads <= b_sh_stride) - { - b_gl_rd = threadIdx.x; - } - else - { - b_gl_rd = b_gl_stride * (threadIdx.x / b_sh_stride) + (threadIdx.x % b_sh_stride); - } - - b_gl_rd += b_sh_stride * slice_col; - b_gl_rd += b_gl_rd_delta_o * slice_row; - auto b_sh_rd = threadIdx.x * b_thread_vecs; - b_sh_rd += b_sh_rd / b_sh_stride * (b_sh_stride * (b_sh_wr_iters - 1)); - - // For act_order - int slice_k_start = tb_k * slice_row; - int slice_k_finish = slice_k_start + tb_k * slice_iters; - int slice_k_start_shared_fetch = slice_k_start; - int slice_n_offset = act_s_col_tb_stride * slice_col; - - // No act_order - int s_gl_rd; - if constexpr (!has_act_order) - { - if constexpr (group_blocks == -1) - { - s_gl_rd = s_sh_stride * slice_col + threadIdx.x; - } - else if constexpr (group_blocks >= thread_k_blocks) - { - s_gl_rd - = s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) + s_sh_stride * slice_col + threadIdx.x; - } - else - { - s_gl_rd = s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks + threadIdx.x / s_sh_stride) - + s_sh_stride * slice_col + threadIdx.x % s_sh_stride; - } - } - auto s_sh_wr = threadIdx.x; - bool s_sh_wr_pred = threadIdx.x < s_sh_stage; - - // Zero-points - int zp_gl_rd; - if constexpr (has_zp) - { - if constexpr (group_blocks == -1) - { - zp_gl_rd = zp_sh_stride * slice_col + threadIdx.x; - } - else if constexpr (group_blocks >= thread_k_blocks) - { - zp_gl_rd = zp_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) + zp_sh_stride * slice_col - + threadIdx.x; - } - else - { - zp_gl_rd = zp_gl_stride * ((thread_k_blocks * slice_row) / group_blocks + threadIdx.x / zp_sh_stride) - + zp_sh_stride * slice_col + threadIdx.x % zp_sh_stride; - } - } - auto zp_sh_wr = threadIdx.x; - bool zp_sh_wr_pred = zp_sh_stage > 0 && threadIdx.x < zp_sh_stage; - - // We use a different scale layout for grouped and column-wise quantization as - // we scale a `half2` tile in column-major layout in the former and in - // row-major in the latter case. - int s_sh_rd; - if constexpr (is_a_8bit) - { - s_sh_rd = 4 * ((threadIdx.x / 32) % tb_n_warps) + (threadIdx.x % 4); - } - else if constexpr (group_blocks != -1) - s_sh_rd = 8 * ((threadIdx.x / 32) % tb_n_warps) + (threadIdx.x % 32) / 4; - else if constexpr (group_blocks == -1 && (m_block_size_8 || (has_zp && !dequant_skip_flop))) - s_sh_rd = 8 * ((threadIdx.x / 32) % tb_n_warps) + (threadIdx.x % 32) / 8; - else - s_sh_rd = 8 * ((threadIdx.x / 32) % tb_n_warps) + (threadIdx.x % 32) % 4; - - int bias_sh_rd; - if constexpr (m_block_size_8) - { - bias_sh_rd = 8 * ((threadIdx.x / 32) % tb_n_warps) + (threadIdx.x % 32) / 8; - } - else - { - bias_sh_rd = (is_a_8bit ? 4 : 8) * ((threadIdx.x / 32) % tb_n_warps) + (threadIdx.x % 32) % 4; - } - - int bias_sh_wr = threadIdx.x; - int bias_gl_rd = (thread_n_blocks * 16 / 8) * slice_col + threadIdx.x; - - // Zero-points have the same read layout as the scales - // (without column-wise case) - constexpr int num_col_threads = 8; - constexpr int num_row_threads = 4; - constexpr int num_ints_per_thread = 8 / pack_factor; - int zp_sh_rd; - if constexpr (has_zp) - { - if (is_a_8bit) - { - zp_sh_rd = num_ints_per_thread * num_col_threads * ((threadIdx.x / 32) % tb_n_warps / 2) - + num_ints_per_thread * ((threadIdx.x % 32) / num_row_threads); - } - else - { - zp_sh_rd = num_ints_per_thread * num_col_threads * ((threadIdx.x / 32) % tb_n_warps) - + num_ints_per_thread * ((threadIdx.x % 32) / num_row_threads); - } - } - - // Precompute which thread should not read memory in which iterations; this is - // needed if there are more threads than required for a certain tilesize or - // when the batchsize is not a multiple of 16. - bool a_sh_wr_pred[a_sh_wr_iters]; -#pragma unroll - for (int i = 0; i < a_sh_wr_iters; i++) - a_sh_wr_pred[i] = a_sh_wr_delta * i + a_sh_wr < a_sh_stride * prob_m; - - // To ensure that writing and reading A tiles to/from shared memory, the - // latter in fragment format, is fully bank conflict free, we need to use a - // rather fancy XOR-based layout. The key here is that neither reads nor - // writes of the 16-byte `int4` blocks of 8 consecutive threads involve the - // same shared memory banks. Further, it seems (based on NSight-Compute) that - // each warp must also write a consecutive memory segment? - auto transform_a = [&](int i) - { - int row = i / a_gl_rd_delta_o; - return a_gl_rd_delta_o * row + (i % a_gl_rd_delta_o) ^ (row % 8); - }; - // Since the computation of this remapping is non-trivial and, due to our main - // loop unrolls, all shared memory accesses are static, we simply precompute - // both transformed reads and writes. - int a_sh_wr_trans[a_sh_wr_iters]; -#pragma unroll - for (int i = 0; i < a_sh_wr_iters; i++) - a_sh_wr_trans[i] = transform_a(a_sh_wr_delta * i + a_sh_wr); - int a_sh_rd_trans[b_sh_wr_iters][thread_m_blocks]; -#pragma unroll - for (int i = 0; i < b_sh_wr_iters; i++) - { -#pragma unroll - for (int j = 0; j < thread_m_blocks; j++) - a_sh_rd_trans[i][j] = transform_a(2 * i + a_sh_rd_delta_i * j + a_sh_rd); - } - - // Since B-accesses have non-constant stride they have to be computed at - // runtime; we break dependencies between subsequent accesses with a tile by - // maintining multiple pointers (we have enough registers), a tiny - // optimization. - - // Shared memory storage for global fetch pipelines. - constexpr int sh_red_size = (2 * thread_n_blocks + 1) * 16 * thread_m_blocks; - constexpr int sh_b_size = stages * b_sh_stage; - int4* sh_b = sh_new; - int4* sh_red = sh_new; - constexpr int sh_size_b_red_min = (sh_red_size < sh_b_size ? sh_red_size : sh_b_size); - constexpr int sh_size_b_red_max = (sh_red_size > sh_b_size ? sh_red_size : sh_b_size); - constexpr int sh_bias_size = (thread_n_blocks * 16 / 8); - constexpr int sh_b_red_bias_size = sh_size_b_red_max > (sh_size_b_red_min + sh_bias_size) - ? sh_size_b_red_max - : (sh_size_b_red_min + sh_bias_size); - - int4* sh_bias = sh_new + sh_size_b_red_min; - int4* sh_g_idx = sh_new + sh_b_red_bias_size; - int4* sh_zp = sh_g_idx + (stages * g_idx_stage); - constexpr int sh_s_size = has_act_order ? (act_s_max_num_groups * s_sh_stride) : (stages * s_sh_stage); - int4* sh_s = sh_zp + (stages * zp_sh_stage); - int4* sh_a = sh_s + sh_s_size; - - // Register storage for double buffer of shared memory reads. - FragA frag_a[2][thread_m_blocks]; - I4 frag_b_quant[2][b_thread_vecs]; - FragC frag_c[thread_m_blocks][is_a_8bit ? 2 : 4][2]; - FragC frag_c_tmp[thread_m_blocks][is_a_8bit ? 2 : 4][2]; - FragS frag_s[2][4]; // No act-order - FragS frag_bias[2][4]; - FragS act_frag_s[2][4][4]; // For act-order - int frag_qzp[2][num_ints_per_thread]; // Zero-points - FragZP frag_zp; // Zero-points in fp16 - FragZP frag_zpf[2]; // Zero-points in fp16 in HQQ - - if constexpr (is_a_8bit) - { -#pragma unroll - for (int j = 0; j < 2; j++) - { -#pragma unroll - for (int i = 0; i < thread_m_blocks; i++) - { -#pragma unroll - for (int g = 0; g < 4; g++) - { - frag_c_tmp[i][j][0][g] = 0.0f; - } - -#pragma unroll - for (int g = 0; g < 4; g++) - { - frag_c_tmp[i][j][1][g] = 0.0f; - } - } - } - } - - // Zero accumulators. - auto zero_accums = [&]() - { -#pragma unroll - for (int i = 0; i < thread_m_blocks * 4 * 2 * 4; i++) - reinterpret_cast(frag_c)[i] = 0; - }; - - int sh_first_group_id = -1; - int sh_num_groups = -1; - - auto fetch_act_order_scales_to_shared = [&](bool is_async, int first_group_id, int last_group_id) - { - sh_first_group_id = first_group_id; - sh_num_groups = last_group_id - first_group_id + 1; - - if (sh_num_groups > act_s_max_num_groups) - { - sh_num_groups = act_s_max_num_groups; - } - - if (sh_first_group_id + sh_num_groups > num_groups) - { - sh_num_groups = num_groups - sh_first_group_id; - } - - int row_offset = first_group_id * s_gl_stride; - - if (is_async) - { - for (int i = 0; i < sh_num_groups; i++) - { - if (threadIdx.x < s_sh_stride) - { - cp_async4_pred(&sh_s[(i * s_sh_stride) + threadIdx.x], - &scales_ptr[row_offset + (i * s_gl_stride) + slice_n_offset + threadIdx.x]); - } - } - } - else - { - for (int i = 0; i < sh_num_groups; i++) - { - if (threadIdx.x < s_sh_stride) - { - sh_s[(i * s_sh_stride) + threadIdx.x] - = scales_ptr[row_offset + (i * s_gl_stride) + slice_n_offset + threadIdx.x]; - } - } - } - }; - // Asynchronously fetch the next A, B and s tile from global to the next - // shared memory pipeline location. - auto fetch_to_shared = [&](int pipe, int a_off, bool pred = true) - { - if (pred) - { - int4* sh_a_stage = sh_a + a_sh_stage * pipe; -#pragma unroll - for (int i = 0; i < a_sh_wr_iters; i++) - { - cp_async4_pred(&sh_a_stage[a_sh_wr_trans[i]], - &A[a_gl_rd_delta_i * i + a_gl_rd + a_gl_rd_delta_o * a_off], a_sh_wr_pred[i]); - } - int4* sh_b_stage = sh_b + b_sh_stage * pipe; -#pragma unroll - for (int i = 0; i < (b_sh_wr_iters * b_thread_vecs); i++) - { - constexpr int count = div_ceil(b_sh_stride, threads); - int b_gl_idx - = b_gl_rd + (i % count) * threads + b_gl_stride * (i / count) * div_ceil(threads, b_sh_stride); - - cp_async4(&sh_b_stage[threads * i + threadIdx.x], &B[b_gl_idx]); - } - - b_gl_rd += b_gl_rd_delta_o; - - if constexpr (has_act_order) - { - // Fetch g_idx thread-block portion - int full_pipe = a_off; - int cur_k = slice_k_start_shared_fetch + tb_k * full_pipe; - if (cur_k < prob_k && cur_k < slice_k_finish) - { - int4* sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe; - - int4 const* cur_g_idx_stage_ptr = reinterpret_cast(&g_idx[cur_k]); - - if (threadIdx.x < g_idx_stage) - { - cp_async4_pred(&sh_g_idx_stage[threadIdx.x], &cur_g_idx_stage_ptr[threadIdx.x]); - } - } - } - else - { - if constexpr (group_blocks != -1) - { - int4* sh_s_stage = sh_s + s_sh_stage * pipe; - - // Only fetch scales if this tile starts a new group - if (pipe % div_ceil(group_blocks, thread_k_blocks) == 0) - { - if (s_sh_wr_pred) - { - cp_async4(&sh_s_stage[s_sh_wr], &scales_ptr[s_gl_rd]); - } - s_gl_rd += s_gl_rd_delta * s_tb_groups; - } - } - - if constexpr (has_zp && group_blocks != -1) - { - int4* sh_zp_stage = sh_zp + zp_sh_stage * pipe; - - // Only fetch zero points if this tile starts a new group - if (pipe % div_ceil(group_blocks, thread_k_blocks) == 0) - { - if (zp_sh_wr_pred) - { - cp_async4(&sh_zp_stage[zp_sh_wr], &zp_ptr[zp_gl_rd]); - } - zp_gl_rd += zp_gl_rd_delta * zp_tb_groups; - } - } - } - } - // Insert a fence even when we are winding down the pipeline to ensure that - // waiting is also correct at this point. - cp_async_fence(); - }; - - auto fetch_col_zp_to_shared = [&]() - { - if (zp_sh_wr_pred) - { - cp_async4(&sh_zp[zp_sh_wr], &zp_ptr[zp_gl_rd]); - } - }; - - auto fetch_col_scale_to_shared = [&]() - { - if (s_sh_wr_pred) - { - cp_async4(&sh_s[s_sh_wr], &scales_ptr[s_gl_rd]); - } - }; - - // Wait until the next thread tile has been loaded to shared memory. - auto wait_for_stage = [&]() - { - // We only have `stages - 2` active fetches since we are double buffering - // and can only issue the next fetch when it is guaranteed that the previous - // shared memory load is fully complete (as it may otherwise be - // overwritten). - cp_async_wait(); - __syncthreads(); - }; - - // Load the next sub-tile from the current location in the shared memory pipe - // into the current register buffer. - auto fetch_to_registers = [&](int k, int pipe) - { - int4* sh_a_stage = sh_a + a_sh_stage * pipe; -#pragma unroll - for (int i = 0; i < thread_m_blocks; i++) - ldsm(frag_a[k % 2][i], &sh_a_stage[a_sh_rd_trans[k % b_sh_wr_iters][i]]); - int4* sh_b_stage = sh_b + b_sh_stage * pipe; - -#pragma unroll - for (int i = 0; i < b_thread_vecs; i++) - { - frag_b_quant[k % 2][i] - = *reinterpret_cast(&sh_b_stage[b_sh_stride * (k % b_sh_wr_iters) + b_sh_rd + i]); - } - }; - - bool is_same_group[stages]; - int same_group_id[stages]; - - auto init_same_group = [&](int pipe) - { - if constexpr (!has_act_order) - { - return; - } - - int4* sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe; - int* sh_g_idx_int_ptr = reinterpret_cast(sh_g_idx_stage); - - int group_id_1 = sh_g_idx_int_ptr[0]; - int group_id_2 = sh_g_idx_int_ptr[tb_k - 1]; - - is_same_group[pipe] = group_id_1 == group_id_2; - same_group_id[pipe] = group_id_1; - }; - - auto fetch_scales_to_registers = [&](int k, int full_pipe) - { - int pipe = full_pipe % stages; - using IT1 = typename std::conditional_t; - using IT0 = typename std::conditional_t; - constexpr int group_blocks2 = div_ceil(group_blocks, is_a_8bit ? 2 : 1); - - if constexpr (!has_act_order) - { - // No act-order case - if constexpr (group_blocks == -1) - { - // load only when starting a new slice - if (k == 0 && full_pipe == 0 && dequant_skip_flop) - { - reinterpret_cast(&frag_s)[0] = sh_s[s_sh_rd]; - reinterpret_cast(&frag_s)[1] = sh_s[s_sh_rd + 4]; - } - } - else if constexpr (group_blocks != -1) - { - if constexpr (group_blocks >= thread_k_blocks) - { - constexpr int g = group_blocks / thread_k_blocks; - if (pipe % g == 0) - { - if (k % b_sh_wr_iters == 0) - { - int4* sh_s_stage = sh_s + s_sh_stage * (g * (pipe / g)); - reinterpret_cast(&frag_s[k % 2])[0] = sh_s_stage[s_sh_rd]; - } - else - { - reinterpret_cast(&frag_s[1])[0] = reinterpret_cast(&frag_s[0])[0]; - } - } - } - else if (group_blocks2 < b_sh_wr_iters || k % b_sh_wr_iters == 0) - { - auto warp_id = threadIdx.x / 32; - int warp_row = warp_id / tb_n_warps; - - int k_blocks = b_sh_wr_iters * warp_row + k % b_sh_wr_iters; - int cur_group_id = k_blocks / group_blocks2; - - int4* sh_s_stage = sh_s + s_sh_stage * pipe; - - reinterpret_cast(&frag_s[k % 2])[0] - = reinterpret_cast(sh_s_stage)[s_sh_rd + cur_group_id * (2 * s_sh_stride)]; - } - else if (group_blocks >= b_sh_wr_iters) - { - reinterpret_cast(&frag_s[1])[0] = reinterpret_cast(&frag_s[0])[0]; - } - } - - return; - } - - // Act-order case - - // Determine K of the "current" thread-block - int cur_k = slice_k_start + tb_k * full_pipe; - if (cur_k >= prob_k || cur_k >= slice_k_finish) - { - return; - } - - // Reset (to current thread-block) since we read g_idx portion from the - // shared memory - cur_k = 0; - - // Progress to current iteration - cur_k += k % b_sh_wr_iters; - - // Determine "position" inside the thread-block (based on warp and - // thread-id) - auto warp_id = threadIdx.x / 32; - int warp_row = warp_id / tb_n_warps; - int warp_col = warp_id % tb_n_warps; - - cur_k += warp_row * 16 * b_sh_wr_iters; - - auto th_id = threadIdx.x % 32; - cur_k += (th_id % 4) * 2; // Due to tensor-core layout for fp16 B matrix - - int s_col_shift = - /*slice_n_offset +*/ (act_s_col_warp_stride * warp_col) + (th_id / 4) * act_s_col_stride; - - if (is_same_group[pipe]) - { - if (k % 2 == 0) - { - *(reinterpret_cast(&(act_frag_s[k % 2][0][0]))) - = sh_s[(same_group_id[pipe] - sh_first_group_id) * s_sh_stride + s_col_shift]; - } - else - { - *(reinterpret_cast(&(act_frag_s[k % 2][0][0]))) - = *(reinterpret_cast(&(act_frag_s[(k - 1) % 2][0][0]))); - } - - for (int i = 1; i < 4; i++) - { - *(reinterpret_cast(&(act_frag_s[k % 2][i][0]))) - = *(reinterpret_cast(&(act_frag_s[k % 2][0][0]))); - } - return; - } - - int4* sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe; - int* sh_g_idx_int_ptr = reinterpret_cast(sh_g_idx_stage); - - constexpr int k_frag_offsets[4] = {0, 1, 8, 9}; // Tensor core offsets per thread - -#pragma unroll - for (int i = 0; i < 4; i++) - { - int actual_k = cur_k + k_frag_offsets[i]; - - int group_id = sh_g_idx_int_ptr[actual_k]; - int rel_group_id = group_id - sh_first_group_id; - - *(reinterpret_cast(&(act_frag_s[k % 2][i][0]))) = sh_s[rel_group_id * s_sh_stride + s_col_shift]; - } - }; - - auto fetch_zp_to_registers = [&](int k, int full_pipe) - { - // This code does not handle group_blocks == 0, - // which signifies act_order. - // has_zp implies AWQ, which doesn't have act_order, - static_assert(!has_zp || group_blocks != 0); - - if constexpr (has_zp) - { - int pipe = full_pipe % stages; - - if constexpr (group_blocks == -1) - { - // load only when starting a new slice - if (k == 0 && full_pipe == 0 || is_a_8bit) - { -#pragma unroll - for (int i = 0; i < num_ints_per_thread; i++) - { - frag_qzp[k % 2][i] = (reinterpret_cast(sh_zp))[zp_sh_rd + i]; - } - } - } - else if constexpr (group_blocks >= thread_k_blocks) - { - constexpr int g = group_blocks / thread_k_blocks; - if (pipe % g == 0 && k % b_sh_wr_iters == 0 || is_a_8bit) - { - int4* sh_zp_stage = sh_zp + zp_sh_stage * (g * (pipe / g)); -#pragma unroll - for (int i = 0; i < num_ints_per_thread; i++) - { - frag_qzp[k % 2][i] = (reinterpret_cast(sh_zp_stage))[zp_sh_rd + i]; - } - } - } - else - { - auto warp_id = threadIdx.x / 32; - - int warp_row = warp_id / tb_n_warps; - - int k_blocks = b_sh_wr_iters * warp_row + k % b_sh_wr_iters; - int cur_group_id = k_blocks / div_ceil(group_blocks, is_a_8bit ? 2 : 1); - - int4* sh_zp_stage = sh_zp + zp_sh_stage * pipe; - - sh_zp_stage += cur_group_id * zp_sh_stride; - -#pragma unroll - for (int i = 0; i < num_ints_per_thread; i++) - { - frag_qzp[k % 2][i] = (reinterpret_cast(sh_zp_stage))[zp_sh_rd + i]; - } - } - } - }; - - auto dequant_data = [&](int q, scalar_32bit_t* frag_b_ptr, int zp = 0) - { - // 16-bit != 4-bit, always dequant - dequant_fp4(q, frag_b_ptr); - }; - - // Execute the actual tensor core matmul of a sub-tile. - bool is_first_matmul_in_slice = true; - auto matmul = [&](int k, int pipe) - { - if (is_a_8bit) - return; - int k2 = k % 2; - constexpr int g = group_blocks > 0 ? div_ceil(group_blocks, thread_k_blocks) : 1; - const bool is_new_zp = (group_blocks == 0) - || ((group_blocks > 0) && (group_blocks < b_sh_wr_iters || k == 0)) && (pipe % g == 0) - || (group_blocks == -1 && is_first_matmul_in_slice); - if constexpr (has_zp) - { - if (is_new_zp) - { - if constexpr (group_blocks == -1) - is_first_matmul_in_slice = false; - int zp_quant_0, zp_quant_1; - - // FP4: 4-bit zero-points - zp_quant_0 = frag_qzp[k2][0]; - zp_quant_1 = zp_quant_0 >> 8; - - dequant_data(zp_quant_0, reinterpret_cast(&frag_zp)); - dequant_data(zp_quant_1, reinterpret_cast(&frag_zp) + 2); - } - } - - // NVFP4: dequant FP8 scales to BF16 - { - int s_quant_0 = reinterpret_cast(frag_s[k2])[0]; - int s_quant_1 = reinterpret_cast(frag_s[k2])[1]; - - dequant_fp8_scales(s_quant_0, reinterpret_cast(&frag_s[k2])); - dequant_fp8_scales(s_quant_1, reinterpret_cast(&frag_s[k2]) + 2); - } - -// We have the m dimension as the inner loop in order to encourage overlapping -// dequantization and matmul operations. -#pragma unroll - for (int j = 0; j < 4; j++) - { - FragB frag_b0; - FragB frag_b1; - int b_quant_0, b_quant_1; - - // NVFP4 (FE2M1f): shift to extract two halves - b_quant_1 = frag_b_quant[k2][0][j]; - b_quant_0 = b_quant_1 << 8; - - dequant_data(b_quant_0, reinterpret_cast(&frag_b0)); - dequant_data(b_quant_1, reinterpret_cast(&frag_b1)); - - if constexpr (dequant_skip_flop && has_zp && !is_a_8bit) - { - sub_zp(frag_b0, frag_zp[j], 0); - sub_zp(frag_b1, frag_zp[j], 1); - } - - // Apply scale to frag_b0 - if constexpr (has_act_order && !is_a_8bit) - { - static_assert(group_blocks != -1); - scale4( - frag_b0, act_frag_s[k2][0][j], act_frag_s[k2][1][j], act_frag_s[k2][2][j], act_frag_s[k2][3][j], 0); - scale4( - frag_b1, act_frag_s[k2][0][j], act_frag_s[k2][1][j], act_frag_s[k2][2][j], act_frag_s[k2][3][j], 1); - } - else if constexpr (!dequant_skip_flop && has_zp && group_blocks == -1 && !is_a_8bit) - { - int idx = (threadIdx.x / 4) % 2; - scalar_t2 s2 - = MarlinType::nums2num2(reinterpret_cast(&frag_s[j / 2][j % 2 * 2 + 0])[idx], - reinterpret_cast(&frag_s[j / 2][j % 2 * 2 + 1])[idx]); - if (is_new_zp) - frag_zp[j] = __hmul2(frag_zp[j], s2); - scale_and_sub(frag_b0, s2.x, frag_zp[j].x); - scale_and_sub(frag_b1, s2.y, frag_zp[j].y); - } - else if constexpr (!dequant_skip_flop && has_zp && group_blocks != -1 && !is_a_8bit) - { - if (is_new_zp) - frag_zp[j] = __hmul2(frag_zp[j], *reinterpret_cast(&frag_s[k2][j])); - scale_and_sub(frag_b0, frag_s[k2][j][0].x, frag_zp[j].x); - scale_and_sub(frag_b1, frag_s[k2][j][0].y, frag_zp[j].y); - } - else if constexpr (group_blocks != -1 && !is_a_8bit) - { - scale(frag_b0, frag_s[k2][j], 0); - scale(frag_b1, frag_s[k2][j], 1); - } - -#pragma unroll - for (int i = 0; i < thread_m_blocks; i++) - { - if constexpr (m_block_size_8) - { - mma_trans(frag_a[k2][i], frag_b0, frag_b1, frag_c[i][j][0]); - } - else - { - mma(frag_a[k2][i], frag_b0, frag_c[i][j][0]); - mma(frag_a[k2][i], frag_b1, frag_c[i][j][1]); - } - } - } - }; - - auto matmul_a8 = [&](int k) - { - int k2 = k % 2; -#pragma unroll - for (int j = 0; j < 2; j++) - { - FragB frag_b[2]; - - if (is_a_8bit && !has_zp) - { - dequant_data(frag_b_quant[k2][0][j * 2], reinterpret_cast(&frag_b)); - dequant_data(frag_b_quant[k2][0][j * 2 + 1], reinterpret_cast(&frag_b) + 2); - } - else if (is_a_8bit && has_zp) - { - int off = (threadIdx.x / 32) % 2 * 2 + j; - int zp = (frag_qzp[k2][0] >> (off * 8)) & 0xF; - dequant_data(frag_b_quant[k2][0][j * 2], reinterpret_cast(&frag_b), zp); - zp = (frag_qzp[k2][0] >> (off * 8 + 4)) & 0xF; - dequant_data(frag_b_quant[k2][0][j * 2 + 1], reinterpret_cast(&frag_b) + 2, zp); - } - else - { - reinterpret_cast(&frag_b)[0] = reinterpret_cast(&frag_b_quant[k2][j])[0]; - reinterpret_cast(&frag_b)[1] = reinterpret_cast(&frag_b_quant[k2][j])[1]; - } - -#pragma unroll - for (int i = 0; i < thread_m_blocks; i++) - { - mma(frag_a[k2][i], frag_b[0], (group_blocks == -1 ? frag_c : frag_c_tmp)[i][j][0]); - mma(frag_a[k2][i], frag_b[1], (group_blocks == -1 ? frag_c : frag_c_tmp)[i][j][1]); - } - - if constexpr (group_blocks != -1) - { - if (group_blocks == 2 || k == 1) - { - { - float2 s_vals[2]; - s_vals[0] = MarlinType::num22float2(frag_s[k2][j * 2][0]); - s_vals[1] = MarlinType::num22float2(frag_s[k2][j * 2 + 1][0]); - -#pragma unroll - for (int i = 0; i < thread_m_blocks; i++) - { -#pragma unroll - for (int g = 0; g < 4; g++) - { - float scale = reinterpret_cast(&s_vals[0])[g % 2]; - frag_c[i][j][0][g] += frag_c_tmp[i][j][0][g] * scale; - frag_c_tmp[i][j][0][g] = 0.0f; - } - -#pragma unroll - for (int g = 0; g < 4; g++) - { - float scale = reinterpret_cast(&s_vals[1])[g % 2]; - frag_c[i][j][1][g] += frag_c_tmp[i][j][1][g] * scale; - frag_c_tmp[i][j][1][g] = 0.0f; - } - } - } - } - } - } - }; - - // Since we slice across the k dimension of a tile in order to increase the - // number of warps while keeping the n dimension of a tile reasonable, we have - // multiple warps that accumulate their partial sums of the same output - // location; which we have to reduce over in the end. We do in shared memory. - auto thread_block_reduce = [&]() - { - constexpr int red_off = threads / b_sh_stride_threads / 2; - if (red_off >= 1) - { - auto red_idx = threadIdx.x / b_sh_stride_threads; - constexpr int red_sh_stride = b_sh_stride_threads * (is_a_8bit ? 2 : 4) * 2; - constexpr int red_sh_delta = b_sh_stride_threads; - int red_sh_rd = red_sh_stride * (threadIdx.x / b_sh_stride_threads) + (threadIdx.x % b_sh_stride_threads); - - // Parallel logarithmic shared memory reduction. We make sure to avoid any - // unnecessary read or write iterations, e.g., for two warps we write only - // once by warp 1 and read only once by warp 0. - -#pragma unroll - for (int m_block = 0; m_block < thread_m_blocks; m_block++) - { -#pragma unroll - for (int i = red_off; i > 0; i /= 2) - { - if (i <= red_idx && red_idx < 2 * i) - { -#pragma unroll - for (int j = 0; j < (is_a_8bit ? 2 : 4) * 2; j += (m_block_size_8 ? 2 : 1)) - { - int red_sh_wr = red_sh_delta * j + (red_sh_rd - red_sh_stride * i); - if (i < red_off) - { - float* c_rd = reinterpret_cast(&sh_red[red_sh_delta * j + red_sh_rd]); - float* c_wr = reinterpret_cast(&sh_red[red_sh_wr]); -#pragma unroll - for (int k = 0; k < 4; k++) - reinterpret_cast(frag_c)[(is_a_8bit ? 2 : 4) * 2 * m_block + j][k] - += c_rd[k] + c_wr[k]; - } - sh_red[red_sh_wr] = reinterpret_cast(&frag_c)[(is_a_8bit ? 2 : 4) * 2 * m_block + j]; - } - } - __syncthreads(); - } - if (red_idx == 0) - { -#pragma unroll - for (int i = 0; i < (is_a_8bit ? 2 : 4) * 2; i += (m_block_size_8 ? 2 : 1)) - { - float* c_rd = reinterpret_cast(&sh_red[red_sh_delta * i + red_sh_rd]); -#pragma unroll - for (int j = 0; j < 4; j++) - reinterpret_cast(frag_c)[(is_a_8bit ? 2 : 4) * 2 * m_block + i][j] += c_rd[j]; - } - } - __syncthreads(); - } - } - }; - - // Since multiple threadblocks may process parts of the same column slice, we - // finally have to globally reduce over the results. As the striped - // partitioning minimizes the number of such reductions and our outputs are - // usually rather small, we perform this reduction serially in L2 cache. - auto global_reduce_fp16 = [&](bool first = false, bool last = false) - { - // We are very careful here to reduce directly in the output buffer to - // maximize L2 cache utilization in this step. To do this, we write out - // results in FP16 (but still reduce with FP32 compute). - constexpr int active_threads = 32 * tb_n_warps; - if (threadIdx.x < active_threads) - { - int c_gl_stride = prob_n / 8; - int c_gl_wr_delta_o = 8 * c_gl_stride * (is_a_8bit ? 2 : 1); - int c_gl_wr_delta_i = 4 * (active_threads / 32); - int c_gl_wr; - if constexpr (m_block_size_8) - { - c_gl_wr = c_gl_stride * ((threadIdx.x % 4) * 2) + 4 * (threadIdx.x / 32) + (threadIdx.x % 32) / 8; - c_gl_wr += (2 * thread_n_blocks) * slice_col; - } - else - { - c_gl_wr = c_gl_stride * ((threadIdx.x % 32) / 4) * (is_a_8bit ? 2 : 1) + 4 * (threadIdx.x / 32) - + threadIdx.x % 4; - c_gl_wr += (2 * thread_n_blocks) * slice_col * (is_a_8bit ? 2 : 1); - } - constexpr int c_sh_wr_delta = active_threads; - auto c_sh_wr = threadIdx.x; - - int row = (threadIdx.x % 32) / 4; - - if (!first) - { -// Interestingly, doing direct global accesses here really seems to mess up -// the compiler and lead to slowdowns, hence we also use async-copies even -// though these fetches are not actually asynchronous. -#pragma unroll - for (int i = 0; i < (m_block_size_8 ? 2 : thread_m_blocks * 4); i++) - { - if constexpr (m_block_size_8) - { - cp_async4_pred(&sh_red[c_sh_wr + c_sh_wr_delta * i], - &C[c_gl_wr + i * c_gl_stride + (threadIdx.x % 8) / 4 * c_gl_wr_delta_i], - (threadIdx.x % 4) * 2 + i < prob_m); - } - else if constexpr (is_a_8bit) - { - int2* sh_red_int2 = reinterpret_cast(sh_red); - int2* c_int2 = reinterpret_cast(C); - cp_async2_ca_pred(&sh_red_int2[c_sh_wr + c_sh_wr_delta * i], - &c_int2[c_gl_wr + c_gl_wr_delta_o * (i / 2) + c_gl_wr_delta_i * (i % 2)], - i < (thread_m_blocks - 1) * 4 || 8 * (i / 2) + row < prob_m); - } - else - { - cp_async4_pred(&sh_red[c_sh_wr + c_sh_wr_delta * i], - &C[c_gl_wr + c_gl_wr_delta_o * (i / 2) + c_gl_wr_delta_i * (i % 2)], - i < (thread_m_blocks - 1) * 4 || 8 * (i / 2) + row < prob_m); - } - } - cp_async_fence(); - cp_async_wait<0>(); - } - -#pragma unroll - for (int i = 0; i < (m_block_size_8 ? 2 : thread_m_blocks * 4); i++) - { - bool mask = (!m_block_size_8) && (i < (thread_m_blocks - 1) * 4 || 8 * (i / 2) + row < prob_m) - || (m_block_size_8) && ((threadIdx.x % 4) * 2 + i < prob_m); - if (mask) - { - if (!first) - { - c_scalar_t* c_red_f16; - if constexpr (is_a_8bit) - { - int2 tmp = reinterpret_cast(sh_red)[c_sh_wr + i * c_sh_wr_delta]; - c_red_f16 = reinterpret_cast(&tmp); - } - else - { - int4 tmp = sh_red[c_sh_wr + i * c_sh_wr_delta]; - c_red_f16 = reinterpret_cast(&tmp); - } -#pragma unroll - for (int j = 0; j < 2 * (is_a_8bit ? 2 : 4); j++) - { - int delta = 0; - if constexpr (m_block_size_8) - { - delta = j % 2 == 1 ? -2 : 0; - } - reinterpret_cast( - &frag_c)[(is_a_8bit ? 2 : 4) * 2 * 4 * (i / 4) + 4 * j + (i % 4) + delta] - += MarlinType::num2float(c_red_f16[j]); - } - } - if (!last) - { - c_scalar_t c_f16[is_a_8bit ? 4 : 8]; -#pragma unroll - for (int j = 0; j < 2 * (is_a_8bit ? 2 : 4); j++) - { - int delta = 0; - if constexpr (m_block_size_8) - { - delta = j % 2 == 1 ? -2 : 0; - } - c_f16[j] = MarlinType::float2num(reinterpret_cast( - &frag_c)[(is_a_8bit ? 2 : 4) * 2 * 4 * (i / 4) + 4 * j + (i % 4) + delta]); - } - if constexpr (m_block_size_8) - { - C[c_gl_wr + i * c_gl_stride + (threadIdx.x % 8) / 4 * c_gl_wr_delta_i] - = *reinterpret_cast(c_f16); - } - else if constexpr (is_a_8bit) - { - int2* c_int2 = reinterpret_cast(C); - c_int2[c_gl_wr + c_gl_wr_delta_o * (i / 2) + c_gl_wr_delta_i * (i % 2)] - = *reinterpret_cast(c_f16); - } - else - { - C[c_gl_wr + c_gl_wr_delta_o * (i / 2) + c_gl_wr_delta_i * (i % 2)] - = *reinterpret_cast(c_f16); - } - } - } - } - } - }; - - // Globally reduce over threadblocks that compute the same column block. - // We use a tmp C buffer to reduce in full fp32 precision. - auto global_reduce_fp32 = [&](bool first = false, bool last = false) - { - constexpr int tb_m = thread_m_blocks * 16; - constexpr int tb_n = thread_n_blocks * 16; - - constexpr int c_size = tb_m * tb_n * sizeof(float) / 16; - - constexpr int active_threads = 32 * tb_n_warps; - bool is_th_active = threadIdx.x < active_threads; - - constexpr int num_floats = thread_m_blocks * (is_a_8bit ? 2 : 4) * 2 * 4; - constexpr int th_size = num_floats * sizeof(float) / 16; - - int c_cur_offset = locks_off * c_size; - - if (!is_th_active) - { - return; - } - - if (!first) - { - float* frag_c_ptr = reinterpret_cast(&frag_c); -#pragma unroll - for (int k = 0; k < th_size; k += (m_block_size_8 ? 2 : 1)) - { - sh_red[threadIdx.x] = C_tmp[c_cur_offset + active_threads * k + threadIdx.x]; - - float* sh_c_ptr = reinterpret_cast(&sh_red[threadIdx.x]); -#pragma unroll - for (int f = 0; f < 4; f++) - { - frag_c_ptr[k * 4 + f] += sh_c_ptr[f]; - } - } - } - - if (!last) - { - int4* frag_c_ptr = reinterpret_cast(&frag_c); -#pragma unroll - for (int k = 0; k < th_size; k += (m_block_size_8 ? 2 : 1)) - { - C_tmp[c_cur_offset + active_threads * k + threadIdx.x] = frag_c_ptr[k]; - } - } - }; - - // Write out the reduce final result in the correct layout. We only actually - // reshuffle matrix fragments in this step, the reduction above is performed - // in fragment layout. - auto write_result = [&](bool last) - { - int c_gl_stride = prob_n / 8; - constexpr int c_sh_stride = 2 * thread_n_blocks + 1; - int c_gl_wr_delta = c_gl_stride * (threads / (2 * thread_n_blocks)); - constexpr int c_sh_rd_delta = c_sh_stride * (threads / (2 * thread_n_blocks)); - - int c_gl_wr = c_gl_stride * (threadIdx.x / (2 * thread_n_blocks)) + (threadIdx.x % (2 * thread_n_blocks)); - c_gl_wr += (2 * thread_n_blocks) * slice_col; - int c_sh_wr; - if constexpr (m_block_size_8) - { - c_sh_wr = (8 * c_sh_stride) * ((threadIdx.x % 32) % 4 * 2) + (threadIdx.x % 32) / 4; - c_sh_wr += 64 * (threadIdx.x / 32); - } - else - { - c_sh_wr = (4 * c_sh_stride) * ((threadIdx.x % 32) / 4) + (threadIdx.x % 32) % 4; - c_sh_wr += (is_a_8bit ? 16 : 32) * (threadIdx.x / 32); - } - - int c_sh_rd = c_sh_stride * (threadIdx.x / (2 * thread_n_blocks)) + (threadIdx.x % (2 * thread_n_blocks)); - - int c_gl_wr_end = c_gl_stride * prob_m; - // We first reorder in shared memory to guarantee the most efficient final - // global write patterns - auto write = [&](int idx, float c0, float c1, FragS& s, FragS& b_bias) - { - c_scalar_t2 res = MarlinType::nums2num2( - MarlinType::float2num(c0), MarlinType::float2num(c1)); - - // For per-column quantization we finally apply the scale here (only for - // 4-bit) - if constexpr (!has_act_order && group_blocks == -1 && !is_a_8bit - && (has_zp && dequant_skip_flop || !has_zp)) - { - c_scalar_t2 tmp_scale = s[0]; - if constexpr (m_block_size_8) - { - tmp_scale - = MarlinType::num2num2(reinterpret_cast(&s[0])[(threadIdx.x % 8) / 4]); - } - res = __hmul2(res, tmp_scale); - } - - // NVFP4 with FP8 (E4M3) scales - res = __hmul2(res, global_scale); - if (has_bias && last) - { - c_scalar_t2 tmp_bias = b_bias[0]; - if constexpr (m_block_size_8) - { - tmp_bias = MarlinType::num2num2( - reinterpret_cast(&b_bias[0])[(threadIdx.x % 8) / 4]); - } - res = __hadd2(res, tmp_bias); - } - - if constexpr (m_block_size_8) - { - ((c_scalar_t*) sh_red)[idx] = res.x; - ((c_scalar_t*) sh_red)[idx + 8 * c_sh_stride] = res.y; - } - else - { - ((c_scalar_t2*) sh_red)[idx] = res; - } - }; - - if (threadIdx.x / 32 < tb_n_warps) - { -#pragma unroll - for (int i = 0; i < thread_m_blocks; i++) - { -#pragma unroll - for (int j = 0; j < (is_a_8bit ? 2 : 4); j++) - { - if constexpr (m_block_size_8) - { - int wr = c_sh_wr + 16 * j; - write(wr, frag_c[i][j][0][0], frag_c[i][j][0][1], frag_s[j / 2][2 * (j % 2) + 0], - frag_bias[j / 2][2 * (j % 2) + 0]); - write(wr + 8, frag_c[i][j][0][2], frag_c[i][j][0][3], frag_s[j / 2][2 * (j % 2) + 1], - frag_bias[j / 2][2 * (j % 2) + 1]); - } - else - { - int wr = c_sh_wr + 8 * j; - write(wr + (4 * c_sh_stride) * 0 + 0, frag_c[i][j][0][0], frag_c[i][j][0][1], - frag_s[j / 2][2 * (j % 2) + 0], frag_bias[j / 2][2 * (j % 2) + 0]); - write(wr + (4 * c_sh_stride) * 8 + 0, frag_c[i][j][0][2], frag_c[i][j][0][3], - frag_s[j / 2][2 * (j % 2) + 0], frag_bias[j / 2][2 * (j % 2) + 0]); - write(wr + (4 * c_sh_stride) * 0 + 4, frag_c[i][j][1][0], frag_c[i][j][1][1], - frag_s[j / 2][2 * (j % 2) + 1], frag_bias[j / 2][2 * (j % 2) + 1]); - write(wr + (4 * c_sh_stride) * 8 + 4, frag_c[i][j][1][2], frag_c[i][j][1][3], - frag_s[j / 2][2 * (j % 2) + 1], frag_bias[j / 2][2 * (j % 2) + 1]); - } - } - c_sh_wr += 16 * (4 * c_sh_stride); - } - } - __syncthreads(); - -#pragma unroll - for (int i = 0; i < div_ceil(16 * thread_m_blocks, threads / (2 * thread_n_blocks)); i++) - { - if (c_gl_wr < c_gl_wr_end) - { - if (use_atomic_add && slice_count > 1) - { - c_scalar_t2* C_half2 = reinterpret_cast(&C[c_gl_wr]); - c_scalar_t2* sh_red_half2 = reinterpret_cast(&sh_red[c_sh_rd]); -#pragma unroll - for (int a = 0; a < 4; a++) - { - atomicAdd(&C_half2[a], sh_red_half2[a]); - } - } - else - { - C[c_gl_wr] = sh_red[c_sh_rd]; - } - c_gl_wr += c_gl_wr_delta; - c_sh_rd += c_sh_rd_delta; - } - } - __syncthreads(); - }; - - // Start global fetch and register load pipelines. - auto start_pipes = [&]() - { - -#pragma unroll - for (int i = 0; i < stages - 1; i++) - { - if (has_act_order && i == 0) - { - int last_g_idx = slice_k_start + stages * tb_k * 2; - if (last_g_idx >= prob_k) - { - last_g_idx = prob_k - 1; - } - fetch_act_order_scales_to_shared(true, g_idx[slice_k_start], g_idx[last_g_idx]); - } - - if constexpr (has_zp && group_blocks == -1) - { - if (i == 0) - { - fetch_col_zp_to_shared(); - if constexpr (!dequant_skip_flop) - { - fetch_col_scale_to_shared(); - } - } - } - fetch_to_shared(i, i, i < slice_iters); - } - - zero_accums(); - wait_for_stage(); - init_same_group(0); - fetch_to_registers(0, 0); - fetch_scales_to_registers(0, 0); - fetch_zp_to_registers(0, 0); - a_gl_rd += a_gl_rd_delta_o * (stages - 1); - if constexpr (has_act_order) - { - slice_k_start_shared_fetch += tb_k * (stages - 1); - } - }; - if (slice_iters) - { - start_pipes(); - } - - // Main loop. - while (slice_iters) - { - // We unroll over both the global fetch and the register load pipeline to - // ensure all shared memory accesses are static. Note that both pipelines - // have even length meaning that the next iteration will always start at - // index 0. - -#pragma unroll - for (int pipe = 0; pipe < stages;) - { -#pragma unroll - for (int k = 0; k < b_sh_wr_iters; k++) - { - fetch_to_registers(k + 1, pipe % stages); - fetch_scales_to_registers(k + 1, pipe); - fetch_zp_to_registers(k + 1, pipe); - if (k == b_sh_wr_iters - 2) - { - fetch_to_shared((pipe + stages - 1) % stages, pipe, slice_iters >= stages); - pipe++; - wait_for_stage(); - init_same_group(pipe % stages); - } - - if constexpr (!is_a_8bit) - { - matmul(k, pipe - (k >= b_sh_wr_iters - 2 ? 1 : 0)); - } - else - { - static_assert(group_blocks != 0 && group_blocks != 1); - matmul_a8(k); - } - } - slice_iters--; - if (slice_iters == 0) - { - break; - } - } - - a_gl_rd += a_gl_rd_delta_o * stages; - - if constexpr (has_act_order) - { - slice_k_start += tb_k * stages; - - if (slice_k_start < prob_k) - { - slice_k_start_shared_fetch += tb_k * stages; - int first_group_id = g_idx[slice_k_start]; - int last_g_idx = slice_k_start + stages * tb_k * 2; - if (last_g_idx >= prob_k) - { - last_g_idx = prob_k - 1; - } - int last_group_id = g_idx[last_g_idx]; - if (last_group_id >= sh_first_group_id + sh_num_groups) - { - fetch_act_order_scales_to_shared(false, first_group_id, last_group_id); - __syncthreads(); - } - } - } - - // Process results and, if necessary, proceed to the next column slice. - // While this pattern may not be the most readable, other ways of writing - // the loop seemed to noticeably worse performance after compilation. - if (slice_iters == 0) - { - // convert fp16 accum to fp32 for reduction - if constexpr (use_fp16_accum) - { -#pragma unroll - for (int i = 0; i < (thread_m_blocks * (is_a_8bit ? 2 : 4) * 2); i++) - { - float* frag_c_part_float = reinterpret_cast(frag_c) + i * 4; - scalar_t* frag_c_part_half = reinterpret_cast(frag_c_part_float); - -#pragma unroll - for (int i = 3; i >= 0; i--) - { - frag_c_part_float[i] = MarlinType::num2float(frag_c_part_half[i]); - } - } - } - - if constexpr (is_a_8bit) - { - float frag_a_s[2 * thread_m_blocks]; - - for (int i = 0; i < 2 * thread_m_blocks; i++) - frag_a_s[i] = sh_a_s[i * 8 + (threadIdx.x % 32) / 4]; - -#pragma unroll - for (int j = 0; j < 2; j++) - { -#pragma unroll - for (int i = 0; i < thread_m_blocks; i++) - { -#pragma unroll - for (int g = 0; g < 4; g++) - { - float c_val = frag_c[i][j][0][g]; - float s_val = frag_a_s[i * 2 + g / 2]; - frag_c[i][j][0][g] = c_val * s_val; - } -#pragma unroll - for (int g = 0; g < 4; g++) - { - float c_val = frag_c[i][j][1][g]; - float s_val = frag_a_s[i * 2 + g / 2]; - frag_c[i][j][1][g] = c_val * s_val; - } - } - } - } - - cp_async_wait<0>(); - bool last = slice_idx == slice_count - 1; - // For per-column scales, we only fetch them here in the final step before - // write-out - if constexpr (!has_act_order && group_blocks == -1 && (has_zp && dequant_skip_flop || !has_zp)) - { - if ((last || use_atomic_add) || is_a_8bit) - { - if (s_sh_wr_pred) - { - cp_async4(&sh_s[s_sh_wr], &scales_ptr[s_gl_rd]); - } - cp_async_fence(); - } - } - - thread_block_reduce(); - - if (has_bias && last) - { - __syncthreads(); - cp_async4_pred(&sh_bias[bias_sh_wr], &b_bias_ptr[bias_gl_rd], threadIdx.x < 16 * thread_n_blocks / 8); - cp_async_fence(); - } - - if constexpr (!has_act_order && group_blocks == -1 && (has_zp && dequant_skip_flop || !has_zp || is_a_8bit)) - { - if constexpr (is_a_8bit) - { - cp_async_wait<0>(); - __syncthreads(); - if (threadIdx.x / 32 < tb_n_warps) - { - reinterpret_cast(&frag_s)[0] = sh_s[s_sh_rd + 0]; - } - } - else if (last || use_atomic_add) - { - cp_async_wait<0>(); - __syncthreads(); - if (threadIdx.x / 32 < tb_n_warps) - { - reinterpret_cast(&frag_s)[0] = sh_s[s_sh_rd + 0]; - reinterpret_cast(&frag_s)[1] = sh_s[s_sh_rd + 4]; - if constexpr (m_block_size_8) - { - int idx = (threadIdx.x / 4) % 2; - c_scalar_t2* frag_s_half2 = reinterpret_cast(frag_s); -#pragma unroll - for (int i = 0; i < 8; i++) - { - frag_s_half2[i] = MarlinType::num2num2( - reinterpret_cast(&frag_s_half2[i])[idx]); - } - } - } - } - } - - // For 8-bit channelwise, we apply the scale before the global reduction - // that converts the fp32 results to fp16 (so that we avoid possible - // overflow in fp16) - if constexpr (!has_act_order && group_blocks == -1 && is_a_8bit) - { -#pragma unroll - for (int j = 0; j < 2; j++) - { - float2 aa[2]; - aa[0] = MarlinType::num22float2(frag_s[0][j * 2][0]); - aa[1] = MarlinType::num22float2(frag_s[0][j * 2 + 1][0]); - -#pragma unroll - for (int i = 0; i < thread_m_blocks; i++) - { -#pragma unroll - for (int g = 0; g < 4; g++) - { - float scale = reinterpret_cast(&aa[0])[g % 2]; - frag_c[i][j][0][g] *= scale; - } - -#pragma unroll - for (int g = 0; g < 4; g++) - { - float scale = reinterpret_cast(&aa[1])[g % 2]; - frag_c[i][j][1][g] *= scale; - } - } - } - } - - if (slice_count > 1 && !use_atomic_add) - { - // only globally reduce if there is more than one block in a slice - barrier_acquire(&locks[locks_off], slice_idx); - if (use_fp32_reduce) - { - global_reduce_fp32(slice_idx == 0, last); - } - else - { - global_reduce_fp16(slice_idx == 0, last); - } - barrier_release(&locks[locks_off], last); - } - - if (has_bias && last) - { - cp_async_wait<0>(); - __syncthreads(); - reinterpret_cast(&frag_bias)[0] = sh_bias[bias_sh_rd]; - if constexpr (!is_a_8bit) - reinterpret_cast(&frag_bias)[1] = sh_bias[bias_sh_rd + 4]; - __syncthreads(); - } - - if (use_atomic_add && slice_count > 1 && slice_idx != 0) - wait_negative_and_add(&locks[locks_off]); - if (last || use_atomic_add) - // only the last block in a slice actually writes the result - write_result(last); - slice_row = 0; - if (!in_part2) - { - slice_col_par += gridDim.x; - } - else - { - slice_col_par++; - slice_col++; - } - is_first_matmul_in_slice = true; - init_slice(); - - if (slice_iters) - { - a_gl_rd = a_gl_stride * (threadIdx.x / a_gl_rd_delta_o) + (threadIdx.x % a_gl_rd_delta_o); - a_gl_rd += a_gl_rd_delta_o * slice_row; - b_gl_rd = b_gl_stride * (threadIdx.x / b_sh_stride) + (threadIdx.x % b_sh_stride); - b_gl_rd += b_sh_stride * slice_col + b_gl_rd_delta_o * slice_row; - - bias_gl_rd = (thread_n_blocks * 16 / 8) * slice_col + threadIdx.x; - // Update slice k/n for scales loading - if constexpr (has_act_order) - { - slice_k_start = tb_k * slice_row; - slice_k_finish = slice_k_start + tb_k * slice_iters; - slice_k_start_shared_fetch = slice_k_start; - slice_n_offset = act_s_col_tb_stride * slice_col; - } - else - { - if constexpr (group_blocks == -1) - { - s_gl_rd = s_sh_stride * slice_col + threadIdx.x; - zp_gl_rd = zp_sh_stride * slice_col + threadIdx.x; - } - else if constexpr (group_blocks >= thread_k_blocks) - { - s_gl_rd = s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) + s_sh_stride * slice_col - + threadIdx.x; - zp_gl_rd = zp_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) - + zp_sh_stride * slice_col + threadIdx.x; - } - else - { - s_gl_rd - = s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks + threadIdx.x / s_sh_stride) - + s_sh_stride * slice_col + threadIdx.x % s_sh_stride; - zp_gl_rd - = zp_gl_stride * ((thread_k_blocks * slice_row) / group_blocks + threadIdx.x / zp_sh_stride) - + zp_sh_stride * slice_col + threadIdx.x % zp_sh_stride; - } - } - start_pipes(); - } - } - } -} - -} // namespace MARLIN_NAMESPACE_NAME - -#endif diff --git a/cpp/tensorrt_llm/kernels/marlin/marlin_repack.cu b/cpp/tensorrt_llm/kernels/marlin/marlin_repack.cu deleted file mode 100644 index 03996eebbf68..000000000000 --- a/cpp/tensorrt_llm/kernels/marlin/marlin_repack.cu +++ /dev/null @@ -1,350 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2022-2025 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. - */ - -#include "marlin.cuh" -#include "marlin_nvfp4.h" -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/cudaUtils.h" - -namespace marlin -{ - -template -__global__ void gptq_marlin_repack_kernel(uint32_t const* __restrict__ b_q_weight_ptr, - uint32_t const* __restrict__ perm_ptr, uint32_t* __restrict__ out_ptr, int size_k, int size_n) -{ - constexpr int pack_factor = 32 / num_bits; - - constexpr int target_tile_n_size = tile_n_size / (is_a_8bit ? 2 : 1); - constexpr int target_tile_k_size = tile_k_size * (is_a_8bit ? 2 : 1); - int k_tiles = size_k / target_tile_k_size; - int n_tiles = size_n / target_tile_n_size; - int block_k_tiles = div_ceil(k_tiles, gridDim.x); - - auto start_k_tile = blockIdx.x * block_k_tiles; - if (start_k_tile >= k_tiles) - { - return; - } - - int finish_k_tile = min(start_k_tile + block_k_tiles, k_tiles); - - // Wait until the next thread tile has been loaded to shared memory. - auto wait_for_stage = [&]() - { - // We only have `stages - 2` active fetches since we are double buffering - // and can only issue the next fetch when it is guaranteed that the previous - // shared memory load is fully complete (as it may otherwise be - // overwritten). - cp_async_wait(); - __syncthreads(); - }; - - extern __shared__ int4 sh[]; - - constexpr int perm_size = target_tile_k_size / 4; - - int4* sh_perm_ptr = sh; - int4* sh_pipe_ptr = sh_perm_ptr; - if constexpr (has_perm) - { - sh_pipe_ptr += perm_size; - } - - constexpr int tile_ints = target_tile_k_size / pack_factor; - - constexpr int stage_n_threads = target_tile_n_size / 4; - constexpr int stage_k_threads = has_perm ? target_tile_k_size : tile_ints; - constexpr int stage_size = stage_k_threads * stage_n_threads; - - auto load_perm_to_shared = [&](int k_tile_id) - { - int first_k_int4 = (k_tile_id * target_tile_k_size) / 4; - - int4 const* perm_int4_ptr = reinterpret_cast(perm_ptr); - - if (threadIdx.x < perm_size) - { - sh_perm_ptr[threadIdx.x] = perm_int4_ptr[first_k_int4 + threadIdx.x]; - } - __syncthreads(); - }; - - auto fetch_to_shared = [&](int pipe, int k_tile_id, int n_tile_id) - { - if (n_tile_id >= n_tiles) - { - cp_async_fence(); - return; - } - - int first_n = n_tile_id * target_tile_n_size; - - int4* sh_ptr = sh_pipe_ptr + stage_size * pipe; - - if constexpr (has_perm) - { - if (threadIdx.x < stage_size) - { - auto k_id = threadIdx.x / stage_n_threads; - auto n_id = threadIdx.x % stage_n_threads; - - uint32_t const* sh_perm_int_ptr = reinterpret_cast(sh_perm_ptr); - - int src_k = sh_perm_int_ptr[k_id]; - int src_k_packed = src_k / pack_factor; - - cp_async4(&sh_ptr[k_id * stage_n_threads + n_id], - reinterpret_cast(&(b_q_weight_ptr[src_k_packed * size_n + first_n + (n_id * 4)]))); - } - } - else - { - if (threadIdx.x < stage_size) - { - auto k_id = threadIdx.x / stage_n_threads; - auto n_id = threadIdx.x % stage_n_threads; - - int first_k = k_tile_id * target_tile_k_size; - int first_k_packed = first_k / pack_factor; - - cp_async4(&sh_ptr[k_id * stage_n_threads + n_id], - reinterpret_cast( - &(b_q_weight_ptr[(first_k_packed + k_id) * size_n + first_n + (n_id * 4)]))); - } - } - - cp_async_fence(); - }; - - auto repack_tile = [&](int pipe, int k_tile_id, int n_tile_id) - { - if (n_tile_id >= n_tiles) - { - return; - } - - auto warp_id = threadIdx.x / 32; - auto th_id = threadIdx.x % 32; - - if (warp_id >= 4) - { - return; - } - - int tc_col = th_id / 4; - int tc_row = (th_id % 4) * (is_a_8bit ? 4 : 2); - - constexpr int tc_offsets[4] = {0, 1, 8, 9}; - - int cur_n = (warp_id / (is_a_8bit ? 2 : 1)) * 16 + tc_col; - - constexpr int sh_stride = target_tile_n_size; - constexpr uint32_t mask = (1 << num_bits) - 1; - - int4* sh_stage_ptr = sh_pipe_ptr + stage_size * pipe; - uint32_t* sh_stage_int_ptr = reinterpret_cast(sh_stage_ptr); - - uint32_t* sh_perm_int_ptr = reinterpret_cast(sh_perm_ptr); - - uint32_t vals[8]; - - if constexpr (has_perm) - { - static_assert(!is_a_8bit); - for (int i = 0; i < 4; i++) - { - int k_idx = tc_row + tc_offsets[i]; - - uint32_t src_k = sh_perm_int_ptr[k_idx]; - uint32_t src_k_pos = src_k % pack_factor; - - uint32_t b1_val = sh_stage_int_ptr[k_idx * sh_stride + cur_n]; - uint32_t b1_cur_val = (b1_val >> (src_k_pos * num_bits)) & mask; - - uint32_t b2_val = sh_stage_int_ptr[k_idx * sh_stride + cur_n + 8]; - uint32_t b2_cur_val = (b2_val >> (src_k_pos * num_bits)) & mask; - - vals[i] = b1_cur_val; - vals[4 + i] = b2_cur_val; - } - } - else - { - uint32_t b1_vals[tile_ints]; - uint32_t b2_vals[tile_ints]; - -#pragma unroll - for (int i = 0; i < tile_ints; i++) - { - if constexpr (is_a_8bit) - { - b1_vals[i] = sh_stage_int_ptr[cur_n + sh_stride * i + (warp_id % 2) * 8]; - } - else - { - b1_vals[i] = sh_stage_int_ptr[cur_n + sh_stride * i]; - b2_vals[i] = sh_stage_int_ptr[cur_n + 8 + sh_stride * i]; - } - } - -#pragma unroll - for (int i = 0; i < 4; i++) - { - int cur_elem = tc_row + (is_a_8bit ? i : tc_offsets[i]); - int cur_int = cur_elem / pack_factor; - int cur_pos = cur_elem % pack_factor; - - vals[i] = (b1_vals[cur_int] >> (cur_pos * num_bits)) & mask; - if constexpr (is_a_8bit) - vals[4 + i] = (b1_vals[cur_int + tile_ints / 2] >> (cur_pos * num_bits)) & mask; - else - vals[4 + i] = (b2_vals[cur_int] >> (cur_pos * num_bits)) & mask; - } - } - - constexpr int tile_size = target_tile_k_size * target_tile_n_size / pack_factor; - int out_offset = (k_tile_id * n_tiles + n_tile_id) * tile_size; - - // Result of: - // https://github.com/NVIDIA/FasterTransformer/blob/main/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h - if constexpr (!is_a_8bit && num_bits == 4) - { - int pack_idx[8] = {0, 2, 4, 6, 1, 3, 5, 7}; - - uint32_t res = 0; -#pragma unroll - for (int i = 0; i < 8; i++) - { - res |= vals[pack_idx[i]] << (i * 4); - } - - out_ptr[out_offset + th_id * 4 + warp_id] = res; - } - else if constexpr (is_a_8bit && num_bits == 4) - { - int pack_idx[8] = {0, 4, 1, 5, 2, 6, 3, 7}; - - uint32_t res = 0; -#pragma unroll - for (int i = 0; i < 8; i++) - { - res |= vals[pack_idx[i]] << (i * 4); - } - - out_ptr[out_offset + th_id * 4 + warp_id] = res; - } - else - { - constexpr int pack_idx[4] = {0, 2, 1, 3}; - - uint32_t res1 = 0; - uint32_t res2 = 0; -#pragma unroll - for (int i = 0; i < 4; i++) - { - const int ii = is_a_8bit ? i : pack_idx[i]; - res1 |= vals[ii] << (i * 8); - res2 |= vals[4 + ii] << (i * 8); - } - - out_ptr[out_offset + th_id * 8 + (warp_id * 2) + 0] = res1; - out_ptr[out_offset + th_id * 8 + (warp_id * 2) + 1] = res2; - } - }; - - auto start_pipes = [&](int k_tile_id, int n_tile_id) - { -#pragma unroll - for (int pipe = 0; pipe < repack_stages - 1; pipe++) - { - fetch_to_shared(pipe, k_tile_id, n_tile_id + pipe); - } - - wait_for_stage(); - }; -#pragma unroll - for (int k_tile_id = start_k_tile; k_tile_id < finish_k_tile; k_tile_id++) - { - int n_tile_id = 0; - - if constexpr (has_perm) - { - load_perm_to_shared(k_tile_id); - } - - start_pipes(k_tile_id, n_tile_id); - - while (n_tile_id < n_tiles) - { -#pragma unroll - for (int pipe = 0; pipe < repack_stages; pipe++) - { - fetch_to_shared( - (pipe + repack_stages - 1) % repack_stages, k_tile_id, n_tile_id + pipe + repack_stages - 1); - repack_tile(pipe, k_tile_id, n_tile_id + pipe); - wait_for_stage(); - } - n_tile_id += repack_stages; - } - } -} - -} // namespace marlin - -#define CALL_IF(NUM_BITS, HAS_PERM, IS_A_8BIT) \ - else if (num_bits == NUM_BITS && has_perm == HAS_PERM && is_a_8bit == IS_A_8BIT) \ - { \ - cudaFuncSetAttribute(marlin::gptq_marlin_repack_kernel, \ - cudaFuncAttributeMaxDynamicSharedMemorySize, max_shared_mem); \ - marlin::gptq_marlin_repack_kernel \ - <<>>( \ - b_q_weight_ptr, perm_ptr, out_ptr, size_k, size_n); \ - } - -namespace marlin_nvfp4 -{ - -void gptq_marlin_repack_dispatch(uint32_t const* b_q_weight_ptr, uint32_t const* perm_ptr, uint32_t* out_ptr, - int size_k, int size_n, int num_bits, bool has_perm, bool is_a_8bit, cudaStream_t stream) -{ - int const sm = tensorrt_llm::common::getSMVersion(); - TLLM_CHECK_WITH_INFO( - sm >= 90 && sm < 100, "Marlin NVFP4 repack is only supported on Hopper (SM 9.x); current SM = %d", sm); - - int blocks; - int dev; - cudaGetDevice(&dev); - cudaDeviceGetAttribute(&blocks, cudaDevAttrMultiProcessorCount, dev); - - int max_shared_mem = 0; - cudaDeviceGetAttribute(&max_shared_mem, cudaDevAttrMaxSharedMemoryPerBlockOptin, dev); - - if (false) - { - } - CALL_IF(4, false, false) - CALL_IF(4, true, false) - CALL_IF(8, false, false) - CALL_IF(8, true, false) - CALL_IF(4, false, true) - CALL_IF(8, false, true) -} - -} // namespace marlin_nvfp4 - -#undef CALL_IF diff --git a/cpp/tensorrt_llm/kernels/megaMoePrepareKernel.cu b/cpp/tensorrt_llm/kernels/megaMoePrepareKernel.cu deleted file mode 100644 index 60a07269f8f1..000000000000 --- a/cpp/tensorrt_llm/kernels/megaMoePrepareKernel.cu +++ /dev/null @@ -1,158 +0,0 @@ -/* - * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. - * - * 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. - */ - -#include "tensorrt_llm/kernels/megaMoePrepareKernel.h" -#include "tensorrt_llm/kernels/quantization.cuh" - -#include -#include -#include -#include -#include - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels -{ -namespace -{ - -constexpr int kSfVecSize = 32; - -template -__device__ __forceinline__ float toFloat(T value) -{ - return static_cast(value); -} - -template <> -__device__ __forceinline__ float toFloat(half value) -{ - return __half2float(value); -} - -template <> -__device__ __forceinline__ float toFloat<__nv_bfloat16>(__nv_bfloat16 value) -{ - return __bfloat162float(value); -} - -template -__global__ void megaMoePrepareKernel(__nv_bfloat16 const* __restrict__ input, - ExpertT const* __restrict__ tokenSelectedExperts, ScaleT const* __restrict__ tokenFinalScales, - uint32_t* __restrict__ xOut, uint32_t* __restrict__ xSfOut, int64_t* __restrict__ topkIdxOut, - float* __restrict__ topkWeightsOut, int numTokens, int hiddenSize, int topK) -{ -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) - static constexpr int kEltsPerThread = CVT_ELTS_PER_THREAD; - static constexpr int kThreadsPerSf = kSfVecSize / kEltsPerThread; - using InputVec = PackedVec<__nv_bfloat16>; - - int const numColThreads = hiddenSize / kEltsPerThread; - int const numSfCols = hiddenSize / kSfVecSize; - - for (int rowIdx = blockIdx.x; rowIdx < numTokens; rowIdx += gridDim.x) - { - for (int topk = threadIdx.x; topk < topK; topk += blockDim.x) - { - int64_t const offset = static_cast(rowIdx) * topK + topk; - topkIdxOut[offset] = static_cast(tokenSelectedExperts[offset]); - topkWeightsOut[offset] = toFloat(tokenFinalScales[offset]); - } - - for (int colIdx = threadIdx.x; colIdx < numColThreads; colIdx += blockDim.x) - { - uint8_t* sfOut = nullptr; - if (colIdx % kThreadsPerSf == 0) - { - int64_t const sfOffset = static_cast(rowIdx) * numSfCols + colIdx / kThreadsPerSf; - sfOut = reinterpret_cast(xSfOut) + sfOffset; - } - - int64_t const inOffset = static_cast(rowIdx) * numColThreads + colIdx; - InputVec inVec = reinterpret_cast(input)[inOffset]; - reinterpret_cast(xOut)[inOffset] - = cvt_warp_fp16_to_mxfp8<__nv_bfloat16, kSfVecSize>(inVec, sfOut); - } - } -#endif -} - -template -void launchMegaMoePrepare(void const* input, void const* tokenSelectedExperts, void const* tokenFinalScales, void* xOut, - void* xSfOut, int64_t* topkIdxOut, float* topkWeightsOut, int numTokens, int hiddenSize, int topK, - int multiProcessorCount, cudaStream_t stream) -{ - int const blockX = std::min(hiddenSize / CVT_ELTS_PER_THREAD, 512); - int const numBlocksPerSm = std::max(1, 2048 / blockX); - int const gridX = std::min(numTokens, multiProcessorCount * numBlocksPerSm); - - megaMoePrepareKernel<<>>(static_cast<__nv_bfloat16 const*>(input), - static_cast(tokenSelectedExperts), static_cast(tokenFinalScales), - static_cast(xOut), static_cast(xSfOut), topkIdxOut, topkWeightsOut, numTokens, hiddenSize, - topK); -} - -template -void dispatchMegaMoePrepareScale(void const* input, void const* tokenSelectedExperts, void const* tokenFinalScales, - void* xOut, void* xSfOut, int64_t* topkIdxOut, float* topkWeightsOut, int numTokens, int hiddenSize, int topK, - MegaMoePrepareScaleType scaleType, int multiProcessorCount, cudaStream_t stream) -{ - switch (scaleType) - { - case MegaMoePrepareScaleType::FP32: - launchMegaMoePrepare(input, tokenSelectedExperts, tokenFinalScales, xOut, xSfOut, topkIdxOut, - topkWeightsOut, numTokens, hiddenSize, topK, multiProcessorCount, stream); - break; - case MegaMoePrepareScaleType::FP16: - launchMegaMoePrepare(input, tokenSelectedExperts, tokenFinalScales, xOut, xSfOut, topkIdxOut, - topkWeightsOut, numTokens, hiddenSize, topK, multiProcessorCount, stream); - break; - case MegaMoePrepareScaleType::BF16: - launchMegaMoePrepare(input, tokenSelectedExperts, tokenFinalScales, xOut, xSfOut, - topkIdxOut, topkWeightsOut, numTokens, hiddenSize, topK, multiProcessorCount, stream); - break; - } -} - -} // namespace - -void invokeMegaMoePrepare(void const* input, void const* tokenSelectedExperts, void const* tokenFinalScales, void* xOut, - void* xSfOut, int64_t* topkIdxOut, float* topkWeightsOut, int numTokens, int hiddenSize, int topK, - MegaMoePrepareExpertType expertType, MegaMoePrepareScaleType scaleType, int multiProcessorCount, - cudaStream_t stream) -{ - if (numTokens == 0) - { - return; - } - - switch (expertType) - { - case MegaMoePrepareExpertType::INT32: - dispatchMegaMoePrepareScale(input, tokenSelectedExperts, tokenFinalScales, xOut, xSfOut, topkIdxOut, - topkWeightsOut, numTokens, hiddenSize, topK, scaleType, multiProcessorCount, stream); - break; - case MegaMoePrepareExpertType::INT64: - dispatchMegaMoePrepareScale(input, tokenSelectedExperts, tokenFinalScales, xOut, xSfOut, topkIdxOut, - topkWeightsOut, numTokens, hiddenSize, topK, scaleType, multiProcessorCount, stream); - break; - } -} - -} // namespace kernels - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/megaMoePrepareKernel.h b/cpp/tensorrt_llm/kernels/megaMoePrepareKernel.h deleted file mode 100644 index bca7bd723ca2..000000000000 --- a/cpp/tensorrt_llm/kernels/megaMoePrepareKernel.h +++ /dev/null @@ -1,63 +0,0 @@ -/* - * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. - * - * 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. - */ -#pragma once - -#include "tensorrt_llm/common/config.h" - -#include -#include - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels -{ - -//! Dtype of the token-selected expert index tensor consumed by invokeMegaMoePrepare. -enum class MegaMoePrepareExpertType -{ - INT32, - INT64, -}; - -//! Dtype of the token-final scale tensor consumed by invokeMegaMoePrepare. -enum class MegaMoePrepareScaleType -{ - FP32, - FP16, - BF16, -}; - -//! Prepare DeepGEMM MegaMoE inputs by quantizing activations and copying routing metadata. -//! -//! Expected tensor-like arguments: -//! - input: [numTokens, hiddenSize] BF16 activations. -//! - tokenSelectedExperts: [numTokens, topK] INT32 or INT64 expert/slot ids. -//! - tokenFinalScales: [numTokens, topK] FP32, FP16, or BF16 routing scales. -//! - xOut: [>=numTokens, hiddenSize] FP8 E4M3 output activations. -//! - xSfOut: [>=numTokens, hiddenSize / 128] INT32 packed UE8M0 scales. -//! - topkIdxOut: [>=numTokens, topK] INT64 expert/slot ids. -//! - topkWeightsOut: [>=numTokens, topK] FP32 routing scales. -//! -//! hiddenSize must be divisible by 128. All pointers must refer to contiguous -//! CUDA buffers on the same device, and the kernel requires SM100 or newer. -void invokeMegaMoePrepare(void const* input, void const* tokenSelectedExperts, void const* tokenFinalScales, void* xOut, - void* xSfOut, int64_t* topkIdxOut, float* topkWeightsOut, int numTokens, int hiddenSize, int topK, - MegaMoePrepareExpertType expertType, MegaMoePrepareScaleType scaleType, int multiProcessorCount, - cudaStream_t stream); - -} // namespace kernels - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/mhcKernels/CMakeLists.txt b/cpp/tensorrt_llm/kernels/mhcKernels/CMakeLists.txt deleted file mode 100644 index eb07c6f165d3..000000000000 --- a/cpp/tensorrt_llm/kernels/mhcKernels/CMakeLists.txt +++ /dev/null @@ -1,38 +0,0 @@ -# -# 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. -# - -# Library sources are listed explicitly. Profiling/benchmark/test harnesses in -# this directory (bench_*.cu, probe_*.cu, profile_*.cu, test_*.cu) are -# standalone binaries with their own main() and must NOT be compiled into the -# trtllm shared library. -set(SRC_CU mhcKernels.cu mhcFusedHcKernel.cu) -add_library(mhcKernels_src OBJECT ${SRC_CU}) - -set_property(TARGET mhcKernels_src PROPERTY POSITION_INDEPENDENT_CODE ON) -set_property(TARGET mhcKernels_src PROPERTY CUDA_RESOLVE_DEVICE_SYMBOLS ON) -target_compile_options(mhcKernels_src - PRIVATE $<$:--use_fast_math>) - -if(BUILD_DEEP_GEMM) - # Expose DeepGEMM helper headers (sm100_utils.cuh, utils.cuh, tma_utils.cuh, - # reduction.cuh) used by the tcgen05.mma-based fused post-mapping + GEMM - # kernel on SM100. - target_include_directories( - mhcKernels_src - PRIVATE ${CMAKE_BINARY_DIR}/_deps/deepgemm-src/deep_gemm/include) - target_compile_definitions(mhcKernels_src PRIVATE TRTLLM_MHC_ENABLE_FUSED_HC) -endif() diff --git a/cpp/tensorrt_llm/kernels/mhcKernels/fused_tf32_pmap_gemm.cuh b/cpp/tensorrt_llm/kernels/mhcKernels/fused_tf32_pmap_gemm.cuh deleted file mode 100644 index ed2841f63ce2..000000000000 --- a/cpp/tensorrt_llm/kernels/mhcKernels/fused_tf32_pmap_gemm.cuh +++ /dev/null @@ -1,1602 +0,0 @@ -/* - * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. - * - * 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. - */ - -// Fused post-mapping + TF32 HC prenorm GEMM on B200 (SM100). -// -// Mathematical formula: -// new_r[i, hc, h] = post_mix[i, hc] * x[i, h] -// + sum_j comb_mix[i, j, hc] * residual[i, j, h] -// D[i, n] = sum_{hc, h} new_r[i, hc, h] * W[n, hc, h] -// sqr[i] = sum_{hc, h} new_r[i, hc, h]^2 (bf16-rounded) -// -// Shape: -// residual [M, HC_MULT, hidden] bf16 -// x [M, hidden] bf16 -// post_mix [M, HC_MULT] fp32 -// comb_mix [M, HC_MULT, HC_MULT] fp32 -// W [N, HC_MULT*hidden] fp32 (TF32) -// D [M, N] fp32 -// sqr [M] fp32 -// -// Kernel architecture (derived from DeepGEMM sm100_tf32_hc_prenorm_gemm): -// - BLOCK_M x BLOCK_N x BLOCK_K = 64 x 32 x 64, TF32 MMA on tcgen05 -// - 256 threads per CTA (warps 0..3 = MMA group, warps 4..7 = pmap group) -// - new_r is NEVER materialized in GMEM - it is computed on-chip directly -// into TMEM (fp32) where UMMA consumes it. -// - Iteration order: outer h_tile (slow), inner hc_idx (fast). residual+x -// SMEM is reused for HC_MULT=4 consecutive MMA stages. -// -// Barriers: -// full_B[N_B_STAGES] TMA -> MMA (B arrived in SMEM) -// empty_B[N_B_STAGES] MMA -> TMA (B slot empty) -// full_input[N_INPUT] TMA -> pmap (residual+x arrived) -// empty_input[N_INPUT] pmap -> TMA (input slot empty) -// full_cast[2] pmap -> MMA (A ready in TMEM) -// empty_cast[2] MMA -> pmap (TMEM slot empty) -// tmem_full[1] MMA -> epilogue - -#pragma once -#pragma clang diagnostic push -#pragma clang diagnostic ignored "-Wunknown-attributes" - -#include -#include -#include - -namespace deep_gemm::sm90 -{ -using cuda::std::swap; -} - -namespace deep_gemm::sm100 -{ -using cuda::std::swap; -} - -#include -#include -#include -#include -#include -#include -#include - -namespace fused_mhc -{ - -// Reuse DeepGEMM's swizzle helper -using deep_gemm::sm100::make_umma_desc; -using deep_gemm::sm100::get_num_aligned_tmem_cols; -using deep_gemm::sm100::tcgen05_before_thread_sync; -using deep_gemm::sm100::tcgen05_after_thread_sync; -using deep_gemm::sm100::advance_umma_desc_lo; -using deep_gemm::tma_copy; -using deep_gemm::utils::PatternVisitor; -using deep_gemm::ptx::get_lane_idx; - -template -__device__ __forceinline__ uint32_t get_swizzled_smem_offset(uint32_t const& offset, uint32_t const& lane_idx) -{ - auto const& bank_group_idx = offset + lane_idx * (kSwizzleMode / kSwizzleBase); - constexpr uint32_t kNumBankGroups = 128 / kSwizzleBase; - constexpr bool kHasShortcut = (kSwizzleMode / kSwizzleBase) == kNumBankGroups; - auto row = kHasShortcut ? (offset / kNumBankGroups + lane_idx) : (bank_group_idx / kNumBankGroups); - auto col = kHasShortcut ? (offset) : (bank_group_idx % kNumBankGroups); - col ^= row % (kSwizzleMode / kSwizzleBase); - return row * 128 + col * kSwizzleBase; -} - -__device__ __forceinline__ void stsm_x4_b16_rout(void* smem_dst, uint32_t a, uint32_t b, uint32_t c, uint32_t d) -{ - asm volatile( - "stmatrix.sync.aligned.x4.m8n8.shared.b16 [%0], {%1, %2, %3, %4};\n" ::"l"(__cvta_generic_to_shared(smem_dst)), - "r"(a), "r"(b), "r"(c), "r"(d)); -} - -template -__global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) fused_tf32_pmap_gemm_rout_atomic_impl( - const uint32_t shape_m, const __grid_constant__ cute::TmaDescriptor tensor_map_residual, - const __grid_constant__ cute::TmaDescriptor tensor_map_x, const __grid_constant__ cute::TmaDescriptor tensor_map_b, - const __grid_constant__ cute::TmaDescriptor tensor_map_residual_out, - float* __restrict__ D, // [M, SHAPE_N] (caller memsets to 0) - float const* __restrict__ post_mix, float const* __restrict__ comb_mix, float* __restrict__ sqr_sum) -{ // [M] (caller memsets to 0) -#if (defined(__CUDA_ARCH__) and (__CUDA_ARCH__ >= 1000) and (__CUDA_ARCH__ < 1100)) or defined(__CLION_IDE__) - using Barrier = cutlass::arch::ClusterTransactionBarrier; - - constexpr uint32_t SHAPE_K = HC_MULT * HIDDEN; - constexpr uint32_t H_TILES_PER_HC = HIDDEN / BLOCK_K; - static_assert(H_TILES_PER_HC % kNumSplits == 0, "H_TILES_PER_HC must be divisible by kNumSplits"); - constexpr uint32_t H_TILES_PER_SPLIT = H_TILES_PER_HC / kNumSplits; - constexpr uint32_t kNumCastStages = 4; - constexpr uint32_t kSwizzleAMode = cute::min(BLOCK_K * sizeof(nv_bfloat16), 128); - constexpr uint32_t kSwizzleBMode = cute::min(BLOCK_K * sizeof(float), 128); - constexpr uint32_t kSwizzleXMode = kSwizzleAMode; - constexpr uint32_t kSwizzleResMode = kSwizzleAMode; - constexpr uint32_t kSwizzleRoutMode = kSwizzleAMode; - constexpr auto kMajorA = cute::UMMA::Major::K; - constexpr auto kMajorB = cute::UMMA::Major::K; - static_assert(HIDDEN % BLOCK_K == 0, "HIDDEN must be multiple of BLOCK_K"); - static_assert(N_B_STAGES >= HC_MULT, "N_B_STAGES must be >= HC_MULT"); - static_assert(kSwizzleCDMode / sizeof(float) == BLOCK_N, "Invalid block N"); - static_assert(kNumMMAThreads == 128, "Invalid MMA threads"); - static_assert(kNumPmapThreads == 128, "Invalid pmap threads"); - static_assert(BLOCK_M == 64, "Invalid block M"); - static_assert(HC_MULT == 4, "Only HC_MULT=4 supported"); - static_assert(kSwizzleCDMode == 128, "Atomic variant expects kSwizzleCDMode=128"); - static_assert(SHAPE_N <= BLOCK_N, "SHAPE_N must fit within BLOCK_N"); - - auto const warp_idx = cutlass::canonical_warp_idx_sync(); - auto const lane_idx = get_lane_idx(); - - extern __shared__ __align__(1024) uint8_t smem_buffer[]; - - constexpr uint32_t SMEM_CD_SIZE = BLOCK_M * kSwizzleCDMode; - constexpr uint32_t SMEM_B_PER_STAGE = BLOCK_N * BLOCK_K * sizeof(float); - constexpr uint32_t SMEM_RES_PER_ISTG = BLOCK_M * HC_MULT * BLOCK_K * sizeof(nv_bfloat16); - constexpr uint32_t SMEM_X_PER_ISTG = BLOCK_M * BLOCK_K * sizeof(nv_bfloat16); - constexpr uint32_t SMEM_POST_SIZE = BLOCK_M * HC_MULT * sizeof(float); - constexpr uint32_t SMEM_COMB_SIZE = BLOCK_M * HC_MULT * HC_MULT * sizeof(float); - constexpr uint32_t SMEM_RC_PER_HC = BLOCK_M * BLOCK_K * sizeof(nv_bfloat16); // 8 KB - constexpr uint32_t SMEM_RC_SIZE = HC_MULT * SMEM_RC_PER_HC; // 32 KB - - constexpr uint32_t kNumTmemCols = get_num_aligned_tmem_cols(); - - // Prefetch TMA descriptors - if (warp_idx == 0 and cute::elect_one_sync()) - { - cute::prefetch_tma_descriptor(&tensor_map_residual); - cute::prefetch_tma_descriptor(&tensor_map_x); - cute::prefetch_tma_descriptor(&tensor_map_b); - cute::prefetch_tma_descriptor(&tensor_map_residual_out); - } - - // SMEM layout: [cd, B stages, res stages, x stages, post, comb, rc (HC_MULT slices)] - auto smem_cd = reinterpret_cast(smem_buffer); - uint8_t* cursor = smem_buffer + SMEM_CD_SIZE; - auto smem_b = PatternVisitor( - [&, base = cursor](uint32_t const& i) { return reinterpret_cast(base + i * SMEM_B_PER_STAGE); }); - cursor += N_B_STAGES * SMEM_B_PER_STAGE; - auto smem_res = PatternVisitor( - [&, base = cursor](uint32_t const& i) { return reinterpret_cast(base + i * SMEM_RES_PER_ISTG); }); - cursor += N_INPUT_STAGES * SMEM_RES_PER_ISTG; - auto smem_x_stg = PatternVisitor( - [&, base = cursor](uint32_t const& i) { return reinterpret_cast(base + i * SMEM_X_PER_ISTG); }); - cursor += N_INPUT_STAGES * SMEM_X_PER_ISTG; - auto smem_post = reinterpret_cast(cursor); - cursor += SMEM_POST_SIZE; - auto smem_comb = reinterpret_cast(cursor); - cursor += SMEM_COMB_SIZE; - auto smem_rc = reinterpret_cast(cursor); // [HC_MULT][BLOCK_M][BLOCK_K] bf16, single-buffered - cursor += SMEM_RC_SIZE; - - cursor = reinterpret_cast((reinterpret_cast(cursor) + 7) & ~uintptr_t(7)); - auto barrier_start_ptr = reinterpret_cast(cursor); - auto full_B = PatternVisitor([=](uint32_t const& i) { return barrier_start_ptr + i; }); - auto empty_B = PatternVisitor([=](uint32_t const& i) { return barrier_start_ptr + N_B_STAGES + i; }); - auto full_input = PatternVisitor([=](uint32_t const& i) { return barrier_start_ptr + 2 * N_B_STAGES + i; }); - auto empty_input - = PatternVisitor([=](uint32_t const& i) { return barrier_start_ptr + 2 * N_B_STAGES + N_INPUT_STAGES + i; }); - auto full_cast = PatternVisitor( - [=](uint32_t const& i) { return barrier_start_ptr + 2 * N_B_STAGES + 2 * N_INPUT_STAGES + i; }); - auto empty_cast = PatternVisitor([=](uint32_t const& i) - { return barrier_start_ptr + 2 * N_B_STAGES + 2 * N_INPUT_STAGES + kNumCastStages + i; }); - auto tmem_full_barrier = barrier_start_ptr + 2 * N_B_STAGES + 2 * N_INPUT_STAGES + 2 * kNumCastStages; - - cursor += (2 * N_B_STAGES + 2 * N_INPUT_STAGES + 2 * kNumCastStages + 1) * sizeof(Barrier); - auto tmem_ptr_in_smem = reinterpret_cast(cursor); - - if (warp_idx == 1 and cute::elect_one_sync()) - { -#pragma unroll - for (uint32_t i = 0; i < N_B_STAGES; ++i) - { - full_B[i]->init(1); - empty_B[i]->init(1); - } -#pragma unroll - for (uint32_t i = 0; i < N_INPUT_STAGES; ++i) - { - full_input[i]->init(1); - empty_input[i]->init(kNumPmapThreads); - } -#pragma unroll - for (uint32_t i = 0; i < kNumCastStages; ++i) - { - full_cast[i]->init(kNumPmapThreads); - empty_cast[i]->init(1); - } - tmem_full_barrier->init(1); - cutlass::arch::fence_barrier_init(); - } - else if (warp_idx == 2) - { - cute::TMEM::Allocator1Sm().allocate(kNumTmemCols, tmem_ptr_in_smem); - } - __syncthreads(); - - const uint32_t block_idx = __shfl_sync(0xffffffff, blockIdx.x, 0); - const uint32_t m_block_idx = block_idx / kNumSplits; - const uint32_t k_split_idx = block_idx % kNumSplits; - const uint32_t m_offset = m_block_idx * BLOCK_M; - const uint32_t h_tile_start = k_split_idx * H_TILES_PER_SPLIT; - constexpr uint32_t num_total_stages = H_TILES_PER_SPLIT * HC_MULT; - - // Prologue: pmap warp group loads post_mix, comb_mix into SMEM - if (warp_idx >= kNumMMAThreads / 32) - { - const uint32_t pmap_tid = threadIdx.x - kNumMMAThreads; -#pragma unroll - for (uint32_t t = 0; t < 2; ++t) - { - uint32_t idx = pmap_tid + t * kNumPmapThreads; - if (idx < BLOCK_M * HC_MULT) - { - uint32_t m = idx / HC_MULT; - uint32_t hc = idx % HC_MULT; - uint32_t gmem_m = m_offset + m; - float v = (gmem_m < shape_m) ? post_mix[gmem_m * HC_MULT + hc] : 0.f; - smem_post[idx] = v; - } - } -#pragma unroll - for (uint32_t t = 0; t < 8; ++t) - { - uint32_t idx = pmap_tid + t * kNumPmapThreads; - if (idx < BLOCK_M * HC_MULT * HC_MULT) - { - uint32_t m = idx / (HC_MULT * HC_MULT); - uint32_t jk = idx % (HC_MULT * HC_MULT); - uint32_t gmem_m = m_offset + m; - float v = (gmem_m < shape_m) ? comb_mix[gmem_m * HC_MULT * HC_MULT + jk] : 0.f; - smem_comb[idx] = v; - } - } - } - __syncthreads(); - - if (warp_idx < kNumMMAThreads / 32) - { - // ----- TMA warp (warp 0) ----- - if (warp_idx == 0 and cute::elect_one_sync()) - { - uint32_t b_stage = 0; - uint32_t i_stage = 0; - uint32_t s = 0; - for (uint32_t ht = 0; ht < H_TILES_PER_SPLIT; ++ht) - { - const uint32_t h_tile = h_tile_start + ht; - empty_input[i_stage]->wait(((ht / N_INPUT_STAGES) & 1) ^ 1); - uint32_t m_idx = m_block_idx * BLOCK_M; - uint32_t h_idx = h_tile * BLOCK_K; -#pragma unroll - for (uint32_t j = 0; j < HC_MULT; ++j) - { - tma_copy(&tensor_map_residual, full_input[i_stage], - smem_res[i_stage] + j * BLOCK_M * BLOCK_K, j * HIDDEN + h_idx, m_idx); - } - tma_copy( - &tensor_map_x, full_input[i_stage], smem_x_stg[i_stage], h_idx, m_idx); - constexpr uint32_t kInputBytes = SMEM_RES_PER_ISTG + SMEM_X_PER_ISTG; - full_input[i_stage]->arrive_and_expect_tx(kInputBytes); - -#pragma unroll - for (uint32_t hc = 0; hc < HC_MULT; ++hc) - { - empty_B[b_stage]->wait(((s / N_B_STAGES) & 1) ^ 1); - uint32_t k_idx = hc * HIDDEN + h_idx; - tma_copy( - &tensor_map_b, full_B[b_stage], smem_b[b_stage], k_idx, 0); - full_B[b_stage]->arrive_and_expect_tx(SMEM_B_PER_STAGE); - b_stage = (b_stage + 1) % N_B_STAGES; - ++s; - } - i_stage = (i_stage + 1) % N_INPUT_STAGES; - } - } - - // ----- MMA issue warp (warp 1) ----- - if (warp_idx == 1) - { - constexpr uint32_t UMMA_M = BLOCK_M; - constexpr uint32_t UMMA_N = BLOCK_N; - constexpr uint32_t UMMA_K = 32 / sizeof(float); - constexpr uint32_t BLOCK_SWIZZLED_BK = kSwizzleBMode / sizeof(float); - using umma_t = cute::SM100_MMA_TF32_TS; - auto instr_desc = cute::UMMA::make_instr_desc(); - auto const& runtime_instr_desc = cute::UMMA::make_runtime_instr_desc(instr_desc); - static_assert(N_B_STAGES <= 32, "Too many B stages"); - auto b_desc = make_umma_desc(smem_b[0], 0, 0); - uint32_t const& b_desc_lo = lane_idx < N_B_STAGES ? b_desc.lo + lane_idx * SMEM_B_PER_STAGE / 16 : 0u; - - for (uint32_t s = 0; s < num_total_stages; ++s) - { - const uint32_t b_stage = s % N_B_STAGES; - const uint32_t cast_stage_idx = s % kNumCastStages; - full_cast[cast_stage_idx]->wait((s / kNumCastStages) & 1); - full_B[b_stage]->wait((s / N_B_STAGES) & 1); - tcgen05_after_thread_sync(); - auto const& b_desc_base_lo = __shfl_sync(0xffffffff, b_desc_lo, static_cast(b_stage)); -#pragma unroll - for (uint32_t k = 0; k < BLOCK_K / UMMA_K; ++k) - { - uint32_t const& atom_idx = (k * UMMA_K) / BLOCK_SWIZZLED_BK; - uint32_t const& in_atom_idx = (k * UMMA_K) % BLOCK_SWIZZLED_BK; - uint32_t const& offset = atom_idx * BLOCK_N * BLOCK_SWIZZLED_BK; - b_desc.lo = advance_umma_desc_lo( - b_desc_base_lo, offset, in_atom_idx); - umma_t::fma(BLOCK_K * cast_stage_idx + k * UMMA_K, b_desc, BLOCK_K * kNumCastStages, s > 0 or k > 0, - runtime_instr_desc); - } - cutlass::arch::umma_arrive(reinterpret_cast(empty_cast[cast_stage_idx])); - cutlass::arch::umma_arrive(reinterpret_cast(empty_B[b_stage])); - } - cutlass::arch::umma_arrive(reinterpret_cast(tmem_full_barrier)); - } - - // ----- Epilogue (warps 0..3, 128 threads) ----- - constexpr uint32_t kNumBankGroupBytes = 16; - constexpr uint32_t kNumElemsPerBankGroup = kNumBankGroupBytes / sizeof(float); - static_assert(BLOCK_N % kNumElemsPerBankGroup == 0, "Invalid swizzling"); - - tmem_full_barrier->wait(0); - tcgen05_after_thread_sync(); - -#pragma unroll - for (uint32_t i = 0; i < BLOCK_N / kNumElemsPerBankGroup; ++i) - { - uint32_t tmem_addr = BLOCK_K * kNumCastStages + i * kNumElemsPerBankGroup; - auto smem_ptr = reinterpret_cast(smem_cd) + warp_idx * BLOCK_M / 4 * kSwizzleCDMode - + get_swizzled_smem_offset(i, lane_idx); - uint32_t values[kNumElemsPerBankGroup]; - static_assert(kNumElemsPerBankGroup == 4, "Invalid type"); - cute::SM100_TMEM_LOAD_32dp32b4x::copy(tmem_addr, values[0], values[1], values[2], values[3]); - cutlass::arch::fence_view_async_tmem_load(); - if (BLOCK_M == 128 or (BLOCK_M == 64 and lane_idx < 16)) - deep_gemm::ptx::st_shared(smem_ptr, values[0], values[1], values[2], values[3]); - if constexpr (BLOCK_M == 64) - __syncwarp(); - } - cutlass::arch::NamedBarrier::sync(kNumMMAThreads, 0); - - constexpr uint32_t kTotalOut = BLOCK_M * SHAPE_N; - const uint32_t tid = threadIdx.x; -#pragma unroll - for (uint32_t k = tid; k < kTotalOut; k += kNumMMAThreads) - { - uint32_t m = k / SHAPE_N; - uint32_t n = k - m * SHAPE_N; - uint32_t gm = m_block_idx * BLOCK_M + m; - if (gm < shape_m) - { - uint32_t col_group = n >> 2; - uint32_t in_group = n & 3; - uint32_t phys_col_grp = col_group ^ (m & 7); - uint32_t byte = m * kSwizzleCDMode + phys_col_grp * kNumBankGroupBytes + in_group * sizeof(float); - float val = *reinterpret_cast(reinterpret_cast(smem_cd) + byte); - if constexpr (kNumSplits == 1) - { - // Single-CTA-per-(m,n) at KS=1 → safe to store directly. - // Caller no longer needs to pre-zero D (see launcher). - D[gm * SHAPE_N + n] = val; - } - else - { - atomicAdd(&D[gm * SHAPE_N + n], val); - } - } - } - - if (warp_idx == 1) - cute::TMEM::Allocator1Sm().free(0, kNumTmemCols); - } - else - { - // ----- Pmap warp group (warps 4..7, 128 threads) ----- - const uint32_t sub_warp_idx = warp_idx - kNumMMAThreads / 32; - const uint32_t upper_row = sub_warp_idx * 16 + lane_idx / 4; - const uint32_t lower_row = upper_row + 8; - const uint32_t col_lane = lane_idx % 4; - - float pm_u[HC_MULT], pm_l[HC_MULT]; - float cm_u[HC_MULT][HC_MULT], cm_l[HC_MULT][HC_MULT]; -#pragma unroll - for (uint32_t hc = 0; hc < HC_MULT; ++hc) - { - pm_u[hc] = smem_post[upper_row * HC_MULT + hc]; - pm_l[hc] = smem_post[lower_row * HC_MULT + hc]; - } -#pragma unroll - for (uint32_t j = 0; j < HC_MULT; ++j) - { -#pragma unroll - for (uint32_t hc = 0; hc < HC_MULT; ++hc) - { - cm_u[j][hc] = smem_comb[upper_row * HC_MULT * HC_MULT + j * HC_MULT + hc]; - cm_l[j][hc] = smem_comb[lower_row * HC_MULT * HC_MULT + j * HC_MULT + hc]; - } - } - - float sqr_u = 0.f, sqr_l = 0.f; - constexpr uint32_t kNumBankGroupBytes = 16; - constexpr uint32_t kNumElemsPerBankGroup = kNumBankGroupBytes / sizeof(nv_bfloat16); - constexpr uint32_t kNumLoads = BLOCK_K / kNumElemsPerBankGroup; - constexpr uint32_t BLOCK_M_PER_WARP = BLOCK_M / 4; - static_assert(BLOCK_K * sizeof(nv_bfloat16) == kSwizzleAMode, "BLOCK_K must match swizzle A mode"); - static_assert(kNumLoads % 2 == 0, "kNumLoads must be even for LDSM.x4"); - - uint32_t s = 0; - for (uint32_t ht = 0; ht < H_TILES_PER_SPLIT; ++ht) - { - const uint32_t i_stage = ht % N_INPUT_STAGES; - full_input[i_stage]->wait((ht / N_INPUT_STAGES) & 1); - - uint32_t x_vals[2][kNumLoads]; - { - uint8_t const* x_base - = reinterpret_cast(smem_x_stg[i_stage]) + sub_warp_idx * BLOCK_M_PER_WARP * kSwizzleXMode; -#pragma unroll - for (uint32_t i = 0; i < kNumLoads; i += 2) - { - auto smem_ptr = x_base + get_swizzled_smem_offset(i + lane_idx / 16, lane_idx % 16); - deep_gemm::sm90::SM90_U32x4_LDSM_N::copy(x_vals[0][i + 0], x_vals[1][i + 0], x_vals[0][i + 1], - x_vals[1][i + 1], const_cast(smem_ptr)); - } - } - - uint32_t r_vals[HC_MULT][2][kNumLoads]; -#pragma unroll - for (uint32_t j = 0; j < HC_MULT; ++j) - { - uint8_t const* r_base = reinterpret_cast(smem_res[i_stage]) - + j * BLOCK_M * BLOCK_K * sizeof(nv_bfloat16) + sub_warp_idx * BLOCK_M_PER_WARP * kSwizzleResMode; -#pragma unroll - for (uint32_t i = 0; i < kNumLoads; i += 2) - { - auto smem_ptr - = r_base + get_swizzled_smem_offset(i + lane_idx / 16, lane_idx % 16); - deep_gemm::sm90::SM90_U32x4_LDSM_N::copy(r_vals[j][0][i + 0], r_vals[j][1][i + 0], - r_vals[j][0][i + 1], r_vals[j][1][i + 1], const_cast(smem_ptr)); - } - } - - float2 xf[2][kNumLoads]; -#pragma unroll - for (uint32_t u = 0; u < 2; ++u) - { -#pragma unroll - for (uint32_t i = 0; i < kNumLoads; ++i) - { - xf[u][i] = __bfloat1622float2(*reinterpret_cast(&x_vals[u][i])); - } - } - - if constexpr (kEarlyRelease) - { - empty_input[i_stage]->arrive(); - } - - // Wait for previous ht's residual_out TMA_STOREs to drain before we - // overwrite single-buffered smem_rc with new hc values. - if (ht > 0) - { - cute::tma_store_wait<0>(); - } - -#pragma unroll - for (uint32_t hc = 0; hc < HC_MULT; ++hc) - { - const uint32_t cast_stage_idx = s % kNumCastStages; - empty_cast[cast_stage_idx]->wait(((s / kNumCastStages) & 1) ^ 1); - - uint32_t rc_u_buf[kNumLoads], rc_l_buf[kNumLoads]; -#pragma unroll - for (uint32_t i = 0; i < kNumLoads; ++i) - { - float2 nu{pm_u[hc] * xf[0][i].x, pm_u[hc] * xf[0][i].y}; - float2 nl{pm_l[hc] * xf[1][i].x, pm_l[hc] * xf[1][i].y}; -#pragma unroll - for (uint32_t j = 0; j < HC_MULT; ++j) - { - float2 ruj = __bfloat1622float2(*reinterpret_cast(&r_vals[j][0][i])); - float2 rlj = __bfloat1622float2(*reinterpret_cast(&r_vals[j][1][i])); - nu.x = fmaf(cm_u[j][hc], ruj.x, nu.x); - nu.y = fmaf(cm_u[j][hc], ruj.y, nu.y); - nl.x = fmaf(cm_l[j][hc], rlj.x, nl.x); - nl.y = fmaf(cm_l[j][hc], rlj.y, nl.y); - } - nv_bfloat162 b_up = __float22bfloat162_rn(nu); - nv_bfloat162 b_lo = __float22bfloat162_rn(nl); - uint32_t b_up_bits = *reinterpret_cast(&b_up); - uint32_t b_lo_bits = *reinterpret_cast(&b_lo); - rc_u_buf[i] = b_up_bits; - rc_l_buf[i] = b_lo_bits; - float2 ru = __bfloat1622float2(b_up); - float2 rl = __bfloat1622float2(b_lo); - sqr_u = fmaf(ru.x, ru.x, sqr_u); - sqr_u = fmaf(ru.y, ru.y, sqr_u); - sqr_l = fmaf(rl.x, rl.x, sqr_l); - sqr_l = fmaf(rl.y, rl.y, sqr_l); - cute::SM100_TMEM_STORE_16dp256b1x::copy(*reinterpret_cast(&ru.x), - *reinterpret_cast(&ru.y), *reinterpret_cast(&rl.x), - *reinterpret_cast(&rl.y), cast_stage_idx * BLOCK_K + i * 8); - } - cutlass::arch::fence_view_async_tmem_store(); - tcgen05_before_thread_sync(); - full_cast[cast_stage_idx]->arrive(); - ++s; - - // STSM bf16 new_r values into smem_rc[hc] sub-region for this warp. - uint8_t* rc_base = reinterpret_cast(smem_rc) + hc * SMEM_RC_PER_HC - + sub_warp_idx * BLOCK_M_PER_WARP * kSwizzleRoutMode; -#pragma unroll - for (uint32_t i = 0; i < kNumLoads; i += 2) - { - auto smem_ptr - = rc_base + get_swizzled_smem_offset(i + lane_idx / 16, lane_idx % 16); - stsm_x4_b16_rout(smem_ptr, rc_u_buf[i + 0], rc_l_buf[i + 0], rc_u_buf[i + 1], rc_l_buf[i + 1]); - } - } - if constexpr (!kEarlyRelease) - { - empty_input[i_stage]->arrive(); - } - - // Emit HC_MULT TMA_STOREs of residual_cur: one per hc slice, per-warp rows. - cute::tma_store_fence(); - if (cute::elect_one_sync()) - { - const uint32_t h_idx = (h_tile_start + ht) * BLOCK_K; -#pragma unroll - for (uint32_t hc = 0; hc < HC_MULT; ++hc) - { - uint8_t* rc_base = reinterpret_cast(smem_rc) + hc * SMEM_RC_PER_HC - + sub_warp_idx * BLOCK_M_PER_WARP * kSwizzleRoutMode; - cute::SM90_TMA_STORE_2D::copy(&tensor_map_residual_out, rc_base, hc * HIDDEN + h_idx, - m_offset + sub_warp_idx * BLOCK_M_PER_WARP); - cute::tma_store_arrive(); - } - } - } - - // Drain any in-flight residual_out TMA stores before exit. - cute::tma_store_wait<0>(); - - // Warp-reduce sqr across 4 col_lanes then atomicAdd to global. - sqr_u += __shfl_xor_sync(0xffffffff, sqr_u, 1); - sqr_u += __shfl_xor_sync(0xffffffff, sqr_u, 2); - sqr_l += __shfl_xor_sync(0xffffffff, sqr_l, 1); - sqr_l += __shfl_xor_sync(0xffffffff, sqr_l, 2); - if (col_lane == 0) - { - uint32_t gm_u = m_block_idx * BLOCK_M + upper_row; - uint32_t gm_l = m_block_idx * BLOCK_M + lower_row; - if constexpr (kNumSplits == 1) - { - // KS=1 → only this CTA writes (gm_u, gm_l), no race possible. - if (gm_u < shape_m) - sqr_sum[gm_u] = sqr_u; - if (gm_l < shape_m) - sqr_sum[gm_l] = sqr_l; - } - else - { - if (gm_u < shape_m) - atomicAdd(&sqr_sum[gm_u], sqr_u); - if (gm_l < shape_m) - atomicAdd(&sqr_sum[gm_l], sqr_l); - } - } - } -#else - if (blockIdx.x == 0 and threadIdx.x == 0) - DG_DEVICE_ASSERT(false and "This kernel only supports sm_100a"); -#endif -} - -// ============================================================================ -// ALL-IN-ONE variant (Path D, tf32 tcgen05 MMA analogue of Path F). -// -// Single-kernel fusion of: -// 1) post_mapping : residual_cur[j,h] = post_mix_prev[j]*x_prev[h] -// + sum_k comb_mix_prev[k,j]*residual_prev[k,h] -// 2) pre_GEMM : D[i,n] = sum_{hc,h} residual_cur[i,hc,h] * W_T[n, hc*HIDDEN+h] -// sqr[i] = sum_{hc,h} residual_cur[i,hc,h]^2 -// 3) bigFuse : rmsnorm+sigmoid+sinkhorn on D/sqr -> post_mix_out, -// comb_mix_out, pre_mix; layer_input = pre_mix @ residual_cur. -// -// Semantics match Path F (fused_pmap_gemm_fma_allinone) exactly. The only -// algorithmic difference vs Path F is that we use tf32 tcgen05 MMA for the -// residual_cur @ W_T GEMM (instead of CUDA-core FMA). -// -// Pipelining, warp layout, and SMEM layout inherit from Path B -// (fused_tf32_pmap_gemm_rout_atomic_impl): the pmap warp group computes -// residual_cur, TMA-stores it to GMEM, and STSM-casts it into TMEM; the MMA -// warp group consumes TMEM for the GEMM; the epilogue atomicAdd's into -// y_acc / sqr_sum. Phase 3 elects the last-home CTA (per m_block) via -// atomicAdd on done_counter; Phase 4 runs bigfuse inline on that CTA only, -// reloading residual_cur from GMEM and writing layer_input + post_mix_out + -// comb_mix_out. -// -// Caller MUST zero D (y_acc), sqr_sum (r_acc), done_counter before launch. -// ============================================================================ - -template -__global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) - fused_allinone_tf32_pmap_gemm_atomic_impl(const uint32_t shape_m, - const __grid_constant__ cute::TmaDescriptor tensor_map_residual, // residual_prev, bf16 - const __grid_constant__ cute::TmaDescriptor tensor_map_x, // x_prev, bf16 - const __grid_constant__ cute::TmaDescriptor tensor_map_b, // W_T, tf32 - const __grid_constant__ cute::TmaDescriptor tensor_map_residual_out, // residual_cur, bf16 (TMA store) - __nv_bfloat16 const* __restrict__ residual_cur_ptr, // same buffer as TMA target - __nv_bfloat16* __restrict__ layer_input_out, // [M, HIDDEN] bf16 - float* __restrict__ D, // [M, SHAPE_N] fp32 (y_acc, caller zeros) - float* __restrict__ sqr_sum, // [M] fp32 (r_acc, caller zeros) - int* __restrict__ done_counter, // [ceil(M/BLOCK_M)] int (caller zeros) - float const* __restrict__ post_mix_prev, // [M, HC_MULT] - float const* __restrict__ comb_mix_prev, // [M, HC_MULT, HC_MULT] - float const* __restrict__ hc_scale, // [3] - float const* __restrict__ hc_base, // [HC_MULT*(2+HC_MULT)] - float* __restrict__ post_mix_out, // [M, HC_MULT] - float* __restrict__ comb_mix_out, // [M, HC_MULT, HC_MULT] - // When kFuseNorm: layer_input_out receives the RMSNorm-normalized - // values: out[t,h] = bf16(li[t,h] * rsqrt(mean(li²)+norm_eps) * w[h]). - // norm_weight must be bf16 [HIDDEN]; norm_eps is the RMSNorm epsilon. - // When kFuseNorm is false, norm_weight/norm_eps are ignored. - __nv_bfloat16 const* __restrict__ norm_weight, float norm_eps, float rms_eps, float hc_pre_eps, - float hc_sinkhorn_eps, float hc_post_mult_value, uint32_t sinkhorn_repeat) -{ -#if (defined(__CUDA_ARCH__) and (__CUDA_ARCH__ >= 1000) and (__CUDA_ARCH__ < 1100)) or defined(__CLION_IDE__) - using Barrier = cutlass::arch::ClusterTransactionBarrier; - - constexpr uint32_t HC_MULT2 = HC_MULT * HC_MULT; - constexpr uint32_t HC_MULT3 = HC_MULT * (2 + HC_MULT); - constexpr uint32_t SHAPE_K = HC_MULT * HIDDEN; - constexpr uint32_t H_TILES_PER_HC = HIDDEN / BLOCK_K; - static_assert(H_TILES_PER_HC % kNumSplits == 0, "H_TILES_PER_HC must be divisible by kNumSplits"); - constexpr uint32_t H_TILES_PER_SPLIT = H_TILES_PER_HC / kNumSplits; - constexpr uint32_t kNumCastStages = 4; - constexpr uint32_t kSwizzleAMode = cute::min(BLOCK_K * sizeof(nv_bfloat16), 128); - constexpr uint32_t kSwizzleBMode = cute::min(BLOCK_K * sizeof(float), 128); - constexpr uint32_t kSwizzleXMode = kSwizzleAMode; - constexpr uint32_t kSwizzleResMode = kSwizzleAMode; - constexpr uint32_t kSwizzleRoutMode = kSwizzleAMode; - constexpr auto kMajorA = cute::UMMA::Major::K; - constexpr auto kMajorB = cute::UMMA::Major::K; - static_assert(HIDDEN % BLOCK_K == 0, "HIDDEN must be multiple of BLOCK_K"); - static_assert(N_B_STAGES >= HC_MULT, "N_B_STAGES must be >= HC_MULT"); - static_assert(kSwizzleCDMode / sizeof(float) == BLOCK_N, "Invalid block N"); - static_assert(kNumMMAThreads == 128, "Invalid MMA threads"); - static_assert(kNumPmapThreads == 128, "Invalid pmap threads"); - static_assert(BLOCK_M == 64, "Invalid block M"); - static_assert(HC_MULT == 4, "Only HC_MULT=4 supported"); - static_assert(kSwizzleCDMode == 128, "Atomic variant expects kSwizzleCDMode=128"); - static_assert(SHAPE_N <= BLOCK_N, "SHAPE_N must fit within BLOCK_N"); - static_assert(SHAPE_N == HC_MULT3, "Path D expects SHAPE_N == HC_MULT*(2+HC_MULT)=24"); - - auto const warp_idx = cutlass::canonical_warp_idx_sync(); - auto const lane_idx = get_lane_idx(); - - extern __shared__ __align__(1024) uint8_t smem_buffer[]; - - constexpr uint32_t SMEM_CD_SIZE = BLOCK_M * kSwizzleCDMode; - constexpr uint32_t SMEM_B_PER_STAGE = BLOCK_N * BLOCK_K * sizeof(float); - constexpr uint32_t SMEM_RES_PER_ISTG = BLOCK_M * HC_MULT * BLOCK_K * sizeof(nv_bfloat16); - constexpr uint32_t SMEM_X_PER_ISTG = BLOCK_M * BLOCK_K * sizeof(nv_bfloat16); - constexpr uint32_t SMEM_POST_SIZE = BLOCK_M * HC_MULT * sizeof(float); - constexpr uint32_t SMEM_COMB_SIZE = BLOCK_M * HC_MULT * HC_MULT * sizeof(float); - constexpr uint32_t SMEM_RC_PER_HC = BLOCK_M * BLOCK_K * sizeof(nv_bfloat16); // 8 KB - constexpr uint32_t SMEM_RC_SIZE = HC_MULT * SMEM_RC_PER_HC; // 32 KB - - constexpr uint32_t kNumTmemCols = get_num_aligned_tmem_cols(); - - // Prefetch TMA descriptors - if (warp_idx == 0 and cute::elect_one_sync()) - { - cute::prefetch_tma_descriptor(&tensor_map_residual); - cute::prefetch_tma_descriptor(&tensor_map_x); - cute::prefetch_tma_descriptor(&tensor_map_b); - cute::prefetch_tma_descriptor(&tensor_map_residual_out); - } - - // SMEM layout: [cd, B stages, res stages, x stages, post, comb, rc (HC_MULT slices)] - auto smem_cd = reinterpret_cast(smem_buffer); - uint8_t* cursor = smem_buffer + SMEM_CD_SIZE; - auto smem_b = PatternVisitor( - [&, base = cursor](uint32_t const& i) { return reinterpret_cast(base + i * SMEM_B_PER_STAGE); }); - cursor += N_B_STAGES * SMEM_B_PER_STAGE; - auto smem_res = PatternVisitor( - [&, base = cursor](uint32_t const& i) { return reinterpret_cast(base + i * SMEM_RES_PER_ISTG); }); - cursor += N_INPUT_STAGES * SMEM_RES_PER_ISTG; - auto smem_x_stg = PatternVisitor( - [&, base = cursor](uint32_t const& i) { return reinterpret_cast(base + i * SMEM_X_PER_ISTG); }); - cursor += N_INPUT_STAGES * SMEM_X_PER_ISTG; - auto smem_post = reinterpret_cast(cursor); - cursor += SMEM_POST_SIZE; - auto smem_comb = reinterpret_cast(cursor); - cursor += SMEM_COMB_SIZE; - auto smem_rc = reinterpret_cast(cursor); // [HC_MULT][BLOCK_M][BLOCK_K] bf16 - cursor += SMEM_RC_SIZE; - - cursor = reinterpret_cast((reinterpret_cast(cursor) + 7) & ~uintptr_t(7)); - auto barrier_start_ptr = reinterpret_cast(cursor); - auto full_B = PatternVisitor([=](uint32_t const& i) { return barrier_start_ptr + i; }); - auto empty_B = PatternVisitor([=](uint32_t const& i) { return barrier_start_ptr + N_B_STAGES + i; }); - auto full_input = PatternVisitor([=](uint32_t const& i) { return barrier_start_ptr + 2 * N_B_STAGES + i; }); - auto empty_input - = PatternVisitor([=](uint32_t const& i) { return barrier_start_ptr + 2 * N_B_STAGES + N_INPUT_STAGES + i; }); - auto full_cast = PatternVisitor( - [=](uint32_t const& i) { return barrier_start_ptr + 2 * N_B_STAGES + 2 * N_INPUT_STAGES + i; }); - auto empty_cast = PatternVisitor([=](uint32_t const& i) - { return barrier_start_ptr + 2 * N_B_STAGES + 2 * N_INPUT_STAGES + kNumCastStages + i; }); - auto tmem_full_barrier = barrier_start_ptr + 2 * N_B_STAGES + 2 * N_INPUT_STAGES + 2 * kNumCastStages; - - cursor += (2 * N_B_STAGES + 2 * N_INPUT_STAGES + 2 * kNumCastStages + 1) * sizeof(Barrier); - auto tmem_ptr_in_smem = reinterpret_cast(cursor); - - if (warp_idx == 1 and cute::elect_one_sync()) - { -#pragma unroll - for (uint32_t i = 0; i < N_B_STAGES; ++i) - { - full_B[i]->init(1); - empty_B[i]->init(1); - } -#pragma unroll - for (uint32_t i = 0; i < N_INPUT_STAGES; ++i) - { - full_input[i]->init(1); - empty_input[i]->init(kNumPmapThreads); - } -#pragma unroll - for (uint32_t i = 0; i < kNumCastStages; ++i) - { - full_cast[i]->init(kNumPmapThreads); - empty_cast[i]->init(1); - } - tmem_full_barrier->init(1); - cutlass::arch::fence_barrier_init(); - } - else if (warp_idx == 2) - { - cute::TMEM::Allocator1Sm().allocate(kNumTmemCols, tmem_ptr_in_smem); - } - __syncthreads(); - - const uint32_t block_idx = __shfl_sync(0xffffffff, blockIdx.x, 0); - const uint32_t m_block_idx = block_idx / kNumSplits; - const uint32_t k_split_idx = block_idx % kNumSplits; - const uint32_t m_offset = m_block_idx * BLOCK_M; - const uint32_t h_tile_start = k_split_idx * H_TILES_PER_SPLIT; - constexpr uint32_t num_total_stages = H_TILES_PER_SPLIT * HC_MULT; - - // Prologue: pmap warp group loads post_mix_prev, comb_mix_prev into SMEM - if (warp_idx >= kNumMMAThreads / 32) - { - const uint32_t pmap_tid = threadIdx.x - kNumMMAThreads; -#pragma unroll - for (uint32_t t = 0; t < 2; ++t) - { - uint32_t idx = pmap_tid + t * kNumPmapThreads; - if (idx < BLOCK_M * HC_MULT) - { - uint32_t m = idx / HC_MULT; - uint32_t hc = idx % HC_MULT; - uint32_t gmem_m = m_offset + m; - float v = (gmem_m < shape_m) ? post_mix_prev[gmem_m * HC_MULT + hc] : 0.f; - smem_post[idx] = v; - } - } -#pragma unroll - for (uint32_t t = 0; t < 8; ++t) - { - uint32_t idx = pmap_tid + t * kNumPmapThreads; - if (idx < BLOCK_M * HC_MULT * HC_MULT) - { - uint32_t m = idx / (HC_MULT * HC_MULT); - uint32_t jk = idx % (HC_MULT * HC_MULT); - uint32_t gmem_m = m_offset + m; - float v = (gmem_m < shape_m) ? comb_mix_prev[gmem_m * HC_MULT * HC_MULT + jk] : 0.f; - smem_comb[idx] = v; - } - } - } - __syncthreads(); - - if (warp_idx < kNumMMAThreads / 32) - { - // ----- TMA warp (warp 0) ----- - if (warp_idx == 0 and cute::elect_one_sync()) - { - uint32_t b_stage = 0; - uint32_t i_stage = 0; - uint32_t s = 0; - for (uint32_t ht = 0; ht < H_TILES_PER_SPLIT; ++ht) - { - const uint32_t h_tile = h_tile_start + ht; - empty_input[i_stage]->wait(((ht / N_INPUT_STAGES) & 1) ^ 1); - uint32_t m_idx = m_block_idx * BLOCK_M; - uint32_t h_idx = h_tile * BLOCK_K; -#pragma unroll - for (uint32_t j = 0; j < HC_MULT; ++j) - { - tma_copy(&tensor_map_residual, full_input[i_stage], - smem_res[i_stage] + j * BLOCK_M * BLOCK_K, j * HIDDEN + h_idx, m_idx); - } - tma_copy( - &tensor_map_x, full_input[i_stage], smem_x_stg[i_stage], h_idx, m_idx); - constexpr uint32_t kInputBytes = SMEM_RES_PER_ISTG + SMEM_X_PER_ISTG; - full_input[i_stage]->arrive_and_expect_tx(kInputBytes); - -#pragma unroll - for (uint32_t hc = 0; hc < HC_MULT; ++hc) - { - empty_B[b_stage]->wait(((s / N_B_STAGES) & 1) ^ 1); - uint32_t k_idx = hc * HIDDEN + h_idx; - tma_copy( - &tensor_map_b, full_B[b_stage], smem_b[b_stage], k_idx, 0); - full_B[b_stage]->arrive_and_expect_tx(SMEM_B_PER_STAGE); - b_stage = (b_stage + 1) % N_B_STAGES; - ++s; - } - i_stage = (i_stage + 1) % N_INPUT_STAGES; - } - } - - // ----- MMA issue warp (warp 1) ----- - if (warp_idx == 1) - { - constexpr uint32_t UMMA_M = BLOCK_M; - constexpr uint32_t UMMA_N = BLOCK_N; - constexpr uint32_t UMMA_K = 32 / sizeof(float); - constexpr uint32_t BLOCK_SWIZZLED_BK = kSwizzleBMode / sizeof(float); - using umma_t = cute::SM100_MMA_TF32_TS; - auto instr_desc = cute::UMMA::make_instr_desc(); - auto const& runtime_instr_desc = cute::UMMA::make_runtime_instr_desc(instr_desc); - static_assert(N_B_STAGES <= 32, "Too many B stages"); - auto b_desc = make_umma_desc(smem_b[0], 0, 0); - uint32_t const& b_desc_lo = lane_idx < N_B_STAGES ? b_desc.lo + lane_idx * SMEM_B_PER_STAGE / 16 : 0u; - - for (uint32_t s = 0; s < num_total_stages; ++s) - { - const uint32_t b_stage = s % N_B_STAGES; - const uint32_t cast_stage_idx = s % kNumCastStages; - full_cast[cast_stage_idx]->wait((s / kNumCastStages) & 1); - full_B[b_stage]->wait((s / N_B_STAGES) & 1); - tcgen05_after_thread_sync(); - auto const& b_desc_base_lo = __shfl_sync(0xffffffff, b_desc_lo, static_cast(b_stage)); -#pragma unroll - for (uint32_t k = 0; k < BLOCK_K / UMMA_K; ++k) - { - uint32_t const& atom_idx = (k * UMMA_K) / BLOCK_SWIZZLED_BK; - uint32_t const& in_atom_idx = (k * UMMA_K) % BLOCK_SWIZZLED_BK; - uint32_t const& offset = atom_idx * BLOCK_N * BLOCK_SWIZZLED_BK; - b_desc.lo = advance_umma_desc_lo( - b_desc_base_lo, offset, in_atom_idx); - umma_t::fma(BLOCK_K * cast_stage_idx + k * UMMA_K, b_desc, BLOCK_K * kNumCastStages, s > 0 or k > 0, - runtime_instr_desc); - } - cutlass::arch::umma_arrive(reinterpret_cast(empty_cast[cast_stage_idx])); - cutlass::arch::umma_arrive(reinterpret_cast(empty_B[b_stage])); - } - cutlass::arch::umma_arrive(reinterpret_cast(tmem_full_barrier)); - } - - // ----- Epilogue (warps 0..3, 128 threads) ----- - constexpr uint32_t kNumBankGroupBytes = 16; - constexpr uint32_t kNumElemsPerBankGroup = kNumBankGroupBytes / sizeof(float); - static_assert(BLOCK_N % kNumElemsPerBankGroup == 0, "Invalid swizzling"); - - tmem_full_barrier->wait(0); - tcgen05_after_thread_sync(); - -#pragma unroll - for (uint32_t i = 0; i < BLOCK_N / kNumElemsPerBankGroup; ++i) - { - uint32_t tmem_addr = BLOCK_K * kNumCastStages + i * kNumElemsPerBankGroup; - auto smem_ptr = reinterpret_cast(smem_cd) + warp_idx * BLOCK_M / 4 * kSwizzleCDMode - + get_swizzled_smem_offset(i, lane_idx); - uint32_t values[kNumElemsPerBankGroup]; - static_assert(kNumElemsPerBankGroup == 4, "Invalid type"); - cute::SM100_TMEM_LOAD_32dp32b4x::copy(tmem_addr, values[0], values[1], values[2], values[3]); - cutlass::arch::fence_view_async_tmem_load(); - if (BLOCK_M == 128 or (BLOCK_M == 64 and lane_idx < 16)) - deep_gemm::ptx::st_shared(smem_ptr, values[0], values[1], values[2], values[3]); - if constexpr (BLOCK_M == 64) - __syncwarp(); - } - cutlass::arch::NamedBarrier::sync(kNumMMAThreads, 0); - - constexpr uint32_t kTotalOut = BLOCK_M * SHAPE_N; - const uint32_t tid = threadIdx.x; -#pragma unroll - for (uint32_t k = tid; k < kTotalOut; k += kNumMMAThreads) - { - uint32_t m = k / SHAPE_N; - uint32_t n = k - m * SHAPE_N; - uint32_t gm = m_block_idx * BLOCK_M + m; - if (gm < shape_m) - { - uint32_t col_group = n >> 2; - uint32_t in_group = n & 3; - uint32_t phys_col_grp = col_group ^ (m & 7); - uint32_t byte = m * kSwizzleCDMode + phys_col_grp * kNumBankGroupBytes + in_group * sizeof(float); - float val = *reinterpret_cast(reinterpret_cast(smem_cd) + byte); - if constexpr (kNumSplits == 1) - { - // Single-CTA-per-(m,n) at KS=1 → safe to store directly. - // Caller no longer needs to pre-zero D (see launcher). - D[gm * SHAPE_N + n] = val; - } - else - { - atomicAdd(&D[gm * SHAPE_N + n], val); - } - } - } - - if (warp_idx == 1) - cute::TMEM::Allocator1Sm().free(0, kNumTmemCols); - } - else - { - // ----- Pmap warp group (warps 4..7, 128 threads) ----- - const uint32_t sub_warp_idx = warp_idx - kNumMMAThreads / 32; - const uint32_t upper_row = sub_warp_idx * 16 + lane_idx / 4; - const uint32_t lower_row = upper_row + 8; - const uint32_t col_lane = lane_idx % 4; - - float pm_u[HC_MULT], pm_l[HC_MULT]; - float cm_u[HC_MULT][HC_MULT], cm_l[HC_MULT][HC_MULT]; -#pragma unroll - for (uint32_t hc = 0; hc < HC_MULT; ++hc) - { - pm_u[hc] = smem_post[upper_row * HC_MULT + hc]; - pm_l[hc] = smem_post[lower_row * HC_MULT + hc]; - } -#pragma unroll - for (uint32_t j = 0; j < HC_MULT; ++j) - { -#pragma unroll - for (uint32_t hc = 0; hc < HC_MULT; ++hc) - { - cm_u[j][hc] = smem_comb[upper_row * HC_MULT * HC_MULT + j * HC_MULT + hc]; - cm_l[j][hc] = smem_comb[lower_row * HC_MULT * HC_MULT + j * HC_MULT + hc]; - } - } - - float sqr_u = 0.f, sqr_l = 0.f; - constexpr uint32_t kNumBankGroupBytes = 16; - constexpr uint32_t kNumElemsPerBankGroup = kNumBankGroupBytes / sizeof(nv_bfloat16); - constexpr uint32_t kNumLoads = BLOCK_K / kNumElemsPerBankGroup; - constexpr uint32_t BLOCK_M_PER_WARP = BLOCK_M / 4; - static_assert(BLOCK_K * sizeof(nv_bfloat16) == kSwizzleAMode, "BLOCK_K must match swizzle A mode"); - static_assert(kNumLoads % 2 == 0, "kNumLoads must be even for LDSM.x4"); - - uint32_t s = 0; - for (uint32_t ht = 0; ht < H_TILES_PER_SPLIT; ++ht) - { - const uint32_t i_stage = ht % N_INPUT_STAGES; - full_input[i_stage]->wait((ht / N_INPUT_STAGES) & 1); - - uint32_t x_vals[2][kNumLoads]; - { - uint8_t const* x_base - = reinterpret_cast(smem_x_stg[i_stage]) + sub_warp_idx * BLOCK_M_PER_WARP * kSwizzleXMode; -#pragma unroll - for (uint32_t i = 0; i < kNumLoads; i += 2) - { - auto smem_ptr = x_base + get_swizzled_smem_offset(i + lane_idx / 16, lane_idx % 16); - deep_gemm::sm90::SM90_U32x4_LDSM_N::copy(x_vals[0][i + 0], x_vals[1][i + 0], x_vals[0][i + 1], - x_vals[1][i + 1], const_cast(smem_ptr)); - } - } - - uint32_t r_vals[HC_MULT][2][kNumLoads]; -#pragma unroll - for (uint32_t j = 0; j < HC_MULT; ++j) - { - uint8_t const* r_base = reinterpret_cast(smem_res[i_stage]) - + j * BLOCK_M * BLOCK_K * sizeof(nv_bfloat16) + sub_warp_idx * BLOCK_M_PER_WARP * kSwizzleResMode; -#pragma unroll - for (uint32_t i = 0; i < kNumLoads; i += 2) - { - auto smem_ptr - = r_base + get_swizzled_smem_offset(i + lane_idx / 16, lane_idx % 16); - deep_gemm::sm90::SM90_U32x4_LDSM_N::copy(r_vals[j][0][i + 0], r_vals[j][1][i + 0], - r_vals[j][0][i + 1], r_vals[j][1][i + 1], const_cast(smem_ptr)); - } - } - - float2 xf[2][kNumLoads]; -#pragma unroll - for (uint32_t u = 0; u < 2; ++u) - { -#pragma unroll - for (uint32_t i = 0; i < kNumLoads; ++i) - { - xf[u][i] = __bfloat1622float2(*reinterpret_cast(&x_vals[u][i])); - } - } - - // Wait for previous ht's residual_out TMA_STOREs to drain before we - // overwrite single-buffered smem_rc with new hc values. - if (ht > 0) - { - cute::tma_store_wait<0>(); - } - -#pragma unroll - for (uint32_t hc = 0; hc < HC_MULT; ++hc) - { - const uint32_t cast_stage_idx = s % kNumCastStages; - empty_cast[cast_stage_idx]->wait(((s / kNumCastStages) & 1) ^ 1); - - uint32_t rc_u_buf[kNumLoads], rc_l_buf[kNumLoads]; -#pragma unroll - for (uint32_t i = 0; i < kNumLoads; ++i) - { - float2 nu{pm_u[hc] * xf[0][i].x, pm_u[hc] * xf[0][i].y}; - float2 nl{pm_l[hc] * xf[1][i].x, pm_l[hc] * xf[1][i].y}; -#pragma unroll - for (uint32_t j = 0; j < HC_MULT; ++j) - { - float2 ruj = __bfloat1622float2(*reinterpret_cast(&r_vals[j][0][i])); - float2 rlj = __bfloat1622float2(*reinterpret_cast(&r_vals[j][1][i])); - nu.x = fmaf(cm_u[j][hc], ruj.x, nu.x); - nu.y = fmaf(cm_u[j][hc], ruj.y, nu.y); - nl.x = fmaf(cm_l[j][hc], rlj.x, nl.x); - nl.y = fmaf(cm_l[j][hc], rlj.y, nl.y); - } - nv_bfloat162 b_up = __float22bfloat162_rn(nu); - nv_bfloat162 b_lo = __float22bfloat162_rn(nl); - uint32_t b_up_bits = *reinterpret_cast(&b_up); - uint32_t b_lo_bits = *reinterpret_cast(&b_lo); - rc_u_buf[i] = b_up_bits; - rc_l_buf[i] = b_lo_bits; - float2 ru = __bfloat1622float2(b_up); - float2 rl = __bfloat1622float2(b_lo); - sqr_u = fmaf(ru.x, ru.x, sqr_u); - sqr_u = fmaf(ru.y, ru.y, sqr_u); - sqr_l = fmaf(rl.x, rl.x, sqr_l); - sqr_l = fmaf(rl.y, rl.y, sqr_l); - cute::SM100_TMEM_STORE_16dp256b1x::copy(*reinterpret_cast(&ru.x), - *reinterpret_cast(&ru.y), *reinterpret_cast(&rl.x), - *reinterpret_cast(&rl.y), cast_stage_idx * BLOCK_K + i * 8); - } - cutlass::arch::fence_view_async_tmem_store(); - tcgen05_before_thread_sync(); - full_cast[cast_stage_idx]->arrive(); - ++s; - - // STSM bf16 new_r values into smem_rc[hc] sub-region for this warp. - uint8_t* rc_base = reinterpret_cast(smem_rc) + hc * SMEM_RC_PER_HC - + sub_warp_idx * BLOCK_M_PER_WARP * kSwizzleRoutMode; -#pragma unroll - for (uint32_t i = 0; i < kNumLoads; i += 2) - { - auto smem_ptr - = rc_base + get_swizzled_smem_offset(i + lane_idx / 16, lane_idx % 16); - stsm_x4_b16_rout(smem_ptr, rc_u_buf[i + 0], rc_l_buf[i + 0], rc_u_buf[i + 1], rc_l_buf[i + 1]); - } - } - empty_input[i_stage]->arrive(); - - // Emit HC_MULT TMA_STOREs of residual_cur: one per hc slice, per-warp rows. - cute::tma_store_fence(); - if (cute::elect_one_sync()) - { - const uint32_t h_idx = (h_tile_start + ht) * BLOCK_K; -#pragma unroll - for (uint32_t hc = 0; hc < HC_MULT; ++hc) - { - uint8_t* rc_base = reinterpret_cast(smem_rc) + hc * SMEM_RC_PER_HC - + sub_warp_idx * BLOCK_M_PER_WARP * kSwizzleRoutMode; - cute::SM90_TMA_STORE_2D::copy(&tensor_map_residual_out, rc_base, hc * HIDDEN + h_idx, - m_offset + sub_warp_idx * BLOCK_M_PER_WARP); - cute::tma_store_arrive(); - } - } - } - - // Drain any in-flight residual_out TMA stores before exit. - cute::tma_store_wait<0>(); - - // Warp-reduce sqr across 4 col_lanes then atomicAdd to global. - sqr_u += __shfl_xor_sync(0xffffffff, sqr_u, 1); - sqr_u += __shfl_xor_sync(0xffffffff, sqr_u, 2); - sqr_l += __shfl_xor_sync(0xffffffff, sqr_l, 1); - sqr_l += __shfl_xor_sync(0xffffffff, sqr_l, 2); - if (col_lane == 0) - { - uint32_t gm_u = m_block_idx * BLOCK_M + upper_row; - uint32_t gm_l = m_block_idx * BLOCK_M + lower_row; - if constexpr (kNumSplits == 1) - { - // KS=1 → only this CTA writes (gm_u, gm_l), no race possible. - if (gm_u < shape_m) - sqr_sum[gm_u] = sqr_u; - if (gm_l < shape_m) - sqr_sum[gm_l] = sqr_l; - } - else - { - if (gm_u < shape_m) - atomicAdd(&sqr_sum[gm_u], sqr_u); - if (gm_l < shape_m) - atomicAdd(&sqr_sum[gm_l], sqr_l); - } - } - } - - // ======================================================================== - // Phase 3: cross-split barrier. - // For kNumSplits == 1 we only need a block-scope fence + __syncthreads. - // For kNumSplits > 1 ALL splits participate in Phase 4: each CTA - // increments done_counter, then spin-waits until the counter reaches - // kNumSplits. Phase 4 work is then partitioned across CTAs by - // k_split_idx — CTA i processes tokens - // [i * TOKS_PER_CTA, (i+1) * TOKS_PER_CTA) - // where TOKS_PER_CTA = BLOCK_M / kNumSplits. This replaces the old - // "last-home CTA does all Phase 4 work serially" design that bottlenecked - // Path D at BLOCK_M=64 (1 CTA's Phase 4 = ~28 µs regardless of M). - // ======================================================================== - if constexpr (kNumSplits == 1) - { - __threadfence_block(); - __syncthreads(); - } - else - { - __threadfence(); - __syncthreads(); - if (threadIdx.x == 0) - { - atomicAdd(&done_counter[m_block_idx], 1); - // Spin-wait until all kNumSplits CTAs finish Phase 2. The - // atomicAdd(..., 0) is a zero-increment load with full device - // coherence — cheap on B200 and avoids an extra flag allocation. - while (atomicAdd(&done_counter[m_block_idx], 0) < static_cast(kNumSplits)) - { - /* spin */ - } - } - __syncthreads(); - } - - // ======================================================================== - // Phase 4: inline bigFuse for this CTA's subset of BLOCK_M tokens. - // y_acc[tok, 0..HC_MULT) -> pre_mix (sigmoid) - // y_acc[tok, HC_MULT..2*HC_MULT) -> post_mix_out - // y_acc[tok, 2*HC_MULT..HC_MULT3) -> comb_mix_out (sinkhorn) - // layer_input[tok, h] = sum_j pre_mix[tok, j] * residual_cur[tok, j, h] - // - // Token subset: CTA k_split_idx handles tokens - // [k_split_idx * TOKS_PER_CTA, (k_split_idx + 1) * TOKS_PER_CTA), - // where TOKS_PER_CTA = ceil(BLOCK_M / kNumSplits). At kNumSplits == 1 - // this is the whole m_block (64 tokens); at kNumSplits == 8 it's 8 tokens - // spread over 8 CTAs running Phase 4 concurrently — ~8× the Phase-4 - // throughput of the old single-last-home-CTA design. - // - // Within a CTA, tokens are parallelized across warps: each warp handles - // max(1, TOKS_PER_CTA / NUM_WARPS_BF) tokens serially. Within a warp, - // lanes 0..HC_MULT-1 compute rmsnorm/sigmoid/sinkhorn; pre_mix is - // broadcast via __shfl_sync so all 32 lanes run the layer_input dot - // product across HIDDEN=4096. - // ======================================================================== - constexpr uint32_t BLOCK_SIZE_BF = kNumMMAThreads + kNumPmapThreads; // 256 - constexpr uint32_t WARP_SIZE_BF = 32; - constexpr uint32_t NUM_WARPS_BF = BLOCK_SIZE_BF / WARP_SIZE_BF; // 8 - constexpr uint32_t TOKS_PER_CTA = (BLOCK_M + kNumSplits - 1) / kNumSplits; - // When TOKS_PER_CTA < NUM_WARPS_BF (large kNumSplits / small BLOCK_M case), - // have WARPS_PER_TOK warps cooperate on the HIDDEN-stride layer_input loop - // so all 8 warps stay active. Example at BLOCK_M=64, kNumSplits=16: - // TOKS_PER_CTA=4, WARPS_PER_TOK=2, TOKS_PER_PASS=4, TOKEN_PASSES=1 — - // 2 warps per token each sweep HIDDEN/2, zero idle warps. - constexpr uint32_t WARPS_PER_TOK = (NUM_WARPS_BF > TOKS_PER_CTA) ? (NUM_WARPS_BF / TOKS_PER_CTA) : 1u; - constexpr uint32_t TOKS_PER_PASS = NUM_WARPS_BF / WARPS_PER_TOK; - constexpr uint32_t TOKEN_PASSES = (TOKS_PER_CTA + TOKS_PER_PASS - 1) / TOKS_PER_PASS; - constexpr uint32_t BF16_VEC_LI = 8; - // The vectorized loop covers floor(HIDDEN / H_STRIDE) * H_STRIDE elements. - // Anything past that is handled by a scalar-vec tail loop below — relax the - // old static_assert so HIDDEN values like 7168 (which 8-warp teams cannot - // cleanly span) are also valid. - static_assert(HIDDEN % BF16_VEC_LI == 0, "HIDDEN must be a multiple of BF16_VEC_LI=8"); - const uint32_t tid_bf = threadIdx.x; - const uint32_t lane_bf = tid_bf % WARP_SIZE_BF; - const uint32_t warp_bf = tid_bf / WARP_SIZE_BF; - const uint32_t warp_tok_pos = warp_bf / WARPS_PER_TOK; // which token in a pass - const uint32_t warp_in_team = warp_bf % WARPS_PER_TOK; // which warp inside team - const uint32_t cta_tok_base = k_split_idx * TOKS_PER_CTA; - -#pragma unroll 1 - for (uint32_t pass = 0; pass < TOKEN_PASSES; ++pass) - { - const uint32_t t_in_cta = pass * TOKS_PER_PASS + warp_tok_pos; - if (t_in_cta >= TOKS_PER_CTA) - continue; - const uint32_t t = cta_tok_base + t_in_cta; - if (t >= BLOCK_M) - continue; - const uint32_t tok = m_offset + t; - if (tok >= shape_m) - continue; - - // Lanes 0..HC_MULT-1 compute rmsnorm / sigmoid / sinkhorn; pre_mix is - // held in `pre_mix_local` on lanes 0..HC_MULT-1 and later broadcast to - // all 32 lanes via __shfl_sync. All warps in a team redundantly run - // these ~tens of FLOPs (cheap) to avoid a cross-warp SMEM sync; only - // warp_in_team==0 writes comb_mix_out / post_mix_out to GMEM. - float pre_mix_local = 0.f; - if (lane_bf < HC_MULT) - { - float const r_val = sqr_sum[tok]; - float y_local[HC_MULT3]; - float const* y_row = D + static_cast(tok) * SHAPE_N; -#pragma unroll - for (uint32_t c = 0; c < HC_MULT3; ++c) - y_local[c] = y_row[c]; - - float const rstd = rsqrtf(r_val / static_cast(HC_MULT * HIDDEN) + rms_eps); - float const s0 = hc_scale[0]; - float const s1 = hc_scale[1]; - float const s2 = hc_scale[2]; - - float v = y_local[lane_bf] * rstd * s0 + hc_base[lane_bf]; - pre_mix_local = 1.0f / (1.0f + __expf(-v)) + hc_pre_eps; - - v = y_local[HC_MULT + lane_bf] * rstd * s1 + hc_base[HC_MULT + lane_bf]; - float post_val = 1.0f / (1.0f + __expf(-v)) * hc_post_mult_value; - if (warp_in_team == 0) - { - post_mix_out[tok * HC_MULT + lane_bf] = post_val; - } - - float cm_vals[HC_MULT]; -#pragma unroll - for (uint32_t k = 0; k < HC_MULT; ++k) - cm_vals[k] = y_local[2 * HC_MULT + lane_bf * HC_MULT + k] * rstd * s2 - + hc_base[2 * HC_MULT + lane_bf * HC_MULT + k]; - - constexpr unsigned LANE_MASK = (1u << HC_MULT) - 1; - float const rowMax = fmaxf(fmaxf(cm_vals[0], cm_vals[1]), fmaxf(cm_vals[2], cm_vals[3])); -#pragma unroll - for (uint32_t k = 0; k < HC_MULT; ++k) - cm_vals[k] = __expf(cm_vals[k] - rowMax); - // Reciprocal-multiply for sinkhorn: 1 fdiv + 4 fmul instead of 4 - // fdivs per row-normalize. Equivalent under fp32 round-off. - float inv_rs = 1.0f / (cm_vals[0] + cm_vals[1] + cm_vals[2] + cm_vals[3]); -#pragma unroll - for (uint32_t k = 0; k < HC_MULT; ++k) - cm_vals[k] = cm_vals[k] * inv_rs + hc_sinkhorn_eps; -#pragma unroll - for (uint32_t k = 0; k < HC_MULT; ++k) - { - float cs = cm_vals[k]; - cs += __shfl_xor_sync(LANE_MASK, cs, 1); - cs += __shfl_xor_sync(LANE_MASK, cs, 2); - cm_vals[k] *= 1.0f / (cs + hc_sinkhorn_eps); - } - for (uint32_t it = 1; it < sinkhorn_repeat; ++it) - { - inv_rs = 1.0f / (cm_vals[0] + cm_vals[1] + cm_vals[2] + cm_vals[3] + hc_sinkhorn_eps); -#pragma unroll - for (uint32_t k = 0; k < HC_MULT; ++k) - cm_vals[k] *= inv_rs; -#pragma unroll - for (uint32_t k = 0; k < HC_MULT; ++k) - { - float cs = cm_vals[k]; - cs += __shfl_xor_sync(LANE_MASK, cs, 1); - cs += __shfl_xor_sync(LANE_MASK, cs, 2); - cm_vals[k] *= 1.0f / (cs + hc_sinkhorn_eps); - } - } - if (warp_in_team == 0) - { - float* cm_out_ptr = comb_mix_out + tok * HC_MULT2; -#pragma unroll - for (uint32_t k = 0; k < HC_MULT; ++k) - cm_out_ptr[lane_bf * HC_MULT + k] = cm_vals[k]; - } - } - - // Broadcast pre_mix[j] from lane j to all 32 lanes (intra-warp shfl). - float pm[HC_MULT]; -#pragma unroll - for (uint32_t j = 0; j < HC_MULT; ++j) - pm[j] = __shfl_sync(0xffffffff, pre_mix_local, j); - - // Layer_input[tok, h] = sum_j pm[j] * residual_cur[tok, j, h]. - // When WARPS_PER_TOK>1, warp_in_team 0..WARPS_PER_TOK-1 together cover - // HIDDEN in strides of WARPS_PER_TOK * 32 * 8. When WARPS_PER_TOK==1, - // each warp sweeps HIDDEN alone (same as the original single-warp case). - __nv_bfloat16 const* rbase = residual_cur_ptr + static_cast(tok) * HC_MULT * HIDDEN; - __nv_bfloat16* obase = layer_input_out + static_cast(tok) * HIDDEN; - - constexpr uint32_t H_STRIDE = WARPS_PER_TOK * WARP_SIZE_BF * BF16_VEC_LI; - // Largest multiple of H_STRIDE that fits in HIDDEN — after this comes - // the scalar-vec tail (only relevant when WARPS_PER_TOK*32*8 > HIDDEN - // residue, e.g. KS=112 / HIDDEN=7168 / H_STRIDE=2048 → tail = 1024). - constexpr uint32_t H_VEC_END = (HIDDEN / H_STRIDE) * H_STRIDE; - const uint32_t h_start = warp_in_team * WARP_SIZE_BF * BF16_VEC_LI + lane_bf * BF16_VEC_LI; - - if constexpr (!kFuseNorm) - { -#pragma unroll - for (uint32_t h = h_start; h < H_VEC_END; h += H_STRIDE) - { - // Issue all HC_MULT=4 residual_cur reads first so their L2 latency - // is hidden by the bf16→fp32 arithmetic that follows. The compiler - // schedules the 4 independent __ldg's in parallel. - uint4 raws[HC_MULT]; -#pragma unroll - for (uint32_t j = 0; j < HC_MULT; ++j) - { - raws[j] = __ldg(reinterpret_cast(&rbase[j * HIDDEN + h])); - } - float acc_li[BF16_VEC_LI] = {}; -#pragma unroll - for (uint32_t j = 0; j < HC_MULT; ++j) - { - __nv_bfloat162 const* pairs = reinterpret_cast<__nv_bfloat162 const*>(&raws[j]); -#pragma unroll - for (uint32_t v = 0; v < BF16_VEC_LI / 2; ++v) - { - float2 f = __bfloat1622float2(pairs[v]); - acc_li[2 * v + 0] += pm[j] * f.x; - acc_li[2 * v + 1] += pm[j] * f.y; - } - } - uint4 out_raw; - __nv_bfloat162* opairs = reinterpret_cast<__nv_bfloat162*>(&out_raw); -#pragma unroll - for (uint32_t v = 0; v < BF16_VEC_LI / 2; ++v) - opairs[v] = __float22bfloat162_rn(make_float2(acc_li[2 * v], acc_li[2 * v + 1])); - *reinterpret_cast(&obase[h]) = out_raw; - } - - // Scalar-vec tail: hidden residue [H_VEC_END, HIDDEN) is shorter than a - // full team stride. Distribute leftover BF16_VEC_LI-sized chunks across - // the team's threads in lane-major order (skip threads whose chunk - // falls past HIDDEN). Each chunk is still a single uint4 LDG/STG, so - // tail throughput matches the main loop's per-thread bandwidth — the - // only loss is that some lanes/warps in the team idle. - if constexpr (H_VEC_END < HIDDEN) - { - constexpr uint32_t TAIL_CHUNKS = (HIDDEN - H_VEC_END) / BF16_VEC_LI; - const uint32_t my_chunk = warp_in_team * WARP_SIZE_BF + lane_bf; - if (my_chunk < TAIL_CHUNKS) - { - const uint32_t h = H_VEC_END + my_chunk * BF16_VEC_LI; - uint4 raws[HC_MULT]; -#pragma unroll - for (uint32_t j = 0; j < HC_MULT; ++j) - { - raws[j] = __ldg(reinterpret_cast(&rbase[j * HIDDEN + h])); - } - float acc_li[BF16_VEC_LI] = {}; -#pragma unroll - for (uint32_t j = 0; j < HC_MULT; ++j) - { - __nv_bfloat162 const* pairs = reinterpret_cast<__nv_bfloat162 const*>(&raws[j]); -#pragma unroll - for (uint32_t v = 0; v < BF16_VEC_LI / 2; ++v) - { - float2 f = __bfloat1622float2(pairs[v]); - acc_li[2 * v + 0] += pm[j] * f.x; - acc_li[2 * v + 1] += pm[j] * f.y; - } - } - uint4 out_raw; - __nv_bfloat162* opairs = reinterpret_cast<__nv_bfloat162*>(&out_raw); -#pragma unroll - for (uint32_t v = 0; v < BF16_VEC_LI / 2; ++v) - opairs[v] = __float22bfloat162_rn(make_float2(acc_li[2 * v], acc_li[2 * v + 1])); - *reinterpret_cast(&obase[h]) = out_raw; - } - } - } - else - { - // ----------------------------------------------------------------- - // Fused RMSNorm path: layer_input[t,h] = bf16(li[t,h] * rsqrt( - // mean(li²)+norm_eps) * norm_weight[h]). - // - // Pass 1: identical to the un-fused path — compute li per chunk - // and STG to layer_input_out as bf16. The bf16 store lands in L2; - // pass 2 re-reads from L2 (no extra HBM read in the steady case). - // In parallel we accumulate per-thread `sum_sq_local`. - // Reduce: intra-warp __shfl_xor; cross-warp via SMEM + per-team - // named PTX barrier when WARPS_PER_TOK > 1 (KS ≥ 16 instances). - // Inactive teams `continue`'d above and never reach this barrier. - // Pass 2: re-LDG layer_input_out from L2, multiply by - // rsqrt * norm_weight, STG normalized bf16 back to the same - // address. Avoids the FMA recompute that doubling pass 1 would - // require (Path D Phase 4 is already FMA-heavy). - // - // Saves the separate RMSNorm kernel launch + its HBM round-trip - // (~14 KB read + ~14 KB write per token) vs the un-fused path. - // ----------------------------------------------------------------- - float sum_sq_local = 0.f; - -#pragma unroll - for (uint32_t h = h_start; h < H_VEC_END; h += H_STRIDE) - { - uint4 raws[HC_MULT]; -#pragma unroll - for (uint32_t j = 0; j < HC_MULT; ++j) - raws[j] = __ldg(reinterpret_cast(&rbase[j * HIDDEN + h])); - float acc_li[BF16_VEC_LI] = {}; -#pragma unroll - for (uint32_t j = 0; j < HC_MULT; ++j) - { - __nv_bfloat162 const* pairs = reinterpret_cast<__nv_bfloat162 const*>(&raws[j]); -#pragma unroll - for (uint32_t v = 0; v < BF16_VEC_LI / 2; ++v) - { - float2 f = __bfloat1622float2(pairs[v]); - acc_li[2 * v + 0] += pm[j] * f.x; - acc_li[2 * v + 1] += pm[j] * f.y; - } - } - // Round-to-bf16 *before* squaring so sum_sq matches the value - // we actually store (a separate RMSNorm kernel would also - // compute sum_sq from the bf16-rounded layer_input). - uint4 out_raw; - __nv_bfloat162* opairs = reinterpret_cast<__nv_bfloat162*>(&out_raw); -#pragma unroll - for (uint32_t v = 0; v < BF16_VEC_LI / 2; ++v) - opairs[v] = __float22bfloat162_rn(make_float2(acc_li[2 * v], acc_li[2 * v + 1])); - *reinterpret_cast(&obase[h]) = out_raw; -#pragma unroll - for (uint32_t v = 0; v < BF16_VEC_LI / 2; ++v) - { - float2 b = __bfloat1622float2(opairs[v]); - sum_sq_local += b.x * b.x + b.y * b.y; - } - } - if constexpr (H_VEC_END < HIDDEN) - { - constexpr uint32_t TAIL_CHUNKS = (HIDDEN - H_VEC_END) / BF16_VEC_LI; - const uint32_t my_chunk = warp_in_team * WARP_SIZE_BF + lane_bf; - if (my_chunk < TAIL_CHUNKS) - { - const uint32_t h = H_VEC_END + my_chunk * BF16_VEC_LI; - uint4 raws[HC_MULT]; -#pragma unroll - for (uint32_t j = 0; j < HC_MULT; ++j) - raws[j] = __ldg(reinterpret_cast(&rbase[j * HIDDEN + h])); - float acc_li[BF16_VEC_LI] = {}; -#pragma unroll - for (uint32_t j = 0; j < HC_MULT; ++j) - { - __nv_bfloat162 const* pairs = reinterpret_cast<__nv_bfloat162 const*>(&raws[j]); -#pragma unroll - for (uint32_t v = 0; v < BF16_VEC_LI / 2; ++v) - { - float2 f = __bfloat1622float2(pairs[v]); - acc_li[2 * v + 0] += pm[j] * f.x; - acc_li[2 * v + 1] += pm[j] * f.y; - } - } - uint4 out_raw; - __nv_bfloat162* opairs = reinterpret_cast<__nv_bfloat162*>(&out_raw); -#pragma unroll - for (uint32_t v = 0; v < BF16_VEC_LI / 2; ++v) - opairs[v] = __float22bfloat162_rn(make_float2(acc_li[2 * v], acc_li[2 * v + 1])); - *reinterpret_cast(&obase[h]) = out_raw; -#pragma unroll - for (uint32_t v = 0; v < BF16_VEC_LI / 2; ++v) - { - float2 b = __bfloat1622float2(opairs[v]); - sum_sq_local += b.x * b.x + b.y * b.y; - } - } - } - - // Intra-warp reduce sum_sq → one value broadcast to all 32 lanes. - sum_sq_local += __shfl_xor_sync(0xffffffff, sum_sq_local, 16); - sum_sq_local += __shfl_xor_sync(0xffffffff, sum_sq_local, 8); - sum_sq_local += __shfl_xor_sync(0xffffffff, sum_sq_local, 4); - sum_sq_local += __shfl_xor_sync(0xffffffff, sum_sq_local, 2); - sum_sq_local += __shfl_xor_sync(0xffffffff, sum_sq_local, 1); - - // Cross-warp reduce when WARPS_PER_TOK > 1. Use a per-team named - // PTX barrier (bar.sync %0, %1) instead of __syncthreads so that - // inactive teams (which `continue`'d above) don't deadlock. - // Barrier IDs 1..TOKS_PER_PASS are private to each team. - if constexpr (WARPS_PER_TOK > 1) - { - __shared__ float team_sumsq[TOKS_PER_PASS][WARPS_PER_TOK]; - if (lane_bf == 0) - team_sumsq[warp_tok_pos][warp_in_team] = sum_sq_local; - asm volatile("bar.sync %0, %1;" ::"r"(warp_tok_pos + 1u), - "n"(static_cast(WARPS_PER_TOK * WARP_SIZE_BF))); - float team_sum = 0.f; -#pragma unroll - for (uint32_t w = 0; w < WARPS_PER_TOK; ++w) - team_sum += team_sumsq[warp_tok_pos][w]; - sum_sq_local = team_sum; - } - - float const rsqrt_val = rsqrtf(sum_sq_local / static_cast(HIDDEN) + norm_eps); - - // Pass 2: re-LDG the un-normalized layer_input we just wrote - // (L2-hot), LDG norm_weight, normalize, STG back. No FMA recompute. - __nv_bfloat16 const* nbase = norm_weight; -#pragma unroll - for (uint32_t h = h_start; h < H_VEC_END; h += H_STRIDE) - { - uint4 li_raw = __ldg(reinterpret_cast(&obase[h])); - uint4 nw_raw = __ldg(reinterpret_cast(&nbase[h])); - __nv_bfloat162 const* li_pairs = reinterpret_cast<__nv_bfloat162 const*>(&li_raw); - __nv_bfloat162 const* nw_pairs = reinterpret_cast<__nv_bfloat162 const*>(&nw_raw); - uint4 out_raw; - __nv_bfloat162* opairs = reinterpret_cast<__nv_bfloat162*>(&out_raw); -#pragma unroll - for (uint32_t v = 0; v < BF16_VEC_LI / 2; ++v) - { - float2 li_f = __bfloat1622float2(li_pairs[v]); - float2 nw_f = __bfloat1622float2(nw_pairs[v]); - opairs[v] - = __float22bfloat162_rn(make_float2(li_f.x * rsqrt_val * nw_f.x, li_f.y * rsqrt_val * nw_f.y)); - } - *reinterpret_cast(&obase[h]) = out_raw; - } - if constexpr (H_VEC_END < HIDDEN) - { - constexpr uint32_t TAIL_CHUNKS = (HIDDEN - H_VEC_END) / BF16_VEC_LI; - const uint32_t my_chunk = warp_in_team * WARP_SIZE_BF + lane_bf; - if (my_chunk < TAIL_CHUNKS) - { - const uint32_t h = H_VEC_END + my_chunk * BF16_VEC_LI; - uint4 li_raw = __ldg(reinterpret_cast(&obase[h])); - uint4 nw_raw = __ldg(reinterpret_cast(&nbase[h])); - __nv_bfloat162 const* li_pairs = reinterpret_cast<__nv_bfloat162 const*>(&li_raw); - __nv_bfloat162 const* nw_pairs = reinterpret_cast<__nv_bfloat162 const*>(&nw_raw); - uint4 out_raw; - __nv_bfloat162* opairs = reinterpret_cast<__nv_bfloat162*>(&out_raw); -#pragma unroll - for (uint32_t v = 0; v < BF16_VEC_LI / 2; ++v) - { - float2 li_f = __bfloat1622float2(li_pairs[v]); - float2 nw_f = __bfloat1622float2(nw_pairs[v]); - opairs[v] = __float22bfloat162_rn( - make_float2(li_f.x * rsqrt_val * nw_f.x, li_f.y * rsqrt_val * nw_f.y)); - } - *reinterpret_cast(&obase[h]) = out_raw; - } - } - } - } -#else - if (blockIdx.x == 0 and threadIdx.x == 0) - DG_DEVICE_ASSERT(false and "This kernel only supports sm_100a"); -#endif -} - -} // namespace fused_mhc - -#pragma clang diagnostic pop diff --git a/cpp/tensorrt_llm/kernels/mhcKernels/mhcFusedHcKernel.cu b/cpp/tensorrt_llm/kernels/mhcKernels/mhcFusedHcKernel.cu deleted file mode 100644 index 234318b769f8..000000000000 --- a/cpp/tensorrt_llm/kernels/mhcKernels/mhcFusedHcKernel.cu +++ /dev/null @@ -1,865 +0,0 @@ -/* - * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. - * - * 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. - */ - -// Single-launch fused hyper-connection boundary op. -// -// This implementation wraps the SM100 tcgen05-based residual-out variant -// (fused_tf32_pmap_gemm_rout_atomic_impl) to produce residual_cur, D, and -// sqr_sum in a single kernel launch; the big-fuse postlogue kernel then -// consumes (D, sqr_sum, residual_cur) to emit (post_mix_cur, comb_mix_cur, -// layer_input_cur). Semantically identical to -// -// residual_cur = prev_mHC.post_mapping(x_prev, residual_prev, ...) -// post_mix_cur, comb_mix_cur, layer_input_cur = self.pre_mapping(residual_cur) -// -// but exposed as a single entry point. - -#ifdef TRTLLM_MHC_ENABLE_FUSED_HC -#include "fused_tf32_pmap_gemm.cuh" -#endif -#include "mhcKernels.h" -#include "mhc_fused_fma.cuh" - -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/cudaUtils.h" -#include "tensorrt_llm/common/envUtils.h" - -#include -#include -#include -#include -#include - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels::mhc -{ - -// ---- Single-launch workspace zero kernel ----------------------------------- -// -// Replaces 2-3 separate cudaMemsetAsync calls for the atomic accumulator -// workspaces (y_acc, r_acc, optional done_counter). Avoids per-memset launch -// latency that is visible at small M / high-frequency inference. -namespace -{ - -__global__ void fhcZeroWorkspacesKernel(float* __restrict__ y_acc, uint32_t y_elems, float* __restrict__ r_acc, - uint32_t r_elems, int* __restrict__ done_counter, uint32_t done_elems) -{ - uint32_t const tid = blockIdx.x * blockDim.x + threadIdx.x; - uint32_t const stride = gridDim.x * blockDim.x; - for (uint32_t i = tid; i < y_elems; i += stride) - { - y_acc[i] = 0.0f; - } - for (uint32_t i = tid; i < r_elems; i += stride) - { - r_acc[i] = 0.0f; - } - if (done_counter != nullptr) - { - for (uint32_t i = tid; i < done_elems; i += stride) - { - done_counter[i] = 0; - } - } -} - -inline void fhcZeroWorkspaces(float* y_acc, uint32_t y_elems, float* r_acc, uint32_t r_elems, int* done_counter, - uint32_t done_elems, cudaStream_t stream) -{ - uint32_t const total = y_elems + r_elems + (done_counter != nullptr ? done_elems : 0u); - if (total == 0u) - { - return; - } - constexpr uint32_t kBlock = 256; - // Cap grid so we don't over-launch for small workspaces; 148 SMs on B200. - uint32_t const num_blocks = min(static_cast((total + kBlock - 1) / kBlock), 148u * 8u); - fhcZeroWorkspacesKernel<<>>( - y_acc, y_elems, r_acc, r_elems, done_counter, done_elems); -} - -} // namespace - -// ---- mHC fused kernel shape constants (mirrors the Python module) ---- -// HC_MULT * (2 + HC_MULT) = 4 * 6 = 24. -static constexpr uint32_t FHC_SHAPE_N = 24; -static constexpr uint32_t FHC_HC_MULT = 4; -static constexpr uint32_t FHC_BLOCK_K = 64; - -#ifdef TRTLLM_MHC_ENABLE_FUSED_HC -static constexpr uint32_t FHC_HIDDEN_FLASH = 4096; -static constexpr uint32_t FHC_HIDDEN_PRO = 7168; -static constexpr uint32_t FHC_BLOCK_M = 64; -static constexpr uint32_t FHC_BLOCK_N = 32; -static constexpr uint32_t FHC_SWIZZLE_CD = 128; -// Rebalanced from N_B=12 / N_INPUT=2: SASS PC-sampling on M=4096 KS=2 showed -// ~25% of all stalls landed on a single NANOSLEEP.SYNCS (mbarrier wait) — the -// pmap warp was input-buffer-starved with only 2 TMA stages for 56 h_tiles. -// Trade 5 B-stages (-40 KiB) for +1 input stage (+40 KiB), keeping total SMEM -// at 226 KiB. N_B=7 still leaves >= HC_MULT slack for the MMA pipeline. -static constexpr uint32_t FHC_N_B_STAGES = 7; -static constexpr uint32_t FHC_N_INPUT_STG = 3; -static constexpr uint32_t FHC_NUM_MMA_TH = 128; -static constexpr uint32_t FHC_NUM_PMAP_TH = 128; - -template